ic

packages/mongo-memory/banortegpt/database/mongo_memory/crud.py

@@ -0,0 +1,35 @@
+from uuid import UUID
+
+from .models import Conversation, Message
+
+
+async def create_conversation(
+    conversation_id: UUID, system_prompt: str
+) -> Conversation:
+    conversation = Conversation(
+        conversation_id=conversation_id,
+        messages=[Message(role="system", content=system_prompt)],
+    )
+    await conversation.create()
+
+    return conversation
+
+
+async def get_conversation(conversation_id: UUID) -> Conversation | None:
+    conversation = await Conversation.find_one(
+        Conversation.conversation_id == conversation_id
+    )
+    return conversation
+
+
+async def get_or_create_conversation(
+    conversation_id: UUID, system_prompt: str
+) -> Conversation:
+    conversation = await get_conversation(conversation_id)
+
+    if not conversation:
+        conversation = await create_conversation(conversation_id, system_prompt)
+    else:
+        conversation.messages[0].content = system_prompt
+
+    return conversation

packages/mongo-memory/banortegpt/database/mongo_memory/models.py

@@ -0,0 +1,67 @@
+from typing import Annotated, Any, Literal
+from uuid import UUID
+
+from beanie import Document, Indexed
+from pydantic import BaseModel
+
+type Role = Literal["user", "assistant", "system", "tool"]
+
+
+class Function(BaseModel):
+    name: str
+    arguments: str
+
+
+class Tool(BaseModel):
+    id: str
+    type: Literal["function"]
+    function: Function
+
+
+class Message(BaseModel):
+    role: Role
+    content: str | None = None
+    tool_call_id: str | None = None
+    tool_calls: list[Tool] | None = None
+
+
+class Conversation(Document):
+    conversation_id: Annotated[UUID, Indexed(unique=True)]
+    messages: list[Message]
+
+    class Settings:
+        name = "conversations"
+
+    def to_openai_format(self, limit: int, langchain_compat: bool = False):
+        history = [m.model_dump(exclude_none=True) for m in self.messages]
+
+        if langchain_compat:
+            for msg in history:
+                if "tool_calls" in msg:
+                    msg["additional_kwargs"] = {"tool_calls": msg.pop("tool_calls")}
+                if "content" not in msg:
+                    msg["content"] = ""
+
+        if len(history) > limit:
+            # Truncate the history to the last `limit` messages
+            # Always keep the first message (system prompt)
+            history = history[:1] + history[-limit:]
+
+        return history
+
+    def add(
+        self,
+        role: Role,
+        *,
+        content: str | None = None,
+        tool_call_id: str | None = None,
+        tool_calls: list[Any] | None = None,
+    ):
+        self.messages.append(
+            Message(
+                role=role,
+                content=content,
+                tool_call_id=tool_call_id,
+                tool_calls=tool_calls,
+            )
+        )