Fix missing space in the stream listener (#8360)

Author: chenmoneygithub

dspy/streaming/streaming_listener.py

@@ -99,7 +99,7 @@ def receive(self, chunk: ModelResponseStream):
             # We keep appending the tokens to the queue until we have a full identifier or the concanated
             # tokens no longer match our expected identifier.
             self.field_start_queue.append(chunk_message)
-            concat_message = "".join(self.field_start_queue).strip()
+            concat_message = "".join(self.field_start_queue)
 
             if start_identifier in concat_message:
                 # We have a full identifier, we can start the stream.
@@ -113,8 +113,9 @@ def receive(self, chunk: ModelResponseStream):
                     # because there could be a few splitters between ':' and '"', e.g., '"name": "value"'.
                     chunk_message = chunk_message[1:]
 
-            elif self._buffered_message_end_with_start_identifier(concat_message, start_identifier):
-                # If the buffered message ends with part of the start_identifier, we can start the stream.
+            elif self._buffered_message_end_with_start_identifier(concat_message.strip(), start_identifier):
+                # If the buffered message ends with part of the start_identifier, we keep looking for the
+                # start_identifier from the token stream.
                 return
             else:
                 # Doesn't match the expected identifier, reset the queue.

tests/streaming/test_streaming.py

@@ -229,6 +229,36 @@ def __call__(self, x: str, **kwargs):
     assert all_chunks[-1].signature_field_name == "judgement"
 
 
+@pytest.mark.anyio
+async def test_streaming_handles_space_correctly():
+    dspy.settings.configure(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter())
+    my_program = dspy.Predict("question->answer")
+    program = dspy.streamify(
+        my_program, stream_listeners=[dspy.streaming.StreamListener(signature_field_name="answer")]
+    )
+
+    async def gpt_4o_mini_stream(*args, **kwargs):
+        yield ModelResponseStream(
+            model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[[ ## answer ## ]]\n"))]
+        )
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="How "))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="are "))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="you "))])
+        yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="doing?"))])
+        yield ModelResponseStream(
+            model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="\n\n[[ ## completed ## ]]"))]
+        )
+
+    with mock.patch("litellm.acompletion", side_effect=gpt_4o_mini_stream):
+        output = program(question="What is the capital of France?")
+        all_chunks = []
+        async for value in output:
+            if isinstance(value, dspy.streaming.StreamResponse):
+                all_chunks.append(value)
+
+    assert all_chunks[0].chunk == "How are you doing?"
+
+
 @pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OpenAI API key not found in environment variables")
 def test_sync_streaming():
     class MyProgram(dspy.Module):