mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
fix(proxy): keep extra chunk fields on the streaming fast path
_fast_serialize_simple_model_response_stream writes a fixed set of fields, so a field set on the chunk outside its schema (model_extra), for example by an async_post_call_streaming_iterator_hook, was silently dropped from plain text chunks. Fall back to the canonical serializer when the chunk carries an extra the fast payload does not write; service_tier is still served on the fast path.
This commit is contained in:
parent
10444df3a0
commit
cccfe196e5
2 changed files with 23 additions and 0 deletions
|
|
@ -9565,6 +9565,14 @@ def _fast_serialize_simple_model_response_stream(
|
|||
):
|
||||
return None
|
||||
|
||||
# Fields set on the chunk outside its schema (e.g. by a streaming iterator
|
||||
# hook) live in ``model_extra``. The payload below only writes
|
||||
# ``service_tier``, so any other extra would be silently dropped here
|
||||
# while the canonical serializer keeps it.
|
||||
model_extra: Final = getattr(chunk, "model_extra", None)
|
||||
if model_extra and any(key != "service_tier" for key in model_extra):
|
||||
return None
|
||||
|
||||
choices: Final = getattr(chunk, "choices", None)
|
||||
if not isinstance(choices, list) or len(choices) != 1:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -550,6 +550,21 @@ def test_fast_serialize_simple_model_response_stream_with_usage_returns_none_inv
|
|||
assert _fast_serialize_simple_model_response_stream(chunk) is None
|
||||
|
||||
|
||||
def test_fast_serialize_simple_model_response_stream_with_extra_field_returns_none():
|
||||
"""A field a callback attaches to the chunk is not in the fast payload, so
|
||||
the fast path must bail rather than drop it."""
|
||||
chunk = _simple_chunk()
|
||||
chunk.gateway_metadata = {"served_by": "fallback-model"}
|
||||
assert _fast_serialize_simple_model_response_stream(chunk) is None
|
||||
|
||||
|
||||
def test_serialize_streaming_chunk_keeps_extra_field():
|
||||
chunk = _simple_chunk()
|
||||
chunk.gateway_metadata = {"served_by": "fallback-model"}
|
||||
payload = json.loads(_serialize_streaming_chunk(chunk))
|
||||
assert payload["gateway_metadata"] == {"served_by": "fallback-model"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _serialize_streaming_chunk
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue