pydantic_ai.ag_ui

Provides an AG-UI protocol adapter for the Pydantic AI agent.

This package provides seamless integration between pydantic-ai agents and ag-ui for building interactive AI applications with streaming event-based communication.

SSE_CONTENT_TYPE module-attribute

SSE_CONTENT_TYPE: Final[str] = 'text/event-stream'

Content type header value for Server-Sent Events (SSE).

AGUIApp

Bases: Generic[AgentDepsT, OutputDataT], Starlette

ASGI application for running Pydantic AI agents with AG-UI protocol support.

Source code in pydantic_ai_slim/pydantic_ai/ag_ui.py

class AGUIApp(Generic[AgentDepsT, OutputDataT], Starlette):
    """ASGI application for running Pydantic AI agents with AG-UI protocol support."""

    def __init__(
        self,
        agent: Agent[AgentDepsT, OutputDataT],
        *,
        # Agent.iter parameters.
        output_type: OutputSpec[Any] | None = None,
        model: Model | KnownModelName | str | None = None,
        deps: AgentDepsT = None,
        model_settings: ModelSettings | None = None,
        usage_limits: UsageLimits | None = None,
        usage: Usage | None = None,
        infer_name: bool = True,
        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
        # Starlette parameters.
        debug: bool = False,
        routes: Sequence[BaseRoute] | None = None,
        middleware: Sequence[Middleware] | None = None,
        exception_handlers: Mapping[Any, ExceptionHandler] | None = None,
        on_startup: Sequence[Callable[[], Any]] | None = None,
        on_shutdown: Sequence[Callable[[], Any]] | None = None,
        lifespan: Lifespan[AGUIApp[AgentDepsT, OutputDataT]] | None = None,
    ) -> None:
        """An ASGI application that handles every AG-UI request by running the agent.

        Note that the `deps` will be the same for each request, with the exception of the AG-UI state that's
        injected into the `state` field of a `deps` object that implements the [`StateHandler`][pydantic_ai.ag_ui.StateHandler] protocol.
        To provide different `deps` for each request (e.g. based on the authenticated user),
        use [`pydantic_ai.ag_ui.run_ag_ui`][pydantic_ai.ag_ui.run_ag_ui] or
        [`pydantic_ai.ag_ui.handle_ag_ui_request`][pydantic_ai.ag_ui.handle_ag_ui_request] instead.

        Args:
            agent: The agent to run.

            output_type: Custom output type to use for this run, `output_type` may only be used if the agent has
                no output validators since output validators would expect an argument that matches the agent's
                output type.
            model: Optional model to use for this run, required if `model` was not set when creating the agent.
            deps: Optional dependencies to use for this run.
            model_settings: Optional settings to use for this model's request.
            usage_limits: Optional limits on model request count or token usage.
            usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
            infer_name: Whether to try to infer the agent name from the call frame if it's not set.
            toolsets: Optional additional toolsets for this run.

            debug: Boolean indicating if debug tracebacks should be returned on errors.
            routes: A list of routes to serve incoming HTTP and WebSocket requests.
            middleware: A list of middleware to run for every request. A starlette application will always
                automatically include two middleware classes. `ServerErrorMiddleware` is added as the very
                outermost middleware, to handle any uncaught errors occurring anywhere in the entire stack.
                `ExceptionMiddleware` is added as the very innermost middleware, to deal with handled
                exception cases occurring in the routing or endpoints.
            exception_handlers: A mapping of either integer status codes, or exception class types onto
                callables which handle the exceptions. Exception handler callables should be of the form
                `handler(request, exc) -> response` and may be either standard functions, or async functions.
            on_startup: A list of callables to run on application startup. Startup handler callables do not
                take any arguments, and may be either standard functions, or async functions.
            on_shutdown: A list of callables to run on application shutdown. Shutdown handler callables do
                not take any arguments, and may be either standard functions, or async functions.
            lifespan: A lifespan context function, which can be used to perform startup and shutdown tasks.
                This is a newer style that replaces the `on_startup` and `on_shutdown` handlers. Use one or
                the other, not both.
        """
        super().__init__(
            debug=debug,
            routes=routes,
            middleware=middleware,
            exception_handlers=exception_handlers,
            on_startup=on_startup,
            on_shutdown=on_shutdown,
            lifespan=lifespan,
        )

        async def endpoint(request: Request) -> Response:
            """Endpoint to run the agent with the provided input data."""
            return await handle_ag_ui_request(
                agent,
                request,
                output_type=output_type,
                model=model,
                deps=deps,
                model_settings=model_settings,
                usage_limits=usage_limits,
                usage=usage,
                infer_name=infer_name,
                toolsets=toolsets,
            )

        self.router.add_route('/', endpoint, methods=['POST'], name='run_agent')

__init__

__init__(
    agent: Agent[AgentDepsT, OutputDataT],
    *,
    output_type: OutputSpec[Any] | None = None,
    model: Model | KnownModelName | str | None = None,
    deps: AgentDepsT = None,
    model_settings: ModelSettings | None = None,
    usage_limits: UsageLimits | None = None,
    usage: Usage | None = None,
    infer_name: bool = True,
    toolsets: (
        Sequence[AbstractToolset[AgentDepsT]] | None
    ) = None,
    debug: bool = False,
    routes: Sequence[BaseRoute] | None = None,
    middleware: Sequence[Middleware] | None = None,
    exception_handlers: (
        Mapping[Any, ExceptionHandler] | None
    ) = None,
    on_startup: Sequence[Callable[[], Any]] | None = None,
    on_shutdown: Sequence[Callable[[], Any]] | None = None,
    lifespan: (
        Lifespan[AGUIApp[AgentDepsT, OutputDataT]] | None
    ) = None
) -> None

An ASGI application that handles every AG-UI request by running the agent.

Note that the deps will be the same for each request, with the exception of the AG-UI state that's injected into the state field of a deps object that implements the StateHandler protocol. To provide different deps for each request (e.g. based on the authenticated user), use pydantic_ai.ag_ui.run_ag_ui or pydantic_ai.ag_ui.handle_ag_ui_request instead.

Parameters:

Name	Type	Description	Default
agent	Agent[AgentDepsT, OutputDataT]	The agent to run.	required
output_type	OutputSpec[Any] | None	Custom output type to use for this run, output_type may only be used if the agent has no output validators since output validators would expect an argument that matches the agent's output type.	None
model	Model | KnownModelName | str | None	Optional model to use for this run, required if model was not set when creating the agent.	None
deps	AgentDepsT	Optional dependencies to use for this run.	None
model_settings	ModelSettings | None	Optional settings to use for this model's request.	None
usage_limits	UsageLimits | None	Optional limits on model request count or token usage.	None
usage	Usage | None	Optional usage to start with, useful for resuming a conversation or agents used in tools.	None
infer_name	bool	Whether to try to infer the agent name from the call frame if it's not set.	True
toolsets	Sequence[AbstractToolset[AgentDepsT]] | None	Optional additional toolsets for this run.	None
debug	bool	Boolean indicating if debug tracebacks should be returned on errors.	False
routes	Sequence[BaseRoute] | None	A list of routes to serve incoming HTTP and WebSocket requests.	None
middleware	Sequence[Middleware] | None	A list of middleware to run for every request. A starlette application will always automatically include two middleware classes. ServerErrorMiddleware is added as the very outermost middleware, to handle any uncaught errors occurring anywhere in the entire stack. ExceptionMiddleware is added as the very innermost middleware, to deal with handled exception cases occurring in the routing or endpoints.	None
exception_handlers	Mapping[Any, ExceptionHandler] | None	A mapping of either integer status codes, or exception class types onto callables which handle the exceptions. Exception handler callables should be of the form handler(request, exc) -> response and may be either standard functions, or async functions.	None
on_startup	Sequence[Callable[[], Any]] | None	A list of callables to run on application startup. Startup handler callables do not take any arguments, and may be either standard functions, or async functions.	None
on_shutdown	Sequence[Callable[[], Any]] | None	A list of callables to run on application shutdown. Shutdown handler callables do not take any arguments, and may be either standard functions, or async functions.	None
lifespan	Lifespan[AGUIApp[AgentDepsT, OutputDataT]] | None	A lifespan context function, which can be used to perform startup and shutdown tasks. This is a newer style that replaces the on_startup and on_shutdown handlers. Use one or the other, not both.	None

Source code in pydantic_ai_slim/pydantic_ai/ag_ui.py

def __init__(
    self,
    agent: Agent[AgentDepsT, OutputDataT],
    *,
    # Agent.iter parameters.
    output_type: OutputSpec[Any] | None = None,
    model: Model | KnownModelName | str | None = None,
    deps: AgentDepsT = None,
    model_settings: ModelSettings | None = None,
    usage_limits: UsageLimits | None = None,
    usage: Usage | None = None,
    infer_name: bool = True,
    toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
    # Starlette parameters.
    debug: bool = False,
    routes: Sequence[BaseRoute] | None = None,
    middleware: Sequence[Middleware] | None = None,
    exception_handlers: Mapping[Any, ExceptionHandler] | None = None,
    on_startup: Sequence[Callable[[], Any]] | None = None,
    on_shutdown: Sequence[Callable[[], Any]] | None = None,
    lifespan: Lifespan[AGUIApp[AgentDepsT, OutputDataT]] | None = None,
) -> None:
    """An ASGI application that handles every AG-UI request by running the agent.

    Note that the `deps` will be the same for each request, with the exception of the AG-UI state that's
    injected into the `state` field of a `deps` object that implements the [`StateHandler`][pydantic_ai.ag_ui.StateHandler] protocol.
    To provide different `deps` for each request (e.g. based on the authenticated user),
    use [`pydantic_ai.ag_ui.run_ag_ui`][pydantic_ai.ag_ui.run_ag_ui] or
    [`pydantic_ai.ag_ui.handle_ag_ui_request`][pydantic_ai.ag_ui.handle_ag_ui_request] instead.

    Args:
        agent: The agent to run.

        output_type: Custom output type to use for this run, `output_type` may only be used if the agent has
            no output validators since output validators would expect an argument that matches the agent's
            output type.
        model: Optional model to use for this run, required if `model` was not set when creating the agent.
        deps: Optional dependencies to use for this run.
        model_settings: Optional settings to use for this model's request.
        usage_limits: Optional limits on model request count or token usage.
        usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
        infer_name: Whether to try to infer the agent name from the call frame if it's not set.
        toolsets: Optional additional toolsets for this run.

        debug: Boolean indicating if debug tracebacks should be returned on errors.
        routes: A list of routes to serve incoming HTTP and WebSocket requests.
        middleware: A list of middleware to run for every request. A starlette application will always
            automatically include two middleware classes. `ServerErrorMiddleware` is added as the very
            outermost middleware, to handle any uncaught errors occurring anywhere in the entire stack.
            `ExceptionMiddleware` is added as the very innermost middleware, to deal with handled
            exception cases occurring in the routing or endpoints.
        exception_handlers: A mapping of either integer status codes, or exception class types onto
            callables which handle the exceptions. Exception handler callables should be of the form
            `handler(request, exc) -> response` and may be either standard functions, or async functions.
        on_startup: A list of callables to run on application startup. Startup handler callables do not
            take any arguments, and may be either standard functions, or async functions.
        on_shutdown: A list of callables to run on application shutdown. Shutdown handler callables do
            not take any arguments, and may be either standard functions, or async functions.
        lifespan: A lifespan context function, which can be used to perform startup and shutdown tasks.
            This is a newer style that replaces the `on_startup` and `on_shutdown` handlers. Use one or
            the other, not both.
    """
    super().__init__(
        debug=debug,
        routes=routes,
        middleware=middleware,
        exception_handlers=exception_handlers,
        on_startup=on_startup,
        on_shutdown=on_shutdown,
        lifespan=lifespan,
    )

    async def endpoint(request: Request) -> Response:
        """Endpoint to run the agent with the provided input data."""
        return await handle_ag_ui_request(
            agent,
            request,
            output_type=output_type,
            model=model,
            deps=deps,
            model_settings=model_settings,
            usage_limits=usage_limits,
            usage=usage,
            infer_name=infer_name,
            toolsets=toolsets,
        )

    self.router.add_route('/', endpoint, methods=['POST'], name='run_agent')

handle_ag_ui_request async

handle_ag_ui_request(
    agent: Agent[AgentDepsT, Any],
    request: Request,
    *,
    output_type: OutputSpec[Any] | None = None,
    model: Model | KnownModelName | str | None = None,
    deps: AgentDepsT = None,
    model_settings: ModelSettings | None = None,
    usage_limits: UsageLimits | None = None,
    usage: Usage | None = None,
    infer_name: bool = True,
    toolsets: (
        Sequence[AbstractToolset[AgentDepsT]] | None
    ) = None
) -> Response

Handle an AG-UI request by running the agent and returning a streaming response.

Parameters:

Name	Type	Description	Default
agent	Agent[AgentDepsT, Any]	The agent to run.	required
request	Request	The Starlette request (e.g. from FastAPI) containing the AG-UI run input.	required
output_type	OutputSpec[Any] | None	Custom output type to use for this run, output_type may only be used if the agent has no output validators since output validators would expect an argument that matches the agent's output type.	None
model	Model | KnownModelName | str | None	Optional model to use for this run, required if model was not set when creating the agent.	None
deps	AgentDepsT	Optional dependencies to use for this run.	None
model_settings	ModelSettings | None	Optional settings to use for this model's request.	None
usage_limits	UsageLimits | None	Optional limits on model request count or token usage.	None
usage	Usage | None	Optional usage to start with, useful for resuming a conversation or agents used in tools.	None
infer_name	bool	Whether to try to infer the agent name from the call frame if it's not set.	True
toolsets	Sequence[AbstractToolset[AgentDepsT]] | None	Optional additional toolsets for this run.	None

Returns:

Type	Description
Response	A streaming Starlette response with AG-UI protocol events.

Source code in pydantic_ai_slim/pydantic_ai/ag_ui.py

async def handle_ag_ui_request(
    agent: Agent[AgentDepsT, Any],
    request: Request,
    *,
    output_type: OutputSpec[Any] | None = None,
    model: Model | KnownModelName | str | None = None,
    deps: AgentDepsT = None,
    model_settings: ModelSettings | None = None,
    usage_limits: UsageLimits | None = None,
    usage: Usage | None = None,
    infer_name: bool = True,
    toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
) -> Response:
    """Handle an AG-UI request by running the agent and returning a streaming response.

    Args:
        agent: The agent to run.
        request: The Starlette request (e.g. from FastAPI) containing the AG-UI run input.

        output_type: Custom output type to use for this run, `output_type` may only be used if the agent has no
            output validators since output validators would expect an argument that matches the agent's output type.
        model: Optional model to use for this run, required if `model` was not set when creating the agent.
        deps: Optional dependencies to use for this run.
        model_settings: Optional settings to use for this model's request.
        usage_limits: Optional limits on model request count or token usage.
        usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
        infer_name: Whether to try to infer the agent name from the call frame if it's not set.
        toolsets: Optional additional toolsets for this run.

    Returns:
        A streaming Starlette response with AG-UI protocol events.
    """
    accept = request.headers.get('accept', SSE_CONTENT_TYPE)
    try:
        input_data = RunAgentInput.model_validate(await request.json())
    except ValidationError as e:  # pragma: no cover
        return Response(
            content=json.dumps(e.json()),
            media_type='application/json',
            status_code=HTTPStatus.UNPROCESSABLE_ENTITY,
        )

    return StreamingResponse(
        run_ag_ui(
            agent,
            input_data,
            accept,
            output_type=output_type,
            model=model,
            deps=deps,
            model_settings=model_settings,
            usage_limits=usage_limits,
            usage=usage,
            infer_name=infer_name,
            toolsets=toolsets,
        ),
        media_type=accept,
    )

run_ag_ui async

run_ag_ui(
    agent: Agent[AgentDepsT, Any],
    run_input: RunAgentInput,
    accept: str = SSE_CONTENT_TYPE,
    *,
    output_type: OutputSpec[Any] | None = None,
    model: Model | KnownModelName | str | None = None,
    deps: AgentDepsT = None,
    model_settings: ModelSettings | None = None,
    usage_limits: UsageLimits | None = None,
    usage: Usage | None = None,
    infer_name: bool = True,
    toolsets: (
        Sequence[AbstractToolset[AgentDepsT]] | None
    ) = None
) -> AsyncIterator[str]

Run the agent with the AG-UI run input and stream AG-UI protocol events.

Parameters:

Name	Type	Description	Default
agent	Agent[AgentDepsT, Any]	The agent to run.	required
run_input	RunAgentInput	The AG-UI run input containing thread_id, run_id, messages, etc.	required
accept	str	The accept header value for the run.	SSE_CONTENT_TYPE
output_type	OutputSpec[Any] | None	Custom output type to use for this run, output_type may only be used if the agent has no output validators since output validators would expect an argument that matches the agent's output type.	None
model	Model | KnownModelName | str | None	Optional model to use for this run, required if model was not set when creating the agent.	None
deps	AgentDepsT	Optional dependencies to use for this run.	None
model_settings	ModelSettings | None	Optional settings to use for this model's request.	None
usage_limits	UsageLimits | None	Optional limits on model request count or token usage.	None
usage	Usage | None	Optional usage to start with, useful for resuming a conversation or agents used in tools.	None
infer_name	bool	Whether to try to infer the agent name from the call frame if it's not set.	True
toolsets	Sequence[AbstractToolset[AgentDepsT]] | None	Optional additional toolsets for this run.	None

Yields:

Type	Description
AsyncIterator[str]	Streaming event chunks encoded as strings according to the accept header value.

Source code in pydantic_ai_slim/pydantic_ai/ag_ui.py

async def run_ag_ui(
    agent: Agent[AgentDepsT, Any],
    run_input: RunAgentInput,
    accept: str = SSE_CONTENT_TYPE,
    *,
    output_type: OutputSpec[Any] | None = None,
    model: Model | KnownModelName | str | None = None,
    deps: AgentDepsT = None,
    model_settings: ModelSettings | None = None,
    usage_limits: UsageLimits | None = None,
    usage: Usage | None = None,
    infer_name: bool = True,
    toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
) -> AsyncIterator[str]:
    """Run the agent with the AG-UI run input and stream AG-UI protocol events.

    Args:
        agent: The agent to run.
        run_input: The AG-UI run input containing thread_id, run_id, messages, etc.
        accept: The accept header value for the run.

        output_type: Custom output type to use for this run, `output_type` may only be used if the agent has no
            output validators since output validators would expect an argument that matches the agent's output type.
        model: Optional model to use for this run, required if `model` was not set when creating the agent.
        deps: Optional dependencies to use for this run.
        model_settings: Optional settings to use for this model's request.
        usage_limits: Optional limits on model request count or token usage.
        usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
        infer_name: Whether to try to infer the agent name from the call frame if it's not set.
        toolsets: Optional additional toolsets for this run.

    Yields:
        Streaming event chunks encoded as strings according to the accept header value.
    """
    encoder = EventEncoder(accept=accept)
    if run_input.tools:
        # AG-UI tools can't be prefixed as that would result in a mismatch between the tool names in the
        # Pydantic AI events and actual AG-UI tool names, preventing the tool from being called. If any
        # conflicts arise, the AG-UI tool should be renamed or a `PrefixedToolset` used for local toolsets.
        toolset = DeferredToolset[AgentDepsT](
            [
                ToolDefinition(
                    name=tool.name,
                    description=tool.description,
                    parameters_json_schema=tool.parameters,
                )
                for tool in run_input.tools
            ]
        )
        toolsets = [*toolsets, toolset] if toolsets else [toolset]

    try:
        yield encoder.encode(
            RunStartedEvent(
                thread_id=run_input.thread_id,
                run_id=run_input.run_id,
            ),
        )

        if not run_input.messages:
            raise _NoMessagesError

        raw_state: dict[str, Any] = run_input.state or {}
        if isinstance(deps, StateHandler):
            if isinstance(deps.state, BaseModel):
                try:
                    state = type(deps.state).model_validate(raw_state)
                except ValidationError as e:  # pragma: no cover
                    raise _InvalidStateError from e
            else:
                state = raw_state

            deps = replace(deps, state=state)
        elif raw_state:
            raise UserError(
                f'AG-UI state is provided but `deps` of type `{type(deps).__name__}` does not implement the `StateHandler` protocol: it needs to be a dataclass with a non-optional `state` field.'
            )
        else:
            # `deps` not being a `StateHandler` is OK if there is no state.
            pass

        messages = _messages_from_ag_ui(run_input.messages)

        async with agent.iter(
            user_prompt=None,
            output_type=[output_type or agent.output_type, DeferredToolCalls],
            message_history=messages,
            model=model,
            deps=deps,
            model_settings=model_settings,
            usage_limits=usage_limits,
            usage=usage,
            infer_name=infer_name,
            toolsets=toolsets,
        ) as run:
            async for event in _agent_stream(run):
                yield encoder.encode(event)
    except _RunError as e:
        yield encoder.encode(
            RunErrorEvent(message=e.message, code=e.code),
        )
    except Exception as e:
        yield encoder.encode(
            RunErrorEvent(message=str(e)),
        )
        raise e
    else:
        yield encoder.encode(
            RunFinishedEvent(
                thread_id=run_input.thread_id,
                run_id=run_input.run_id,
            ),
        )

StateHandler

Bases: Protocol

Protocol for state handlers in agent runs. Requires the class to be a dataclass with a state field.

Source code in pydantic_ai_slim/pydantic_ai/ag_ui.py

@runtime_checkable
class StateHandler(Protocol):
    """Protocol for state handlers in agent runs. Requires the class to be a dataclass with a `state` field."""

    # Has to be a dataclass so we can use `replace` to update the state.
    # From https://github.com/python/typeshed/blob/9ab7fde0a0cd24ed7a72837fcb21093b811b80d8/stdlib/_typeshed/__init__.pyi#L352
    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]

    @property
    def state(self) -> State:
        """Get the current state of the agent run."""
        ...

    @state.setter
    def state(self, state: State) -> None:
        """Set the state of the agent run.

        This method is called to update the state of the agent run with the
        provided state.

        Args:
            state: The run state.

        Raises:
            InvalidStateError: If `state` does not match the expected model.
        """
        ...

state property writable

state: State

Get the current state of the agent run.

StateDeps dataclass

Bases: Generic[StateT]

Provides AG-UI state management.

This class is used to manage the state of an agent run. It allows setting the state of the agent run with a specific type of state model, which must be a subclass of BaseModel.

The state is set using the state setter by the Adapter when the run starts.

Implements the StateHandler protocol.

Source code in pydantic_ai_slim/pydantic_ai/ag_ui.py

@dataclass
class StateDeps(Generic[StateT]):
    """Provides AG-UI state management.

    This class is used to manage the state of an agent run. It allows setting
    the state of the agent run with a specific type of state model, which must
    be a subclass of `BaseModel`.

    The state is set using the `state` setter by the `Adapter` when the run starts.

    Implements the `StateHandler` protocol.
    """

    state: StateT