Merge pull request #38458 from BerriAI/litellm_streaming_flex_service_tier

fix(streaming): preserve provider service-tier metadata so Vertex flex streams bill at flex rates

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

tests/test_litellm/litellm_core_utils/test_streaming_handler.py

@@ -4460,3 +4460,51 @@ 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():
+    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):
+    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"