tests

Author: DouweM

pydantic_ai_slim/pydantic_ai/ag_ui.py

@@ -20,6 +20,7 @@
 from .usage import RunUsage, UsageLimits
 
 try:
+    from ag_ui.core import BaseEvent
     from ag_ui.core.types import RunAgentInput
 
     from .ui import OnCompleteFunc, StateDeps, StateHandler
@@ -171,7 +172,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.
 
@@ -226,7 +227,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.
 

pydantic_ai_slim/pydantic_ai/ui/adapter.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
+from collections.abc import AsyncIterator, Mapping, Sequence
 from dataclasses import Field, dataclass, replace
 from functools import cached_property
 from http import HTTPStatus
@@ -17,24 +17,23 @@
     ClassVar,
     Generic,
     Protocol,
-    TypeAlias,
     TypeVar,
     runtime_checkable,
 )
 
 from pydantic import BaseModel, ValidationError
 
-from .. import DeferredToolRequests, DeferredToolResults, _utils
-from ..agent import AbstractAgent, AgentDepsT, AgentRunResult
+from .. import DeferredToolRequests, DeferredToolResults
+from ..agent import AbstractAgent, AgentDepsT
 from ..builtin_tools import AbstractBuiltinTool
 from ..exceptions import UserError
 from ..messages import ModelMessage
 from ..models import KnownModelName, Model
-from ..output import OutputSpec
+from ..output import OutputDataT, OutputSpec
 from ..settings import ModelSettings
 from ..toolsets import AbstractToolset
 from ..usage import RunUsage, UsageLimits
-from .event_stream import BaseEventStream, SourceEvent
+from .event_stream import BaseEventStream, OnCompleteFunc, SourceEvent
 
 if TYPE_CHECKING:
     from starlette.requests import Request
@@ -55,9 +54,6 @@
 EventT = TypeVar('EventT')
 """Type variable for protocol-specific event types."""
 
-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."""
-
 
 # State management types
 
@@ -111,10 +107,10 @@ class StateDeps(Generic[StateT]):
 
 
 @dataclass
-class BaseAdapter(ABC, Generic[RunRequestT, MessageT, EventT, AgentDepsT]):
+class BaseAdapter(ABC, Generic[RunRequestT, MessageT, EventT, AgentDepsT, OutputDataT]):
     """TODO (DouwM): Docstring."""
 
-    agent: AbstractAgent[AgentDepsT]
+    agent: AbstractAgent[AgentDepsT, OutputDataT]
     """The Pydantic AI agent to run."""
 
     request: RunRequestT
@@ -134,7 +130,7 @@ def load_messages(cls, messages: Sequence[MessageT]) -> list[ModelMessage]:
 
     @property
     @abstractmethod
-    def event_stream(self) -> BaseEventStream[RunRequestT, EventT, AgentDepsT]:
+    def event_stream(self) -> BaseEventStream[RunRequestT, EventT, AgentDepsT, OutputDataT]:
         """Create an event stream for the adapter."""
         raise NotImplementedError
 
@@ -178,24 +174,17 @@ def encode_stream(self, stream: AsyncIterator[EventT], accept: str | None = None
     async def process_stream(
         self,
         stream: AsyncIterator[SourceEvent],
-        on_complete: OnCompleteFunc | None = None,
+        on_complete: OnCompleteFunc[EventT] | None = None,
     ) -> AsyncIterator[EventT]:
         """Process a stream of events and return a stream of events.
 
         Args:
             stream: The stream of events to process.
             on_complete: Optional callback function called when the agent run completes successfully.
         """
-        event_stream = self.event_stream
-        async for event in event_stream.handle_stream(stream):
+        async for event in self.event_stream.handle_stream(stream, on_complete=on_complete):
             yield event
 
-            if (result := event_stream.result) and on_complete is not None:
-                if _utils.is_async_callable(on_complete):
-                    await on_complete(result)
-                else:
-                    await _utils.run_in_executor(on_complete, result)
-
     async def run_stream(
         self,
         *,
@@ -210,7 +199,7 @@ async def run_stream(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
-        on_complete: OnCompleteFunc | None = None,
+        on_complete: OnCompleteFunc[EventT] | None = None,
     ) -> AsyncIterator[EventT]:
         """Run the agent with the AG-UI run input and stream AG-UI protocol events.
 
@@ -298,7 +287,7 @@ async def stream_response(self, stream: AsyncIterator[EventT], accept: str | Non
     @classmethod
     async def dispatch_request(
         cls,
-        agent: AbstractAgent[AgentDepsT, Any],
+        agent: AbstractAgent[AgentDepsT, OutputDataT],
         request: Request,
         *,
         message_history: Sequence[ModelMessage] | None = None,
@@ -312,7 +301,7 @@ async def dispatch_request(
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
-        on_complete: OnCompleteFunc | None = None,
+        on_complete: OnCompleteFunc[EventT] | None = None,
     ) -> Response:
         """Handle an AG-UI request and return a streaming response.
 

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_adapter.py

@@ -25,6 +25,7 @@
     ToolReturnPart,
     UserPromptPart,
 )
+from ...output import OutputDataT
 from ...toolsets import AbstractToolset
 
 try:
@@ -92,7 +93,7 @@ def label(self) -> str:
         return 'the AG-UI frontend tools'  # pragma: no cover
 
 
-class AGUIAdapter(BaseAdapter[RunAgentInput, Message, BaseEvent, AgentDepsT]):
+class AGUIAdapter(BaseAdapter[RunAgentInput, Message, BaseEvent, AgentDepsT, OutputDataT]):
     """TODO (DouwM): Docstring."""
 
     @classmethod
@@ -101,7 +102,7 @@ async def validate_request(cls, request: Request) -> RunAgentInput:
         return RunAgentInput.model_validate(await request.json())
 
     @property
-    def event_stream(self) -> BaseEventStream[RunAgentInput, BaseEvent, AgentDepsT]:
+    def event_stream(self) -> BaseEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]:
         """Create an event stream for the adapter."""
         return AGUIEventStream(self.request)
 

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_event_stream.py

@@ -8,6 +8,7 @@
 
 import json
 from collections.abc import AsyncIterator, Iterable
+from dataclasses import dataclass, field
 from typing import Final
 
 from ...messages import (
@@ -23,6 +24,7 @@
     ToolCallPartDelta,
     ToolReturnPart,
 )
+from ...output import OutputDataT
 from ...tools import AgentDepsT
 from .. import BaseEventStream
 
@@ -69,14 +71,13 @@
 BUILTIN_TOOL_CALL_ID_PREFIX: Final[str] = 'pyd_ai_builtin'
 
 
-class AGUIEventStream(BaseEventStream[RunAgentInput, BaseEvent, AgentDepsT]):
+@dataclass
+class AGUIEventStream(BaseEventStream[RunAgentInput, BaseEvent, AgentDepsT, OutputDataT]):
     """TODO (DouwM): Docstring."""
 
-    def __init__(self, request: RunAgentInput) -> None:
-        """Initialize AG-UI event stream state."""
-        super().__init__(request)
-        self._thinking_text = False
-        self._builtin_tool_call_ids: dict[str, str] = {}
+    _thinking_text: bool = False
+    _builtin_tool_call_ids: dict[str, str] = field(default_factory=dict)
+    _error: bool = False
 
     def encode_event(self, event: BaseEvent, accept: str | None = None) -> str:
         """Encode an AG-UI event as SSE.
@@ -100,13 +101,15 @@ async def before_stream(self) -> AsyncIterator[BaseEvent]:
 
     async def after_stream(self) -> AsyncIterator[BaseEvent]:
         """Handle an AgentRunResultEvent, cleaning up any pending state."""
-        yield RunFinishedEvent(
-            thread_id=self.request.thread_id,
-            run_id=self.request.run_id,
-        )
+        if not self._error:
+            yield RunFinishedEvent(
+                thread_id=self.request.thread_id,
+                run_id=self.request.run_id,
+            )
 
     async def on_error(self, error: Exception) -> AsyncIterator[BaseEvent]:
         """Handle errors during streaming."""
+        self._error = True
         yield RunErrorEvent(message=str(error))
 
     async def handle_text_start(self, part: TextPart, follows_text: bool = False) -> AsyncIterator[BaseEvent]: