feat: Move to async cozo client

agents-api/agents_api/activities/embed_docs.py

@@ -38,7 +38,7 @@ async def embed_batch(snippets):
         await asyncio.gather(*[embed_batch(snippets) for snippets in batched_snippets]),
     )
 
-    embed_snippets_query(
+    await embed_snippets_query(
         developer_id=payload.developer_id,
         doc_id=payload.doc_id,
         snippet_indices=indices,

agents-api/agents_api/activities/execute_system.py

@@ -78,7 +78,7 @@ async def execute_system(
 
         # Handle chat operations
         if system.operation == "chat" and system.resource == "session":
-            developer = get_developer(developer_id=arguments.pop("developer_id"))
+            developer = await get_developer(developer_id=arguments.pop("developer_id"))
             session_id = arguments.pop("session_id")
             x_custom_api_key = arguments.pop("x_custom_api_key", None)
             chat_input = ChatInput(**arguments)

agents-api/agents_api/activities/task_steps/prompt_step.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Callable
+from typing import Awaitable, Callable
 
 from anthropic import AsyncAnthropic  # Import AsyncAnthropic client
 from anthropic.types.beta.beta_message import BetaMessage
@@ -58,7 +58,9 @@ def format_tool(tool: Tool) -> dict:
     }
 
     if tool.type == "system":
-        handler: Callable = get_handler_with_filtered_params(tool.system)
+        handler: Callable[..., Awaitable] = get_handler_with_filtered_params(
+            tool.system
+        )
 
         lc_tool: BaseTool = tool_decorator(handler)
 

agents-api/agents_api/activities/utils.py

@@ -9,7 +9,7 @@
 import string
 import time
 import urllib.parse
-from typing import Any, Callable, ParamSpec, TypeVar
+from typing import Any, Callable, ParamSpec, TypeVar, Awaitable
 
 import re2
 import zoneinfo
@@ -270,7 +270,7 @@ def filtered_handler(*args, **kwargs):
     return filtered_handler
 
 
-def get_handler(system: SystemDef) -> Callable:
+def get_handler(system: SystemDef) -> Callable[..., Awaitable]:
     """
     Internal function to get the base handler without parameter filtering.
 

agents-api/agents_api/clients/cozo.py

@@ -1,6 +1,6 @@
 from typing import Dict
 
-from pycozo.client import Client
+from pycozo_async.client import Client
 
 from ..env import cozo_auth, cozo_host
 from ..web import app

agents-api/agents_api/common/utils/cozo.py

@@ -6,7 +6,7 @@
 from uuid import UUID
 
 from beartype import beartype
-from pycozo import Client
+from pycozo_async import Client
 
 # Define a mock client for testing purposes, simulating Cozo API client behavior.
 _fake_client: SimpleNamespace = SimpleNamespace()

agents-api/agents_api/dependencies/developer_id.py

@@ -24,7 +24,7 @@ async def get_developer_id(
         except ValueError as e:
             raise InvalidHeaderFormat("X-Developer-Id must be a valid UUID") from e
 
-    verify_developer(developer_id=x_developer_id)
+    await verify_developer(developer_id=x_developer_id)
 
     return x_developer_id
 
@@ -36,7 +36,7 @@ async def get_developer_data(
         assert (
             not x_developer_id
         ), "X-Developer-Id header not allowed in multi-tenant mode"
-        return get_developer(developer_id=UUID("00000000-0000-0000-0000-000000000000"))
+        return await get_developer(developer_id=UUID("00000000-0000-0000-0000-000000000000"))
 
     if not x_developer_id:
         raise InvalidHeaderFormat("X-Developer-Id header required")
@@ -47,6 +47,6 @@ async def get_developer_data(
         except ValueError as e:
             raise InvalidHeaderFormat("X-Developer-Id must be a valid UUID") from e
 
-    developer = get_developer(developer_id=x_developer_id)
+    developer = await get_developer(developer_id=x_developer_id)
 
     return developer

agents-api/agents_api/metrics/counters.py

@@ -1,5 +1,6 @@
+import inspect
 from functools import wraps
-from typing import Callable, ParamSpec, TypeVar
+from typing import Awaitable, Callable, ParamSpec, TypeVar
 
 from prometheus_client import Counter
 
@@ -8,17 +9,21 @@
 
 
 def increase_counter(metric_label: str, id_field_name: str = "developer_id"):
-    def decor(func: Callable[P, T]):
+    def decor(func: Callable[P, T] | Callable[P, Awaitable[T]]):
         metric = Counter(
             metric_label,
             f"Number of {metric_label} calls",
             labelnames=(id_field_name,),
         )
 
         @wraps(func)
-        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             metric.labels(kwargs.get(id_field_name, "not_set")).inc()
-            return func(*args, **kwargs)
+            return (
+                await func(*args, **kwargs)
+                if inspect.iscoroutinefunction(func)
+                else func(*args, **kwargs)
+            )
 
         return wrapper
 

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

@@ -8,7 +8,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import Agent, CreateAgentRequest

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

@@ -8,7 +8,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import Agent, CreateOrUpdateAgentRequest

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

@@ -7,7 +7,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import ResourceDeletedResponse

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

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import Agent

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

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import Agent

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

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import PatchAgentRequest, ResourceUpdatedResponse

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

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateAgentRequest

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

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.Chat import ChatInput
@@ -42,7 +42,7 @@ async def gather_messages(
     assert len(new_raw_messages) > 0
 
     # Get the session history
-    history: History = get_history(
+    history: History = await get_history(
         developer_id=developer.id,
         session_id=session_id,
         allowed_sources=["api_request", "api_response", "tool_response", "summarizer"],
@@ -75,7 +75,7 @@ async def gather_messages(
     user_ids = [user.id for user in chat_context.users]
     owners = [("user", user_id) for user_id in user_ids] + [("agent", active_agent_id)]
 
-    doc_references: list[DocReference] = search_docs_hybrid(
+    doc_references: list[DocReference] = await search_docs_hybrid(
         developer_id=developer.id,
         owners=owners,
         query=query_text,

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

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...common.protocol.sessions import ChatContext, make_session

agents-api/agents_api/models/developer/get_developer.py

@@ -5,7 +5,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...common.protocol.developers import Developer

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

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import CreateDocRequest, Doc

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

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import ResourceDeletedResponse

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

@@ -5,7 +5,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import ResourceUpdatedResponse

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

@@ -5,7 +5,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import Doc

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

@@ -6,7 +6,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import Doc

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

@@ -6,7 +6,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import DocReference

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

@@ -7,7 +7,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import DocReference

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

@@ -90,7 +90,7 @@ def dbsf_fuse(
 
 
 @beartype
-def search_docs_hybrid(
+async def search_docs_hybrid(
     *,
     developer_id: UUID,
     owners: list[tuple[Literal["user", "agent"], UUID]],
@@ -103,7 +103,7 @@ def search_docs_hybrid(
     metadata_filter: dict[str, Any] = {},
 ) -> list[DocReference]:
     # TODO: We should probably parallelize these queries
-    text_results = search_docs_by_text(
+    text_results = await search_docs_by_text(
         developer_id=developer_id,
         owners=owners,
         query=query,
@@ -112,7 +112,7 @@ def search_docs_hybrid(
         **text_search_options,
     )
 
-    embedding_results = search_docs_by_embedding(
+    embedding_results = await search_docs_by_embedding(
         developer_id=developer_id,
         owners=owners,
         query_embedding=query_embedding,

agents-api/agents_api/models/entry/create_entries.py

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import CreateEntryRequest, Entry, Relation

agents-api/agents_api/models/entry/delete_entries.py

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import ResourceDeletedResponse

agents-api/agents_api/models/entry/get_history.py

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import History

agents-api/agents_api/models/entry/list_entries.py

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import Entry

agents-api/agents_api/models/execution/create_execution.py

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import CreateExecutionRequest, Execution

agents-api/agents_api/models/execution/create_execution_transition.py

@@ -2,7 +2,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import (

agents-api/agents_api/models/execution/create_temporal_lookup.py

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 from temporalio.client import WorkflowHandle
 

agents-api/agents_api/models/execution/get_execution.py

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from pycozo.client import QueryException
+from pycozo_async.client import QueryException
 from pydantic import ValidationError
 
 from ...autogen.openapi_model import Execution