From 86414f90f6cf5355c1ed58a3a25538cfc765d677 Mon Sep 17 00:00:00 2001
From: Diwank Singh Tomer <diwank.singh@gmail.com>
Date: Thu, 18 Apr 2024 11:33:10 +0530
Subject: [PATCH] refactor(agents-api): Add decorator to wrap cozo queries
 inside and execute

Signed-off-by: Diwank Singh Tomer <diwank.singh@gmail.com>
---
 .../agents_api/models/agent/create_agent.py   | 47 ++++++++---------
 .../agents_api/models/agent/create_tools.py   |  9 ++--
 .../agents_api/models/agent/delete_agent.py   | 16 ++----
 .../agents_api/models/agent/get_agent.py      | 13 ++---
 .../agents_api/models/agent/list_agents.py    | 13 +++--
 .../agents_api/models/agent/patch_agent.py    |  8 +--
 .../agents_api/models/agent/patch_tool.py     |  9 ++--
 .../agents_api/models/agent/update_agent.py   | 39 +++++++-------
 .../agents_api/models/agent/update_tool.py    |  9 ++--
 .../agents_api/models/docs/create_docs.py     | 11 ++--
 .../agents_api/models/docs/delete_docs.py     | 10 ++--
 .../agents_api/models/docs/embed_docs.py      | 34 ++++++-------
 agents-api/agents_api/models/docs/get_docs.py | 35 ++++++-------
 .../agents_api/models/docs/list_docs.py       | 16 +++---
 .../agents_api/models/docs/search_docs.py     | 40 +++++++--------
 .../agents_api/models/entry/add_entries.py    | 13 ++---
 .../agents_api/models/entry/delete_entries.py |  8 +--
 .../models/entry/entries_summarization.py     | 51 +++++++++----------
 .../agents_api/models/entry/get_entries.py    | 16 ++----
 .../models/entry/proc_mem_context.py          | 10 ++--
 .../models/session/create_session.py          | 42 +++++++--------
 .../models/session/delete_session.py          | 11 ++--
 .../agents_api/models/session/get_session.py  | 11 ++--
 .../models/session/list_sessions.py           | 14 +++--
 .../models/session/patch_session.py           |  8 +--
 .../agents_api/models/session/session_data.py | 15 +++---
 .../models/session/update_session.py          |  8 +--
 .../agents_api/models/tools/create_tools.py   | 16 +++---
 .../agents_api/models/tools/delete_tools.py   | 10 ++--
 .../agents_api/models/tools/embed_tools.py    | 10 ++--
 .../agents_api/models/tools/get_tools.py      | 10 ++--
 .../agents_api/models/tools/list_tools.py     | 10 ++--
 .../agents_api/models/tools/search_tools.py   | 10 ++--
 .../agents_api/models/user/create_user.py     | 41 +++++++--------
 agents-api/agents_api/models/user/get_user.py | 10 ++--
 .../agents_api/models/user/list_users.py      | 15 +++---
 .../agents_api/models/user/patch_user.py      | 29 ++++++-----
 .../agents_api/models/user/update_user.py     | 12 ++---
 agents-api/agents_api/models/utils.py         | 23 +++++++++
 39 files changed, 326 insertions(+), 386 deletions(-)
 create mode 100644 agents-api/agents_api/models/utils.py

diff --git a/agents-api/agents_api/models/agent/create_agent.py b/agents-api/agents_api/models/agent/create_agent.py
index 48418a6e1..d065cee77 100644
--- a/agents-api/agents_api/models/agent/create_agent.py
+++ b/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,
diff --git a/agents-api/agents_api/models/agent/create_tools.py b/agents-api/agents_api/models/agent/create_tools.py
index 4496b3747..42332b8dd 100644
--- a/agents-api/agents_api/models/agent/create_tools.py
+++ b/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})
diff --git a/agents-api/agents_api/models/agent/delete_agent.py b/agents-api/agents_api/models/agent/delete_agent.py
index a5deca6dc..6557c63d6 100644
--- a/agents-api/agents_api/models/agent/delete_agent.py
+++ b/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)})
diff --git a/agents-api/agents_api/models/agent/get_agent.py b/agents-api/agents_api/models/agent/get_agent.py
index cac39c290..325df7930 100644
--- a/agents-api/agents_api/models/agent/get_agent.py
+++ b/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)})
diff --git a/agents-api/agents_api/models/agent/list_agents.py b/agents-api/agents_api/models/agent/list_agents.py
index 8ae442122..838b7acdb 100644
--- a/agents-api/agents_api/models/agent/list_agents.py
+++ b/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},
     )
diff --git a/agents-api/agents_api/models/agent/patch_agent.py b/agents-api/agents_api/models/agent/patch_agent.py
index e4f995041..ee2ad7103 100644
--- a/agents-api/agents_api/models/agent/patch_agent.py
+++ b/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,
diff --git a/agents-api/agents_api/models/agent/patch_tool.py b/agents-api/agents_api/models/agent/patch_tool.py
index fe284b6fb..a565d9f7b 100644
--- a/agents-api/agents_api/models/agent/patch_tool.py
+++ b/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)
diff --git a/agents-api/agents_api/models/agent/update_agent.py b/agents-api/agents_api/models/agent/update_agent.py
index dea381279..c1b950b79 100644
--- a/agents-api/agents_api/models/agent/update_agent.py
+++ b/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,
diff --git a/agents-api/agents_api/models/agent/update_tool.py b/agents-api/agents_api/models/agent/update_tool.py
index 5a58f9ce2..3310a29d1 100644
--- a/agents-api/agents_api/models/agent/update_tool.py
+++ b/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)
diff --git a/agents-api/agents_api/models/docs/create_docs.py b/agents-api/agents_api/models/docs/create_docs.py
index ba48c4fdd..389b4e5b7 100644
--- a/agents-api/agents_api/models/docs/create_docs.py
+++ b/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),
diff --git a/agents-api/agents_api/models/docs/delete_docs.py b/agents-api/agents_api/models/docs/delete_docs.py
index 05b7b67b4..99539a9c5 100644
--- a/agents-api/agents_api/models/docs/delete_docs.py
+++ b/agents-api/agents_api/models/docs/delete_docs.py
@@ -1,18 +1,16 @@
 from typing import Literal
 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
 
 
+@cozo_query
 def delete_docs_by_id_query(
     owner_type: Literal["user", "agent"],
     owner_id: UUID,
     doc_id: UUID,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """Constructs and returns a datalog query for deleting documents and associated information snippets.
 
     This function targets the 'cozodb' database, allowing for the removal of documents and their related information snippets based on the provided document ID and owner (user or agent).
@@ -64,4 +62,4 @@ def delete_docs_by_id_query(
         :returning
     }}"""
 
-    return client.run(query, {"doc_id": doc_id, "owner_id": owner_id})
+    return (query, {"doc_id": doc_id, "owner_id": owner_id})
diff --git a/agents-api/agents_api/models/docs/embed_docs.py b/agents-api/agents_api/models/docs/embed_docs.py
index 5d407577f..3ece66300 100644
--- a/agents-api/agents_api/models/docs/embed_docs.py
+++ b/agents-api/agents_api/models/docs/embed_docs.py
@@ -2,31 +2,27 @@
 
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
-
-
-"""Embeds document snippets in the cozodb database.
-
-Parameters:
-doc_id (UUID): The unique identifier for the document.
-snippet_indices (list[int]): Indices of the snippets in the document.
-embeddings (list[list[float]]): Embedding vectors for the snippets.
-client (CozoClient, optional): The Cozo client to interact with the database. Defaults to a pre-configured client instance.
-
-Returns:
-pd.DataFrame: A DataFrame containing the results of the embedding operation.
-"""
+from ..utils import cozo_query
 
 
+@cozo_query
 def embed_docs_snippets_query(
     doc_id: UUID,
     snippet_indices: list[int],
     embeddings: list[list[float]],
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
+    """Embeds document snippets in the cozodb database.
+
+    Parameters:
+    doc_id (UUID): The unique identifier for the document.
+    snippet_indices (list[int]): Indices of the snippets in the document.
+    embeddings (list[list[float]]): Embedding vectors for the snippets.
+
+    Returns:
+    tuple[str, dict]: A DataFrame containing the results of the embedding operation.
+    """
+
     doc_id = str(doc_id)
     # Ensure the number of snippet indices matches the number of embeddings.
     assert len(snippet_indices) == len(embeddings)
@@ -50,4 +46,4 @@ def embed_docs_snippets_query(
         :returning
     }"""
 
-    return client.run(query, {"records": records})
+    return (query, {"records": records})
diff --git a/agents-api/agents_api/models/docs/get_docs.py b/agents-api/agents_api/models/docs/get_docs.py
index aaf09cda1..ce46cc068 100644
--- a/agents-api/agents_api/models/docs/get_docs.py
+++ b/agents-api/agents_api/models/docs/get_docs.py
@@ -3,30 +3,27 @@
 from typing import Literal
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
-
-
-"""
-Retrieves snippets of documents by their ID from the CozoDB.
-
-Parameters:
-    owner_type (Literal["user", "agent"]): The type of the owner of the document.
-    doc_id (UUID): The unique identifier of the document.
-    client (CozoClient, optional): The CozoDB client instance. Defaults to a pre-configured client.
-
-Returns:
-    pd.DataFrame: A DataFrame containing the document snippets and related metadata.
-"""
+from ..utils import cozo_query
 
 
+@cozo_query
 def get_docs_snippets_by_id_query(
     owner_type: Literal["user", "agent"],
     doc_id: UUID,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
+    """
+    Retrieves snippets of documents by their ID from the CozoDB.
+
+    Parameters:
+        owner_type (Literal["user", "agent"]): The type of the owner of the document.
+        doc_id (UUID): The unique identifier of the document.
+        client (CozoClient, optional): The CozoDB client instance. Defaults to a pre-configured client.
+
+    Returns:
+        pd.DataFrame: A DataFrame containing the document snippets and related metadata.
+    """
+
     doc_id = str(doc_id)
 
     query = f"""
@@ -58,4 +55,4 @@ def get_docs_snippets_by_id_query(
             }}
     }}"""
 
-    return client.run(query, {"doc_id": doc_id})
+    return (query, {"doc_id": doc_id})
diff --git a/agents-api/agents_api/models/docs/list_docs.py b/agents-api/agents_api/models/docs/list_docs.py
index c10489d7c..ea53381c9 100644
--- a/agents-api/agents_api/models/docs/list_docs.py
+++ b/agents-api/agents_api/models/docs/list_docs.py
@@ -3,17 +3,15 @@
 from typing import Literal
 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
 
 
+@cozo_query
 def ensure_owner_exists_query(
     owner_type: Literal["user", "agent"],
     owner_id: UUID,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     owner_id = str(owner_id)
 
     # Query to check if an owner (user or agent) exists in the database
@@ -30,14 +28,14 @@ def ensure_owner_exists_query(
             }}
     }}"""
 
-    return client.run(query, {"owner_id": owner_id})
+    return (query, {"owner_id": owner_id})
 
 
+@cozo_query
 def list_docs_snippets_by_owner_query(
     owner_type: Literal["user", "agent"],
     owner_id: UUID,
-    client: CozoClient = client,
-):
+) -> tuple[str, dict]:
     owner_id = str(owner_id)
 
     # Query to retrieve document snippets by owner (user or agent)
@@ -70,4 +68,4 @@ def list_docs_snippets_by_owner_query(
             }}
     }}"""
 
-    return client.run(query, {"owner_id": owner_id})
+    return (query, {"owner_id": owner_id})
diff --git a/agents-api/agents_api/models/docs/search_docs.py b/agents-api/agents_api/models/docs/search_docs.py
index 96c1af321..2e9c9d07c 100644
--- a/agents-api/agents_api/models/docs/search_docs.py
+++ b/agents-api/agents_api/models/docs/search_docs.py
@@ -3,36 +3,32 @@
 from typing import Literal
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
-
-
-"""
-Searches for document snippets in CozoDB by embedding query.
-
-Parameters:
-- owner_type (Literal["user", "agent"]): The type of the owner of the documents.
-- owner_id (UUID): The unique identifier of the owner.
-- query_embedding (list[float]): The embedding vector of the query.
-- k (int, optional): The number of nearest neighbors to retrieve. Defaults to 3.
-- confidence (float, optional): The confidence threshold for filtering results. Defaults to 0.8.
-- client (CozoClient, optional): The CozoDB client instance. Defaults to a pre-configured client.
-
-Returns:
-- pd.DataFrame: A DataFrame containing the search results.
-"""
+from ..utils import cozo_query
 
 
+@cozo_query
 def search_docs_snippets_by_embedding_query(
     owner_type: Literal["user", "agent"],
     owner_id: UUID,
     query_embedding: list[float],
     k: int = 3,
     confidence: float = 0.8,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
+    """
+    Searches for document snippets in CozoDB by embedding query.
+
+    Parameters:
+    - owner_type (Literal["user", "agent"]): The type of the owner of the documents.
+    - owner_id (UUID): The unique identifier of the owner.
+    - query_embedding (list[float]): The embedding vector of the query.
+    - k (int, optional): The number of nearest neighbors to retrieve. Defaults to 3.
+    - confidence (float, optional): The confidence threshold for filtering results. Defaults to 0.8.
+
+    Returns:
+    - pd.DataFrame: A DataFrame containing the search results.
+    """
+
     owner_id = str(owner_id)
     # Calculate the search radius based on confidence level
     radius: float = 1.0 - confidence
@@ -83,7 +79,7 @@ def search_docs_snippets_by_embedding_query(
         :sort distance
     }}"""
 
-    return client.run(
+    return (
         query,
         {
             "owner_id": owner_id,
diff --git a/agents-api/agents_api/models/entry/add_entries.py b/agents-api/agents_api/models/entry/add_entries.py
index 4643e8f4c..e8119ec6b 100644
--- a/agents-api/agents_api/models/entry/add_entries.py
+++ b/agents-api/agents_api/models/entry/add_entries.py
@@ -1,21 +1,16 @@
 import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
 from ...common.protocol.entries import Entry
+from ..utils import cozo_query
 from ...common.utils import json
 from ...common.utils.datetime import utcnow
 
 
-def add_entries_query(
-    entries: list[Entry], client: CozoClient = client
-) -> pd.DataFrame:
+@cozo_query
+def add_entries_query(entries: list[Entry]) -> tuple[str, dict]:
     """
-    Adds a list of Entry objects and a CozoClient object to a pandas DataFrame.
-
     Parameters:
         entries (list[Entry]): A list of Entry objects to be processed.
-        client (CozoClient): The CozoClient object used for database operations.
 
     Returns:
         pd.DataFrame: A DataFrame containing the entries data ready for insertion into the 'cozodb' database.
@@ -74,4 +69,4 @@ def add_entries_query(
     }
     """
 
-    return client.run(query, {"entries_lst": entries_lst})
+    return (query, {"entries_lst": entries_lst})
diff --git a/agents-api/agents_api/models/entry/delete_entries.py b/agents-api/agents_api/models/entry/delete_entries.py
index afccab17a..8e7aacbcf 100644
--- a/agents-api/agents_api/models/entry/delete_entries.py
+++ b/agents-api/agents_api/models/entry/delete_entries.py
@@ -1,11 +1,11 @@
 from uuid import UUID
 
-import pandas as pd
 
-from ...clients.cozo import client
+from ..utils import cozo_query
 
 
-def delete_entries_query(session_id: UUID) -> pd.DataFrame:
+@cozo_query
+def delete_entries_query(session_id: UUID) -> tuple[str, dict]:
     """
     Constructs and returns a datalog query for deleting entries associated with a given session ID from the 'cozodb' database.
 
@@ -57,4 +57,4 @@ def delete_entries_query(session_id: UUID) -> pd.DataFrame:
         }
     }"""
 
-    return client.run(query, {"session_id": str(session_id)})
+    return (query, {"session_id": str(session_id)})
diff --git a/agents-api/agents_api/models/entry/entries_summarization.py b/agents-api/agents_api/models/entry/entries_summarization.py
index cff1df602..6350968c6 100644
--- a/agents-api/agents_api/models/entry/entries_summarization.py
+++ b/agents-api/agents_api/models/entry/entries_summarization.py
@@ -2,25 +2,24 @@
 
 from uuid import UUID
 
-import pandas as pd
 
-from ...clients.cozo import client
 from ...common.protocol.entries import Entry
+from ..utils import cozo_query
 from ...common.utils import json
 
 
-"""
-Retrieves top-level entries from the database for a given session.
-
-Parameters:
-- session_id (UUID): The unique identifier for the session.
+@cozo_query
+def get_toplevel_entries_query(session_id: UUID) -> tuple[str, dict]:
+    """
+    Retrieves top-level entries from the database for a given session.
 
-Returns:
-- pd.DataFrame: A DataFrame containing the queried top-level entries.
-"""
+    Parameters:
+    - session_id (UUID): The unique identifier for the session.
 
+    Returns:
+    - pd.DataFrame: A DataFrame containing the queried top-level entries.
+    """
 
-def get_toplevel_entries_query(session_id: UUID) -> pd.DataFrame:
     query = """
         # Construct a datalog query to retrieve entries not summarized by any other entry.
     input[session_id] <- [[to_uuid($session_id)]]
@@ -58,25 +57,25 @@ def get_toplevel_entries_query(session_id: UUID) -> pd.DataFrame:
     :sort timestamp
     """
 
-    return client.run(query, {"session_id": str(session_id)})
+    return (query, {"session_id": str(session_id)})
 
 
-"""
-Inserts a new entry and its summarization relations into the database.
-
-Parameters:
-- session_id (UUID): The session identifier.
-- new_entry (Entry): The new entry to be inserted.
-- old_entry_ids (list[UUID]): List of entry IDs that the new entry summarizes.
+@cozo_query
+def entries_summarization_query(
+    session_id: UUID, new_entry: Entry, old_entry_ids: list[UUID]
+) -> tuple[str, dict]:
+    """
+    Inserts a new entry and its summarization relations into the database.
 
-Returns:
-- pd.DataFrame: A DataFrame containing the result of the insertion operation.
-"""
+    Parameters:
+    - session_id (UUID): The session identifier.
+    - new_entry (Entry): The new entry to be inserted.
+    - old_entry_ids (list[UUID]): List of entry IDs that the new entry summarizes.
 
+    Returns:
+    - pd.DataFrame: A DataFrame containing the result of the insertion operation.
+    """
 
-def entries_summarization_query(
-    session_id: UUID, new_entry: Entry, old_entry_ids: list[UUID]
-) -> pd.DataFrame:
     # Prepare relations data for insertion, marking the new entry as a summary of the old entries.
     relations = [
         [str(new_entry.id), "summary_of", str(old_id)] for old_id in old_entry_ids
@@ -131,4 +130,4 @@ def entries_summarization_query(
     }
     """
 
-    return client.run(query, {"relations": relations, "entries": entries})
+    return (query, {"relations": relations, "entries": entries})
diff --git a/agents-api/agents_api/models/entry/get_entries.py b/agents-api/agents_api/models/entry/get_entries.py
index c5959c51b..c8be17daa 100644
--- a/agents-api/agents_api/models/entry/get_entries.py
+++ b/agents-api/agents_api/models/entry/get_entries.py
@@ -1,14 +1,13 @@
 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
 
 
+@cozo_query
 def get_entries_query(
-    session_id: UUID, limit: int = 100, offset: int = 0, client: CozoClient = client
-) -> pd.DataFrame:
+    session_id: UUID, limit: int = 100, offset: int = 0
+) -> tuple[str, dict]:
     """
     Constructs and executes a query to retrieve entries from the 'cozodb' database.
 
@@ -55,9 +54,4 @@ def get_entries_query(
         :offset $offset
     }"""
 
-    results = client.run(
-        query, {"session_id": str(session_id), "limit": limit, "offset": offset}
-    )
-
-    # Execute the constructed query against the 'cozodb' database using the provided parameters and return the results.
-    return results
+    return (query, {"session_id": str(session_id), "limit": limit, "offset": offset})
diff --git a/agents-api/agents_api/models/entry/proc_mem_context.py b/agents-api/agents_api/models/entry/proc_mem_context.py
index 3e9c9287c..4a9a9649a 100644
--- a/agents-api/agents_api/models/entry/proc_mem_context.py
+++ b/agents-api/agents_api/models/entry/proc_mem_context.py
@@ -1,11 +1,10 @@
 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
 
 
+@cozo_query
 def proc_mem_context_query(
     session_id: UUID,
     tool_query_embedding: list[float],
@@ -14,8 +13,7 @@ def proc_mem_context_query(
     docs_confidence: float = 0.7,
     k_tools: int = 3,
     k_docs: int = 2,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """Executes a complex query to retrieve memory context based on session ID, tool and document embeddings.
 
     Parameters:
@@ -312,7 +310,7 @@ def proc_mem_context_query(
     }}
     """
 
-    return client.run(
+    return (
         query,
         {
             "session_id": session_id,
diff --git a/agents-api/agents_api/models/session/create_session.py b/agents-api/agents_api/models/session/create_session.py
index b65b8291f..62dd926cc 100644
--- a/agents-api/agents_api/models/session/create_session.py
+++ b/agents-api/agents_api/models/session/create_session.py
@@ -5,29 +5,11 @@
 
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
-
-
-"""
-Constructs and executes a datalog query to create a new session in the database.
-
-Parameters:
-- session_id (UUID): The unique identifier for the session.
-- developer_id (UUID): The unique identifier for the developer.
-- agent_id (UUID): The unique identifier for the agent.
-- user_id (UUID | None): The unique identifier for the user, if applicable.
-- situation (str | None): The situation/context of the session.
-- metadata (dict): Additional metadata for the session.
-- client (CozoClient): The database client used to execute the query.
-
-Returns:
-- pd.DataFrame: The result of the query execution.
-"""
+from ..utils import cozo_query
 
 
+@cozo_query
 def create_session_query(
     session_id: UUID,
     developer_id: UUID,
@@ -35,8 +17,22 @@ def create_session_query(
     user_id: UUID | None,
     situation: str | None,
     metadata: dict = {},
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
+    """
+    Constructs and executes a datalog query to create a new session in the database.
+
+    Parameters:
+    - session_id (UUID): The unique identifier for the session.
+    - developer_id (UUID): The unique identifier for the developer.
+    - agent_id (UUID): The unique identifier for the agent.
+    - user_id (UUID | None): The unique identifier for the user, if applicable.
+    - situation (str | None): The situation/context of the session.
+    - metadata (dict): Additional metadata for the session.
+
+    Returns:
+    - pd.DataFrame: The result of the query execution.
+    """
+
     situation: str = situation or ""
 
     # Construct the datalog query for creating a new session and its lookup.
@@ -74,7 +70,7 @@ def create_session_query(
      }"""
 
     # Execute the constructed query with the provided parameters and return the result.
-    return client.run(
+    return (
         query,
         {
             "session_id": str(session_id),
diff --git a/agents-api/agents_api/models/session/delete_session.py b/agents-api/agents_api/models/session/delete_session.py
index fa864eec6..050d0ce82 100644
--- a/agents-api/agents_api/models/session/delete_session.py
+++ b/agents-api/agents_api/models/session/delete_session.py
@@ -2,24 +2,21 @@
 
 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
 
 
+@cozo_query
 def delete_session_query(
     developer_id: UUID,
     session_id: UUID,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """
     Deletes a session and its related data from the 'cozodb' database.
 
     Parameters:
     - developer_id (UUID): The unique identifier for the developer.
     - session_id (UUID): The unique identifier for the session to be deleted.
-    - client (CozoClient): The database client used to execute the query.
 
     Returns:
     - pd.DataFrame: A DataFrame containing the result of the deletion query.
@@ -79,4 +76,4 @@ def delete_session_query(
     }
     """
 
-    return client.run(query, {"session_id": session_id, "developer_id": developer_id})
+    return (query, {"session_id": session_id, "developer_id": developer_id})
diff --git a/agents-api/agents_api/models/session/get_session.py b/agents-api/agents_api/models/session/get_session.py
index f38e6db6d..ef8af6ce2 100644
--- a/agents-api/agents_api/models/session/get_session.py
+++ b/agents-api/agents_api/models/session/get_session.py
@@ -1,23 +1,20 @@
 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
 
 
+@cozo_query
 def get_session_query(
     developer_id: UUID,
     session_id: UUID,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """
     Constructs and executes a datalog query to retrieve session information from the 'cozodb' database.
 
     Parameters:
         developer_id (UUID): The developer's unique identifier.
         session_id (UUID): The session's unique identifier.
-        client (CozoClient): Allows specifying a custom CozoClient instance for the database connection.
 
     Returns:
         pd.DataFrame: The result of the query as a pandas DataFrame.
@@ -60,4 +57,4 @@ def get_session_query(
             session_id: id,
         }, updated_at = to_int(validity)"""
 
-    return client.run(query, {"session_id": session_id, "developer_id": developer_id})
+    return (query, {"session_id": session_id, "developer_id": developer_id})
diff --git a/agents-api/agents_api/models/session/list_sessions.py b/agents-api/agents_api/models/session/list_sessions.py
index ae61cbbd3..7460a8e3f 100644
--- a/agents-api/agents_api/models/session/list_sessions.py
+++ b/agents-api/agents_api/models/session/list_sessions.py
@@ -3,20 +3,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_sessions_query(
     developer_id: UUID,
     limit: int = 100,
     offset: int = 0,
     metadata_filter: dict[str, Any] = {},
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """Lists sessions from the 'cozodb' database based on the provided filters.
 
     Parameters:
@@ -24,7 +22,6 @@ def list_sessions_query(
         limit (int): The maximum number of sessions to return.
         offset (int): The offset from which to start listing sessions.
         metadata_filter (dict[str, Any]): A dictionary of metadata fields to filter sessions by.
-        client (CozoClient): The database client instance to use for the query.
 
     Returns:
         pd.DataFrame: A DataFrame containing the queried session data.
@@ -76,6 +73,7 @@ def list_sessions_query(
     """
 
     # Execute the datalog query and return the results as a pandas DataFrame.
-    return client.run(
-        query, {"developer_id": str(developer_id), "limit": limit, "offset": offset}
+    return (
+        query,
+        {"developer_id": str(developer_id), "limit": limit, "offset": offset},
     )
diff --git a/agents-api/agents_api/models/session/patch_session.py b/agents-api/agents_api/models/session/patch_session.py
index 60c5885e2..36cab0790 100644
--- a/agents-api/agents_api/models/session/patch_session.py
+++ b/agents-api/agents_api/models/session/patch_session.py
@@ -2,10 +2,9 @@
 
 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
 
 
 _fields = [
@@ -17,11 +16,12 @@
 ]
 
 
+@cozo_query
 def patch_session_query(
     session_id: UUID,
     developer_id: UUID,
     **update_data,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """Patch session data in the 'cozodb' database.
 
     Parameters:
@@ -85,7 +85,7 @@ def patch_session_query(
 
     combined_query = "{" + assertion_query + "}" + session_update_query
 
-    return client.run(
+    return (
         combined_query,
         {
             "session_update_vals": session_update_vals,
diff --git a/agents-api/agents_api/models/session/session_data.py b/agents-api/agents_api/models/session/session_data.py
index 3043b3d51..ea765371a 100644
--- a/agents-api/agents_api/models/session/session_data.py
+++ b/agents-api/agents_api/models/session/session_data.py
@@ -1,23 +1,22 @@
 from uuid import UUID
 
-import pandas as pd
 from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
+from ...clients.cozo import client as cozo_client
 from ...common.protocol.sessions import SessionData
+from ..utils import cozo_query
 
 
+@cozo_query
 def session_data_query(
     developer_id: UUID,
     session_id: UUID,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """Constructs and executes a datalog query to retrieve session data from the 'cozodb' database.
 
     Parameters:
         developer_id (UUID): The developer's unique identifier.
         session_id (UUID): The session's unique identifier.
-        client (CozoClient): The database client used to execute the query.
 
     Returns:
         pd.DataFrame: A DataFrame containing the query results.
@@ -98,13 +97,11 @@ def session_data_query(
         }
     """
 
-    return client.run(
-        query, {"developer_id": str(developer_id), "session_id": str(session_id)}
-    )
+    return (query, {"developer_id": str(developer_id), "session_id": str(session_id)})
 
 
 def get_session_data(
-    developer_id: UUID, session_id: UUID, client: CozoClient = client
+    developer_id: UUID, session_id: UUID, client: CozoClient = cozo_client
 ) -> SessionData | None:
     """Calls `session_data_query` to get session data and processes the result.
 
diff --git a/agents-api/agents_api/models/session/update_session.py b/agents-api/agents_api/models/session/update_session.py
index c4227afa4..906dfeb34 100644
--- a/agents-api/agents_api/models/session/update_session.py
+++ b/agents-api/agents_api/models/session/update_session.py
@@ -1,9 +1,8 @@
 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
 
 
 _fields = [
@@ -16,11 +15,12 @@
 ]
 
 
+@cozo_query
 def update_session_query(
     session_id: UUID,
     developer_id: UUID,
     **update_data,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     # Process the update data to prepare it for the query.
     assertion_query = """
     ?[session_id, developer_id] := 
@@ -71,7 +71,7 @@ def update_session_query(
 
     combined_query = "{" + assertion_query + "}" + session_update_query
 
-    return client.run(
+    return (
         combined_query,
         {
             "session_update_vals": session_update_vals,
diff --git a/agents-api/agents_api/models/tools/create_tools.py b/agents-api/agents_api/models/tools/create_tools.py
index 7207e1a1f..f0bc1f317 100644
--- a/agents-api/agents_api/models/tools/create_tools.py
+++ b/agents-api/agents_api/models/tools/create_tools.py
@@ -1,20 +1,18 @@
 from uuid import uuid4, UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
 from ...autogen.openapi_model import FunctionDef
-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_function_query(
     agent_id: UUID,
     id: UUID,
     function: FunctionDef,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     created_at = utcnow().timestamp()
 
     # Process function definitions
@@ -42,14 +40,14 @@ def create_function_query(
         :returning
     }}"""
 
-    return client.run(query, {"function_rows": function_rows})
+    return (query, {"function_rows": function_rows})
 
 
+@cozo_query
 def create_multiple_functions_query(
     agent_id: UUID,
     functions: list[FunctionDef],
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     agent_id = str(agent_id)
     created_at = utcnow().timestamp()
 
@@ -82,4 +80,4 @@ def create_multiple_functions_query(
         :returning
     }}"""
 
-    return client.run(query, {"function_rows": function_rows})
+    return (query, {"function_rows": function_rows})
diff --git a/agents-api/agents_api/models/tools/delete_tools.py b/agents-api/agents_api/models/tools/delete_tools.py
index 6f751a136..b028310ac 100644
--- a/agents-api/agents_api/models/tools/delete_tools.py
+++ b/agents-api/agents_api/models/tools/delete_tools.py
@@ -1,16 +1,14 @@
 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
 
 
+@cozo_query
 def delete_function_by_id_query(
     agent_id: UUID,
     tool_id: UUID,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     agent_id = str(agent_id)
     tool_id = str(tool_id)
 
@@ -29,4 +27,4 @@ def delete_function_by_id_query(
         :returning
     }"""
 
-    return client.run(query, {"tool_id": tool_id, "agent_id": agent_id})
+    return (query, {"tool_id": tool_id, "agent_id": agent_id})
diff --git a/agents-api/agents_api/models/tools/embed_tools.py b/agents-api/agents_api/models/tools/embed_tools.py
index ab9867998..b83972bc5 100644
--- a/agents-api/agents_api/models/tools/embed_tools.py
+++ b/agents-api/agents_api/models/tools/embed_tools.py
@@ -1,17 +1,15 @@
 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
 
 
+@cozo_query
 def embed_functions_query(
     agent_id: UUID,
     tool_ids: list[UUID],
     embeddings: list[list[float]],
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     agent_id = str(agent_id)
     tool_ids = [str(id) for id in tool_ids]
     assert len(tool_ids) == len(embeddings)
@@ -31,4 +29,4 @@ def embed_functions_query(
         :returning
     }"""
 
-    return client.run(query, {"records": records})
+    return (query, {"records": records})
diff --git a/agents-api/agents_api/models/tools/get_tools.py b/agents-api/agents_api/models/tools/get_tools.py
index 0cc96e8f2..c9c68e675 100644
--- a/agents-api/agents_api/models/tools/get_tools.py
+++ b/agents-api/agents_api/models/tools/get_tools.py
@@ -1,16 +1,14 @@
 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
 
 
+@cozo_query
 def get_function_by_id_query(
     agent_id: UUID,
     tool_id: UUID,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     agent_id = str(agent_id)
     tool_id = str(tool_id)
 
@@ -38,4 +36,4 @@ def get_function_by_id_query(
             }
     }"""
 
-    return client.run(query, {"agent_id": agent_id, "tool_id": tool_id})
+    return (query, {"agent_id": agent_id, "tool_id": tool_id})
diff --git a/agents-api/agents_api/models/tools/list_tools.py b/agents-api/agents_api/models/tools/list_tools.py
index 0ecf182b8..a59fa6674 100644
--- a/agents-api/agents_api/models/tools/list_tools.py
+++ b/agents-api/agents_api/models/tools/list_tools.py
@@ -1,17 +1,15 @@
 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
 
 
+@cozo_query
 def list_functions_by_agent_query(
     agent_id: UUID,
     limit: int = 100,
     offset: int = 0,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     agent_id = str(agent_id)
 
     query = """
@@ -41,4 +39,4 @@ def list_functions_by_agent_query(
         :offset $offset
     }"""
 
-    return client.run(query, {"agent_id": agent_id, "limit": limit, "offset": offset})
+    return (query, {"agent_id": agent_id, "limit": limit, "offset": offset})
diff --git a/agents-api/agents_api/models/tools/search_tools.py b/agents-api/agents_api/models/tools/search_tools.py
index db0faf016..4a1a5e57f 100644
--- a/agents-api/agents_api/models/tools/search_tools.py
+++ b/agents-api/agents_api/models/tools/search_tools.py
@@ -1,18 +1,16 @@
 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
 
 
+@cozo_query
 def search_functions_by_embedding_query(
     agent_id: UUID,
     query_embedding: list[float],
     k: int = 3,
     confidence: float = 0.8,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     agent_id = str(agent_id)
     radius: float = 1.0 - confidence
 
@@ -64,7 +62,7 @@ def search_functions_by_embedding_query(
         :sort distance
     """
 
-    return client.run(
+    return (
         query,
         {
             "agent_id": agent_id,
diff --git a/agents-api/agents_api/models/user/create_user.py b/agents-api/agents_api/models/user/create_user.py
index 3c1f70fa0..8e7d1b9c0 100644
--- a/agents-api/agents_api/models/user/create_user.py
+++ b/agents-api/agents_api/models/user/create_user.py
@@ -5,36 +5,33 @@
 
 from uuid import UUID
 
-import pandas as pd
-from pycozo.client import Client as CozoClient
 
-from ...clients.cozo import client
-
-
-"""
-Constructs and executes a datalog query to create a new user in the CozoDB database.
-
-Parameters:
-    user_id (UUID): The unique identifier for the user.
-    developer_id (UUID): The unique identifier for the developer creating the user.
-    name (str): The name of the user.
-    about (str): A brief description about the user.
-    metadata (dict, optional): Additional metadata about the user. Defaults to an empty dict.
-    client (CozoClient, optional): The CozoDB client instance to run the query. Defaults to a pre-configured client instance.
-
-Returns:
-    pd.DataFrame: A DataFrame containing the result of the query execution.
-"""
+from ..utils import cozo_query
 
 
+@cozo_query
 def create_user_query(
     user_id: UUID,
     developer_id: UUID,
     name: str,
     about: str,
     metadata: dict = {},
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
+    """
+    Constructs and executes a datalog query to create a new user in the CozoDB database.
+
+    Parameters:
+        user_id (UUID): The unique identifier for the user.
+        developer_id (UUID): The unique identifier for the developer creating the user.
+        name (str): The name of the user.
+        about (str): A brief description about the user.
+        metadata (dict, optional): Additional metadata about the user. Defaults to an empty dict.
+        client (CozoClient, optional): The CozoDB client instance to run the query. Defaults to a pre-configured client instance.
+
+    Returns:
+        pd.DataFrame: A DataFrame containing the result of the query execution.
+    """
+
     query = """
     {
         # Then create the user
@@ -52,7 +49,7 @@ def create_user_query(
         :returning
     }"""
 
-    return client.run(
+    return (
         query,
         {
             "user_id": str(user_id),
diff --git a/agents-api/agents_api/models/user/get_user.py b/agents-api/agents_api/models/user/get_user.py
index da300ada6..d4ed6aa14 100644
--- a/agents-api/agents_api/models/user/get_user.py
+++ b/agents-api/agents_api/models/user/get_user.py
@@ -1,16 +1,14 @@
 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
 
 
+@cozo_query
 def get_user_query(
     developer_id: UUID,
     user_id: UUID,
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     # Convert UUIDs to strings for query compatibility.
     user_id = str(user_id)
     developer_id = str(developer_id)
@@ -36,4 +34,4 @@ def get_user_query(
             metadata,
         }"""
 
-    return client.run(query, {"developer_id": developer_id, "user_id": user_id})
+    return (query, {"developer_id": developer_id, "user_id": user_id})
diff --git a/agents-api/agents_api/models/user/list_users.py b/agents-api/agents_api/models/user/list_users.py
index 4de472131..6ef82e960 100644
--- a/agents-api/agents_api/models/user/list_users.py
+++ b/agents-api/agents_api/models/user/list_users.py
@@ -1,21 +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 ..utils import cozo_query
 from ...common.utils import json
 
 
+@cozo_query
 def list_users_query(
     developer_id: UUID,
     limit: int = 100,
     offset: int = 0,
     metadata_filter: dict[str, Any] = {},
-    client: CozoClient = client,
-) -> pd.DataFrame:
+) -> tuple[str, dict]:
     """
     Queries the 'cozodb' database to list users associated with a specific developer.
 
@@ -24,7 +21,6 @@ def list_users_query(
     - limit (int): The maximum number of users to return. Defaults to 100.
     - offset (int): The number of users to skip before starting to collect the result set. Defaults to 0.
     - metadata_filter (dict[str, Any]): A dictionary representing filters to apply on user metadata.
-    - client (CozoClient): The database client used to run the query. Defaults to an instance of CozoClient.
 
     Returns:
     - pd.DataFrame: A DataFrame containing the queried user data.
@@ -67,6 +63,7 @@ def list_users_query(
     """
 
     # Execute the datalog query with the specified parameters and return the results as a DataFrame.
-    return client.run(
-        query, {"developer_id": str(developer_id), "limit": limit, "offset": offset}
+    return (
+        query,
+        {"developer_id": str(developer_id), "limit": limit, "offset": offset},
     )
diff --git a/agents-api/agents_api/models/user/patch_user.py b/agents-api/agents_api/models/user/patch_user.py
index f0c4429ee..a9f20d7bb 100644
--- a/agents-api/agents_api/models/user/patch_user.py
+++ b/agents-api/agents_api/models/user/patch_user.py
@@ -2,27 +2,28 @@
 
 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
 
 
-"""
-Generates a datalog query for updating a user's information.
-
-Parameters:
-- developer_id (UUID): The UUID of the developer.
-- user_id (UUID): The UUID of the user to be updated.
-- **update_data: Arbitrary keyword arguments representing the data to be updated.
+@cozo_query
+def patch_user_query(
+    developer_id: UUID, user_id: UUID, **update_data
+) -> tuple[str, dict]:
+    """
+    Generates a datalog query for updating a user's information.
 
-Returns:
-- pd.DataFrame: A pandas DataFrame containing the results of the query execution.
-"""
+    Parameters:
+    - developer_id (UUID): The UUID of the developer.
+    - user_id (UUID): The UUID of the user to be updated.
+    - **update_data: Arbitrary keyword arguments representing the data to be updated.
 
+    Returns:
+    - tuple[str, dict]: A pandas DataFrame containing the results of the query execution.
+    """
 
-def patch_user_query(developer_id: UUID, user_id: UUID, **update_data) -> pd.DataFrame:
     # Prepare data for mutation by filtering out None values and adding system-generated fields.
     metadata = update_data.pop("metadata", {}) or {}
     user_update_cols, user_update_vals = cozo_process_mutate_data(
@@ -53,7 +54,7 @@ def patch_user_query(developer_id: UUID, user_id: UUID, **update_data) -> pd.Dat
         :returning
     """
 
-    return client.run(
+    return (
         query,
         {
             "user_update_vals": user_update_vals,
diff --git a/agents-api/agents_api/models/user/update_user.py b/agents-api/agents_api/models/user/update_user.py
index 3c75b9681..bec0c39dd 100644
--- a/agents-api/agents_api/models/user/update_user.py
+++ b/agents-api/agents_api/models/user/update_user.py
@@ -1,16 +1,14 @@
 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
 from ...common.utils.cozo import cozo_process_mutate_data
 
 
+@cozo_query
 def update_user_query(
-    developer_id: UUID, user_id: UUID, client: CozoClient = client, **update_data
-) -> pd.DataFrame:
+    developer_id: UUID, user_id: UUID, **update_data
+) -> tuple[str, dict]:
     """Updates user information in the 'cozodb' database.
 
     Parameters:
@@ -73,7 +71,7 @@ def update_user_query(
     query = "{" + assertion_query + "} {" + query + "}"
     # Combines the assertion and update queries.
 
-    return client.run(
+    return (
         query,
         {
             "user_update_vals": user_update_vals,
diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py
new file mode 100644
index 000000000..a6fb71c86
--- /dev/null
+++ b/agents-api/agents_api/models/utils.py
@@ -0,0 +1,23 @@
+from functools import wraps
+from typing import Callable
+
+import pandas as pd
+
+from ..clients.cozo import client as cozo_client
+
+
+def cozo_query(func: Callable[..., tuple[str, dict]]):
+    """
+    Decorator that wraps a function that takes arbitrary arguments, and
+    returns a (query string, variables) tuple.
+
+    The wrapped function should additionally take a client keyword argument
+    and then run the query using the client, returning a DataFrame.
+    """
+
+    @wraps(func)
+    def wrapper(*args, client=cozo_client, **kwargs) -> pd.DataFrame:
+        query, variables = func(*args, **kwargs)
+        return client.run(query, variables)
+
+    return wrapper