scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(streaming): keep the n>1 rebuild inside litellm's lint budgets

Browse source 
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.
This commit is contained in:
nazsats 2026-08-25 18:33:32 +05:30
parent f943c86f31
commit 6694c373a7
2 changed files with 81 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -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:

110
litellm/main.py
View file

@ -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