diff --git a/litellm/main.py b/litellm/main.py index 8ee102f5d07..cc131f6947f 100644 --- a/litellm/main.py +++ b/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): diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 3eea47bcd5a..bc14451b732 100644 --- a/tests/test_litellm/test_main.py +++ b/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