From f7228a46702b19e384626b25a0fa84d6d7cb5ed9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 21:47:15 -0700 Subject: [PATCH 1/2] fix(streaming): preserve parsed-chunk provider_specific_fields so Vertex flex streams bill at flex rates --- .../litellm_core_utils/streaming_handler.py | 24 +++++++- .../test_streaming_handler.py | 59 +++++++++++++++++++ 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f6340426c1b..0f46f1b718c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -8,11 +8,12 @@ import time import traceback from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import Any, Final, NoReturn, Protocol, TypeVar, cast import anyio import httpx -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from typing_extensions import NotRequired, TypedDict import litellm @@ -182,6 +183,23 @@ class _VertexChunkLike(Protocol): candidates: Sequence[_VertexCandidateLike] +class _ParsedChunkHiddenParams(BaseModel): + provider_specific_fields: Mapping[str, object] | None = None + + +def _provider_hidden_params(chunk: object) -> Mapping[str, object] | None: + hidden: Final[object] = getattr(chunk, "_hidden_params", None) + if not isinstance(hidden, dict): + return None + try: + parsed: Final = _ParsedChunkHiddenParams.model_validate(hidden) + except ValidationError: + return None + if not parsed.provider_specific_fields: + return None + return MappingProxyType({"provider_specific_fields": dict(parsed.provider_specific_fields)}) + + class CustomStreamWrapper: def __init__( self, @@ -801,7 +819,7 @@ class CustomStreamWrapper: except Exception as e: raise e - def model_response_creator(self, chunk: dict | None = None, hidden_params: dict | None = None): + def model_response_creator(self, chunk: dict | None = None, hidden_params: Mapping[str, object] | None = None): _model: Final = self._cached_model_name _logging_obj_llm_provider: Final = self._cached_logging_llm_provider @@ -1504,7 +1522,7 @@ class CustomStreamWrapper: def chunk_creator(self, chunk: Any): if hasattr(chunk, "id"): self.response_id = chunk.id - model_response = self.model_response_creator() + model_response = self.model_response_creator(hidden_params=_provider_hidden_params(chunk)) response_obj: dict[str, Any] = {} try: # return this for all models diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index b5e33a4e421..d39c84ce5e7 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4460,3 +4460,62 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp finally: trace_id_var.set("") session_id_var.set("") + + +def test_chunk_creator_preserves_hidden_provider_specific_fields_from_parsed_chunk(): + """ + Vertex/Gemini chunk_parser stores usageMetadata.trafficType in the parsed + chunk's _hidden_params["provider_specific_fields"], but chunk_creator builds + a fresh ModelResponseStream per outgoing chunk. Before the fix those hidden + provider fields were dropped, so streaming flex traffic was billed at + standard rates (LIT-6292). + """ + wrapper = CustomStreamWrapper( + completion_stream=None, + model="gemini-3.5-flash", + logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + ) + parsed_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="hello", role="assistant"), finish_reason=None)], + ) + parsed_chunk._hidden_params["provider_specific_fields"] = {"traffic_type": "ON_DEMAND_FLEX"} + + result = wrapper.chunk_creator(chunk=parsed_chunk) + + assert result is not None + assert result._hidden_params["provider_specific_fields"] == {"traffic_type": "ON_DEMAND_FLEX"} + + +@pytest.mark.asyncio +async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging_obj: Logging): + """ + End-to-end through the async wrapper: the assembled response handed to cost + tracking must carry traffic_type so flex/priority tiers price correctly. + """ + content_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="hello", role="assistant"), finish_reason=None)], + ) + final_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + setattr(final_chunk, "usage", Usage(prompt_tokens=7, completion_tokens=5, total_tokens=12)) + final_chunk._hidden_params["provider_specific_fields"] = {"traffic_type": "ON_DEMAND_FLEX"} + + async def _stream(): + yield content_chunk + yield final_chunk + + wrapper = CustomStreamWrapper( + completion_stream=_stream(), + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + stream_options={"include_usage": True}, + ) + + received = [chunk async for chunk in wrapper] + + assembled = litellm.stream_chunk_builder(chunks=received, messages=[{"role": "user", "content": "hi"}]) + assert assembled is not None + assert assembled._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND_FLEX" From ab71807985237352965250bc3c3e9be73ef803da Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 22:23:17 -0700 Subject: [PATCH 2/2] test(streaming): drop redundant docstrings from flex-tier regression tests --- .../litellm_core_utils/test_streaming_handler.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index d39c84ce5e7..5329edce47e 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4463,13 +4463,6 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp def test_chunk_creator_preserves_hidden_provider_specific_fields_from_parsed_chunk(): - """ - Vertex/Gemini chunk_parser stores usageMetadata.trafficType in the parsed - chunk's _hidden_params["provider_specific_fields"], but chunk_creator builds - a fresh ModelResponseStream per outgoing chunk. Before the fix those hidden - provider fields were dropped, so streaming flex traffic was billed at - standard rates (LIT-6292). - """ wrapper = CustomStreamWrapper( completion_stream=None, model="gemini-3.5-flash", @@ -4489,10 +4482,6 @@ def test_chunk_creator_preserves_hidden_provider_specific_fields_from_parsed_chu @pytest.mark.asyncio async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging_obj: Logging): - """ - End-to-end through the async wrapper: the assembled response handed to cost - tracking must carry traffic_type so flex/priority tiers price correctly. - """ content_chunk = ModelResponseStream( choices=[StreamingChoices(index=0, delta=Delta(content="hello", role="assistant"), finish_reason=None)], )