add lazy implementation

Author: ajac-zero

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -586,12 +586,13 @@ 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.CollectedRunResult[AgentDepsT, Any]]:
-        """Run the agent with a user prompt in collected streaming mode.
+    ) -> Iterator[result.SyncStreamedRunResult[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
         runs the graph until the model produces output matching the `output_type`, for example text or structured data.
-        At this point, a streaming run result object is collected and -- once this output has completed streaming -- you can iterate over the complete output, message history, and usage.
+        At this point, a streaming run result object is yielded from which you can stream the output as it comes in,
+        and -- once this output has completed streaming -- get the complete output, message history, and usage.
 
         As this method will consider the first output matching the `output_type` to be the final output,
         it will stop running the agent graph and will not execute any tool calls made by the model after this "final" output.
@@ -647,7 +648,7 @@ def main():
             event_stream_handler=event_stream_handler,
         )
         async_result = get_event_loop().run_until_complete(async_cm.__aenter__())
-        yield result.CollectedRunResult.from_streamed_result(async_result)  # type: ignore[reportReturnType]
+        yield result.SyncStreamedRunResult.from_streamed_result(async_result)  # type: ignore[reportReturnType]
 
     @overload
     def run_stream_events(

pydantic_ai_slim/pydantic_ai/result.py

@@ -546,14 +546,14 @@ async def _marked_completed(self, message: _messages.ModelResponse | None = None
 
 
 @dataclass(init=False)
-class CollectedRunResult(StreamedRunResult[AgentDepsT, OutputDataT]):
-    """Provides a synchronous API over 'StreamedRunResult' by eagerly loading the stream."""
+class SyncStreamedRunResult(StreamedRunResult[AgentDepsT, OutputDataT]):
+    """Provides a synchronous API over 'StreamedRunResult'."""
 
     @classmethod
     def from_streamed_result(
         cls, streamed_run_result: StreamedRunResult[AgentDepsT, OutputDataT]
-    ) -> CollectedRunResult[AgentDepsT, OutputDataT]:
-        """Create a CollectedRunResult from an existing StreamedRunResult."""
+    ) -> SyncStreamedRunResult[AgentDepsT, OutputDataT]:
+        """Create a 'SyncStreamedRunResult' from an existing 'StreamedRunResult'."""
         instance = cls.__new__(cls)
 
         instance._all_messages = streamed_run_result._all_messages
@@ -565,14 +565,19 @@ def from_streamed_result(
 
         return instance
 
-    def _collect_async_iterator(self, async_iter: AsyncIterator[T]) -> list[T]:
-        async def collect():
-            return [item async for item in async_iter]
+    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()
 
-        return get_event_loop().run_until_complete(collect())
+        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]
-        """Collect and stream the output as an iterable.
+        """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)
@@ -587,10 +592,10 @@ def stream_output(self, *, debounce_by: float | None = 0.1) -> Iterator[OutputDa
             An iterable of the response data.
         """
         async_stream = super().stream_output(debounce_by=debounce_by)
-        yield from self._collect_async_iterator(async_stream)
+        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]
-        """Collect and stream the text result as an iterable.
+        """Stream the text result as an iterable.
 
         !!! note
             Result validators will NOT be called on the text result if `delta=True`.
@@ -603,10 +608,10 @@ def stream_text(self, *, delta: bool = False, debounce_by: float | None = 0.1) -
                 performing validation as each token is received.
         """
         async_stream = super().stream_text(delta=delta, debounce_by=debounce_by)
-        yield from self._collect_async_iterator(async_stream)
+        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]
-        """Collect and stream the response as an iterable of Structured LLM Messages.
+        """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.
@@ -617,7 +622,7 @@ def stream_responses(self, *, debounce_by: float | None = 0.1) -> Iterator[tuple
             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._collect_async_iterator(async_stream)
+        yield from self._lazy_async_iterator(async_stream)
 
     def get_output(self) -> OutputDataT:  # type: ignore[reportIncompatibleMethodOverride]
         """Stream the whole response, validate and return it."""