Merge 6694c373a7 into 176b2e5eb8

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)

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)
 

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"