Repository: stanfordnlp/dspy
Pull request #8461 (2025-06-25 20:45:45)

Allow reusing the stream listener
Add `allow_reuse` to stream listeners so that the same listener can be used multiple times. This feature is useful when streaming out the ReAct agent (or any programs that have a loop)

We turn it off by default because we don't want the stream chunk go to all listeners, which is not efficient. This downside is negligible with <5 listeners, but can be pretty bad if there are like 20 listeners configured in a complex system.


Patch:
diff --git a/dspy/streaming/streamify.py b/dspy/streaming/streamify.py
index 76861d1d..ab2201d4 100644
--- a/dspy/streaming/streamify.py
+++ b/dspy/streaming/streamify.py
@@ -183,7 +183,7 @@ def streamify(
                         # No listeners are configured, yield the chunk directly for backwards compatibility.
                         yield value
                     else:
-                        # We are receiving a chunk from the LM's response stream, delgate it to the listeners to
+                        # We are receiving a chunk from the LM's response stream, delegate it to the listeners to
                         # determine if we should yield a value to the user.
                         output = None
                         for listener in predict_id_to_listener[value.predict_id]:
diff --git a/dspy/streaming/streaming_listener.py b/dspy/streaming/streaming_listener.py
index 88d0925c..1cedeeae 100644
--- a/dspy/streaming/streaming_listener.py
+++ b/dspy/streaming/streaming_listener.py
@@ -17,7 +17,13 @@ if TYPE_CHECKING:
 class StreamListener:
     """Class that listens to the stream to capture the streeaming of a specific output field of a predictor."""
 
-    def __init__(self, signature_field_name: str, predict: Any = None, predict_name: str | None = None):
+    def __init__(
+        self,
+        signature_field_name: str,
+        predict: Any = None,
+        predict_name: str | None = None,
+        allow_reuse: bool = False,
+    ):
         """
         Args:
             signature_field_name: The name of the field to listen to.
@@ -25,6 +31,8 @@ class StreamListener:
                 the predictor that has the `signature_field_name` in its signature.
             predict_name: The name of the predictor to listen to. If None, when calling `streamify()` it will
                 automatically look for the predictor that has the `signature_field_name` in its signature.
+            allow_reuse: If True, the stream listener can be reused for multiple streams. Please note that this could
+                hurt the performance because the same stream chunk is sent to multiple listeners.
         """
         self.signature_field_name = signature_field_name
         self.predict = predict
@@ -35,6 +43,7 @@ class StreamListener:
         self.stream_start = False
         self.stream_end = False
         self.cache_hit = False
+        self.allow_reuse = allow_reuse
 
         self.json_adapter_start_identifier = f'"{self.signature_field_name}":'
         self.json_adapter_end_identifier = re.compile(r"\w*\"(,|\s*})")
@@ -53,7 +62,7 @@ class StreamListener:
             start_identifier = self.json_adapter_start_identifier
             end_identifier = self.json_adapter_end_identifier
 
-            start_indicator = "{"
+            start_indicator = '"'
         elif isinstance(settings.adapter, ChatAdapter) or settings.adapter is None:
             start_identifier = self.chat_adapter_start_identifier
             end_identifier = self.chat_adapter_end_identifier
@@ -66,7 +75,15 @@ class StreamListener:
             )
 
         if self.stream_end:
-            return
+            if self.allow_reuse:
+                # Clear up the state for the next stream.
+                self.stream_end = False
+                self.cache_hit = False
+                self.field_start_queue = []
+                self.field_end_queue = Queue()
+                self.stream_start = False
+            else:
+                return
 
         try:
             chunk_message = chunk.choices[0].delta.content


Test patch:
diff --git a/tests/streaming/test_streaming.py b/tests/streaming/test_streaming.py
index e9befaa3..d79402d2 100644
--- a/tests/streaming/test_streaming.py
+++ b/tests/streaming/test_streaming.py
@@ -707,3 +707,59 @@ async def test_status_message_non_blocking_async_program():
     # 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_stream_listener_allow_reuse():
+    class MyProgram(dspy.Module):
+        def __init__(self):
+            super().__init__()
+            self.predict = dspy.Predict("question->answer")
+
+        def forward(self, question, **kwargs):
+            self.predict(question=question, **kwargs)
+            return self.predict(question=question, **kwargs)
+
+    program = dspy.streamify(
+        MyProgram(),
+        stream_listeners=[
+            dspy.streaming.StreamListener(signature_field_name="answer", allow_reuse=True),
+        ],
+    )
+
+    async def gpt_4o_mini_stream(*args, **kwargs):
+        # Recorded streaming from openai/gpt-4o-mini
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[["))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" answer"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ]]\n\n"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="To"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" get"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" to"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" the"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" other"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" side"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="!"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="\n\n"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[[ ##"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" completed"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ##"))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ]]"))])
+
+    stream_generators = [gpt_4o_mini_stream, gpt_4o_mini_stream]
+
+    async def completion_side_effect(*args, **kwargs):
+        return stream_generators.pop(0)()  # return new async generator instance
+
+    with mock.patch("litellm.acompletion", side_effect=completion_side_effect):
+        with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False)):
+            output = program(question="why did a chicken cross the kitchen?")
+            all_chunks = []
+            async for value in output:
+                if isinstance(value, dspy.streaming.StreamResponse):
+                    all_chunks.append(value)
+
+    concat_message = "".join([chunk.chunk for chunk in all_chunks])
+    # The listener functions twice.
+    assert concat_message == "To get to the other side!To get to the other side!"