From 7fae31903858f5e4a5a79a15413a591fa7122567 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 16:58:52 -0700 Subject: [PATCH] fix(hosted_vllm): send the first choice's finish chunk before another choice's later token-id chunks --- .../litellm_core_utils/streaming_handler.py | 58 ++++++++++++++++--- .../test_streaming_handler.py | 27 +++++++++ 2 files changed, 77 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index ac24859d4b4..7875cdb0df1 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 38400c5960f..f3656783e7e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -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)