diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 8fb602ebeb0..6bd0e45fbc0 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2500,7 +2500,9 @@ class CustomStreamWrapper: return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) def _new_usage_chunk(self) -> ModelResponseStream: - usage_chunk: Final = self.model_response_creator() + usage_chunk: Final = self.model_response_creator( + hidden_params=_provider_hidden_params(None, self._provider_response_model) + ) if self.caller_requested_stream_usage: usage_chunk.choices = [] return usage_chunk diff --git a/litellm/proxy/usage_cost_egress.py b/litellm/proxy/usage_cost_egress.py index 817c8372610..1dcdd69706f 100644 --- a/litellm/proxy/usage_cost_egress.py +++ b/litellm/proxy/usage_cost_egress.py @@ -10,6 +10,7 @@ from typing import Final, TypeAlias, TypeVar from pydantic import BaseModel +from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.utils import ModelResponse, ModelResponseStream, Usage @@ -47,7 +48,8 @@ def with_client_stream_usage_cost(chunk: BaseModel, price: Pricer) -> BaseModel: usage: Final = getattr(chunk, "usage", None) if not isinstance(usage, Usage): return chunk - priced: Final = ModelResponse(usage=usage, service_tier=getattr(chunk, "service_tier", None)) + priced: Final = ModelResponse(model=chunk.model, usage=usage, service_tier=getattr(chunk, "service_tier", None)) + priced.hidden_params = dict(get_hidden_params(chunk) or {}) cost: Final = _client_usage_cost(getattr(usage, "cost", None), price(priced)) return chunk if cost is None else _with_usage_cost(chunk, usage, cost) if getattr(chunk, "type", None) not in _RESPONSES_USAGE_EVENTS: diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 8d2b4b51164..6ccb4b8828c 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6118,6 +6118,28 @@ def test_chunk_creator_keeps_provider_model_private_across_stream(): ) +def test_the_usage_chunk_carries_the_served_model_so_its_price_matches_the_assembled_response(): + from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict + + wrapper = CustomStreamWrapper( + completion_stream=None, + model="requested-route", + logging_obj=MagicMock(), + custom_llm_provider="openai", + ) + wrapper.chunk_creator( + chunk=ModelResponseStream( + id="chunk-1", + model="selected-model", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="hello"))], + ) + ) + + usage_chunk = wrapper._new_usage_chunk() + + assert get_hidden_params_dict(usage_chunk)["provider_response_model"] == "selected-model" + + def test_assembled_stream_uses_later_provider_model_for_cost( monkeypatch: pytest.MonkeyPatch, ): diff --git a/tests/unit/proxy/test_usage_cost_egress.py b/tests/unit/proxy/test_usage_cost_egress.py index cab33035443..070139c5019 100644 --- a/tests/unit/proxy/test_usage_cost_egress.py +++ b/tests/unit/proxy/test_usage_cost_egress.py @@ -103,16 +103,26 @@ def test_a_content_chunk_is_not_priced() -> None: assert priced == [] -def test_the_stream_pricer_sees_the_upstream_cost_and_the_served_tier_so_the_bill_matches_the_spend_row() -> None: - seen: list[tuple[object, object]] = [] +def test_the_stream_pricer_sees_the_served_model_tier_and_upstream_cost_so_the_bill_matches_the_spend_row() -> None: + seen: list[tuple[object, object, object, object]] = [] def price(reply: object) -> float | None: - seen.append((getattr(getattr(reply, "usage", None), "cost", None), getattr(reply, "service_tier", None))) + seen.append( + ( + getattr(reply, "model", None), + getattr(reply, "hidden_params", {}).get("provider_response_model"), + getattr(reply, "service_tier", None), + getattr(getattr(reply, "usage", None), "cost", None), + ) + ) return OPERATOR_BILL - with_client_stream_usage_cost(_usage_chunk(UPSTREAM_COST, service_tier="default"), price) + chunk: Final = _usage_chunk(UPSTREAM_COST, service_tier="default") + chunk.hidden_params = {"provider_response_model": "gpt-4o-mini"} - assert seen == [(UPSTREAM_COST, "default")] + with_client_stream_usage_cost(chunk, price) + + assert seen == [("omni-model", "gpt-4o-mini", "default", UPSTREAM_COST)] @pytest.mark.parametrize("upstream_cost", [UPSTREAM_COST, None])