turn run_stream_sync into normal method

Author: ajac-zero

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -4,7 +4,7 @@
 import inspect
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
-from contextlib import AbstractAsyncContextManager, AbstractContextManager, asynccontextmanager, contextmanager
+from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from types import FrameType
 from typing import TYPE_CHECKING, Any, Generic, TypeAlias, cast, overload
 
@@ -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,
-    ) -> AbstractContextManager[result.StreamedRunResult[AgentDepsT, OutputDataT]]: ...
+    ) -> result.StreamedRunResult[AgentDepsT, OutputDataT]: ...
 
     @overload
     def run_stream_sync(
@@ -617,9 +617,8 @@ def run_stream_sync(
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         builtin_tools: Sequence[AbstractBuiltinTool] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
-    ) -> AbstractContextManager[result.StreamedRunResult[AgentDepsT, RunOutputDataT]]: ...
+    ) -> result.StreamedRunResult[AgentDepsT, RunOutputDataT]: ...
 
-    @contextmanager
     def run_stream_sync(
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
@@ -636,7 +635,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.StreamedRunResult[AgentDepsT, Any]]:
+    ) -> result.StreamedRunResult[AgentDepsT, Any]:
         """Run the agent with a user prompt in sync streaming mode.
 
         This is a convenience method that wraps [`self.run_stream`][pydantic_ai.agent.AbstractAgent.run_stream] with `loop.run_until_complete(...)`.
@@ -652,6 +651,18 @@ def run_stream_sync(
         If you want to always run the agent graph to completion and stream events and output at the same time,
         use [`agent.run()`][pydantic_ai.agent.AbstractAgent.run] with an `event_stream_handler` or [`agent.iter()`][pydantic_ai.agent.AbstractAgent.iter] instead.
 
+        Example:
+        ```python
+        from pydantic_ai import Agent
+
+        agent = Agent('openai:gpt-4o')
+
+        def main():
+            response = agent.run_stream_sync('What is the capital of the UK?')
+            print(response.get_output_sync())
+            #> The capital of the UK is London.
+        ```
+
         Args:
             user_prompt: User input to start/continue the conversation.
             output_type: Custom output type to use for this run, `output_type` may only be used if the agent has no
@@ -673,22 +684,26 @@ def run_stream_sync(
         Returns:
             The result of the run.
         """
-        async_cm = self.run_stream(
-            user_prompt,
-            output_type=output_type,
-            message_history=message_history,
-            deferred_tool_results=deferred_tool_results,
-            model=model,
-            deps=deps,
-            model_settings=model_settings,
-            usage_limits=usage_limits,
-            usage=usage,
-            infer_name=infer_name,
-            toolsets=toolsets,
-            builtin_tools=builtin_tools,
-            event_stream_handler=event_stream_handler,
-        )
-        yield get_event_loop().run_until_complete(async_cm.__aenter__())
+
+        async def _consume_stream():
+            async with self.run_stream(
+                user_prompt,
+                output_type=output_type,
+                message_history=message_history,
+                deferred_tool_results=deferred_tool_results,
+                model=model,
+                deps=deps,
+                model_settings=model_settings,
+                usage_limits=usage_limits,
+                usage=usage,
+                infer_name=infer_name,
+                toolsets=toolsets,
+                builtin_tools=builtin_tools,
+                event_stream_handler=event_stream_handler,
+            ) as stream_result:
+                yield stream_result
+
+        return get_event_loop().run_until_complete(_consume_stream().__anext__())
 
     @overload
     def run_stream_events(

pydantic_ai_slim/pydantic_ai/result.py

@@ -425,7 +425,7 @@ async def stream_output(self, *, debounce_by: float | None = 0.1) -> AsyncIterat
     def stream_output_sync(self, *, debounce_by: float | None = 0.1) -> Iterator[OutputDataT]:
         """Stream the output as an iterable.
 
-        This is a convenience method that wraps [`self.stream_output`][pydantic_ai.result.StreamedRunResult.stream_output] with `loop.run_until_complete(...)`.
+        This is a convenience method that wraps [`stream_output()`][pydantic_ai.result.StreamedRunResult.stream_output] with `loop.run_until_complete(...)`.
         You therefore can't use this method inside async code or if there's an active event loop.
 
         The pydantic validator for structured data will be called in
@@ -530,7 +530,7 @@ def stream_responses_sync(
     ) -> Iterator[tuple[_messages.ModelResponse, bool]]:
         """Stream the response as an iterable of Structured LLM Messages.
 
-        This is a convenience method that wraps [`self.stream_responses`][pydantic_ai.result.StreamedRunResult.stream_responses] with `loop.run_until_complete(...)`.
+        This is a convenience method that wraps [`stream_responses()`][pydantic_ai.result.StreamedRunResult.stream_responses] with `loop.run_until_complete(...)`.
         You therefore can't use this method inside async code or if there's an active event loop.
 
         Args:
@@ -560,7 +560,7 @@ async def get_output(self) -> OutputDataT:
     def get_output_sync(self) -> OutputDataT:
         """Stream the whole response, validate and return it.
 
-        This is a convenience method that wraps [`self.get_output`][pydantic_ai.result.StreamedRunResult.get_output] with `loop.run_until_complete(...)`.
+        This is a convenience method that wraps [`get_output()`][pydantic_ai.result.StreamedRunResult.get_output] with `loop.run_until_complete(...)`.
         You therefore can't use this method inside async code or if there's an active event loop.
         """
         return get_event_loop().run_until_complete(self.get_output())
@@ -642,11 +642,9 @@ class FinalResult(Generic[OutputDataT]):
 
 def _blocking_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
+            yield loop.run_until_complete(async_iter.__anext__())
         except StopAsyncIteration:
             break