chore: remove cozo completely, and integrate postgres

agents-api/Dockerfile

@@ -30,4 +30,4 @@ COPY . ./
 ENV PYTHONUNBUFFERED=1
 ENV GUNICORN_CMD_ARGS="--capture-output --enable-stdio-inheritance"
 
-ENTRYPOINT ["uv", "run", "gunicorn", "agents_api.web:app", "-c", "gunicorn_conf.py"]
+ENTRYPOINT ["uv", "run", "--offline", "--no-sync", "gunicorn", "agents_api.web:app", "-c", "gunicorn_conf.py"]

agents-api/Dockerfile.worker

@@ -30,4 +30,4 @@ COPY . ./
 ENV PYTHONUNBUFFERED=1
 ENV GUNICORN_CMD_ARGS="--capture-output --enable-stdio-inheritance"
 
-ENTRYPOINT ["uv", "run", "python", "-m", "agents_api.worker"]
+ENTRYPOINT ["uv", "run", "--offline", "--no-sync", "python", "-m", "agents_api.worker"]

agents-api/agents_api/__init__.py

@@ -9,3 +9,6 @@
 
 with workflow.unsafe.imports_passed_through():
     import msgpack as msgpack
+
+import os
+

agents-api/agents_api/activities/execute_integration.py

@@ -8,7 +8,7 @@
 from ..common.exceptions.tools import IntegrationExecutionException
 from ..common.protocol.tasks import ExecutionInput, StepContext
 from ..env import testing
-from ..models.tools import get_tool_args_from_metadata
+from ..queries.tools import get_tool_args_from_metadata
 
 
 @beartype

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

@@ -2,7 +2,7 @@
 
 from .base_evaluate import base_evaluate
 
-# from .cozo_query_step import cozo_query_step
+from .pg_query_step import pg_query_step
 from .evaluate_step import evaluate_step
 from .for_each_step import for_each_step
 from .get_value_step import get_value_step

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

@@ -6,7 +6,7 @@
 
 from ... import queries
 from ...clients.pg import create_db_pool
-from ...env import db_dsn, testing
+from ...env import pg_dsn, testing
 
 
 @alru_cache(maxsize=1)
@@ -18,7 +18,7 @@ async def get_db_pool(dsn: str):
 async def pg_query_step(
     query_name: str,
     values: dict[str, Any],
-    dsn: str = db_dsn,
+    dsn: str = pg_dsn,
 ) -> Any:
     pool = await get_db_pool(dsn=dsn)
 

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

@@ -15,7 +15,7 @@
 )
 from ...exceptions import LastErrorInput, TooManyRequestsError
 from ...queries.executions.create_execution_transition import (
-    create_execution_transition_async,
+    create_execution_transition,
 )
 from ..utils import RateLimiter
 
@@ -52,7 +52,7 @@ async def transition_step(
 
     # Create transition
     try:
-        transition = await create_execution_transition_async(
+        transition = await create_execution_transition(
             developer_id=context.execution_input.developer_id,
             execution_id=context.execution_input.execution.id,
             task_id=context.execution_input.task.id,

agents-api/agents_api/app.py

@@ -17,10 +17,10 @@
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # INIT POSTGRES #
-    db_dsn = os.environ.get("DB_DSN")
+    pg_dsn = os.environ.get("PG_DSN")
 
     if not getattr(app.state, "postgres_pool", None):
-        app.state.postgres_pool = await create_db_pool(db_dsn)
+        app.state.postgres_pool = await create_db_pool(pg_dsn)
 
     # INIT S3 #
     s3_access_key = os.environ.get("S3_ACCESS_KEY")
@@ -67,7 +67,8 @@ async def lifespan(app: FastAPI):
     lifespan=lifespan,
     #
     # Global dependencies
-    dependencies=[Depends(valid_content_length)],
+    # FIXME: This is blocking access to scalar
+    # dependencies=[Depends(valid_content_length)],
 )
 
 # Enable metrics
@@ -92,19 +93,20 @@ async def scalar_html():
 
 
 # content-length validation
+# FIXME: This is blocking access to scalar
 # NOTE: This relies on client reporting the correct content-length header
 # TODO: We should use streaming for large payloads
-@app.middleware("http")
-async def validate_content_length(
-    request: Request,
-    call_next: Callable[[Request], Coroutine[Any, Any, Response]],
-):
-    content_length = request.headers.get("content-length")
+# @app.middleware("http")
+# async def validate_content_length(
+#     request: Request,
+#     call_next: Callable[[Request], Coroutine[Any, Any, Response]],
+# ):
+#     content_length = request.headers.get("content-length")
 
-    if not content_length:
-        return Response(status_code=411, content="Content-Length header is required")
+#     if not content_length:
+#         return Response(status_code=411, content="Content-Length header is required")
 
-    if int(content_length) > max_payload_size:
-        return Response(status_code=413, content="Payload too large")
+#     if int(content_length) > max_payload_size:
+#         return Response(status_code=413, content="Payload too large")
 
-    return await call_next(request)
+#     return await call_next(request)

agents-api/agents_api/clients/pg.py

@@ -2,7 +2,7 @@
 
 import asyncpg
 
-from ..env import db_dsn
+from ..env import pg_dsn
 
 
 async def _init_conn(conn):
@@ -16,5 +16,5 @@ async def _init_conn(conn):
 
 async def create_db_pool(dsn: str | None = None):
     return await asyncpg.create_pool(
-        dsn if dsn is not None else db_dsn, init=_init_conn
+        dsn if dsn is not None else pg_dsn, init=_init_conn
     )

agents-api/agents_api/common/protocol/tasks.py

@@ -139,7 +139,7 @@ class PartialTransition(create_partial_model(CreateTransitionRequest)):
 
 class ExecutionInput(BaseModel):
     developer_id: UUID
-    execution: Execution
+    execution: Execution | None = None
     task: TaskSpecDef
     agent: Agent
     agent_tools: list[Tool | CreateToolRequest]

agents-api/agents_api/env.py

@@ -51,24 +51,15 @@
 s3_secret_key: str | None = env.str("S3_SECRET_KEY", default=None)
 
 
-# Cozo
-# ----
-cozo_host: str = env.str("COZO_HOST", default="http://127.0.0.1:9070")
-cozo_auth: str = env.str("COZO_AUTH_TOKEN", default=None)
-summarization_model_name: str = env.str(
-    "SUMMARIZATION_MODEL_NAME", default="gpt-4-turbo"
-)
-do_verify_developer: bool = env.bool("DO_VERIFY_DEVELOPER", default=True)
-do_verify_developer_owns_resource: bool = env.bool(
-    "DO_VERIFY_DEVELOPER_OWNS_RESOURCE", default=True
-)
-
 # PostgreSQL
 # ----
-db_dsn: str = env.str(
-    "DB_DSN",
+pg_dsn: str = env.str(
+    "PG_DSN",
     default="postgres://postgres:[email protected]:5432/postgres?sslmode=disable",
 )
+summarization_model_name: str = env.str(
+    "SUMMARIZATION_MODEL_NAME", default="gpt-4-turbo"
+)
 
 query_timeout: float = env.float("QUERY_TIMEOUT", default=90.0)
 
@@ -156,8 +147,6 @@ def _parse_optional_int(val: str | None) -> int | None:
 environment: Dict[str, Any] = dict(
     debug=debug,
     multi_tenant_mode=multi_tenant_mode,
-    cozo_host=cozo_host,
-    cozo_auth=cozo_auth,
     sentry_dsn=sentry_dsn,
     temporal_endpoint=temporal_endpoint,
     temporal_task_queue=temporal_task_queue,

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

@@ -1,11 +1,9 @@
 from typing import Any, TypeVar
 from uuid import UUID
 
-import sqlvalidator
 from beartype import beartype
 
 from ...common.protocol.sessions import ChatContext, make_session
-from ...exceptions import InvalidSQLQuery
 from ..utils import (
     pg_query,
     wrap_in_class,
@@ -15,8 +13,8 @@
 T = TypeVar("T")
 
 
-sql_query = sqlvalidator.parse(
-    """SELECT * FROM 
+sql_query ="""
+SELECT * FROM 
 (
     SELECT jsonb_agg(u) AS users FROM (
         SELECT 
@@ -103,9 +101,6 @@
             session_lookup.participant_type = 'agent'
     ) r
 ) AS toolsets"""
-)
-if not sql_query.is_valid():
-    raise InvalidSQLQuery("prepare_chat_context")
 
 
 def _transform(d):

agents-api/agents_api/queries/executions/prepare_execution_input.py

@@ -14,14 +14,17 @@
 
 sql_query = """SELECT * FROM 
 (
-    SELECT to_jsonb(a) AS agents FROM (
+    SELECT to_jsonb(a) AS agent FROM (
         SELECT * FROM agents
         WHERE
             developer_id = $1  AND
-            agent_id = $4
+            agent_id = (
+                SELECT agent_id FROM tasks 
+                WHERE developer_id = $1 AND task_id = $2
+            )
         LIMIT 1
     ) a
-) AS agents,
+) AS agent,
 (
     SELECT jsonb_agg(r) AS tools FROM (
         SELECT * FROM tools
@@ -31,25 +34,25 @@
     ) r
 ) AS tools,
 (
-    SELECT to_jsonb(t) AS tasks FROM (
+    SELECT to_jsonb(t) AS task FROM (
         SELECT * FROM tasks
         WHERE
             developer_id = $1 AND 
             task_id = $2
         LIMIT 1
     ) t
-) AS tasks,
-(
-    SELECT to_jsonb(e) AS executions FROM (
-        SELECT * FROM executions
-        WHERE
-            developer_id = $1 AND 
-            task_id = $2 AND
-            execution_id = $3
-        LIMIT 1
-    ) e
-) AS executions;
+) AS task;
 """
+# (
+#     SELECT to_jsonb(e) AS execution FROM (
+#         SELECT * FROM latest_executions
+#         WHERE
+#             developer_id = $1 AND 
+#             task_id = $2 AND
+#             execution_id = $3
+#         LIMIT 1
+#     ) e
+# ) AS execution;
 
 
 # @rewrap_exceptions(
@@ -70,7 +73,7 @@
     transform=lambda d: {
         **d,
         "task": {
-            "tools": [*d["task"].pop("tools")],
+            "tools": d["tools"],
             **d["task"],
         },
         "agent_tools": [
@@ -86,14 +89,11 @@ async def prepare_execution_input(
     task_id: UUID,
     execution_id: UUID,
 ) -> tuple[str, list]:
-    dummy_agent_id = UUID(int=0)
-
     return (
         sql_query,
         [
             str(developer_id),
             str(task_id),
-            str(execution_id),
-            str(dummy_agent_id),
+            # str(execution_id),
         ],
     )

agents-api/agents_api/queries/tasks/create_or_update_task.py

@@ -192,13 +192,13 @@ async def create_or_update_task(
             tool.type,
             tool.name,
             tool.description,
-            getattr(tool, tool.type),  # spec
+            getattr(tool, tool.type) and getattr(tool, tool.type).model_dump(mode="json"),  # spec
         ]
         for tool in data.tools or []
     ]
 
     # Generate workflows from task data using task_to_spec
-    workflows_spec = task_to_spec(data).model_dump(exclude_none=True, mode="json")
+    workflows_spec = task_to_spec(data).model_dump(mode="json")
     workflow_params = []
     for workflow in workflows_spec.get("workflows", []):
         workflow_name = workflow.get("name")
@@ -211,7 +211,7 @@ async def create_or_update_task(
                     workflow_name,  # $3
                     step_idx,  # $4
                     step["kind_"],  # $5
-                    step[step["kind_"]],  # $6
+                    step,  # $6
                 ]
             )
 

agents-api/agents_api/queries/tasks/create_task.py

@@ -167,13 +167,13 @@ async def create_task(
             tool.type,
             tool.name,
             tool.description,
-            getattr(tool, tool.type),  # spec
+            getattr(tool, tool.type) and getattr(tool, tool.type).model_dump(mode="json"),  # spec
         ]
         for tool in data.tools or []
     ]
 
     # Generate workflows from task data using task_to_spec
-    workflows_spec = task_to_spec(data).model_dump(exclude_none=True, mode="json")
+    workflows_spec = task_to_spec(data).model_dump(mode="json")
     workflow_params = []
     for workflow in workflows_spec.get("workflows", []):
         workflow_name = workflow.get("name")
@@ -187,7 +187,7 @@ async def create_task(
                     workflow_name,  # $4
                     step_idx,  # $5
                     step["kind_"],  # $6
-                    step[step["kind_"]],  # $7
+                    step,  # $7
                 ]
             )
 

agents-api/agents_api/queries/tasks/get_task.py

@@ -17,12 +17,7 @@
             CASE WHEN w.name IS NOT NULL THEN
                 jsonb_build_object(
                     'name', w.name,
-                    'steps', jsonb_build_array(
-                        jsonb_build_object(
-                            w.step_type, w.step_definition,
-                            'step_idx', w.step_idx           -- Not sure if this is needed
-                        )
-                    )
+                    'steps', jsonb_build_array(w.step_definition)
                 )
             END
         ) FILTER (WHERE w.name IS NOT NULL),

agents-api/agents_api/queries/tasks/patch_task.py

@@ -198,7 +198,7 @@ async def patch_task(
     else:
         workflow_query = new_workflows_query
         workflow_params = []
-        workflows_spec = task_to_spec(data).model_dump(exclude_none=True, mode="json")
+        workflows_spec = task_to_spec(data).model_dump(mode="json")
         for workflow in workflows_spec.get("workflows", []):
             workflow_name = workflow.get("name")
             steps = workflow.get("steps", [])

agents-api/agents_api/queries/tasks/update_task.py

@@ -137,7 +137,7 @@ async def update_task(
     ]
 
     # Generate workflows from task data
-    workflows_spec = task_to_spec(data).model_dump(exclude_none=True, mode="json")
+    workflows_spec = task_to_spec(data).model_dump(mode="json")
     workflow_params = []
     for workflow in workflows_spec.get("workflows", []):
         workflow_name = workflow.get("name")