diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index ee0518c4aec..9237ec61a7f 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -173,6 +173,50 @@ 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)} # mutable-ok: a streamed chunk is a plain dict downstream. + if hasattr(chunk, "model_copy"): + return chunk.model_copy(update={"choices": list(kept)}) # mutable-ok: `choices` is a declared `list` field. + return chunk + + +def chunks_for_choice(chunks: Sequence[object], index: int) -> list: # mutable-ok: stream_chunk_builder wants a 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. + """ + narrowed_chunks: Final = (narrowed for chunk in chunks if (narrowed := _narrow(chunk, index)) is not None) + return list(narrowed_chunks) # mutable-ok: see the return annotation. + + 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 98f92e50599..aa2a22352c4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -185,7 +185,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 @@ -8587,6 +8591,75 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N return TextCompletionResponse(**response) +def _reindexed_choice(choice: Choices, index: int) -> Choices: + reindexed: Final = choice.model_copy() + reindexed.index = index + return reindexed + + +def _rebuild_each_choice( + *, + chunks: Sequence[object], + indices: tuple[int, ...], + model: str, + messages: list | None, # mutable-ok: forwarded verbatim to stream_chunk_builder, which declares `list | None`. + processor: ChunkProcessor, + start_time: datetime.datetime | None, + end_time: datetime.datetime | None, + logging_obj: Optional["Logging"], +) -> ModelResponse | None: + """Assemble an `n>1` streamed response one choice at a time. + + Everything in stream_chunk_builder assembles a single choice, reading + `chunks[...]["choices"][0]` throughout, so a stream carrying several choices + would otherwise have every choice's deltas concatenated into one message and + the rest dropped. Slicing per index and reusing that assembly keeps `n=1` on + exactly the path it takes today. + + Returns None when nothing could be assembled, leaving the caller on its + original path. + """ + 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 not assembled: + return None + + merged: Final = assembled[0][1] + if not isinstance(merged, ModelResponse): + return None + + reindexed: Final = tuple(_reindexed_choice(one.choices[0], index) for index, one in assembled) + merged.choices = list(reindexed) # mutable-ok: ModelResponse.choices is declared `list[Choices]`. + completion_output: Final[str] = get_content_from_model_response(merged) + merged.usage = processor.calculate_usage( # pyright: ignore[reportAttributeAccessIssue] # set in __init__, not declared + chunks=chunks, + model=model, + completion_output=completion_output, + messages=messages, + reasoning_tokens=processor.count_reasoning_tokens(merged), + ) + if litellm.include_cost_in_streaming_usage and logging_obj is not None: + merged.usage.cost = logging_obj._response_cost_calculator( # pyright: ignore[reportAttributeAccessIssue] # same as above + result=merged + ) + processor.apply_provider_assembled_streaming_metadata(merged, chunks, logging_obj) + return merged + + def stream_chunk_builder( chunks: list, messages: list | None = None, @@ -8619,6 +8692,22 @@ 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: + merged: Final = _rebuild_each_choice( + chunks=chunks, + indices=indices, + model=model, + messages=messages, + processor=processor, + start_time=start_time, + end_time=end_time, + logging_obj=logging_obj, + ) + if merged is not None: + 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"