wip

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

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

@@ -3,23 +3,23 @@
 It includes functions to construct and execute datalog queries for inserting new agent records.
 """
 
-from uuid import UUID
+from uuid import UUID, uuid4
 
+from beartype import beartype
 
+from ...autogen.openapi_model import Agent, CreateAgentRequest
 from ...common.utils.cozo import cozo_process_mutate_data
-from ..utils import cozo_query
+from ..utils import cozo_query, verify_developer_id_query, wrap_in_class
 
 
+@wrap_in_class(Agent, one=True, transform=lambda d: {"id": UUID(d.pop("agent_id")), **d})
 @cozo_query
+@beartype
 def create_agent_query(
-    agent_id: UUID,
+    *,
     developer_id: UUID,
-    name: str,
-    about: str,
-    model: str,
-    instructions: list[str] = [],
-    metadata: dict = {},
-    default_settings: dict = {},
+    agent_id: UUID | None = None,
+    create_agent: CreateAgentRequest,
 ) -> tuple[str, dict]:
     """
     Constructs and executes a datalog query to create a new agent in the database.
@@ -36,11 +36,18 @@ def create_agent_query(
     - client (CozoClient, optional): The CozoDB client instance to use for the query. Defaults to a preconfigured client instance.
 
     Returns:
-    pd.DataFrame: A DataFrame containing the results of the query execution.
+    Agent: The newly created agent record.
     """
 
-    preset = default_settings.get("preset")
-    default_settings["preset"] = getattr(preset, "value", preset)
+    agent_id = agent_id or uuid4()
+
+    # Extract the agent data from the payload
+    create_agent.metadata = create_agent.metadata or {}
+    create_agent.instructions = create_agent.instructions if isinstance(create_agent.instructions, list) else [create_agent.instructions]
+    create_agent.default_settings = create_agent.default_settings or {}
+
+    agent_data = create_agent.model_dump()
+    default_settings = agent_data.pop("default_settings")
 
     settings_cols, settings_vals = cozo_process_mutate_data(
         {
@@ -61,8 +68,8 @@ def create_agent_query(
     # create the agent
     # Construct a query to insert the new agent record into the agents table
     agent_query = """
-        ?[agent_id, developer_id, model, name, about, metadata, instructions] <- [
-            [$agent_id, $developer_id, $model, $name, $about, $metadata, $instructions]
+        ?[agent_id, developer_id, model, name, about, metadata, instructions, created_at, updated_at] <- [
+            [$agent_id, $developer_id, $model, $name, $about, $metadata, $instructions, now(), now()]
         ]
 
         :insert agents {
@@ -73,11 +80,14 @@ def create_agent_query(
             about,
             metadata,
             instructions,
+            created_at,
+            updated_at,
         }
         :returning
     """
 
     queries = [
+        verify_developer_id_query(developer_id),
         default_settings_query,
         agent_query,
     ]
@@ -91,10 +101,6 @@ def create_agent_query(
             "settings_vals": settings_vals,
             "agent_id": str(agent_id),
             "developer_id": str(developer_id),
-            "model": model,
-            "name": name,
-            "about": about,
-            "metadata": metadata,
-            "instructions": instructions,
+            **agent_data,
         },
     )

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

@@ -2,65 +2,68 @@
 
 from uuid import UUID, uuid4
 
+from beartype import beartype
 
-from ...autogen.openapi_model import FunctionDef
+from ...autogen.openapi_model import CreateToolRequest, Tool
 
-from ..utils import cozo_query
+from ..utils import cozo_query, verify_developer_id_query, verify_developer_owns_resource_query, wrap_in_class
 
 
+@wrap_in_class(
+    Tool,
+    transform=lambda d: {"id": UUID(d.pop("tool_id")), d["tool_type"]: d.pop("spec"), "type": d.pop("tool_type"), **d},
+)
 @cozo_query
+@beartype
 def create_tools_query(
+    *,
+    developer_id: UUID,
     agent_id: UUID,
-    functions: list[FunctionDef],
-    embeddings: list[list[float]],
+    create_tools: list[CreateToolRequest],
 ) -> tuple[str, dict]:
     """
     Constructs a datalog query for inserting tool records into the 'agent_functions' relation in the CozoDB.
 
     Parameters:
     - agent_id (UUID): The unique identifier for the agent.
-    - functions (list[FunctionDef]): A list of function definitions to be inserted.
-    - embeddings (list[list[float]]): A list of embeddings corresponding to each function.
+    - create_tools (list[CreateToolRequest]): A list of function definitions to be inserted.
 
     Returns:
-    - pd.DataFrame: A DataFrame containing the results of the query execution.
+    list[Tool]
     """
-    # Ensure the number of functions matches the number of embeddings
-    assert len(functions) == len(embeddings)
 
-    # Construct the input records for the datalog query
-    functions_input: list[list] = []
-
-    for function, embedding in zip(functions, embeddings):
-        parameters = function.parameters.model_dump()
-        functions_input.append(
-            [
-                str(agent_id),
-                str(uuid4()),
-                function.name,
-                function.description or "",
-                parameters,
-                embedding,
-            ]
-        )
+    tools_data = [
+        [
+            str(agent_id),
+            str(uuid4()),
+            tool.type,
+            tool.name,
+            getattr(tool, tool.type).dict(),
+        ]
+        for tool in create_tools
+    ]
 
     # Datalog query for inserting new tool records into the 'agent_functions' relation
-    query = """
-        input[agent_id, tool_id, name, description, parameters, embedding] <- $records
-        ?[agent_id, tool_id, name, description, parameters, embedding, updated_at] :=
-            input[agent_id, tool_id, name, description, parameters, embedding],
-            updated_at = now(),
+    create_tools_query = """
+        ?[agent_id, tool_id, tool_type, name, spec] <- $records
 
-        :insert agent_functions {
+        :insert tools {
             agent_id,
             tool_id,
+            tool_type,
             name,
-            description,
-            parameters,
-            embedding,
-            updated_at,
+            spec,
         }
         :returning
     """
 
-    return (query, {"records": functions_input})
+    queries = [
+        verify_developer_id_query(developer_id),
+        verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id),
+        create_tools_query,
+    ]
+
+    query = "}\n\n{\n".join(queries)
+    query = f"{{ {query} }}"
+
+    return (query, {"records": tools_data})

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

@@ -4,40 +4,91 @@
 
 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,
+)
 
-"""
-Constructs and returns a datalog query for deleting an agent and its default settings from the database.
 
-Parameters:
-- developer_id (UUID): The UUID of the developer owning the agent.
-- agent_id (UUID): The UUID of the agent to be deleted.
-- client (CozoClient, optional): An instance of the CozoClient to execute the query.
+@wrap_in_class(
+    ResourceDeletedResponse,
+    one=True,
+    transform=lambda d: {"id": UUID(d.pop("agent_id")), "deleted_at": utcnow(), "jobs": []},
+)
+@cozo_query
+@beartype
+def delete_agent_query(*, developer_id: UUID, agent_id: UUID) -> tuple[str, dict]:
+    """
+    Constructs and returns a datalog query for deleting an agent and its default settings from the database.
 
-Returns:
-- tuple[str, dict]: A DataFrame containing the results of the deletion query.
-"""
+    Parameters:
+    - developer_id (UUID): The UUID of the developer owning the agent.
+    - agent_id (UUID): The UUID of the agent to be deleted.
+    - client (CozoClient, optional): An instance of the CozoClient to execute the query.
 
+    Returns:
+    - ResourceDeletedResponse: The response indicating the deletion of the agent.
+    """
 
-@cozo_query
-def delete_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]:
-    query = """
-    {
+    queries = [
+        verify_developer_id_query(developer_id),
+        verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id),
+        """
+        # Delete docs
+        ?[agent_id, doc_id] :=
+            *agent_docs{
+                agent_id,
+                doc_id,
+            }, agent_id = to_uuid($agent_id)
+
+        :delete agent_docs {
+            agent_id,
+            doc_id
+        }
+        :returning
+        """,
+        """
+        # Delete tools
+        ?[agent_id, tool_id] :=
+            *tools{
+                agent_id,
+                tool_id,
+            }, agent_id = to_uuid($agent_id)
+
+        :delete tools {
+            agent_id,
+            tool_id
+        }
+        :returning
+        """,
+        """
         # Delete default agent settings
         ?[agent_id] <- [[$agent_id]]
 
         :delete agent_default_settings {
             agent_id
         }
-    } {
+        :returning
+        """,
+        """
         # Delete the agent
         ?[agent_id, developer_id] <- [[$agent_id, $developer_id]]
 
         :delete agents {
             developer_id,
             agent_id
         }
-    }"""
+        :returning
+        """,
+    ]
+
+    query = "}\n\n{\n".join(queries)
+    query = f"{{ {query} }}"
 
     return (query, {"agent_id": str(agent_id), "developer_id": str(developer_id)})

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

@@ -1,11 +1,21 @@
 from uuid import UUID
 
+from beartype import beartype
 
-from ..utils import cozo_query
 
+from ...autogen.openapi_model import Agent
+from ..utils import (
+    cozo_query,
+    verify_developer_id_query,
+    verify_developer_owns_resource_query,
+    wrap_in_class,
+)
 
+
+@wrap_in_class(Agent, one=True)
 @cozo_query
-def get_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]:
+@beartype
+def get_agent_query(*, developer_id: UUID, agent_id: UUID) -> tuple[str, dict]:
     """
     Fetches agent details and default settings from the database.
 
@@ -17,13 +27,12 @@ def get_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]:
     - client (CozoClient, optional): The database client used to execute the query.
 
     Returns:
-    - pd.DataFrame: A DataFrame containing the agent details and default settings.
+    - Agent
     """
     # Constructing a datalog query to retrieve agent details and default settings.
     # The query uses input parameters for agent_id and developer_id to filter the results.
     # It joins the 'agents' and 'agent_default_settings' relations to fetch comprehensive details.
-    query = """
-    {
+    get_query = """
         input[agent_id, developer_id] <- [[to_uuid($agent_id), to_uuid($developer_id)]]
 
         ?[
@@ -69,9 +78,17 @@ def get_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]:
                 "min_p": min_p,
                 "preset": preset,
             }
-    }
     """
 
+    queries = [
+        verify_developer_id_query(developer_id),
+        verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id),
+        get_query,
+    ]
+
+    query = "}\n\n{\n".join(queries)
+    query = f"{{ {query} }}"
+
     # Execute the constructed datalog query using the provided CozoClient.
     # The result is returned as a pandas DataFrame.
     return (query, {"agent_id": str(agent_id), "developer_id": str(developer_id)})