Merge 4dd061f29e into ca9007be39

litellm/main.py

@@ -8587,6 +8587,87 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N
     return TextCompletionResponse(**response)
 
 
+def _streaming_choice_index(choice: StreamingChoices | Mapping[str, Any]) -> int | None:
+    index: Final = choice.get("index") if isinstance(choice, dict) else getattr(choice, "index", None)
+    return index if isinstance(index, int) else None
+
+
+def _streaming_chunk_choices(chunk: ModelResponseStream | Mapping[str, Any]) -> Sequence[Any]:
+    choices: Final = chunk.get("choices") if isinstance(chunk, dict) else getattr(chunk, "choices", None)
+    return choices or []  # mutable-ok: the caller iterates it; an empty list is the natural absent value
+
+
+def _distinct_streaming_choice_indices(chunks: Sequence[Any]) -> Sequence[int]:
+    indices: Final = {  # mutable-ok: a local set, discarded after sorting
+        _streaming_choice_index(choice) for chunk in chunks for choice in _streaming_chunk_choices(chunk)
+    }
+    return sorted(index for index in indices if index is not None)
+
+
+def _replace_chunk_choices(
+    chunk: ModelResponseStream | Mapping[str, Any], choices: Sequence[Any]
+) -> ModelResponseStream | Mapping[str, Any]:
+    if not choices:
+        return chunk
+    # Narrow on the model rather than on dict, so the mapping branch stays a Mapping.
+    if isinstance(chunk, ModelResponseStream):
+        return chunk.model_copy(update={"choices": list(choices)})  # mutable-ok: model_copy stores what it is given, and ModelResponse.choices is a list
+    return {**chunk, "choices": choices}  # mutable-ok: a new mapping so the caller's chunk is not mutated
+
+
+def _chunks_for_streaming_choice_index(
+    chunks: Sequence[Any], index: int
+) -> list[Any]:  # mutable-ok: fed straight to stream_chunk_builder, whose `chunks` parameter is a list
+    matched: Final = (
+        (chunk, [c for c in _streaming_chunk_choices(chunk) if _streaming_choice_index(c) == index])  # mutable-ok: per-chunk filtered choices, consumed by the comprehension below
+        for chunk in chunks
+    )
+    return [  # mutable-ok: stream_chunk_builder takes chunks as a list
+        _replace_chunk_choices(chunk, choices)
+        for chunk, choices in matched
+        if choices or not _streaming_chunk_choices(chunk)
+    ]
+
+
+def _renumbered_first_choice(response: object, index: int) -> Choices | None:
+    """The single choice a per-index build produced, renumbered to that index.
+
+    None when the build did not yield a usable non-streaming response: `stream_chunk_builder`
+    is typed to admit a streaming one, which leaves nothing to merge.
+    """
+    if not isinstance(response, ModelResponse) or not response.choices:
+        return None
+    # ModelResponse.choices is declared list[Choices], so the first entry needs no
+    # further narrowing; only the response itself can be the streaming variant.
+    return response.choices[0].model_copy(update={"index": index})  # mutable-ok: model_copy stores what it is given
+
+
+def _build_streaming_choices_per_index(
+    chunks: Sequence[Any],
+    indices: Sequence[int],
+    messages: list | None,  # mutable-ok: forwarded verbatim to stream_chunk_builder, whose `messages` is a list
+    start_time: datetime.datetime | None,
+    end_time: datetime.datetime | None,
+) -> list[Choices] | None:  # mutable-ok: assigned to ModelResponse.choices, which is a list
+    built: Final = [  # mutable-ok: one choice per index, returned as ModelResponse.choices
+        _renumbered_first_choice(
+            stream_chunk_builder(
+                chunks=_chunks_for_streaming_choice_index(chunks, index),
+                messages=messages,
+                start_time=start_time,
+                end_time=end_time,
+            ),
+            index,
+        )
+        for index in indices
+    ]
+    # One unusable build means there is nothing to merge, so the caller falls back to
+    # the single-response path rather than being handed a partial list.
+    if any(choice is None for choice in built):
+        return None
+    return [choice for choice in built if choice is not None]  # mutable-ok: assigned to ModelResponse.choices, which is a list
+
+
 def stream_chunk_builder(
     chunks: list,
     messages: list | None = None,
@@ -8618,6 +8699,13 @@ def stream_chunk_builder(
         ):  # route to the text completion logic
             return stream_chunk_builder_text_completion(chunks=chunks, messages=messages)
 
+        choice_indices: Final = _distinct_streaming_choice_indices(chunks)
+        per_choice: Final = (
+            _build_streaming_choices_per_index(chunks, choice_indices, messages, start_time, end_time)
+            if len(choice_indices) > 1
+            else None
+        )
+
         model: Final = chunks[0]["model"]
         # Initialize the response dictionary
         response: Final = processor.build_base_response(chunks)
@@ -8669,6 +8757,9 @@ def stream_chunk_builder(
             )
             setattr(response, "usage", usage)
 
+            if per_choice is not None:
+                response.choices = per_choice
+
             # Propagate provider_specific_fields from chunk hidden params when present.
             for chunk in reversed(chunks):
                 if isinstance(chunk, dict):
@@ -8842,6 +8933,9 @@ def stream_chunk_builder(
 
         setattr(response, "usage", usage)
 
+        if per_choice is not None:
+            response.choices = per_choice
+
         # Propagate provider_specific_fields from the last chunk (contains provider
         # metadata like traffic_type set during streaming)
         for chunk in reversed(chunks):

tests/test_litellm/test_main.py

@@ -1485,6 +1485,62 @@ async def test_async_mock_delay():
     assert delay >= 0.01
 
 
+def _n_choice_stream_chunk(index, content, finish_reason=None, role=None):
+    from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+    delta = {"content": content}
+    if role is not None:
+        delta["role"] = role
+    return ModelResponseStream(
+        id="chatcmpl-n2",
+        created=1,
+        model="gpt-4o",
+        object="chat.completion.chunk",
+        choices=[StreamingChoices(index=index, delta=Delta(**delta), finish_reason=finish_reason)],
+    )
+
+
+def test_stream_chunk_builder_keeps_n_greater_than_one_choices_separate():
+    from litellm import stream_chunk_builder
+
+    chunks = [
+        _n_choice_stream_chunk(0, "AAA", role="assistant"),
+        _n_choice_stream_chunk(1, "BBB", role="assistant"),
+        _n_choice_stream_chunk(0, " aaa"),
+        _n_choice_stream_chunk(1, " bbb"),
+        _n_choice_stream_chunk(0, None, finish_reason="stop"),
+        _n_choice_stream_chunk(1, None, finish_reason="length"),
+    ]
+
+    response = stream_chunk_builder(chunks, messages=[{"role": "user", "content": "hi"}])
+
+    assert response is not None
+    assert len(response.choices) == 2
+    assert [choice.index for choice in response.choices] == [0, 1]
+    assert response.choices[0].message.content == "AAA aaa"
+    assert response.choices[1].message.content == "BBB bbb"
+    assert response.choices[0].finish_reason == "stop"
+    assert response.choices[1].finish_reason == "length"
+
+
+def test_stream_chunk_builder_single_choice_is_unchanged():
+    from litellm import stream_chunk_builder
+
+    chunks = [
+        _n_choice_stream_chunk(0, "AAA", role="assistant"),
+        _n_choice_stream_chunk(0, " aaa"),
+        _n_choice_stream_chunk(0, None, finish_reason="stop"),
+    ]
+
+    response = stream_chunk_builder(chunks, messages=[{"role": "user", "content": "hi"}])
+
+    assert response is not None
+    assert len(response.choices) == 1
+    assert response.choices[0].index == 0
+    assert response.choices[0].message.content == "AAA aaa"
+    assert response.choices[0].finish_reason == "stop"
+
+
 def test_stream_chunk_builder_thinking_blocks():
     from litellm import stream_chunk_builder
     from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices