diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 091e05610ba..635b0c2bacd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10130,10 +10130,11 @@ def _fast_serialize_simple_model_response_stream( 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 + choice_dict: Final = { + "index": getattr(choice, "index", 0), + "delta": delta_dict, + "finish_reason": getattr(choice, "finish_reason", None), + } # Match the canonical ``model_dump_json(exclude_none=True)`` shape — if a # field is None, omit it entirely rather than emitting ``"key": null``. @@ -10159,13 +10160,23 @@ def _fast_serialize_simple_model_response_stream( return orjson.dumps(payload) -def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes: - if isinstance(chunk, ModelResponseStream): - serialized_chunk: Final = _fast_serialize_simple_model_response_stream(chunk) - if serialized_chunk is not None: - return serialized_chunk +_STREAMING_CHOICE_PAYLOADS: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) - return chunk.model_dump_json(exclude_none=True, exclude_unset=True) + +def _serialize_model_response_stream(chunk: ModelResponseStream) -> bytes: + payload: Final[Mapping[str, object]] = BaseModel.model_dump( + chunk, mode="json", exclude_none=True, exclude_unset=True + ) + choices: Final = _STREAMING_CHOICE_PAYLOADS.validate_python(payload["choices"]) + return _SSE_EVENT_OBJECT.dump_json( + {**payload, "choices": tuple({**choice, "finish_reason": choice.get("finish_reason")} for choice in choices)} + ) + + +def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes: + if not isinstance(chunk, ModelResponseStream): + return chunk.model_dump_json(exclude_none=True, exclude_unset=True) + return _fast_serialize_simple_model_response_stream(chunk) or _serialize_model_response_stream(chunk) def _is_injected_stream_usage_artifact(chunk: object) -> bool: diff --git a/tests/integration/streaming/test_openai_sdk_plain_chunk_fields.py b/tests/integration/streaming/test_openai_sdk_plain_chunk_fields.py index 87dd7a16a41..19c13a58f5a 100644 --- a/tests/integration/streaming/test_openai_sdk_plain_chunk_fields.py +++ b/tests/integration/streaming/test_openai_sdk_plain_chunk_fields.py @@ -120,11 +120,12 @@ def test_plain_chunks_keep_text_role_and_provider_fields(gateway: Gateway, provi 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") + finish_reasons: Final = 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",) + if event["choices"] + ) + assert finish_reasons == (*(None,) * (len(finish_reasons) - 1), "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), diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 9a3516d20d6..73435887282 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -539,6 +539,7 @@ def test_fast_serialize_simple_model_response_stream_returns_bytes_payload(): { "index": 0, "delta": {"role": "assistant", "content": "hi"}, + "finish_reason": None, } ], } @@ -571,11 +572,61 @@ def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): { "index": 0, "delta": {"role": "assistant", "content": "hi"}, + "finish_reason": None, } ], } +def _stream_chunk(delta: Delta, finish_reason: str | None = None, usage: Usage | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-test", + choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=delta)], + created=0, + model="gpt-4", + object="chat.completion.chunk", + usage=usage, + ) + + +@pytest.mark.parametrize( + "delta", + [ + pytest.param(Delta(content="hi", role="assistant"), id="text"), + pytest.param( + Delta( + tool_calls=[ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": ""}, "index": 0} + ] + ), + id="tool-call", + ), + pytest.param(Delta(reasoning_content="plan"), id="reasoning"), + ], +) +def test_serialize_streaming_chunk_sends_a_null_finish_reason_on_an_in_progress_chunk(delta: Delta) -> None: + payload: Final = json.loads(_serialize_streaming_chunk(_stream_chunk(delta))) + choice: Final = payload["choices"][0] + + assert "finish_reason" in choice + assert choice["finish_reason"] is None + assert "logprobs" not in choice + assert "system_fingerprint" not in payload + + +@pytest.mark.parametrize( + "usage", + [ + pytest.param(None, id="text"), + pytest.param(Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5), id="with-usage"), + ], +) +def test_serialize_streaming_chunk_keeps_the_finish_reason_of_the_final_chunk(usage: Usage | None) -> None: + payload: Final = json.loads(_serialize_streaming_chunk(_stream_chunk(Delta(), finish_reason="stop", usage=usage))) + + assert payload["choices"][0]["finish_reason"] == "stop" + + def test_serialize_streaming_chunk_invalid_input_raises_attribute_error(): """The helper is typed as ``BaseModel`` — handing it a plain dict trips the attribute-access path (no ``model_dump_json``).""" @@ -713,6 +764,7 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): { "index": 0, "delta": {"role": "assistant", "content": "hello"}, + "finish_reason": None, } ], } diff --git a/tests/unit/proxy/test_response_model_sanitization.py b/tests/unit/proxy/test_response_model_sanitization.py index f52736453c1..76cb619b3ae 100644 --- a/tests/unit/proxy/test_response_model_sanitization.py +++ b/tests/unit/proxy/test_response_model_sanitization.py @@ -64,6 +64,11 @@ def _make_model_response_stream_chunk(model: str) -> litellm.ModelResponseStream return litellm.ModelResponseStream(**chunk_dict) +def _canonical_dump_with_null_finish_reason(chunk: litellm.ModelResponseStream) -> dict[str, object]: + canonical = json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True)) + return {**canonical, "choices": [{**choice, "finish_reason": None} for choice in canonical["choices"]]} + + def _decode_sse_chunk(chunk) -> str: return chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk @@ -98,14 +103,12 @@ 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_fast_serialize_simple_streaming_chunk_matches_model_dump_json_with_a_null_finish_reason(): 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) - ) + assert json.loads(_serialize_streaming_chunk(chunk)) == _canonical_dump_with_null_finish_reason(chunk) def test_fast_serialize_returns_none_when_model_field_is_missing(): @@ -133,7 +136,7 @@ def test_fast_serialize_returns_none_when_model_field_is_missing(): 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 json.loads(payload_str) == _canonical_dump_with_null_finish_reason(chunk) def test_proxy_chat_completion_does_not_return_provider_prefixed_model(tmp_path, monkeypatch):