Fix tests

.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: |

agents-api/agents_api/activities/demo.py

@@ -5,7 +5,8 @@
 
 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:

agents-api/agents_api/activities/excecute_api_call.py

@@ -1,5 +1,5 @@
 import base64
-from typing import Any, Optional, TypedDict, Union
+from typing import Any, TypedDict
 
 import httpx
 from beartype import beartype
@@ -10,13 +10,13 @@
 
 
 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
 
 
 @beartype

agents-api/agents_api/activities/execute_integration.py

@@ -23,7 +23,8 @@ 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
@@ -45,9 +46,7 @@ async def execute_integration(
         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
 
@@ -62,10 +61,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"],
@@ -78,9 +74,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
 

agents-api/agents_api/activities/execute_system.py

@@ -9,7 +9,7 @@
 from fastapi.background import BackgroundTasks
 from temporalio import activity
 
-from ..app import lifespan
+from ..app import app, lifespan
 from ..autogen.openapi_model import (
     ChatInput,
     CreateDocRequest,
@@ -29,7 +29,7 @@
 process_pool_executor = ProcessPoolExecutor()
 
 
-@lifespan(container)
+@lifespan(app, container)  # Both are needed because we are using the routes
 @beartype
 async def execute_system(
     context: StepContext,
@@ -39,7 +39,8 @@ async def execute_system(
     arguments: dict[str, Any] = system.arguments or {}
 
     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)
 
     arguments["developer_id"] = context.execution_input.developer_id
 
@@ -131,9 +132,7 @@ async def execute_system(
 
         # Run the synchronous function in another process
         loop = asyncio.get_running_loop()
-        return await loop.run_in_executor(
-            process_pool_executor, partial(handler, **arguments)
-        )
+        return await loop.run_in_executor(process_pool_executor, partial(handler, **arguments))
     except BaseException as e:
         if activity.in_activity():
             activity.logger.error(f"Error in execute_system_call: {e}")
@@ -151,19 +150,20 @@ def _create_search_request(arguments: dict) -> Any:
             confidence=arguments.pop("confidence", 0.5),
             limit=arguments.get("limit", 10),
         )
-    elif "text" in arguments:
+    if "text" in arguments:
         return TextOnlyDocSearchRequest(
             text=arguments.pop("text"),
             mmr_strength=arguments.pop("mmr_strength", 0),
             limit=arguments.get("limit", 10),
         )
-    elif "vector" in arguments:
+    if "vector" in arguments:
         return VectorDocSearchRequest(
             vector=arguments.pop("vector"),
             mmr_strength=arguments.pop("mmr_strength", 0),
             confidence=arguments.pop("confidence", 0.7),
             limit=arguments.get("limit", 10),
         )
+    return None
 
 
 # Keep the existing mock and activity definition