mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
perf(streaming): stop raising caught AttributeErrors inside pydantic on every streamed chunk (#49166)
* perf(streaming): stop raising AttributeErrors inside pydantic on every streamed chunk Every later plain chunk on openai and azure streams raised and caught 12 AttributeErrors inside pydantic. Assigning an undeclared extra goes through a class lookup that raises in the model metaclass and is never memoized, and the hidden_params getter went through the private attribute fallback set_extra_attribute writes the extra and fields_set directly when a cached per-class check says pydantic would do the same, and falls back to setattr otherwise. The hidden_params getter reads the private storage directly. Emitted chunk JSON is unchanged * fix(streaming): match pydantic's own extra-field bookkeeping on every supported version Take the direct extra write only when the installed pydantic itself adds an assigned extra to model_fields_set (2.13 and later) and the class defines no attribute of that name, so pydantic 2.11, 2.12 and 2.14 see the same model state as a plain setattr. A deleted _hidden_params raises AttributeError again, and the mutable test models carry frozen-ok reasons. * fix(streaming): keep the normal write for models that do not allow extra fields openai's construct gives every model an extras dict, so a forbid model built that way still needs setattr's ValueError instead of a silent extra write. * test(streaming): cover a ClassVar declared without a value in the extra-field parity test * fix(streaming): type the streamed choice's logprobs as the parsed value main now stores --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tavren27
GitHub
parent
8ac9b75737
commit
9432bf6b8e
6 changed files with 388 additions and 53 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue