diff --git a/dspy/streaming/streaming_listener.py b/dspy/streaming/streaming_listener.py index 7d5650e543..99fb84ca47 100644 --- a/dspy/streaming/streaming_listener.py +++ b/dspy/streaming/streaming_listener.py @@ -345,9 +345,9 @@ def flush(self) -> str: boundary_index = len(last_tokens) return last_tokens[:boundary_index] elif isinstance(settings.adapter, ChatAdapter) or settings.adapter is None: - boundary_index = last_tokens.find("[[") - if boundary_index == -1: - boundary_index = len(last_tokens) + end_identifier = self.adapter_identifiers["ChatAdapter"]["end_identifier"] + boundary = end_identifier.search(last_tokens) + boundary_index = boundary.start() if boundary else len(last_tokens) return last_tokens[:boundary_index] else: raise ValueError( diff --git a/tests/streaming/test_streaming.py b/tests/streaming/test_streaming.py index 1233ae7dfe..2c9fa06932 100644 --- a/tests/streaming/test_streaming.py +++ b/tests/streaming/test_streaming.py @@ -1655,6 +1655,28 @@ async def multi_field_stream(*args, **kwargs): assert "second response" in second_content +@pytest.mark.parametrize("adapter", [None, dspy.ChatAdapter()]) +@pytest.mark.parametrize( + "content_chunks", + [ + ["Matrix: [[1, 2], [3, 4]]", "\n\n[[ ## completed ## ]]"], + ["Matrix: [", "[", "1, 2], [3, 4]]", "\n\n[[ ## completed ## ]]"], + ["Matrix: [[1, 2], [3, 4]]\n\n[[ ## completed ## ]]"], + ], +) +def test_stream_listener_preserves_literal_double_brackets(adapter, content_chunks): + listener = dspy.streaming.StreamListener(signature_field_name="answer") + chunks = ["[[ ## answer", " ## ]]\n\n", *content_chunks] + outputs = [] + with dspy.context(adapter=adapter): + for text in chunks: + chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content=text))]) + if output := listener.receive(chunk): + outputs.append(output) + assert "".join(output.chunk for output in outputs) == "Matrix: [[1, 2], [3, 4]]" + assert outputs[-1].is_last_chunk + + def test_stream_listener_could_form_end_identifier_chat_adapter(): listener = dspy.streaming.StreamListener(signature_field_name="answer")