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 #<issue-number>

* 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)

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__",
 ]

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():

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",
+    )