Merge pull request #959 from julep-ai/f/switch-to-pg

f/switch to pg

.github/workflows/lint-agents-api-pr.yml

@@ -23,6 +23,11 @@ jobs:
         uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
+
+      - name: Install Go migrate
+        uses: jaxxstorm/[email protected]
+        with: # Grab the latest version
+          repo: golang-migrate/migrate
 
       - name: Set up python and install dependencies
         run: |

.github/workflows/test-agents-api-pr.yml

@@ -23,6 +23,11 @@ jobs:
         uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
+
+      - name: Install Go migrate
+        uses: jaxxstorm/[email protected]
+        with: # Grab the latest version
+          repo: golang-migrate/migrate
 
       - name: Set up python and install dependencies
         run: |

.github/workflows/typecheck-agents-api-pr.yml

@@ -31,6 +31,11 @@ jobs:
         uses: astral-sh/setup-uv@v4
         with:
           enable-cache: true
+
+      - name: Install Go migrate
+        uses: jaxxstorm/[email protected]
+        with: # Grab the latest version
+          repo: golang-migrate/migrate
 
       - name: Set up python and install dependencies
         run: |

.gitignore

@@ -10,3 +10,4 @@ ngrok*
 */node_modules/
 .aider*
 .vscode/
+schema.sql

agents-api/.gitignore

@@ -1,5 +1,4 @@
 # Local database files
-cozo.db
 temporal.db
 *.bak
 *.dat

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/activities/container.py

@@ -0,0 +1,26 @@
+from typing import Any
+
+from aiobotocore.client import AioBaseClient
+from asyncpg.pool import Pool
+
+
+class State:
+    postgres_pool: Pool | None
+    s3_client: AioBaseClient | None
+
+    def __init__(self):
+        self.postgres_pool = None
+        self.s3_client = None
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        super().__setattr__(name, value)
+
+
+class Container:
+    state: State
+
+    def __init__(self):
+        self.state = State()
+
+
+container = Container()

agents-api/agents_api/activities/demo.py

@@ -1,19 +1,18 @@
-from typing import Callable
-
 from temporalio import activity
 
 from ..env import testing
 
 
 async def demo_activity(a: int, b: int) -> int:
     # Should throw an error if testing is not enabled
-    raise Exception("This should not be called in production")
+    msg = "This should not be called in production"
+    raise Exception(msg)
 
 
 async def mock_demo_activity(a: int, b: int) -> int:
     return a + b
 
 
-demo_activity: Callable[[int, int], int] = activity.defn(name="demo_activity")(
+demo_activity = activity.defn(name="demo_activity")(
     demo_activity if not testing else mock_demo_activity
 )

agents-api/agents_api/activities/excecute_api_call.py

@@ -1,26 +1,24 @@
 import base64
-from typing import Any, Optional, TypedDict, Union
+from typing import Any, TypedDict
 
 import httpx
 from beartype import beartype
 from temporalio import activity
 
 from ..autogen.openapi_model import ApiCallDef
-from ..common.storage_handler import auto_blob_store
 from ..env import testing
 
 
 class RequestArgs(TypedDict):
-    content: Optional[str]
-    data: Optional[dict[str, Any]]
-    json_: Optional[dict[str, Any]]
-    cookies: Optional[dict[str, str]]
-    params: Optional[Union[str, dict[str, Any]]]
-    url: Optional[str]
-    headers: Optional[dict[str, str]]
+    content: str | None
+    data: dict[str, Any] | None
+    json_: dict[str, Any] | None
+    cookies: dict[str, str] | None
+    params: str | dict[str, Any] | None
+    url: str | None
+    headers: dict[str, str] | None
 
 
-@auto_blob_store(deep=True)
 @beartype
 async def execute_api_call(
     api_call: ApiCallDef,

agents-api/agents_api/activities/execute_integration.py

@@ -3,16 +3,17 @@
 from beartype import beartype
 from temporalio import activity
 
+from ..app import lifespan
 from ..autogen.openapi_model import BaseIntegrationDef
 from ..clients import integrations
 from ..common.exceptions.tools import IntegrationExecutionException
 from ..common.protocol.tasks import ExecutionInput, StepContext
-from ..common.storage_handler import auto_blob_store
 from ..env import testing
-from ..models.tools import get_tool_args_from_metadata
+from ..queries import tools
+from .container import container
 
 
-@auto_blob_store(deep=True)
+@lifespan(container)
 @beartype
 async def execute_integration(
     context: StepContext,
@@ -22,23 +23,35 @@ async def execute_integration(
     setup: dict[str, Any] = {},
 ) -> Any:
     if not isinstance(context.execution_input, ExecutionInput):
-        raise TypeError("Expected ExecutionInput type for context.execution_input")
+        msg = "Expected ExecutionInput type for context.execution_input"
+        raise TypeError(msg)
 
     developer_id = context.execution_input.developer_id
     agent_id = context.execution_input.agent.id
+
+    if context.execution_input.task is None:
+        msg = "Task cannot be None in execution_input"
+        raise ValueError(msg)
+
     task_id = context.execution_input.task.id
 
-    merged_tool_args = get_tool_args_from_metadata(
-        developer_id=developer_id, agent_id=agent_id, task_id=task_id, arg_type="args"
+    merged_tool_args = await tools.get_tool_args_from_metadata(
+        developer_id=developer_id,
+        agent_id=agent_id,
+        task_id=task_id,
+        arg_type="args",
+        connection_pool=container.state.postgres_pool,
     )
 
-    merged_tool_setup = get_tool_args_from_metadata(
-        developer_id=developer_id, agent_id=agent_id, task_id=task_id, arg_type="setup"
+    merged_tool_setup = await tools.get_tool_args_from_metadata(
+        developer_id=developer_id,
+        agent_id=agent_id,
+        task_id=task_id,
+        arg_type="setup",
+        connection_pool=container.state.postgres_pool,
     )
 
-    arguments = (
-        merged_tool_args.get(tool_name, {}) | (integration.arguments or {}) | arguments
-    )
+    arguments = merged_tool_args.get(tool_name, {}) | (integration.arguments or {}) | arguments
 
     setup = merged_tool_setup.get(tool_name, {}) | (integration.setup or {}) | setup
 
@@ -53,10 +66,7 @@ async def execute_integration(
             arguments=arguments,
         )
 
-        if (
-            "error" in integration_service_response
-            and integration_service_response["error"]
-        ):
+        if integration_service_response.get("error"):
             raise IntegrationExecutionException(
                 integration=integration,
                 error=integration_service_response["error"],
@@ -69,9 +79,7 @@ async def execute_integration(
             integration_str = integration.provider + (
                 "." + integration.method if integration.method else ""
             )
-            activity.logger.error(
-                f"Error in execute_integration {integration_str}: {e}"
-            )
+            activity.logger.error(f"Error in execute_integration {integration_str}: {e}")
 
         raise