diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index b58f40af155..555a566497d 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -41,7 +41,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) _REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter( list[dict[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True) diff --git a/litellm/constants.py b/litellm/constants.py index c17e776f54b..0c594ca1b86 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -694,6 +694,9 @@ 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",) +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 58fe5b4297f..6bb036e9080 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 Generator, Iterator, Mapping, Sequence +from collections.abc import Generator, Iterable, Iterator, 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, @@ -170,6 +176,28 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]: return _REPLAY_CONTENT_SLICE_RE.findall(content) +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]: + 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({}) + + _REPLAY_REASONING_FIELDS: Final = frozenset({"reasoning_content", "thinking_blocks", "reasoning_items"}) _REPLAY_LEAD_FIELDS: Final = _REPLAY_REASONING_FIELDS | {"role", "provider_specific_fields"} _REPLAY_TOOL_FIELDS: Final = frozenset({"tool_calls", "function_call"}) @@ -220,8 +248,9 @@ def _replay_frame( choice.finish_reason = finish_reason if not keeps_choice_metadata: choice.logprobs = None - if hasattr(choice, "enhancements"): - del choice.enhancements + _drop_choice_fields(choice, ("enhancements", *STREAM_CHUNK_CHOICE_FIELDS)) + if finish_reason is None: + _drop_choice_fields(choice, STREAM_FINISH_CHOICE_FIELDS) return frame @@ -325,7 +354,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 @@ -388,6 +419,7 @@ def convert_to_streaming_response( delta=delta, logprobs=logprobs, enhancements=enhancements, + **_vllm_choice_fields(choice), ) choice_list.append(choice) @@ -872,6 +904,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/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 981859c620b..ca508fc9d08 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -196,6 +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)): + 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 @@ -227,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]: @@ -251,6 +267,7 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ "compaction_blocks", "compaction_start", "compaction_delta", + "token_ids", ) if key in fields }, @@ -422,6 +439,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/stream_usage_chunks.py b/litellm/litellm_core_utils/stream_usage_chunks.py index eda57c7ea12..d073ddc2724 100644 --- a/litellm/litellm_core_utils/stream_usage_chunks.py +++ b/litellm/litellm_core_utils/stream_usage_chunks.py @@ -1,6 +1,7 @@ from collections.abc import Iterable from typing import Final +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.types.utils import ModelResponseStream, StreamingChoices @@ -9,12 +10,19 @@ def is_empty_streaming_choice(choice: StreamingChoices) -> bool: return False if getattr(choice, "logprobs", None) is not None: return False + if not _all_values_none(_vllm_choice_values(choice)): + return False delta: Final = choice.delta if hasattr(choice, "delta") else None if delta is None: return True return _all_values_none(delta.model_dump().values()) +def _vllm_choice_values(choice: StreamingChoices) -> Iterable[object]: + choice_extras: Final = choice.model_extra or {} + return (choice_extras.get(field) for field in VLLM_CHOICE_FIELDS) + + def _all_values_none(values: Iterable[object]) -> bool: return all(value is None for value in values) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 4857973d51d..6b2557399ff 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 ( PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, UPSTREAM_USAGE_COST, @@ -434,6 +435,8 @@ class ChunkProcessor: chunks, unreported="tool_calls" if _carries_tool_calls(chunks) else "stop" ), message={"role": ChunkProcessor._get_role_from_chunks(chunks), "content": ""}, + provider_specific_fields=None, + **_streamed_vllm_choice_fields(chain.from_iterable(chunk.get("choices") or () for chunk in chunks)), ) @staticmethod @@ -504,6 +507,7 @@ class ChunkProcessor: "index": 0, "message": {"role": role, "content": ""}, "finish_reason": finish_reason, + **_streamed_vllm_choice_fields(_first_choices_of_index_zero(chunks)), } ] ), @@ -1371,6 +1375,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), ) @@ -1504,6 +1509,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 484edd869cb..94169bd2d6e 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 ( UPSTREAM_USAGE_COST, @@ -37,6 +38,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, ) @@ -137,6 +139,62 @@ def _prepend_upstream_choice_extras( choice.__pydantic_fields_set__.update(upstream_extras) +_NO_UPSTREAM_CHOICE_FIELDS: Final[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: + match getattr(upstream_chunk, "choices", None): + case [BaseModel() as first_choice, *_]: + return first_choice + case _: + 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 _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 + if not upstream_extras: + return _NO_UPSTREAM_CHOICE_FIELDS + return MappingProxyType( + {field: upstream_extras[field] for field in fields if upstream_extras.get(field) 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 _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) @@ -375,6 +433,8 @@ class CustomStreamWrapper: self._provider_response_model: str | None = None self._upstream_provider: str | None = None self.received_finish_reason: str | None = None + 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|>", @@ -819,6 +879,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 @@ -843,7 +904,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 @@ -856,11 +917,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": choice_index, + "sibling_finish_reason": finish_reason if choice_index != 0 else None, } except Exception as e: @@ -1171,6 +1236,11 @@ class CustomStreamWrapper: ) is_chunk_non_empty: Final = self.is_chunk_non_empty(completion_obj, model_response, response_obj) + 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) + 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() @@ -1258,16 +1328,40 @@ 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 + _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 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( + 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: @@ -1305,12 +1399,20 @@ 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") + and _response_choice_index(response_obj) == model_response.choices[0].index + ): + _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 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) @@ -1533,6 +1635,8 @@ class CustomStreamWrapper: completion_obj["content"] = response_obj["text"] print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: + 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"] @@ -1596,6 +1700,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 @@ -1623,6 +1728,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._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) @@ -1966,6 +2072,75 @@ 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_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: + 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 = _response_choice_index(response_obj) + 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, + response_obj: Mapping[str, object], + chunk_choice_fields: Mapping[str, object], + ) -> ModelResponseStream | None: + 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 + 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} @@ -1985,6 +2160,7 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = _finish_reason else: model_response.choices[0].finish_reason = "stop" + self._apply_finish_choice_fields(model_response.choices[0]) ## if tool use if ( 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/proxy_server.py b/litellm/proxy/proxy_server.py index b9b15d332e9..2ad940d7c52 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -88,6 +88,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.access_windows import access_windows_config_error from litellm.litellm_core_utils.add_retry_fallback_headers import ( @@ -10152,6 +10153,11 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) - return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending +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( chunk: ModelResponseStream, ) -> bytes | None: @@ -10172,7 +10178,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: + 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) @@ -10201,7 +10211,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 @@ -10236,7 +10246,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)) async def _apply_streaming_chunk_hooks( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9e237ffcb57..787c5d3694b 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 ( @@ -165,6 +166,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.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -174,6 +176,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_field, stream_item_items from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, @@ -702,6 +705,42 @@ 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) + + +_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)) @@ -4676,7 +4715,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. @@ -4710,20 +4749,24 @@ 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) + 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 + ) + 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 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 ( @@ -4781,6 +4824,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) @@ -10148,7 +10193,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 61e47e9d7af..6bba56a5afe 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 VLLM_CHOICE_FIELDS from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse @@ -37,16 +38,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"]: - delta: Final = chunk["choices"][0]["delta"] + 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"] = 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"] response["choices"] = [text_choices] else: response["choices"] = [] diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index c4727cbeff0..20fe57d898e 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -472,6 +472,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/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 d0f5bfc6b9f..3e197740e98 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 @@ -3089,6 +3089,73 @@ def test_convert_to_model_response_object_function_output(): 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, + ] + + +@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 + + @pytest.mark.parametrize( "reported_cost", [ diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 9e385c588ae..b523ab71f54 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -561,6 +561,64 @@ 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_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/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index bbf1897b3a7..4ea8578d6b1 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 @@ -2757,6 +2757,80 @@ def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it 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_each_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 [ + {field: choice[field] for field in ("index", "stop_reason", "token_ids")} + for choice in built.model_dump()["choices"] + ] == [ + {"index": 0, "stop_reason": ",5", "token_ids": [16, 11, 17, 11, 18, 19]}, + {"index": 1, "stop_reason": None, "token_ids": [22, 11]}, + ] + + +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])] + + @pytest.mark.parametrize( ("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")] ) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 4d1699c7fe4..38c048755c4 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -10,6 +10,7 @@ import threading import time from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import chain from pathlib import Path from types import CodeType, FrameType, MappingProxyType, TracebackType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast @@ -27,12 +28,17 @@ 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, _ProviderChunkParsed, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( @@ -6673,6 +6679,293 @@ 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) +_CHOICE_INDEX: Final = TypeAdapter(int) +_DELTA: Final = TypeAdapter(Mapping[str, object]) + + +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], ...]: + 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 [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [stop_reason] + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice["stop_reason"]) + for choice in emitted + 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) + + +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, + ( + {"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}, + {"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 + ] + + +_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) + ), + ) + + +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) + ] == [] + + +@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 _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]) + for index, content in zip((first_finished, 1 - first_finished), finish_contents) + if content + ) + + 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()) + + +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) + + +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["finish_reason"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None + ] == [("stop", ", 5")] + 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/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/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..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,14 +1,16 @@ import json -from unittest.mock import MagicMock, patch +from typing import Final +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 +399,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: Final = { + "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: Final = 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: Final = _transform_vllm_chat_response(choice_extras) + + 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: Final = _transform_vllm_chat_response({}) + + 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 c37475bd539..048af8611fd 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, route_entry_points_to_stubs @@ -576,6 +583,89 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): } +@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 = 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/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_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index 182ef833069..7373589d9df 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -347,6 +347,121 @@ 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]) + + +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/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 3647146eb4e..76d94628d81 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13900,7 +13900,7 @@ def _stream_usage_generator_chunks(): def test_is_injected_stream_usage_artifact(): from litellm.litellm_core_utils.stream_usage_chunks 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 @@ -13923,6 +13923,21 @@ def test_is_injected_stream_usage_artifact(): assert is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False + 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: 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 + + 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 e5459a5a0cd..571451164d9 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -1,11 +1,15 @@ -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Sequence +from itertools import chain 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 +from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, TextCompletionResponse, Usage def test_reasoning_content_preserved_in_text_completion_wrapper(): @@ -85,6 +89,146 @@ 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}], + ) + + +def _text_completion_stream(upstream: Sequence[Completion]) -> list[TextCompletionResponse]: + 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", + ) + 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")) + + +@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", + "index": 0, + "stop_reason": ", 7", + } + + +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( + ( + _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_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_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])} + + +@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")