From f943c86f312cdc39da7940382cf9174b3ead3960 Mon Sep 17 00:00:00 2001 From: nazsats Date: Tue, 25 Aug 2026 13:04:01 +0530 Subject: [PATCH] fix(streaming): rebuild each n>1 choice separately in stream_chunk_builder Reassembling a streamed response concatenated the deltas of every choice into one message and dropped the rest, so an n>1 request was logged, billed and guardrail-scanned as a single completion holding both answers. Split the chunks per choice index and reuse the existing assembly on each slice. n=1 keeps the original path unchanged. Usage stays request-level rather than per choice. Fixes #38180 --- .../streaming_chunk_builder_utils.py | 43 ++++++ litellm/main.py | 49 +++++- .../test_streaming_chunk_builder_utils.py | 146 ++++++++++++++++++ 3 files changed, 237 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index ee0518c4aec..5a566033540 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -173,6 +173,49 @@ def attach_cache_creation_token_details( return prompt_tokens_details.model_copy(update={"cache_creation_token_details": cache_creation_token_details}) +def _choices_of(chunk: object) -> Sequence[object]: + choices: Final = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None) + return choices or () + + +def _index_of(choice: object) -> int: + raw: Final = choice.get("index", 0) if isinstance(choice, dict) else getattr(choice, "index", 0) + return 0 if raw is None else int(raw) + + +def choice_indices(chunks: Sequence[object]) -> tuple[int, ...]: + """Distinct `choices[].index` across a streamed response, ascending. + + An absent index means 0, per the OpenAI schema. + """ + found: Final = frozenset(_index_of(choice) for chunk in chunks for choice in _choices_of(chunk)) + return tuple(sorted(found)) or (0,) + + +def _narrow(chunk: object, index: int) -> object | None: + choices: Final = _choices_of(chunk) + if not choices: + return chunk + + kept: Final = tuple(choice for choice in choices if _index_of(choice) == index) + if not kept: + return None + if isinstance(chunk, dict): + return {**chunk, "choices": list(kept)} + if hasattr(chunk, "model_copy"): + return chunk.model_copy(update={"choices": list(kept)}) + return chunk + + +def chunks_for_choice(chunks: Sequence[object], index: int) -> list: + """The same stream carrying only the choices for `index`. + + Chunks without choices, such as the usage-only chunk many providers send + last, are kept because usage belongs to the request rather than a choice. + """ + return [narrowed for chunk in chunks if (narrowed := _narrow(chunk, index)) is not None] + + class ChunkProcessor: def __init__(self, chunks: list, messages: list | None = None): self.chunks = self._sort_chunks(chunks) diff --git a/litellm/main.py b/litellm/main.py index d3967473f99..daad9df0582 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -184,7 +184,11 @@ from .litellm_core_utils.prompt_templates.factory import ( prompt_factory, stringify_json_tool_call_content, ) -from .litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor +from .litellm_core_utils.streaming_chunk_builder_utils import ( + ChunkProcessor, + choice_indices, + chunks_for_choice, +) from .llms.anthropic.chat import AnthropicChatCompletion from .llms.azure.audio_transcriptions import AzureAudioTranscription from .llms.azure.azure import AzureChatCompletion, _check_dynamic_azure_params @@ -8543,6 +8547,11 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N return TextCompletionResponse(**response) +def _reindexed_choice(choice: Choices, index: int) -> Choices: + choice.index = index + return choice + + def stream_chunk_builder( chunks: list, messages: list | None = None, @@ -8575,6 +8584,44 @@ def stream_chunk_builder( return stream_chunk_builder_text_completion(chunks=chunks, messages=messages) model: Final = chunks[0]["model"] + + indices: Final = choice_indices(chunks) + if len(indices) > 1: + rebuilt: Final = tuple( + ( + index, + stream_chunk_builder( + chunks=chunks_for_choice(chunks, index), + messages=messages, + start_time=start_time, + end_time=end_time, + logging_obj=None, + ), + ) + for index in indices + ) + assembled: Final = tuple( + (index, one) for index, one in rebuilt if one is not None and getattr(one, "choices", None) + ) + if assembled: + merged: Final = assembled[0][1] + merged.choices = [_reindexed_choice(cast(Choices, one.choices[0]), index) for index, one in assembled] + completion_output: Final[str] = "".join( + (getattr(choice.message, "content", None) or "") for choice in cast(list[Choices], merged.choices) + ) + merged_usage: Final = processor.calculate_usage( + chunks=chunks, + model=model, + completion_output=completion_output, + messages=messages, + reasoning_tokens=0, + ) + setattr(merged, "usage", merged_usage) + if litellm.include_cost_in_streaming_usage and logging_obj is not None: + setattr(merged_usage, "cost", logging_obj._response_cost_calculator(result=merged)) + processor.apply_provider_assembled_streaming_metadata(merged, chunks, logging_obj) + return merged + # Initialize the response dictionary response: Final = processor.build_base_response(chunks) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 4b5b51cb4b8..c7c24363c73 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1363,3 +1363,149 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate( assert usage.completion_tokens == 100 assert usage.completion_tokens_details.reasoning_tokens == expected_reasoning_tokens assert usage.completion_tokens_details.text_tokens == expected_text_tokens + + +N_CHOICE_MESSAGES = [{"role": "user", "content": "Name one colour. One word only."}] + + +def _n_choice_chunk(content, index, role=None, finish_reason=None, usage=None): + delta = {} + if role is not None: + delta["role"] = role + if content is not None: + delta["content"] = content + + payload = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [{"index": index, "delta": delta, "finish_reason": finish_reason}], + } + if usage is not None: + payload["usage"] = usage + return payload + + +def test_two_choices_rebuild_separately(): + chunks = [ + _n_choice_chunk("Azure", 0, role="assistant"), + _n_choice_chunk("Blue", 1, role="assistant"), + _n_choice_chunk(".", 0), + _n_choice_chunk(".", 1), + _n_choice_chunk(None, 0, finish_reason="stop"), + _n_choice_chunk(None, 1, finish_reason="stop"), + ] + + response = stream_chunk_builder(chunks, messages=N_CHOICE_MESSAGES) + + assert len(response.choices) == 2 + assert [c.index for c in response.choices] == [0, 1] + assert response.choices[0].message.content == "Azure." + assert response.choices[1].message.content == "Blue." + + +def test_choices_are_not_concatenated(): + """The specific regression: one choice holding both completions.""" + chunks = [ + _n_choice_chunk("first", 0, role="assistant"), + _n_choice_chunk("second", 1, role="assistant"), + _n_choice_chunk(None, 0, finish_reason="stop"), + _n_choice_chunk(None, 1, finish_reason="stop"), + ] + + response = stream_chunk_builder(chunks, messages=N_CHOICE_MESSAGES) + + contents = [c.message.content for c in response.choices] + # Each choice separately, not the join: two correct choices legitimately + # concatenate to "firstsecond" when you glue them together, so asserting on + # the joined string tests nothing. + assert all("firstsecond" != c for c in contents) + assert set(contents) == {"first", "second"} + + +def test_four_choices(): + chunks = [_n_choice_chunk(f"answer{i}", i, role="assistant") for i in range(4)] + chunks += [_n_choice_chunk(None, i, finish_reason="stop") for i in range(4)] + + response = stream_chunk_builder(chunks, messages=N_CHOICE_MESSAGES) + + assert len(response.choices) == 4 + assert [c.message.content for c in response.choices] == [ + "answer0", + "answer1", + "answer2", + "answer3", + ] + + +def test_usage_is_counted_once_for_the_request(): + """Usage belongs to the request, not to each choice. + + Each per-choice slice sees the provider's usage chunk, so a naive + implementation either reports one slice's figure or sums them. Neither is + right — the provider already reported the total for every completion. + """ + usage = {"prompt_tokens": 15, "completion_tokens": 4, "total_tokens": 19} + chunks = [ + _n_choice_chunk("Azure.", 0, role="assistant"), + _n_choice_chunk("Blue.", 1, role="assistant"), + _n_choice_chunk(None, 0, finish_reason="stop"), + _n_choice_chunk(None, 1, finish_reason="stop", usage=usage), + ] + + response = stream_chunk_builder(chunks, messages=N_CHOICE_MESSAGES) + + assert len(response.choices) == 2 + assert response.usage.prompt_tokens == 15 + assert response.usage.completion_tokens == 4 + assert response.usage.total_tokens == 19 + + +def test_finish_reason_is_per_choice(): + chunks = [ + _n_choice_chunk("short", 0, role="assistant"), + _n_choice_chunk("truncated", 1, role="assistant"), + _n_choice_chunk(None, 0, finish_reason="stop"), + _n_choice_chunk(None, 1, finish_reason="length"), + ] + + response = stream_chunk_builder(chunks, messages=N_CHOICE_MESSAGES) + + assert [c.finish_reason for c in response.choices] == ["stop", "length"] + + +@pytest.mark.parametrize("content", ["one", ""]) +def test_single_choice_is_unchanged(content): + """n=1 must take the original path and behave exactly as before.""" + chunks = [ + _n_choice_chunk(content, 0, role="assistant"), + _n_choice_chunk(None, 0, finish_reason="stop"), + ] + + response = stream_chunk_builder(chunks, messages=N_CHOICE_MESSAGES) + + assert len(response.choices) == 1 + assert response.choices[0].index == 0 + assert response.choices[0].message.content == content + assert response.choices[0].finish_reason == "stop" + + +def test_missing_index_counts_as_zero(): + """Some providers omit `index` on a chunk; it must not be dropped.""" + chunks = [ + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [{"delta": {"role": "assistant", "content": "hello"}, "finish_reason": None}], + }, + _n_choice_chunk(" world", 0), + _n_choice_chunk(None, 0, finish_reason="stop"), + ] + + response = stream_chunk_builder(chunks, messages=N_CHOICE_MESSAGES) + + assert len(response.choices) == 1 + assert response.choices[0].message.content == "hello world"