feat(agents-api): Add support for reading setup args from metadata an…

…d Upgrade to python 3.12

Signed-off-by: Diwank Singh Tomer <[email protected]>

agents-api/.tool-versions

@@ -1 +1 @@
-python 3.11.9
+python 3.12.5

agents-api/Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.11-slim
+FROM python:3.12-slim
 
 ENV PYTHONUNBUFFERED True
 ENV POETRY_CACHE_DIR=/tmp/poetry_cache

agents-api/Dockerfile.worker

@@ -1,4 +1,4 @@
-FROM python:3.11-slim
+FROM python:3.12-slim
 
 ENV PYTHONUNBUFFERED True
 ENV POETRY_CACHE_DIR=/tmp/poetry_cache

agents-api/agents_api/activities/execute_integration.py

@@ -16,26 +16,33 @@ async def execute_integration(
     tool_name: str,
     integration: IntegrationDef,
     arguments: dict[str, Any],
+    setup: dict[str, Any] = {},
 ) -> Any:
     developer_id = context.execution_input.developer_id
     agent_id = context.execution_input.agent.id
     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
+        developer_id=developer_id, agent_id=agent_id, task_id=task_id, arg_type="args"
+    )
+
+    merged_tool_setup = get_tool_args_from_metadata(
+        developer_id=developer_id, agent_id=agent_id, task_id=task_id, arg_type="setup"
     )
 
     arguments = (
         merged_tool_args.get(tool_name, {}) | (integration.arguments or {}) | arguments
     )
 
+    setup = merged_tool_setup.get(tool_name, {}) | (integration.setup or {}) | setup
+
     try:
         if integration.provider == "dummy":
             return arguments
 
         return await integrations.run_integration_service(
             provider=integration.provider,
-            setup=integration.setup,
+            setup=setup,
             method=integration.method,
             arguments=arguments,
         )

agents-api/agents_api/models/tools/get_tool_args_from_metadata.py

@@ -1,3 +1,4 @@
+from typing import Literal
 from uuid import UUID
 
 from beartype import beartype
@@ -20,29 +21,30 @@ def tool_args_for_task(
     developer_id: UUID,
     agent_id: UUID,
     task_id: UUID,
+    arg_type: Literal["args", "setup"] = "args",
 ) -> tuple[list[str], dict]:
     agent_id = str(agent_id)
     task_id = str(task_id)
 
-    get_query = """
+    get_query = f"""
         input[agent_id, task_id] <- [[to_uuid($agent_id), to_uuid($task_id)]]
 
-        ?[args] :=
+        ?[values] :=
             input[agent_id, task_id],
-            *tasks {
+            *tasks {{
                 task_id,
                 metadata: task_metadata,
-            },
-            *agents {
+            }},
+            *agents {{
                 agent_id,
                 metadata: agent_metadata,
-            },
-            task_args = get(task_metadata, "x-tool-args", {}),
-            agent_args = get(agent_metadata, "x-tool-args", {}),
+            }},
+            task_{arg_type} = get(task_metadata, "x-tool-{arg_type}", {{}}),
+            agent_{arg_type} = get(agent_metadata, "x-tool-{arg_type}", {{}}),
 
             # Right values overwrite left values
             # See: https://docs.cozodb.org/en/latest/functions.html#Func.Vector.concat
-            args = concat(agent_args, task_args),
+            values = concat(agent_{arg_type}, task_{arg_type}),
     """
 
     queries = [
@@ -61,28 +63,29 @@ def tool_args_for_session(
     developer_id: UUID,
     session_id: UUID,
     agent_id: UUID,
+    arg_type: Literal["args", "setup"] = "args",
 ) -> tuple[list[str], dict]:
     session_id = str(session_id)
 
-    get_query = """
+    get_query = f"""
         input[session_id, agent_id] <- [[to_uuid($session_id), to_uuid($agent_id)]]
 
-        ?[args] :=
+        ?[values] :=
             input[session_id, agent_id],
-            *sessions {
+            *sessions {{
                 session_id,
                 metadata: session_metadata,
-            },
-            *agents {
+            }},
+            *agents {{
                 agent_id,
                 metadata: agent_metadata,
-            },
-            session_args = get(session_metadata, "x-tool-args"),
-            agent_args = get(agent_metadata, "x-tool-args"),
+            }},
+            session_{arg_type} = get(session_metadata, "x-tool-{arg_type}", {{}}),
+            agent_{arg_type} = get(agent_metadata, "x-tool-{arg_type}", {{}}),
 
             # Right values overwrite left values
             # See: https://docs.cozodb.org/en/latest/functions.html#Func.Vector.concat
-            args = concat(agent_args, session_args),  
+            values = concat(agent_{arg_type}, session_{arg_type}),
     """
 
     queries = [
@@ -103,7 +106,7 @@ def tool_args_for_session(
         TypeError: partialclass(HTTPException, status_code=400),
     }
 )
-@wrap_in_class(dict, transform=lambda x: x["args"], one=True)
+@wrap_in_class(dict, transform=lambda x: x["values"], one=True)
 @cozo_query
 @beartype
 def get_tool_args_from_metadata(
@@ -112,15 +115,22 @@ def get_tool_args_from_metadata(
     agent_id: UUID,
     session_id: UUID | None = None,
     task_id: UUID | None = None,
+    arg_type: Literal["args", "setup"] = "args",
 ) -> tuple[list[str], dict]:
     match session_id, task_id:
         case (None, task_id) if task_id is not None:
             return tool_args_for_task(
-                developer_id=developer_id, agent_id=agent_id, task_id=task_id
+                developer_id=developer_id,
+                agent_id=agent_id,
+                task_id=task_id,
+                arg_type=arg_type,
             )
         case (session_id, None) if session_id is not None:
             return tool_args_for_session(
-                developer_id=developer_id, agent_id=agent_id, session_id=session_id
+                developer_id=developer_id,
+                agent_id=agent_id,
+                session_id=session_id,
+                arg_type=arg_type,
             )
         case (_, _):
             raise ValueError("Either session_id or task_id must be provided")