Clean up Pydantic AI message building

Author: DouweM

pydantic_ai_slim/pydantic_ai/ag_ui.py

@@ -23,9 +23,8 @@
     from ag_ui.core import BaseEvent
     from ag_ui.core.types import RunAgentInput
 
-    from .ui import OnCompleteFunc, StateDeps, StateHandler
+    from .ui import SSE_CONTENT_TYPE, OnCompleteFunc, StateDeps, StateHandler
     from .ui.ag_ui import (
-        SSE_CONTENT_TYPE,
         AGUIAdapter,
         AGUIApp,
     )

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -755,7 +755,6 @@ async def run_agent() -> AgentRunResult[Any]:
                 yield message
 
         result = await task
-        # TODO (DouweM): Consider adding this to every event stream, if we're adding new events anyway
         yield AgentRunResultEvent(result)
 
     @overload

pydantic_ai_slim/pydantic_ai/ui/__init__.py

@@ -9,6 +9,7 @@
 from .adapter import OnCompleteFunc, StateDeps, StateHandler, UIAdapter
 from .app import UIApp
 from .event_stream import SSE_CONTENT_TYPE, UIEventStream
+from .messages_builder import MessagesBuilder
 
 __all__ = [
     'UIAdapter',
@@ -18,4 +19,5 @@
     'StateHandler',
     'OnCompleteFunc',
     'UIApp',
+    'MessagesBuilder',
 ]

pydantic_ai_slim/pydantic_ai/ui/adapter.py

@@ -106,7 +106,7 @@ class StateDeps(Generic[StateT]):
 
 @dataclass
 class UIAdapter(ABC, Generic[RunInputT, MessageT, EventT, AgentDepsT, OutputDataT]):
-    """TODO (DouwM): Docstring."""
+    """TODO (DouweM): Docstring."""
 
     agent: AbstractAgent[AgentDepsT, OutputDataT]
     """The Pydantic AI agent to run."""
@@ -123,7 +123,7 @@ class UIAdapter(ABC, Generic[RunInputT, MessageT, EventT, AgentDepsT, OutputData
     async def from_request(
         cls, request: Request, *, agent: AbstractAgent[AgentDepsT, OutputDataT]
     ) -> UIAdapter[RunInputT, MessageT, EventT, AgentDepsT, OutputDataT]:
-        """Create an adapter from a protocol-specific request."""
+        """Create an adapter from a protocol-specific run input."""
         return cls(
             agent=agent,
             run_input=await cls.build_run_input(request),
@@ -150,17 +150,17 @@ def build_event_stream(self) -> UIEventStream[RunInputT, EventT, AgentDepsT, Out
     @cached_property
     @abstractmethod
     def messages(self) -> list[ModelMessage]:
-        """Pydantic AI messages from the protocol-specific request."""
+        """Pydantic AI messages from the protocol-specific run input."""
         raise NotImplementedError
 
     @cached_property
     def toolset(self) -> AbstractToolset[AgentDepsT] | None:
-        """Toolset representing frontend tools from the protocol-specific request."""
+        """Toolset representing frontend tools from the protocol-specific run input."""
         return None
 
     @cached_property
     def state(self) -> dict[str, Any] | None:
-        """Run state from the protocol-specific request."""
+        """Run state from the protocol-specific run input."""
         return None
 
     def transform_stream(
@@ -210,7 +210,7 @@ def run_stream_native(
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
     ) -> AsyncIterator[NativeEvent]:
-        """Run the agent with the protocol-specific request as input and stream Pydantic AI events.
+        """Run the agent with the protocol-specific run input and stream Pydantic AI events.
 
         Args:
             output_type: Custom output type to use for this run, `output_type` may only be used if the agent has no
@@ -276,7 +276,7 @@ def run_stream(
         builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
         on_complete: OnCompleteFunc[EventT] | None = None,
     ) -> AsyncIterator[EventT]:
-        """Run the agent with the protocol-specific request as input and stream protocol-specific events.
+        """Run the agent with the protocol-specific run input and stream protocol-specific events.
 
         Args:
             output_type: Custom output type to use for this run, `output_type` may only be used if the agent has no
@@ -333,8 +333,8 @@ async def dispatch_request(
         """Handle an protocol-specific HTTP request by running the agent and return a streaming response of protocol-specific events.
 
         Args:
-            agent: The agent to run.
             request: The incoming Starlette/FastAPI request.
+            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.

pydantic_ai_slim/pydantic_ai/ui/ag_ui/__init__.py

@@ -8,12 +8,11 @@
 
 from .. import UIApp
 from ._adapter import AGUIAdapter
-from ._event_stream import SSE_CONTENT_TYPE, AGUIEventStream
+from ._event_stream import AGUIEventStream
 
 __all__ = [
     'AGUIAdapter',
     'AGUIEventStream',
-    'SSE_CONTENT_TYPE',
     'AGUIApp',
 ]
 

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_adapter.py

@@ -15,10 +15,6 @@
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     ModelMessage,
-    ModelRequest,
-    ModelRequestPart,
-    ModelResponse,
-    ModelResponsePart,
     SystemPromptPart,
     TextPart,
     ToolCallPart,
@@ -43,6 +39,7 @@
 
     from ..adapter import UIAdapter
     from ..event_stream import UIEventStream
+    from ..messages_builder import MessagesBuilder
     from ._event_stream import BUILTIN_TOOL_CALL_ID_PREFIX, AGUIEventStream
 except ImportError as e:  # pragma: no cover
     raise ImportError(
@@ -94,7 +91,7 @@ def label(self) -> str:
 
 
 class AGUIAdapter(UIAdapter[RunAgentInput, Message, BaseEvent, AgentDepsT, OutputDataT]):
-    """TODO (DouwM): Docstring."""
+    """TODO (DouweM): Docstring."""
 
     @classmethod
     async def build_run_input(cls, request: Request) -> RunAgentInput:
@@ -132,31 +129,24 @@ def messages(self) -> list[ModelMessage]:
     @classmethod
     def load_messages(cls, messages: Sequence[Message]) -> list[ModelMessage]:
         """Load messages from the request and return the loaded messages."""
-        result: list[ModelMessage] = []
+        builder = MessagesBuilder()
         tool_calls: dict[str, str] = {}  # Tool call ID to tool name mapping.
-        request_parts: list[ModelRequestPart] | None = None
-        response_parts: list[ModelResponsePart] | None = None
 
         for msg in messages:
             if isinstance(msg, UserMessage | SystemMessage | DeveloperMessage) or (
                 isinstance(msg, ToolMessage) and not msg.tool_call_id.startswith(BUILTIN_TOOL_CALL_ID_PREFIX)
             ):
-                if request_parts is None:
-                    request_parts = []
-                    result.append(ModelRequest(parts=request_parts))
-                    response_parts = None
-
                 if isinstance(msg, UserMessage):
-                    request_parts.append(UserPromptPart(content=msg.content))
+                    builder.add(UserPromptPart(content=msg.content))
                 elif isinstance(msg, SystemMessage | DeveloperMessage):
-                    request_parts.append(SystemPromptPart(content=msg.content))
+                    builder.add(SystemPromptPart(content=msg.content))
                 else:
                     tool_call_id = msg.tool_call_id
                     tool_name = tool_calls.get(tool_call_id)
                     if tool_name is None:  # pragma: no cover
                         raise ValueError(f'Tool call with ID {tool_call_id} not found in the history.')
 
-                    request_parts.append(
+                    builder.add(
                         ToolReturnPart(
                             tool_name=tool_name,
                             content=msg.content,
@@ -167,14 +157,9 @@ def load_messages(cls, messages: Sequence[Message]) -> list[ModelMessage]:
             elif isinstance(msg, AssistantMessage) or (  # pragma: no branch
                 isinstance(msg, ToolMessage) and msg.tool_call_id.startswith(BUILTIN_TOOL_CALL_ID_PREFIX)
             ):
-                if response_parts is None:
-                    response_parts = []
-                    result.append(ModelResponse(parts=response_parts))
-                    request_parts = None
-
                 if isinstance(msg, AssistantMessage):
                     if msg.content:
-                        response_parts.append(TextPart(content=msg.content))
+                        builder.add(TextPart(content=msg.content))
 
                     if msg.tool_calls:
                         for tool_call in msg.tool_calls:
@@ -184,7 +169,7 @@ def load_messages(cls, messages: Sequence[Message]) -> list[ModelMessage]:
 
                             if tool_call_id.startswith(BUILTIN_TOOL_CALL_ID_PREFIX):
                                 _, provider_name, tool_call_id = tool_call_id.split('|', 2)
-                                response_parts.append(
+                                builder.add(
                                     BuiltinToolCallPart(
                                         tool_name=tool_name,
                                         args=tool_call.function.arguments,
@@ -193,7 +178,7 @@ def load_messages(cls, messages: Sequence[Message]) -> list[ModelMessage]:
                                     )
                                 )
                             else:
-                                response_parts.append(
+                                builder.add(
                                     ToolCallPart(
                                         tool_name=tool_name,
                                         tool_call_id=tool_call_id,
@@ -207,7 +192,7 @@ def load_messages(cls, messages: Sequence[Message]) -> list[ModelMessage]:
                         raise ValueError(f'Tool call with ID {tool_call_id} not found in the history.')
                     _, provider_name, tool_call_id = tool_call_id.split('|', 2)
 
-                    response_parts.append(
+                    builder.add(
                         BuiltinToolReturnPart(
                             tool_name=tool_name,
                             content=msg.content,
@@ -216,4 +201,4 @@ def load_messages(cls, messages: Sequence[Message]) -> list[ModelMessage]:
                         )
                     )
 
-        return result
+        return builder.messages

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_event_stream.py

@@ -69,7 +69,7 @@
 
 @dataclass
 class AGUIEventStream(UIEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]):
-    """TODO (DouwM): Docstring."""
+    """TODO (DouweM): Docstring."""
 
     _thinking_text: bool = False
     _builtin_tool_call_ids: dict[str, str] = field(default_factory=dict)

pydantic_ai_slim/pydantic_ai/ui/app.py

@@ -1,18 +1,21 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Mapping, Sequence
-from typing import Any, Generic, Self
+from typing import Any, Generic
+
+from typing_extensions import Self
 
 from .. import DeferredToolResults
 from ..agent import AbstractAgent
+from ..builtin_tools import AbstractBuiltinTool
 from ..messages import ModelMessage
 from ..models import KnownModelName, Model
 from ..output import OutputDataT, OutputSpec
 from ..settings import ModelSettings
 from ..tools import AgentDepsT
 from ..toolsets import AbstractToolset
 from ..usage import RunUsage, UsageLimits
-from .adapter import UIAdapter
+from .adapter import OnCompleteFunc, UIAdapter
 
 try:
     from starlette.applications import Starlette
@@ -36,7 +39,7 @@ def __init__(
         adapter_type: type[UIAdapter[Any, Any, Any, AgentDepsT, OutputDataT]],
         agent: AbstractAgent[AgentDepsT, OutputDataT],
         *,
-        # Agent.iter parameters.
+        # UIAdapter.dispatch_request parameters
         output_type: OutputSpec[Any] | None = None,
         message_history: Sequence[ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
@@ -47,7 +50,9 @@ def __init__(
         usage: RunUsage | None = None,
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
-        # Starlette parameters.
+        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,
@@ -58,12 +63,11 @@ def __init__(
     ) -> None:
         """An ASGI application that handles every request by running the agent and streaming the response.
 
-        # TODO (DouweM): Docstring
-        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.
+        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 [`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.
+        use [`UIAdapter.run_stream()`][pydantic_ai.ui.UIAdapter.run_stream] or
+        [`UIAdapter.dispatch_request()`][pydantic_ai.ui.UIAdapter.dispatch_request] instead.
 
         Args:
             adapter_type: The type of the UI adapter to use.
@@ -81,6 +85,9 @@ def __init__(
             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.
@@ -125,6 +132,8 @@ async def run_agent(request: Request) -> Response:
                 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/event_stream.py

@@ -69,7 +69,7 @@
 
 @dataclass
 class UIEventStream(ABC, Generic[RunInputT, EventT, AgentDepsT, OutputDataT]):
-    """TODO (DouwM): Docstring."""
+    """TODO (DouweM): Docstring."""
 
     run_input: RunInputT
 
@@ -164,8 +164,6 @@ async def transform_stream(  # noqa: C901
 
         try:
             async for event in stream:
-                # TODO (DouweM): Introduce, possibly, MessageStartEvent, MessageEndEvent with ModelRequest/Response?
-                # People have requested these before. We can store Request and Response
                 if isinstance(event, PartStartEvent):
                     async for e in self._turn_to('response'):
                         yield e

pydantic_ai_slim/pydantic_ai/ui/messages_builder.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass, field
+from typing import cast
+
+from pydantic_ai._utils import get_union_args
+from pydantic_ai.messages import ModelMessage, ModelRequest, ModelRequestPart, ModelResponse, ModelResponsePart
+
+
+@dataclass
+class MessagesBuilder:
+    """Helper class to build Pydantic AI messages from protocol-specific messages."""
+
+    messages: list[ModelMessage] = field(default_factory=list)
+
+    def add(self, part: ModelRequest | ModelResponse | ModelRequestPart | ModelResponsePart) -> None:
+        """Add a new part, creating a new request or response message if necessary."""
+        last_message = self.messages[-1] if self.messages else None
+        if isinstance(part, get_union_args(ModelRequestPart)):
+            if isinstance(last_message, ModelRequest):
+                last_message.parts = [*last_message.parts, cast(ModelRequestPart, part)]
+            else:
+                self.messages.append(ModelRequest(parts=[part]))
+        else:
+            if isinstance(last_message, ModelResponse):
+                last_message.parts = [*last_message.parts, part]
+            else:
+                self.messages.append(ModelResponse(parts=[part]))