From 3e13909e59447791d63dc560b67c649e4acbf769 Mon Sep 17 00:00:00 2001 From: claude89757 <138977524+claude89757@users.noreply.github.com> Date: Mon, 12 Jan 2026 23:07:34 +0800 Subject: [PATCH] Python: Add Pydantic request model and OpenAPI tags support to AG-UI FastAPI endpoint (#2522) * feat(ag-ui): Add Pydantic request model and OpenAPI tags support - Add AGUIRequest Pydantic model in _types.py with field descriptions - Update add_agent_framework_fastapi_endpoint() to accept tags parameter - Use AGUIRequest model for automatic validation and OpenAPI schema generation - Export AGUIRequest and DEFAULT_TAGS in __init__.py - Update test_endpoint.py to expect 422 for invalid requests - Add tests for OpenAPI schema, default tags, custom tags, and validation Benefits: - Better API documentation with complete request schema in Swagger UI - Automatic request validation with Pydantic - Organized endpoints under 'AG-UI' tag instead of 'default' - Improved developer experience and type safety Fixes # * test(ag-ui): Add test for internal error handling to achieve 100% coverage - Add test_endpoint_internal_error_handling() to cover exception handling code - Mock copy.deepcopy to simulate internal error during default_state processing - Add type: ignore for FastAPI tags parameter (known pyright compatibility issue) - Achieves 100% test coverage for _endpoint.py (previously missing lines 103-105) --- .../ag-ui/agent_framework_ag_ui/__init__.py | 6 + .../ag-ui/agent_framework_ag_ui/_endpoint.py | 11 +- .../ag-ui/agent_framework_ag_ui/_types.py | 23 ++++ python/packages/ag-ui/tests/test_endpoint.py | 125 +++++++++++++++++- 4 files changed, 156 insertions(+), 9 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui/__init__.py index 143f2499a0..631168067c 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/__init__.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/__init__.py @@ -16,22 +16,28 @@ from ._confirmation_strategies import ( from ._endpoint import add_agent_framework_fastapi_endpoint from ._event_converters import AGUIEventConverter from ._http_service import AGUIHttpService +from ._types import AGUIRequest try: __version__ = importlib.metadata.version(__name__) except importlib.metadata.PackageNotFoundError: __version__ = "0.0.0" +# Default OpenAPI tags for AG-UI endpoints +DEFAULT_TAGS = ["AG-UI"] + __all__ = [ "AgentFrameworkAgent", "add_agent_framework_fastapi_endpoint", "AGUIChatClient", "AGUIEventConverter", "AGUIHttpService", + "AGUIRequest", "ConfirmationStrategy", "DefaultConfirmationStrategy", "TaskPlannerConfirmationStrategy", "RecipeConfirmationStrategy", "DocumentWriterConfirmationStrategy", + "DEFAULT_TAGS", "__version__", ] diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py index eedf88db14..dda045b3fb 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py @@ -8,10 +8,11 @@ from typing import Any from ag_ui.encoder import EventEncoder from agent_framework import AgentProtocol -from fastapi import FastAPI, Request +from fastapi import FastAPI from fastapi.responses import StreamingResponse from ._agent import AgentFrameworkAgent +from ._types import AGUIRequest logger = logging.getLogger(__name__) @@ -24,6 +25,7 @@ def add_agent_framework_fastapi_endpoint( predict_state_config: dict[str, dict[str, str]] | None = None, allow_origins: list[str] | None = None, default_state: dict[str, Any] | None = None, + tags: list[str] | None = None, ) -> None: """Add an AG-UI endpoint to a FastAPI app. @@ -36,6 +38,7 @@ def add_agent_framework_fastapi_endpoint( Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}} allow_origins: CORS origins (not yet implemented) default_state: Optional initial state to seed when the client does not provide state keys + tags: OpenAPI tags for endpoint categorization (defaults to ["AG-UI"]) """ if isinstance(agent, AgentProtocol): wrapped_agent = AgentFrameworkAgent( @@ -46,15 +49,15 @@ def add_agent_framework_fastapi_endpoint( else: wrapped_agent = agent - @app.post(path) - async def agent_endpoint(request: Request): # type: ignore[misc] + @app.post(path, tags=tags or ["AG-UI"]) # type: ignore[arg-type] + async def agent_endpoint(request_body: AGUIRequest): # type: ignore[misc] """Handle AG-UI agent requests. Note: Function is accessed via FastAPI's decorator registration, despite appearing unused to static analysis. """ try: - input_data = await request.json() + input_data = request_body.model_dump(exclude_none=True) if default_state: state = input_data.setdefault("state", {}) for key, value in default_state.items(): diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_types.py b/python/packages/ag-ui/agent_framework_ag_ui/_types.py index da7d80ea66..b404f6d57e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_types.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_types.py @@ -4,6 +4,8 @@ from typing import Any, TypedDict +from pydantic import BaseModel, Field + class PredictStateConfig(TypedDict): """Configuration for predictive state updates.""" @@ -25,3 +27,24 @@ class AgentState(TypedDict): """Base state for AG-UI agents.""" messages: list[Any] | None + + +class AGUIRequest(BaseModel): + """Request model for AG-UI endpoints.""" + + messages: list[dict[str, Any]] = Field( + ..., + description="AG-UI format messages array", + ) + run_id: str | None = Field( + None, + description="Optional run identifier for tracking", + ) + thread_id: str | None = Field( + None, + description="Optional thread identifier for conversation context", + ) + state: dict[str, Any] | None = Field( + None, + description="Optional shared state for agentic generative UI", + ) diff --git a/python/packages/ag-ui/tests/test_endpoint.py b/python/packages/ag-ui/tests/test_endpoint.py index 36c9e3bc32..e8b4c86a49 100644 --- a/python/packages/ag-ui/tests/test_endpoint.py +++ b/python/packages/ag-ui/tests/test_endpoint.py @@ -176,11 +176,8 @@ async def test_endpoint_error_handling(): # Send invalid JSON to trigger parsing error before streaming response = client.post("/failing", data=b"invalid json", headers={"content-type": "application/json"}) # type: ignore - # The exception handler catches it and returns JSON error - assert response.status_code == 200 - content = json.loads(response.content) - assert "error" in content - assert content["error"] == "An internal error has occurred." + # Pydantic validation now returns 422 for invalid request body + assert response.status_code == 422 async def test_endpoint_multiple_paths(): @@ -266,3 +263,121 @@ async def test_endpoint_complex_input(): ) assert response.status_code == 200 + + +async def test_endpoint_openapi_schema(): + """Test that endpoint generates proper OpenAPI schema with request model.""" + app = FastAPI() + agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client()) + + add_agent_framework_fastapi_endpoint(app, agent, path="/schema-test") + + client = TestClient(app) + response = client.get("/openapi.json") + + assert response.status_code == 200 + openapi_spec = response.json() + + # Verify the endpoint exists in the schema + assert "/schema-test" in openapi_spec["paths"] + endpoint_spec = openapi_spec["paths"]["/schema-test"]["post"] + + # Verify request body schema is defined + assert "requestBody" in endpoint_spec + request_body = endpoint_spec["requestBody"] + assert "content" in request_body + assert "application/json" in request_body["content"] + + # Verify schema references AGUIRequest model + schema_ref = request_body["content"]["application/json"]["schema"] + assert "$ref" in schema_ref + assert "AGUIRequest" in schema_ref["$ref"] + + # Verify AGUIRequest model is in components + assert "components" in openapi_spec + assert "schemas" in openapi_spec["components"] + assert "AGUIRequest" in openapi_spec["components"]["schemas"] + + # Verify AGUIRequest has required fields + agui_request_schema = openapi_spec["components"]["schemas"]["AGUIRequest"] + assert "properties" in agui_request_schema + assert "messages" in agui_request_schema["properties"] + assert "run_id" in agui_request_schema["properties"] + assert "thread_id" in agui_request_schema["properties"] + assert "state" in agui_request_schema["properties"] + assert "required" in agui_request_schema + assert "messages" in agui_request_schema["required"] + + +async def test_endpoint_default_tags(): + """Test that endpoint uses default 'AG-UI' tag.""" + app = FastAPI() + agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client()) + + add_agent_framework_fastapi_endpoint(app, agent, path="/default-tags") + + client = TestClient(app) + response = client.get("/openapi.json") + + assert response.status_code == 200 + openapi_spec = response.json() + + endpoint_spec = openapi_spec["paths"]["/default-tags"]["post"] + assert "tags" in endpoint_spec + assert endpoint_spec["tags"] == ["AG-UI"] + + +async def test_endpoint_custom_tags(): + """Test that endpoint accepts custom tags.""" + app = FastAPI() + agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client()) + + add_agent_framework_fastapi_endpoint(app, agent, path="/custom-tags", tags=["Custom", "Agent"]) + + client = TestClient(app) + response = client.get("/openapi.json") + + assert response.status_code == 200 + openapi_spec = response.json() + + endpoint_spec = openapi_spec["paths"]["/custom-tags"]["post"] + assert "tags" in endpoint_spec + assert endpoint_spec["tags"] == ["Custom", "Agent"] + + +async def test_endpoint_missing_required_field(): + """Test that endpoint validates required fields with Pydantic.""" + app = FastAPI() + agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client()) + + add_agent_framework_fastapi_endpoint(app, agent, path="/validation") + + client = TestClient(app) + + # Missing required 'messages' field should trigger validation error + response = client.post("/validation", json={"run_id": "test-123"}) + + assert response.status_code == 422 + error_detail = response.json() + assert "detail" in error_detail + + +async def test_endpoint_internal_error_handling(): + """Test endpoint error handling when an exception occurs before streaming starts.""" + from unittest.mock import patch + + app = FastAPI() + agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client()) + + # Use default_state to trigger the code path that can raise an exception + add_agent_framework_fastapi_endpoint(app, agent, path="/error-test", default_state={"key": "value"}) + + client = TestClient(app) + + # Mock copy.deepcopy to raise an exception during default_state processing + with patch("agent_framework_ag_ui._endpoint.copy.deepcopy") as mock_deepcopy: + mock_deepcopy.side_effect = Exception("Simulated internal error") + response = client.post("/error-test", json={"messages": [{"role": "user", "content": "Hello"}]}) + + assert response.status_code == 200 + assert response.json() == {"error": "An internal error has occurred."}