Add UIApp, AGUIApp, VercelAIApp

Author: DouweM

pydantic_ai_slim/pydantic_ai/ag_ui.py

@@ -6,14 +6,14 @@
 
 from __future__ import annotations
 
-from collections.abc import AsyncIterator, Callable, Mapping, Sequence
-from typing import Any, Generic
+from collections.abc import AsyncIterator, Sequence
+from typing import Any
 
 from . import DeferredToolResults
 from .agent import AbstractAgent
 from .messages import ModelMessage
 from .models import KnownModelName, Model
-from .output import OutputDataT, OutputSpec
+from .output import OutputSpec
 from .settings import ModelSettings
 from .tools import AgentDepsT
 from .toolsets import AbstractToolset
@@ -27,6 +27,7 @@
     from .ui.ag_ui import (
         SSE_CONTENT_TYPE,
         AGUIAdapter,
+        AGUIApp,
     )
 except ImportError as e:  # pragma: no cover
     raise ImportError(
@@ -35,12 +36,8 @@
     ) from e
 
 try:
-    from starlette.applications import Starlette
-    from starlette.middleware import Middleware
     from starlette.requests import Request
     from starlette.responses import Response
-    from starlette.routing import BaseRoute
-    from starlette.types import ExceptionHandler, Lifespan
 except ImportError as e:  # pragma: no cover
     raise ImportError(
         'Please install the `starlette` package to use `Agent.to_ag_ui()` method, '
@@ -59,105 +56,6 @@
 ]
 
 
-class AGUIApp(Generic[AgentDepsT, OutputDataT], Starlette):
-    """ASGI application for running Pydantic AI agents with AG-UI protocol support."""
-
-    def __init__(
-        self,
-        agent: AbstractAgent[AgentDepsT, OutputDataT],
-        *,
-        # Agent.iter parameters.
-        output_type: OutputSpec[Any] | None = None,
-        message_history: Sequence[ModelMessage] | None = None,
-        deferred_tool_results: DeferredToolResults | None = None,
-        model: Model | KnownModelName | str | None = None,
-        deps: AgentDepsT = None,
-        model_settings: ModelSettings | None = None,
-        usage_limits: UsageLimits | None = None,
-        usage: RunUsage | 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.
-            message_history: History of the conversation so far.
-            deferred_tool_results: Optional results for deferred tool calls in the message history.
-            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,
-                message_history=message_history,
-                deferred_tool_results=deferred_tool_results,
-                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')
-
-
 async def handle_ag_ui_request(
     agent: AbstractAgent[AgentDepsT, Any],
     request: Request,
@@ -202,6 +100,8 @@ async def handle_ag_ui_request(
         request,
         deps=deps,
         output_type=output_type,
+        message_history=message_history,
+        deferred_tool_results=deferred_tool_results,
         model=model,
         model_settings=model_settings,
         usage_limits=usage_limits,

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -993,6 +993,7 @@ async def __aenter__(self) -> AbstractAgent[AgentDepsT, OutputDataT]:
     async def __aexit__(self, *args: Any) -> bool | None:
         raise NotImplementedError
 
+    # TODO (v2): Remove in favor of using `AGUIApp` directly -- we don't have `to_temporal()` or `to_vercel_ai()` either.
     def to_ag_ui(
         self,
         *,

pydantic_ai_slim/pydantic_ai/ui/__init__.py

@@ -6,14 +6,16 @@
 
 from __future__ import annotations
 
-from .adapter import BaseAdapter, OnCompleteFunc, StateDeps, StateHandler
-from .event_stream import SSE_CONTENT_TYPE, BaseEventStream
+from .adapter import OnCompleteFunc, StateDeps, StateHandler, UIAdapter
+from .app import UIApp
+from .event_stream import SSE_CONTENT_TYPE, UIEventStream
 
 __all__ = [
-    'BaseAdapter',
-    'BaseEventStream',
+    'UIAdapter',
+    'UIEventStream',
     'SSE_CONTENT_TYPE',
     'StateDeps',
     'StateHandler',
     'OnCompleteFunc',
+    'UIApp',
 ]

pydantic_ai_slim/pydantic_ai/ui/adapter.py

@@ -33,15 +33,15 @@
 from ..settings import ModelSettings
 from ..toolsets import AbstractToolset
 from ..usage import RunUsage, UsageLimits
-from .event_stream import BaseEventStream, OnCompleteFunc, SourceEvent
+from .event_stream import OnCompleteFunc, SourceEvent, UIEventStream
 
 if TYPE_CHECKING:
     from starlette.requests import Request
     from starlette.responses import Response
 
 
 __all__ = [
-    'BaseAdapter',
+    'UIAdapter',
 ]
 
 
@@ -107,7 +107,7 @@ class StateDeps(Generic[StateT]):
 
 
 @dataclass
-class BaseAdapter(ABC, Generic[RunRequestT, MessageT, EventT, AgentDepsT, OutputDataT]):
+class UIAdapter(ABC, Generic[RunRequestT, MessageT, EventT, AgentDepsT, OutputDataT]):
     """TODO (DouwM): Docstring."""
 
     agent: AbstractAgent[AgentDepsT, OutputDataT]
@@ -131,7 +131,7 @@ def load_messages(cls, messages: Sequence[MessageT]) -> list[ModelMessage]:
     @abstractmethod
     def build_event_stream(
         self, accept: str | None = None
-    ) -> BaseEventStream[RunRequestT, EventT, AgentDepsT, OutputDataT]:
+    ) -> UIEventStream[RunRequestT, EventT, AgentDepsT, OutputDataT]:
         """Create an event stream for the adapter.
 
         Args:
@@ -235,7 +235,7 @@ async def run_stream(
         toolset = self.toolset
         if toolset:
             output_type = [output_type or self.agent.output_type, DeferredToolRequests]
-            toolsets = [*toolsets, toolset] if toolsets else [toolset]
+            toolsets = [*(toolsets or []), toolset]
 
         if isinstance(deps, StateHandler):
             raw_state = self.state or {}

pydantic_ai_slim/pydantic_ai/ui/ag_ui/__init__.py

@@ -1,10 +1,25 @@
 """AG-UI protocol integration for Pydantic AI agents."""
 
+from typing import Any
+
+from pydantic_ai.agent import AbstractAgent
+from pydantic_ai.output import OutputDataT
+from pydantic_ai.tools import AgentDepsT
+
+from .. import UIApp
 from ._adapter import AGUIAdapter
 from ._event_stream import SSE_CONTENT_TYPE, AGUIEventStream
 
 __all__ = [
     'AGUIAdapter',
     'AGUIEventStream',
     'SSE_CONTENT_TYPE',
+    'AGUIApp',
 ]
+
+
+class AGUIApp(UIApp[AgentDepsT, OutputDataT]):
+    """ASGI application for running Pydantic AI agents with AG-UI protocol support."""
+
+    def __init__(self, agent: AbstractAgent[AgentDepsT, OutputDataT], **kwargs: Any) -> None:
+        super().__init__(AGUIAdapter[AgentDepsT, OutputDataT], agent, **kwargs)

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_adapter.py

@@ -41,8 +41,8 @@
         UserMessage,
     )
 
-    from ..adapter import BaseAdapter
-    from ..event_stream import BaseEventStream
+    from ..adapter import UIAdapter
+    from ..event_stream import UIEventStream
     from ._event_stream import BUILTIN_TOOL_CALL_ID_PREFIX, AGUIEventStream
 except ImportError as e:  # pragma: no cover
     raise ImportError(
@@ -93,7 +93,7 @@ def label(self) -> str:
         return 'the AG-UI frontend tools'  # pragma: no cover
 
 
-class AGUIAdapter(BaseAdapter[RunAgentInput, Message, BaseEvent, AgentDepsT, OutputDataT]):
+class AGUIAdapter(UIAdapter[RunAgentInput, Message, BaseEvent, AgentDepsT, OutputDataT]):
     """TODO (DouwM): Docstring."""
 
     @classmethod
@@ -103,7 +103,7 @@ async def validate_request(cls, request: Request) -> RunAgentInput:
 
     def build_event_stream(
         self, accept: str | None = None
-    ) -> BaseEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]:
+    ) -> UIEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]:
         """Create an event stream for the adapter.
 
         Args:

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_event_stream.py

@@ -26,7 +26,7 @@
 )
 from ...output import OutputDataT
 from ...tools import AgentDepsT
-from .. import SSE_CONTENT_TYPE, BaseEventStream
+from .. import SSE_CONTENT_TYPE, UIEventStream
 
 try:
     from ag_ui.core import (
@@ -68,7 +68,7 @@
 
 
 @dataclass
-class AGUIEventStream(BaseEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]):
+class AGUIEventStream(UIEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]):
     """TODO (DouwM): Docstring."""
 
     _thinking_text: bool = False