Pavlo Paliychuk bca838f61d
fix: Make groupIds option in search query dto (#154)
* fix: Make groupIds option in search query dto

* chore: Format
2024-09-24 19:05:12 -04:00

46 lines
1.3 KiB
Python

from datetime import datetime, timezone
from pydantic import BaseModel, Field
from graph_service.dto.common import Message
class SearchQuery(BaseModel):
group_ids: list[str] | None = Field(
None, description='The group ids for the memories to search'
)
query: str
max_facts: int = Field(default=10, description='The maximum number of facts to retrieve')
class FactResult(BaseModel):
uuid: str
name: str
fact: str
valid_at: datetime | None
invalid_at: datetime | None
created_at: datetime
expired_at: datetime | None
class Config:
json_encoders = {datetime: lambda v: v.astimezone(timezone.utc).isoformat()}
class SearchResults(BaseModel):
facts: list[FactResult]
class GetMemoryRequest(BaseModel):
group_id: str = Field(..., description='The group id of the memory to get')
max_facts: int = Field(default=10, description='The maximum number of facts to retrieve')
center_node_uuid: str | None = Field(
..., description='The uuid of the node to center the retrieval on'
)
messages: list[Message] = Field(
..., description='The messages to build the retrieval query from '
)
class GetMemoryResponse(BaseModel):
facts: list[FactResult] = Field(..., description='The facts that were retrieved from the graph')