From 8a1c185da54e7d1661c2b6a4a1934071179f7fcd Mon Sep 17 00:00:00 2001 From: Aaron Deng <84045863+aarongdeng@users.noreply.github.com> Date: Sun, 25 Oct 2026 14:08:31 +0800 Subject: [PATCH] fix(hosted_vllm): preserve vLLM choice fields --- .../llms/hosted_vllm/chat/transformation.py | 60 ++++++++++++++++++- .../test_hosted_vllm_chat_transformation.py | 46 +++++++++++++- 2 files changed, 102 insertions(+), 4 deletions(-) diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 63e6cb2b5da..f578a638d76 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -2,9 +2,15 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions` """ +from __future__ import annotations + import json -from collections.abc import Coroutine -from typing import Final, Literal, cast, overload +from collections.abc import Coroutine, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, cast, overload + +import httpx +from pydantic import TypeAdapter from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_image_mime_type_from_url, @@ -22,12 +28,62 @@ 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[tuple[str, ...]] = ("stop_reason", "token_ids") +_PROVIDER_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + +def _preserve_vllm_choice_fields(choice: Choices) -> Choices: + raw_provider_fields: Final[object] = getattr(choice, "provider_specific_fields", None) + if not isinstance(raw_provider_fields, Mapping): + return choice + provider_fields: Final = _PROVIDER_FIELDS.validate_python(raw_provider_fields) + direct_fields: Final[Mapping[str, object]] = MappingProxyType( + {field: provider_fields[field] for field in _VLLM_CHOICE_FIELDS if field in provider_fields} + ) + return choice.model_copy(update=direct_fields) if direct_fields else choice + class HostedVLLMChatConfig(OpenAIGPTConfig): + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: LiteLLMLoggingObj, + request_data: Mapping[str, object], + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + encoding: Tokenizer | None, + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + response: Final = super().transform_response( # pyright: ignore[reportUnknownMemberType] # base contract uses unparameterized dicts + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=dict(request_data), + messages=list(messages), + optional_params=dict(optional_params), + litellm_params=dict(litellm_params), + encoding=encoding, + api_key=api_key, + json_mode=json_mode, + ) + choices: Final = [_preserve_vllm_choice_fields(choice) for choice in response.choices] + return response.model_copy(update=MappingProxyType({"choices": choices})) + def _convert_custom_tools_to_function_tools(self, tools: list[dict[str, object]]) -> list[dict[str, object]]: """ vLLM chat completions currently accepts only OpenAI function tools. diff --git a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index c792d5dffcd..1a6af3b3915 100644 --- a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -1,14 +1,56 @@ import json -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock +import httpx import pytest - from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, ) from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig +from litellm.types.utils import ModelResponse + + +def test_hosted_vllm_preserves_choice_fields(): + config = HostedVLLMChatConfig() + response = config.transform_response( + model="hosted_vllm/qwen3", + raw_response=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-vllm", + "model": "qwen3", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "done"}, + "finish_reason": "stop", + "stop_reason": "", + "token_ids": [13, 42], + "routed_experts": [1, 3], + } + ], + }, + ), + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "answer"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + choice = response.model_dump()["choices"][0] + assert choice["stop_reason"] == "" + assert choice["token_ids"] == [13, 42] + assert "routed_experts" not in choice + assert choice["provider_specific_fields"] == { + "stop_reason": "", + "token_ids": [13, 42], + "routed_experts": [1, 3], + } def test_hosted_vllm_chat_transformation_file_url():