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

fix(hosted_vllm): send the first choice's finish chunk before another choice's later token-id chunks

Browse source
This commit is contained in:
mateo-berri 2026-10-29 16:58:52 -07:00
parent 3c1c9cc43d
commit 7fae319038
2 changed files with 77 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
litellm/litellm_core_utils/streaming_handler.py
View file

@ -32,6 +32,7 @@ from litellm.litellm_core_utils.hidden_params import (
set_hidden_param,
stored_attribute,
)
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_items
from litellm.litellm_core_utils.model_response_utils import (
is_model_response_stream_empty,
)
@ -127,7 +128,7 @@ def _prepend_upstream_choice_extras(
_NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
_NO_FINISH_CHOICE_FIELDS: Final[Mapping[int, Mapping[str, object]]] = MappingProxyType({})
_NO_CHOICE_FIELDS_BY_INDEX: Final[Mapping[int, Mapping[str, object]]] = MappingProxyType({})
def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None:
@ -399,7 +400,8 @@ class CustomStreamWrapper:
self.system_fingerprint: str | None = None
self._provider_response_model: str | None = None
self.received_finish_reason: str | None = None
self.received_finish_choice_fields: Mapping[int, Mapping[str, object]] = _NO_FINISH_CHOICE_FIELDS
self.received_finish_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX
self.carried_chunk_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX
self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream
self.special_tokens = [
"<|assistant|>",
@ -1198,9 +1200,10 @@ class CustomStreamWrapper:
)
is_chunk_non_empty: Final = self.is_chunk_non_empty(completion_obj, model_response, response_obj)
chunk_choice_fields: Final[Mapping[str, object]] = response_obj.get(
upstream_chunk_choice_fields: Final[Mapping[str, object]] = response_obj.get(
"chunk_choice_fields", _NO_UPSTREAM_CHOICE_FIELDS
)
chunk_choice_fields: Final = self._take_carried_chunk_choice_fields(response_obj, upstream_chunk_choice_fields)
if is_chunk_non_empty: # cannot set content of an OpenAI Object to be an empty string
self.raise_on_model_repetition()
@ -1300,11 +1303,19 @@ class CustomStreamWrapper:
else:
return
elif self.received_finish_reason is not None:
another_choice_fields_chunk: Final = self._another_choice_fields_chunk(
model_response, response_obj, chunk_choice_fields
)
if another_choice_fields_chunk is not None:
return another_choice_fields_chunk
chunk_upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk"))
if chunk_choice_fields and self._streams_past_the_pending_finish(
model_response, response_obj, chunk_upstream_index
):
self.carried_chunk_choice_fields = MappingProxyType(
{**self.carried_chunk_choice_fields, chunk_upstream_index: chunk_choice_fields}
)
else:
another_choice_fields_chunk: Final = self._another_choice_fields_chunk(
model_response, response_obj, chunk_choice_fields
)
if another_choice_fields_chunk is not None:
return another_choice_fields_chunk
if self.sent_last_chunk is True:
# Bedrock returns the guardrail trace in the last chunk - we want to return this here
if self.custom_llm_provider == "bedrock" and "trace" in model_response:
@ -2022,6 +2033,37 @@ class CustomStreamWrapper:
def _apply_finish_choice_fields(self, choice: StreamingChoices) -> None:
_apply_choice_fields(choice, self.received_finish_choice_fields.get(choice.index, _NO_UPSTREAM_CHOICE_FIELDS))
def _streams_past_the_pending_finish(
self, model_response: ModelResponseStream, response_obj: Mapping[str, object], upstream_index: int
) -> bool:
finishing_index: Final = model_response.choices[0].index
return (
not self.sent_last_chunk
and not response_obj.get("is_finished")
and upstream_index != finishing_index
and finishing_index in self.received_finish_choice_fields
)
def _take_carried_chunk_choice_fields(
self, response_obj: Mapping[str, object], chunk_choice_fields: Mapping[str, object]
) -> 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"))
carried: Final = self.carried_chunk_choice_fields.get(upstream_index)
if carried is None:
return chunk_choice_fields
self.carried_chunk_choice_fields = MappingProxyType(
{index: fields for index, fields in self.carried_chunk_choice_fields.items() if index != upstream_index}
)
return MappingProxyType(
{
field: [*object_items(carried.get(field)), *object_items(chunk_choice_fields.get(field))]
for field in STREAM_CHUNK_CHOICE_FIELDS
if field in carried or field in chunk_choice_fields
}
)
def _another_choice_fields_chunk(
self,
model_response: ModelResponseStream,

27
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -6372,6 +6372,33 @@ def test_hosted_vllm_stream_with_two_choices_delivers_every_token_id_and_stop_re
) == frozenset(_TWO_CHOICE_STOP_REASONS.items())
def test_hosted_vllm_stream_with_two_choices_finishes_the_first_choice_before_the_other_choice_streams_on(
logging_obj: Logging,
) -> None:
emitted: Final = _hosted_vllm_stream_choices(
logging_obj,
(
*({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)),
{"index": 0, "delta": {"content": "1"}, "finish_reason": "stop", "stop_reason": ",2", "token_ids": [16]},
{"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [26]},
{"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [11]},
{"index": 1, "delta": {"content": "1,"}, "finish_reason": None, "token_ids": [27]},
{"index": 1, "delta": {"content": "3"}, "finish_reason": "stop", "stop_reason": 151645, "token_ids": [28]},
),
)
first_choice_finish: Final = next(
position for position, choice in enumerate(emitted) if choice["index"] == 0 and choice["finish_reason"]
)
assert emitted[first_choice_finish]["stop_reason"] == ",2"
assert _joined_token_ids([choice for choice in emitted[first_choice_finish:] if choice["index"] == 1]) == (
26,
11,
27,
28,
)
_CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19)