From 1c117a534e4ef0619caaebbdc9ca5dcdb87da1d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:22:10 -0700 Subject: [PATCH] fix(streaming): count every choice a chunk finishes when an n>1 stream drops before [DONE] --- .../litellm_core_utils/streaming_handler.py | 33 ++++++++-- .../test_streaming_handler.py | 61 +++++++++++++++++++ 2 files changed, 89 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2a95d02c866..8b7a4817a10 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -109,6 +109,27 @@ def _first_choice_index(chunk: ChatCompletionChunk | ModelResponseStream | Compl return 0 +def _choice_finish_reasons( + chunk: ChatCompletionChunk | ModelResponseStream | Completion | None, +) -> tuple[tuple[int, str | None], ...]: + match chunk: + case ChatCompletionChunk(choices=choices): + return tuple((choice.index, choice.finish_reason) for choice in choices) + case ModelResponseStream(choices=choices): + return tuple((choice.index, choice.finish_reason) for choice in choices) + case Completion(choices=choices): + return tuple((choice.index, choice.finish_reason) for choice in choices) + case _: + return () + + +def _indices_of_finished_choices( + chunk: ChatCompletionChunk | ModelResponseStream | Completion | None, +) -> frozenset[int]: + finished: Final = frozenset(index for index, finish_reason in _choice_finish_reasons(chunk) if finish_reason) + return finished or frozenset({_first_choice_index(chunk)}) + + def _stamp_usage_cost( usage_chunk: ModelResponse | TextCompletionResponse | ModelResponseStream, complete_response: ModelResponse | TextCompletionResponse | ModelResponseStream, @@ -1526,7 +1547,7 @@ class CustomStreamWrapper: print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] - self._record_finished_choice(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None) + self._record_finished_choices(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None) if response_obj["usage"] is not None: _text_completion_usage: Final[Usage] = response_obj["usage"] setattr( @@ -1566,7 +1587,7 @@ class CustomStreamWrapper: print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] - self._record_finished_choice(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None) + self._record_finished_choices(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None) elif self.custom_llm_provider == "cached_response": cached_chunk: Final = cast(ModelResponseStream, chunk) cached_choice: Final = cached_chunk.choices[0] if cached_chunk.choices else None @@ -1617,7 +1638,7 @@ class CustomStreamWrapper: f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}" ) self.received_finish_reason = response_obj["finish_reason"] - self._record_finished_choice(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None) + self._record_finished_choices(chunk if isinstance(chunk, _CHOICE_CHUNK_TYPES) else None) if response_obj.get("original_chunk", None) is not None: if hasattr(response_obj["original_chunk"], "id"): model_response = self.set_model_id(response_obj["original_chunk"].id, model_response) @@ -2276,8 +2297,10 @@ class CustomStreamWrapper: raise raise StopIteration from e - def _record_finished_choice(self, chunk: ChatCompletionChunk | ModelResponseStream | Completion | None) -> None: - self._finished_choice_indices = (self._finished_choice_indices or frozenset()) | {_first_choice_index(chunk)} + def _record_finished_choices(self, chunk: ChatCompletionChunk | ModelResponseStream | Completion | None) -> None: + self._finished_choice_indices = (self._finished_choice_indices or frozenset()) | _indices_of_finished_choices( + chunk + ) def _every_choice_finished(self) -> bool: if self.received_finish_reason is None: diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b1a10a9b0e2..c0383ffbd4f 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -5284,6 +5284,67 @@ async def test_two_choice_text_completion_stream_dropped_before_done_ends_cleanl assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "HiHel" +def _chat_choices_finishing_in_one_chunk() -> tuple[ModelResponseStream, ...]: + return ( + _reset_test_chunk(content="Hi", index=0), + _reset_test_chunk(content="Hel", index=1), + ModelResponseStream( + id="chatcmpl-reset-test", + created=1783458104, + model="stub-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop"), + StreamingChoices(index=1, delta=Delta(content=None), finish_reason="stop"), + ], + ), + ) + + +def _text_choices_finishing_in_one_chunk() -> tuple[Completion, ...]: + return ( + _text_completion_chunk(0, "Hi"), + _text_completion_chunk(1, "Hel"), + Completion.model_construct( + id="cmpl-two-choices", + object="text_completion", + created=1, + model="gpt-3.5-turbo-instruct", + choices=[ + CompletionChoice.model_construct(index=0, text="", finish_reason="stop", logprobs=None), + CompletionChoice.model_construct(index=1, text="", finish_reason="stop", logprobs=None), + ], + ), + ) + + +@pytest.mark.asyncio +@_SYNC_AND_ASYNC +@pytest.mark.parametrize( + ("provider", "chunks"), + (("openai", _chat_choices_finishing_in_one_chunk), ("text-completion-openai", _text_choices_finishing_in_one_chunk)), + ids=("chat", "text_completion"), +) +async def test_one_chunk_finishing_both_requested_choices_then_dropped_before_done_ends_cleanly( + logging_obj: Logging, + sync_mode: bool, + provider: str, + chunks: Callable[[], Sequence[ModelResponseStream | Completion]], +): + logging_obj.update_environment_variables(litellm_params={}, optional_params={"n": 2}) + stream: Final = _chunks_then_failure(chunks(), _httpx2_read_error()) + wrapper: Final = CustomStreamWrapper( + completion_stream=stream if sync_mode else _as_async(stream), + model="stub-model", + custom_llm_provider=provider, + logging_obj=logging_obj, + ) + + collected: Final = await _collect(wrapper, sync_mode) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in collected if chunk.choices) == "HiHel" + + @pytest.mark.asyncio @_SYNC_AND_ASYNC @_SDK3_DROPPED_CONNECTIONS