mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(proxy): keep provider fields on streamed chat completion chunks (#49398)
* fix(proxy): keep provider fields on streamed chat completion chunks The hand-written serializer for plain text chunks dropped every field it did not list, so Azure's per-choice content_filter_results, the chunk's obfuscation, and delta reasoning_details and reasoning_items never reached clients since v1.86.0. Every chunk now serializes with model_dump_json, which keeps them and is faster per chunk. * fix(proxy): send streamed chunk frames as bytes so plain text frames keep their type * fix(proxy): send null for chat chunk values JSON cannot encode instead of failing the stream * fix(proxy): drop chat chunk fields JSON cannot encode instead of failing the stream A callback or provider value pydantic cannot encode, like invalid UTF-8 bytes, a lone surrogate, a circular dict, or an unknown type, now leaves only that field out of the frame with a warning naming it. Responses API events keep raising. * fix(proxy): drop only the nested chunk value JSON cannot encode and keep lone surrogates Main removed the StreamingChoices import, so the field drop helpers now take any BaseModel. A value JSON cannot encode is left out down to the nested key or list item that holds it, and a lone surrogate goes out as an ASCII JSON escape. * fix(proxy): keep streaming when an unencodable value sits under non-string dict keys --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
33457d3ac9
commit
16ee2c71a7
3 changed files with 286 additions and 174 deletions
|
|
@ -61,6 +61,8 @@ import websockets
|
|||
import websockets.exceptions
|
||||
from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config
|
||||
from pydantic.fields import FieldInfo, PydanticUndefined
|
||||
from pydantic.main import IncEx
|
||||
from pydantic_core import PydanticSerializationError, to_jsonable_python
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never
|
||||
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -10186,6 +10188,8 @@ def _event_with_restamped_nested_model(event: Mapping[str, object], target_model
|
|||
_SSE_DATA_FIELD: Final = b"data:"
|
||||
_SSE_NESTED_RESPONSE_MARKER: Final = b'"response"'
|
||||
_SSE_EVENT_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[object, object])
|
||||
_JSON_ARRAY: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def _restamp_nested_response_model_in_sse_frame(frame: bytes, target_model: str) -> bytes:
|
||||
|
|
@ -10211,91 +10215,64 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) -
|
|||
return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending
|
||||
|
||||
|
||||
def _fast_serialize_simple_model_response_stream(
|
||||
chunk: ModelResponseStream,
|
||||
) -> bytes | None:
|
||||
"""
|
||||
Serialize the common OpenAI text streaming chunk without the full Pydantic
|
||||
serializer. Fall back for richer chunks so tool calls, logprobs, usage, and
|
||||
provider-specific fields keep the canonical model_dump_json behavior.
|
||||
"""
|
||||
if (
|
||||
getattr(chunk, "provider_specific_fields", None) is not None
|
||||
or getattr(chunk, "system_fingerprint", None) is not None
|
||||
or getattr(chunk, "usage", None) is not None
|
||||
):
|
||||
return None
|
||||
|
||||
choices: Final = getattr(chunk, "choices", None)
|
||||
if not isinstance(choices, list) or len(choices) != 1:
|
||||
return None
|
||||
|
||||
choice: Final = choices[0]
|
||||
if getattr(choice, "logprobs", None) is not None or getattr(choice, "enhancements", None) is not None:
|
||||
return None
|
||||
|
||||
delta: Final[object] = getattr(choice, "delta", None)
|
||||
if delta is None:
|
||||
return None
|
||||
|
||||
unsupported_delta_fields: Final = (
|
||||
"function_call",
|
||||
"tool_calls",
|
||||
"audio",
|
||||
"images",
|
||||
"annotations",
|
||||
"reasoning_content",
|
||||
"thinking_blocks",
|
||||
"provider_specific_fields",
|
||||
"refusal",
|
||||
)
|
||||
if any(getattr(delta, field, None) is not None for field in unsupported_delta_fields):
|
||||
return None
|
||||
|
||||
delta_dict: Final[dict] = {}
|
||||
role: Final = getattr(delta, "role", None)
|
||||
content: Final = getattr(delta, "content", None)
|
||||
if role is not None:
|
||||
delta_dict["role"] = role
|
||||
if content is not None:
|
||||
delta_dict["content"] = content
|
||||
|
||||
choice_dict: Final = {"index": getattr(choice, "index", 0), "delta": delta_dict}
|
||||
finish_reason: Final = getattr(choice, "finish_reason", None)
|
||||
if finish_reason is not None:
|
||||
choice_dict["finish_reason"] = finish_reason
|
||||
|
||||
# Match the canonical ``model_dump_json(exclude_none=True)`` shape — if a
|
||||
# field is None, omit it entirely rather than emitting ``"key": null``.
|
||||
# Strict OpenAI-compatible clients reject ``null`` for optional fields like
|
||||
# ``model``, so diverging here would surface as a client-side regression
|
||||
# only on the fast path. Fall back to the slow path if a required-looking
|
||||
# top-level identifier is missing.
|
||||
model: Final = getattr(chunk, "model", None)
|
||||
if model is None:
|
||||
return None
|
||||
|
||||
payload: Final[dict] = {
|
||||
"id": getattr(chunk, "id", None),
|
||||
"object": getattr(chunk, "object", None),
|
||||
"created": getattr(chunk, "created", None),
|
||||
"model": model,
|
||||
"service_tier": getattr(chunk, "service_tier", None),
|
||||
"choices": [choice_dict],
|
||||
}
|
||||
for top_level_key in ("id", "object", "created", "service_tier"):
|
||||
if payload[top_level_key] is None:
|
||||
payload.pop(top_level_key)
|
||||
return orjson.dumps(payload)
|
||||
def _json_encodes(value: object) -> bool:
|
||||
try:
|
||||
to_jsonable_python(value)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes:
|
||||
def _unencodable_part(value: object, ancestors: frozenset[int]) -> IncEx | bool:
|
||||
if id(value) in ancestors:
|
||||
return True
|
||||
lineage: Final = ancestors | {id(value)}
|
||||
if isinstance(value, BaseModel):
|
||||
return _unencodable_keys(iter(value), lineage) or True
|
||||
if isinstance(value, dict):
|
||||
return _unencodable_entries(_JSON_OBJECT.validate_python(value), lineage)
|
||||
if isinstance(value, (list, tuple)):
|
||||
return _unencodable_indexes(enumerate(_JSON_ARRAY.validate_python(value)), lineage) or True
|
||||
return True
|
||||
|
||||
|
||||
def _unencodable_entries(entries: Mapping[object, object], lineage: frozenset[int]) -> IncEx | bool:
|
||||
by_name: Final = {key: item for key, item in entries.items() if isinstance(key, str)}
|
||||
if len(by_name) == len(entries):
|
||||
return _unencodable_keys(by_name.items(), lineage) or True
|
||||
by_index: Final = {key: item for key, item in entries.items() if isinstance(key, int)}
|
||||
if len(by_index) == len(entries):
|
||||
return _unencodable_indexes(by_index.items(), lineage) or True
|
||||
return True
|
||||
|
||||
|
||||
def _unencodable_keys(items: Iterable[tuple[str, object]], lineage: frozenset[int]) -> Mapping[str, IncEx | bool]:
|
||||
return {key: _unencodable_part(item, lineage) for key, item in items if not _json_encodes(item)}
|
||||
|
||||
|
||||
def _unencodable_indexes(items: Iterable[tuple[int, object]], lineage: frozenset[int]) -> Mapping[int, IncEx | bool]:
|
||||
return {index: _unencodable_part(item, lineage) for index, item in items if not _json_encodes(item)}
|
||||
|
||||
|
||||
def _serialize_chat_chunk_leniently(chunk: ModelResponseStream) -> bytes:
|
||||
exclude: Final = _unencodable_keys(iter(chunk), frozenset({id(chunk)}))
|
||||
if exclude:
|
||||
verbose_proxy_logger.warning("Dropping streamed chunk fields that JSON cannot encode: %s", exclude)
|
||||
jsonable: Final = chunk.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=exclude)
|
||||
return json.dumps(jsonable, ensure_ascii=True, separators=(",", ":")).encode()
|
||||
|
||||
|
||||
def _serialize_chat_chunk(chunk: ModelResponseStream) -> bytes:
|
||||
try:
|
||||
return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode()
|
||||
except PydanticSerializationError:
|
||||
return _serialize_chat_chunk_leniently(chunk)
|
||||
|
||||
|
||||
def _serialize_streaming_chunk(chunk: BaseModel) -> bytes:
|
||||
if isinstance(chunk, ModelResponseStream):
|
||||
serialized_chunk: Final = _fast_serialize_simple_model_response_stream(chunk)
|
||||
if serialized_chunk is not None:
|
||||
return serialized_chunk
|
||||
|
||||
return chunk.model_dump_json(exclude_none=True, exclude_unset=True)
|
||||
return _serialize_chat_chunk(chunk)
|
||||
return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode()
|
||||
|
||||
|
||||
async def _apply_streaming_chunk_hooks(
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ Pins covered:
|
|||
- ``async_assistants_data_generator``
|
||||
- ``_get_client_requested_model_for_streaming``
|
||||
- ``_restamp_streaming_chunk_model``
|
||||
- ``_fast_serialize_simple_model_response_stream``
|
||||
- ``_serialize_streaming_chunk``
|
||||
- ``_apply_streaming_chunk_hooks``
|
||||
- ``_format_streaming_sse_chunk``
|
||||
|
|
@ -16,8 +15,11 @@ Pins covered:
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -36,7 +38,6 @@ from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
|||
from litellm.proxy.common_request_processing import create_response
|
||||
from litellm.proxy.proxy_server import (
|
||||
_apply_streaming_chunk_hooks,
|
||||
_fast_serialize_simple_model_response_stream,
|
||||
_format_fallback_metadata_sse_event,
|
||||
_format_streaming_sse_chunk,
|
||||
_get_client_requested_model_for_streaming,
|
||||
|
|
@ -55,7 +56,14 @@ from litellm.types.llms.openai import (
|
|||
ResponseFailedEvent,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionDeltaToolCall,
|
||||
Delta,
|
||||
Function,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
Usage,
|
||||
)
|
||||
from tests._support.clock_loops import AutojumpClockLoop
|
||||
|
||||
from .conftest import normalize, route_entry_points_to_stubs
|
||||
|
|
@ -81,6 +89,10 @@ def _simple_chunk(model: str = "gpt-4", content: str = "hi") -> ModelResponseStr
|
|||
)
|
||||
|
||||
|
||||
def _sse_text(frame: str | bytes) -> str:
|
||||
return frame.decode() if isinstance(frame, bytes) else frame
|
||||
|
||||
|
||||
async def _async_iter(items):
|
||||
for it in items:
|
||||
yield it
|
||||
|
|
@ -520,48 +532,13 @@ def test_format_fallback_metadata_sse_event():
|
|||
assert isinstance(payload["created"], int)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _fast_serialize_simple_model_response_stream
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_fast_serialize_simple_model_response_stream_returns_bytes_payload():
|
||||
chunk = _simple_chunk()
|
||||
result = _fast_serialize_simple_model_response_stream(chunk)
|
||||
assert isinstance(result, bytes)
|
||||
payload = json.loads(result)
|
||||
assert normalize(payload) == {
|
||||
"id": "<VOLATILE>",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": "<VOLATILE>",
|
||||
"model": "gpt-4",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "content": "hi"},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def test_fast_serialize_simple_model_response_stream_with_usage_returns_none_invalid():
|
||||
"""Fast path bails (returns None) when ``usage`` is populated — the slow
|
||||
path is required to preserve usage fields. Returning None here is the
|
||||
"I cannot handle this" sentinel, not a hard error."""
|
||||
chunk = _simple_chunk()
|
||||
chunk.usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2)
|
||||
assert _fast_serialize_simple_model_response_stream(chunk) is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _serialize_streaming_chunk
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_serialize_streaming_chunk_simple_uses_fast_path_bytes():
|
||||
result = _serialize_streaming_chunk(_simple_chunk())
|
||||
assert isinstance(result, bytes)
|
||||
payload = json.loads(result)
|
||||
def test_serialize_streaming_chunk_simple_payload():
|
||||
payload = json.loads(_serialize_streaming_chunk(_simple_chunk()))
|
||||
assert normalize(payload) == {
|
||||
"id": "<VOLATILE>",
|
||||
"object": "chat.completion.chunk",
|
||||
|
|
@ -700,10 +677,7 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
|
|||
out.append(line)
|
||||
|
||||
assert out[-1] == "data: [DONE]\n\n"
|
||||
# First chunk is bytes (fast path) wrapped via _format_streaming_sse_chunk.
|
||||
first = out[0]
|
||||
assert isinstance(first, bytes)
|
||||
payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n"))
|
||||
payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
|
||||
assert normalize(payload) == {
|
||||
"id": "<VOLATILE>",
|
||||
"object": "chat.completion.chunk",
|
||||
|
|
@ -718,6 +692,205 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
|
|||
}
|
||||
|
||||
|
||||
_AZURE_CONTENT_FILTER_RESULTS: Final = {
|
||||
"hate": {"filtered": False, "severity": "safe"},
|
||||
"protected_material_text": {"detected": False, "filtered": False},
|
||||
}
|
||||
|
||||
|
||||
def _azure_chunk(choice_fields: dict[str, object], chunk_fields: dict[str, object]) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-azure",
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content="Red"), **choice_fields)],
|
||||
created=0,
|
||||
model="gpt-4",
|
||||
object="chat.completion.chunk",
|
||||
**chunk_fields,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("chunk", "read_field", "sent"),
|
||||
[
|
||||
pytest.param(
|
||||
_azure_chunk({"content_filter_results": _AZURE_CONTENT_FILTER_RESULTS}, {}),
|
||||
lambda payload: payload["choices"][0]["content_filter_results"],
|
||||
_AZURE_CONTENT_FILTER_RESULTS,
|
||||
id="choice-content-filter-results",
|
||||
),
|
||||
pytest.param(
|
||||
_azure_chunk({}, {"obfuscation": "Ab3"}),
|
||||
lambda payload: payload["obfuscation"],
|
||||
"Ab3",
|
||||
id="chunk-obfuscation",
|
||||
),
|
||||
pytest.param(
|
||||
ModelResponseStream(
|
||||
id="chatcmpl-or",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(content="Red", reasoning_details=[{"type": "reasoning.text", "text": "colors"}]),
|
||||
)
|
||||
],
|
||||
created=0,
|
||||
model="gpt-4",
|
||||
object="chat.completion.chunk",
|
||||
),
|
||||
lambda payload: payload["choices"][0]["delta"]["reasoning_details"],
|
||||
[{"type": "reasoning.text", "text": "colors"}],
|
||||
id="delta-reasoning-details",
|
||||
),
|
||||
pytest.param(
|
||||
ModelResponseStream(
|
||||
id="chatcmpl-resp",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(content="Red", reasoning_items=[{"type": "reasoning", "id": "rs_1", "summary": []}]),
|
||||
)
|
||||
],
|
||||
created=0,
|
||||
model="gpt-4",
|
||||
object="chat.completion.chunk",
|
||||
),
|
||||
lambda payload: payload["choices"][0]["delta"]["reasoning_items"],
|
||||
[{"type": "reasoning", "id": "rs_1", "summary": []}],
|
||||
id="delta-reasoning-items",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_async_data_generator_sse_frame_keeps_provider_fields_on_a_text_chunk(
|
||||
monkeypatch, chunk: ModelResponseStream, read_field: Callable[[dict[str, object]], object], sent: object
|
||||
):
|
||||
_patch_logging_flags(monkeypatch)
|
||||
|
||||
out = [
|
||||
line
|
||||
async for line in async_data_generator(
|
||||
response=_async_iter([chunk]),
|
||||
user_api_key_dict=_user_auth(),
|
||||
request_data={"model": "gpt-4"},
|
||||
)
|
||||
]
|
||||
|
||||
assert read_field(json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))) == sent
|
||||
|
||||
|
||||
def _circular_dict() -> dict[str, object]:
|
||||
looped: Final[dict[str, object]] = {"keep": 1} # mutable-ok: a circular dict has to be inserted into itself
|
||||
looped["self"] = looped
|
||||
return looped
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"unencodable",
|
||||
[
|
||||
pytest.param(object(), id="unknown-type"),
|
||||
pytest.param(hashlib.sha256(b"trace").digest(), id="invalid-utf8-bytes"),
|
||||
],
|
||||
)
|
||||
async def test_async_data_generator_streams_a_text_chunk_a_hook_tagged_with_a_value_json_cannot_encode(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, unencodable: object
|
||||
) -> None:
|
||||
_patch_logging_flags(monkeypatch)
|
||||
chunk: Final = _azure_chunk({"content_filter_results": _AZURE_CONTENT_FILTER_RESULTS}, {"obfuscation": "Ab3"})
|
||||
setattr(chunk, "trace_chunk", unencodable)
|
||||
setattr(chunk.choices[0], "trace_choice", unencodable)
|
||||
setattr(chunk.choices[0].delta, "trace_delta", unencodable)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out: Final = [
|
||||
line
|
||||
async for line in async_data_generator(
|
||||
response=_async_iter([chunk]),
|
||||
user_api_key_dict=_user_auth(),
|
||||
request_data={"model": "gpt-4"},
|
||||
)
|
||||
]
|
||||
|
||||
payload: Final = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
|
||||
assert payload["obfuscation"] == "Ab3"
|
||||
assert payload["choices"][0]["content_filter_results"] == _AZURE_CONTENT_FILTER_RESULTS
|
||||
assert payload["choices"][0]["delta"] == {"content": "Red"}
|
||||
assert "trace_chunk" not in payload
|
||||
assert "trace_choice" not in payload["choices"][0]
|
||||
assert all(name in caplog.text for name in ("trace_chunk", "trace_choice", "trace_delta"))
|
||||
|
||||
|
||||
def test_serialize_streaming_chunk_keeps_lone_surrogates_a_provider_split_across_frames(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
tool_call: Final = ChatCompletionDeltaToolCall(
|
||||
id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Par\udc00is"}')
|
||||
)
|
||||
chunk: Final = ModelResponseStream(
|
||||
id="chatcmpl-split",
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content="Hi \ud83d", tool_calls=[tool_call]))],
|
||||
created=0,
|
||||
model="gpt-4",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
delta: Final = json.loads(_serialize_streaming_chunk(chunk))["choices"][0]["delta"]
|
||||
|
||||
assert caplog.text == ""
|
||||
assert delta["content"] == "Hi \ud83d"
|
||||
assert delta["tool_calls"][0]["id"] == "call_1"
|
||||
assert delta["tool_calls"][0]["function"] == {"name": "get_weather", "arguments": '{"city": "Par\udc00is"}'}
|
||||
|
||||
|
||||
def test_serialize_streaming_chunk_drops_only_the_nested_value_json_cannot_encode() -> None:
|
||||
digest: Final = hashlib.sha256(b"trace").digest()
|
||||
chunk: Final = ModelResponseStream(
|
||||
id="chatcmpl-nested",
|
||||
choices=[
|
||||
StreamingChoices(index=0, delta=Delta(content="a")),
|
||||
StreamingChoices(
|
||||
index=1, delta=Delta(content="b", provider_specific_fields={"citations": ["ok"], "raw": digest})
|
||||
),
|
||||
],
|
||||
created=0,
|
||||
model="gpt-4",
|
||||
object="chat.completion.chunk",
|
||||
usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5),
|
||||
)
|
||||
setattr(chunk.usage, "raw", digest)
|
||||
setattr(chunk, "trace_chunk", _circular_dict())
|
||||
|
||||
payload: Final = json.loads(_serialize_streaming_chunk(chunk))
|
||||
|
||||
assert payload["choices"][0]["delta"] == {"content": "a"}
|
||||
assert payload["choices"][1]["delta"] == {"content": "b", "provider_specific_fields": {"citations": ["ok"]}}
|
||||
assert payload["usage"] == {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}
|
||||
assert payload["trace_chunk"] == {"keep": 1}
|
||||
|
||||
|
||||
def test_serialize_streaming_chunk_keeps_streaming_when_unencodable_values_sit_under_non_string_keys() -> None:
|
||||
chunk: Final = ModelResponseStream(
|
||||
id="chatcmpl-keys",
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content="still streaming"))],
|
||||
created=0,
|
||||
model="gpt-4",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
setattr(chunk, "by_index", {0: "kept", 1: object()})
|
||||
setattr(chunk, "mixed_keys", {"name": "x", 2: object()})
|
||||
setattr(chunk, "bytes_keys", {b"raw": object()})
|
||||
setattr(chunk, "read_only_view", MappingProxyType({"name": object()}))
|
||||
|
||||
payload: Final = json.loads(_serialize_streaming_chunk(chunk))
|
||||
|
||||
assert payload["choices"][0]["delta"] == {"content": "still streaming"}
|
||||
assert payload["by_index"] == {"0": "kept"}
|
||||
assert "mixed_keys" not in payload
|
||||
assert "bytes_keys" not in payload
|
||||
assert "read_only_view" not in payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch):
|
||||
_patch_logging_flags(monkeypatch)
|
||||
|
|
@ -739,9 +912,7 @@ async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch)
|
|||
):
|
||||
out.append(line)
|
||||
|
||||
first = out[0]
|
||||
assert isinstance(first, bytes)
|
||||
payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n"))
|
||||
payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
|
||||
assert payload["model"] == "fallback-model"
|
||||
|
||||
|
||||
|
|
@ -764,9 +935,7 @@ async def test_async_data_generator_uses_chunk_fallback_metadata(monkeypatch):
|
|||
):
|
||||
out.append(line)
|
||||
|
||||
first = out[0]
|
||||
assert isinstance(first, bytes)
|
||||
payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n"))
|
||||
payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
|
||||
assert payload["model"] == "fallback-model"
|
||||
|
||||
|
||||
|
|
@ -793,8 +962,8 @@ async def test_async_data_generator_switches_model_mid_stream_on_fallback(monkey
|
|||
):
|
||||
out.append(line)
|
||||
|
||||
first_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n"))
|
||||
second_payload = json.loads(out[1].removeprefix(b"data: ").removesuffix(b"\n\n"))
|
||||
first_payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
|
||||
second_payload = json.loads(_sse_text(out[1]).removeprefix("data: ").removesuffix("\n\n"))
|
||||
assert first_payload["model"] == "primary-model"
|
||||
assert second_payload["model"] == "fallback-model"
|
||||
|
||||
|
|
@ -830,8 +999,7 @@ async def test_async_data_generator_emits_fallback_error_metadata_event(monkeypa
|
|||
):
|
||||
out.append(line)
|
||||
|
||||
assert isinstance(out[0], bytes)
|
||||
chunk_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n"))
|
||||
chunk_payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
|
||||
assert chunk_payload["model"] == "fallback-model"
|
||||
assert isinstance(out[1], str)
|
||||
assert out[1].startswith("data: ")
|
||||
|
|
@ -879,8 +1047,7 @@ async def test_async_data_generator_skips_fallback_error_event_without_opt_in(
|
|||
):
|
||||
out.append(line)
|
||||
|
||||
assert isinstance(out[0], bytes)
|
||||
payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n"))
|
||||
payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))
|
||||
assert payload["model"] == "fallback-model"
|
||||
|
||||
|
||||
|
|
@ -2434,14 +2601,11 @@ async def test_queue_request_stream_is_untouched_while_keepalives_are_unconfigur
|
|||
assert chunks[-1] == b"data: [DONE]\n\n"
|
||||
|
||||
|
||||
def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier():
|
||||
def test_serialize_streaming_chunk_keeps_served_service_tier():
|
||||
chunk = _simple_chunk()
|
||||
chunk.service_tier = "priority"
|
||||
|
||||
result = _fast_serialize_simple_model_response_stream(chunk)
|
||||
|
||||
assert result is not None
|
||||
assert json.loads(result)["service_tier"] == "priority"
|
||||
assert json.loads(_serialize_streaming_chunk(chunk))["service_tier"] == "priority"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -98,42 +98,13 @@ def test_restamp_streaming_chunk_skips_matching_model():
|
|||
assert model_mismatch_logged is False
|
||||
|
||||
|
||||
def test_fast_serialize_simple_streaming_chunk_matches_model_dump_json():
|
||||
def test_serialize_streaming_chunk_omits_a_missing_model_instead_of_null():
|
||||
from litellm.proxy.proxy_server import _serialize_streaming_chunk
|
||||
|
||||
chunk = _make_model_response_stream_chunk("client-model")
|
||||
|
||||
assert json.loads(_serialize_streaming_chunk(chunk)) == json.loads(
|
||||
chunk.model_dump_json(exclude_none=True, exclude_unset=True)
|
||||
)
|
||||
|
||||
|
||||
def test_fast_serialize_returns_none_when_model_field_is_missing():
|
||||
"""
|
||||
The fast path must mirror ``model_dump_json(exclude_none=True)``: when
|
||||
``chunk.model`` is ``None`` the slow path omits the field entirely.
|
||||
Emitting ``"model": null`` would diverge and trip strict OpenAI-
|
||||
compatible clients that reject ``null`` for optional string fields.
|
||||
Falling back to ``None`` lets the canonical serializer handle the edge.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
_fast_serialize_simple_model_response_stream,
|
||||
_serialize_streaming_chunk,
|
||||
)
|
||||
|
||||
chunk = _make_model_response_stream_chunk("client-model")
|
||||
chunk.model = None # type: ignore[assignment]
|
||||
|
||||
assert _fast_serialize_simple_model_response_stream(chunk) is None
|
||||
|
||||
# Going through the public ``_serialize_streaming_chunk`` should still
|
||||
# produce a serialized result via the slow-path fallback, and it must
|
||||
# not contain ``"model": null``.
|
||||
serialized = _serialize_streaming_chunk(chunk)
|
||||
payload_str = serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized
|
||||
assert '"model": null' not in payload_str
|
||||
assert '"model":null' not in payload_str
|
||||
assert json.loads(payload_str) == json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True))
|
||||
assert "model" not in json.loads(_serialize_streaming_chunk(chunk))
|
||||
|
||||
|
||||
def test_proxy_chat_completion_does_not_return_provider_prefixed_model(tmp_path, monkeypatch):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue