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