From ce4a6cbf5b7aa90d7ea25d1b313142b8b92777ff Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 19 Oct 2026 19:42:32 -0700 Subject: [PATCH] perf(streaming): cut per-chunk pydantic round trips on openai-compatible streams (#47030) * perf(streaming): skip full-chunk dump and delta rebuild on openai-compatible chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * perf(streaming): reuse the built choice and delta on plain openai-compatible chunks Co-authored-by: yassin Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(streaming): copy upstream chunk attributes in upstream key order The scalar fast path iterated a frozenset of keys, so the SSE key order of chunk extras like service_tier and obfuscation changed with each worker's hash seed * test(streaming): count chunk_creator calls with a profiler instead of patching classes * test(streaming): check service_tier still passes through chunk_creator * fix(streaming): reuse the built choice for chunks with upstream choice extras * test(streaming): cover ModelResponseStream chunks on the choice reuse path * fix(streaming): keep main's delta rebuild for deltas with nested values Clearing role in place on a delta that holds tool calls left its nested fields unset, so chunk.model_dump() dropped keys like a custom tool call's name on later chunks. Deltas holding only scalars keep the in-place path. * perf(streaming): build each plain OpenAI chunk's delta once and skip pydantic attribute misses Later plain ChatCompletionChunks now build one Delta and one StreamingChoices, with no full-chunk model_dump, and read stored values instead of attribute probes that raise and catch AttributeError. Nested-value chunks no longer pay the dict(model) misses that made them slower than the merge base. * fix(streaming): read hidden params a chunk carries as a pydantic extra * test(streaming): pin logprobs, Azure filter extras, and audio and image deltas on later OpenAI chunks * test(streaming): cover role, reasoning, choice extras and billing on plain OpenAI SDK chunks --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 26 +- litellm/litellm_core_utils/hidden_params.py | 20 +- .../litellm_core_utils/streaming_handler.py | 166 ++++-- .../test_openai_sdk_plain_chunk_fields.py | 142 ++++++ .../litellm_core_utils/test_core_helpers.py | 45 +- .../litellm_core_utils/test_hidden_params.py | 37 ++ .../test_streaming_handler.py | 474 +++++++++++++++++- 7 files changed, 871 insertions(+), 39 deletions(-) create mode 100644 tests/integration/streaming/test_openai_sdk_plain_chunk_fields.py diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index b7bfd5eeeb3..1e5bc7ecba4 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -612,10 +612,30 @@ def preserve_upstream_non_openai_attributes( Preserve non-OpenAI attributes from the original chunk. """ # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings - expected_keys: Final = set(type(model_response).model_fields.keys()).union({"usage"}) - for key, value in original_chunk.model_dump().items(): - if key not in expected_keys: + expected_keys: Final = frozenset(type(model_response).model_fields) | frozenset({"usage"}) + candidate_keys: Final = tuple( + key + for key in (*type(original_chunk).model_fields, *(original_chunk.__pydantic_extra__ or {})) + if key not in expected_keys + ) + if not candidate_keys: + return + candidate_values: Final[tuple[tuple[str, object], ...]] = tuple( + ( + key, + cast( # cast-ok: Pydantic extras are scalar-checked before assignment + object, + getattr(original_chunk, key), + ), + ) + for key in candidate_keys + ) + if all(value is None or isinstance(value, (str, int, float, bool)) for _, value in candidate_values): + for key, value in candidate_values: setattr(model_response, key, value) + return + for key, value in original_chunk.model_dump(include=frozenset(candidate_keys)).items(): + setattr(model_response, key, value) def safe_deep_copy(data): diff --git a/litellm/litellm_core_utils/hidden_params.py b/litellm/litellm_core_utils/hidden_params.py index 4cb9dcd53f1..f8d44d44cdc 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -1,6 +1,8 @@ -from collections.abc import Iterator, MutableMapping +from collections.abc import Iterator, Mapping, MutableMapping from typing import Final, Protocol, cast +from pydantic import BaseModel + from litellm.types.llms.base import HiddenParams _HIDDEN_PARAMS_ATTR: Final = "_hidden_params" @@ -47,8 +49,22 @@ class HiddenParamsModelView(MutableMapping[str, object]): return frozenset(self._hidden_params.model_fields_set) | frozenset(self._hidden_params.model_extra or {}) +def stored_attribute(obj: object, name: str) -> object: + if not isinstance(obj, BaseModel): + return getattr(obj, name, None) + values: Final[Mapping[str, object]] = vars(obj) + if name in values: + return values[name] + extras: Final[Mapping[str, object]] = obj.__pydantic_extra__ or {} + return extras.get(name) + + def _get_hidden_params_storage(obj: object) -> object | None: - return obj.get(_HIDDEN_PARAMS_ATTR) if isinstance(obj, dict) else getattr(obj, _HIDDEN_PARAMS_ATTR, None) + if isinstance(obj, dict): + return obj.get(_HIDDEN_PARAMS_ATTR) + if isinstance(obj, BaseModel) and _HIDDEN_PARAMS_ATTR in type(obj).__private_attributes__: + return getattr(obj, _HIDDEN_PARAMS_ATTR, None) + return stored_attribute(obj, _HIDDEN_PARAMS_ATTR) def get_hidden_params_storage(obj: object) -> dict[str, object] | HiddenParams | None: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index e0155f88de1..a627b26b9af 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -15,6 +15,8 @@ from typing import Any, Final, NoReturn, Protocol, TypeVar, cast import anyio import httpx from openai.types.chat import ChatCompletionChunk +from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice +from openai.types.chat.chat_completion_chunk import ChoiceDelta from pydantic import BaseModel, ValidationError from typing_extensions import NotRequired, TypedDict @@ -22,7 +24,7 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params +from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, stored_attribute from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) @@ -88,6 +90,30 @@ def _apply_provider_fields(response: ModelResponseStream, fields: Mapping[str, o setattr(response, key, value) +def _prepend_upstream_choice_extras( + choice: StreamingChoices, upstream_choice: ChunkChoice, extra_keys: frozenset[str] +) -> None: + if not extra_keys: + return + upstream_extras: Final[Mapping[str, object]] = upstream_choice.model_dump(include=set(extra_keys)) + choice.__pydantic_extra__ = {**upstream_extras, **(choice.__pydantic_extra__ or {})} # rebind-ok: reuse choice + choice.__pydantic_fields_set__.update(upstream_extras) + + +def _holds_only_scalars(model: BaseModel) -> bool: + fields: Final[Iterator[tuple[str, object]]] = iter(model) + return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields) + + +def _delta_from_upstream(upstream_delta: BaseModel | None) -> Delta | None: + if upstream_delta is None: + return Delta() + try: + return Delta(**dict(iter(upstream_delta))) + except Exception: + return None + + def _is_transport_timeout(error: BaseException) -> bool: return isinstance(error, httpx.TimeoutException) or ( isinstance(error, Timeout) and isinstance(error.__cause__, httpx.TimeoutException) @@ -252,6 +278,7 @@ class CustomStreamWrapper: self.completion_stream = completion_stream self.sent_first_chunk = False self.sent_last_chunk = False + self._openai_handler_receives_chat_chunks = False self._failure_logging: asyncio.Task[None] | None = None self._closed = False self._stream_created_time: float = time.time() @@ -789,7 +816,10 @@ class CustomStreamWrapper: raise e def model_response_creator( - self, chunk: dict | None = None, hidden_params: Mapping[str, object] | None = None + self, + chunk: dict | None = None, + hidden_params: Mapping[str, object] | None = None, + delta: Delta | None = None, ) -> ModelResponseStream: _model: Final = self._cached_model_name _logging_obj_llm_provider: Final = self._cached_logging_llm_provider @@ -802,7 +832,11 @@ class CustomStreamWrapper: if chunk: args.update({k: v for k, v in chunk.items() if k != "stream"}) - model_response: Final = ModelResponseStream(**args) + model_response: Final = ( + ModelResponseStream(**args) + if delta is None + else ModelResponseStream(**args, choices=[StreamingChoices(delta=delta)]) + ) if self.response_id is not None: model_response.id = self.response_id elif model_response.id: @@ -943,6 +977,39 @@ class CustomStreamWrapper: model_response.choices[0].delta = Delta(**_initial_delta) return model_response + def _prebuilt_plain_delta(self, chunk: object) -> Delta | None: + if ( + not self._openai_handler_receives_chat_chunks + or not isinstance(chunk, ChatCompletionChunk) + or not chunk.choices + ): + return None + upstream_delta: Final = cast( # cast-ok: [LIT006] runtime delta may be absent + ChoiceDelta | None, + chunk.choices[0].delta, + ) + if ( + upstream_delta is None + or upstream_delta.function_call is not None + or upstream_delta.tool_calls is not None + or self._has_any_special_delta_attributes(upstream_delta) + ): + return None + return _delta_from_upstream(upstream_delta) + + def _set_role_on_fresh_delta(self, model_response: ModelResponseStream) -> None: + delta: Final = model_response.choices[0].delta + if self.sent_first_chunk is False: + delta.role = "assistant" + self.sent_first_chunk = True + return + if not _holds_only_scalars(delta): + self.strip_role_from_delta(model_response) + return + delta.role = None + if "provider_specific_fields" not in delta.model_fields_set: + delta.provider_specific_fields = None + def _has_special_delta_content(self, model_response: ModelResponseStream) -> bool: """ Check if the delta contains special content types (tool_calls, function_call, audio, or image). @@ -975,11 +1042,11 @@ class CustomStreamWrapper: """ return self.strip_role_from_delta(model_response) - def _has_special_delta_attribute(self, delta, attribute_name: str) -> bool: + def _has_special_delta_attribute(self, delta: object, attribute_name: str) -> bool: """ Check if delta has a specific attribute and it's not None. """ - return delta is not None and getattr(delta, attribute_name, None) is not None + return delta is not None and stored_attribute(delta, attribute_name) is not None def _copy_delta_attribute(self, source_delta, target_delta, attribute_name: str) -> None: """ @@ -987,7 +1054,7 @@ class CustomStreamWrapper: """ setattr(target_delta, attribute_name, getattr(source_delta, attribute_name)) - def _has_any_special_delta_attributes(self, delta) -> bool: + def _has_any_special_delta_attributes(self, delta: object) -> bool: """ Check if delta has any special attributes (audio, image). """ @@ -1011,6 +1078,7 @@ class CustomStreamWrapper: completion_obj: dict[str, Any], model_response: ModelResponseStream, response_obj: dict[str, Any], + plain_delta_created: bool = False, ): from litellm.litellm_core_utils.core_helpers import ( preserve_upstream_non_openai_attributes, @@ -1030,32 +1098,66 @@ class CustomStreamWrapper: original_chunk: Final = response_obj.get("original_chunk", None) if original_chunk: if len(original_chunk.choices) > 0: - choices: Final = [] - for choice in original_chunk.choices: - try: - if isinstance(choice, BaseModel): - choice_json = choice.model_dump() - choice_json.pop( - "finish_reason", None - ) # for mistral etc. which return a value in their last chunk (not-openai compatible). - choices.append(StreamingChoices(**choice_json)) - except Exception: - choices.append(StreamingChoices()) - setattr(model_response, "choices", choices) + original_chunk_for_reuse: Final[object] = cast( # cast-ok: runtime type-check before access + object, + original_chunk, + ) + upstream_choice: Final = ( + original_chunk_for_reuse.choices[0] + if plain_delta_created + and isinstance(original_chunk_for_reuse, ChatCompletionChunk) + and len(original_chunk_for_reuse.choices) == 1 + else None + ) + upstream_delta: Final = cast( # cast-ok: [LIT006] runtime delta may be absent + ChoiceDelta | None, + getattr(upstream_choice, "delta", None), + ) + upstream_choice_extra_keys: Final = frozenset( + (upstream_choice.__pydantic_extra__ or {}) if upstream_choice is not None else () + ) + can_reuse_choice: Final = ( + upstream_choice is not None + and upstream_delta is not None + and "enhancements" not in upstream_choice_extra_keys + and upstream_choice.logprobs is None + and _holds_only_scalars(upstream_delta) + ) + if can_reuse_choice and upstream_choice is not None: + model_response.choices[0].finish_reason = None # rebind-ok: keep reused choice + model_response.choices[0].index = upstream_choice.index # rebind-ok: keep reused choice + _prepend_upstream_choice_extras( + model_response.choices[0], + upstream_choice, + upstream_choice_extra_keys - {"message"}, + ) + else: + choices: Final = [] + for choice in original_chunk.choices: + try: + if isinstance(choice, BaseModel): + choice_json = choice.model_dump() + choice_json.pop( + "finish_reason", None + ) # for mistral etc. which return a value in their last chunk (not-openai compatible). + choices.append(StreamingChoices(**choice_json)) + except Exception: + choices.append(StreamingChoices()) + setattr(model_response, "choices", choices) else: return model_response.system_fingerprint = original_chunk.system_fingerprint setattr( model_response, "citations", - getattr(original_chunk, "citations", None), + stored_attribute(original_chunk, "citations"), ) preserve_upstream_non_openai_attributes( model_response=model_response, original_chunk=original_chunk, ) - model_response = self.strip_role_from_delta(model_response) + self._set_role_on_fresh_delta(model_response) if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug( "model_response.choices[0].delta: %s", @@ -1404,6 +1506,8 @@ class CustomStreamWrapper: if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] else: # openai / azure chat model + if isinstance(chunk, ChatCompletionChunk): + self._openai_handler_receives_chat_chunks = True if self.custom_llm_provider in [ LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value, @@ -1468,8 +1572,11 @@ class CustomStreamWrapper: provider_response_model: Final = _provider_response_model(chunk) if provider_response_model is not None: self._provider_response_model = provider_response_model + prebuilt_delta: Final = self._prebuilt_plain_delta(chunk) + plain_delta_created: bool = prebuilt_delta is not None # rebind-ok: the first chunk builds its delta below model_response = self.model_response_creator( - hidden_params=_provider_hidden_params(chunk, self._provider_response_model) + hidden_params=_provider_hidden_params(chunk, self._provider_response_model), + delta=prebuilt_delta, ) response_obj: dict[str, Any] = {} try: @@ -1493,7 +1600,7 @@ class CustomStreamWrapper: # enter this branch when no content has been passed in response if hasattr(original_chunk, "id"): model_response = self.set_model_id(original_chunk.id, model_response) - if hasattr(original_chunk, "provider_specific_fields"): + if stored_attribute(original_chunk, "provider_specific_fields") is not None: model_response = self.copy_model_response_level_provider_specific_fields( original_chunk, model_response ) @@ -1545,16 +1652,10 @@ class CustomStreamWrapper: model_response.choices[0].delta = Delta() elif self._has_any_special_delta_attributes(delta): self._handle_special_delta_attributes(delta, model_response) - else: - try: - delta = ( - dict() - if original_chunk.choices[0].delta is None - else dict(original_chunk.choices[0].delta) - ) - model_response.choices[0].delta = Delta(**delta) - except Exception: - model_response.choices[0].delta = Delta() + elif prebuilt_delta is None: + plain_delta: Final = _delta_from_upstream(delta) + model_response.choices[0].delta = plain_delta if plain_delta is not None else Delta() + plain_delta_created = plain_delta is not None # rebind-ok: the first chunk builds its delta else: if self.send_stream_usage is True: model_response.choices = [] @@ -1581,6 +1682,7 @@ class CustomStreamWrapper: completion_obj=completion_obj, model_response=model_response, response_obj=response_obj, + plain_delta_created=plain_delta_created, ) return result diff --git a/tests/integration/streaming/test_openai_sdk_plain_chunk_fields.py b/tests/integration/streaming/test_openai_sdk_plain_chunk_fields.py new file mode 100644 index 00000000000..87dd7a16a41 --- /dev/null +++ b/tests/integration/streaming/test_openai_sdk_plain_chunk_fields.py @@ -0,0 +1,142 @@ +import json +import uuid +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, list_value, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +PROVIDER_FIELDS: Final = MappingProxyType({"tier": "gold"}) + + +def _frame(identity: str, choice: Mapping[str, JsonValue] | None, **chunk_fields: JsonValue) -> bytes: + value: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [] if choice is None else [{"index": 0, "finish_reason": None, **choice}], + **chunk_fields, + } + return b"data: " + json.dumps(value).encode() + b"\n\n" + + +def _upstream_stream(identity: str) -> tuple[bytes, ...]: + return ( + _frame(identity, {"delta": {"role": "assistant", "content": ""}}, system_fingerprint="fp_audit"), + _frame( + identity, + { + "delta": {"content": "Hel", "reasoning_content": "plan"}, + "content_filter_results": {"hate": {"filtered": False, "severity": "safe"}}, + }, + ), + _frame(identity, {"delta": {"role": "assistant", "content": "lo"}}), + _frame(identity, {"delta": {"content": " there"}}), + _frame(identity, {"delta": {"content": "!", "provider_specific_fields": dict(PROVIDER_FIELDS)}}), + _frame(identity, {"delta": {}, "finish_reason": "stop"}), + _frame(identity, None, usage={"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}), + b"data: [DONE]\n\n", + ) + + +def _deployment(provider: str, api_base: str) -> dict[str, JsonValue]: + if provider == "azure": + return { + "model": "azure/gpt-4.1-mini", + "api_base": api_base, + "api_key": "synthetic-azure-key", + "api_version": "2024-10-21", + } + return {"model": "openai/gpt-4o-mini", "api_base": f"{api_base}/v1"} + + +def _stream_events(gateway: Gateway, model: str, prompt: str) -> tuple[dict[str, JsonValue], ...]: + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": prompt}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read().decode() + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.iter_lines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _deltas(events: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + object_value(object_value(list_value(event["choices"])[0])["delta"]) for event in events if event["choices"] + ) + + +def _event_with_content(events: tuple[dict[str, JsonValue], ...], content: str) -> dict[str, JsonValue]: + matches: Final = tuple( + event + for event in events + if event["choices"] + and object_value(object_value(list_value(event["choices"])[0])["delta"]).get("content") == content + ) + assert len(matches) == 1, f"{content!r} arrived in {len(matches)} events" + return matches[0] + + +@pytest.mark.parametrize("provider", ["openai", "azure"]) +def test_plain_chunks_keep_text_role_and_provider_fields(gateway: Gateway, provider: str) -> None: + identity: Final = f"plain-chunks-{uuid.uuid4().hex}" + + def upstream(request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=_upstream_stream(identity)) + + with wire_server(upstream) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + **_deployment(provider, wire.url), input_cost_per_token=0.001, output_cost_per_token=0.002 + ) + events: Final = _stream_events(gateway, model, identity) + received: Final = wire.drain() + + assert len(received) == 1 + assert json.loads(received[0].body)["stream"] is True + deltas: Final = _deltas(events) + assert "".join(str(delta.get("content") or "") for delta in deltas) == "Hello there!" + assert tuple(delta.get("role") for delta in deltas if delta.get("role") is not None) == ("assistant",) + assert deltas[0].get("role") == "assistant" + assert {event["id"] for event in events} == {identity} + assert events[0].get("system_fingerprint") == "fp_audit" + hel_choice: Final = object_value(list_value(_event_with_content(events, "Hel")["choices"])[0]) + assert object_value(hel_choice["delta"]).get("reasoning_content") == "plan" + assert hel_choice.get("content_filter_results") == {"hate": {"filtered": False, "severity": "safe"}} + bang_delta: Final = object_value(object_value(list_value(_event_with_content(events, "!")["choices"])[0])["delta"]) + assert bang_delta.get("provider_specific_fields") == dict(PROVIDER_FIELDS) + assert tuple( + object_value(list_value(event["choices"])[0]).get("finish_reason") + for event in events + if event["choices"] and object_value(list_value(event["choices"])[0]).get("finish_reason") is not None + ) == ("stop",) + usages: Final = tuple(object_value(event["usage"]) for event in events if event.get("usage")) + assert tuple((usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) for usage in usages) == ( + (11, 4, 15), + ) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + database_url=gateway.database_url, + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (11, 4) + assert float(str(rows[0]["spend"])) == pytest.approx(11 * 0.001 + 4 * 0.002) diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index c3a2f032651..6459aa98e70 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -1,5 +1,6 @@ """Tests for litellm_core_utils.core_helpers module.""" +import json import logging import threading from collections.abc import Iterator, Mapping @@ -7,6 +8,7 @@ from typing import Final import httpx import pytest +from openai.types.chat import ChatCompletionChunk from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, @@ -22,13 +24,14 @@ from litellm.litellm_core_utils.core_helpers import ( map_finish_reason, normalize_drop_params, process_response_headers, + preserve_upstream_non_openai_attributes, reconstruct_model_name, redact_nested_match_and_regex_keys, set_provider_response_headers_in_hidden_params, unbind_budget_reservation_from_callbacks, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.utils import ImageResponse, TranscriptionResponse +from litellm.types.utils import Delta, ImageResponse, ModelResponseStream, StreamingChoices, TranscriptionResponse def test_parent_otel_span_reads_dict_like_metadata_by_index(): @@ -57,6 +60,46 @@ def test_native_request_id_survives_stream_header_processing(header: str) -> Non assert processed[header if header.startswith("llm_provider-") else "llm_provider-" + header] == "req_native" +def test_preserve_upstream_non_openai_attributes_keeps_scalar_and_structured_values() -> None: + original_chunk: Final = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="hello"))], + scalar_attribute="scalar", + structured_attribute={"items": ["value"]}, + ) + model_response: Final = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="hello"))] + ) + + preserve_upstream_non_openai_attributes(model_response, original_chunk) + + assert model_response.scalar_attribute == "scalar" + assert model_response.structured_attribute == {"items": ["value"]} + + +@pytest.mark.parametrize("last_extra_value", [9, {"nested": 9}], ids=["all_scalar", "one_structured"]) +def test_preserve_upstream_non_openai_attributes_keeps_upstream_key_order(last_extra_value: object) -> None: + extra_names: Final = ("zeta", "alpha", "mid", "beta", "omega", "delta", "kappa", "gamma", "theta", "iota") + extras: Final = { + **{f"extra_{name}": index for index, name in enumerate(extra_names[:-1])}, + f"extra_{extra_names[-1]}": last_extra_value, + } + original_chunk: Final = ChatCompletionChunk( + id="chatcmpl-1", + created=0, + model="gpt-x", + object="chat.completion.chunk", + choices=[], + service_tier="default", + **extras, + ) + model_response: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="hello"))]) + + preserve_upstream_non_openai_attributes(model_response, original_chunk) + + serialized_keys: Final = tuple(json.loads(model_response.model_dump_json(exclude_none=True, exclude_unset=True))) + assert serialized_keys[-len(extras) - 1 :] == ("service_tier", *extras) + + class TestBudgetReservationBinding: """The request-end release skips a reservation a cost callback has claimed, so the claim must land on the one dict auth stamped, through whichever metadata field or auth object diff --git a/tests/unit/litellm_core_utils/test_hidden_params.py b/tests/unit/litellm_core_utils/test_hidden_params.py index 53b6c0370fa..c0485a07696 100644 --- a/tests/unit/litellm_core_utils/test_hidden_params.py +++ b/tests/unit/litellm_core_utils/test_hidden_params.py @@ -1,6 +1,7 @@ from typing import Final import pytest +from openai.types.chat import ChatCompletionChunk from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, @@ -61,6 +62,42 @@ def test_get_hidden_params_storage_preserves_supported_storage_identity() -> Non assert get_hidden_params_storage(model_response) is model_storage +_UPSTREAM_CHUNK_FIELDS: Final = { + "id": "chunk", + "choices": [], + "created": 1, + "model": "m", + "object": "chat.completion.chunk", +} + + +def _chunk_with_assigned_hidden_params(hidden_params: dict[str, object]) -> ChatCompletionChunk: + chunk: Final = ChatCompletionChunk(**_UPSTREAM_CHUNK_FIELDS) + setattr(chunk, HIDDEN_PARAMS_ATTR, hidden_params) + return chunk + + +@pytest.mark.parametrize( + "build_chunk", + [ + _chunk_with_assigned_hidden_params, + lambda hidden_params: ChatCompletionChunk(**_UPSTREAM_CHUNK_FIELDS, _hidden_params=hidden_params), + lambda hidden_params: ChatCompletionChunk.model_construct( + **_UPSTREAM_CHUNK_FIELDS, _hidden_params=hidden_params + ), + ], + ids=["assigned", "constructor_extra", "model_construct_extra"], +) +def test_get_hidden_params_storage_reads_upstream_chunk_hidden_params(build_chunk) -> None: + hidden_params: Final[dict[str, object]] = {"is_finished": True} + + assert get_hidden_params_storage(build_chunk(hidden_params)) is hidden_params + + +def test_get_hidden_params_storage_is_none_for_upstream_chunk_without_hidden_params() -> None: + assert get_hidden_params_storage(ChatCompletionChunk(**_UPSTREAM_CHUNK_FIELDS)) is None + + def test_get_or_create_hidden_params_sets_empty_dict_on_plain_object() -> None: class PlainResponse: pass diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 50bcc14b681..691442e80d0 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,12 +1,18 @@ import asyncio import json +import sys import threading import time -from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import CodeType, FrameType, MappingProxyType +from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest +from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel import litellm from litellm import verbose_logger @@ -5066,6 +5072,472 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp session_id_var.set("") +@dataclass(frozen=True, slots=True) +class _ProfiledCall: + code: CodeType + arguments: Mapping[str, object] + + +def _chunk_creator_calls( + wrapper: CustomStreamWrapper, + chunks: Sequence[ChatCompletionChunk | ModelResponseStream], + codes: frozenset[CodeType], +) -> tuple[tuple[ModelResponseStream | None, ...], tuple[_ProfiledCall, ...]]: + calls: Final[list[_ProfiledCall]] = [] + + def record(frame: FrameType, event: str, _arg: object) -> None: + if event == "call" and frame.f_code in codes: + arguments: Final = dict(cast(Mapping[str, object], frame.f_locals)) + calls.append(_ProfiledCall(code=frame.f_code, arguments=MappingProxyType(arguments))) + + previous_profiler: Final = sys.getprofile() + sys.setprofile(record) + try: + outputs: Final = tuple(wrapper.chunk_creator(chunk=chunk) for chunk in chunks) + finally: + sys.setprofile(previous_profiler) + return outputs, tuple(calls) + + +def test_openai_chunk_creator_skips_full_chunk_dump_and_delta_rebuild( + logging_obj: Logging, +) -> None: + chunks: Final[list[ChatCompletionChunk]] = [ + ChatCompletionChunk( + id="chunk-1", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[{"index": 0, "delta": {"role": "assistant", "content": "hello"}, "finish_reason": None}], + ), + ChatCompletionChunk( + id="chunk-2", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[{"index": 0, "delta": {"content": " world"}, "finish_reason": None}], + service_tier="default", + citations=["https://a"], + custom_chunk_attribute="custom value", + ), + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider="openai", + logging_obj=logging_obj, + ) + chunk_model_dump: Final = ChatCompletionChunk.model_dump.__code__ + delta_init: Final = Delta.__init__.__code__ + streaming_choices_init: Final = StreamingChoices.__init__.__code__ + + outputs, calls = _chunk_creator_calls( + wrapper, chunks, frozenset({chunk_model_dump, delta_init, streaming_choices_init}) + ) + emitted_chunks: Final = tuple(output for output in outputs if output is not None) + choice_dumps: Final = tuple( + call + for call in calls + if call.code is chunk_model_dump and not isinstance(call.arguments["self"], ChatCompletionChunk) + ) + full_chunk_dumps: Final = tuple( + call + for call in calls + if call.code is chunk_model_dump + and isinstance(call.arguments["self"], ChatCompletionChunk) + and call.arguments["include"] is None + ) + delta_constructions: Final = tuple(call for call in calls if call.code is delta_init) + streaming_choice_constructions: Final = tuple(call for call in calls if call.code is streaming_choices_init) + + assert len(emitted_chunks) == 2 + assert full_chunk_dumps == () + assert choice_dumps == () + assert len(delta_constructions) <= len(emitted_chunks) * 2 + assert len(streaming_choice_constructions) <= len(emitted_chunks) + assert [chunk.choices[0].delta.content for chunk in emitted_chunks] == ["hello", " world"] + assert emitted_chunks[0].choices[0].delta.role == "assistant" + assert emitted_chunks[1].choices[0].delta.role is None + assert getattr(emitted_chunks[1], "citations") == ["https://a"] + assert getattr(emitted_chunks[1], "custom_chunk_attribute") == "custom value" + + assert getattr(emitted_chunks[1], "service_tier") == "default" + chunk2_delta: Final = emitted_chunks[1].model_dump()["choices"][0]["delta"] + assert "provider_specific_fields" in chunk2_delta + assert chunk2_delta["provider_specific_fields"] is None + + finish_output: Final = wrapper.chunk_creator( + chunk=ChatCompletionChunk( + id="chunk-3", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}], + ) + ) + + assert finish_output is not None + assert finish_output.choices[0].finish_reason == "stop" + assert finish_output.choices[0].delta.content is None + + +def _openai_chunk(chunk_id: str, delta: Mapping[str, object]) -> ChatCompletionChunk: + return ChatCompletionChunk( + id=chunk_id, + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[{"index": 0, "delta": delta, "finish_reason": None}], + ) + + +def _missed_attributes(wrapper: CustomStreamWrapper, chunks: Sequence[ChatCompletionChunk]) -> tuple[str, ...]: + attribute_lookup: Final = BaseModel.__getattr__.__code__ + misses: Final[list[str]] = [] + + def record_miss(frame: FrameType, event: str, arg: object) -> Callable[[FrameType, str, object], object]: + if event == "exception" and isinstance(arg, tuple) and isinstance(arg[1], AttributeError): + misses.append(str(frame.f_locals["item"])) + return record_miss + + def trace_lookups(frame: FrameType, event: str, _arg: object) -> Callable[[FrameType, str, object], object] | None: + return record_miss if event == "call" and frame.f_code is attribute_lookup else None + + previous_tracer: Final = sys.gettrace() + sys.settrace(trace_lookups) + try: + tuple(wrapper.chunk_creator(chunk=chunk) for chunk in chunks) + finally: + sys.settrace(previous_tracer) + return tuple(misses) + + +def test_openai_chunk_creator_builds_one_delta_and_misses_no_attribute_per_later_chunk( + logging_obj: Logging, +) -> None: + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider="openai", + logging_obj=logging_obj, + ) + wrapper.chunk_creator(chunk=_openai_chunk("chunk-1", {"role": "assistant", "content": ""})) + plain_chunks: Final = ( + _openai_chunk("chunk-2", {"content": "hello"}), + _openai_chunk("chunk-3", {"content": " world"}), + ) + nested_chunk: Final = _openai_chunk( + "chunk-4", {"content": "!", "reasoning_details": [{"type": "reasoning.text", "text": "thought"}]} + ) + delta_init: Final = Delta.__init__.__code__ + streaming_choices_init: Final = StreamingChoices.__init__.__code__ + + outputs, calls = _chunk_creator_calls(wrapper, plain_chunks, frozenset({delta_init, streaming_choices_init})) + + assert sum(call.code is delta_init for call in calls) == len(plain_chunks) + assert sum(call.code is streaming_choices_init for call in calls) == len(plain_chunks) + assert [output.choices[0].delta.content if output else None for output in outputs] == ["hello", " world"] + assert _missed_attributes(wrapper, (*plain_chunks, nested_chunk)) == () + + +def test_chunk_creator_skips_full_dump_of_model_response_stream_chunks( + logging_obj: Logging, +) -> None: + chunks: Final = ( + ModelResponseStream( + id="chunk-1", + created=1, + model="claude-model", + choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="hello"))], + ), + ModelResponseStream( + id="chunk-2", + created=1, + model="claude-model", + choices=[StreamingChoices(index=0, delta=Delta(content=" world"))], + custom_chunk_attribute="custom value", + ), + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="claude-model", + custom_llm_provider="anthropic", + logging_obj=logging_obj, + ) + model_dump: Final = BaseModel.model_dump.__code__ + + outputs, calls = _chunk_creator_calls(wrapper, chunks, frozenset({model_dump})) + emitted_chunks: Final = tuple(output for output in outputs if output is not None) + full_upstream_dumps: Final = tuple( + call + for call in calls + if any(call.arguments["self"] is chunk for chunk in chunks) and call.arguments["include"] is None + ) + + assert full_upstream_dumps == () + assert [chunk.choices[0].delta.content for chunk in emitted_chunks] == ["hello", " world"] + assert emitted_chunks[1].choices[0].delta.role is None + assert getattr(emitted_chunks[1], "custom_chunk_attribute") == "custom value" + + + +def test_openai_chunk_creator_keeps_multiple_choice_logprobs_and_extras( + logging_obj: Logging, +) -> None: + chunk: Final = ChatCompletionChunk( + id="chunk-with-multiple-choices", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + citations=["https://a"], + choices=[ + { + "index": 3, + "delta": {"content": "first"}, + "finish_reason": None, + "logprobs": {"content": []}, + }, + { + "index": 7, + "delta": {"content": "second"}, + "finish_reason": None, + "custom_choice_attribute": "choice extra", + }, + ], + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider="openai", + logging_obj=logging_obj, + ) + + result: Final = wrapper.chunk_creator(chunk=chunk) + + assert result is not None + assert [choice.index for choice in result.choices] == [3, 7] + result_logprobs = result.choices[0].logprobs + chunk_logprobs = chunk.choices[0].logprobs + assert result_logprobs is not None + assert chunk_logprobs is not None + assert result_logprobs.model_dump() == chunk_logprobs.model_dump() + assert getattr(result.choices[1], "custom_choice_attribute") == "choice extra" + assert getattr(result, "citations") == ["https://a"] + + +@pytest.mark.parametrize( + ("choice_extras", "expected_choice_keys", "expected_choice_constructions"), + [ + pytest.param( + { + "content_filter_results": {"hate": {"filtered": False, "severity": "safe"}}, + "message": {"role": "assistant", "content": "hello"}, + }, + ["content_filter_results", "index", "delta"], + 1, + id="azure-content-filter-results", + ), + pytest.param( + {"enhancements": {"grounding": {"lines": []}}}, + ["index", "delta", "enhancements"], + 2, + id="enhancements", + ), + ], +) +def test_chunk_creator_keeps_upstream_choice_extras_in_upstream_order( + logging_obj: Logging, + choice_extras: dict[str, object], + expected_choice_keys: list[str], + expected_choice_constructions: int, +) -> None: + chunk: Final = ChatCompletionChunk( + id="chunk-with-choice-extras", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[{"index": 0, "delta": {"content": "hello"}, "finish_reason": None, **choice_extras}], + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider="azure", + logging_obj=logging_obj, + ) + streaming_choices_init: Final = StreamingChoices.__init__.__code__ + + outputs, calls = _chunk_creator_calls(wrapper, [chunk], frozenset({streaming_choices_init})) + + result: Final = outputs[0] + assert result is not None + choice_dump: Final = result.choices[0].model_dump(exclude_none=True) + assert list(choice_dump) == expected_choice_keys + assert choice_dump["delta"]["content"] == "hello" + assert {key: choice_dump[key] for key in choice_extras if key != "message"} == { + key: value for key, value in choice_extras.items() if key != "message" + } + assert len(calls) == expected_choice_constructions + + +def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj: Logging) -> None: + chunk: Final = ChatCompletionChunk( + id="chunk-with-index-1", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[{"index": 1, "delta": {"content": "hello"}, "finish_reason": None}], + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider="openai", + logging_obj=logging_obj, + ) + + result: Final = wrapper.chunk_creator(chunk=chunk) + + assert result is not None + assert result.choices[0].index == 1 + assert result.choices[0].delta.content == "hello" + + +def _later_chunk_outputs( + custom_llm_provider: str, later_chunks: Sequence[ChatCompletionChunk], logging_obj: Logging +) -> tuple[ModelResponseStream, ...]: + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) + wrapper.chunk_creator(chunk=_openai_chunk("chunk-0", {"role": "assistant", "content": ""})) + outputs: Final = tuple(wrapper.chunk_creator(chunk=chunk) for chunk in later_chunks) + return tuple(output for output in outputs if output is not None) + + +def _wire_dump(chunk: ModelResponseStream) -> dict[str, object]: + return chunk.model_dump(exclude_none=True, exclude_unset=True) + + +def test_openai_chunk_creator_keeps_single_choice_logprobs_on_later_chunks(logging_obj: Logging) -> None: + logprobs: Final = {"content": [{"token": "hello", "logprob": -0.25, "bytes": [104], "top_logprobs": []}]} + chunk: Final = ChatCompletionChunk( + id="chunk-1", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[{"index": 0, "delta": {"content": "hello"}, "finish_reason": None, "logprobs": logprobs}], + ) + + outputs: Final = _later_chunk_outputs("openai", [chunk], logging_obj) + + assert [_wire_dump(output)["choices"][0]["logprobs"] for output in outputs] == [logprobs] + + +def test_azure_chunk_creator_streams_choice_content_filter_results_on_later_chunks(logging_obj: Logging) -> None: + filter_results: Final = {"hate": {"filtered": False, "severity": "safe"}} + chunk: Final = ChatCompletionChunk( + id="chunk-1", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[ + {"index": 0, "delta": {"content": "hello"}, "finish_reason": None, "content_filter_results": filter_results} + ], + ) + + outputs: Final = _later_chunk_outputs("azure", [chunk], logging_obj) + + assert [_wire_dump(output)["choices"][0]["content_filter_results"] for output in outputs] == [filter_results] + assert ["content_filter_results" in output.choices[0].model_fields_set for output in outputs] == [True] + + +@pytest.mark.parametrize( + "special_delta", + [ + pytest.param({"audio": {"id": "audio_1", "data": "AAAA", "transcript": "hi"}}, id="audio"), + pytest.param({"images": [{"type": "image_url", "image_url": {"url": "data:,"}, "index": 0}]}, id="images"), + ], +) +def test_openai_chunk_creator_dumps_special_delta_attributes_like_other_deltas( + logging_obj: Logging, special_delta: dict[str, object] +) -> None: + chunk: Final = ChatCompletionChunk.model_construct( + id="chunk-1", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[{"index": 0, "delta": special_delta, "finish_reason": None}], + ) + + outputs: Final = _later_chunk_outputs("openai", [chunk], logging_obj) + + assert [output.model_dump()["choices"][0]["delta"] for output in outputs] == [ + { + "audio": None, + "content": None, + "function_call": None, + "provider_specific_fields": None, + "role": None, + "tool_calls": None, + **special_delta, + } + ] + + +def test_openai_chunk_creator_keeps_custom_tool_name_key_on_later_chunks( + logging_obj: Logging, +) -> None: + chunks: Final = ( + ChatCompletionChunk.model_construct( + id="chunk-1", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[ + { + "index": 0, + "delta": { + "role": "assistant", + "tool_calls": [ + {"index": 0, "id": "call_1", "type": "custom", "custom": {"name": "run_python", "input": ""}} + ], + }, + "finish_reason": None, + } + ], + ), + ChatCompletionChunk.model_construct( + id="chunk-2", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[ + { + "index": 0, + "delta": {"tool_calls": [{"index": 0, "custom": {"input": "print"}}]}, + "finish_reason": None, + } + ], + ), + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider="openai", + logging_obj=logging_obj, + ) + + emitted_chunks: Final = tuple( + output for output in (wrapper.chunk_creator(chunk=chunk) for chunk in chunks) if output is not None + ) + + assert [chunk.model_dump()["choices"][0]["delta"]["tool_calls"][0]["custom"] for chunk in emitted_chunks] == [ + {"name": "run_python", "input": ""}, + {"name": None, "input": "print"}, + ] + + def test_chunk_creator_preserves_hidden_provider_specific_fields_from_parsed_chunk(): wrapper = CustomStreamWrapper( completion_stream=None,