From efefc437cef3cf3d8e45d5ce1e4b74271acdba16 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 02:12:04 -0700 Subject: [PATCH] fix(hosted_vllm): keep vLLM choice fields for null, cache hits, n>1, and streamed completions Response conversion keeps stop_reason and token_ids on the choice for any OpenAI compatible response, cache replays carry them, finish fields are recorded per choice index, text completion streams keep token_ids, and the proxy passes only these two keys through, nulls included. --- litellm/constants.py | 1 + .../convert_dict_to_response.py | 33 ++++++-- .../litellm_core_utils/streaming_handler.py | 50 ++++++++---- .../llms/hosted_vllm/chat/transformation.py | 53 +------------ litellm/proxy/proxy_server.py | 20 ++--- litellm/proxy/utils.py | 3 +- .../text_completions/streaming_iterator.py | 4 +- .../test_convert_dict_to_response.py | 36 +++++++++ .../test_streaming_handler.py | 78 ++++++++++++++++++- .../proxy_server/test_streaming_helpers.py | 22 +++++- .../test_model_dump_with_preserved_fields.py | 28 +++++++ ...test_proxy_server_endpoints_and_startup.py | 5 ++ .../test_streaming_iterator.py | 33 ++++++-- 13 files changed, 272 insertions(+), 94 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6774348bd65..3722d2a9e6a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -688,6 +688,7 @@ STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8") STREAM_FINISH_CHOICE_FIELDS: Final = ("stop_reason",) STREAM_CHUNK_CHOICE_FIELDS: Final = ("token_ids",) +VLLM_CHOICE_FIELDS: Final = STREAM_FINISH_CHOICE_FIELDS + STREAM_CHUNK_CHOICE_FIELDS ### SPEND TRACKING ### DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = get_env_float( "DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400 diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 8b94c34942c..abb063f2b90 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -3,7 +3,8 @@ import json import re import time import traceback -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence +from types import MappingProxyType from typing import Final, Literal, cast from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config @@ -11,7 +12,12 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger -from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.constants import ( + RESPONSE_FORMAT_TOOL_NAME, + STREAM_CHUNK_CHOICE_FIELDS, + STREAM_FINISH_CHOICE_FIELDS, + VLLM_CHOICE_FIELDS, +) from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_reasoning_content, @@ -154,8 +160,17 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None: # later can't silently re-introduce the duplication. choice.delta = Delta(content=choice.delta.content) choice.logprobs = None - if hasattr(choice, "enhancements"): - del choice.enhancements + _drop_choice_fields(choice, ("enhancements", *STREAM_CHUNK_CHOICE_FIELDS)) + + +def _drop_choice_fields(choice: StreamingChoices, fields: Iterable[str]) -> None: + for field in fields: + if hasattr(choice, field): + delattr(choice, field) + + +def _vllm_choice_fields(choice: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({field: choice[field] for field in VLLM_CHOICE_FIELDS if field in choice}) def _first_replay_delta( @@ -242,7 +257,9 @@ async def convert_to_streaming_response_async( logprobs = choice.get("logprobs", None) - choice = StreamingChoices(finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs) + choice = StreamingChoices( + finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs, **_vllm_choice_fields(choice) + ) choice_list.append(choice) model_response_object.choices = choice_list @@ -292,6 +309,8 @@ async def convert_to_streaming_response_async( if i > 0: _clear_later_replay_slice_metadata(slice_chunk.choices[0]) slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None + if i < last_idx: + _drop_choice_fields(slice_chunk.choices[0], STREAM_FINISH_CHOICE_FIELDS) if i == last_idx and original_usage is not None: setattr(slice_chunk, "usage", original_usage) yield slice_chunk @@ -332,6 +351,7 @@ def convert_to_streaming_response( delta=delta, logprobs=logprobs, enhancements=enhancements, + **_vllm_choice_fields(choice), ) choice_list.append(choice) @@ -377,6 +397,8 @@ def convert_to_streaming_response( if i > 0: _clear_later_replay_slice_metadata(slice_chunk.choices[0]) slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None + if i < last_idx: + _drop_choice_fields(slice_chunk.choices[0], STREAM_FINISH_CHOICE_FIELDS) if i == last_idx and original_usage is not None: setattr(slice_chunk, "usage", original_usage) yield slice_chunk @@ -834,6 +856,7 @@ def convert_to_model_response_object( logprobs=logprobs, enhancements=enhancements, provider_specific_fields=provider_specific_fields, + **_vllm_choice_fields(choice), ) if provider_sent_logprobs and logprobs is None: choice.logprobs = None diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 10947a85489..054825373f9 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -115,6 +115,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({}) def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None: @@ -139,8 +140,12 @@ def _upstream_choice_index(upstream_chunk: object) -> int: 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 + if not upstream_extras: + return _NO_UPSTREAM_CHOICE_FIELDS return MappingProxyType( - {field: value for field, value in _upstream_choice_fields(upstream_chunk, fields).items() if value is not None} + {field: upstream_extras[field] for field in fields if upstream_extras.get(field) is not None} ) @@ -373,7 +378,7 @@ 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[str, object] = _NO_UPSTREAM_CHOICE_FIELDS + self.received_finish_choice_fields: Mapping[int, Mapping[str, object]] = _NO_FINISH_CHOICE_FIELDS self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.special_tokens = [ "<|assistant|>", @@ -818,6 +823,7 @@ class CustomStreamWrapper: "logprobs": logprobs, "original_chunk": str_line, "usage": usage, + "chunk_choice_fields": _non_null_upstream_choice_fields(str_line, STREAM_CHUNK_CHOICE_FIELDS), } except Exception as e: raise e @@ -842,7 +848,7 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_openai_text_completion_chunk(self, chunk): + def handle_openai_text_completion_chunk(self, chunk: object): try: text = "" is_finished = False @@ -860,6 +866,7 @@ class CustomStreamWrapper: "is_finished": is_finished, "finish_reason": finish_reason, "usage": usage, + "chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS), } except Exception as e: @@ -1170,6 +1177,9 @@ 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( + "chunk_choice_fields", _NO_UPSTREAM_CHOICE_FIELDS + ) if is_chunk_non_empty: # cannot set content of an OpenAI Object to be an empty string self.raise_on_model_repetition() @@ -1262,6 +1272,7 @@ class CustomStreamWrapper: if _index is not None: model_response.choices[0].index = _index + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) self._optional_combine_thinking_block_in_choices(model_response=model_response) return model_response @@ -1305,25 +1316,20 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = map_finish_reason( finish_reason=self.received_finish_reason ) # ensure consistent output to openai - _apply_choice_fields(model_response.choices[0], self.received_finish_choice_fields) + self._apply_finish_choice_fields(model_response.choices[0]) if response_obj.get("is_finished"): - _apply_choice_fields( - model_response.choices[0], - _non_null_upstream_choice_fields(_original_chunk, STREAM_CHUNK_CHOICE_FIELDS), - ) + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) self.sent_last_chunk = True return model_response elif self._has_special_delta_content(model_response): return self._handle_special_delta_content(model_response) - elif upstream_chunk_fields := _non_null_upstream_choice_fields( - response_obj.get("original_chunk"), STREAM_CHUNK_CHOICE_FIELDS - ): + elif chunk_choice_fields: upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk")) model_response.choices[0].delta = Delta(content=None) # 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], upstream_chunk_fields) + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) return model_response else: if hasattr(model_response, "usage"): @@ -1543,7 +1549,7 @@ class CustomStreamWrapper: print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] - self.received_finish_choice_fields = _upstream_choice_fields(chunk, STREAM_FINISH_CHOICE_FIELDS) + self._record_finish_choice_fields(chunk) if response_obj["usage"] is not None: _text_completion_usage: Final[Usage] = response_obj["usage"] setattr( @@ -1606,6 +1612,7 @@ class CustomStreamWrapper: self.system_fingerprint = cached_chunk.system_fingerprint if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] + self._record_finish_choice_fields(cached_chunk) else: # openai / azure chat model if isinstance(chunk, ChatCompletionChunk): self._openai_handler_receives_chat_chunks = True @@ -1633,7 +1640,7 @@ class CustomStreamWrapper: f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}" ) self.received_finish_reason = response_obj["finish_reason"] - self.received_finish_choice_fields = _upstream_choice_fields(chunk, STREAM_FINISH_CHOICE_FIELDS) + self._record_finish_choice_fields(chunk) if response_obj.get("original_chunk", None) is not None: if hasattr(response_obj["original_chunk"], "id"): model_response = self.set_model_id(response_obj["original_chunk"].id, model_response) @@ -1970,6 +1977,19 @@ class CustomStreamWrapper: if self.logging_obj.is_sync_litellm_request(litellm_params): self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) + def _record_finish_choice_fields(self, upstream_chunk: object) -> None: + self.received_finish_choice_fields = MappingProxyType( + { + **self.received_finish_choice_fields, + _upstream_choice_index(upstream_chunk): _upstream_choice_fields( + upstream_chunk, STREAM_FINISH_CHOICE_FIELDS + ), + } + ) + + 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 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} @@ -1979,7 +1999,7 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = _finish_reason else: model_response.choices[0].finish_reason = "stop" - _apply_choice_fields(model_response.choices[0], self.received_finish_choice_fields) + self._apply_finish_choice_fields(model_response.choices[0]) ## if tool use if ( diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 4a1635f38d1..63e6cb2b5da 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -3,11 +3,8 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions` """ import json -from collections.abc import Coroutine, Mapping -from typing import TYPE_CHECKING, Final, Literal, cast, overload - -import httpx -from pydantic import JsonValue +from collections.abc import Coroutine +from typing import Final, Literal, cast, overload from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_image_mime_type_from_url, @@ -25,24 +22,10 @@ from litellm.types.llms.openai import ( ChatCompletionVideoObject, ChatCompletionVideoUrlObject, ) -from litellm.types.utils import Choices, ModelResponse from ....utils import remove_additional_properties, remove_strict_from_schema from ...openai.chat.gpt_transformation import OpenAIGPTConfig -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer - -_VLLM_CHOICE_FIELDS: Final = ("stop_reason", "token_ids") - - -def _vllm_choice_fields(choice: Choices) -> Mapping[str, object]: - provider_fields: Final[Mapping[str, object]] = ( - choice.provider_specific_fields or {} if "provider_specific_fields" in choice else {} - ) - return {field: provider_fields[field] for field in _VLLM_CHOICE_FIELDS if field in provider_fields} - class HostedVLLMChatConfig(OpenAIGPTConfig): def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, object]]) -> list[dict[str, object]]: @@ -123,38 +106,6 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): return super().map_openai_params(non_default_params, optional_params, model, drop_params) - def transform_response( - self, - model: str, - raw_response: httpx.Response, - model_response: ModelResponse, - logging_obj: "LiteLLMLoggingObj", - request_data: dict[str, JsonValue], # mutable-ok: BaseConfig signature - messages: list[AllMessageValues], # mutable-ok: BaseConfig signature - optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature - litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature - encoding: "Tokenizer | None", - api_key: str | None = None, - json_mode: bool | None = None, - ) -> ModelResponse: - response: Final = super().transform_response( # pyright: ignore[reportUnknownMemberType] # OpenAIGPTConfig.transform_response takes bare dicts - model=model, - raw_response=raw_response, - model_response=model_response, - logging_obj=logging_obj, - request_data=request_data, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - encoding=encoding, - api_key=api_key, - json_mode=json_mode, - ) - for choice in response.choices: - for field, value in _vllm_choice_fields(choice).items(): - choice[field] = value - return response - def _get_openai_compatible_provider_info( self, api_base: str | None, api_key: str | None ) -> tuple[str | None, str | None]: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1774a3fcd5b..121ba30a92c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -87,6 +87,7 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, + VLLM_CHOICE_FIELDS, ) from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.bug_report import ( @@ -10065,7 +10066,9 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) - return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending -_STANDARD_STREAM_CHOICE_FIELDS: Final = frozenset({"finish_reason", "index", "delta", "logprobs", "enhancements"}) +def _vllm_choice_fields(choice: BaseModel) -> Mapping[str, object]: + choice_extras: Final = choice.model_extra or {} + return MappingProxyType({field: choice_extras[field] for field in VLLM_CHOICE_FIELDS if field in choice_extras}) def _fast_serialize_simple_model_response_stream( @@ -10088,9 +10091,11 @@ def _fast_serialize_simple_model_response_stream( return None choice: Final = choices[0] - if getattr(choice, "logprobs", None) is not None or getattr(choice, "enhancements", None) is not None: - return None - if not isinstance(choice, BaseModel) or (choice.model_extra or {}).keys() - _STANDARD_STREAM_CHOICE_FIELDS: + if ( + not isinstance(choice, BaseModel) + or getattr(choice, "logprobs", None) is not None + or getattr(choice, "enhancements", None) is not None + ): return None delta: Final[object] = getattr(choice, "delta", None) @@ -10119,7 +10124,7 @@ def _fast_serialize_simple_model_response_stream( if content is not None: delta_dict["content"] = content - choice_dict: Final = {"index": getattr(choice, "index", 0), "delta": delta_dict} + choice_dict: Final = {"index": getattr(choice, "index", 0), "delta": delta_dict, **_vllm_choice_fields(choice)} finish_reason: Final = getattr(choice, "finish_reason", None) if finish_reason is not None: choice_dict["finish_reason"] = finish_reason @@ -10170,10 +10175,7 @@ def _is_empty_streaming_choice(choice: StreamingChoices) -> bool: return False if getattr(choice, "logprobs", None) is not None: return False - provider_choice_values: Final = ( - value for field, value in (choice.model_extra or {}).items() if field not in _STANDARD_STREAM_CHOICE_FIELDS - ) - if not _all_values_none(provider_choice_values): + if not _all_values_none(_vllm_choice_fields(choice).values()): return False delta: Final = choice.delta if hasattr(choice, "delta") else None if delta is None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0ddfc5ce62f..4b9a83bf48b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -68,6 +68,7 @@ from litellm.constants import ( SPEND_LOG_QUEUE_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, + VLLM_CHOICE_FIELDS, ) from litellm.cost_calculator import BaseTokenUsageProcessor from litellm.litellm_core_utils.bug_report import ( @@ -9917,7 +9918,7 @@ _PRESERVED_NONE_FIELDS: Final[list[tuple[str, str]]] = [ ("message", "refusal"), ("delta", "content"), # null in streaming chunks ] -_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs",) +_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs", *VLLM_CHOICE_FIELDS) _PRESERVED_NONE_RESPONSE_FIELDS: Final = ("system_fingerprint",) diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index 0e27b90480f..d3739c90bd3 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,6 +1,6 @@ from typing import Final -from litellm.constants import STREAM_FINISH_CHOICE_FIELDS +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.types.utils import Choices, ModelResponse, TextChoices, TextCompletionResponse @@ -41,7 +41,7 @@ class TextCompletionStreamWrapper: text_choices["reasoning_content"] = delta.get("reasoning_content") text_choices["index"] = chat_choice["index"] text_choices["finish_reason"] = chat_choice["finish_reason"] - for field in STREAM_FINISH_CHOICE_FIELDS: + for field in VLLM_CHOICE_FIELDS: if field in chat_choice: text_choices[field] = chat_choice[field] response["choices"] = [text_choices] diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index 0c601a6b2c8..d9ad1d74712 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -2968,3 +2968,39 @@ def test_convert_to_model_response_object_function_output(): assert result.usage.completion_tokens_details == CompletionTokensDetailsWrapper(reasoning_tokens=0) assert result._hidden_params is not None + + +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": 151645}, {"stop_reason": None, "token_ids": None}], + ids=["stop-string-with-token-ids", "stop-token", "natural-stop"], +) +def test_vllm_choice_fields_stay_on_the_choice_through_a_cache_round_trip( + vllm_choice_fields: dict[str, object], +) -> None: + upstream: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + **vllm_choice_fields, + } + ], + } + response: Final = convert_to_model_response_object(response_object=upstream, model_response_object=ModelResponse()) + assert isinstance(response, ModelResponse) + cache_hit: Final = convert_to_model_response_object( + response_object=json.loads(response.model_dump_json()), model_response_object=ModelResponse() + ) + assert isinstance(cache_hit, ModelResponse) + + dumped_choices: Final = [json.loads(rebuilt.model_dump_json())["choices"][0] for rebuilt in (response, cache_hit)] + assert [{field: choice.get(field, "missing") for field in vllm_choice_fields} for choice in dumped_choices] == [ + vllm_choice_fields, + vllm_choice_fields, + ] diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index e901a8d27b7..3226df7d237 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -25,6 +25,10 @@ from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, +) from litellm.litellm_core_utils.streaming_handler import ( CustomStreamWrapper, _ProviderChunkEarlyReturn, @@ -6048,6 +6052,13 @@ def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj: assert result.choices[0].delta.content == "hello" +_TOKEN_IDS: Final = TypeAdapter(list[int] | None) + + +def _joined_token_ids(choices: Sequence[Mapping[str, object]]) -> tuple[int, ...]: + return tuple(chain.from_iterable(_TOKEN_IDS.validate_python(choice.get("token_ids")) or () for choice in choices)) + + def _hosted_vllm_stream_choices( logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]] ) -> tuple[dict[str, object], ...]: @@ -6091,7 +6102,7 @@ def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id( ) assert (emitted[-1]["finish_reason"], emitted[-1].get("stop_reason")) == ("stop", stop_reason) - assert list(chain.from_iterable(choice.get("token_ids") or () for choice in emitted)) == [16, 11, 17, 20] + assert _joined_token_ids(emitted) == (16, 11, 17, 20) def test_hosted_vllm_stream_keeps_held_back_token_ids_on_their_own_choice(logging_obj: Logging) -> None: @@ -6131,6 +6142,71 @@ def test_hosted_vllm_stream_with_null_token_ids_emits_the_same_chunks_as_without ] +def test_hosted_vllm_stream_finish_chunk_keeps_the_first_choice_stop_reason(logging_obj: Logging) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 1, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": "1, 2"}, "finish_reason": "stop", "stop_reason": ", 3"}, + {"index": 1, "delta": {"content": "1, 2, 3"}, "finish_reason": "stop", "stop_reason": 151645}, + ), + ) + + assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [", 3"] + + +_CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19) + + +def _cached_vllm_completion() -> dict[str, object]: + return { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + "stop_reason": ", 5", + "token_ids": list(_CACHED_VLLM_TOKEN_IDS), + } + ], + } + + +async def _replayed_cache_hit_choices(logging_obj: Logging, replay: str) -> tuple[dict[str, object], ...]: + if replay == "async": + async_wrapper: Final = CustomStreamWrapper( + completion_stream=convert_to_streaming_response_async(_cached_vllm_completion()), + model="m", + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + return tuple([chunk.choices[0].model_dump() async for chunk in async_wrapper if chunk.choices]) + sync_wrapper: Final = CustomStreamWrapper( + completion_stream=convert_to_streaming_response(_cached_vllm_completion()), + model="m", + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + return tuple(chunk.choices[0].model_dump() for chunk in sync_wrapper if chunk.choices) + + +@pytest.mark.parametrize("replay", ["sync", "async"]) +@pytest.mark.asyncio +async def test_a_streamed_cache_hit_keeps_the_vllm_stop_reason_and_every_token_id( + logging_obj: Logging, replay: str +) -> None: + 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[:-2]} == {None} + assert _joined_token_ids(emitted) == _CACHED_VLLM_TOKEN_IDS + + def _later_chunk_outputs( custom_llm_provider: str, later_chunks: Sequence[ChatCompletionChunk], logging_obj: Logging ) -> tuple[ModelResponseStream, ...]: diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 3128f678e09..3f291d1abd6 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -576,13 +576,27 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): } -def test_serialize_streaming_chunk_keeps_provider_choice_fields(): +@pytest.mark.parametrize( + ("choice_extras", "serialized_extras"), + [ + ({"stop_reason": ",5", "token_ids": [20]}, {"stop_reason": ",5", "token_ids": [20]}), + ({"text": "1,2,3,4"}, {}), + ( + {"finish_reason": "stop", "stop_reason": None, "text": "1,2,3,4"}, + {"finish_reason": "stop", "stop_reason": None}, + ), + ], + ids=["vllm-fields", "unrelated-extra", "natural-stop"], +) +def test_serialize_streaming_chunk_keeps_only_the_vllm_choice_fields( + choice_extras: dict[str, object], serialized_extras: dict[str, object] +): chunk: Final = _simple_chunk(content="1,2,3,4") - chunk.choices[0]["stop_reason"] = ",5" - chunk.choices[0]["token_ids"] = [20] + for field, value in choice_extras.items(): + chunk.choices[0][field] = value assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [ - {"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, "stop_reason": ",5", "token_ids": [20]} + {"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, **serialized_extras} ] diff --git a/tests/unit/proxy/test_model_dump_with_preserved_fields.py b/tests/unit/proxy/test_model_dump_with_preserved_fields.py index 31531dccd16..92c7e58126b 100644 --- a/tests/unit/proxy/test_model_dump_with_preserved_fields.py +++ b/tests/unit/proxy/test_model_dump_with_preserved_fields.py @@ -545,3 +545,31 @@ def test_provider_refusal_string_is_preserved_on_message() -> None: "object": "chat.completion", } assert model_response.choices[0].message.refusal == "I can't help with that." + + +def test_vllm_null_stop_reason_is_kept_only_on_the_choice_that_carries_it() -> None: + model_response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-vllm-natural-stop", + "created": 4, + "model": "Qwen/Qwen3-0.6B", + "object": "chat.completion", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"role": "assistant", "content": "1, 2"}, + "stop_reason": None, + }, + {"finish_reason": "stop", "index": 1, "message": {"role": "assistant", "content": "1, 2"}}, + ], + }, + model_response_object=ModelResponse(), + ) + + dumped_choices: Final = model_dump_with_preserved_fields(model_response, exclude_unset=True)["choices"] + + assert [(choice["index"], "stop_reason" in choice, choice.get("stop_reason")) for choice in dumped_choices] == [ + (0, True, None), + (1, False, None), + ] diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 7c405b63220..481add758bd 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13615,6 +13615,11 @@ def test_is_injected_stream_usage_artifact(): ) assert _is_injected_stream_usage_artifact(null_token_ids_chunk) is True + unrelated_extra_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), text="")] + ) + assert _is_injected_stream_usage_artifact(unrelated_extra_chunk) is True + async def _collect_async_data_generator_frames(request_data: dict) -> list: from litellm.proxy.proxy_server import async_data_generator diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index 37d079eab26..6118f6becd8 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -1,4 +1,5 @@ -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Sequence +from itertools import chain from typing import Final import pytest @@ -8,7 +9,7 @@ from litellm.exceptions import APIConnectionError from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper -from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, Usage +from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, TextCompletionResponse, Usage def test_reasoning_content_preserved_in_text_completion_wrapper(): @@ -98,8 +99,7 @@ def _completion_chunk(text: str, finish_reason: str | None, **choice_extras: obj ) -@pytest.mark.parametrize("final_text", [", 5, 6", ""], ids=["text-then-finish", "finish-only"]) -def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text: str) -> None: +def _text_completion_stream(upstream: Sequence[Completion]) -> list[TextCompletionResponse]: logging_obj: Final = Logging( model="m", messages=[{"role": "user", "content": "1, 2, 3, 4"}], @@ -109,15 +109,20 @@ def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text litellm_call_id="call-1", function_id="function-1", ) - upstream: Final = (_completion_chunk(", 5", None), _completion_chunk(final_text, "stop", stop_reason=", 7")) chat_stream: Final = CustomStreamWrapper( completion_stream=iter(upstream), model="m", custom_llm_provider="text-completion-openai", logging_obj=logging_obj, ) + return list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m")) - responses: Final = list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m")) + +@pytest.mark.parametrize("final_text", [", 5, 6", ""], ids=["text-then-finish", "finish-only"]) +def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text: str) -> None: + responses: Final = _text_completion_stream( + (_completion_chunk(", 5", None), _completion_chunk(final_text, "stop", stop_reason=", 7")) + ) assert responses[-1].choices[0].model_dump(exclude_none=True) == { "finish_reason": "stop", @@ -126,6 +131,22 @@ def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text } +@pytest.mark.parametrize("final_text", [", 6", ""], ids=["text-then-finish", "finish-only"]) +def test_text_completion_stream_keeps_held_back_token_ids_and_the_stop_reason(final_text: str) -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk("", None, token_ids=[11]), + _completion_chunk("", None, token_ids=[220]), + _completion_chunk(", 5", None, token_ids=[20]), + _completion_chunk(final_text, "stop", stop_reason=", 7", token_ids=[11, 220, 21]), + ) + ) + choices: Final = [response.choices[0] for response in responses] + + assert list(chain.from_iterable(choice.get("token_ids") or () for choice in choices)) == [11, 220, 20, 11, 220, 21] + assert [choice.get("stop_reason") for choice in choices if choice.finish_reason is not None] == [", 7"] + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model")