From dedd49dbd535b43698dcdb2676360a391c74ac21 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:30:48 -0700 Subject: [PATCH] fix(streaming): keep each streamed text completion choice on its own index --- .../litellm_core_utils/streaming_handler.py | 22 ++++++++++++------- .../test_streaming_iterator.py | 22 +++++++++++++++++++ 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7875cdb0df1..5c1f354fe4b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -152,6 +152,13 @@ def _upstream_choice_index(upstream_chunk: object) -> int: return upstream_index if isinstance(upstream_index, int) else 0 +def _response_choice_index(response_obj: Mapping[str, object]) -> int: + choice_index: Final = response_obj.get("choice_index") + if isinstance(choice_index, int): + return choice_index + return _upstream_choice_index(response_obj.get("original_chunk")) + + def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: first_choice: Final = _upstream_first_choice(upstream_chunk) upstream_extras: Final = first_choice.model_extra if first_choice is not None else None @@ -890,6 +897,7 @@ class CustomStreamWrapper: "finish_reason": finish_reason, "usage": usage, "chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS), + "choice_index": _upstream_choice_index(chunk), } except Exception as e: @@ -1292,7 +1300,7 @@ class CustomStreamWrapper: if response_obj.get("provider_specific_fields") is not None: completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] model_response.choices[0].delta = Delta(**completion_obj) - _index: Final[int | None] = completion_obj.get("index") + _index: Final[int | None] = completion_obj.get("index", response_obj.get("choice_index")) if _index is not None: model_response.choices[0].index = _index @@ -1303,7 +1311,7 @@ class CustomStreamWrapper: else: return elif self.received_finish_reason is not None: - chunk_upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk")) + chunk_upstream_index: Final = _response_choice_index(response_obj) if chunk_choice_fields and self._streams_past_the_pending_finish( model_response, response_obj, chunk_upstream_index ): @@ -1356,7 +1364,7 @@ class CustomStreamWrapper: self._apply_finish_choice_fields(model_response.choices[0]) if ( response_obj.get("is_finished") - and _upstream_choice_index(_original_chunk) == model_response.choices[0].index + and _response_choice_index(response_obj) == model_response.choices[0].index ): _apply_choice_fields(model_response.choices[0], chunk_choice_fields) @@ -1366,9 +1374,7 @@ class CustomStreamWrapper: elif self._has_special_delta_content(model_response): return self._handle_special_delta_content(model_response) elif chunk_choice_fields: - return _choice_fields_chunk( - model_response, _upstream_choice_index(response_obj.get("original_chunk")), chunk_choice_fields - ) + return _choice_fields_chunk(model_response, _response_choice_index(response_obj), chunk_choice_fields) else: if hasattr(model_response, "usage"): self.chunks.append(model_response) @@ -2049,7 +2055,7 @@ class CustomStreamWrapper: ) -> Mapping[str, object]: if not self.carried_chunk_choice_fields: return chunk_choice_fields - upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk")) + upstream_index: Final = _response_choice_index(response_obj) carried: Final = self.carried_chunk_choice_fields.get(upstream_index) if carried is None: return chunk_choice_fields @@ -2070,7 +2076,7 @@ class CustomStreamWrapper: response_obj: Mapping[str, object], chunk_choice_fields: Mapping[str, object], ) -> ModelResponseStream | None: - upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk")) + upstream_index: Final = _response_choice_index(response_obj) is_finished: Final = bool(response_obj.get("is_finished")) if is_finished and upstream_index == model_response.choices[0].index and not self.sent_last_chunk: return None diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index 38d7f99e449..402e3302eeb 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -165,6 +165,28 @@ def test_text_completion_stream_sends_empty_text_on_token_id_only_chunks() -> No ] +def test_text_completion_stream_keeps_each_choices_text_and_token_ids_on_its_own_index() -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk(" 5", None, index=0, token_ids=[20]), + _completion_chunk(" 7", None, index=1, token_ids=[22]), + _completion_chunk(", 6", None, index=0, token_ids=[11, 21]), + _completion_chunk("", None, index=1, token_ids=[11]), + _completion_chunk("", "stop", index=0, stop_reason=", 9", token_ids=[11, 220]), + _completion_chunk(", 8", "length", index=1, token_ids=[220, 23]), + ) + ) + choices: Final = tuple(chain.from_iterable(response.choices for response in responses)) + + assert { + index: ( + "".join(choice.text or "" for choice in choices if choice.index == index), + list(chain.from_iterable(choice.get("token_ids") or () for choice in choices if choice.index == index)), + ) + for index in (0, 1) + } == {0: (" 5, 6", [20, 11, 21, 11, 220]), 1: (" 7, 8", [22, 11, 220, 23])} + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model")