diff --git a/dspy/streaming/streamify.py b/dspy/streaming/streamify.py index 90a055baab..f32114c82c 100644 --- a/dspy/streaming/streamify.py +++ b/dspy/streaming/streamify.py @@ -177,6 +177,8 @@ async def generator(args, kwargs, stream: MemoryObjectSendStream): await stream.send(prediction) async def async_streamer(*args, **kwargs): + for listener in stream_listeners: + listener.reset() send_stream, receive_stream = create_memory_object_stream(16) async with create_task_group() as tg, send_stream, receive_stream: tg.start_soon(generator, args, kwargs, send_stream) diff --git a/dspy/streaming/streaming_listener.py b/dspy/streaming/streaming_listener.py index d38112ab63..1a9755f296 100644 --- a/dspy/streaming/streaming_listener.py +++ b/dspy/streaming/streaming_listener.py @@ -80,6 +80,22 @@ def __init__( }, } + def reset(self) -> None: + """Reset the run-time state of the listener so it can process a fresh stream. + + This clears the per-stream buffers and flags (start/end/cache) while preserving the + listener's configuration (signature field, predictor, and ``allow_reuse``). It is invoked + at the start of each ``streamify`` call so that a streamer callable can be invoked + multiple times; without it, listeners retain ``stream_end=True`` from the previous call + and silently drop all incremental chunks on reuse. + """ + self.stream_start = False + self.stream_end = False + self.cache_hit = False + self.field_start_queue = [] + self.field_end_queue = Queue() + self.json_adapter_state["field_accumulated_messages"] = "" + def _buffered_message_end_with_start_identifier(self, concat_message: str, start_identifier: str) -> str: for i in range(len(concat_message)): if start_identifier.startswith(concat_message[len(concat_message) - i - 1 :]): @@ -128,13 +144,7 @@ def receive(self, chunk: ModelResponseStream): if self.stream_end: 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.json_adapter_state["field_accumulated_messages"] = "" - self.stream_start = False + self.reset() else: return diff --git a/tests/streaming/test_streaming.py b/tests/streaming/test_streaming.py index 315aa5b1c7..2245b55fbb 100644 --- a/tests/streaming/test_streaming.py +++ b/tests/streaming/test_streaming.py @@ -949,6 +949,202 @@ async def completion_side_effect(*args, **kwargs): assert concat_message == "To get to the other side!To get to the other side!" +def _gpt_4o_mini_answer_stream(): + # Recorded streaming from openai/gpt-4o-mini for a single `question->answer` predict. + toks = [ + "[[", + " ##", + " answer", + " ##", + " ]]\n\n", + "To", + " get", + " to", + " the", + " other", + " side", + "!\n\n[[ ##", + " completed", + " ##", + " ]]", + ] + + async def gen(*args, **kwargs): + for t in toks: + yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=t))]) + + return gen + + +def _patched_acompletion(n_streams): + streams = [_gpt_4o_mini_answer_stream() for _ in range(n_streams)] + + async def side_effect(*args, **kwargs): + return streams.pop(0)() + + return mock.patch("litellm.acompletion", side_effect=side_effect) + + +@pytest.mark.anyio +async def test_streamer_reusable_default_config_streams_every_call(): + # Regression test: invoking a streamify-returned callable more than once must keep emitting + # incremental StreamResponse chunks on every call. Before the fix, the captured listeners + # retained stream_end=True from finalize() of the first call, so receive() early-returned + # (allow_reuse defaults to False) and all incremental chunks were silently dropped on call 2+. + class MyProgram(dspy.Module): + def __init__(self): + super().__init__() + self.predict = dspy.Predict("question->answer") + + def forward(self, question, **kwargs): + return self.predict(question=question, **kwargs) + + program = dspy.streamify( + MyProgram(), + stream_listeners=[dspy.streaming.StreamListener(signature_field_name="answer")], + ) + with ( + _patched_acompletion(3), + dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter()), + ): + results = [] + for _ in range(3): + n_chunks = 0 + has_prediction = False + async for value in program(question="why did a chicken cross the kitchen?"): + if isinstance(value, dspy.streaming.StreamResponse): + n_chunks += 1 + elif isinstance(value, dspy.Prediction): + has_prediction = True + results.append((n_chunks, has_prediction)) + assert results == [(7, True), (7, True), (7, True)] + + +@pytest.mark.anyio +async def test_streamer_reusable_with_final_prediction_suppressed(): + # Regression test for the worst-case tier: when include_final_prediction_in_output_stream=False + # is set at streamify() construction, reusing the streamer used to yield an EMPTY stream on call + # 2+ (no chunks, no Prediction) because the final-prediction guard depended on stale listener + # state. After the fix, every call should stream the field chunks and the final Prediction. + class MyProgram(dspy.Module): + def __init__(self): + super().__init__() + self.predict = dspy.Predict("question->answer") + + def forward(self, question, **kwargs): + return self.predict(question=question, **kwargs) + + program = dspy.streamify( + MyProgram(), + stream_listeners=[dspy.streaming.StreamListener(signature_field_name="answer")], + include_final_prediction_in_output_stream=False, + ) + with ( + _patched_acompletion(3), + dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter()), + ): + counts = [] + for _ in range(3): + n_values = 0 + async for _ in program(question="why did a chicken cross the kitchen?"): + n_values += 1 + counts.append(n_values) + assert counts == [7, 7, 7] + + +@pytest.mark.anyio +async def test_streamer_reusable_consecutive_calls_match_content(): + # Regression test: not only the count of chunks but their concatenated content must be + # identical across reused calls, ensuring listener buffers/queues are fully reset (not just + # the flags) so no stale tokens leak between calls. + class MyProgram(dspy.Module): + def __init__(self): + super().__init__() + self.predict = dspy.Predict("question->answer") + + def forward(self, question, **kwargs): + return self.predict(question=question, **kwargs) + + program = dspy.streamify( + MyProgram(), + stream_listeners=[dspy.streaming.StreamListener(signature_field_name="answer")], + ) + with ( + _patched_acompletion(2), + dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter()), + ): + messages = [] + for _ in range(2): + concat = [] + async for value in program(question="why did a chicken cross the kitchen?"): + if isinstance(value, dspy.streaming.StreamResponse): + concat.append(value.chunk) + messages.append("".join(concat)) + assert messages[0] == "To get to the other side!" + assert messages[1] == messages[0] + + +@pytest.mark.anyio +async def test_stream_listener_reset_clears_runtime_state(): + # Unit test for StreamListener.reset(): after a listener has been driven to stream_end, reset() + # must restore the same run-time state as a freshly constructed listener (excluding config). + listener = dspy.streaming.StreamListener(signature_field_name="answer") + listener.stream_start = True + listener.stream_end = True + listener.cache_hit = True + listener.field_start_queue = ["stale"] + listener.field_end_queue.put("stale") + listener.json_adapter_state["field_accumulated_messages"] = "stale" + + listener.reset() + + assert listener.stream_start is False + assert listener.stream_end is False + assert listener.cache_hit is False + assert listener.field_start_queue == [] + assert listener.field_end_queue.qsize() == 0 + assert listener.json_adapter_state["field_accumulated_messages"] == "" + # Configuration must be preserved across reset. + assert listener.signature_field_name == "answer" + assert listener.allow_reuse is False + + +@pytest.mark.anyio +async def test_streamer_reusable_allow_reuse_true_streams_every_call(): + # Regression test: allow_reuse=True (the documented intra-run reuse opt-in) must keep working + # for inter-call reuse too, and reset() at the start of each call must not regress intra-run + # reuse behavior. + class MyProgram(dspy.Module): + def __init__(self): + super().__init__() + self.predict = dspy.Predict("question->answer") + + def forward(self, question, **kwargs): + return self.predict(question=question, **kwargs) + + program = dspy.streamify( + MyProgram(), + stream_listeners=[ + dspy.streaming.StreamListener(signature_field_name="answer", allow_reuse=True), + ], + ) + with ( + _patched_acompletion(3), + dspy.context(lm=dspy.LM("openai/gpt-4o-mini", cache=False), adapter=dspy.ChatAdapter()), + ): + results = [] + for _ in range(3): + n_chunks = 0 + has_prediction = False + async for value in program(question="why did a chicken cross the kitchen?"): + if isinstance(value, dspy.streaming.StreamResponse): + n_chunks += 1 + elif isinstance(value, dspy.Prediction): + has_prediction = True + results.append((n_chunks, has_prediction)) + assert results == [(7, True), (7, True), (7, True)] + + @pytest.mark.anyio async def test_stream_listener_returns_correct_chunk_xml_adapter(): class MyProgram(dspy.Module): @@ -1305,9 +1501,7 @@ async def chat_stream(*args, **kwargs): yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="[[ ##"))]) yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" response"))]) yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ## ]]\n\n"))]) - yield ModelResponseStream( - model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="1"))] - ) + yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="1"))]) yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content="\n\n[[ ##"))]) yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" completed"))]) yield ModelResponseStream(model="gpt-4o-mini", choices=[StreamingChoices(delta=Delta(content=" ## ]]"))])