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[OutputDataT] | 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:
        """Initialise the AG-UI application.

        Args:
            agent: The Pydantic AI `Agent` to adapt.

            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 list of toolsets to use for this agent, defaults to the agent's toolset.

            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,
        )
        adapter = _Adapter(agent=agent)

        async def endpoint(request: Request) -> Response | StreamingResponse:
            """Endpoint to run the agent with the provided input data."""
            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(
                adapter.run(
                    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=SSE_CONTENT_TYPE,
            )

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

__init__

__init__(
    agent: Agent[AgentDepsT, OutputDataT],
    *,
    output_type: OutputSpec[OutputDataT] | 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

Initialise the AG-UI application.

Parameters:

Name	Type	Description	Default
agent	Agent[AgentDepsT, OutputDataT]	The Pydantic AI Agent to adapt.	required
output_type	OutputSpec[OutputDataT] | 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 list of toolsets to use for this agent, defaults to the agent's toolset.	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[OutputDataT] | 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:
    """Initialise the AG-UI application.

    Args:
        agent: The Pydantic AI `Agent` to adapt.

        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 list of toolsets to use for this agent, defaults to the agent's toolset.

        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,
    )
    adapter = _Adapter(agent=agent)

    async def endpoint(request: Request) -> Response | StreamingResponse:
        """Endpoint to run the agent with the provided input data."""
        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(
            adapter.run(
                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=SSE_CONTENT_TYPE,
        )

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

StateHandler

Bases: Protocol

Protocol for state handlers in agent runs.

Source code in pydantic_ai_slim/pydantic_ai/ag_ui.py

@runtime_checkable
class StateHandler(Protocol):
    """Protocol for state handlers in agent runs."""

    @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

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

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.
    """

    def __init__(self, default: StateT) -> None:
        """Initialize the state with the provided state type."""
        self._state = default

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

        Returns:
            The current run state.
        """
        return self._state

    @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.

        Implements the `StateHandler` protocol.

        Args:
            state: The run state, which must be `None` or model validate for the state type.

        Raises:
            InvalidStateError: If `state` does not validate.
        """
        if state is None:
            # If state is None, we keep the current state, which will be the default state.
            return

        try:
            self._state = type(self._state).model_validate(state)
        except ValidationError as e:  # pragma: no cover
            raise _InvalidStateError from e

__init__

__init__(default: StateT) -> None

Initialize the state with the provided state type.

Source code in pydantic_ai_slim/pydantic_ai/ag_ui.py

def __init__(self, default: StateT) -> None:
    """Initialize the state with the provided state type."""
    self._state = default

state property writable

state: StateT

Get the current state of the agent run.

Returns:

Type	Description
StateT	The current run state.