From 4e6d1840feb885634aa9fe325e3f688802cefcb8 Mon Sep 17 00:00:00 2001 From: Prabal Pratap Singh Date: Sun, 25 Oct 2026 11:39:01 +0530 Subject: [PATCH 01/12] feat(hosted_vllm): keep vLLM's stop_reason and token_ids on chat choices Non-streamed hosted_vllm chat completions moved vLLM's choice-level stop_reason and token_ids into provider_specific_fields, so clients written against vLLM could not find them where vLLM puts them. They are now also set on the choice itself, and the copy in provider_specific_fields stays for existing readers. Fixes #47945 --- .../llms/hosted_vllm/chat/transformation.py | 53 +++++++++++++++- .../test_hosted_vllm_chat_transformation.py | 63 ++++++++++++++++++- 2 files changed, 112 insertions(+), 4 deletions(-) diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 63e6cb2b5da..4a1635f38d1 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -3,8 +3,11 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions` """ import json -from collections.abc import Coroutine -from typing import Final, Literal, cast, overload +from collections.abc import Coroutine, Mapping +from typing import TYPE_CHECKING, Final, Literal, cast, overload + +import httpx +from pydantic import JsonValue from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_image_mime_type_from_url, @@ -22,10 +25,24 @@ 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]]: @@ -106,6 +123,38 @@ 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/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..d7baf44bec7 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,15 @@ 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_chat_transformation_file_url(): @@ -397,3 +398,61 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema(): assert tools[0]["function"]["name"] == "search" assert tools[0]["function"]["description"] == "Search docs" assert tools[0]["function"]["parameters"] == input_schema + + +def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResponse: + payload = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "glm", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "logprobs": None, + "finish_reason": "stop", + **choice_extras, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + } + raw_response = httpx.Response( + 200, json=payload, request=httpx.Request("POST", "http://vllm.local/v1/chat/completions") + ) + return HostedVLLMChatConfig().transform_response( + model="hosted_vllm/glm", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + "choice_extras", + [ + {"stop_reason": ", 5", "token_ids": [16, 11, 220]}, + {"stop_reason": 128009}, + {"stop_reason": None}, + ], +) +def test_hosted_vllm_keeps_vllm_choice_fields_on_the_choice(choice_extras: dict[str, object]): + response = _transform_vllm_chat_response(choice_extras) + + dumped_choice: dict[str, dict[str, object]] = json.loads(response.model_dump_json())["choices"][0] + for field, value in choice_extras.items(): + assert dumped_choice[field] == value + assert dumped_choice["provider_specific_fields"][field] == value + + +def test_hosted_vllm_does_not_invent_vllm_choice_fields(): + response = _transform_vllm_chat_response({}) + + dumped_choice: dict[str, object] = json.loads(response.model_dump_json())["choices"][0] + assert "stop_reason" not in dumped_choice + assert "token_ids" not in dumped_choice From 97c6616945ce6b682500dff9ed35ba1cf19a7514 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:41:46 -0700 Subject: [PATCH 02/12] feat(hosted_vllm): keep stop_reason and token_ids on streamed chunks and proxy stream output Streamed chat finish chunks and streamed /v1/completions finish chunks now carry vLLM's stop_reason. Each stream chunk keeps the token_ids vLLM sent with it, including held-back empty chunks, and the proxy no longer strips those chunks as empty usage chunks or drops the fields in its fast serializer. --- litellm/constants.py | 2 + .../litellm_core_utils/streaming_handler.py | 40 +++++++++++ litellm/proxy/proxy_server.py | 10 +++ .../text_completions/streaming_iterator.py | 11 +++- .../test_streaming_handler.py | 66 +++++++++++++++++++ .../proxy_server/test_streaming_helpers.py | 10 +++ ...test_proxy_server_endpoints_and_startup.py | 12 +++- .../test_streaming_iterator.py | 41 ++++++++++++ 8 files changed, 188 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 2b083b03766..1cc6161b2ed 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -685,6 +685,8 @@ STREAM_SSE_DONE_STRING: Final[str] = "[DONE]" STREAM_SSE_DATA_PREFIX: Final[str] = "data: " STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n' 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",) ### 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/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b64e146ca08..e13446a6e66 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_FIELDS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import ( @@ -113,6 +114,29 @@ def _prepend_upstream_choice_extras( choice.__pydantic_fields_set__.update(upstream_extras) +_NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: + match getattr(upstream_chunk, "choices", None): + case [BaseModel() as first_choice, *_]: + upstream_extras: Final[Mapping[str, object]] = first_choice.model_extra or {} + return MappingProxyType({field: upstream_extras[field] for field in fields if field in upstream_extras}) + case _: + return _NO_UPSTREAM_CHOICE_FIELDS + + +def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: + return MappingProxyType( + {field: value for field, value in _upstream_choice_fields(upstream_chunk, fields).items() if value is not None} + ) + + +def _apply_choice_fields(choice: StreamingChoices, fields: Mapping[str, object]) -> None: + for field, value in fields.items(): + setattr(choice, field, value) + + 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) @@ -337,6 +361,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.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.special_tokens = [ "<|assistant|>", @@ -1268,12 +1293,24 @@ 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) + if response_obj.get("is_finished"): + _apply_choice_fields( + model_response.choices[0], + _non_null_upstream_choice_fields(_original_chunk, STREAM_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 + ): + model_response.choices[0].delta = Delta(content=None) # rebind-ok: chunks are built in place here + _apply_choice_fields(model_response.choices[0], upstream_chunk_fields) + return model_response else: if hasattr(model_response, "usage"): self.chunks.append(model_response) @@ -1492,6 +1529,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) if response_obj["usage"] is not None: _text_completion_usage: Final[Usage] = response_obj["usage"] setattr( @@ -1575,6 +1613,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) 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) @@ -1920,6 +1959,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) ## if tool use if ( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 64fb2791913..d27b43f7c0c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10061,6 +10061,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 _fast_serialize_simple_model_response_stream( chunk: ModelResponseStream, ) -> bytes | None: @@ -10083,6 +10086,8 @@ def _fast_serialize_simple_model_response_stream( 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: + return None delta: Final[object] = getattr(choice, "delta", None) if delta is None: @@ -10161,6 +10166,11 @@ 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): + return False delta: Final = choice.delta if hasattr(choice, "delta") else None if delta is None: return True diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index ec1b96fd432..0e27b90480f 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,5 +1,6 @@ from typing import Final +from litellm.constants import STREAM_FINISH_CHOICE_FIELDS from litellm.types.utils import Choices, ModelResponse, TextChoices, TextCompletionResponse @@ -34,11 +35,15 @@ class TextCompletionStreamWrapper: raise Exception if chunk["choices"]: - delta: Final = chunk["choices"][0]["delta"] + chat_choice: Final = chunk["choices"][0] + delta: Final = chat_choice["delta"] text_choices["text"] = delta["content"] text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chunk["choices"][0]["index"] - text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] + text_choices["index"] = chat_choice["index"] + text_choices["finish_reason"] = chat_choice["finish_reason"] + for field in STREAM_FINISH_CHOICE_FIELDS: + if field in chat_choice: + text_choices[field] = chat_choice[field] response["choices"] = [text_choices] else: response["choices"] = [] diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6d700bae514..0b0b72a4545 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -7,6 +7,7 @@ import threading import time from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import chain from types import CodeType, FrameType, MappingProxyType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -28,6 +29,7 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkParsed, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -5901,6 +5903,70 @@ def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj: assert result.choices[0].delta.content == "hello" +def _hosted_vllm_stream_choices( + logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]] +) -> tuple[dict[str, object], ...]: + sse_lines: Final = ( + *( + "data: " + json.dumps({"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [c]}) + for c in upstream_choices + ), + "data: [DONE]", + ) + completion_stream: Final = HostedVLLMChatConfig().get_model_response_iterator( + streaming_response=iter(sse_lines), sync_stream=True, json_mode=False + ) + return tuple( + chunk.choices[0].model_dump() + for chunk in CustomStreamWrapper( + completion_stream=completion_stream, model="m", custom_llm_provider="hosted_vllm", logging_obj=logging_obj + ) + ) + + +@pytest.mark.parametrize("stop_reason", [",5", 151645]) +@pytest.mark.parametrize("final_content", [",3", ""], ids=["content-then-finish", "finish-only"]) +def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id( + logging_obj: Logging, final_content: str, stop_reason: str | int +) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, "token_ids": [11, 17]}, + { + "index": 0, + "delta": {"content": final_content}, + "finish_reason": "stop", + "stop_reason": stop_reason, + "token_ids": [20], + }, + ), + ) + + 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] + + +def _held_back_stop_string_stream(token_ids: Mapping[str, None]) -> tuple[Mapping[str, object], ...]: + return ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": ""}, "finish_reason": "stop", "stop_reason": ",3", **token_ids}, + ) + + +def test_hosted_vllm_stream_with_null_token_ids_emits_the_same_chunks_as_without_them(logging_obj: Logging) -> None: + with_null_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({"token_ids": None})) + without_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({})) + + assert [(c["delta"], c["finish_reason"]) for c in with_null_ids] == [ + (c["delta"], c["finish_reason"]) for c in without_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 9a3516d20d6..458a48fff76 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -576,6 +576,16 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): } +def test_serialize_streaming_chunk_keeps_provider_choice_fields(): + chunk = _simple_chunk(content="1,2,3,4") + chunk.choices[0]["stop_reason"] = ",5" + chunk.choices[0]["token_ids"] = [20] + + assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [ + {"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, "stop_reason": ",5", "token_ids": [20]} + ] + + def test_serialize_streaming_chunk_invalid_input_raises_attribute_error(): """The helper is typed as ``BaseModel`` — handing it a plain dict trips the attribute-access path (no ``model_dump_json``).""" 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 55d1587669a..c157825a039 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13582,7 +13582,7 @@ def _stream_usage_generator_chunks(): def test_is_injected_stream_usage_artifact(): from litellm.proxy.proxy_server import _is_injected_stream_usage_artifact - from litellm.types.utils import ModelResponseStream, Usage + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage content_chunk, finish_chunk, empty_choices_usage_chunk = _stream_usage_test_chunks() assert _is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True @@ -13605,6 +13605,16 @@ def test_is_injected_stream_usage_artifact(): assert _is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False + held_back_token_ids_chunk = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=[16])] + ) + assert _is_injected_stream_usage_artifact(held_back_token_ids_chunk) is False + + null_token_ids_chunk = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=None)] + ) + assert _is_injected_stream_usage_artifact(null_token_ids_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 b7edffc0753..37d079eab26 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -2,8 +2,11 @@ from collections.abc import AsyncIterator, Iterator from typing import Final import pytest +from openai.types.completion import Completion 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 @@ -85,6 +88,44 @@ async def test_text_completion_stream_usage_is_opt_in( ] +def _completion_chunk(text: str, finish_reason: str | None, **choice_extras: object) -> Completion: + return Completion.construct( + id="cmpl-1", + created=1, + model="m", + object="text_completion", + choices=[{"index": 0, "text": text, "finish_reason": finish_reason, "logprobs": None, **choice_extras}], + ) + + +@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: + logging_obj: Final = Logging( + model="m", + messages=[{"role": "user", "content": "1, 2, 3, 4"}], + stream=True, + call_type="text_completion", + start_time=0.0, + 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, + ) + + responses: Final = list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m")) + + assert responses[-1].choices[0].model_dump(exclude_none=True) == { + "finish_reason": "stop", + "index": 0, + "stop_reason": ", 7", + } + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model") From a7ac7ad54ed66e00510381f6127ea905d681ef8b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:28:32 -0700 Subject: [PATCH 03/12] fix(hosted_vllm): keep the upstream choice index on token-id-only stream chunks --- .../litellm_core_utils/streaming_handler.py | 22 +++++++++++++++---- .../test_streaming_handler.py | 19 ++++++++++++++++ .../test_hosted_vllm_chat_transformation.py | 13 ++++++----- .../proxy_server/test_streaming_helpers.py | 2 +- ...test_proxy_server_endpoints_and_startup.py | 4 ++-- 5 files changed, 47 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 433a9ca460f..10947a85489 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -117,13 +117,25 @@ def _prepend_upstream_choice_extras( _NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) -def _upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: +def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None: match getattr(upstream_chunk, "choices", None): case [BaseModel() as first_choice, *_]: - upstream_extras: Final[Mapping[str, object]] = first_choice.model_extra or {} - return MappingProxyType({field: upstream_extras[field] for field in fields if field in upstream_extras}) + return first_choice case _: - return _NO_UPSTREAM_CHOICE_FIELDS + return None + + +def _upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: + first_choice: Final = _upstream_first_choice(upstream_chunk) + if first_choice is None: + return _NO_UPSTREAM_CHOICE_FIELDS + upstream_extras: Final[Mapping[str, object]] = first_choice.model_extra or {} + return MappingProxyType({field: upstream_extras[field] for field in fields if field in upstream_extras}) + + +def _upstream_choice_index(upstream_chunk: object) -> int: + upstream_index: Final[object] = getattr(_upstream_first_choice(upstream_chunk), "index", 0) + return upstream_index if isinstance(upstream_index, int) else 0 def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: @@ -1308,7 +1320,9 @@ class CustomStreamWrapper: elif upstream_chunk_fields := _non_null_upstream_choice_fields( response_obj.get("original_chunk"), STREAM_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) return model_response else: diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b25757af3dc..e901a8d27b7 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6094,6 +6094,25 @@ def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id( assert list(chain.from_iterable(choice.get("token_ids") or () for choice in emitted)) == [16, 11, 17, 20] +def test_hosted_vllm_stream_keeps_held_back_token_ids_on_their_own_choice(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": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [17]}, + {"index": 1, "delta": {"content": "2"}, "finish_reason": None, "token_ids": [18]}, + ), + ) + + assert [(choice["index"], choice.get("token_ids")) for choice in emitted if choice.get("token_ids")] == [ + (1, [16]), + (0, [17]), + (1, [18]), + ] + + def _held_back_stop_string_stream(token_ids: Mapping[str, None]) -> tuple[Mapping[str, object], ...]: return ( {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, 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 d7baf44bec7..21c9179773c 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,4 +1,5 @@ import json +from typing import Final from unittest.mock import MagicMock import httpx @@ -401,7 +402,7 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema(): def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResponse: - payload = { + payload: Final = { "id": "chatcmpl-1", "object": "chat.completion", "created": 1, @@ -417,7 +418,7 @@ def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResp ], "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, } - raw_response = httpx.Response( + raw_response: Final = httpx.Response( 200, json=payload, request=httpx.Request("POST", "http://vllm.local/v1/chat/completions") ) return HostedVLLMChatConfig().transform_response( @@ -442,17 +443,17 @@ def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResp ], ) def test_hosted_vllm_keeps_vllm_choice_fields_on_the_choice(choice_extras: dict[str, object]): - response = _transform_vllm_chat_response(choice_extras) + response: Final = _transform_vllm_chat_response(choice_extras) - dumped_choice: dict[str, dict[str, object]] = json.loads(response.model_dump_json())["choices"][0] + dumped_choice: Final[dict[str, dict[str, object]]] = json.loads(response.model_dump_json())["choices"][0] for field, value in choice_extras.items(): assert dumped_choice[field] == value assert dumped_choice["provider_specific_fields"][field] == value def test_hosted_vllm_does_not_invent_vllm_choice_fields(): - response = _transform_vllm_chat_response({}) + response: Final = _transform_vllm_chat_response({}) - dumped_choice: dict[str, object] = json.loads(response.model_dump_json())["choices"][0] + dumped_choice: Final[dict[str, object]] = json.loads(response.model_dump_json())["choices"][0] assert "stop_reason" not in dumped_choice assert "token_ids" not in dumped_choice diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 458a48fff76..3128f678e09 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -577,7 +577,7 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): def test_serialize_streaming_chunk_keeps_provider_choice_fields(): - chunk = _simple_chunk(content="1,2,3,4") + chunk: Final = _simple_chunk(content="1,2,3,4") chunk.choices[0]["stop_reason"] = ",5" chunk.choices[0]["token_ids"] = [20] 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 c157825a039..7c405b63220 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13605,12 +13605,12 @@ def test_is_injected_stream_usage_artifact(): assert _is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False - held_back_token_ids_chunk = ModelResponseStream( + held_back_token_ids_chunk: Final = ModelResponseStream( model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=[16])] ) assert _is_injected_stream_usage_artifact(held_back_token_ids_chunk) is False - null_token_ids_chunk = ModelResponseStream( + null_token_ids_chunk: Final = ModelResponseStream( model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=None)] ) assert _is_injected_stream_usage_artifact(null_token_ids_chunk) is True 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 04/12] 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") From 4b1069df6168f525a8354279f74527cc72c60369 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:48:40 -0700 Subject: [PATCH 05/12] fix(hosted_vllm): send empty text on token-id chunks, keep fields on their own choice, and keep null stop_reason on slow-path streams Token-id-only chunks now carry an empty content or text string, a finish chunk only takes stop_reason and token_ids from its own upstream choice, content chunks no longer repeat stop_reason, the slow stream serializer keeps a null stop_reason or token_ids, and fields cached under provider_specific_fields come back on the choice --- .../convert_dict_to_response.py | 14 +++- .../litellm_core_utils/streaming_handler.py | 13 +-- litellm/proxy/proxy_server.py | 28 ++++++- .../text_completions/streaming_iterator.py | 7 +- .../test_convert_dict_to_response.py | 31 +++++++ .../test_streaming_handler.py | 83 ++++++++++++++++--- .../proxy_server/test_streaming_helpers.py | 76 +++++++++++++++-- .../test_streaming_iterator.py | 18 ++++ 8 files changed, 242 insertions(+), 28 deletions(-) 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 abb063f2b90..15eaca50e58 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 @@ -170,7 +170,19 @@ def _drop_choice_fields(choice: StreamingChoices, fields: Iterable[str]) -> None 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}) + cached_fields: Final = _choice_provider_specific_fields(choice) + return MappingProxyType( + { + field: choice[field] if field in choice else cached_fields[field] + for field in VLLM_CHOICE_FIELDS + if field in choice or field in cached_fields + } + ) + + +def _choice_provider_specific_fields(choice: Mapping[str, object]) -> Mapping[str, object]: + fields: Final = choice.get("provider_specific_fields") + return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else MappingProxyType({}) def _first_replay_delta( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 054825373f9..403fa947c8f 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"}, + upstream_choice_extra_keys - {"message", *STREAM_FINISH_CHOICE_FIELDS}, ) else: choices: Final = [] for choice in upstream_chunk.choices: try: if isinstance(choice, BaseModel): - choice_json = choice.model_dump() + choice_json = choice.model_dump(exclude=set(STREAM_FINISH_CHOICE_FIELDS)) choice_json.pop( "finish_reason", None ) # for mistral etc. which return a value in their last chunk (not-openai compatible). @@ -1316,9 +1316,10 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = map_finish_reason( finish_reason=self.received_finish_reason ) # ensure consistent output to openai - 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) + 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.sent_last_chunk = True @@ -1327,7 +1328,7 @@ class CustomStreamWrapper: 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=None) # rebind-ok: chunks are built in place here + 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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 121ba30a92c..bf24e877cb1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10159,7 +10159,33 @@ def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes: if serialized_chunk is not None: return serialized_chunk - return chunk.model_dump_json(exclude_none=True, exclude_unset=True) + choices: Final = _streaming_chunk_choices(chunk) + if not any(_null_vllm_choice_fields(choice) for choice in choices): + return chunk.model_dump_json(exclude_none=True, exclude_unset=True) + return json.dumps( + { + **_dump_set_fields(chunk, exclude=frozenset({"choices"})), + "choices": [{**_dump_set_fields(choice), **_null_vllm_choice_fields(choice)} for choice in choices], + }, + separators=(",", ":"), + ensure_ascii=False, + ) + + +def _streaming_chunk_choices(chunk: BaseModel) -> Sequence[BaseModel]: + match chunk: + case ModelResponseStream() | TextCompletionResponse(): + return chunk.choices + case _: + return () + + +def _null_vllm_choice_fields(choice: BaseModel) -> Mapping[str, None]: + return MappingProxyType({field: None for field, value in _vllm_choice_fields(choice).items() if value is None}) + + +def _dump_set_fields(model: BaseModel, exclude: AbstractSet[str] = frozenset()) -> Mapping[str, object]: + return model.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=set(exclude)) def _is_injected_stream_usage_artifact(chunk: object) -> bool: diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index d3739c90bd3..7b65bfbd249 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -30,20 +30,19 @@ class TextCompletionStreamWrapper: response["object"] = "text_completion" response["created"] = chunk.get("created", None) response["model"] = chunk.get("model", None) - text_choices: Final = TextChoices() if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices raise Exception if chunk["choices"]: chat_choice: Final = chunk["choices"][0] + text_choices: Final = TextChoices( + **{field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice} + ) delta: Final = chat_choice["delta"] text_choices["text"] = delta["content"] 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 VLLM_CHOICE_FIELDS: - if field in chat_choice: - text_choices[field] = chat_choice[field] response["choices"] = [text_choices] else: response["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 d9ad1d74712..eb20f675868 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 @@ -3004,3 +3004,34 @@ def test_vllm_choice_fields_stay_on_the_choice_through_a_cache_round_trip( vllm_choice_fields, vllm_choice_fields, ] + + +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": None, "token_ids": None}], + ids=["stop-string-with-token-ids", "natural-stop"], +) +def test_vllm_choice_fields_cached_only_under_provider_specific_fields_come_back_on_the_choice( + vllm_choice_fields: dict[str, object], +) -> None: + cached_before_upgrade: 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", + "provider_specific_fields": vllm_choice_fields, + } + ], + } + cache_hit: Final = convert_to_model_response_object( + response_object=cached_before_upgrade, model_response_object=ModelResponse() + ) + assert isinstance(cache_hit, ModelResponse) + + dumped_choice: Final = json.loads(cache_hit.model_dump_json())["choices"][0] + assert {field: dumped_choice.get(field, "missing") for field in 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 3226df7d237..446082d57af 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6053,6 +6053,8 @@ def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj: _TOKEN_IDS: Final = TypeAdapter(list[int] | None) +_CHOICE_INDEX: Final = TypeAdapter(int) +_DELTA: Final = TypeAdapter(Mapping[str, object]) def _joined_token_ids(choices: Sequence[Mapping[str, object]]) -> tuple[int, ...]: @@ -6101,10 +6103,32 @@ 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 [ + (choice["finish_reason"], choice.get("stop_reason")) + for choice in emitted + if choice["finish_reason"] is not None or choice.get("stop_reason") is not None + ] == [("stop", stop_reason)] assert _joined_token_ids(emitted) == (16, 11, 17, 20) +def test_hosted_vllm_stream_sends_empty_content_on_token_id_only_chunks(logging_obj: Logging) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [11]}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [220]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [16]}, + ), + ) + + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice.get("token_ids")) + for choice in emitted + if choice.get("token_ids") + ] == [("", [11]), ("", [220]), ("1", [16])] + + def test_hosted_vllm_stream_keeps_held_back_token_ids_on_their_own_choice(logging_obj: Logging) -> None: emitted: Final = _hosted_vllm_stream_choices( logging_obj, @@ -6142,18 +6166,55 @@ 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}, +_TWO_CHOICE_STOP_REASONS: Final[Mapping[int, str | int]] = MappingProxyType({0: ", 3", 1: 151645}) +_TWO_CHOICE_TOKEN_IDS: Final[Mapping[int, tuple[int, int]]] = MappingProxyType({0: (16, 17), 1: (26, 27)}) + + +def _two_choice_stream(first_finished: int, finish_contents: tuple[str, str]) -> tuple[Mapping[str, object], ...]: + return ( + *({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)), + *( + {"index": index, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [ids[0]]} + for index, ids in _TWO_CHOICE_TOKEN_IDS.items() + ), + *( + { + "index": index, + "delta": {"content": content}, + "finish_reason": "stop", + "stop_reason": _TWO_CHOICE_STOP_REASONS[index], + "token_ids": [_TWO_CHOICE_TOKEN_IDS[index][1]], + } + for index, content in zip((first_finished, 1 - first_finished), finish_contents) ), ) - assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [", 3"] + +def _carries_another_choices_fields(choice: Mapping[str, object]) -> bool: + own_index: Final = _CHOICE_INDEX.validate_python(choice["index"]) + token_ids: Final = _TOKEN_IDS.validate_python(choice.get("token_ids")) or () + return not set(token_ids) <= set(_TWO_CHOICE_TOKEN_IDS[own_index]) or choice.get("stop_reason") not in ( + None, + _TWO_CHOICE_STOP_REASONS[own_index], + ) + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", + [("", ""), ("2", ""), ("2", "2")], + ids=["both-finish-only", "content-then-finish-only", "both-content-then-finish"], +) +def test_hosted_vllm_stream_with_two_choices_keeps_ids_and_stop_reason_on_their_own_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 [ + (choice["index"], choice.get("token_ids"), choice.get("stop_reason")) + for choice in emitted + if _carries_another_choices_fields(choice) + ] == [] _CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19) @@ -6203,7 +6264,7 @@ 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[:-2]} == {None} + 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 diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 3f291d1abd6..9bca937033d 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -17,7 +17,7 @@ from __future__ import annotations import asyncio import json -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Callable, Mapping from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -55,7 +55,14 @@ from litellm.types.llms.openai import ( ResponseFailedEvent, ResponsesAPIResponse, ) -from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage +from litellm.types.utils import ( + Delta, + ModelResponseStream, + StreamingChoices, + TextChoices, + TextCompletionResponse, + Usage, +) from tests._support.clock_loops import AutojumpClockLoop from .conftest import normalize @@ -591,15 +598,74 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): 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") - for field, value in choice_extras.items(): - chunk.choices[0][field] = value + chunk: Final = ModelResponseStream( + id="chatcmpl-test", + choices=[StreamingChoices(index=0, delta=Delta(content="1,2,3,4", role="assistant"), **choice_extras)], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [ {"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, **serialized_extras} ] +def _text_completion_chunk(choice_extras: Mapping[str, object]) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", + choices=[TextChoices(finish_reason="length", index=0, text="", **choice_extras)], + created=1, + model="Qwen/Qwen3-0.6B", + ) + + +def _chat_chunk_with_fingerprint(choice_extras: Mapping[str, object]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + choices=[StreamingChoices(finish_reason="length", index=0, delta=Delta(content=None), **choice_extras)], + created=1, + model="Qwen/Qwen3-0.6B", + object="chat.completion.chunk", + system_fingerprint="fp_1", + ) + + +_SLOW_PATH_CHUNKS: Final = pytest.mark.parametrize( + "build_chunk", [_text_completion_chunk, _chat_chunk_with_fingerprint], ids=["text-completion", "chat-fingerprint"] +) + + +@_SLOW_PATH_CHUNKS +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": None}, {"stop_reason": None, "token_ids": [21]}, {"stop_reason": None, "token_ids": None}], + ids=["null-stop-reason", "null-stop-reason-with-token-ids", "null-stop-reason-and-token-ids"], +) +def test_serialize_streaming_chunk_slow_path_keeps_null_vllm_choice_fields( + build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream], + vllm_choice_fields: dict[str, object], +) -> None: + chunk: Final = build_chunk(vllm_choice_fields) + plain_chunk: Final = json.loads(build_chunk({}).model_dump_json(exclude_none=True, exclude_unset=True)) + + assert json.loads(_serialize_streaming_chunk(chunk)) == { + **plain_chunk, + "choices": [{**plain_chunk["choices"][0], **vllm_choice_fields}], + } + + +@_SLOW_PATH_CHUNKS +@pytest.mark.parametrize("choice_extras", [{}, {"stop_reason": ",5", "token_ids": [20]}], ids=["no-vllm", "set-vllm"]) +def test_serialize_streaming_chunk_slow_path_without_null_vllm_fields_matches_model_dump_json( + build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream], + choice_extras: dict[str, object], +) -> None: + chunk: Final = build_chunk(choice_extras) + + assert _serialize_streaming_chunk(chunk) == chunk.model_dump_json(exclude_none=True, exclude_unset=True) + + def test_serialize_streaming_chunk_invalid_input_raises_attribute_error(): """The helper is typed as ``BaseModel`` — handing it a plain dict trips the attribute-access path (no ``model_dump_json``).""" diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index 6118f6becd8..fbcb7f6ae3f 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -147,6 +147,24 @@ def test_text_completion_stream_keeps_held_back_token_ids_and_the_stop_reason(fi assert [choice.get("stop_reason") for choice in choices if choice.finish_reason is not None] == [", 7"] +def test_text_completion_stream_sends_empty_text_on_token_id_only_chunks() -> 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]), + ) + ) + + choices: Final = [response.choices[0] for response in responses] + + assert [(choice.text, choice.get("token_ids")) for choice in choices if choice.get("token_ids")] == [ + ("", [11]), + ("", [220]), + (", 5", [20]), + ] + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model") 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 06/12] 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 From a284371030d06c227d079920ec369c47fb6a5f38 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:43:47 -0700 Subject: [PATCH 07/12] fix(hosted_vllm): keep another choice's final token_ids and stop_reason on their own stream chunk With n above 1, a finish-only vLLM chunk for a choice other than the one the single finish chunk carries was dropped or taken as choice 0's finish, losing that choice's last token_ids. It now goes out as an id-only chunk on its own choice index. --- .../litellm_core_utils/streaming_handler.py | 43 ++++++++++++++++--- .../test_streaming_handler.py | 23 +++++++++- 2 files changed, 60 insertions(+), 6 deletions(-) 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) 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 08/12] 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) From dedd49dbd535b43698dcdb2676360a391c74ac21 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:30:48 -0700 Subject: [PATCH 09/12] fix(streaming): keep each streamed text completion choice on its own index --- .../litellm_core_utils/streaming_handler.py | 22 ++++++++++++------- .../test_streaming_iterator.py | 22 +++++++++++++++++++ 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7875cdb0df1..5c1f354fe4b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -152,6 +152,13 @@ def _upstream_choice_index(upstream_chunk: object) -> int: return upstream_index if isinstance(upstream_index, int) else 0 +def _response_choice_index(response_obj: Mapping[str, object]) -> int: + choice_index: Final = response_obj.get("choice_index") + if isinstance(choice_index, int): + return choice_index + return _upstream_choice_index(response_obj.get("original_chunk")) + + 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 @@ -890,6 +897,7 @@ class CustomStreamWrapper: "finish_reason": finish_reason, "usage": usage, "chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS), + "choice_index": _upstream_choice_index(chunk), } except Exception as e: @@ -1292,7 +1300,7 @@ class CustomStreamWrapper: if response_obj.get("provider_specific_fields") is not None: completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] model_response.choices[0].delta = Delta(**completion_obj) - _index: Final[int | None] = completion_obj.get("index") + _index: Final[int | None] = completion_obj.get("index", response_obj.get("choice_index")) if _index is not None: model_response.choices[0].index = _index @@ -1303,7 +1311,7 @@ class CustomStreamWrapper: else: return elif self.received_finish_reason is not None: - chunk_upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk")) + chunk_upstream_index: Final = _response_choice_index(response_obj) if chunk_choice_fields and self._streams_past_the_pending_finish( model_response, response_obj, chunk_upstream_index ): @@ -1356,7 +1364,7 @@ class CustomStreamWrapper: 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 + and _response_choice_index(response_obj) == model_response.choices[0].index ): _apply_choice_fields(model_response.choices[0], chunk_choice_fields) @@ -1366,9 +1374,7 @@ class CustomStreamWrapper: elif self._has_special_delta_content(model_response): return self._handle_special_delta_content(model_response) elif chunk_choice_fields: - return _choice_fields_chunk( - model_response, _upstream_choice_index(response_obj.get("original_chunk")), chunk_choice_fields - ) + return _choice_fields_chunk(model_response, _response_choice_index(response_obj), chunk_choice_fields) else: if hasattr(model_response, "usage"): self.chunks.append(model_response) @@ -2049,7 +2055,7 @@ class CustomStreamWrapper: ) -> 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")) + upstream_index: Final = _response_choice_index(response_obj) carried: Final = self.carried_chunk_choice_fields.get(upstream_index) if carried is None: return chunk_choice_fields @@ -2070,7 +2076,7 @@ class CustomStreamWrapper: response_obj: Mapping[str, object], chunk_choice_fields: Mapping[str, object], ) -> ModelResponseStream | None: - upstream_index: Final = _upstream_choice_index(response_obj.get("original_chunk")) + upstream_index: Final = _response_choice_index(response_obj) 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 diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index 38d7f99e449..402e3302eeb 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -165,6 +165,28 @@ def test_text_completion_stream_sends_empty_text_on_token_id_only_chunks() -> No ] +def test_text_completion_stream_keeps_each_choices_text_and_token_ids_on_its_own_index() -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk(" 5", None, index=0, token_ids=[20]), + _completion_chunk(" 7", None, index=1, token_ids=[22]), + _completion_chunk(", 6", None, index=0, token_ids=[11, 21]), + _completion_chunk("", None, index=1, token_ids=[11]), + _completion_chunk("", "stop", index=0, stop_reason=", 9", token_ids=[11, 220]), + _completion_chunk(", 8", "length", index=1, token_ids=[220, 23]), + ) + ) + choices: Final = tuple(chain.from_iterable(response.choices for response in responses)) + + assert { + index: ( + "".join(choice.text or "" for choice in choices if choice.index == index), + list(chain.from_iterable(choice.get("token_ids") or () for choice in choices if choice.index == index)), + ) + for index in (0, 1) + } == {0: (" 5, 6", [20, 11, 21, 11, 220]), 1: (" 7, 8", [22, 11, 220, 23])} + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model") From ca2249429347361e8af9d7cfb34c0b17886c071b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 02:51:06 -0700 Subject: [PATCH 10/12] fix(streaming): finish each streamed text completion choice on its own index and null token_ids in redacted logs --- litellm/litellm_core_utils/redact_messages.py | 8 +++++ .../litellm_core_utils/streaming_handler.py | 20 +++++++++++-- .../test_redact_messages.py | 28 +++++++++++++++++ .../test_streaming_iterator.py | 30 +++++++++++++++++++ 4 files changed, 84 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 548d088d833..91f4d144b47 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -196,6 +196,11 @@ def _redact_function_call(function_call) -> None: def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" + if ( + isinstance(choice, (litellm.Choices, litellm.utils.StreamingChoices)) + and getattr(choice, "token_ids", None) is not None + ): + setattr(choice, "token_ids", None) # noqa: B010 # pydantic extra, not a declared field if isinstance(choice, litellm.Choices): if choice.message.content is not None: choice.message.content = REDACTED_BY_LITELLM @@ -250,6 +255,7 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ "citations", "compaction_blocks", "compaction_delta", + "token_ids", ) if key in fields }, @@ -421,6 +427,8 @@ def _redact_model_response_dict_choices(choices, redacted_str: str): for choice in choices: if isinstance(choice, dict): _redact_provider_specific_fields_dict(choice) + if choice.get("token_ids") is not None: + choice["token_ids"] = None if "message" in choice and isinstance(choice["message"], dict): if choice["message"].get("content") is not None: choice["message"]["content"] = redacted_str diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 5c1f354fe4b..b248aebd121 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -891,13 +891,15 @@ class CustomStreamWrapper: is_finished = True finish_reason = choices[0].finish_reason usage = getattr(chunk, "usage", None) + choice_index: Final = _upstream_choice_index(chunk) return { "text": text, "is_finished": is_finished, "finish_reason": finish_reason, "usage": usage, "chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS), - "choice_index": _upstream_choice_index(chunk), + "choice_index": choice_index, + "sibling_finish_reason": finish_reason if choice_index != 0 else None, } except Exception as e: @@ -1212,6 +1214,7 @@ class CustomStreamWrapper: "chunk_choice_fields", _NO_UPSTREAM_CHOICE_FIELDS ) chunk_choice_fields: Final = self._take_carried_chunk_choice_fields(response_obj, upstream_chunk_choice_fields) + sibling_finish_reason: Final[str | None] = response_obj.get("sibling_finish_reason") if is_chunk_non_empty: # cannot set content of an OpenAI Object to be an empty string self.raise_on_model_repetition() @@ -1305,11 +1308,19 @@ class CustomStreamWrapper: model_response.choices[0].index = _index _apply_choice_fields(model_response.choices[0], chunk_choice_fields) + if sibling_finish_reason is not None: + self._finish_sibling_choice(model_response.choices[0], sibling_finish_reason) self._optional_combine_thinking_block_in_choices(model_response=model_response) return model_response else: return + elif sibling_finish_reason is not None: + sibling_chunk: Final = _choice_fields_chunk( + model_response, _response_choice_index(response_obj), chunk_choice_fields + ) + self._finish_sibling_choice(sibling_chunk.choices[0], sibling_finish_reason) + return sibling_chunk elif self.received_finish_reason is not None: chunk_upstream_index: Final = _response_choice_index(response_obj) if chunk_choice_fields and self._streams_past_the_pending_finish( @@ -1597,8 +1608,9 @@ class CustomStreamWrapper: completion_obj["content"] = response_obj["text"] print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] self._record_finish_choice_fields(chunk) + if response_obj["is_finished"] and response_obj["sibling_finish_reason"] is None: + self.received_finish_reason = response_obj["finish_reason"] if response_obj["usage"] is not None: _text_completion_usage: Final[Usage] = response_obj["usage"] setattr( @@ -2039,6 +2051,10 @@ 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 _finish_sibling_choice(self, choice: StreamingChoices, finish_reason: str) -> None: + choice.finish_reason = map_finish_reason(finish_reason) # rebind-ok: chunks are built in place here + self._apply_finish_choice_fields(choice) + def _streams_past_the_pending_finish( self, model_response: ModelResponseStream, response_obj: Mapping[str, object], upstream_index: int ) -> bool: diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 5b64e42d8a8..387289eeea6 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -560,6 +560,34 @@ class TestPerformRedaction: {"text": "redacted-by-litellm", "finish_reason": "length", "index": 1}, ] + def test_redacts_token_ids_that_decode_to_the_completion(self): + details = { + "stream": True, + "complete_streaming_response": litellm.ModelResponse( + choices=[ + litellm.Choices(message=litellm.Message(content="secret", role="assistant"), token_ids=[20, 11]) + ] + ), + "standard_logging_object": { + "response": { + "choices": [ + { + "token_ids": [20, 11], + "message": {"content": "secret", "provider_specific_fields": {"token_ids": [20, 11]}}, + }, + {"text": " Paris.", "token_ids": [20, 11], "index": 1}, + ] + } + }, + } + + perform_redaction(details, None) + + logged_choices = details["standard_logging_object"]["response"]["choices"] + assert [choice["token_ids"] for choice in logged_choices] == [None, None] + assert logged_choices[0]["message"]["provider_specific_fields"]["token_ids"] is None + assert details["complete_streaming_response"].choices[0].get("token_ids") is None + def test_redacts_object_choices_inside_model_response_dict(self): result = { "choices": [ diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index 402e3302eeb..8f7fd6fa327 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -187,6 +187,36 @@ def test_text_completion_stream_keeps_each_choices_text_and_token_ids_on_its_own } == {0: (" 5, 6", [20, 11, 21, 11, 220]), 1: (" 7, 8", [22, 11, 220, 23])} +@pytest.mark.parametrize( + "upstream", + [ + ( + _completion_chunk(" a", None, index=0), + _completion_chunk(" b", None, index=1), + _completion_chunk("", "length", index=2), + _completion_chunk(" red", "stop", index=1, stop_reason=","), + _completion_chunk(" s", "stop", index=0, stop_reason=","), + ), + ( + _completion_chunk(" a", None, index=0), + _completion_chunk(" b", None, index=1), + _completion_chunk(" s", "stop", index=0, stop_reason=","), + _completion_chunk("", "length", index=2), + _completion_chunk(" red", "stop", index=1, stop_reason=","), + ), + ], + ids=["other-choices-finish-first", "choice-0-finishes-first"], +) +def test_text_completion_stream_finishes_each_choice_on_its_own_index(upstream: Sequence[Completion]) -> None: + finished: Final = [ + (choice.index, choice.finish_reason, choice.get("stop_reason")) + for choice in chain.from_iterable(response.choices for response in _text_completion_stream(upstream)) + if choice.finish_reason is not None + ] + + assert sorted(finished) == [(0, "stop", ","), (1, "stop", ","), (2, "length", None)] + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model") From 864ad9e71dcf977e96b15a3995b24af5fe3c1f94 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 05:14:32 -0700 Subject: [PATCH 11/12] fix(proxy): drop token_ids from guarded streams, redact choice-level provider fields, and keep one message_delta per Anthropic stream --- litellm/litellm_core_utils/redact_messages.py | 29 ++++++---- .../messages/adapters/streaming_iterator.py | 15 +++++- litellm/proxy/utils.py | 31 +++++++---- .../test_redact_messages.py | 30 +++++++++++ .../test_streaming_iterator_combined_chunk.py | 14 +++++ .../test_proxy_logging_hook_detection.py | 53 +++++++++++++++++++ 6 files changed, 153 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 91f4d144b47..8cc8b2a8ef6 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -196,11 +196,9 @@ def _redact_function_call(function_call) -> None: def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" - if ( - isinstance(choice, (litellm.Choices, litellm.utils.StreamingChoices)) - and getattr(choice, "token_ids", None) is not None - ): - setattr(choice, "token_ids", None) # noqa: B010 # pydantic extra, not a declared field + if isinstance(choice, (litellm.Choices, litellm.utils.StreamingChoices)): + drop_token_ids(choice) + _redact_provider_specific_fields(choice) if isinstance(choice, litellm.Choices): if choice.message.content is not None: choice.message.content = REDACTED_BY_LITELLM @@ -232,14 +230,27 @@ def _redact_provider_specific_fields(message: object) -> None: setattr(message, "provider_specific_fields", _redacted_provider_specific_fields(fields)) # noqa: B010 # typed as object +def _drop_provider_specific_token_ids(holder: object) -> None: + fields: Final = _as_mapping(getattr(holder, "provider_specific_fields", None)) + if fields is not None and fields.get("token_ids") is not None: + setattr(holder, "provider_specific_fields", {**fields, "token_ids": None}) # noqa: B010 # typed as object + + +def drop_token_ids(choice: object) -> None: + if getattr(choice, "token_ids", None) is not None: + setattr(choice, "token_ids", None) # noqa: B010 # pydantic extra, not a declared field + for holder in (choice, getattr(choice, "delta", None), getattr(choice, "message", None)): + _drop_provider_specific_token_ids(holder) + + def _redact_provider_specific_fields_dict(choice: Mapping[str, object]) -> None: - for message in (choice.get("message"), choice.get("delta")): + for holder in (choice, choice.get("message"), choice.get("delta")): if ( - (readable := _as_mapping(message)) is not None + (readable := _as_mapping(holder)) is not None and (fields := _as_mapping(readable.get("provider_specific_fields"))) is not None - and isinstance(message, dict) + and isinstance(holder, dict) ): - message["provider_specific_fields"] = _redacted_provider_specific_fields(fields) + holder["provider_specific_fields"] = _redacted_provider_specific_fields(fields) def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[str, object]: diff --git a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py index b19c0d8e2a9..0b778eaa09c 100644 --- a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py @@ -494,10 +494,23 @@ class _CombinedChunkSplitter: for position, block in enumerate(blocks) ) + @staticmethod + def _without_sibling_finish(chunk: "ModelResponseStream") -> "ModelResponseStream": + choices: Final = _optional_attr_sequence(chunk, "choices") + if ( + len(choices) != 1 + or _optional_attr(choices[0], "index") in (None, 0) + or _optional_attr(choices[0], "finish_reason") is None + ): + return chunk + unfinished: Final = copy.deepcopy(chunk) + unfinished.choices[0].finish_reason = None + return unfinished + def _expand(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]": split_chunks: Final = tuple( split_chunk - for combined_chunk in self._split(chunk) + for combined_chunk in self._split(self._without_sibling_finish(chunk)) for split_chunk in self._split_by_payload_kind(combined_chunk) # comprehension-ok: flattens nested splits ) block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks))) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ca8d5bcc741..c77cf0c2aa9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -163,6 +163,7 @@ from litellm.litellm_core_utils.duration_parser import get_budget_window_start from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.litellm_core_utils.redact_messages import drop_token_ids from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.served_output_texts import ( @@ -171,6 +172,7 @@ from litellm.litellm_core_utils.served_output_texts import ( ) from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model from litellm.llms import load_guardrail_translation_mappings +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_items from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, @@ -695,6 +697,11 @@ def _pipeline_step_streams_live(step: "PipelineStep") -> bool: return not unified_guardrail.resolve_streaming_flag(callback, "streaming_buffer_until_moderated", True) +def _drop_stream_chunk_token_ids(chunk: object) -> None: + for choice in stream_item_items(chunk, "choices"): + drop_token_ids(choice) + + def _pipelines_stream_live(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> bool: return all(_pipeline_step_streams_live(step) for step in _pipeline_steps(pipelines)) @@ -4569,16 +4576,20 @@ class ProxyLogging: stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines) + running_overrides: Final = tuple( + (callback, kind) + for callback, kind in caps.iterator_overrides + if not isinstance(callback, CustomGuardrail) + or ( + callback.guardrail_name not in pipeline_gated_names + and callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is True + ) + ) + drops_token_ids: Final = bool(post_call_pipelines) or any( + isinstance(callback, CustomGuardrail) for callback, _ in running_overrides + ) guarded_layers: Final[list[AsyncGenerator[object, None]]] = [] # mutable-ok: closed on disconnect - for resolved_callback, kind in caps.iterator_overrides: - if isinstance(resolved_callback, CustomGuardrail): - if resolved_callback.guardrail_name in pipeline_gated_names: - continue - if ( - resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) - is not True - ): - continue + for resolved_callback, kind in running_overrides: effective_kind = ( "apply_guardrail" if ( @@ -4636,6 +4647,8 @@ class ProxyLogging: try: async for chunk in current_response: served_chunks.append(chunk) + if drops_token_ids: + _drop_stream_chunk_token_ids(served_chunks[-1]) yield chunk except (GeneratorExit, asyncio.CancelledError): await ProxyLogging._close_guarded_layers(guarded_layers) diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 387289eeea6..30a1a21e369 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -588,6 +588,36 @@ class TestPerformRedaction: assert logged_choices[0]["message"]["provider_specific_fields"]["token_ids"] is None assert details["complete_streaming_response"].choices[0].get("token_ids") is None + def test_redacts_token_ids_a_cache_hit_keeps_on_the_choice_provider_specific_fields(self): + result = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret", role="assistant"), + provider_specific_fields={"token_ids": [20, 11], "stop_reason": "END"}, + ) + ] + ) + details = { + "standard_logging_object": { + "response": { + "choices": [ + { + "message": {"content": "secret"}, + "provider_specific_fields": {"token_ids": [20, 11], "stop_reason": "END"}, + } + ] + } + } + } + + redacted = perform_redaction(details, result) + + logged_fields = details["standard_logging_object"]["response"]["choices"][0]["provider_specific_fields"] + assert logged_fields == {"token_ids": None, "stop_reason": "END"} + assert redacted.choices[0].provider_specific_fields["token_ids"] is None + assert redacted.choices[0].provider_specific_fields["stop_reason"] == "END" + assert result.choices[0].provider_specific_fields["token_ids"] == [20, 11] + def test_redacts_object_choices_inside_model_response_dict(self): result = { "choices": [ diff --git a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py index d4911427e73..a57011b6410 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py @@ -204,6 +204,20 @@ def test_splitter_splits_combined_chunk_into_content_then_finish(): assert finish_chunk.choices[0].delta.content is None +def test_a_finish_on_another_choice_index_does_not_end_the_message(): + chunks: Final = [ + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Hello"))]), + ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content="Hi"))]), + ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content=""), finish_reason="length")]), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")]), + ] + + events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="qwen3")) + + stop_reasons: Final = [event["delta"]["stop_reason"] for event in events if event.get("type") == "message_delta"] + assert stop_reasons == ["end_turn"] + + def test_is_combined_false_when_choices_empty(): """A metadata-only chunk with no choices is never treated as combined.""" assert _CombinedChunkSplitter._is_combined(SimpleNamespace(choices=[])) is False diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index 182ef833069..5895ca4882b 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -347,6 +347,59 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions assert delivered_text != "" +async def _served_token_ids(monkeypatch, guardrail_default_on: bool) -> tuple[object, ...]: + from litellm.caching.caching import DualCache + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = ContentFilterGuardrail( + guardrail_name="output-filter", + blocked_words=[BlockedWord(keyword="zebra", action=ContentFilterAction("MASK"))], + event_hook="post_call", + default_on=guardrail_default_on, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + def chunk(content: str, token_ids: list[int], finish_reason: str | None = None) -> ModelResponseStream: + streamed = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)] + ) + setattr(streamed.choices[0], "token_ids", token_ids) + return streamed + + async def fake_stream(): + yield chunk("the zebra runs", [1, 2, 3]) + yield chunk("", [4], finish_reason="stop") + + served = [ + served_chunk + async for served_chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + request_data={"model": "gpt-4o-mini", "metadata": {}}, + ) + ] + return tuple(getattr(served_chunk.choices[0], "token_ids", None) for served_chunk in served) + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_drops_token_ids_that_decode_to_the_unmasked_text(monkeypatch): + served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=True) + + assert served_token_ids + assert all(token_ids is None for token_ids in served_token_ids) + + +@pytest.mark.asyncio +async def test_post_call_stream_keeps_token_ids_when_no_guardrail_runs_on_the_request(monkeypatch): + served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=False) + + assert served_token_ids == ([1, 2, 3], [4]) + + @pytest.mark.asyncio async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch): from litellm.caching.caching import DualCache From 1d2858c4492bc20aac45fdb52f659f219e750757 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:15:47 -0700 Subject: [PATCH 12/12] fix(proxy): skip token-id-only chunks on guarded streams and put stop_reason on the final text chunk --- .../litellm_core_utils/streaming_handler.py | 2 + litellm/proxy/utils.py | 37 ++++++++++- .../test_proxy_logging_hook_detection.py | 62 +++++++++++++++++++ .../test_streaming_iterator.py | 12 ++++ 4 files changed, 110 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 15a6986ff4f..d9d7efa02d4 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1314,6 +1314,8 @@ class CustomStreamWrapper: _apply_choice_fields(model_response.choices[0], chunk_choice_fields) if sibling_finish_reason is not None: self._finish_sibling_choice(model_response.choices[0], sibling_finish_reason) + elif response_obj.get("is_finished"): + self._apply_finish_choice_fields(model_response.choices[0]) self._optional_combine_thinking_block_in_choices(model_response=model_response) return model_response diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index fba15eef00e..18b95c2058a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -172,7 +172,7 @@ from litellm.litellm_core_utils.served_output_texts import ( ) from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model from litellm.llms import load_guardrail_translation_mappings -from litellm.llms.base_llm.guardrail_translation.utils import stream_item_items +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, @@ -706,6 +706,37 @@ def _drop_stream_chunk_token_ids(chunk: object) -> None: drop_token_ids(choice) +_PLAIN_STREAM_FIELDS: Final = TypeAdapter(dict[str, object]) +_EMPTY_STREAM_VALUES: Final[tuple[object, ...]] = (None, "", [], {}) + + +def _holds_stream_payload(value: object) -> bool: + if isinstance(value, dict): + nested: Final = _PLAIN_STREAM_FIELDS.validate_python(value).values() + return any(item not in _EMPTY_STREAM_VALUES for item in nested) + return value not in _EMPTY_STREAM_VALUES + + +def _choice_carries_only_token_ids(choice: object) -> bool: + if stream_item_field(choice, "token_ids") is None or not isinstance(choice, BaseModel): + return False + fields: Final = _PLAIN_STREAM_FIELDS.validate_python(choice.model_dump()) + return not any(_holds_stream_payload(value) for key, value in fields.items() if key not in ("index", "token_ids")) + + +def _is_token_id_only_chunk(chunk: object) -> bool: + choices: Final = stream_item_items(chunk, "choices") + return bool(choices) and all(_choice_carries_only_token_ids(choice) for choice in choices) + + +async def _without_token_id_only_chunks(response: AsyncIterable[object]) -> AsyncGenerator[object, None]: + async for chunk in response: + if _is_token_id_only_chunk(chunk): + _drop_stream_chunk_token_ids(chunk) + continue + yield chunk + + def _pipelines_stream_live(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> bool: return all(_pipeline_step_streams_live(step) for step in _pipeline_steps(pipelines)) @@ -4540,7 +4571,7 @@ class ProxyLogging: self, response, user_api_key_dict: UserAPIKeyAuth, - request_data: dict, + request_data: dict[str, object], ): """ Allow user to modify outgoing streaming data -> Given a whole response iterator. @@ -4574,7 +4605,6 @@ class ProxyLogging: # Merge model-level guardrails before checking which guardrails to run request_data = check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router) - current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines) @@ -4590,6 +4620,7 @@ class ProxyLogging: drops_token_ids: Final = bool(post_call_pipelines) or any( isinstance(callback, CustomGuardrail) for callback, _ in running_overrides ) + current_response = _without_token_id_only_chunks(response) if drops_token_ids else response guarded_layers: Final[list[AsyncGenerator[object, None]]] = [] # mutable-ok: closed on disconnect for resolved_callback, kind in running_overrides: effective_kind = ( diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index 5895ca4882b..7373589d9df 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -400,6 +400,68 @@ async def test_post_call_stream_keeps_token_ids_when_no_guardrail_runs_on_the_re assert served_token_ids == ([1, 2, 3], [4]) +def _completion_chunk(text: str, token_ids: list[int], finish_reason: str | None = None) -> object: + from litellm.types.utils import TextChoices, TextCompletionResponse + + return TextCompletionResponse( + object="text_completion", + choices=[TextChoices(index=0, text=text, finish_reason=finish_reason, token_ids=token_ids)], + ) + + +def _guarded_completions_stream(monkeypatch, action: str, held_back: tuple[object, ...]): + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + + monkeypatch.setattr( + litellm, + "callbacks", + [ + ContentFilterGuardrail( + guardrail_name="output-filter", + blocked_words=[BlockedWord(keyword="7", action=ContentFilterAction(action))], + event_hook="post_call", + default_on=True, + ) + ], + ) + + async def upstream(): + for chunk in held_back: + yield chunk + yield _completion_chunk(", 6, 7", [21, 11, 22], finish_reason="stop") + + return ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=upstream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/completions"), + request_data={"model": "qwen3", "metadata": {}}, + ) + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_blocks_completions_before_serving_token_id_only_chunks(monkeypatch): + stream = _guarded_completions_stream( + monkeypatch, "BLOCK", (_completion_chunk("", [11]), _completion_chunk("", [220])) + ) + + with pytest.raises(HTTPException): + await stream.__anext__() + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_serves_no_chunk_left_empty_by_the_token_id_drop(monkeypatch): + held_back = (_completion_chunk("", [11]), _completion_chunk("", [220])) + + served = [chunk async for chunk in _guarded_completions_stream(monkeypatch, "MASK", held_back)] + + assert [(chunk.choices[0].text, chunk.choices[0].token_ids) for chunk in served] == [ + (", 6, [KEYWORD_REDACTED]", None) + ] + assert [chunk.choices[0].token_ids for chunk in held_back] == [None, None] + + @pytest.mark.asyncio async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch): from litellm.caching.caching import DualCache diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index 8f7fd6fa327..571451164d9 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -131,6 +131,18 @@ def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text } +def test_text_completion_stream_puts_stop_reason_on_the_final_text_chunk() -> None: + responses: Final = _text_completion_stream( + (_completion_chunk(", 5", None), _completion_chunk(", 6", "stop", stop_reason=", 7")) + ) + + assert [(response.choices[0].get("text"), response.choices[0].get("stop_reason")) for response in responses] == [ + (", 5", None), + (", 6", ", 7"), + (None, ", 7"), + ] + + @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(