chore(agents-api): misc fixes + refators

agents-api/agents_api/activities/execute_system.py

@@ -78,14 +78,10 @@ async def execute_system(
         # Handle special cases for doc operations
         if system.operation == "create" and system.subresource == "doc":
             arguments["x_developer_id"] = arguments.pop("developer_id")
-            bg_runner = BackgroundTasks()
-            res = await handler(
+            return await handler(
                 data=CreateDocRequest(**arguments.pop("data")),
-                background_tasks=bg_runner,
                 **arguments,
             )
-            await bg_runner()
-            return res
 
         # Handle search operations
         if system.operation == "search" and system.subresource == "doc":

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

@@ -48,7 +48,7 @@ def _recursive_evaluate(expr, evaluator: SimpleEval):
                 }
 
                 activity.logger.error(
-                    f"Error in base_evaluate: {evaluate_error}\nVariables accessed: {variables_accessed}"
+                    f"Error in base_evaluate: {evaluate_error}\n"
                 )
             raise evaluate_error from e
     elif isinstance(expr, list):

agents-api/agents_api/app.py

@@ -1,8 +1,10 @@
 import os
 from contextlib import asynccontextmanager
-from typing import Any, Protocol
+from typing import Protocol
 
+from aiobotocore.client import AioBaseClient
 from aiobotocore.session import get_session
+from asyncpg.pool import Pool
 from fastapi import APIRouter, FastAPI
 from prometheus_fastapi_instrumentator import Instrumentator
 from scalar_fastapi import get_scalar_api_reference
@@ -11,17 +13,18 @@
 from .env import api_prefix, hostname, protocol, public_port
 
 
-class Assignable(Protocol):
-    def __setattr__(self, name: str, value: Any) -> None: ...
+class State(Protocol):
+    postgres_pool: Pool | None
+    s3_client: AioBaseClient | None
 
 
 class ObjectWithState(Protocol):
-    state: Assignable
+    state: State
 
 
 # TODO: This currently doesn't use env.py, we should move to using them
 @asynccontextmanager
-async def lifespan(*containers: list[FastAPI | ObjectWithState]):
+async def lifespan(*containers: FastAPI | ObjectWithState):
     # INIT POSTGRES #
     pg_dsn = os.environ.get("PG_DSN")
 
@@ -50,13 +53,17 @@ async def lifespan(*containers: list[FastAPI | ObjectWithState]):
         # CLOSE POSTGRES #
         for container in containers:
             if hasattr(container, "state") and getattr(container.state, "postgres_pool", None):
-                await container.state.postgres_pool.close()
+                pool = getattr(container.state, "postgres_pool", None)
+                if pool:
+                    await pool.close()
                 container.state.postgres_pool = None
 
         # CLOSE S3 #
         for container in containers:
             if hasattr(container, "state") and getattr(container.state, "s3_client", None):
-                await container.state.s3_client.close()
+                s3_client = getattr(container.state, "s3_client", None)
+                if s3_client:
+                    await s3_client.close()
                 container.state.s3_client = None
 
 

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

@@ -60,7 +60,7 @@ def _transform(d):
     Transition,
     transform=_transform,
 )
-@pg_query(debug=True)
+@pg_query
 @beartype
 async def list_execution_transitions(
     *,