pydantic-ai-slim 1.8.0__py3-none-any.whl → 1.9.1__py3-none-any.whl

pydantic_ai/ag_ui.py

@@ -4,107 +4,35 @@ This package provides seamless integration between pydantic-ai agents and ag-ui
 for building interactive AI applications with streaming event-based communication.
 """
 
-from __future__ import annotations
+# TODO (v2): Remove this module in favor of `pydantic_ai.ui.ag_ui`
 
-import json
-import uuid
-from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping, Sequence
-from dataclasses import Field, dataclass, field, replace
-from http import HTTPStatus
-from typing import (
-    Any,
-    ClassVar,
-    Final,
-    Generic,
-    Protocol,
-    TypeAlias,
-    TypeVar,
-    runtime_checkable,
-)
+from __future__ import annotations
 
-from pydantic import BaseModel, ValidationError
+from collections.abc import AsyncIterator, Sequence
+from typing import Any
 
-from . import _utils
-from ._agent_graph import CallToolsNode, ModelRequestNode
-from .agent import AbstractAgent, AgentRun, AgentRunResult
-from .exceptions import UserError
-from .messages import (
-    BaseToolCallPart,
-    BuiltinToolCallPart,
-    BuiltinToolReturnPart,
-    FunctionToolResultEvent,
-    ModelMessage,
-    ModelRequest,
-    ModelRequestPart,
-    ModelResponse,
-    ModelResponsePart,
-    ModelResponseStreamEvent,
-    PartDeltaEvent,
-    PartStartEvent,
-    SystemPromptPart,
-    TextPart,
-    TextPartDelta,
-    ThinkingPart,
-    ThinkingPartDelta,
-    ToolCallPart,
-    ToolCallPartDelta,
-    ToolReturnPart,
-    UserPromptPart,
-)
+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, DeferredToolRequests, ToolDefinition
+from .tools import AgentDepsT
 from .toolsets import AbstractToolset
-from .toolsets.external import ExternalToolset
 from .usage import RunUsage, UsageLimits
 
 try:
-    from ag_ui.core import (
-        AssistantMessage,
-        BaseEvent,
-        DeveloperMessage,
-        EventType,
-        Message,
-        RunAgentInput,
-        RunErrorEvent,
-        RunFinishedEvent,
-        RunStartedEvent,
-        State,
-        SystemMessage,
-        TextMessageContentEvent,
-        TextMessageEndEvent,
-        TextMessageStartEvent,
-        ThinkingEndEvent,
-        ThinkingStartEvent,
-        ThinkingTextMessageContentEvent,
-        ThinkingTextMessageEndEvent,
-        ThinkingTextMessageStartEvent,
-        Tool as AGUITool,
-        ToolCallArgsEvent,
-        ToolCallEndEvent,
-        ToolCallResultEvent,
-        ToolCallStartEvent,
-        ToolMessage,
-        UserMessage,
-    )
-    from ag_ui.encoder import EventEncoder
-except ImportError as e:  # pragma: no cover
-    raise ImportError(
-        'Please install the `ag-ui-protocol` package to use `Agent.to_ag_ui()` method, '
-        'you can use the `ag-ui` optional group — `pip install "pydantic-ai-slim[ag-ui]"`'
-    ) from e
-
-try:
-    from starlette.applications import Starlette
-    from starlette.middleware import Middleware
+    from ag_ui.core import BaseEvent
+    from ag_ui.core.types import RunAgentInput
     from starlette.requests import Request
-    from starlette.responses import Response, StreamingResponse
-    from starlette.routing import BaseRoute
-    from starlette.types import ExceptionHandler, Lifespan
+    from starlette.responses import Response
+
+    from .ui import SSE_CONTENT_TYPE, OnCompleteFunc, StateDeps, StateHandler
+    from .ui.ag_ui import AGUIAdapter
+    from .ui.ag_ui.app import AGUIApp
 except ImportError as e:  # pragma: no cover
     raise ImportError(
-        'Please install the `starlette` package to use `Agent.to_ag_ui()` method, '
+        'Please install the `ag-ui-protocol` and `starlette` packages to use `AGUIAdapter`, '
         'you can use the `ag-ui` optional group — `pip install "pydantic-ai-slim[ag-ui]"`'
     ) from e
 
@@ -119,113 +47,14 @@ __all__ = [
     'run_ag_ui',
 ]
 
-SSE_CONTENT_TYPE: Final[str] = 'text/event-stream'
-"""Content type header value for Server-Sent Events (SSE)."""
-
-OnCompleteFunc: TypeAlias = Callable[[AgentRunResult[Any]], None] | Callable[[AgentRunResult[Any]], Awaitable[None]]
-"""Callback function type that receives the `AgentRunResult` of the completed run. Can be sync or async."""
-
-_BUILTIN_TOOL_CALL_ID_PREFIX: Final[str] = 'pyd_ai_builtin'
-
-
-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,
-        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.
-            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')
-
 
 async def handle_ag_ui_request(
     agent: AbstractAgent[AgentDepsT, Any],
     request: Request,
     *,
     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,
@@ -233,7 +62,7 @@ async def handle_ag_ui_request(
     usage: RunUsage | None = None,
     infer_name: bool = True,
     toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
-    on_complete: OnCompleteFunc | None = None,
+    on_complete: OnCompleteFunc[BaseEvent] | None = None,
 ) -> Response:
     """Handle an AG-UI request by running the agent and returning a streaming response.
 
@@ -243,6 +72,8 @@ async def handle_ag_ui_request(
 
         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.
@@ -256,41 +87,31 @@ async def handle_ag_ui_request(
     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,
-            on_complete=on_complete,
-        ),
-        media_type=accept,
+    return await AGUIAdapter[AgentDepsT].dispatch_request(
+        request,
+        agent=agent,
+        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,
+        usage=usage,
+        infer_name=infer_name,
+        toolsets=toolsets,
+        on_complete=on_complete,
     )
 
 
-async def run_ag_ui(
+def run_ag_ui(
     agent: AbstractAgent[AgentDepsT, Any],
     run_input: RunAgentInput,
     accept: str = SSE_CONTENT_TYPE,
     *,
     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,
@@ -298,7 +119,7 @@ async def run_ag_ui(
     usage: RunUsage | None = None,
     infer_name: bool = True,
     toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
-    on_complete: OnCompleteFunc | None = None,
+    on_complete: OnCompleteFunc[BaseEvent] | None = None,
 ) -> AsyncIterator[str]:
     """Run the agent with the AG-UI run input and stream AG-UI protocol events.
 
@@ -309,6 +130,8 @@ async def run_ag_ui(
 
         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.
@@ -322,50 +145,12 @@ async def run_ag_ui(
     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 = _AGUIFrontendToolset[AgentDepsT](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, DeferredToolRequests],
-            message_history=messages,
+    adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
+    return adapter.encode_stream(
+        adapter.run_stream(
+            output_type=output_type,
+            message_history=message_history,
+            deferred_tool_results=deferred_tool_results,
             model=model,
             deps=deps,
             model_settings=model_settings,
@@ -373,437 +158,6 @@ async def run_ag_ui(
             usage=usage,
             infer_name=infer_name,
             toolsets=toolsets,
-        ) as run:
-            async for event in _agent_stream(run):
-                yield encoder.encode(event)
-
-        if on_complete is not None and run.result is not None:
-            if _utils.is_async_callable(on_complete):
-                await on_complete(run.result)
-            else:
-                await _utils.run_in_executor(on_complete, run.result)
-    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,
-            ),
-        )
-
-
-async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEvent]:
-    """Run the agent streaming responses using AG-UI protocol events.
-
-    Args:
-        run: The agent run to process.
-
-    Yields:
-        AG-UI Server-Sent Events (SSE).
-    """
-    async for node in run:
-        stream_ctx = _RequestStreamContext()
-        if isinstance(node, ModelRequestNode):
-            async with node.stream(run.ctx) as request_stream:
-                async for agent_event in request_stream:
-                    async for msg in _handle_model_request_event(stream_ctx, agent_event):
-                        yield msg
-
-                if stream_ctx.part_end:  # pragma: no branch
-                    yield stream_ctx.part_end
-                    stream_ctx.part_end = None
-                if stream_ctx.thinking:
-                    yield ThinkingEndEvent(
-                        type=EventType.THINKING_END,
-                    )
-                    stream_ctx.thinking = False
-        elif isinstance(node, CallToolsNode):
-            async with node.stream(run.ctx) as handle_stream:
-                async for event in handle_stream:
-                    if isinstance(event, FunctionToolResultEvent):
-                        async for msg in _handle_tool_result_event(stream_ctx, event):
-                            yield msg
-
-
-async def _handle_model_request_event(  # noqa: C901
-    stream_ctx: _RequestStreamContext,
-    agent_event: ModelResponseStreamEvent,
-) -> AsyncIterator[BaseEvent]:
-    """Handle an agent event and yield AG-UI protocol events.
-
-    Args:
-        stream_ctx: The request stream context to manage state.
-        agent_event: The agent event to process.
-
-    Yields:
-        AG-UI Server-Sent Events (SSE) based on the agent event.
-    """
-    if isinstance(agent_event, PartStartEvent):
-        if stream_ctx.part_end:
-            # End the previous part.
-            yield stream_ctx.part_end
-            stream_ctx.part_end = None
-
-        part = agent_event.part
-        if isinstance(part, ThinkingPart):  # pragma: no branch
-            if not stream_ctx.thinking:
-                yield ThinkingStartEvent(
-                    type=EventType.THINKING_START,
-                )
-                stream_ctx.thinking = True
-
-            if part.content:
-                yield ThinkingTextMessageStartEvent(
-                    type=EventType.THINKING_TEXT_MESSAGE_START,
-                )
-                yield ThinkingTextMessageContentEvent(
-                    type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
-                    delta=part.content,
-                )
-                stream_ctx.part_end = ThinkingTextMessageEndEvent(
-                    type=EventType.THINKING_TEXT_MESSAGE_END,
-                )
-        else:
-            if stream_ctx.thinking:
-                yield ThinkingEndEvent(
-                    type=EventType.THINKING_END,
-                )
-                stream_ctx.thinking = False
-
-            if isinstance(part, TextPart):
-                message_id = stream_ctx.new_message_id()
-                yield TextMessageStartEvent(
-                    message_id=message_id,
-                )
-                if part.content:  # pragma: no branch
-                    yield TextMessageContentEvent(
-                        message_id=message_id,
-                        delta=part.content,
-                    )
-                stream_ctx.part_end = TextMessageEndEvent(
-                    message_id=message_id,
-                )
-            elif isinstance(part, BaseToolCallPart):
-                tool_call_id = part.tool_call_id
-                if isinstance(part, BuiltinToolCallPart):
-                    builtin_tool_call_id = '|'.join(
-                        [_BUILTIN_TOOL_CALL_ID_PREFIX, part.provider_name or '', tool_call_id]
-                    )
-                    stream_ctx.builtin_tool_call_ids[tool_call_id] = builtin_tool_call_id
-                    tool_call_id = builtin_tool_call_id
-
-                message_id = stream_ctx.message_id or stream_ctx.new_message_id()
-                yield ToolCallStartEvent(
-                    tool_call_id=tool_call_id,
-                    tool_call_name=part.tool_name,
-                    parent_message_id=message_id,
-                )
-                if part.args:
-                    yield ToolCallArgsEvent(
-                        tool_call_id=tool_call_id,
-                        delta=part.args_as_json_str(),
-                    )
-                stream_ctx.part_end = ToolCallEndEvent(
-                    tool_call_id=tool_call_id,
-                )
-            elif isinstance(part, BuiltinToolReturnPart):  # pragma: no branch
-                tool_call_id = stream_ctx.builtin_tool_call_ids[part.tool_call_id]
-                yield ToolCallResultEvent(
-                    message_id=stream_ctx.new_message_id(),
-                    type=EventType.TOOL_CALL_RESULT,
-                    role='tool',
-                    tool_call_id=tool_call_id,
-                    content=part.model_response_str(),
-                )
-
-    elif isinstance(agent_event, PartDeltaEvent):
-        delta = agent_event.delta
-        if isinstance(delta, TextPartDelta):
-            if delta.content_delta:  # pragma: no branch
-                yield TextMessageContentEvent(
-                    message_id=stream_ctx.message_id,
-                    delta=delta.content_delta,
-                )
-        elif isinstance(delta, ToolCallPartDelta):  # pragma: no branch
-            tool_call_id = delta.tool_call_id
-            assert tool_call_id, '`ToolCallPartDelta.tool_call_id` must be set'
-            if tool_call_id in stream_ctx.builtin_tool_call_ids:
-                tool_call_id = stream_ctx.builtin_tool_call_ids[tool_call_id]
-            yield ToolCallArgsEvent(
-                tool_call_id=tool_call_id,
-                delta=delta.args_delta if isinstance(delta.args_delta, str) else json.dumps(delta.args_delta),
-            )
-        elif isinstance(delta, ThinkingPartDelta):  # pragma: no branch
-            if delta.content_delta:  # pragma: no branch
-                if not isinstance(stream_ctx.part_end, ThinkingTextMessageEndEvent):
-                    yield ThinkingTextMessageStartEvent(
-                        type=EventType.THINKING_TEXT_MESSAGE_START,
-                    )
-                    stream_ctx.part_end = ThinkingTextMessageEndEvent(
-                        type=EventType.THINKING_TEXT_MESSAGE_END,
-                    )
-
-                yield ThinkingTextMessageContentEvent(
-                    type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
-                    delta=delta.content_delta,
-                )
-
-
-async def _handle_tool_result_event(
-    stream_ctx: _RequestStreamContext,
-    event: FunctionToolResultEvent,
-) -> AsyncIterator[BaseEvent]:
-    """Convert a tool call result to AG-UI events.
-
-    Args:
-        stream_ctx: The request stream context to manage state.
-        event: The tool call result event to process.
-
-    Yields:
-        AG-UI Server-Sent Events (SSE).
-    """
-    result = event.result
-    if not isinstance(result, ToolReturnPart):
-        return
-
-    yield ToolCallResultEvent(
-        message_id=stream_ctx.new_message_id(),
-        type=EventType.TOOL_CALL_RESULT,
-        role='tool',
-        tool_call_id=result.tool_call_id,
-        content=result.model_response_str(),
+            on_complete=on_complete,
+        ),
     )
-
-    # Now check for AG-UI events returned by the tool calls.
-    possible_event = result.metadata or result.content
-    if isinstance(possible_event, BaseEvent):
-        yield possible_event
-    elif isinstance(possible_event, str | bytes):  # pragma: no branch
-        # Avoid iterable check for strings and bytes.
-        pass
-    elif isinstance(possible_event, Iterable):  # pragma: no branch
-        for item in possible_event:  # type: ignore[reportUnknownMemberType]
-            if isinstance(item, BaseEvent):  # pragma: no branch
-                yield item
-
-
-def _messages_from_ag_ui(messages: list[Message]) -> list[ModelMessage]:
-    """Convert a AG-UI history to a Pydantic AI one."""
-    result: list[ModelMessage] = []
-    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))
-            elif isinstance(msg, SystemMessage | DeveloperMessage):
-                request_parts.append(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 _ToolCallNotFoundError(tool_call_id=tool_call_id)
-
-                request_parts.append(
-                    ToolReturnPart(
-                        tool_name=tool_name,
-                        content=msg.content,
-                        tool_call_id=tool_call_id,
-                    )
-                )
-
-        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))
-
-                if msg.tool_calls:
-                    for tool_call in msg.tool_calls:
-                        tool_call_id = tool_call.id
-                        tool_name = tool_call.function.name
-                        tool_calls[tool_call_id] = tool_name
-
-                        if tool_call_id.startswith(_BUILTIN_TOOL_CALL_ID_PREFIX):
-                            _, provider_name, tool_call_id = tool_call_id.split('|', 2)
-                            response_parts.append(
-                                BuiltinToolCallPart(
-                                    tool_name=tool_name,
-                                    args=tool_call.function.arguments,
-                                    tool_call_id=tool_call_id,
-                                    provider_name=provider_name,
-                                )
-                            )
-                        else:
-                            response_parts.append(
-                                ToolCallPart(
-                                    tool_name=tool_name,
-                                    tool_call_id=tool_call_id,
-                                    args=tool_call.function.arguments,
-                                )
-                            )
-            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 _ToolCallNotFoundError(tool_call_id=tool_call_id)
-                _, provider_name, tool_call_id = tool_call_id.split('|', 2)
-
-                response_parts.append(
-                    BuiltinToolReturnPart(
-                        tool_name=tool_name,
-                        content=msg.content,
-                        tool_call_id=tool_call_id,
-                        provider_name=provider_name,
-                    )
-                )
-
-    return result
-
-
-@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.
-        """
-        ...
-
-
-StateT = TypeVar('StateT', bound=BaseModel)
-"""Type variable for the state type, which must be a subclass of `BaseModel`."""
-
-
-@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
-
-
-@dataclass(repr=False)
-class _RequestStreamContext:
-    """Data class to hold request stream context."""
-
-    message_id: str = ''
-    part_end: BaseEvent | None = None
-    thinking: bool = False
-    builtin_tool_call_ids: dict[str, str] = field(default_factory=dict)
-
-    def new_message_id(self) -> str:
-        """Generate a new message ID for the request stream.
-
-        Assigns a new UUID to the `message_id` and returns it.
-
-        Returns:
-            A new message ID.
-        """
-        self.message_id = str(uuid.uuid4())
-        return self.message_id
-
-
-@dataclass
-class _RunError(Exception):
-    """Exception raised for errors during agent runs."""
-
-    message: str
-    code: str
-
-    def __str__(self) -> str:  # pragma: no cover
-        return self.message
-
-
-@dataclass
-class _NoMessagesError(_RunError):
-    """Exception raised when no messages are found in the input."""
-
-    message: str = 'no messages found in the input'
-    code: str = 'no_messages'
-
-
-@dataclass
-class _InvalidStateError(_RunError, ValidationError):
-    """Exception raised when an invalid state is provided."""
-
-    message: str = 'invalid state provided'
-    code: str = 'invalid_state'
-
-
-class _ToolCallNotFoundError(_RunError, ValueError):
-    """Exception raised when an tool result is present without a matching call."""
-
-    def __init__(self, tool_call_id: str) -> None:
-        """Initialize the exception with the tool call ID."""
-        super().__init__(  # pragma: no cover
-            message=f'Tool call with ID {tool_call_id} not found in the history.',
-            code='tool_call_not_found',
-        )
-
-
-class _AGUIFrontendToolset(ExternalToolset[AgentDepsT]):
-    def __init__(self, tools: list[AGUITool]):
-        super().__init__(
-            [
-                ToolDefinition(
-                    name=tool.name,
-                    description=tool.description,
-                    parameters_json_schema=tool.parameters,
-                )
-                for tool in tools
-            ]
-        )
-
-    @property
-    def label(self) -> str:
-        return 'the AG-UI frontend tools'  # pragma: no cover