diff --git a/src/utils/db/models/session.py b/src/utils/db/models/session.py new file mode 100644 index 0000000..88fc85c --- /dev/null +++ b/src/utils/db/models/session.py @@ -0,0 +1,21 @@ +from typing import Annotated, List + +from beanie import Document, Indexed +from google.genai.types import Content +from pydantic import BaseModel, Field + + +class SessionBase(BaseModel): + chat_id: Annotated[int, Indexed(unique=True)] + system_prompt: str + history: List[Content] = Field(default_factory=list) + + +class ReviewSession(SessionBase, Document): + class Settings: + name = "review_sessions" + + +class RespondSession(SessionBase, Document): + class Settings: + name = "respond_sessions"