Make it easier to do sync streaming (#8183)

Author: chenmoneygithub

docs/docs/tutorials/streaming/index.md

@@ -393,3 +393,38 @@ Final output:  Prediction(
     sum='9'
 )
 ```
+
+## Synchronous Streaming
+
+By default calling a streamified DSPy program produces an async generator. In order to get back
+a sync generator, you can set the flag `async_streaming=False`:
+
+
+```python
+import os
+
+import dspy
+
+os.environ["OPENAI_API_KEY"] = "your_api_key"
+
+dspy.configure(lm=dspy.LM("openai/gpt-4o-mini"))
+
+predict = dspy.Predict("question->answer")
+
+# Enable streaming for the 'answer' field
+stream_predict = dspy.streamify(
+    predict,
+    stream_listeners=[dspy.streaming.StreamListener(signature_field_name="answer")],
+    async_streaming=False,
+)
+
+output = stream_predict(question="why did a chicken cross the kitchen?")
+
+program_output = None
+for chunk in output:
+    if isinstance(chunk, dspy.streaming.StreamResponse):
+        print(chunk)
+    elif isinstance(chunk, dspy.Prediction):
+        program_output = chunk
+print(f"Program output: {program_output}")
+```

dspy/streaming/streamify.py

@@ -30,6 +30,7 @@ def streamify(
     stream_listeners: Optional[List[StreamListener]] = None,
     include_final_prediction_in_output_stream: bool = True,
     is_async_program: bool = False,
+    async_streaming: bool = True,
 ) -> Callable[[Any, Any], Awaitable[Any]]:
     """
     Wrap a DSPy program so that it streams its outputs incrementally, rather than returning them
@@ -50,6 +51,8 @@ def streamify(
             still be included in the output stream even if this is `False`.
         is_async_program: Whether the program is async. If `False`, the program will be wrapped with `asyncify`,
             otherwise the program will be called with `acall`.
+        async_streaming: Whether to return an async generator or a sync generator. If `False`, the streaming will be
+            converted to a sync generator.
 
     Returns:
         A function that takes the same arguments as the original program, but returns an async
@@ -169,7 +172,7 @@ async def generator(args, kwargs, stream: MemoryObjectSendStream):
 
         await stream.send(prediction)
 
-    async def streamer(*args, **kwargs):
+    async def async_streamer(*args, **kwargs):
         send_stream, receive_stream = create_memory_object_stream(16)
         async with create_task_group() as tg, send_stream, receive_stream:
             tg.start_soon(generator, args, kwargs, send_stream)
@@ -201,7 +204,15 @@ async def streamer(*args, **kwargs):
                         yield value
                     return
 
-    return streamer
+    if async_streaming:
+        return async_streamer
+    else:
+
+        def sync_streamer(*args, **kwargs):
+            output = async_streamer(*args, **kwargs)
+            return apply_sync_streaming(output)
+
+        return sync_streamer
 
 
 def apply_sync_streaming(async_generator: AsyncGenerator) -> Generator:

tests/streaming/test_streaming.py

@@ -249,11 +249,11 @@ def __call__(self, x: str, **kwargs):
             dspy.streaming.StreamListener(signature_field_name="judgement"),
         ],
         include_final_prediction_in_output_stream=False,
+        async_streaming=False,
     )
     output = program(x="why did a chicken cross the kitchen?")
-    sync_output = dspy.streaming.apply_sync_streaming(output)
     all_chunks = []
-    for value in sync_output:
+    for value in output:
         if isinstance(value, dspy.streaming.StreamResponse):
             all_chunks.append(value)