From a4983736bf5efd7d9c1e6c80ace2b2c262541b27 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:12:46 -0700 Subject: [PATCH] fix(streaming): keep the served service tier on the finish chunk the stream adds --- .../litellm_core_utils/streaming_handler.py | 10 ++++++ .../test_streaming_handler.py | 32 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b64e146ca08..bbe5b661d65 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1915,6 +1915,16 @@ class CustomStreamWrapper: model_response: Final = self.model_response_creator( hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id} ) + served_tier: Final = next( + ( + tier + for chunk in reversed(self.chunks) + if isinstance(tier := getattr(chunk, "service_tier", None), str) and tier + ), + None, + ) + if served_tier is not None: + setattr(model_response, "service_tier", served_tier) # noqa: B010 # pydantic extra, not a declared field _finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason if _finish_reason is not None: model_response.choices[0].finish_reason = _finish_reason diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6d700bae514..91a57ceb44e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -7547,6 +7547,38 @@ async def test_openai_stream_relays_the_served_service_tier_on_every_chunk_inclu assert relayed[-1]["usage"]["total_tokens"] == 11 +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_relays_the_served_service_tier_on_the_finish_chunk_it_adds(logging_obj: Logging, sync_mode: bool): + from litellm.utils import ModelResponseListIterator + + def _chunk(content: str, finish_reason: str | None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-tier", + created=1742056047, + model="gpt-4.1-mini", + choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content=content))], + service_tier="priority", + ) + + logging_obj.update_environment_variables( + model="gpt-4.1-mini", optional_params={}, litellm_params={}, custom_llm_provider="openai" + ) + wrapper = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=[_chunk("Hi", None), _chunk("!", "stop")]), + model="gpt-4.1-mini", + custom_llm_provider="openai", + logging_obj=logging_obj, + ) + + relayed = ( + [chunk.model_dump() for chunk in wrapper] if sync_mode else [chunk.model_dump() async for chunk in wrapper] + ) + + assert [chunk["choices"][0]["finish_reason"] for chunk in relayed] == [None, None, "stop"], relayed + assert [chunk.get("service_tier") for chunk in relayed] == ["priority"] * len(relayed), relayed + + def _last_chunk_carries_finish_reason_wrapper( logging_obj: Logging, finish_reason: str, sync_stream: bool ) -> CustomStreamWrapper: