BREAKING CHANGE: Make AgentStreamEvent union of ModelResponseStreamEvent and HandleResponseEvent (#2689)

Author: DouweM

docs/agents.md

@@ -115,15 +115,13 @@ The example below shows how to stream events and text output. You can also [stre
 import asyncio
 from collections.abc import AsyncIterable
 from datetime import date
-from typing import Union
 
 from pydantic_ai import Agent
 from pydantic_ai.messages import (
     AgentStreamEvent,
     FinalResultEvent,
     FunctionToolCallEvent,
     FunctionToolResultEvent,
-    HandleResponseEvent,
     PartDeltaEvent,
     PartStartEvent,
     TextPartDelta,
@@ -152,7 +150,7 @@ output_messages: list[str] = []
 
 async def event_stream_handler(
     ctx: RunContext,
-    event_stream: AsyncIterable[Union[AgentStreamEvent, HandleResponseEvent]],
+    event_stream: AsyncIterable[AgentStreamEvent],
 ):
     async for event in event_stream:
         if isinstance(event, PartStartEvent):

pydantic_ai_slim/pydantic_ai/ag_ui.py

@@ -28,11 +28,11 @@
 from .agent import AbstractAgent, AgentRun
 from .exceptions import UserError
 from .messages import (
-    AgentStreamEvent,
     FunctionToolResultEvent,
     ModelMessage,
     ModelRequest,
     ModelResponse,
+    ModelResponseStreamEvent,
     PartDeltaEvent,
     PartStartEvent,
     SystemPromptPart,
@@ -403,7 +403,7 @@ async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEve
 
 async def _handle_model_request_event(
     stream_ctx: _RequestStreamContext,
-    agent_event: AgentStreamEvent,
+    agent_event: ModelResponseStreamEvent,
 ) -> AsyncIterator[BaseEvent]:
     """Handle an agent event and yield AG-UI protocol events.
 

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -5,7 +5,7 @@
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterator, Mapping, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from types import FrameType
-from typing import TYPE_CHECKING, Any, Callable, Generic, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Generic, cast, overload
 
 from typing_extensions import Self, TypeAlias, TypeIs, TypeVar
 
@@ -53,11 +53,7 @@
 """Type variable for the result data of a run where `output_type` was customized on the run call."""
 
 EventStreamHandler: TypeAlias = Callable[
-    [
-        RunContext[AgentDepsT],
-        AsyncIterable[Union[_messages.AgentStreamEvent, _messages.HandleResponseEvent]],
-    ],
-    Awaitable[None],
+    [RunContext[AgentDepsT], AsyncIterable[_messages.AgentStreamEvent]], Awaitable[None]
 ]
 """A function that receives agent [`RunContext`][pydantic_ai.tools.RunContext] and an async iterable of events from the model's streaming response and the agent's execution of tools."""
 
@@ -445,7 +441,9 @@ async def main():
                     async with node.stream(graph_ctx) as stream:
                         final_result_event = None
 
-                        async def stream_to_final(stream: AgentStream) -> AsyncIterator[_messages.AgentStreamEvent]:
+                        async def stream_to_final(
+                            stream: AgentStream,
+                        ) -> AsyncIterator[_messages.ModelResponseStreamEvent]:
                             nonlocal final_result_event
                             async for event in stream:
                                 yield event

pydantic_ai_slim/pydantic_ai/direct.py

@@ -275,7 +275,9 @@ class StreamedResponseSync:
     """
 
     _async_stream_cm: AbstractAsyncContextManager[StreamedResponse]
-    _queue: queue.Queue[messages.AgentStreamEvent | Exception | None] = field(default_factory=queue.Queue, init=False)
+    _queue: queue.Queue[messages.ModelResponseStreamEvent | Exception | None] = field(
+        default_factory=queue.Queue, init=False
+    )
     _thread: threading.Thread | None = field(default=None, init=False)
     _stream_response: StreamedResponse | None = field(default=None, init=False)
     _exception: Exception | None = field(default=None, init=False)
@@ -295,8 +297,8 @@ def __exit__(
     ) -> None:
         self._cleanup()
 
-    def __iter__(self) -> Iterator[messages.AgentStreamEvent]:
-        """Stream the response as an iterable of [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s."""
+    def __iter__(self) -> Iterator[messages.ModelResponseStreamEvent]:
+        """Stream the response as an iterable of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s."""
         self._check_context_manager_usage()
 
         while True:

pydantic_ai_slim/pydantic_ai/messages.py

@@ -1266,13 +1266,10 @@ class FinalResultEvent:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
-ModelResponseStreamEvent = Annotated[Union[PartStartEvent, PartDeltaEvent], pydantic.Discriminator('event_kind')]
-"""An event in the model response stream, either starting a new part or applying a delta to an existing one."""
-
-AgentStreamEvent = Annotated[
+ModelResponseStreamEvent = Annotated[
     Union[PartStartEvent, PartDeltaEvent, FinalResultEvent], pydantic.Discriminator('event_kind')
 ]
-"""An event in the agent stream."""
+"""An event in the model response stream, starting a new part, applying a delta to an existing one, or indicating the final result."""
 
 
 @dataclass(repr=False)
@@ -1342,3 +1339,6 @@ class BuiltinToolResultEvent:
     pydantic.Discriminator('event_kind'),
 ]
 """An event yielded when handling a model response, indicating tool calls and results."""
+
+AgentStreamEvent = Annotated[Union[ModelResponseStreamEvent, HandleResponseEvent], pydantic.Discriminator('event_kind')]
+"""An event in the agent stream: model response stream events and response-handling events."""

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -25,7 +25,6 @@
 from ..builtin_tools import AbstractBuiltinTool
 from ..exceptions import UserError
 from ..messages import (
-    AgentStreamEvent,
     FileUrl,
     FinalResultEvent,
     ModelMessage,
@@ -555,11 +554,11 @@ class StreamedResponse(ABC):
     final_result_event: FinalResultEvent | None = field(default=None, init=False)
 
     _parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
-    _event_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
+    _event_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
     _usage: RequestUsage = field(default_factory=RequestUsage, init=False)
 
-    def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
-        """Stream the response as an async iterable of [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s.
+    def __aiter__(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        """Stream the response as an async iterable of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s.
 
         This proxies the `_event_iterator()` and emits all events, while also checking for matches
         on the result schema and emitting a [`FinalResultEvent`][pydantic_ai.messages.FinalResultEvent] if/when the
@@ -569,7 +568,7 @@ def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
 
             async def iterator_with_final_event(
                 iterator: AsyncIterator[ModelResponseStreamEvent],
-            ) -> AsyncIterator[AgentStreamEvent]:
+            ) -> AsyncIterator[ModelResponseStreamEvent]:
                 async for event in iterator:
                     yield event
                     if (

pydantic_ai_slim/pydantic_ai/result.py

@@ -22,7 +22,7 @@
     ToolOutputSchema,
 )
 from ._run_context import AgentDepsT, RunContext
-from .messages import AgentStreamEvent
+from .messages import ModelResponseStreamEvent
 from .output import (
     OutputDataT,
     ToolOutput,
@@ -51,7 +51,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
     _usage_limits: UsageLimits | None
     _tool_manager: ToolManager[AgentDepsT]
 
-    _agent_stream_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
+    _agent_stream_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
     _initial_run_ctx_usage: RunUsage = field(init=False)
 
     def __post_init__(self):
@@ -221,8 +221,8 @@ async def _stream_text_deltas() -> AsyncIterator[str]:
                 deltas.append(text)
                 yield ''.join(deltas)
 
-    def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
-        """Stream [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s."""
+    def __aiter__(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        """Stream [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s."""
         if self._agent_stream_iterator is None:
             self._agent_stream_iterator = _get_usage_checking_stream_response(
                 self._raw_stream_response, self._usage_limits, self.usage
@@ -426,7 +426,7 @@ def _get_usage_checking_stream_response(
     stream_response: models.StreamedResponse,
     limits: UsageLimits | None,
     get_usage: Callable[[], RunUsage],
-) -> AsyncIterator[AgentStreamEvent]:
+) -> AsyncIterator[ModelResponseStreamEvent]:
     if limits is not None and limits.has_token_limits():
 
         async def _usage_checking_iterator():

tests/test_agent.py

@@ -29,7 +29,6 @@
 from pydantic_ai.messages import (
     AgentStreamEvent,
     BinaryContent,
-    HandleResponseEvent,
     ImageUrl,
     ModelMessage,
     ModelMessagesTypeAdapter,
@@ -4220,9 +4219,7 @@ def foo() -> str:
 
 
 async def test_wrapper_agent():
-    async def event_stream_handler(
-        ctx: RunContext[None], events: AsyncIterable[Union[AgentStreamEvent, HandleResponseEvent]]
-    ):
+    async def event_stream_handler(ctx: RunContext[None], events: AsyncIterable[AgentStreamEvent]):
         pass  # pragma: no cover
 
     foo_toolset = FunctionToolset()

tests/test_streaming.py

@@ -20,7 +20,6 @@
     FinalResultEvent,
     FunctionToolCallEvent,
     FunctionToolResultEvent,
-    HandleResponseEvent,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -1271,11 +1270,9 @@ async def test_run_event_stream_handler():
     async def ret_a(x: str) -> str:
         return f'{x}-apple'
 
-    events: list[AgentStreamEvent | HandleResponseEvent] = []
+    events: list[AgentStreamEvent] = []
 
-    async def event_stream_handler(
-        ctx: RunContext[None], stream: AsyncIterable[AgentStreamEvent | HandleResponseEvent]
-    ):
+    async def event_stream_handler(ctx: RunContext[None], stream: AsyncIterable[AgentStreamEvent]):
         async for event in stream:
             events.append(event)
 
@@ -1314,11 +1311,9 @@ def test_run_sync_event_stream_handler():
     async def ret_a(x: str) -> str:
         return f'{x}-apple'
 
-    events: list[AgentStreamEvent | HandleResponseEvent] = []
+    events: list[AgentStreamEvent] = []
 
-    async def event_stream_handler(
-        ctx: RunContext[None], stream: AsyncIterable[AgentStreamEvent | HandleResponseEvent]
-    ):
+    async def event_stream_handler(ctx: RunContext[None], stream: AsyncIterable[AgentStreamEvent]):
         async for event in stream:
             events.append(event)
 
@@ -1357,11 +1352,9 @@ async def test_run_stream_event_stream_handler():
     async def ret_a(x: str) -> str:
         return f'{x}-apple'
 
-    events: list[AgentStreamEvent | HandleResponseEvent] = []
+    events: list[AgentStreamEvent] = []
 
-    async def event_stream_handler(
-        ctx: RunContext[None], stream: AsyncIterable[AgentStreamEvent | HandleResponseEvent]
-    ):
+    async def event_stream_handler(ctx: RunContext[None], stream: AsyncIterable[AgentStreamEvent]):
         async for event in stream:
             events.append(event)
 

tests/test_temporal.py

@@ -19,7 +19,6 @@
     FinalResultEvent,
     FunctionToolCallEvent,
     FunctionToolResultEvent,
-    HandleResponseEvent,
     ModelMessage,
     ModelRequest,
     PartDeltaEvent,
@@ -196,7 +195,7 @@ class Deps(BaseModel):
 
 async def event_stream_handler(
     ctx: RunContext[Deps],
-    stream: AsyncIterable[AgentStreamEvent | HandleResponseEvent],
+    stream: AsyncIterable[AgentStreamEvent],
 ):
     logfire.info(f'{ctx.run_step=}')
     async for event in stream:
@@ -636,11 +635,11 @@ async def test_complex_agent_run_in_workflow(
 
 
 async def test_complex_agent_run(allow_model_requests: None):
-    events: list[AgentStreamEvent | HandleResponseEvent] = []
+    events: list[AgentStreamEvent] = []
 
     async def event_stream_handler(
         ctx: RunContext[Deps],
-        stream: AsyncIterable[AgentStreamEvent | HandleResponseEvent],
+        stream: AsyncIterable[AgentStreamEvent],
     ):
         async for event in stream:
             events.append(event)
@@ -1161,7 +1160,7 @@ async def test_temporal_agent_iter_in_workflow(allow_model_requests: None, clien
 
 async def simple_event_stream_handler(
     ctx: RunContext[None],
-    stream: AsyncIterable[AgentStreamEvent | HandleResponseEvent],
+    stream: AsyncIterable[AgentStreamEvent],
 ):
     pass