diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index acac00ce345..ac24859d4b4 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -166,6 +166,15 @@ def _apply_choice_fields(choice: StreamingChoices, fields: Mapping[str, object]) setattr(choice, field, value) +def _choice_fields_chunk( + model_response: ModelResponseStream, upstream_index: int, fields: Mapping[str, object] +) -> ModelResponseStream: + model_response.choices[0].delta = Delta(content="") # rebind-ok: chunks are built in place here + model_response.choices[0].index = upstream_index # rebind-ok: chunks are built in place here + _apply_choice_fields(model_response.choices[0], fields) + return model_response + + def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool: return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields) @@ -1291,6 +1300,11 @@ 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 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: @@ -1341,11 +1355,9 @@ class CustomStreamWrapper: elif self._has_special_delta_content(model_response): return self._handle_special_delta_content(model_response) elif chunk_choice_fields: - upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk")) - model_response.choices[0].delta = Delta(content="") # rebind-ok: chunks are built in place here - model_response.choices[0].index = upstream_index # rebind-ok: chunks are built in place here - _apply_choice_fields(model_response.choices[0], chunk_choice_fields) - return model_response + return _choice_fields_chunk( + model_response, _upstream_choice_index(response_obj.get("original_chunk")), chunk_choice_fields + ) else: if hasattr(model_response, "usage"): self.chunks.append(model_response) @@ -2010,6 +2022,27 @@ 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 _another_choice_fields_chunk( + self, + model_response: ModelResponseStream, + response_obj: Mapping[str, object], + chunk_choice_fields: Mapping[str, object], + ) -> ModelResponseStream | None: + upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk")) + 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 + finish_fields: Final = ( + self.received_finish_choice_fields.get(upstream_index, _NO_UPSTREAM_CHOICE_FIELDS) + if is_finished + else _NO_UPSTREAM_CHOICE_FIELDS + ) + if not chunk_choice_fields and all(value is None for value in finish_fields.values()): + return None + return _choice_fields_chunk( + model_response, upstream_index, MappingProxyType({**finish_fields, **chunk_choice_fields}) + ) + def finish_reason_handler(self): model_response: Final = self.model_response_creator( hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id} diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b4b1ecc3c16..977c2b8580b 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6341,7 +6341,7 @@ def test_hosted_vllm_stream_with_two_choices_keeps_each_stop_reason_on_its_final content_chunk_stop_reasons: Final = frozenset( (choice["index"], choice["stop_reason"]) for choice in emitted - if choice["finish_reason"] is None and choice.get("stop_reason") is not None + if _DELTA.validate_python(choice["delta"]).get("content") and choice.get("stop_reason") is not None ) finished_with_content: Final = frozenset( (index, _TWO_CHOICE_STOP_REASONS[index]) @@ -6351,6 +6351,27 @@ def test_hosted_vllm_stream_with_two_choices_keeps_each_stop_reason_on_its_final assert content_chunk_stop_reasons == finished_with_content + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", + [("", ""), ("2", ""), ("", "2"), ("2", "2")], + ids=["both-finish-only", "content-then-finish-only", "finish-only-then-content", "both-content-then-finish"], +) +def test_hosted_vllm_stream_with_two_choices_delivers_every_token_id_and_stop_reason_of_each_choice( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + + assert { + index: _joined_token_ids([choice for choice in emitted if choice["index"] == index]) + for index in _TWO_CHOICE_TOKEN_IDS + } == _TWO_CHOICE_TOKEN_IDS + assert frozenset( + (choice["index"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None + ) == frozenset(_TWO_CHOICE_STOP_REASONS.items()) + + _CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19)