diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 2f81154ab83..20e10f63f83 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -9,12 +9,13 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import TypeIs # noqa: TID251 # narrows a duck-typed value to the Protocol its caller reads import litellm from litellm._logging import verbose_logger from litellm.constants import KUBERNETES_POD_ROUTING_KEY +from litellm.litellm_core_utils.hidden_params import set_extra_attribute from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason @@ -710,9 +711,9 @@ def preserve_upstream_non_openai_attributes( _assign_attributes(model_response, original_chunk.model_dump(include=frozenset(candidate_keys)).items()) -def _assign_attributes(target: object, values: Iterable[tuple[str, object]]) -> None: +def _assign_attributes(target: BaseModel, values: Iterable[tuple[str, object]]) -> None: for key, value in values: - setattr(target, key, value) + set_extra_attribute(target, 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 a74337d2549..c04d7614b90 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -1,8 +1,9 @@ from collections.abc import Iterator, Mapping, MutableMapping +from functools import cache, lru_cache from types import MappingProxyType from typing import Final, Protocol, cast -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, create_model from litellm.types.llms.base import HiddenParams @@ -60,6 +61,38 @@ def stored_attribute(obj: object, name: str) -> object: return extras.get(name) +@cache +def _assigned_extras_join_fields_set() -> bool: + probe: Final = create_model("_ExtraAssignmentProbe", __config__=ConfigDict(extra="allow"))() + setattr(probe, "assigned", None) + return "assigned" in probe.model_fields_set + + +@lru_cache(maxsize=1024) +def _writes_plain_extra(model_class: type[BaseModel], name: str) -> bool: + config: Final = model_class.model_config + return ( + model_class.__setattr__ is BaseModel.__setattr__ + and config.get("extra") == "allow" + and not config.get("frozen") + and not config.get("validate_assignment") + and not name.startswith("_") + and name not in model_class.__pydantic_fields__ + and name not in model_class.__class_vars__ + and not hasattr(model_class, name) + and _assigned_extras_join_fields_set() + ) + + +def set_extra_attribute(model: BaseModel, name: str, value: object) -> None: + extras: Final = model.__pydantic_extra__ + if extras is None or not _writes_plain_extra(type(model), name): + setattr(model, name, value) + return + extras[name] = value + model.__pydantic_fields_set__.add(name) + + def _get_hidden_params_storage(obj: object) -> object | None: if isinstance(obj, dict): return obj.get(_HIDDEN_PARAMS_ATTR) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 098a404ea14..b1b3e20f218 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_report from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, get_hidden_params, + set_extra_attribute, set_hidden_param, stored_attribute, ) @@ -1035,7 +1036,7 @@ class CustomStreamWrapper: _initial_delta: Final = model_response.choices[0].delta.model_dump() _initial_delta.pop("role", None) - model_response.choices[0].delta = Delta(**_initial_delta) + set_extra_attribute(model_response.choices[0], "delta", Delta(**_initial_delta)) return model_response def _prebuilt_plain_delta(self, chunk: object) -> Delta | None: @@ -1061,13 +1062,13 @@ class CustomStreamWrapper: 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" + set_extra_attribute(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 + set_extra_attribute(delta, "role", None) if "provider_specific_fields" not in delta.model_fields_set: delta.provider_specific_fields = None @@ -1186,8 +1187,8 @@ class CustomStreamWrapper: 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 + set_extra_attribute(model_response.choices[0], "finish_reason", None) + set_extra_attribute(model_response.choices[0], "index", upstream_choice.index) _prepend_upstream_choice_extras( model_response.choices[0], upstream_choice, @@ -1212,11 +1213,7 @@ class CustomStreamWrapper: else: return model_response.system_fingerprint = upstream_chunk.system_fingerprint - setattr( - model_response, - "citations", - stored_attribute(upstream_chunk, "citations"), - ) + set_extra_attribute(model_response, "citations", stored_attribute(upstream_chunk, "citations")) preserve_upstream_non_openai_attributes( model_response=model_response, original_chunk=original_chunk, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 55aaf606175..f8115571dfd 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -75,6 +75,7 @@ from litellm.types.workload_identity import ( ) from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers +from ..litellm_core_utils.hidden_params import set_extra_attribute from . import litellm_params as _litellm_params from .agents import LiteLLMSendMessageResponse from .guardrails import GuardrailEventHooks @@ -2149,6 +2150,13 @@ class UnreportedUsage(Usage): class StreamingChoices(OpenAIObject): + if TYPE_CHECKING: + finish_reason: str | None + index: int + delta: Delta + enhancements: object + logprobs: ChoiceLogprobs | JsonValue + def __init__( self, finish_reason=None, @@ -2162,25 +2170,22 @@ class StreamingChoices(OpenAIObject): # https://github.com/BerriAI/litellm/issues/8455 params.pop("message", None) super().__init__(**params) - if finish_reason: - self.finish_reason = map_finish_reason(finish_reason) - else: - self.finish_reason = None - self.index = index + set_extra_attribute(self, "finish_reason", map_finish_reason(finish_reason) if finish_reason else None) + set_extra_attribute(self, "index", index) if delta is not None: if isinstance(delta, Delta): - self.delta = delta + set_extra_attribute(self, "delta", delta) elif isinstance(delta, dict): - self.delta = Delta(**delta) + set_extra_attribute(self, "delta", Delta(**delta)) else: - self.delta = Delta() + set_extra_attribute(self, "delta", Delta()) if enhancements is not None: - self.enhancements = enhancements + set_extra_attribute(self, "enhancements", enhancements) if isinstance(logprobs, (dict, BaseModel)): - self.logprobs = CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs) + set_extra_attribute(self, "logprobs", CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs)) else: - self.logprobs = logprobs + set_extra_attribute(self, "logprobs", logprobs) def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator @@ -2234,7 +2239,10 @@ class ModelResponseBase(OpenAIObject): @property def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation - return self._hidden_params + private: Final = self.__pydantic_private__ + if private is None or "_hidden_params" not in private: + return self._hidden_params + return private["_hidden_params"] @hidden_params.setter def hidden_params(self, hidden_params: dict[str, object]) -> None: # mutable-ok: API requires mutation diff --git a/tests/unit/litellm_core_utils/test_hidden_params.py b/tests/unit/litellm_core_utils/test_hidden_params.py index c0485a07696..b32692c67ba 100644 --- a/tests/unit/litellm_core_utils/test_hidden_params.py +++ b/tests/unit/litellm_core_utils/test_hidden_params.py @@ -1,7 +1,11 @@ -from typing import Final +import pickle +from collections.abc import Callable +from typing import ClassVar, Final import pytest +from openai import BaseModel as OpenAIBaseModel from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, ConfigDict, PrivateAttr from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, @@ -9,13 +13,14 @@ from litellm.litellm_core_utils.hidden_params import ( get_hidden_params_storage, get_or_create_hidden_params, served_from_cache, + set_extra_attribute, set_hidden_param, set_hidden_params, ) from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig from litellm.types.decisions import DecisionsResponse from litellm.types.llms.base import HiddenParams -from litellm.types.utils import ModelResponse, TextChoices, TextCompletionResponse, Usage +from litellm.types.utils import ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse, Usage def test_get_and_set_hidden_params_on_plain_object() -> None: @@ -309,3 +314,133 @@ def test_served_from_cache_is_true_only_for_a_response_the_cache_stamped( assert served_from_cache(model_response) is expected assert served_from_cache({HIDDEN_PARAMS_ATTR: dict(hidden_params)}) is expected assert served_from_cache(object()) is False + + +class _ExtraModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="allow") + + declared: str | None = None + marker: ClassVar[str | None] = None + unset_marker: ClassVar[int] + _secret: str = PrivateAttr(default="hidden") + + @property + def doubled(self) -> str | None: + return self.declared + + @doubled.setter + def doubled(self, value: str) -> None: + self.declared = value * 2 + + +class _FrozenExtraModel(BaseModel): + model_config = ConfigDict(extra="allow", frozen=True) + + +class _ValidatedExtraModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="allow", validate_assignment=True) + + __pydantic_extra__: dict[str, int] + + +class _IgnoredExtraModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="ignore") + + +class _MirroringModel(BaseModel): # frozen-ok: the assignment tests need a model that accepts attribute writes + model_config = ConfigDict(extra="allow") + + def __setattr__(self, name: str, value: object) -> None: + super().__setattr__(name, value) + super().__setattr__(f"{name}_copy", value) + + +class _NoneClassAttribute: + inherited = None + + +class _InheritingExtraModel(_NoneClassAttribute, BaseModel): # frozen-ok: assignment tests need attribute writes + model_config = ConfigDict(extra="allow") + + +class _ForbiddingOpenAIModel(OpenAIBaseModel): # frozen-ok: assignment tests need attribute writes + model_config = ConfigDict(extra="forbid") + + +def _assignment_outcome(model: BaseModel, assign: Callable[[], None], name: str) -> tuple[object, ...]: + try: + assign() + except (AttributeError, ValueError) as error: + return type(error), str(error) + return model.model_dump(), getattr(model, name), model.model_fields_set, model.model_dump_json() + + +@pytest.mark.parametrize( + "build, name, value", + [ + pytest.param(lambda: _ExtraModel(existing=1, later=3), "fresh", 1, id="new-extra"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "existing", 2, id="existing-extra-keeps-order"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "declared", "x", id="declared-field"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "doubled", "ab", id="property"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "_free", 1, id="undeclared-underscore-name"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "_secret", "x", id="private-attribute"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "marker", 1, id="class-var"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "unset_marker", 1, id="class-var-without-value"), + pytest.param(lambda: _ExtraModel(existing=1, later=3), "model_config", 1, id="class-attribute"), + pytest.param(_FrozenExtraModel, "fresh", 1, id="frozen"), + pytest.param(_ValidatedExtraModel, "fresh", "5", id="validate-assignment"), + pytest.param(_IgnoredExtraModel, "fresh", 1, id="extra-ignore"), + pytest.param(_MirroringModel, "fresh", 1, id="setattr-override"), + pytest.param(_InheritingExtraModel, "inherited", 1, id="class-attribute-set-to-none"), + pytest.param( + lambda: _ForbiddingOpenAIModel.construct(existing=1), "fresh", 1, id="extra-forbid-built-by-construct" + ), + ], +) +def test_set_extra_attribute_matches_setattr(build: Callable[[], BaseModel], name: str, value: object) -> None: + through_setattr: Final = build() + through_helper: Final = build() + + expected: Final = _assignment_outcome(through_setattr, lambda: setattr(through_setattr, name, value), name) + actual: Final = _assignment_outcome(through_helper, lambda: set_extra_attribute(through_helper, name, value), name) + + assert actual == expected + + +@pytest.mark.parametrize("response_type", [ModelResponse, ModelResponseStream]) +def test_hidden_params_raises_attribute_error_once_deleted( + response_type: type[ModelResponse | ModelResponseStream], +) -> None: + response: Final = response_type() + del response._hidden_params + + assert getattr(response, "hidden_params", "fallback") == "fallback" + + +def _round_trip(response: ModelResponse | ModelResponseStream) -> ModelResponse | ModelResponseStream: + return pickle.loads(pickle.dumps(response)) + + +@pytest.mark.parametrize("response_type", [ModelResponse, ModelResponseStream]) +@pytest.mark.parametrize( + "rebuild", + [ + pytest.param(lambda response: response, id="init"), + pytest.param(lambda response: response.model_copy(), id="model-copy"), + pytest.param(lambda response: response.model_copy(deep=True), id="deep-copy"), + pytest.param(_round_trip, id="pickle"), + ], +) +def test_hidden_params_is_the_stored_dict_on_every_construction_path( + response_type: type[ModelResponse | ModelResponseStream], + rebuild: Callable[[ModelResponse | ModelResponseStream], ModelResponse | ModelResponseStream], +) -> None: + original: Final = response_type() + original.hidden_params["model_id"] = "deployment-1" + response: Final = rebuild(original) + + response.hidden_params["api_base"] = "https://example.com" + + assert response.hidden_params is response._hidden_params + assert get_hidden_params(response) == response._hidden_params + assert response._hidden_params["api_base"] == "https://example.com" diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b9c8feb223d..148b7193ef6 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -11,7 +11,7 @@ import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path -from types import CodeType, FrameType, MappingProxyType +from types import CodeType, FrameType, MappingProxyType, TracebackType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -6292,53 +6292,76 @@ def test_openai_chunk_creator_skips_full_chunk_dump_and_delta_rebuild( assert finish_output.choices[0].delta.content is None -def _openai_chunk(chunk_id: str, delta: Mapping[str, object]) -> ChatCompletionChunk: +def _openai_chunk( + chunk_id: str, + delta: Mapping[str, object], + choice_extras: Mapping[str, object] = MappingProxyType({}), + finish_reason: str | None = None, + chunk_extras: Mapping[str, object] = MappingProxyType({}), +) -> ChatCompletionChunk: return ChatCompletionChunk( id=chunk_id, created=1, model="gpt-4o", object="chat.completion.chunk", - choices=[{"index": 0, "delta": delta, "finish_reason": None}], + choices=[{"index": 0, "delta": delta, "finish_reason": finish_reason, **choice_extras}], + **chunk_extras, ) -def _missed_attributes(wrapper: CustomStreamWrapper, chunks: Sequence[ChatCompletionChunk]) -> tuple[str, ...]: - attribute_lookup: Final = BaseModel.__getattr__.__code__ - misses: Final[list[str]] = [] +def _raised_attribute_errors(wrapper: CustomStreamWrapper, chunks: Sequence[ChatCompletionChunk]) -> tuple[str, ...]: + raised: 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 + def record_raise(frame: FrameType, event: str, arg: object) -> Callable[[FrameType, str, object], object]: + if event == "exception" and isinstance(arg, tuple): + _, error, traceback = arg + if isinstance(error, AttributeError) and isinstance(traceback, TracebackType) and traceback.tb_next is None: + raised.append(f"{frame.f_code.co_qualname}: {error}") + return record_raise previous_tracer: Final = sys.gettrace() - sys.settrace(trace_lookups) + sys.settrace(record_raise) try: tuple(wrapper.chunk_creator(chunk=chunk) for chunk in chunks) finally: sys.settrace(previous_tracer) - return tuple(misses) + return tuple(raised) -def test_openai_chunk_creator_builds_one_delta_and_misses_no_attribute_per_later_chunk( - logging_obj: Logging, +_AZURE_CHOICE_EXTRAS: Final = MappingProxyType( + {"content_filter_results": {"hate": {"filtered": False, "severity": "safe"}}} +) +_OPENAI_AND_AZURE_CHOICE_EXTRAS: Final = ( + pytest.param("openai", MappingProxyType({}), id="openai"), + pytest.param("azure", _AZURE_CHOICE_EXTRAS, id="azure"), +) + + +@pytest.mark.parametrize(("custom_llm_provider", "choice_extras"), _OPENAI_AND_AZURE_CHOICE_EXTRAS) +def test_openai_chunk_creator_builds_one_delta_and_raises_no_attribute_error_per_later_chunk( + logging_obj: Logging, custom_llm_provider: str, choice_extras: Mapping[str, object] ) -> None: wrapper: Final = CustomStreamWrapper( completion_stream=iter(()), model="gpt-4o", - custom_llm_provider="openai", + custom_llm_provider=custom_llm_provider, 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"}), + + def later_chunk(chunk_id: str, delta: Mapping[str, object]) -> ChatCompletionChunk: + return _openai_chunk( + chunk_id, delta, choice_extras, chunk_extras={"service_tier": "default", "obfuscation": chunk_id} + ) + + nested_delta: Final = {"content": "!", "reasoning_details": [{"type": "reasoning.text", "text": "thought"}]} + first_chunk_of_each_shape: Final = ( + later_chunk("chunk-1", {"role": "assistant", "content": ""}), + later_chunk("chunk-2", nested_delta), ) - nested_chunk: Final = _openai_chunk( - "chunk-4", {"content": "!", "reasoning_details": [{"type": "reasoning.text", "text": "thought"}]} + tuple(wrapper.chunk_creator(chunk=chunk) for chunk in first_chunk_of_each_shape) + plain_chunks: Final = ( + later_chunk("chunk-3", {"content": "hello"}), + later_chunk("chunk-4", {"content": " world"}), ) delta_init: Final = Delta.__init__.__code__ streaming_choices_init: Final = StreamingChoices.__init__.__code__ @@ -6348,7 +6371,145 @@ def test_openai_chunk_creator_builds_one_delta_and_misses_no_attribute_per_later 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)) == () + assert _raised_attribute_errors(wrapper, (*plain_chunks, later_chunk("chunk-5", nested_delta))) == () + + +_OPENAI_STREAM_JSON: Final = ( + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":null,"index":0,' + '"delta":{"provider_specific_fields":null,"refusal":null,"content":"","role":"assistant",' + '"function_call":null,"tool_calls":null,"audio":null},"logprobs":null}],' + '"provider_specific_fields":null,"citations":null,"service_tier":"default","obfuscation":"o0"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":null,"index":0,' + '"delta":{"provider_specific_fields":null,"refusal":null,"content":"one","role":null,' + '"function_call":null,"tool_calls":null,"audio":null},"logprobs":null}],' + '"provider_specific_fields":null,"citations":null,"service_tier":"default","obfuscation":"o1"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":null,"index":0,' + '"delta":{"provider_specific_fields":null,"refusal":null,"content":" two","role":null,' + '"function_call":null,"tool_calls":null,"audio":null},"logprobs":null}],' + '"provider_specific_fields":null,"citations":null,"service_tier":"default","obfuscation":"o2"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":"stop","index":0,' + '"delta":{"provider_specific_fields":null,"content":null,"role":null,"function_call":null,' + '"tool_calls":null,"audio":null},"logprobs":null}],"provider_specific_fields":null,' + '"service_tier":"default","obfuscation":"o3"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[],"provider_specific_fields":null,' + '"usage":{"completion_tokens":3,"prompt_tokens":9,"total_tokens":12,"completion_tokens_details":null,' + '"prompt_tokens_details":null}}' + ), +) +_AZURE_STREAM_JSON: Final = ( + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"content_filter_results":{"hate":{"filtered":false,' + '"severity":"safe"}},"finish_reason":null,"index":0,"delta":{"provider_specific_fields":null,' + '"refusal":null,"content":"","role":"assistant","function_call":null,"tool_calls":null,"audio":null},' + '"logprobs":null}],"provider_specific_fields":null,"citations":null,"service_tier":"default",' + '"obfuscation":"o0"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"content_filter_results":{"hate":{"filtered":false,' + '"severity":"safe"}},"finish_reason":null,"index":0,"delta":{"provider_specific_fields":null,' + '"refusal":null,"content":"one","role":null,"function_call":null,"tool_calls":null,"audio":null},' + '"logprobs":null}],"provider_specific_fields":null,"citations":null,"service_tier":"default",' + '"obfuscation":"o1"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"content_filter_results":{"hate":{"filtered":false,' + '"severity":"safe"}},"finish_reason":null,"index":0,"delta":{"provider_specific_fields":null,' + '"refusal":null,"content":" two","role":null,"function_call":null,"tool_calls":null,"audio":null},' + '"logprobs":null}],"provider_specific_fields":null,"citations":null,"service_tier":"default",' + '"obfuscation":"o2"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[{"finish_reason":"stop","index":0,' + '"delta":{"provider_specific_fields":null,"content":null,"role":null,"function_call":null,' + '"tool_calls":null,"audio":null},"logprobs":null}],"provider_specific_fields":null,' + '"service_tier":"default","obfuscation":"o3"}' + ), + ( + '{"id":"chatcmpl-1","created":1,"model":"gpt-4o","object":"chat.completion.chunk",' + '"system_fingerprint":null,"choices":[],"provider_specific_fields":null,' + '"usage":{"completion_tokens":3,"prompt_tokens":9,"total_tokens":12,"completion_tokens_details":null,' + '"prompt_tokens_details":null}}' + ), +) + + +@pytest.mark.parametrize( + ("custom_llm_provider", "choice_extras", "expected_json"), + [ + pytest.param("openai", MappingProxyType({}), _OPENAI_STREAM_JSON, id="openai"), + pytest.param("azure", _AZURE_CHOICE_EXTRAS, _AZURE_STREAM_JSON, id="azure"), + ], +) +def test_openai_chunk_creator_emits_each_chunk_as_the_same_json( + logging_obj: Logging, + custom_llm_provider: str, + choice_extras: Mapping[str, object], + expected_json: tuple[str | None, ...], +) -> None: + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(()), + model="gpt-4o", + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + stream_options={"include_usage": True}, + ) + chunks: Final = ( + _openai_chunk( + "chatcmpl-1", + {"role": "assistant", "content": ""}, + choice_extras, + chunk_extras={"service_tier": "default", "obfuscation": "o0"}, + ), + _openai_chunk( + "chatcmpl-1", + {"content": "one"}, + choice_extras, + chunk_extras={"service_tier": "default", "obfuscation": "o1"}, + ), + _openai_chunk( + "chatcmpl-1", + {"content": " two"}, + choice_extras, + chunk_extras={"service_tier": "default", "obfuscation": "o2"}, + ), + _openai_chunk( + "chatcmpl-1", {}, choice_extras, "stop", chunk_extras={"service_tier": "default", "obfuscation": "o3"} + ), + ChatCompletionChunk( + id="chatcmpl-1", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[], + service_tier="default", + obfuscation="o4", + usage={"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12}, + ), + ) + + outputs: Final = tuple(wrapper.chunk_creator(chunk=chunk) for chunk in chunks) + + assert tuple( + output.model_copy(update={"created": 1}).model_dump_json() if output else None for output in outputs + ) == expected_json def test_chunk_creator_skips_full_dump_of_model_response_stream_chunks(