Chat queries

agents-api/agents_api/exceptions.py

@@ -49,3 +49,12 @@ class FailedEncodingSentinel:
     """Sentinel object returned when failed to encode payload."""
 
     payload_data: bytes
+
+
+class QueriesBaseException(AgentsBaseException):
+    pass
+
+
+class InvalidSQLQuery(QueriesBaseException):
+    def __init__(self, query_name: str):
+        super().__init__(f"invalid query: {query_name}")

agents-api/agents_api/models/chat/__init__.py → agents-api/agents_api/queries/chat/__init__.py

@@ -17,6 +17,4 @@
 # ruff: noqa: F401, F403, F405
 
 from .gather_messages import gather_messages
-from .get_cached_response import get_cached_response
 from .prepare_chat_context import prepare_chat_context
-from .set_cached_response import set_cached_response

agents-api/agents_api/models/chat/gather_messages.py → agents-api/agents_api/queries/chat/gather_messages.py

@@ -3,7 +3,6 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import ChatInput, DocReference, History
@@ -13,8 +12,8 @@
 from ..docs.search_docs_by_embedding import search_docs_by_embedding
 from ..docs.search_docs_by_text import search_docs_by_text
 from ..docs.search_docs_hybrid import search_docs_hybrid
-from ..entry.get_history import get_history
-from ..session.get_session import get_session
+from ..entries.get_history import get_history
+from ..sessions.get_session import get_session
 from ..utils import (
     partialclass,
     rewrap_exceptions,
@@ -25,7 +24,6 @@
 
 @rewrap_exceptions(
     {
-        QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),
     }

agents-api/agents_api/queries/chat/prepare_chat_context.py

@@ -0,0 +1,169 @@
+from typing import Any, TypeVar
+from uuid import UUID
+
+import sqlvalidator
+from beartype import beartype
+
+from ...common.protocol.sessions import ChatContext, make_session
+from ...exceptions import InvalidSQLQuery
+from ..utils import (
+    pg_query,
+    wrap_in_class,
+)
+
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
+
+sql_query = sqlvalidator.parse(
+    """SELECT * FROM 
+(
+    SELECT jsonb_agg(u) AS users FROM (
+        SELECT 
+            session_lookup.participant_id,
+            users.user_id AS id,
+            users.developer_id,
+            users.name,
+            users.about,
+            users.created_at,
+            users.updated_at,
+            users.metadata
+        FROM session_lookup
+        INNER JOIN users ON session_lookup.participant_id = users.user_id
+        WHERE
+            session_lookup.developer_id = $1 AND
+            session_id = $2 AND
+            session_lookup.participant_type = 'user'
+    ) u
+) AS users,
+(
+    SELECT jsonb_agg(a) AS agents FROM (
+        SELECT 
+            session_lookup.participant_id,
+            agents.agent_id AS id,
+            agents.developer_id,
+            agents.canonical_name,
+            agents.name,
+            agents.about,
+            agents.instructions,
+            agents.model,
+            agents.created_at,
+            agents.updated_at,
+            agents.metadata,
+            agents.default_settings
+        FROM session_lookup
+        INNER JOIN agents ON session_lookup.participant_id = agents.agent_id
+        WHERE
+            session_lookup.developer_id = $1 AND
+            session_id = $2 AND
+            session_lookup.participant_type = 'agent'
+    ) a
+) AS agents,
+(
+    SELECT to_jsonb(s) AS session FROM (
+        SELECT 
+            sessions.session_id AS id,
+            sessions.developer_id,
+            sessions.situation,
+            sessions.system_template,
+            sessions.created_at,
+            sessions.metadata,
+            sessions.render_templates,
+            sessions.token_budget,
+            sessions.context_overflow,
+            sessions.forward_tool_calls,
+            sessions.recall_options
+        FROM sessions
+        WHERE
+            developer_id = $1 AND 
+            session_id = $2
+        LIMIT 1
+    ) s
+) AS session,
+(
+    SELECT jsonb_agg(r) AS toolsets FROM (
+        SELECT 
+            session_lookup.participant_id, 
+            tools.tool_id as id,
+            tools.developer_id,
+            tools.agent_id,
+            tools.task_id,
+            tools.task_version,
+            tools.type,
+            tools.name,
+            tools.description,
+            tools.spec,
+            tools.updated_at,
+            tools.created_at
+        FROM session_lookup
+        INNER JOIN tools ON session_lookup.participant_id = tools.agent_id
+        WHERE
+            session_lookup.developer_id = $1 AND 
+            session_id = $2 AND 
+            session_lookup.participant_type = 'agent'
+    ) r
+) AS toolsets"""
+)
+if not sql_query.is_valid():
+    raise InvalidSQLQuery("prepare_chat_context")
+
+
+def _transform(d):
+    toolsets = {}
+    for tool in d["toolsets"]:
+        agent_id = tool["agent_id"]
+        if agent_id in toolsets:
+            toolsets[agent_id].append(tool)
+        else:
+            toolsets[agent_id] = [tool]
+
+    return {
+        **d,
+        "session": make_session(
+            agents=[a["id"] for a in d["agents"]],
+            users=[u["id"] for u in d["users"]],
+            **d["session"],
+        ),
+        "toolsets": [
+            {
+                "agent_id": agent_id,
+                "tools": [
+                    {
+                        tool["type"]: tool.pop("spec"),
+                        **tool,
+                    }
+                    for tool in tools
+                ],
+            }
+            for agent_id, tools in toolsets.items()
+        ],
+    }
+
+
+# TODO: implement this part
+# @rewrap_exceptions(
+#     {
+#         ValidationError: partialclass(HTTPException, status_code=400),
+#         TypeError: partialclass(HTTPException, status_code=400),
+#     }
+# )
+@wrap_in_class(
+    ChatContext,
+    one=True,
+    transform=_transform,
+)
+@pg_query
+@beartype
+async def prepare_chat_context(
+    *,
+    developer_id: UUID,
+    session_id: UUID,
+) -> tuple[list[str], list]:
+    """
+    Executes a complex query to retrieve memory context based on session ID.
+    """
+
+    return (
+        [sql_query.format()],
+        [developer_id, session_id],
+    )