fix(dspy): keep trailing [ and drop unfinished headers followed by whitespace

Address review: a lone trailing "[" is part of the value, and an unfinished
field header followed by whitespace at the end of the stream is still cut.

Amp-Thread-ID: https://ampcode.com/threads/T-01a12205-cfb6-739c-a9fb-cbe0a0a86db3
Co-authored-by: Isaac Miller <isaac@cmpnd.ai>

Author: ampagent

dspy/streaming/streaming_listener.py

@@ -23,8 +23,9 @@
 
 ADAPTER_SUPPORT_STREAMING = [ChatAdapter, XMLAdapter, JSONAdapter]
 
-# Matches an unfinished "[[ ## <field> ## ]]" at the end of the text, from "[" up to "[[ ## <field> ## ]".
-_CHAT_END_IDENTIFIER_PREFIX = re.compile(r"\[(?:\[(?: (?:#(?:#(?: (?:\w+(?: (?:#(?:#(?: \]?)?)?)?)?)?)?)?)?)?)?\Z")
+# Matches an unfinished "[[ ## <field> ## ]]" at the end of the text, from "[[" up to "[[ ## <field> ## ]", plus any
+# trailing whitespace. A lone "[" is kept because it is more likely part of the value.
+_CHAT_END_IDENTIFIER_PREFIX = re.compile(r"\[\[(?: (?:#(?:#(?: (?:\w+(?: (?:#(?:#(?: \]?)?)?)?)?)?)?)?)?)?\s*\Z")
 
 
 class StreamListener:

tests/streaming/test_streaming.py

@@ -368,6 +368,8 @@ async def gpt_4o_mini_stream(*args, **kwargs):
         ["[[ ## answer ## ]]\n", "See also ", "[[Matrices]]\n\n[[ ## comp", "leted ## ]]"],
         # The end marker is split across chunks.
         ["[[ ## answer ## ]]\n", "m = [[1, 2]]", "\n\n[", "[ #", "# comp", "leted #", "# ]", "]"],
+        # A trailing "[" with no header after it is part of the value.
+        ["[[ ## answer ## ]]\n", "Use ["],
     ],
 )
 async def test_streaming_chunks_join_to_parsed_value_regardless_of_chunking(parts):
@@ -393,6 +395,29 @@ async def stream(*args, **kwargs):
     assert "".join([chunk.chunk for chunk in all_chunks]) == final.answer
 
 
+@pytest.mark.anyio
+@pytest.mark.parametrize("last_part", ["\n\n[[ ## comp", "\n\n[[ ## comp\n", "\n\n[[\n"])
+async def test_streaming_drops_unfinished_field_header_at_end_of_stream(last_part):
+    program = dspy.streamify(
+        dspy.Predict("question->answer"),
+        stream_listeners=[dspy.streaming.StreamListener(signature_field_name="answer")],
+    )
+
+    async def stream(*args, **kwargs):
+        for part in ["[[ ## answer ## ]]\n", "abc", last_part]:
+            yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=part))])
+
+    with mock.patch("litellm.acompletion", side_effect=stream):
+        with dspy.context(lm=dspy.LM("openai/gpt-4o-mini", engine="litellm", cache=False), adapter=dspy.ChatAdapter()):
+            chunks = [
+                value.chunk
+                async for value in program(question="What is the capital of France?")
+                if isinstance(value, dspy.streaming.StreamResponse)
+            ]
+
+    assert "".join(chunks) == "abc"
+
+
 @pytest.mark.llm_call
 def test_sync_streaming(lm_for_test):
     class MyProgram(dspy.Module):