fix(streaming): keep the served service_tier on streamed chunks and spend rows

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -1334,6 +1334,7 @@ def __init__(
     ):
         super().__init__(streaming_response, sync_stream, json_mode)
         self._chat_completion_id: str | None = None
+        self._served_service_tier: str | None = None
         self._tool_call_index_map: dict[int, int] = {}  # mutable-ok: per-stream accumulator state
 
     def _handle_string_chunk(
@@ -1598,6 +1599,7 @@ def translate_responses_chunk_to_openai_stream(
 
                 usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_data.get("usage"))
             provider_metadata: Final = _provider_metadata(response_data)
+            served_service_tier: Final = response_data.get("service_tier")
             return ModelResponseStream(
                 choices=[
                     StreamingChoices(
@@ -1611,6 +1613,11 @@ def translate_responses_chunk_to_openai_stream(
                 ],
                 usage=usage,
                 provider_specific_fields=dict(provider_metadata) or None,  # mutable-ok: field is typed dict
+                **(
+                    MappingProxyType({"service_tier": served_service_tier})
+                    if isinstance(served_service_tier, str)
+                    else MappingProxyType({})
+                ),
             )
         else:
             pass
@@ -1639,12 +1646,28 @@ def chunk_parser(self, chunk: dict[str, object]) -> "ModelResponseStream":
             ModelResponseStream: OpenAI-formatted streaming chunk
         """
         verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk)
-        return self._with_stream_scoped_id(
-            OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
-                chunk, tool_call_index_map=self._tool_call_index_map
+        self._remember_served_service_tier(chunk)
+        return self._with_served_service_tier(
+            self._with_stream_scoped_id(
+                OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
+                    chunk, tool_call_index_map=self._tool_call_index_map
+                )
             )
         )
 
+    def _remember_served_service_tier(self, chunk: dict[str, object]) -> None:
+        response_payload: Final = chunk.get("response")
+        if not isinstance(response_payload, dict):
+            return
+        served_tier: Final = response_payload.get("service_tier")
+        if isinstance(served_tier, str) and served_tier:
+            self._served_service_tier = served_tier
+
+    def _with_served_service_tier(self, chunk: "ModelResponseStream") -> "ModelResponseStream":
+        if self._served_service_tier is not None and chunk.model_dump().get("service_tier") is None:
+            setattr(chunk, "service_tier", self._served_service_tier)  # noqa: B010  # pydantic extra, not a declared field
+        return chunk
+
     def _with_stream_scoped_id(self, chunk: "ModelResponseStream") -> "ModelResponseStream":
         if self._chat_completion_id is None:
             self._chat_completion_id = chunk.id

litellm/cost_calculator.py

@@ -26,6 +26,7 @@
     TranscriptionUsageObjectTransformation,
 )
 from litellm.litellm_core_utils.llm_cost_calc.utils import (
+    _SERVICE_TIER_TO_COST_KEY_SUFFIX,
     BilledTokenRates,
     CostCalculatorUtils,
     _generic_cost_per_character,
@@ -696,7 +697,7 @@ def cost_per_token(
             data_residency=data_residency,
         )
     elif custom_llm_provider == "databricks":
-        return databricks_cost_per_token(model=model, usage=usage_block)
+        return databricks_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
     elif custom_llm_provider == "fireworks_ai":
         return fireworks_ai_cost_per_token(model=model, usage=usage_block)
     elif custom_llm_provider == "azure":
@@ -959,6 +960,37 @@ def _normalize_service_tier(service_tier: object) -> str | None:
     return service_tier
 
 
+_BASE_PRICING_SERVICE_TIERS: Final[frozenset[str]] = frozenset({"default", "standard"})
+
+
+def _resolve_billable_service_tier(requested: object, served: object) -> str | None:
+    """Served tier wins when it names a priced tier or explicitly says base pricing; otherwise the request decides."""
+    served_lower: Final = served.lower() if isinstance(served, str) else None
+    if served_lower is not None and served_lower in _SERVICE_TIER_TO_COST_KEY_SUFFIX:
+        return served_lower
+    if served_lower in _BASE_PRICING_SERVICE_TIERS:
+        return None
+    return _normalize_service_tier(requested)
+
+
+def _served_service_tier(completion_response: object, usage_object: Usage | None) -> str | None:
+    """Find the tier the provider actually served: response, then usage, then Gemini trafficType."""
+    response_tier: Final = _extract_service_tier(completion_response)
+    if isinstance(response_tier, str):
+        return response_tier
+    usage_tier: Final = _extract_service_tier(usage_object)
+    if isinstance(usage_tier, str):
+        return usage_tier
+    hidden_params: Final = getattr(completion_response, "_hidden_params", None)
+    if hidden_params is None:
+        return None
+    provider_specific: Final = hidden_params.get("provider_specific_fields") or {}
+    raw_traffic_type: Final = provider_specific.get("traffic_type")
+    if not raw_traffic_type:
+        return None
+    return _map_traffic_type_to_service_tier(raw_traffic_type) or "default"
+
+
 def _extract_service_tier(source: object) -> str | None:
     """Read a raw ``service_tier`` off a response body or usage object, dict or pydantic model alike."""
     if isinstance(source, BaseModel):
@@ -1378,23 +1410,14 @@ def completion_cost(
         )
         rerank_billed_units: RerankBilledUnits | None = None
 
-        # Extract service_tier from optional_params if not provided directly
-        if service_tier is None and optional_params is not None:
-            service_tier = optional_params.get("service_tier")
-
-        service_tier = _normalize_service_tier(service_tier)
-
-        # Extract service_tier from completion_response if not provided
-        if service_tier is None and completion_response is not None:
-            service_tier = _extract_service_tier(completion_response)
-
-        service_tier = _normalize_service_tier(service_tier)
-
-        # Extract service_tier from usage object if not provided
-        if service_tier is None and cost_per_token_usage_object is not None:
-            service_tier = _extract_service_tier(cost_per_token_usage_object)
-
-        service_tier = _normalize_service_tier(service_tier)
+        explicit_tier: Final = _normalize_service_tier(service_tier)
+        if explicit_tier is not None:
+            service_tier = explicit_tier
+        else:
+            service_tier = _resolve_billable_service_tier(  # rebind-ok: resolved from request then response
+                requested=optional_params.get("service_tier") if optional_params is not None else None,
+                served=_served_service_tier(completion_response, cost_per_token_usage_object),
+            )
 
         explicit_pricing: Final = custom_pricing is True or base_model is not None
         selected_model: Final = _select_model_name_for_cost_calc(
@@ -1484,15 +1507,6 @@ def completion_cost(
                         custom_llm_provider = hidden_params.get("custom_llm_provider", custom_llm_provider or None)
                         region_name = hidden_params.get("region_name", region_name)
 
-                        # For Gemini/Vertex AI responses, trafficType is stored in
-                        # provider_specific_fields.  Map it to the service_tier used
-                        # by the cost key lookup (_priority / _flex suffixes) so that
-                        # ON_DEMAND_PRIORITY requests are billed at priority prices.
-                        if service_tier is None:
-                            provider_specific = hidden_params.get("provider_specific_fields") or {}
-                            raw_traffic_type = provider_specific.get("traffic_type")
-                            if raw_traffic_type:
-                                service_tier = _map_traffic_type_to_service_tier(raw_traffic_type)
                 else:
                     if model is None:
                         raise ValueError(

litellm/litellm_core_utils/litellm_logging.py

@@ -1883,7 +1883,6 @@ def _response_cost_calculator(
                 "standard_built_in_tools_params": self.standard_built_in_tools_params,
                 "router_model_id": router_model_id,
                 "litellm_logging_obj": self,
-                "service_tier": (self.optional_params.get("service_tier") if self.optional_params else None),
                 "data_residency": (
                     self.litellm_params.get("data_residency")
                     if hasattr(self, "litellm_params") and self.litellm_params

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -109,6 +109,7 @@ class _BaseChunk(TypedDict, total=False):
     created: ReadOnly[int]
     model: ReadOnly[str]
     system_fingerprint: ReadOnly[str | None]
+    service_tier: ReadOnly[str | None]
     choices: ReadOnly[Required[Sequence[StreamingChoices]]]
     _hidden_params: ReadOnly[_ChunkHiddenParams]
 
@@ -369,6 +370,13 @@ def _get_model_from_chunks(chunks: Sequence["_BaseChunk"], first_chunk_model: st
         # Fall back to first chunk's model if no different model found
         return first_chunk_model
 
+    @staticmethod
+    def _get_service_tier_from_chunks(chunks: Sequence["_BaseChunk"]) -> str | None:
+        return next(
+            (tier for chunk in reversed(chunks) if isinstance(tier := chunk.get("service_tier"), str) and tier),
+            None,
+        )
+
     def build_base_response(self, chunks: Sequence["_BaseChunk"]) -> ModelResponse:
         chunk = self.first_chunk
         id: Final = ChunkProcessor._get_chunk_id(chunks)
@@ -378,6 +386,7 @@ def build_base_response(self, chunks: Sequence["_BaseChunk"]) -> ModelResponse:
         # Get the actual model - for Azure Model Router, this finds the real model from later chunks
         model: Final = ChunkProcessor._get_model_from_chunks(chunks, first_chunk_model)
         system_fingerprint: Final = chunk.get("system_fingerprint", None)
+        service_tier: Final = ChunkProcessor._get_service_tier_from_chunks(chunks)
 
         role: Final = ChunkProcessor._get_role_from_chunks(chunks)
         finish_reason = "stop"
@@ -399,6 +408,11 @@ def build_base_response(self, chunks: Sequence["_BaseChunk"]) -> ModelResponse:
                 "created": created,
                 "model": model,
                 "system_fingerprint": system_fingerprint,
+                **(
+                    MappingProxyType({"service_tier": service_tier})
+                    if service_tier is not None
+                    else MappingProxyType({})
+                ),
                 "choices": [
                     {
                         "index": 0,

litellm/litellm_core_utils/streaming_handler.py

@@ -72,6 +72,12 @@ def _next_sync_or_exhausted(it: Any) -> object:
         return _SYNC_ITER_EXHAUSTED
 
 
+def _stamp_served_service_tier(response: ModelResponseStream, complete_streaming_response: ModelResponse) -> None:
+    served_tier: Final = complete_streaming_response.model_dump().get("service_tier")
+    if isinstance(served_tier, str) and served_tier:
+        setattr(response, "service_tier", served_tier)  # noqa: B010  # pydantic extra, not a declared field
+
+
 def is_async_iterable(obj: object) -> bool:
     """
     Check if an object is an async iterable (can be used with 'async for').
@@ -1873,6 +1879,7 @@ def __next__(self) -> "ModelResponseStream":
                         "usage",
                         getattr(complete_streaming_response, "usage"),
                     )
+                    _stamp_served_service_tier(response, complete_streaming_response)
                     try:
                         _cache_copy = complete_streaming_response.model_copy(deep=True)
                         _log_copy = complete_streaming_response.model_copy(deep=True)
@@ -2124,6 +2131,7 @@ async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStr
                     "usage",
                     getattr(complete_streaming_response, "usage"),
                 )
+                _stamp_served_service_tier(response, complete_streaming_response)
                 try:
                     _copy = complete_streaming_response.model_copy(deep=True)
                 except RuntimeError:

litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py

@@ -11,6 +11,7 @@
     Final,
     Literal,
     Protocol,
+    cast,
     get_args,
 )
 
@@ -35,6 +36,7 @@
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
+    from litellm.types.llms.openai import AllMessageValues
     from litellm.types.utils import ModelResponseStream
 
 
@@ -115,6 +117,18 @@ def __init__(self, completion_stream: _ChunkStream):
         self._async_iter: AsyncIterator[ModelResponseStream] | None = None
         self._buffer: deque[ModelResponseStream] = deque()
 
+    @property
+    def chunks(self) -> "list[ModelResponseStream] | None":
+        return cast(  # cast-ok: chunks is a list of ModelResponseStream on the inner stream
+            "list[ModelResponseStream] | None", getattr(self._stream, "chunks", None)
+        )
+
+    @property
+    def messages(self) -> "list[AllMessageValues] | None":
+        return cast(  # cast-ok: messages is a list of AllMessageValues on the inner stream
+            "list[AllMessageValues] | None", getattr(self._stream, "messages", None)
+        )
+
     @staticmethod
     def _is_combined(chunk: "ModelResponseStream") -> bool:
         """True if ``chunk`` carries response content AND a finish_reason."""
@@ -351,6 +365,18 @@ def __init__(
             text="",
         )
 
+    @property
+    def chunks(self) -> "list[ModelResponseStream] | None":
+        return cast(  # cast-ok: chunks is a list of ModelResponseStream on the inner stream
+            "list[ModelResponseStream] | None", getattr(self.completion_stream, "chunks", None)
+        )
+
+    @property
+    def messages(self) -> "list[AllMessageValues] | None":
+        return cast(  # cast-ok: messages is a list of AllMessageValues on the inner stream
+            "list[AllMessageValues] | None", getattr(self.completion_stream, "messages", None)
+        )
+
     def _merge_usage_into_held_stop_reason_chunk(self, chunk: Any) -> MessageBlockDelta:
         """Merge usage data from ``chunk`` into the held ``message_delta`` chunk.
 
@@ -1173,3 +1199,37 @@ def _should_start_new_content_block(self, chunk: "ModelResponseStream") -> bool:
                 return True
 
         return False
+
+
+class AnthropicSSEStream(AsyncIterator[bytes]):
+    """
+    AsyncIterator[bytes] view of AnthropicStreamWrapper returned to callers of
+    translate_completion_output_params_streaming. Keeps the wrapper reachable so
+    the proxy's disconnect-time partial billing can read the inner chat stream's
+    collected chunks, messages, and model; a bare async generator would hide them.
+    """
+
+    def __init__(self, anthropic_wrapper: AnthropicStreamWrapper) -> None:
+        self._anthropic_wrapper = anthropic_wrapper
+        self._byte_stream: Final[AsyncIterator[bytes]] = anthropic_wrapper.async_anthropic_sse_wrapper()
+        self._hidden_params: dict[
+            str, object
+        ] = {}  # mutable-ok: the proxy merges provider headers onto _hidden_params in place
+
+    @property
+    def chunks(self) -> "list[ModelResponseStream] | None":
+        return self._anthropic_wrapper.chunks
+
+    @property
+    def messages(self) -> "list[AllMessageValues] | None":
+        return self._anthropic_wrapper.messages
+
+    @property
+    def model(self) -> str:
+        return self._anthropic_wrapper.model
+
+    async def __anext__(self) -> bytes:
+        return await self._byte_stream.__anext__()
+
+    async def aclose(self) -> None:
+        await self._byte_stream.aclose()

litellm/llms/anthropic/pass_through/adapters/transformation.py

@@ -201,7 +201,7 @@ def create_tool_name_mapping(
 from litellm.types.utils import Choices, ModelResponse, StreamingChoices, Usage
 from litellm.utils import supports_mid_conversation_system
 
-from .streaming_iterator import AnthropicStreamWrapper
+from .streaming_iterator import AnthropicSSEStream, AnthropicStreamWrapper
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
@@ -341,7 +341,7 @@ def translate_completion_output_params_streaming(
         )
         # Return the SSE-wrapped version for proper event formatting.
         if is_async:
-            return anthropic_wrapper.async_anthropic_sse_wrapper()
+            return AnthropicSSEStream(anthropic_wrapper)
         return anthropic_wrapper.anthropic_sse_wrapper()
 
 

litellm/llms/anthropic/pass_through/messages/response_cache.py

@@ -1,7 +1,7 @@
 import re
 from collections.abc import AsyncIterator, Mapping, Sequence
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Final
+from typing import TYPE_CHECKING, Final, cast
 
 import litellm
 from litellm._logging import verbose_logger
@@ -17,6 +17,8 @@
 if TYPE_CHECKING:
     from litellm.caching.caching_handler import LLMCachingHandler
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.types.llms.openai import AllMessageValues
+    from litellm.types.utils import ModelResponseStream
 
 CACHED_STREAM_EVENTS_KEY: Final = "litellm_cached_anthropic_sse_events"
 
@@ -51,6 +53,24 @@ def __init__(
     def has_buffered_provider_output(self) -> bool:
         return getattr(self.stream, "has_buffered_provider_output", False) is True
 
+    @property
+    def chunks(self) -> "list[ModelResponseStream] | None":
+        return cast(  # cast-ok: chunks is a list of ModelResponseStream on the inner stream
+            "list[ModelResponseStream] | None", getattr(self.stream, "chunks", None)
+        )
+
+    @property
+    def messages(self) -> "list[AllMessageValues] | None":
+        return cast(  # cast-ok: messages is a list of AllMessageValues on the inner stream
+            "list[AllMessageValues] | None", getattr(self.stream, "messages", None)
+        )
+
+    @property
+    def model(self) -> str | None:
+        return cast(  # cast-ok: model is a str on the inner stream
+            "str | None", getattr(self.stream, "model", None)
+        )
+
     def __aiter__(self) -> "AnthropicMessagesStreamCacheWriter":
         return self