Remove UIApp

Author: DouweM

pydantic_ai_slim/pydantic_ai/ui/__init__.py

@@ -6,16 +6,17 @@
 
 from __future__ import annotations
 
-from .adapter import OnCompleteFunc, StateDeps, StateHandler, UIAdapter
-from .event_stream import SSE_CONTENT_TYPE, UIEventStream
-from .messages_builder import MessagesBuilder
+from ._adapter import StateDeps, StateHandler, UIAdapter
+from ._event_stream import SSE_CONTENT_TYPE, NativeEvent, OnCompleteFunc, UIEventStream
+from ._messages_builder import MessagesBuilder
 
 __all__ = [
     'UIAdapter',
     'UIEventStream',
     'SSE_CONTENT_TYPE',
     'StateDeps',
     'StateHandler',
+    'NativeEvent',
     'OnCompleteFunc',
     'MessagesBuilder',
 ]

pydantic_ai_slim/pydantic_ai/ui/adapter.py renamed to pydantic_ai_slim/pydantic_ai/ui/_adapter.py

@@ -35,7 +35,7 @@
 from pydantic_ai.toolsets import AbstractToolset
 from pydantic_ai.usage import RunUsage, UsageLimits
 
-from .event_stream import NativeEvent, OnCompleteFunc, UIEventStream
+from ._event_stream import NativeEvent, OnCompleteFunc, UIEventStream
 
 if TYPE_CHECKING:
     from starlette.requests import Request

pydantic_ai_slim/pydantic_ai/ui/event_stream.py renamed to pydantic_ai_slim/pydantic_ai/ui/_event_stream.py

@@ -57,7 +57,8 @@
 RunInputT = TypeVar('RunInputT')
 """Type variable for protocol-specific run input types."""
 
-NativeEvent = AgentStreamEvent | AgentRunResultEvent[Any]
+NativeEvent: TypeAlias = AgentStreamEvent | AgentRunResultEvent[Any]
+"""Type alias for the native event type, which is either an `AgentStreamEvent` or an `AgentRunResultEvent`."""
 
 OnCompleteFunc: TypeAlias = (
     Callable[[AgentRunResult[Any]], None]

pydantic_ai_slim/pydantic_ai/ui/messages_builder.py renamed to pydantic_ai_slim/pydantic_ai/ui/_messages_builder.py

@@ -7,7 +7,7 @@
 
 @dataclass
 class MessagesBuilder:
-    """Helper class to build Pydantic AI messages from protocol-specific messages."""
+    """Helper class to build Pydantic AI messages from request/response parts."""
 
     messages: list[ModelMessage] = field(default_factory=list)
 

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_adapter.py

@@ -37,9 +37,7 @@
         UserMessage,
     )
 
-    from ..adapter import UIAdapter
-    from ..event_stream import UIEventStream
-    from ..messages_builder import MessagesBuilder
+    from .. import MessagesBuilder, UIAdapter, UIEventStream
     from ._event_stream import BUILTIN_TOOL_CALL_ID_PREFIX, AGUIEventStream
 except ImportError as e:  # pragma: no cover
     raise ImportError(

pydantic_ai_slim/pydantic_ai/ui/ag_ui/app.py

@@ -1,17 +1,141 @@
 """AG-UI protocol integration for Pydantic AI agents."""
 
-from typing import Any
+from __future__ import annotations
 
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any, Generic
+
+from typing_extensions import Self
+
+from pydantic_ai import DeferredToolResults
 from pydantic_ai.agent import AbstractAgent
-from pydantic_ai.output import OutputDataT
+from pydantic_ai.builtin_tools import AbstractBuiltinTool
+from pydantic_ai.messages import ModelMessage
+from pydantic_ai.models import KnownModelName, Model
+from pydantic_ai.output import OutputDataT, OutputSpec
+from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import AgentDepsT
+from pydantic_ai.toolsets import AbstractToolset
+from pydantic_ai.usage import RunUsage, UsageLimits
 
-from ..app import UIApp
+from .. import OnCompleteFunc
 from ._adapter import AGUIAdapter
 
+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 `AGUIApp`, '
+        'you can use the `ag-ui` optional group — `pip install "pydantic-ai-slim[ag-ui]"`'
+    ) from e
+
 
-class AGUIApp(UIApp[AgentDepsT, OutputDataT]):
+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], **kwargs: Any):
-        super().__init__(AGUIAdapter[AgentDepsT, OutputDataT], agent, **kwargs)
+    def __init__(
+        self,
+        agent: AbstractAgent[AgentDepsT, OutputDataT],
+        *,
+        # AGUIAdapter.dispatch_request 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,
+        builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
+        on_complete: OnCompleteFunc[Any] | 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[Self] | None = None,
+    ) -> None:
+        """An ASGI application that handles every request by running the agent and streaming the response.
+
+        Note that the `deps` will be the same for each request, with the exception of the frontend state that's
+        injected into the `state` field of a `deps` object that implements the [`StateHandler`][pydantic_ai.ui.StateHandler] protocol.
+        To provide different `deps` for each request (e.g. based on the authenticated user),
+        use [`AGUIAdapter.run_stream()`][pydantic_ai.ui.ag_ui.AGUIAdapter.run_stream] or
+        [`AGUIAdapter.dispatch_request()`][pydantic_ai.ui.ag_ui.AGUIAdapter.dispatch_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.
+            builtin_tools: Optional additional builtin tools for this run.
+            on_complete: Optional callback function called when the agent run completes successfully.
+                The callback receives the completed [`AgentRunResult`][pydantic_ai.agent.AgentRunResult] and can access `all_messages()` and other result data.
+
+            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 run_agent(request: Request) -> Response:
+            """Endpoint to run the agent with the provided input data."""
+            return await AGUIAdapter.dispatch_request(
+                request,
+                agent=agent,
+                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,
+                builtin_tools=builtin_tools,
+                on_complete=on_complete,
+            )
+
+        self.router.add_route('/', run_agent, methods=['POST'])

pydantic_ai_slim/pydantic_ai/ui/app.py

@@ -30,9 +30,7 @@
 )
 from ...output import OutputDataT
 from ...tools import AgentDepsT
-from ..adapter import UIAdapter
-from ..event_stream import UIEventStream
-from ..messages_builder import MessagesBuilder
+from .. import MessagesBuilder, UIAdapter, UIEventStream
 from ._event_stream import VercelAIEventStream
 from ._request_types import (
     DataUIPart,

pydantic_ai_slim/pydantic_ai/ui/vercel_ai/_adapter.py

@@ -46,8 +46,7 @@
 from pydantic_ai.run import AgentRunResult, AgentRunResultEvent
 from pydantic_ai.tools import ToolDefinition
 from pydantic_ai.toolsets import AbstractToolset, ExternalToolset
-from pydantic_ai.ui.adapter import UIAdapter
-from pydantic_ai.ui.event_stream import NativeEvent, UIEventStream
+from pydantic_ai.ui import NativeEvent, UIAdapter, UIEventStream
 
 from .conftest import try_import