AG-UI docs

Author: DouweM

docs/ag-ui.md

@@ -33,27 +33,26 @@ pip/uv-add uvicorn
 
 There are three ways to run a Pydantic AI agent based on AG-UI run input with streamed AG-UI events as output, from most to least flexible. If you're using a Starlette-based web framework like FastAPI, you'll typically want to use the second method.
 
-1. [`run_ag_ui()`][pydantic_ai.ag_ui.run_ag_ui] takes an agent and an AG-UI [`RunAgentInput`](https://docs.ag-ui.com/sdk/python/core/types#runagentinput) object, and returns a stream of AG-UI events encoded as strings. It also takes optional [`Agent.iter()`][pydantic_ai.Agent.iter] arguments including `deps`. Use this if you're using a web framework not based on Starlette (e.g. Django or Flask) or want to modify the input or output some way.
-2. [`handle_ag_ui_request()`][pydantic_ai.ag_ui.handle_ag_ui_request] takes an agent and a Starlette request (e.g. from FastAPI) coming from an AG-UI frontend, and returns a streaming Starlette response of AG-UI events that you can return directly from your endpoint. It also takes optional [`Agent.iter()`][pydantic_ai.Agent.iter] arguments including `deps`, that you can vary for each request (e.g. based on the authenticated user).
-3. [`Agent.to_ag_ui()`][pydantic_ai.agent.AbstractAgent.to_ag_ui] returns an ASGI application that handles every AG-UI request by running the agent. It also takes optional [`Agent.iter()`][pydantic_ai.Agent.iter] arguments including `deps`, but these will be the same for each request, with the exception of the AG-UI state that's injected as described under [state management](#state-management). This ASGI app can be [mounted](https://fastapi.tiangolo.com/advanced/sub-applications/) at a given path in an existing FastAPI app.
+1. The [`AGUIAdapter.run_stream()`][pydantic_ai.ui.ag_ui.AGUIAdapter.run_stream] method, when called on an [`AGUIAdapter`][pydantic_ai.ui.ag_ui.AGUIAdapter] instantiated with an agent and an AG-UI [`RunAgentInput`](https://docs.ag-ui.com/sdk/python/core/types#runagentinput) object, will run the agent and return a stream of AG-UI events. It also takes optional [`Agent.iter()`][pydantic_ai.Agent.iter] arguments including `deps`. Use this if you're using a web framework not based on Starlette (e.g. Django or Flask) or want to modify the input or output some way.
+2. The [`AGUIAdapter.dispatch_request()`][pydantic_ai.ui.ag_ui.AGUIAdapter.dispatch_request] class method takes an agent and a Starlette request (e.g. from FastAPI) coming from an AG-UI frontend, and returns a streaming Starlette response of AG-UI events that you can return directly from your endpoint. It also takes optional [`Agent.iter()`][pydantic_ai.Agent.iter] arguments including `deps`, that you can vary for each request (e.g. based on the authenticated user). This is a convenience method that combines [`AGUIAdapter.from_request()`][pydantic_ai.ui.ag_ui.AGUIAdapter.from_request], [`AGUIAdapter.run_stream()`][pydantic_ai.ui.ag_ui.AGUIAdapter.run_stream], and [`AGUIAdapter.streaming_response()`][pydantic_ai.ui.ag_ui.AGUIAdapter.streaming_response].
+3. [`AGUIApp`][pydantic_ai.ui.ag_ui.app.AGUIApp] represents an ASGI application that handles every AG-UI request by running the agent. It also takes optional [`Agent.iter()`][pydantic_ai.Agent.iter] arguments including `deps`, but these will be the same for each request, with the exception of the AG-UI state that's injected as described under [state management](#state-management). This ASGI app can be [mounted](https://fastapi.tiangolo.com/advanced/sub-applications/) at a given path in an existing FastAPI app.
 
 ### Handle run input and output directly
 
-This example uses [`run_ag_ui()`][pydantic_ai.ag_ui.run_ag_ui] and performs its own request parsing and response generation.
+This example uses [`AGUIAdapter.run_stream()`][pydantic_ai.ui.ag_ui.AGUIAdapter.run_stream] and performs its own request parsing and response generation.
 This can be modified to work with any web framework.
 
 ```py {title="run_ag_ui.py"}
 import json
 from http import HTTPStatus
 
-from ag_ui.core import RunAgentInput
 from fastapi import FastAPI
 from fastapi.requests import Request
 from fastapi.responses import Response, StreamingResponse
 from pydantic import ValidationError
 
 from pydantic_ai import Agent
-from pydantic_ai.ag_ui import SSE_CONTENT_TYPE, run_ag_ui
+from pydantic_ai.ui.ag_ui import AGUIAdapter, SSE_CONTENT_TYPE
 
 agent = Agent('openai:gpt-4.1', instructions='Be fun!')
 
@@ -64,19 +63,24 @@ app = FastAPI()
 async def run_agent(request: Request) -> Response:
     accept = request.headers.get('accept', SSE_CONTENT_TYPE)
     try:
-        run_input = RunAgentInput.model_validate(await request.json())
+        run_input = AGUIAdapter.build_run_input(await request.json())  # (1)
     except ValidationError as e:  # pragma: no cover
         return Response(
             content=json.dumps(e.json()),
             media_type='application/json',
             status_code=HTTPStatus.UNPROCESSABLE_ENTITY,
         )
 
-    event_stream = run_ag_ui(agent, run_input, accept=accept)
+    adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
+    events = adapter.run_stream() # (2)
 
-    return StreamingResponse(event_stream, media_type=accept)
+    return StreamingResponse(adapter.encode_stream(events), media_type=accept) # (3)
 ```
 
+1. You can also use the [`AGUIAdapter.from_request()`][pydantic_ai.ui.ag_ui.AGUIAdapter.from_request] class method to build an adapter directly from a request.
+2. You can also use the [`AGUIAdapter.run_stream_native()`][pydantic_ai.ui.ag_ui.AGUIAdapter.run_stream_native] method to run the agent and return a stream of Pydantic AI events instead of AG-UI events. These can then be transformed into AG-UI events using the [`AGUIAdapter.transform_stream()`][pydantic_ai.ui.ag_ui.AGUIAdapter.transform_stream] method.
+3. The [`AGUIAdapter.encode_stream()`][pydantic_ai.ui.ag_ui.AGUIAdapter.encode_stream] method encodes the stream of AG-UI events as strings according to the accept header value. You can also use the [`AGUIAdapter.streaming_response()`][pydantic_ai.ui.ag_ui.AGUIAdapter.streaming_response] method to generate a streaming response directly from the AG-UI event stream returned by `run_stream()`.
+
 Since `app` is an ASGI application, it can be used with any ASGI server:
 
 ```shell
@@ -87,25 +91,27 @@ This will expose the agent as an AG-UI server, and your frontend can start sendi
 
 ### Handle a Starlette request
 
-This example uses [`handle_ag_ui_request()`][pydantic_ai.ag_ui.run_ag_ui] to directly handle a FastAPI request and return a response. Something analogous to this will work with any Starlette-based web framework.
+This example uses [`AGUIAdapter.dispatch_request()`][pydantic_ai.ui.ag_ui.AGUIAdapter.dispatch_request] to directly handle a FastAPI request and return a response. Something analogous to this will work with any Starlette-based web framework.
 
 ```py {title="handle_ag_ui_request.py"}
 from fastapi import FastAPI
 from starlette.requests import Request
 from starlette.responses import Response
 
 from pydantic_ai import Agent
-from pydantic_ai.ag_ui import handle_ag_ui_request
+from pydantic_ai.ui.ag_ui import AGUIAdapter
 
 agent = Agent('openai:gpt-4.1', instructions='Be fun!')
 
 app = FastAPI()
 
 @app.post('/')
 async def run_agent(request: Request) -> Response:
-    return await handle_ag_ui_request(agent, request)
+    return await AGUIAdapter.dispatch_request(request, agent=agent) # (1)
 ```
 
+1. This method essentially does the same as the previous example, but it's more convenient to use when you're already using a Starlette/FastAPI app.
+
 Since `app` is an ASGI application, it can be used with any ASGI server:
 
 ```shell
@@ -116,19 +122,20 @@ This will expose the agent as an AG-UI server, and your frontend can start sendi
 
 ### Stand-alone ASGI app
 
-This example uses [`Agent.to_ag_ui()`][pydantic_ai.agent.AbstractAgent.to_ag_ui] to turn the agent into a stand-alone ASGI application:
+This example uses [`AGUIApp`][pydantic_ai.ui.ag_ui.app.AGUIApp] to turn the agent into a stand-alone ASGI application:
 
-```py {title="agent_to_ag_ui.py" hl_lines="4"}
+```py {title="ag_ui_app.py" hl_lines="4"}
 from pydantic_ai import Agent
+from pydantic_ai.ui.ag_ui.app import AGUIApp
 
 agent = Agent('openai:gpt-4.1', instructions='Be fun!')
-app = agent.to_ag_ui()
+app = AGUIApp(agent)
 ```
 
 Since `app` is an ASGI application, it can be used with any ASGI server:
 
 ```shell
-uvicorn agent_to_ag_ui:app
+uvicorn ag_ui_app:app
 ```
 
 This will expose the agent as an AG-UI server, and your frontend can start sending requests to it.
@@ -174,7 +181,8 @@ validate state contained in [`RunAgentInput.state`](https://docs.ag-ui.com/sdk/j
 from pydantic import BaseModel
 
 from pydantic_ai import Agent
-from pydantic_ai.ag_ui import StateDeps
+from pydantic_ai.ui import StateDeps
+from pydantic_ai.ui.ag_ui import AGUIApp
 
 
 class DocumentState(BaseModel):
@@ -188,7 +196,7 @@ agent = Agent(
     instructions='Be fun!',
     deps_type=StateDeps[DocumentState],
 )
-app = agent.to_ag_ui(deps=StateDeps(DocumentState()))
+app = AGUIApp(agent, deps=StateDeps(DocumentState()))
 ```
 
 Since `app` is an ASGI application, it can be used with any ASGI server:
@@ -214,7 +222,8 @@ from ag_ui.core import CustomEvent, EventType, StateSnapshotEvent
 from pydantic import BaseModel
 
 from pydantic_ai import Agent, RunContext, ToolReturn
-from pydantic_ai.ag_ui import StateDeps
+from pydantic_ai.ui import StateDeps
+from pydantic_ai.ui.ag_ui import AGUIApp
 
 
 class DocumentState(BaseModel):
@@ -228,7 +237,7 @@ agent = Agent(
     instructions='Be fun!',
     deps_type=StateDeps[DocumentState],
 )
-app = agent.to_ag_ui(deps=StateDeps(DocumentState()))
+app = AGUIApp(agent, deps=StateDeps(DocumentState()))
 
 
 @agent.tool
@@ -271,7 +280,7 @@ uvicorn ag_ui_tool_events:app --host 0.0.0.0 --port 9000
 
 ## Examples
 
-For more examples of how to use [`to_ag_ui()`][pydantic_ai.agent.AbstractAgent.to_ag_ui] see
+For more examples of how to use [`AGUIApp`][pydantic_ai.ui.ag_ui.app.AGUIApp] see
 [`pydantic_ai_examples.ag_ui`](https://github.com/pydantic/pydantic-ai/tree/main/examples/pydantic_ai_examples/ag_ui),
 which includes a server for use with the
 [AG-UI Dojo](https://docs.ag-ui.com/tutorials/debugging#the-ag-ui-dojo).

examples/pydantic_ai_examples/ag_ui/api/agentic_chat.py

@@ -6,9 +6,9 @@
 from zoneinfo import ZoneInfo
 
 from pydantic_ai import Agent
+from pydantic_ai.ui.ag_ui.app import AGUIApp
 
 agent = Agent('openai:gpt-4o-mini')
-app = agent.to_ag_ui()
 
 
 @agent.tool_plain
@@ -23,3 +23,6 @@ async def current_time(timezone: str = 'UTC') -> str:
     """
     tz: ZoneInfo = ZoneInfo(timezone)
     return datetime.now(tz=tz).isoformat()
+
+
+app = AGUIApp(agent)

examples/pydantic_ai_examples/ag_ui/api/agentic_generative_ui.py

@@ -9,6 +9,7 @@
 
 from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
 from pydantic_ai import Agent
+from pydantic_ai.ui.ag_ui.app import AGUIApp
 
 StepStatus = Literal['pending', 'completed']
 
@@ -116,4 +117,4 @@ async def update_plan_step(
     )
 
 
-app = agent.to_ag_ui()
+app = AGUIApp(agent)

examples/pydantic_ai_examples/ag_ui/api/human_in_the_loop.py

@@ -8,6 +8,7 @@
 from textwrap import dedent
 
 from pydantic_ai import Agent
+from pydantic_ai.ui.ag_ui.app import AGUIApp
 
 agent = Agent(
     'openai:gpt-4o-mini',
@@ -23,4 +24,4 @@
     ),
 )
 
-app = agent.to_ag_ui()
+app = AGUIApp(agent)

examples/pydantic_ai_examples/ag_ui/api/predictive_state_updates.py

@@ -8,7 +8,8 @@
 
 from ag_ui.core import CustomEvent, EventType
 from pydantic_ai import Agent, RunContext
-from pydantic_ai.ag_ui import StateDeps
+from pydantic_ai.ui import StateDeps
+from pydantic_ai.ui.ag_ui.app import AGUIApp
 
 
 class DocumentState(BaseModel):
@@ -74,4 +75,4 @@ async def story_instructions(ctx: RunContext[StateDeps[DocumentState]]) -> str:
     )
 
 
-app = agent.to_ag_ui(deps=StateDeps(DocumentState()))
+app = AGUIApp(agent, deps=StateDeps(DocumentState()))

examples/pydantic_ai_examples/ag_ui/api/shared_state.py

@@ -9,7 +9,8 @@
 
 from ag_ui.core import EventType, StateSnapshotEvent
 from pydantic_ai import Agent, RunContext
-from pydantic_ai.ag_ui import StateDeps
+from pydantic_ai.ui import StateDeps
+from pydantic_ai.ui.ag_ui.app import AGUIApp
 
 
 class SkillLevel(str, Enum):
@@ -135,4 +136,4 @@ async def recipe_instructions(ctx: RunContext[StateDeps[RecipeSnapshot]]) -> str
     )
 
 
-app = agent.to_ag_ui(deps=StateDeps(RecipeSnapshot()))
+app = AGUIApp(agent, deps=StateDeps(RecipeSnapshot()))

examples/pydantic_ai_examples/ag_ui/api/tool_based_generative_ui.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 from pydantic_ai import Agent
+from pydantic_ai.ui.ag_ui.app import AGUIApp
 
 agent = Agent('openai:gpt-4o-mini')
-app = agent.to_ag_ui()
+app = AGUIApp(agent)

pydantic_ai_slim/pydantic_ai/ag_ui.py

@@ -4,6 +4,8 @@
 for building interactive AI applications with streaming event-based communication.
 """
 
+# TODO (v2): Remove this module in favor of `pydantic_ai.ui.ag_ui`
+
 from __future__ import annotations
 
 from collections.abc import AsyncIterator, Sequence
@@ -22,24 +24,15 @@
 try:
     from ag_ui.core import BaseEvent
     from ag_ui.core.types import RunAgentInput
-
-    from .ui import SSE_CONTENT_TYPE, OnCompleteFunc, StateDeps, StateHandler
-    from .ui.ag_ui import (
-        AGUIAdapter,
-        AGUIApp,
-    )
-except ImportError as e:  # pragma: no cover
-    raise ImportError(
-        'Please install the `ag-ui-protocol` package to use `Agent.to_ag_ui()` method, '
-        'you can use the `ag-ui` optional group — `pip install "pydantic-ai-slim[ag-ui]"`'
-    ) from e
-
-try:
     from starlette.requests import Request
     from starlette.responses import Response
+
+    from .ui import SSE_CONTENT_TYPE, OnCompleteFunc, StateDeps, StateHandler
+    from .ui.ag_ui import AGUIAdapter
+    from .ui.ag_ui.app import AGUIApp
 except ImportError as e:  # pragma: no cover
     raise ImportError(
-        'Please install the `starlette` package to use `Agent.to_ag_ui()` method, '
+        'Please install the `ag-ui-protocol` and `starlette` packages to use `AGUIAdapter`, '
         'you can use the `ag-ui` optional group — `pip install "pydantic-ai-slim[ag-ui]"`'
     ) from e
 

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -49,7 +49,7 @@
     from starlette.routing import BaseRoute, Route
     from starlette.types import ExceptionHandler, Lifespan
 
-    from ..ag_ui import AGUIApp
+    from pydantic_ai.ui.ag_ui.app import AGUIApp
 
 
 T = TypeVar('T')
@@ -1077,7 +1077,7 @@ def to_ag_ui(
         Returns:
             An ASGI application for running Pydantic AI agents with AG-UI protocol support.
         """
-        from ..ag_ui import AGUIApp
+        from pydantic_ai.ui.ag_ui.app import AGUIApp
 
         return AGUIApp(
             agent=self,

pydantic_ai_slim/pydantic_ai/ui/__init__.py

@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 from .adapter import OnCompleteFunc, StateDeps, StateHandler, UIAdapter
-from .app import UIApp
 from .event_stream import SSE_CONTENT_TYPE, UIEventStream
 from .messages_builder import MessagesBuilder
 
@@ -18,6 +17,5 @@
     'StateDeps',
     'StateHandler',
     'OnCompleteFunc',
-    'UIApp',
     'MessagesBuilder',
 ]