wip

Signed-off-by: Diwank Tomer <[email protected]>

agents-api/agents_api/autogen/openapi_model.py

@@ -15,8 +15,11 @@
 
 CreateOrUpdateAgentRequest = UpdateAgentRequest
 CreateOrUpdateUserRequest = UpdateUserRequest
+CreateOrUpdateSessionRequest = CreateSessionRequest
+
 ChatMLRole = Entry.model_fields["role"].annotation
 
+
 def make_session(
     *,
     agents: list[UUID],
@@ -46,4 +49,4 @@ def make_session(
     return cls(
         **data,
         **participants,
-    )
+    )

agents-api/agents_api/models/agent/list_agents.py

@@ -34,9 +34,6 @@ def list_agents_query(
         offset: Number of agents to skip before starting to collect the result set.
         metadata_filter: Dictionary to filter agents based on metadata.
         client: Instance of CozoClient to execute the query.
-
-    Returns:
-        A pandas DataFrame containing the query results.
     """
     # Transforms the metadata_filter dictionary into a string representation for the datalog query.
     metadata_filter_str = ", ".join(

agents-api/agents_api/models/agent/update_agent.py

@@ -57,19 +57,6 @@ def update_agent_query(
         else [update_data["instructions"]]
     )
 
-    assertion_query = """
-        ?[developer_id, agent_id] :=
-            *agents {
-                developer_id,
-                agent_id,
-            },
-            developer_id = to_uuid($developer_id),
-            agent_id = to_uuid($agent_id),
-
-        # Assertion to ensure the agent exists before updating.
-        :assert some
-    """
-
     # Construct the agent update part of the query with dynamic columns and values based on `update_data`.
     # Agent update query
     agent_update_cols, agent_update_vals = cozo_process_mutate_data(
@@ -127,11 +114,9 @@ def update_agent_query(
     if len(default_settings) != 0:
         queries.insert(0, settings_update_query)
 
-    # Combine the assertion query with the update queries
     queries = [
         verify_developer_id_query(developer_id),
         verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id),
-        assertion_query,
         *queries,
     ]
 

agents-api/agents_api/models/session/create_or_update_session.py

@@ -0,0 +1,144 @@
+from beartype import beartype
+
+from uuid import UUID
+
+
+from ...autogen.openapi_model import (
+    CreateOrUpdateSessionRequest,
+    ResourceUpdatedResponse,
+)
+from ...common.utils.cozo import cozo_process_mutate_data
+from ..utils import (
+    cozo_query,
+    verify_developer_id_query,
+    verify_developer_owns_resource_query,
+    wrap_in_class,
+)
+
+
+@wrap_in_class(
+    ResourceUpdatedResponse,
+    one=True,
+    transform=lambda d: {
+        "id": d["session_id"],
+        "updated_at": d.pop("updated_at")[0],
+        "jobs": [],
+        **d,
+    },
+)
+@cozo_query(debug=True)
+@beartype
+def create_or_update_session_query(
+    *,
+    session_id: UUID,
+    developer_id: UUID,
+    create_or_update_session: CreateOrUpdateSessionRequest,
+) -> tuple[str, dict]:
+
+    create_or_update_session.metadata = create_or_update_session.metadata or {}
+    session_data = create_or_update_session.model_dump()
+
+    user = session_data.pop("user")
+    agent = session_data.pop("agent")
+    users = session_data.pop("users")
+    agents = session_data.pop("agents")
+
+    # Only one of agent or agents should be provided.
+    if agent and agents:
+        raise ValueError("Only one of 'agent' or 'agents' should be provided.")
+
+    agents = agents or ([agent] if agent else [])
+    assert len(agents) > 0, "At least one agent must be provided."
+
+    # Users are zero or more, so we default to an empty list if not provided.
+    if not (user or users):
+        users = []
+
+    else:
+        users = users or [user]
+
+    participants = [
+        *[("user", str(user)) for user in users],
+        *[("agent", str(agent)) for agent in agents],
+    ]
+
+    # Construct the datalog query for creating a new session and its lookup.
+    clear_lookup_query = """
+        input[session_id] <- [[$session_id]]
+        ?[session_id, participant_id, participant_type] :=
+            input[session_id],
+            *session_lookup {
+                session_id,
+                participant_type,
+                participant_id,
+            },
+
+        :delete session_lookup {
+            session_id,
+            participant_type,
+            participant_id,
+        }
+    """
+
+    lookup_query = """
+        # This section creates a new session lookup to ensure uniqueness and manage session metadata.
+        session[session_id] <- [[$session_id]]
+        participants[participant_type, participant_id] <- $participants
+        ?[session_id, participant_id, participant_type] :=
+            session[session_id],
+            participants[participant_type, participant_id],
+
+        :put session_lookup {
+            session_id,
+            participant_id,
+            participant_type,
+        }
+    """
+
+    session_update_cols, session_update_vals = cozo_process_mutate_data(
+        {k: v for k, v in session_data.items() if v is not None}
+    )
+
+    # Construct the datalog query for creating or updating session information.
+    update_query = f"""
+        input[{session_update_cols}] <- $session_update_vals
+        ids[session_id, developer_id] <- [[to_uuid($session_id), to_uuid($developer_id)]]
+        
+        ?[{session_update_cols}, session_id, developer_id] :=
+            input[{session_update_cols}],
+            ids[session_id, developer_id],
+
+        :put sessions {{
+           {session_update_cols}, session_id, developer_id 
+        }}
+
+        :returning
+    """
+
+    queries = [
+        verify_developer_id_query(developer_id),
+        *[
+            verify_developer_owns_resource_query(
+                developer_id,
+                f"{participant_type}s",
+                **{f"{participant_type}_id": participant_id},
+            )
+            for participant_type, participant_id in participants
+        ],
+        clear_lookup_query,
+        lookup_query,
+        update_query,
+    ]
+
+    query = "}\n\n{\n".join(queries)
+    query = f"{{ {query} }}"
+
+    return (
+        query,
+        {
+            "session_update_vals": session_update_vals,
+            "session_id": str(session_id),
+            "developer_id": str(developer_id),
+            "participants": participants,
+        },
+    )

agents-api/agents_api/models/session/delete_session.py

@@ -1,16 +1,33 @@
 """This module contains the implementation for deleting sessions from the 'cozodb' database using datalog queries."""
 
-from beartype import beartype
-
 from uuid import UUID
 
+from beartype import beartype
+
 
-from ..utils import cozo_query
+from ...autogen.openapi_model import ResourceDeletedResponse
+from ...common.utils.datetime import utcnow
+from ..utils import (
+    cozo_query,
+    verify_developer_id_query,
+    verify_developer_owns_resource_query,
+    wrap_in_class,
+)
 
 
+@wrap_in_class(
+    ResourceDeletedResponse,
+    one=True,
+    transform=lambda d: {
+        "id": UUID(d.pop("session_id")),
+        "deleted_at": utcnow(),
+        "jobs": [],
+    },
+)
 @cozo_query
 @beartype
 def delete_session_query(
+    *,
     developer_id: UUID,
     session_id: UUID,
 ) -> tuple[str, dict]:
@@ -22,39 +39,40 @@ def delete_session_query(
     - session_id (UUID): The unique identifier for the session to be deleted.
 
     Returns:
-    - pd.DataFrame: A DataFrame containing the result of the deletion query.
+    - ResourceDeletedResponse: The response indicating the deletion of the session.
     """
     session_id = str(session_id)
     developer_id = str(developer_id)
 
     # Constructs and executes a datalog query to delete the specified session and its associated data based on the session_id and developer_id.
-    query = """
-    {
+    delete_lookup_query = """
         # Convert session_id to UUID format
         input[session_id] <- [[
             to_uuid($session_id),
         ]]
 
         # Select sessions based on the session_id provided
         ?[
-            agent_id,
-            user_id,
             session_id,
+            participant_id,
+            participant_type,
         ] :=
             input[session_id],
             *session_lookup{
-                agent_id,
-                user_id,
                 session_id,
+                participant_id,
+                participant_type,
             }
 
         # Delete entries from session_lookup table matching the criteria
         :delete session_lookup {
-            agent_id,
-            user_id,
             session_id,
+            participant_id,
+            participant_type,
         }
-    } {
+    """
+
+    delete_query = """
         # Convert developer_id and session_id to UUID format
         input[developer_id, session_id] <- [[
             to_uuid($developer_id),
@@ -76,7 +94,19 @@ def delete_session_query(
             session_id,
             updated_at,
         }
-    }
+        :returning
     """
 
+    queries = [
+        verify_developer_id_query(developer_id),
+        verify_developer_owns_resource_query(
+            developer_id, "sessions", session_id=session_id
+        ),
+        delete_lookup_query,
+        delete_query,
+    ]
+
+    query = "}\n\n{\n".join(queries)
+    query = f"{{ {query} }}"
+
     return (query, {"session_id": session_id, "developer_id": developer_id})

agents-api/agents_api/models/session/get_session.py

@@ -47,6 +47,14 @@ def get_session_query(
             participant_type,
         }
 
+    # We have to do this dance because users can be zero or more
+    users_p[users] :=
+        participants[users, "user"]
+
+    users_p[users] :=
+        not participants[_, "user"],
+        users = []
+
     ?[
         agents,
         users,
@@ -60,7 +68,7 @@ def get_session_query(
         token_budget,
         context_overflow,
     ] := input[developer_id, id],
-        participants[users, "user"],
+        users_p[users],
         participants[agents, "agent"],
         *sessions{
             developer_id,