scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-08 08:08:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

fix(streaming): keep each streamed text completion choice on its own index

Browse source
This commit is contained in:
mateo-berri 2026-10-29 20:30:48 -07:00
parent bc8635427f
commit dedd49dbd5
2 changed files with 36 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/litellm_core_utils/streaming_handler.py
View file

@ -152,6 +152,13 @@ def _upstream_choice_index(upstream_chunk: object) -> int:
return upstream_index if isinstance(upstream_index, int) else 0
def _response_choice_index(response_obj: Mapping[str, object]) -> int:
choice_index: Final = response_obj.get("choice_index")
if isinstance(choice_index, int):
return choice_index
return _upstream_choice_index(response_obj.get("original_chunk"))
def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]:
first_choice: Final = _upstream_first_choice(upstream_chunk)
upstream_extras: Final = first_choice.model_extra if first_choice is not None else None
@ -890,6 +897,7 @@ class CustomStreamWrapper:
"finish_reason": finish_reason,
"usage": usage,
"chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS),
"choice_index": _upstream_choice_index(chunk),
}
except Exception as e:
@ -1292,7 +1300,7 @@ class CustomStreamWrapper:
if response_obj.get("provider_specific_fields") is not None:
completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"]
model_response.choices[0].delta = Delta(**completion_obj)
_index: Final[int | None] = completion_obj.get("index")
_index: Final[int | None] = completion_obj.get("index", response_obj.get("choice_index"))
if _index is not None:
model_response.choices[0].index = _index
@ -1303,7 +1311,7 @@ class CustomStreamWrapper:
else:
return
elif self.received_finish_reason is not None:
chunk_upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk"))
chunk_upstream_index: Final = _response_choice_index(response_obj)
if chunk_choice_fields and self._streams_past_the_pending_finish(
model_response, response_obj, chunk_upstream_index
):
@ -1356,7 +1364,7 @@ class CustomStreamWrapper:
self._apply_finish_choice_fields(model_response.choices[0])
if (
response_obj.get("is_finished")
and _upstream_choice_index(_original_chunk) == model_response.choices[0].index
and _response_choice_index(response_obj) == model_response.choices[0].index
):
_apply_choice_fields(model_response.choices[0], chunk_choice_fields)
@ -1366,9 +1374,7 @@ class CustomStreamWrapper:
elif self._has_special_delta_content(model_response):
return self._handle_special_delta_content(model_response)
elif chunk_choice_fields:
return _choice_fields_chunk(
model_response, _upstream_choice_index(response_obj.get("original_chunk")), chunk_choice_fields
)
return _choice_fields_chunk(model_response, _response_choice_index(response_obj), chunk_choice_fields)
else:
if hasattr(model_response, "usage"):
self.chunks.append(model_response)
@ -2049,7 +2055,7 @@ class CustomStreamWrapper:
) -> Mapping[str, object]:
if not self.carried_chunk_choice_fields:
return chunk_choice_fields
upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk"))
upstream_index: Final = _response_choice_index(response_obj)
carried: Final = self.carried_chunk_choice_fields.get(upstream_index)
if carried is None:
return chunk_choice_fields
@ -2070,7 +2076,7 @@ class CustomStreamWrapper:
response_obj: Mapping[str, object],
chunk_choice_fields: Mapping[str, object],
) -> ModelResponseStream | None:
upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk"))
upstream_index: Final = _response_choice_index(response_obj)
is_finished: Final = bool(response_obj.get("is_finished"))
if is_finished and upstream_index == model_response.choices[0].index and not self.sent_last_chunk:
return None

22
tests/unit/text_completions/test_streaming_iterator.py
View file

@ -165,6 +165,28 @@ def test_text_completion_stream_sends_empty_text_on_token_id_only_chunks() -> No
]
def test_text_completion_stream_keeps_each_choices_text_and_token_ids_on_its_own_index() -> None:
responses: Final = _text_completion_stream(
(
_completion_chunk(" 5", None, index=0, token_ids=[20]),
_completion_chunk(" 7", None, index=1, token_ids=[22]),
_completion_chunk(", 6", None, index=0, token_ids=[11, 21]),
_completion_chunk("", None, index=1, token_ids=[11]),
_completion_chunk("", "stop", index=0, stop_reason=", 9", token_ids=[11, 220]),
_completion_chunk(", 8", "length", index=1, token_ids=[220, 23]),
)
)
choices: Final = tuple(chain.from_iterable(response.choices for response in responses))
assert {
index: (
"".join(choice.text or "" for choice in choices if choice.index == index),
list(chain.from_iterable(choice.get("token_ids") or () for choice in choices if choice.index == index)),
)
for index in (0, 1)
} == {0: (" 5, 6", [20, 11, 21, 11, 220]), 1: (" 7, 8", [22, 11, 220, 23])}
def test_text_completion_conversion_rejects_non_streaming_choices() -> None:
chunk: Final = Choices()
wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model")