mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
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 <yassin@berri.ai> 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 <mateo@berri.ai> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Yassin Kortam
GitHub
parent
edf4097baf
commit
ce4a6cbf5b
7 changed files with 871 additions and 39 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue