Fix the error that status message streaming was blocking (#8432)

* Fix the status message streaming error

* fix tests

Author: chenmoneygithub

dspy/streaming/messages.py

@@ -1,7 +1,10 @@
 import asyncio
+import concurrent.futures
 from dataclasses import dataclass
 from typing import Any, Dict, Optional
 
+from asyncer import syncify
+
 from dspy.dsp.utils.settings import settings
 from dspy.utils.callback import BaseCallback
 
@@ -21,36 +24,29 @@ class StatusMessage:
 
 
 def sync_send_to_stream(stream, message):
-    """Send message to stream in a sync context, regardless of whether the caller is async or not."""
-    # Try to get current event loop, create one if none exists
-    try:
-        loop = asyncio.get_event_loop()
-    except RuntimeError:
-        # "There is no current event loop in thread" error
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
+    """Send message to stream in a sync context, regardless of event loop state."""
 
-    # If we're in an async context
-    if loop.is_running():
-        # In an async context, we need to use an approach that doesn't block
-        # Create a new thread and run a new event loop there
-        import concurrent.futures
+    async def _send():
+        await stream.send(message)
 
-        def run_async_in_new_loop():
+    try:
+        asyncio.get_running_loop()
+
+        # If we're in an event loop, offload to a new thread with its own event loop
+        def run_in_new_loop():
             new_loop = asyncio.new_event_loop()
             asyncio.set_event_loop(new_loop)
             try:
-                return new_loop.run_until_complete(stream.send(message))
+                return new_loop.run_until_complete(_send())
             finally:
                 new_loop.close()
 
-        # Run the function in a separate thread
         with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
-            future = executor.submit(run_async_in_new_loop)
-            return future.result()  # This shouldn't hang now
-    else:
-        # We're in a sync context, use run_until_complete
-        return loop.run_until_complete(stream.send(message))
+            future = executor.submit(run_in_new_loop)
+            return future.result()
+    except RuntimeError:
+        # Not in an event loop, safe to use a new event loop in this thread
+        return syncify(_send)()
 
 
 class StatusMessageProvider:

tests/streaming/test_streaming.py

@@ -1,4 +1,6 @@
+import asyncio
 import os
+import time
 from unittest import mock
 from unittest.mock import AsyncMock
 
@@ -649,3 +651,59 @@ async def gemini_stream_2(*args, **kwargs):
         assert all_chunks[1].predict_name == "predict2"
         assert all_chunks[1].signature_field_name == "judgement"
         assert all_chunks[1].chunk == "The answer provides a humorous and relevant punchline to the classic joke setup."
+
+
+@pytest.mark.anyio
+async def test_status_message_non_blocking():
+    def dummy_tool():
+        time.sleep(1)
+        return "dummy_tool_output"
+
+    class MyProgram(dspy.Module):
+        def forward(self, question, **kwargs):
+            dspy.Tool(dummy_tool)()
+            return dspy.Prediction(answer="dummy_tool_output")
+
+    program = dspy.streamify(MyProgram(), status_message_provider=StatusMessageProvider())
+
+    with mock.patch("litellm.acompletion", new_callable=AsyncMock, side_effect=[dummy_tool]):
+        with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False)):
+            output = program(question="why did a chicken cross the kitchen?")
+            timestamps = []
+            async for value in output:
+                if isinstance(value, dspy.streaming.StatusMessage):
+                    timestamps.append(time.time())
+
+    # timestamps[0]: tool start message
+    # timestamps[1]: tool end message
+    # There should be ~1 second delay between the tool start and end messages because we explicitly sleep for 1 second
+    # in the tool.
+    assert timestamps[1] - timestamps[0] >= 1
+
+
+@pytest.mark.anyio
+async def test_status_message_non_blocking_async_program():
+    async def dummy_tool():
+        await asyncio.sleep(1)
+        return "dummy_tool_output"
+
+    class MyProgram(dspy.Module):
+        async def aforward(self, question, **kwargs):
+            await dspy.Tool(dummy_tool).acall()
+            return dspy.Prediction(answer="dummy_tool_output")
+
+    program = dspy.streamify(MyProgram(), status_message_provider=StatusMessageProvider(), is_async_program=True)
+
+    with mock.patch("litellm.acompletion", new_callable=AsyncMock, side_effect=[dummy_tool]):
+        with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False)):
+            output = program(question="why did a chicken cross the kitchen?")
+            timestamps = []
+            async for value in output:
+                if isinstance(value, dspy.streaming.StatusMessage):
+                    timestamps.append(time.time())
+
+    # timestamps[0]: tool start message
+    # timestamps[1]: tool end message
+    # There should be ~1 second delay between the tool start and end messages because we explicitly sleep for 1 second
+    # in the tool.
+    assert timestamps[1] - timestamps[0] >= 1