refactor(agents-api): Add decorator to wrap cozo queries inside and e…

…xecute (#273)

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

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

@@ -5,32 +5,12 @@
 
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
 from ...common.utils.cozo import cozo_process_mutate_data
+from ..utils import cozo_query
 
 
-"""
-Constructs and executes a datalog query to create a new agent in the database.
-
-Parameters:
-- agent_id (UUID): The unique identifier for the agent.
-- developer_id (UUID): The unique identifier for the developer creating the agent.
-- name (str): The name of the agent.
-- about (str): A description of the agent.
-- instructions (list[str], optional): A list of instructions for using the agent. Defaults to an empty list.
-- model (str, optional): The model identifier for the agent. Defaults to "julep-ai/samantha-1-turbo".
-- metadata (dict, optional): A dictionary of metadata for the agent. Defaults to an empty dict.
-- default_settings (dict, optional): A dictionary of default settings for the agent. Defaults to an empty dict.
-- 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.
-"""
-
-
+@cozo_query
 def create_agent_query(
     agent_id: UUID,
     developer_id: UUID,
@@ -40,8 +20,25 @@ def create_agent_query(
     instructions: list[str] = [],
     metadata: dict = {},
     default_settings: dict = {},
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
+    """
+    Constructs and executes a datalog query to create a new agent in the database.
+
+    Parameters:
+    - agent_id (UUID): The unique identifier for the agent.
+    - developer_id (UUID): The unique identifier for the developer creating the agent.
+    - name (str): The name of the agent.
+    - about (str): A description of the agent.
+    - instructions (list[str], optional): A list of instructions for using the agent. Defaults to an empty list.
+    - model (str, optional): The model identifier for the agent. Defaults to "julep-ai/samantha-1-turbo".
+    - metadata (dict, optional): A dictionary of metadata for the agent. Defaults to an empty dict.
+    - default_settings (dict, optional): A dictionary of default settings for the agent. Defaults to an empty dict.
+    - 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.
+    """
+
     settings_cols, settings_vals = cozo_process_mutate_data(
         {
             **default_settings,
@@ -85,7 +82,7 @@ def create_agent_query(
     query = "}\n\n{\n".join(queries)
     query = f"{{ {query} }}"
 
-    return client.run(
+    return (
         query,
         {
             "settings_vals": settings_vals,

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

@@ -2,17 +2,18 @@
 
 from uuid import UUID, uuid4
 
-import pandas as pd
 
 from ...autogen.openapi_model import FunctionDef
-from ...clients.cozo import client
 
+from ..utils import cozo_query
 
+
+@cozo_query
 def create_tools_query(
     agent_id: UUID,
     functions: list[FunctionDef],
     embeddings: list[list[float]],
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """
     Constructs a datalog query for inserting tool records into the 'agent_functions' relation in the CozoDB.
 
@@ -62,4 +63,4 @@ def create_tools_query(
         :returning
     """
 
-    return client.run(query, {"records": functions_input})
+    return (query, {"records": functions_input})

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

@@ -4,11 +4,8 @@
 
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
-
-from ...clients.cozo import client
 
+from ..utils import cozo_query
 
 """
 Constructs and returns a datalog query for deleting an agent and its default settings from the database.
@@ -19,13 +16,12 @@
 - client (CozoClient, optional): An instance of the CozoClient to execute the query.
 
 Returns:
-- pd.DataFrame: A DataFrame containing the results of the deletion query.
+- tuple[str, dict]: A DataFrame containing the results of the deletion query.
 """
 
 
-def delete_agent_query(
-    developer_id: UUID, agent_id: UUID, client: CozoClient = client
-) -> pd.DataFrame:
+@cozo_query
+def delete_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]:
     query = """
     {
         # Delete default agent settings
@@ -44,6 +40,4 @@ def delete_agent_query(
         }
     }"""
 
-    return client.run(
-        query, {"agent_id": str(agent_id), "developer_id": str(developer_id)}
-    )
+    return (query, {"agent_id": str(agent_id), "developer_id": str(developer_id)})

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

@@ -1,14 +1,11 @@
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
+from ..utils import cozo_query
 
 
-def get_agent_query(
-    developer_id: UUID, agent_id: UUID, client: CozoClient = client
-) -> pd.DataFrame:
+@cozo_query
+def get_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]:
     """
     Fetches agent details and default settings from the database.
 
@@ -77,6 +74,4 @@ def get_agent_query(
 
     # Execute the constructed datalog query using the provided CozoClient.
     # The result is returned as a pandas DataFrame.
-    return client.run(
-        query, {"agent_id": str(agent_id), "developer_id": str(developer_id)}
-    )
+    return (query, {"agent_id": str(agent_id), "developer_id": str(developer_id)})

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

@@ -1,20 +1,18 @@
 from typing import Any
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
 from ...common.utils import json
+from ..utils import cozo_query
 
 
+@cozo_query
 def list_agents_query(
     developer_id: UUID,
     limit: int = 100,
     offset: int = 0,
     metadata_filter: dict[str, Any] = {},
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """
     Constructs and executes a datalog query to list agents from the 'cozodb' database.
 
@@ -69,6 +67,7 @@ def list_agents_query(
         :sort -created_at
     }}"""
 
-    return client.run(
-        query, {"developer_id": str(developer_id), "limit": limit, "offset": offset}
+    return (
+        query,
+        {"developer_id": str(developer_id), "limit": limit, "offset": offset},
     )

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

@@ -1,18 +1,18 @@
 from uuid import UUID
 
-import pandas as pd
 
-from ...clients.cozo import client
 from ...common.utils.cozo import cozo_process_mutate_data
+from ..utils import cozo_query
 from ...common.utils.datetime import utcnow
 
 
+@cozo_query
 def patch_agent_query(
     agent_id: UUID,
     developer_id: UUID,
     default_settings: dict = {},
     **update_data,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """Patches agent data based on provided updates.
 
     Parameters:
@@ -85,7 +85,7 @@ def patch_agent_query(
 
     combined_query = "\n".join(queries)
 
-    return client.run(
+    return (
         combined_query,
         {
             "agent_update_vals": agent_update_vals,

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

@@ -1,14 +1,15 @@
 from uuid import UUID
 
-import pandas as pd
 
 from ...autogen.openapi_model import FunctionDef
-from ...clients.cozo import client
 
+from ..utils import cozo_query
 
+
+@cozo_query
 def patch_tool_by_id_query(
     agent_id: UUID, tool_id: UUID, function: FunctionDef, embedding: list[float]
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """
     # Execute the datalog query and return the results as a DataFrame
     Updates the tool information for a given agent and tool ID in the 'cozodb' database.
@@ -59,4 +60,4 @@ def patch_tool_by_id_query(
         :returning
     """
 
-    return client.run(query, params)
+    return (query, params)

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

@@ -5,35 +5,32 @@
 
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
 from ...common.utils.cozo import cozo_process_mutate_data
+from ..utils import cozo_query
 
 
-"""
-Constructs and executes a datalog query to update an agent and its default settings in the 'cozodb' database.
-
-Parameters:
-- agent_id (UUID): The unique identifier of the agent to be updated.
-- developer_id (UUID): The unique identifier of the developer associated with the agent.
-- default_settings (dict, optional): A dictionary of default settings to be updated for the agent. Defaults to an empty dict.
-- client (CozoClient, optional): The database client used to execute the query. Defaults to a pre-configured client instance.
-- **update_data: Variable keyword arguments representing additional agent data to be updated.
-
-Returns:
-- pd.DataFrame: A DataFrame containing the result of the update operation.
-"""
-
-
+@cozo_query
 def update_agent_query(
     agent_id: UUID,
     developer_id: UUID,
     default_settings: dict = {},
-    client: CozoClient = client,
     **update_data,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
+    """
+    Constructs and executes a datalog query to update an agent and its default settings in the 'cozodb' database.
+
+    Parameters:
+    - agent_id (UUID): The unique identifier of the agent to be updated.
+    - developer_id (UUID): The unique identifier of the developer associated with the agent.
+    - default_settings (dict, optional): A dictionary of default settings to be updated for the agent. Defaults to an empty dict.
+    - client (CozoClient, optional): The database client used to execute the query. Defaults to a pre-configured client instance.
+    - **update_data: Variable keyword arguments representing additional agent data to be updated.
+
+    Returns:
+    - pd.DataFrame: A DataFrame containing the result of the update operation.
+    """
+
     agent_id = str(agent_id)
     developer_id = str(developer_id)
     update_data["instructions"] = update_data.get("instructions", [])
@@ -115,7 +112,7 @@ def update_agent_query(
     # Combine the assertion query with the update queries
     combined_query = "{" + assertion_query + "} " + "\n".join(queries)
 
-    return client.run(
+    return (
         combined_query,
         {
             "agent_update_vals": agent_update_vals,

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

@@ -1,14 +1,15 @@
 from uuid import UUID
 
-import pandas as pd
 
 from ...autogen.openapi_model import FunctionDef
-from ...clients.cozo import client
 
+from ..utils import cozo_query
 
+
+@cozo_query
 def update_tool_by_id_query(
     agent_id: UUID, tool_id: UUID, function: FunctionDef, embedding: list[float]
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     # Agent update query
     function = function.model_dump()
     params = {
@@ -42,4 +43,4 @@ def update_tool_by_id_query(
         :returning
     """
 
-    return client.run(query, params)
+    return (query, params)

agents-api/agents_api/models/docs/create_docs.py

@@ -1,14 +1,13 @@
 from typing import Callable, Literal
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
 from ...common.utils.cozo import cozo_process_mutate_data
+from ..utils import cozo_query
 from ...common.utils.datetime import utcnow
 
 
+@cozo_query
 def create_docs_query(
     owner_type: Literal["user", "agent"],
     owner_id: UUID,
@@ -17,8 +16,7 @@ def create_docs_query(
     content: str,
     split_fn: Callable[[str], list[str]] = lambda x: x.split("\n\n"),
     metadata: dict = {},
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """
     Constructs and executes a datalog query to create a new document and its associated snippets in the 'cozodb' database.
 
@@ -30,7 +28,6 @@ def create_docs_query(
     - content (str): The content of the document, which will be split into snippets.
     - split_fn (Callable[[str], list[str]]): A function to split the content into snippets. Defaults to splitting by double newlines.
     - metadata (dict): Metadata associated with the document. Defaults to an empty dictionary.
-    - client (CozoClient): The Cozo client instance to execute the query. Defaults to a pre-configured client instance.
 
     Returns:
     pd.DataFrame: A DataFrame containing the results of the query execution.
@@ -86,7 +83,7 @@ def create_docs_query(
     }}"""
 
     # Execute the constructed datalog query and return the results as a DataFrame.
-    return client.run(
+    return (
         query,
         {
             "owner_id": str(owner_id),