add _sync methods to StreamedRunResult

Author: ajac-zero

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -598,7 +598,7 @@ def run_stream_sync(
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
-    ) -> Iterator[result.SyncStreamedRunResult[AgentDepsT, Any]]:
+    ) -> Iterator[result.StreamedRunResult[AgentDepsT, Any]]:
         """Run the agent with a user prompt in sync streaming mode.
 
         This method builds an internal agent graph (using system prompts, tools and output schemas) and then
@@ -659,8 +659,7 @@ def main():
             builtin_tools=builtin_tools,
             event_stream_handler=event_stream_handler,
         )
-        async_result = get_event_loop().run_until_complete(async_cm.__aenter__())
-        yield result.SyncStreamedRunResult.from_streamed_result(async_result)  # type: ignore[reportReturnType]
+        yield get_event_loop().run_until_complete(async_cm.__aenter__())
 
     @overload
     def run_stream_events(

pydantic_ai_slim/pydantic_ai/result.py

@@ -410,6 +410,24 @@ async def stream_output(self, *, debounce_by: float | None = 0.1) -> AsyncIterat
         else:
             raise ValueError('No stream response or run result provided')  # pragma: no cover
 
+    def stream_output_sync(self, *, debounce_by: float | None = 0.1) -> Iterator[OutputDataT]:
+        """Stream the output as an iterable.
+
+        The pydantic validator for structured data will be called in
+        [partial mode](https://docs.pydantic.dev/dev/concepts/experimental/#partial-validation)
+        on each iteration.
+
+        Args:
+            debounce_by: by how much (if at all) to debounce/group the output chunks by. `None` means no debouncing.
+                Debouncing is particularly important for long structured outputs to reduce the overhead of
+                performing validation as each token is received.
+
+        Returns:
+            An iterable of the response data.
+        """
+        async_stream = self.stream_output(debounce_by=debounce_by)
+        yield from _lazy_async_iterator(async_stream)
+
     async def stream_text(self, *, delta: bool = False, debounce_by: float | None = 0.1) -> AsyncIterator[str]:
         """Stream the text result as an async iterable.
 
@@ -438,6 +456,22 @@ async def stream_text(self, *, delta: bool = False, debounce_by: float | None =
         else:
             raise ValueError('No stream response or run result provided')  # pragma: no cover
 
+    def stream_text_sync(self, *, delta: bool = False, debounce_by: float | None = 0.1) -> Iterator[str]:
+        """Stream the text result as an async iterable.
+
+        !!! note
+            Result validators will NOT be called on the text result if `delta=True`.
+
+        Args:
+            delta: if `True`, yield each chunk of text as it is received, if `False` (default), yield the full text
+                up to the current point.
+            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
+                Debouncing is particularly important for long structured responses to reduce the overhead of
+                performing validation as each token is received.
+        """
+        async_stream = self.stream_text(delta=delta, debounce_by=debounce_by)
+        yield from _lazy_async_iterator(async_stream)
+
     @deprecated('`StreamedRunResult.stream_structured` is deprecated, use `stream_responses` instead.')
     async def stream_structured(
         self, *, debounce_by: float | None = 0.1
@@ -473,6 +507,22 @@ async def stream_responses(
         else:
             raise ValueError('No stream response or run result provided')  # pragma: no cover
 
+    def stream_responses_sync(
+        self, *, debounce_by: float | None = 0.1
+    ) -> Iterator[tuple[_messages.ModelResponse, bool]]:
+        """Stream the response as an iterable of Structured LLM Messages.
+
+        Args:
+            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
+                Debouncing is particularly important for long structured responses to reduce the overhead of
+                performing validation as each token is received.
+
+        Returns:
+            An iterable of the structured response message and whether that is the last message.
+        """
+        async_stream = self.stream_responses(debounce_by=debounce_by)
+        yield from _lazy_async_iterator(async_stream)
+
     async def get_output(self) -> OutputDataT:
         """Stream the whole response, validate and return it."""
         if self._run_result is not None:
@@ -486,6 +536,10 @@ async def get_output(self) -> OutputDataT:
         else:
             raise ValueError('No stream response or run result provided')  # pragma: no cover
 
+    def get_output_sync(self) -> OutputDataT:
+        """Stream the whole response, validate and return it."""
+        return get_event_loop().run_until_complete(self.get_output())
+
     @property
     def response(self) -> _messages.ModelResponse:
         """Return the current state of the response."""
@@ -545,90 +599,6 @@ async def _marked_completed(self, message: _messages.ModelResponse | None = None
             await self._on_complete()
 
 
-@dataclass(init=False)
-class SyncStreamedRunResult(StreamedRunResult[AgentDepsT, OutputDataT]):
-    """Provides a synchronous API over 'StreamedRunResult'."""
-
-    @classmethod
-    def from_streamed_result(
-        cls, streamed_run_result: StreamedRunResult[AgentDepsT, OutputDataT]
-    ) -> SyncStreamedRunResult[AgentDepsT, OutputDataT]:
-        """Create a 'SyncStreamedRunResult' from an existing 'StreamedRunResult'."""
-        instance = cls.__new__(cls)
-
-        instance._all_messages = streamed_run_result._all_messages
-        instance._new_message_index = streamed_run_result._new_message_index
-        instance._stream_response = streamed_run_result._stream_response
-        instance._on_complete = streamed_run_result._on_complete
-        instance._run_result = streamed_run_result._run_result
-        instance.is_complete = streamed_run_result.is_complete
-
-        return instance
-
-    def _lazy_async_iterator(self, async_iter: AsyncIterator[T]) -> Iterator[T]:
-        """Lazily yield items from async iterator as they're requested."""
-        loop = get_event_loop()
-
-        while True:
-            try:
-                item = loop.run_until_complete(async_iter.__anext__())
-                yield item
-            except StopAsyncIteration:
-                break
-
-    def stream_output(self, *, debounce_by: float | None = 0.1) -> Iterator[OutputDataT]:  # type: ignore[reportIncompatibleMethodOverride]
-        """Stream the output as an iterable.
-
-        The pydantic validator for structured data will be called in
-        [partial mode](https://docs.pydantic.dev/dev/concepts/experimental/#partial-validation)
-        on each iteration.
-
-        Args:
-            debounce_by: by how much (if at all) to debounce/group the output chunks by. `None` means no debouncing.
-                Debouncing is particularly important for long structured outputs to reduce the overhead of
-                performing validation as each token is received.
-
-        Returns:
-            An iterable of the response data.
-        """
-        async_stream = super().stream_output(debounce_by=debounce_by)
-        yield from self._lazy_async_iterator(async_stream)
-
-    def stream_text(self, *, delta: bool = False, debounce_by: float | None = 0.1) -> Iterator[str]:  # type: ignore[reportIncompatibleMethodOverride]
-        """Stream the text result as an iterable.
-
-        !!! note
-            Result validators will NOT be called on the text result if `delta=True`.
-
-        Args:
-            delta: if `True`, yield each chunk of text as it is received, if `False` (default), yield the full text
-                up to the current point.
-            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
-                Debouncing is particularly important for long structured responses to reduce the overhead of
-                performing validation as each token is received.
-        """
-        async_stream = super().stream_text(delta=delta, debounce_by=debounce_by)
-        yield from self._lazy_async_iterator(async_stream)
-
-    def stream_responses(self, *, debounce_by: float | None = 0.1) -> Iterator[tuple[_messages.ModelResponse, bool]]:  # type: ignore[reportIncompatibleMethodOverride]
-        """Stream the response as an iterable of Structured LLM Messages.
-
-        Args:
-            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
-                Debouncing is particularly important for long structured responses to reduce the overhead of
-                performing validation as each token is received.
-
-        Returns:
-            An iterable of the structured response message and whether that is the last message.
-        """
-        async_stream = super().stream_responses(debounce_by=debounce_by)
-        yield from self._lazy_async_iterator(async_stream)
-
-    def get_output(self) -> OutputDataT:  # type: ignore[reportIncompatibleMethodOverride]
-        """Stream the whole response, validate and return it."""
-        return get_event_loop().run_until_complete(super().get_output())
-
-
 @dataclass(repr=False)
 class FinalResult(Generic[OutputDataT]):
     """Marker class storing the final output of an agent run and associated metadata."""
@@ -645,6 +615,17 @@ class FinalResult(Generic[OutputDataT]):
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
+def _lazy_async_iterator(async_iter: AsyncIterator[T]) -> Iterator[T]:
+    loop = get_event_loop()
+
+    while True:
+        try:
+            item = loop.run_until_complete(async_iter.__anext__())
+            yield item
+        except StopAsyncIteration:
+            break
+
+
 def _get_usage_checking_stream_response(
     stream_response: models.StreamedResponse,
     limits: UsageLimits | None,

tests/test_streaming.py

@@ -174,7 +174,7 @@ async def ret_a(x: str) -> str:
                 tool_calls=1,
             )
         )
-        response = result.get_output()
+        response = result.get_output_sync()
         assert response == snapshot('{"ret_a":"a-apple"}')
         assert result.is_complete
         assert result.timestamp() == IsNow(tz=timezone.utc)
@@ -389,20 +389,20 @@ def test_streamed_text_stream_sync():
 
     with agent.run_stream_sync('Hello') as result:
         # typehint to test (via static typing) that the stream type is correctly inferred
-        chunks: list[str] = [c for c in result.stream_text()]
+        chunks: list[str] = [c for c in result.stream_text_sync()]
         # one chunk with `stream_text()` due to group_by_temporal
         assert chunks == snapshot(['The cat sat on the mat.'])
         assert result.is_complete
 
     with agent.run_stream_sync('Hello') as result:
         # typehint to test (via static typing) that the stream type is correctly inferred
-        chunks: list[str] = [c for c in result.stream_output()]
+        chunks: list[str] = [c for c in result.stream_output_sync()]
         # two chunks with `stream()` due to not-final vs. final
         assert chunks == snapshot(['The cat sat on the mat.', 'The cat sat on the mat.'])
         assert result.is_complete
 
     with agent.run_stream_sync('Hello') as result:
-        assert [c for c in result.stream_text(debounce_by=None)] == snapshot(
+        assert [c for c in result.stream_text_sync(debounce_by=None)] == snapshot(
             [
                 'The ',
                 'The cat ',
@@ -415,20 +415,20 @@ def test_streamed_text_stream_sync():
 
     with agent.run_stream_sync('Hello') as result:
         # with stream_text, there is no need to do partial validation, so we only get the final message once:
-        assert [c for c in result.stream_text(delta=False, debounce_by=None)] == snapshot(
+        assert [c for c in result.stream_text_sync(delta=False, debounce_by=None)] == snapshot(
             ['The ', 'The cat ', 'The cat sat ', 'The cat sat on ', 'The cat sat on the ', 'The cat sat on the mat.']
         )
 
     with agent.run_stream_sync('Hello') as result:
-        assert [c for c in result.stream_text(delta=True, debounce_by=None)] == snapshot(
+        assert [c for c in result.stream_text_sync(delta=True, debounce_by=None)] == snapshot(
             ['The ', 'cat ', 'sat ', 'on ', 'the ', 'mat.']
         )
 
     def upcase(text: str) -> str:
         return text.upper()
 
     with agent.run_stream_sync('Hello', output_type=TextOutput(upcase)) as result:
-        assert [c for c in result.stream_output(debounce_by=None)] == snapshot(
+        assert [c for c in result.stream_output_sync(debounce_by=None)] == snapshot(
             [
                 'THE ',
                 'THE CAT ',
@@ -441,7 +441,7 @@ def upcase(text: str) -> str:
         )
 
     with agent.run_stream_sync('Hello') as result:
-        assert [c for c, _is_last in result.stream_responses(debounce_by=None)] == snapshot(
+        assert [c for c, _is_last in result.stream_responses_sync(debounce_by=None)] == snapshot(
             [
                 ModelResponse(
                     parts=[TextPart(content='The ')],