chore(ag-ui): address PR feedback

* Update examples to use `to_ag_ui`
* Use relative imports where appropriate
* Remove `path` and `logger` parameters from `to_ag_ui` and dependencies
* Clean up import failure messages
* Improve FastAGUI doc comment
* Use `PrefixedToolset` eliminating manual tool prefixing
* Use CallToolsNode to handle tool calls

Also make StateDeps more strict, to avoid runtime failures.

Author: stevenh

examples/pydantic_ai_ag_ui_examples/api/__init__.py

@@ -2,18 +2,18 @@
 
 from __future__ import annotations
 
-from .agentic_chat import router as agentic_chat_router
-from .agentic_generative_ui import router as agentic_generative_ui_router
-from .human_in_the_loop import router as human_in_the_loop_router
-from .predictive_state_updates import router as predictive_state_updates_router
-from .shared_state import router as shared_state_router
-from .tool_based_generative_ui import router as tool_based_generative_ui_router
+from .agentic_chat import app as agentic_chat_app
+from .agentic_generative_ui import app as agentic_generative_ui_app
+from .human_in_the_loop import app as human_in_the_loop_app
+from .predictive_state_updates import app as predictive_state_updates_app
+from .shared_state import app as shared_state_app
+from .tool_based_generative_ui import app as tool_based_generative_ui_app
 
 __all__: list[str] = [
-    'agentic_chat_router',
-    'agentic_generative_ui_router',
-    'human_in_the_loop_router',
-    'predictive_state_updates_router',
-    'shared_state_router',
-    'tool_based_generative_ui_router',
+    'agentic_chat_app',
+    'agentic_generative_ui_app',
+    'human_in_the_loop_app',
+    'predictive_state_updates_app',
+    'shared_state_app',
+    'tool_based_generative_ui_app',
 ]

examples/pydantic_ai_ag_ui_examples/api/agent.py

@@ -2,42 +2,34 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from types import NoneType
-from typing import Generic
-
 from dotenv import load_dotenv
 
 from pydantic_ai import Agent
-from pydantic_ai.ag_ui import Adapter
-from pydantic_ai.result import OutputDataT
+from pydantic_ai.ag_ui import FastAGUI
 from pydantic_ai.tools import AgentDepsT
 
 
-@dataclass(init=False, repr=False)
-class AGUIAgent(Generic[AgentDepsT, OutputDataT]):
-    """Pydantic AI agent with AG-UI adapter."""
-
-    agent: Agent[AgentDepsT, str]
-    adapter: Adapter[AgentDepsT, str]
-    instructions: str | None
-
-    def __init__(
-        self, deps_type: type[AgentDepsT] = NoneType, instructions: str | None = None
-    ) -> None:
-        """Initialize the API agent with AG-UI adapter.
-
-        Args:
-            deps_type: Type annotation for the agent dependencies.
-            instructions: Optional instructions for the agent.
-        """
-        # Ensure environment variables are loaded.
-        load_dotenv()
-
-        self.agent = Agent(
-            'openai:gpt-4o-mini',
-            output_type=str,
-            instructions=instructions,
-            deps_type=deps_type,
-        )
-        self.adapter = Adapter(agent=self.agent)
+def agent(
+    model: str = 'openai:gpt-4o-mini',
+    deps: AgentDepsT = None,
+    instructions: str | None = None,
+) -> FastAGUI[AgentDepsT, str]:
+    """Create a Pydantic AI agent with AG-UI adapter.
+
+    Args:
+        model: The model to use for the agent.
+        deps: Optional dependencies for the agent.
+        instructions: Optional instructions for the agent.
+
+    Returns:
+        An instance of FastAGUI with the agent and adapter.
+    """
+    # Ensure environment variables are loaded.
+    load_dotenv()
+
+    return Agent(
+        model,
+        output_type=str,
+        instructions=instructions,
+        deps_type=type(deps),
+    ).to_ag_ui(deps=deps)

examples/pydantic_ai_ag_ui_examples/api/agentic_chat.py

@@ -3,26 +3,16 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Annotated
 from zoneinfo import ZoneInfo
 
-from ag_ui.core import RunAgentInput
-from fastapi import APIRouter, Header
-from fastapi.responses import StreamingResponse
+from pydantic_ai.ag_ui import FastAGUI
 
-from pydantic_ai.ag_ui import SSE_CONTENT_TYPE
+from .agent import agent
 
-from .agent import AGUIAgent
+app: FastAGUI = agent()
 
-if TYPE_CHECKING:  # pragma: no cover
-    from ag_ui.core import RunAgentInput
 
-
-router: APIRouter = APIRouter(prefix='/agentic_chat')
-agui: AGUIAgent = AGUIAgent()
-
-
-@agui.agent.tool_plain
+@app.adapter.agent.tool_plain
 async def current_time(timezone: str = 'UTC') -> str:
     """Get the current time in ISO format.
 
@@ -34,22 +24,3 @@ async def current_time(timezone: str = 'UTC') -> str:
     """
     tz: ZoneInfo = ZoneInfo(timezone)
     return datetime.now(tz=tz).isoformat()
-
-
-@router.post('')
-async def handler(
-    input_data: RunAgentInput, accept: Annotated[str, Header()] = SSE_CONTENT_TYPE
-) -> StreamingResponse:
-    """Endpoint to handle AG-UI protocol requests and stream responses.
-
-    Args:
-        input_data: The AG-UI run input.
-        accept: The Accept header to specify the response format.
-
-    Returns:
-        A streaming response with event-stream media type.
-    """
-    return StreamingResponse(
-        agui.adapter.run(input_data, accept),
-        media_type=SSE_CONTENT_TYPE,
-    )

examples/pydantic_ai_ag_ui_examples/api/agentic_generative_ui.py

@@ -3,19 +3,17 @@
 from __future__ import annotations
 
 from enum import StrEnum
-from typing import Annotated, Any, Literal
+from typing import Any, Literal
 
-from ag_ui.core import EventType, RunAgentInput, StateDeltaEvent, StateSnapshotEvent
-from fastapi import APIRouter, Header
-from fastapi.responses import StreamingResponse
+from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
 from pydantic import BaseModel, Field
 
-from pydantic_ai.ag_ui import SSE_CONTENT_TYPE
+from pydantic_ai.ag_ui import FastAGUI
 
-from .agent import AGUIAgent
+from .agent import agent
 
-router: APIRouter = APIRouter(prefix='/agentic_generative_ui')
-instructions: str = """When planning use tools only, without any other messages.
+app: FastAGUI = agent(
+    instructions="""When planning use tools only, without any other messages.
 IMPORTANT:
 - Use the `create_plan` tool to set the initial state of the steps
 - Use the `update_plan_step` tool to update the status of each step
@@ -26,7 +24,7 @@
 Only one plan can be active at a time, so do not call the `create_plan` tool
 again until all the steps in current plan are completed.
 """
-agui: AGUIAgent = AGUIAgent(instructions=instructions)
+)
 
 
 class StepStatus(StrEnum):
@@ -73,7 +71,7 @@ class JSONPatchOp(BaseModel):
     )
 
 
-@agui.agent.tool_plain
+@app.adapter.agent.tool_plain
 def create_plan(steps: list[str]) -> StateSnapshotEvent:
     """Create a plan with multiple steps.
 
@@ -92,7 +90,7 @@ def create_plan(steps: list[str]) -> StateSnapshotEvent:
     )
 
 
-@agui.agent.tool_plain
+@app.adapter.agent.tool_plain
 def update_plan_step(
     index: int, description: str | None = None, status: StepStatus | None = None
 ) -> StateDeltaEvent:
@@ -121,22 +119,3 @@ def update_plan_step(
         type=EventType.STATE_DELTA,
         delta=changes,
     )
-
-
-@router.post('')
-async def handler(
-    input_data: RunAgentInput, accept: Annotated[str, Header()] = SSE_CONTENT_TYPE
-) -> StreamingResponse:
-    """Endpoint to handle AG-UI protocol requests and stream responses.
-
-    Args:
-        input_data: The AG-UI run input.
-        accept: The Accept header to specify the response format.
-
-    Returns:
-        A streaming response with event-stream media type.
-    """
-    return StreamingResponse(
-        agui.adapter.run(input_data, accept),
-        media_type=SSE_CONTENT_TYPE,
-    )

examples/pydantic_ai_ag_ui_examples/api/human_in_the_loop.py

@@ -5,45 +5,16 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Annotated
+from pydantic_ai.ag_ui import FastAGUI
 
-from ag_ui.core import RunAgentInput
-from fastapi import APIRouter, Header
-from fastapi.responses import StreamingResponse
+from .agent import agent
 
-from pydantic_ai.ag_ui import SSE_CONTENT_TYPE
-
-from .agent import AGUIAgent
-
-if TYPE_CHECKING:  # pragma: no cover
-    from ag_ui.core import RunAgentInput
-
-
-instructions: str = """When planning tasks use tools only, without any other messages.
+app: FastAGUI = agent(
+    instructions="""When planning tasks use tools only, without any other messages.
 IMPORTANT:
 - Use the `generate_task_steps` tool to display the suggested steps to the user
 - Never repeat the plan, or send a message detailing steps
 - If accepted, confirm the creation of the plan and the number of selected (enabled) steps only
 - If not accepted, ask the user for more information, DO NOT use the `generate_task_steps` tool again
 """
-router: APIRouter = APIRouter(prefix='/human_in_the_loop')
-agui: AGUIAgent = AGUIAgent(instructions=instructions)
-
-
-@router.post('')
-async def handler(
-    input_data: RunAgentInput, accept: Annotated[str, Header()] = SSE_CONTENT_TYPE
-) -> StreamingResponse:
-    """Endpoint to handle AG-UI protocol requests and stream responses.
-
-    Args:
-        input_data: The AG-UI run input.
-        accept: The Accept header to specify the response format.
-
-    Returns:
-        A streaming response with event-stream media type.
-    """
-    return StreamingResponse(
-        agui.adapter.run(input_data, accept),
-        media_type=SSE_CONTENT_TYPE,
-    )
+)

examples/pydantic_ai_ag_ui_examples/api/predictive_state_updates.py

@@ -3,16 +3,14 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Annotated
+from typing import TYPE_CHECKING
 
-from ag_ui.core import CustomEvent, EventType, RunAgentInput
-from fastapi import APIRouter, Header
-from fastapi.responses import StreamingResponse
+from ag_ui.core import CustomEvent, EventType
 from pydantic import BaseModel
 
-from pydantic_ai.ag_ui import SSE_CONTENT_TYPE, StateDeps
+from pydantic_ai.ag_ui import FastAGUI, StateDeps
 
-from .agent import AGUIAgent
+from .agent import agent
 
 if TYPE_CHECKING:  # pragma: no cover
     from pydantic_ai import RunContext
@@ -24,18 +22,15 @@
 class DocumentState(BaseModel):
     """State for the document being written."""
 
-    document: str
+    document: str = ''
 
 
-router: APIRouter = APIRouter(prefix='/predictive_state_updates')
-agui: AGUIAgent[StateDeps[DocumentState]] = AGUIAgent(
-    deps_type=StateDeps[DocumentState]
-)
+app: FastAGUI = agent(deps=StateDeps(DocumentState()))
 
 
 # Tools which return AG-UI events will be sent to the client as part of the
 # event stream, single events and iterables of events are supported.
-@agui.agent.tool_plain
+@app.adapter.agent.tool_plain
 def document_predict_state() -> list[CustomEvent]:
     """Enable document state prediction.
 
@@ -58,7 +53,7 @@ def document_predict_state() -> list[CustomEvent]:
     ]
 
 
-@agui.agent.instructions()
+@app.adapter.agent.instructions()
 def story_instructions(ctx: RunContext[StateDeps[DocumentState]]) -> str:
     """Provide instructions for writing document if present.
 
@@ -86,22 +81,3 @@ def story_instructions(ctx: RunContext[StateDeps[DocumentState]]) -> str:
 
 {ctx.deps.state.document}
 """
-
-
-@router.post('')
-async def handler(
-    input_data: RunAgentInput, accept: Annotated[str, Header()] = SSE_CONTENT_TYPE
-) -> StreamingResponse:
-    """Endpoint to handle AG-UI protocol requests and stream responses.
-
-    Args:
-        input_data: The AG-UI run input.
-        accept: The Accept header to specify the response format.
-
-    Returns:
-        A streaming response with event-stream media type.
-    """
-    return StreamingResponse(
-        agui.adapter.run(input_data, accept, deps=StateDeps(state_type=DocumentState)),
-        media_type=SSE_CONTENT_TYPE,
-    )