scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 359

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

Browse source 
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
This commit is contained in:
mateo-berri 2026-10-29 07:48:40 -07:00
parent efefc437ce
commit 4b1069df61
8 changed files with 242 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
View file

@ -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(

13
litellm/litellm_core_utils/streaming_handler.py
View file

@ -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

28
litellm/proxy/proxy_server.py
View file

@ -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:

7
litellm/text_completions/streaming_iterator.py
View file

@ -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"] = []

31
tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py
View file

@ -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

83
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -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

76
tests/unit/proxy/proxy_server/test_streaming_helpers.py
View file

@ -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``)."""

18
tests/unit/text_completions/test_streaming_iterator.py
View file

@ -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")