From 6694c373a7a246462b5e8dbdd0728b5cc16a79db Mon Sep 17 00:00:00 2001 From: nazsats Date: Tue, 25 Aug 2026 18:33:32 +0530 Subject: [PATCH] fix(streaming): keep the n>1 rebuild inside litellm's lint budgets Extract the multi-choice branch into _rebuild_each_choice so stream_chunk_builder no longer calls itself, drop the setattr calls and casts that the strict-rule and type-discipline budgets flagged, and count tool calls and reasoning tokens in the usage input the way the single-choice path already does. --- .../streaming_chunk_builder_utils.py | 9 +- litellm/main.py | 110 ++++++++++++------ 2 files changed, 81 insertions(+), 38 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 5a566033540..9237ec61a7f 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -201,19 +201,20 @@ def _narrow(chunk: object, index: int) -> object | None: if not kept: return None if isinstance(chunk, dict): - return {**chunk, "choices": list(kept)} + 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)}) + 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: +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. """ - return [narrowed for chunk in chunks if (narrowed := _narrow(chunk, index)) is not None] + 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: diff --git a/litellm/main.py b/litellm/main.py index daad9df0582..20f20ad2e18 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8548,8 +8548,72 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N def _reindexed_choice(choice: Choices, index: int) -> Choices: - choice.index = index - return choice + 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( @@ -8587,39 +8651,17 @@ def stream_chunk_builder( 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 + 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, ) - 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) + if merged is not None: return merged # Initialize the response dictionary