From 11559ed8952d54a39a01efea9595d883ddbdb394 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 28 Oct 2026 13:21:31 +0100 Subject: [PATCH] fix(streaming): exact-match the SSE [DONE] marker in BaseModelResponseIterator (#47832) * fix(streaming): exact-match the SSE [DONE] marker in BaseModelResponseIterator * test(streaming): cover reasoning text and error frames that mention [DONE] --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/base_llm/base_model_iterator.py | 4 +- .../llms/base_llm/test_base_model_iterator.py | 101 +++++++++++++++++- 2 files changed, 103 insertions(+), 2 deletions(-) diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 611051bb2dd..032951be74a 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -3,6 +3,7 @@ from abc import abstractmethod from typing import TYPE_CHECKING, Final, cast import litellm +from litellm.constants import STREAM_SSE_DONE_STRING if TYPE_CHECKING: import httpx @@ -142,7 +143,8 @@ class BaseModelResponseIterator: def _handle_string_chunk(self, str_line: str) -> GenericStreamingChunk | ModelResponseStream: # chunk is a str at this point stripped_json_chunk: Final = BaseModelResponseIterator._string_to_dict_parser(str_line=str_line) - if "[DONE]" in str_line: + stripped_sse_chunk: Final = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(str_line) + if stripped_sse_chunk is not None and stripped_sse_chunk.strip() == STREAM_SSE_DONE_STRING: return GenericStreamingChunk( text="", is_finished=True, diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py index 204b458cadc..2971405885d 100644 --- a/tests/unit/llms/base_llm/test_base_model_iterator.py +++ b/tests/unit/llms/base_llm/test_base_model_iterator.py @@ -4,15 +4,19 @@ and non-string objects (e.g. Pydantic BaseModel events from the Responses API) p """ import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping +from itertools import chain from typing import Final import httpx import pytest from pydantic import BaseModel +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, convert_model_response_to_streaming from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler +from litellm.llms.openai.common_utils import OpenAIError from litellm.types.utils import Choices, GenericStreamingChunk, Message, ModelResponse @@ -336,3 +340,98 @@ def test_fake_stream_chunk_indexes_thinking_blocks_by_position_without_touching_ {"type": "thinking", "thinking": "Step 2.", "signature": "sig-b", "index": 2}, ] assert [block.get("index") for block in response.choices[0].message.thinking_blocks] == [None, None, None] + + +def _sse_chunk(delta: Mapping[str, object], finish_reason: str | None = None) -> str: + return "data: " + json.dumps( + { + "id": "chatcmpl-done", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-small-latest", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + ) + + +async def _aiter_lines(lines: tuple[str, ...]) -> AsyncIterator[str]: + for line in lines: + yield line + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_done_inside_streamed_json_is_content_not_the_end_of_stream(sync_mode: bool) -> None: + lines: Final = ( + _sse_chunk({"role": "assistant", "reasoning_content": "Mark it [DONE] once tests pass"}), + _sse_chunk({"content": "Checklist: "}), + _sse_chunk({"content": "- [DONE] write tests"}), + _sse_chunk({"content": " - [ ] ship it"}), + _sse_chunk( + { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "set_status", "arguments": ""}, + } + ] + } + ), + _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '{"status": '}}]}), + _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": '"[DONE]"'}}]}), + _sse_chunk({"tool_calls": [{"index": 0, "function": {"arguments": ', "n": 1}'}}]}), + _sse_chunk({}, "tool_calls"), + "data: [DONE]", + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=OpenAIChatCompletionStreamingHandler( + streaming_response=iter(lines) if sync_mode else _aiter_lines(lines), sync_stream=sync_mode + ), + model="mistral/mistral-small-latest", + custom_llm_provider="mistral", + logging_obj=Logging( + model="mistral/mistral-small-latest", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=1_790_000_000.0, + litellm_call_id="call-id", + function_id="function-id", + ), + ) + + chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper] + + deltas: Final = [chunk.choices[0].delta for chunk in chunks if chunk.choices] + tool_calls: Final = tuple(chain.from_iterable(delta.tool_calls or () for delta in deltas)) + assert "".join(getattr(delta, "reasoning_content", None) or "" for delta in deltas) == ( + "Mark it [DONE] once tests pass" + ) + assert "".join(delta.content or "" for delta in deltas) == "Checklist: - [DONE] write tests - [ ] ship it" + assert json.loads("".join(tool_call.function.arguments or "" for tool_call in tool_calls)) == { + "status": "[DONE]", + "n": 1, + } + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason] == [ + "tool_calls" + ] + + +@pytest.mark.parametrize("line", ["data: [DONE]", "data:[DONE]", "[DONE]", "data: [DONE]", "data: [DONE] "]) +def test_done_marker_ends_the_stream(line: str) -> None: + chunks: Final = list(BaseModelResponseIterator(streaming_response=iter((line,)), sync_stream=True)) + + assert chunks == [ + GenericStreamingChunk(text="", is_finished=True, finish_reason="stop", usage=None, index=0, tool_use=None) + ] + + +def test_error_frame_mentioning_done_raises_instead_of_ending_the_stream() -> None: + line: Final = 'data: {"error": {"message": "prompt too long [DONE]", "code": 400}}' + + with pytest.raises(OpenAIError, match=r"prompt too long \[DONE\]") as error: + list(OpenAIChatCompletionStreamingHandler(streaming_response=iter((line,)), sync_stream=True)) + + assert error.value.status_code == 400