feat(agents-api): Add YAML support

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

agents-api/agents_api/middleware.py

@@ -0,0 +1,35 @@
+import re
+
+import yaml
+from fastapi import Request
+
+
+class YamlMiddleware:
+    def __init__(self, path_regex: str = r".*"):
+        self.path_regex = re.compile(path_regex)
+
+    async def __call__(self, request: Request, call_next):
+        content_type = request.headers.get("content-type", "").strip().lower()
+
+        # Filter out requests that are not for YAML and not for the specified path
+        if not self.path_regex.match(request.url.path) or content_type not in [
+            "application/x-yaml",
+            "application/yaml",
+            "text/yaml",
+            "text/x-yaml",
+        ]:
+            return await call_next(request)
+
+        # Parse the YAML body into a Python object
+        body = yaml.load(await request.body(), yaml.CSafeLoader)
+        request._json = body
+
+        # Switch headers to JSON
+        headers = request.headers.mutablecopy()
+        headers["content-type"] = "application/json"
+
+        request._headers = headers
+
+        # Continue processing the request
+        response = await call_next(request)
+        return response

agents-api/agents_api/routers/tasks/update_execution.py

@@ -38,7 +38,9 @@ async def update_execution(
             token_data = get_paused_execution_token(
                 developer_id=x_developer_id, execution_id=execution_id
             )
-            act_handle = temporal_client.get_async_activity_handle(token_data["task_token"])
+            act_handle = temporal_client.get_async_activity_handle(
+                token_data["task_token"]
+            )
             await act_handle.complete(data.input)
 
         case _:

agents-api/agents_api/web.py

@@ -11,16 +11,18 @@
 from fastapi import Depends, FastAPI, Request, status
 from fastapi.exceptions import HTTPException, RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.responses import JSONResponse
 from litellm.exceptions import APIError
 from pycozo.client import QueryException
 from temporalio.service import RPCError
 
-from agents_api.common.exceptions import BaseCommonException
-from agents_api.dependencies.auth import get_api_key
-from agents_api.env import sentry_dsn
-from agents_api.exceptions import PromptTooBigError
-from agents_api.routers import (
+from .common.exceptions import BaseCommonException
+from .dependencies.auth import get_api_key
+from .env import sentry_dsn
+from .exceptions import PromptTooBigError
+from .middleware import YamlMiddleware
+from .routers import (
     agents,
     docs,
     jobs,
@@ -89,6 +91,11 @@ def register_exceptions(app: FastAPI) -> None:
     max_age=3600,
 )
 
+app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=3)
+
+# Add yaml middleware
+app.middleware("http")(YamlMiddleware(path_regex=r"/agents/.+/tasks.*"))
+
 register_exceptions(app)
 
 app.include_router(agents.router)

agents-api/pyproject.toml

@@ -32,6 +32,7 @@ pydantic-partial = "^0.5.5"
 simpleeval = "^0.9.13"
 lz4 = "^4.3.3"
 
+pyyaml = "^6.0.2"
 [tool.poetry.group.dev.dependencies]
 ipython = "^8.26.0"
 ruff = "^0.5.5"

agents-api/tests/test_agent_routes.py

@@ -17,7 +17,7 @@ def _(client=client):
     response = client.request(
         method="POST",
         url="/agents",
-        data=data,
+        json=data,
     )
 
     assert response.status_code == 403

agents-api/tests/test_workflow_routes.py

@@ -0,0 +1,88 @@
+# Tests for task queries
+
+from uuid import uuid4
+
+from ward import test
+
+from .fixtures import cozo_client, test_agent, test_developer_id
+from .utils import patch_http_client_with_temporal
+
+
+@test("workflow route: evaluate step single")
+async def _(
+    cozo_client=cozo_client,
+    developer_id=test_developer_id,
+    agent=test_agent,
+):
+    agent_id = str(agent.id)
+    task_id = str(uuid4())
+
+    async with patch_http_client_with_temporal(
+        cozo_client=cozo_client, developer_id=developer_id
+    ) as (
+        make_request,
+        client,
+    ):
+        task_data = {
+            "name": "test task",
+            "description": "test task about",
+            "input_schema": {"type": "object", "additionalProperties": True},
+            "main": [{"evaluate": {"hello": '"world"'}}],
+        }
+
+        make_request(
+            method="POST",
+            url=f"/agents/{agent_id}/tasks/{task_id}",
+            json=task_data,
+        )
+
+        execution_data = dict(input={"test": "input"})
+
+        make_request(
+            method="POST",
+            url=f"/tasks/{task_id}/executions",
+            json=execution_data,
+        )
+
+
+@test("workflow route: evaluate step single with yaml")
+async def _(
+    cozo_client=cozo_client,
+    developer_id=test_developer_id,
+    agent=test_agent,
+):
+    agent_id = str(agent.id)
+    task_id = str(uuid4())
+
+    async with patch_http_client_with_temporal(
+        cozo_client=cozo_client, developer_id=developer_id
+    ) as (
+        make_request,
+        client,
+    ):
+        task_data = """
+        name: test task
+        description: test task about
+        input_schema:
+        type: object
+        additionalProperties: true
+
+        main:
+            - evaluate:
+                hello: '"world"'
+        """
+
+        make_request(
+            method="POST",
+            url=f"/agents/{agent_id}/tasks/{task_id}",
+            content=task_data,
+            headers={"Content-Type": "text/yaml"},
+        )
+
+        execution_data = dict(input={"test": "input"})
+
+        make_request(
+            method="POST",
+            url=f"/tasks/{task_id}/executions",
+            json=execution_data,
+        )

agents-api/tests/utils.py

@@ -3,6 +3,7 @@
 from contextlib import asynccontextmanager
 from unittest.mock import patch
 
+from fastapi.testclient import TestClient
 from temporalio.testing import WorkflowEnvironment
 
 from agents_api.worker.codec import pydantic_data_converter
@@ -41,3 +42,25 @@ async def patch_testing_temporal():
 
     # Reset log level
     logger.setLevel(previous_log_level)
+
+
+@asynccontextmanager
+async def patch_http_client_with_temporal(*, cozo_client, developer_id):
+    async with patch_testing_temporal():
+        from agents_api.env import api_key, api_key_header_name
+        from agents_api.web import app
+
+        client = TestClient(app=app)
+        app.state.cozo_client = cozo_client
+
+        def make_request(method, url, **kwargs):
+            headers = kwargs.pop("headers", {})
+            headers = {
+                **headers,
+                "X-Developer-Id": str(developer_id),
+                api_key_header_name: api_key,
+            }
+
+            return client.request(method, url, headers=headers, **kwargs)
+
+        yield make_request, client