From 31ae498b9b599673133d585b26ba3dcc96a2a9be Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 09:26:14 -0700 Subject: [PATCH] fix(hosted_vllm): keep stop_reason on the last content chunk, keep both fields in built stream responses, and key the cache on return_token_ids --- litellm/caching/caching.py | 2 +- .../streaming_chunk_builder_utils.py | 29 ++++++++ .../litellm_core_utils/streaming_handler.py | 14 ++-- tests/unit/caching/test_caching.py | 11 +++ .../test_streaming_chunk_builder_utils.py | 70 +++++++++++++++++++ .../test_streaming_handler.py | 29 ++++++-- 6 files changed, 144 insertions(+), 11 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 5f5b8210096..86a658791b6 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -42,7 +42,7 @@ from .redis_semantic_cache import RedisSemanticCache from .s3_cache import S3Cache RESPONSE_CACHE_TARGET: Final = "llm_response" -_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"}) +_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token", "return_token_ids"}) _DECODED_LITERAL: Final = TypeAdapter(object) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 8d243e572fa..1a077ff2ad9 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -12,6 +12,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping @@ -454,6 +455,7 @@ class ChunkProcessor: "index": 0, "message": {"role": role, "content": ""}, "finish_reason": finish_reason, + **_streamed_vllm_choice_fields(_first_choices_of_index_zero(chunks)), } ], "usage": { @@ -1306,6 +1308,7 @@ def _rebuilt_text_choice(index: int, choices: Sequence[TextChoices]) -> TextChoi index=index, logprobs=_merged_logprobs(choices), finish_reason=next((choice.finish_reason for choice in reversed(choices) if choice.finish_reason), None), + **_streamed_vllm_choice_fields(choices), ) @@ -1439,6 +1442,32 @@ def _field_of(value: object, name: str) -> object: return getattr(value, name, None) +def _first_choices_of_index_zero(chunks: Iterable[object]) -> tuple[object, ...]: + first_choices: Final = (next(iter(object_items(_field_of(chunk, "choices"))), None) for chunk in chunks) + return tuple(choice for choice in first_choices if _field_of(choice, "index") in (0, None)) + + +def _choice_extras(choice: object) -> Mapping[str, object]: + if isinstance(choice, BaseModel): + return choice.model_extra or {} + return object_mapping(choice) or {} + + +def _streamed_vllm_choice_fields(choices: Iterable[object]) -> Mapping[str, object]: + extras: Final = tuple(map(_choice_extras, choices)) + token_id_runs: Final = tuple(object_items(fields.get("token_ids")) for fields in extras) + stop_reasons: Final = tuple(reason for fields in extras if (reason := fields.get("stop_reason")) is not None) + has_token_ids: Final = any(fields.get("token_ids") is not None for fields in extras) + streamed_values: Final[Mapping[str, object]] = MappingProxyType( + { + "token_ids": list(chain.from_iterable(token_id_runs)) if has_token_ids else None, + "stop_reason": stop_reasons[-1] if stop_reasons else None, + } + ) + present_fields: Final = frozenset(chain.from_iterable(fields.keys() for fields in extras)) + return MappingProxyType({field: streamed_values[field] for field in VLLM_CHOICE_FIELDS if field in present_fields}) + + def _delta_reasoning_items(delta: object) -> Sequence[object]: items: Final = _field_of(delta, "reasoning_items") return cast(Sequence[object], items) if isinstance(items, list) else () # cast-ok: items forwarded verbatim diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 403fa947c8f..b41666428a5 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1225,14 +1225,14 @@ class CustomStreamWrapper: _prepend_upstream_choice_extras( model_response.choices[0], upstream_choice, - upstream_choice_extra_keys - {"message", *STREAM_FINISH_CHOICE_FIELDS}, + upstream_choice_extra_keys - {"message"}, ) else: choices: Final = [] for choice in upstream_chunk.choices: try: if isinstance(choice, BaseModel): - choice_json = choice.model_dump(exclude=set(STREAM_FINISH_CHOICE_FIELDS)) + choice_json = choice.model_dump() choice_json.pop( "finish_reason", None ) # for mistral etc. which return a value in their last chunk (not-openai compatible). @@ -1316,10 +1316,12 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = map_finish_reason( finish_reason=self.received_finish_reason ) # ensure consistent output to openai - if _upstream_choice_index(_original_chunk) == model_response.choices[0].index: - self._apply_finish_choice_fields(model_response.choices[0]) - if response_obj.get("is_finished"): - _apply_choice_fields(model_response.choices[0], chunk_choice_fields) + 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 + ): + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) self.sent_last_chunk = True diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 91083921130..f57ce86507a 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -463,6 +463,17 @@ def test_cache_key_tells_apart_requests_that_differ_only_in_continuation_token() assert cache.get_cache_key(**request, continuation_token="") == first_segment_key +def test_cache_key_tells_apart_requests_that_ask_for_token_ids() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]} + + plain_key: Final = cache.get_cache_key(**request) + + assert cache.get_cache_key(**request, return_token_ids=True) != plain_key + assert cache.get_cache_key(**request, return_token_ids=False) == plain_key + assert cache.get_cache_key(**request, return_token_ids=None) == plain_key + + class PhaseRecordingCache(InMemoryCache): """Records the target and the active span each read / write ran under, as a Redis span would.""" diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 59946d54dc7..41c1f0d8126 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -2519,3 +2519,73 @@ def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] assert entry["guardrail_usage"] == {"topicPolicyUnits": 2} assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015) + + +def _vllm_chat_chunk(index: int, content: str, finish_reason: str | None = None, **fields: object) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + created=1, + model="qwen3", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason, **fields)], + ) + + +@pytest.mark.parametrize("finish_content", ["", "4"], ids=["finish-only", "content-then-finish"]) +def test_stream_chunk_builder_keeps_the_vllm_stop_reason_and_every_token_id_of_the_first_choice( + finish_content: str, +) -> None: + built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(0, "1,", stop_reason=None, token_ids=[16, 11]), + _vllm_chat_chunk(1, "7,", stop_reason=None, token_ids=[22, 11]), + _vllm_chat_chunk(0, "2,3", stop_reason=None, token_ids=[17, 11, 18]), + _vllm_chat_chunk(0, finish_content, "stop", stop_reason=",5", token_ids=[19]), + ] + ) + + assert built is not None + assert built.model_dump()["choices"][0]["stop_reason"] == ",5" + assert built.model_dump()["choices"][0]["token_ids"] == [16, 11, 17, 11, 18, 19] + + +def test_stream_chunk_builder_keeps_null_vllm_fields_and_adds_none_for_other_providers() -> None: + vllm_built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(0, "hi", stop_reason=None, token_ids=None), + _vllm_chat_chunk(0, "", "stop", stop_reason=None, token_ids=None), + ] + ) + other_built: Final = stream_chunk_builder( + chunks=[_vllm_chat_chunk(0, "hi"), _vllm_chat_chunk(0, "", "stop")] + ) + + assert vllm_built is not None and other_built is not None + assert { + field: vllm_built.model_dump()["choices"][0][field] for field in ("stop_reason", "token_ids") + } == {"stop_reason": None, "token_ids": None} + assert {"stop_reason", "token_ids"}.isdisjoint(other_built.model_dump()["choices"][0]) + + +def test_stream_chunk_builder_text_completion_keeps_each_choices_vllm_stop_reason_and_token_ids() -> None: + def text_chunk(index: int, text: str, **fields: object) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", + object="text_completion", + created=1, + model="qwen3", + choices=[TextChoices(index=index, text=text, **fields)], + ) + + built: Final = stream_chunk_builder_text_completion( + chunks=[ + text_chunk(0, "1,", token_ids=[16, 11], stop_reason=None), + text_chunk(1, "7,", token_ids=[22, 11], stop_reason=None), + text_chunk(0, "2", token_ids=[17], stop_reason=", 3"), + text_chunk(1, "8", token_ids=[23], stop_reason=151645), + ] + ) + + assert [ + (choice["index"], choice["stop_reason"], choice["token_ids"]) for choice in built.model_dump()["choices"] + ] == [(0, ", 3", [16, 11, 17]), (1, 151645, [22, 11, 23])] + diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 446082d57af..d2ae6087ef6 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6103,11 +6103,12 @@ def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id( ), ) + assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [stop_reason] assert [ - (choice["finish_reason"], choice.get("stop_reason")) + (_DELTA.validate_python(choice["delta"]).get("content"), choice["stop_reason"]) for choice in emitted - if choice["finish_reason"] is not None or choice.get("stop_reason") is not None - ] == [("stop", stop_reason)] + if choice["finish_reason"] is None and choice.get("stop_reason") is not None + ] == ([(final_content, stop_reason)] if final_content else []) assert _joined_token_ids(emitted) == (16, 11, 17, 20) @@ -6217,6 +6218,27 @@ def test_hosted_vllm_stream_with_two_choices_keeps_ids_and_stop_reason_on_their_ ] == [] +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", [("2", ""), ("2", "2")], ids=["content-then-finish-only", "both-content-then-finish"] +) +def test_hosted_vllm_stream_with_two_choices_keeps_each_stop_reason_on_its_final_content_chunk( + 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)) + 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 + ) + finished_with_content: Final = frozenset( + (index, _TWO_CHOICE_STOP_REASONS[index]) + for index, content in zip((first_finished, 1 - first_finished), finish_contents) + if content + ) + + assert content_chunk_stop_reasons == finished_with_content + _CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19) @@ -6264,7 +6286,6 @@ async def test_a_streamed_cache_hit_keeps_the_vllm_stop_reason_and_every_token_i emitted: Final = await _replayed_cache_hit_choices(logging_obj, replay) assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [", 5"] - assert {choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is None} == {None} assert _joined_token_ids(emitted) == _CACHED_VLLM_TOKEN_IDS