fix(proxy): a streamed usage.cost the client reads is what litellm billed, written once at the egress

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/proxy_server.py

@@ -922,6 +922,7 @@ from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
 from litellm.proxy.ui_crud_endpoints.user_banner_endpoints import (
     router as user_banner_endpoints_router,
 )
+from litellm.proxy.usage_cost_egress import Pricer, with_client_stream_usage_cost
 from litellm.proxy.utils import (  # noqa: F401, RUF100  # legacy module exports
     PrismaClient,
     ProxyLogging,
@@ -10134,6 +10135,22 @@ def _fast_serialize_simple_model_response_stream(
     return orjson.dumps(payload)
 
 
+def _stream_usage_pricer(request_data: Mapping[str, object]) -> Pricer:
+    """Bills a usage-bearing stream event the way the request's logging object bills the spend row."""
+    logging_obj: Final = request_data.get("litellm_logging_obj")
+    model_name: Final = str(request_data.get("model", ""))
+
+    def price(reply: ModelResponse | ResponsesAPIResponse) -> float | None:
+        return ProxyBaseLLMRequestProcessing._streamed_usage_cost(  # pyright: ignore[reportPrivateUsage]  # the proxy's one stream pricer
+            reply,
+            model_name,
+            None,
+            logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
+        )
+
+    return price
+
+
 def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes:
     if isinstance(chunk, ModelResponseStream):
         serialized_chunk: Final = _fast_serialize_simple_model_response_stream(chunk)
@@ -10472,6 +10489,7 @@ async def async_data_generator(
         needs_per_chunk_hook: Final = proxy_logging_obj.needs_per_chunk_streaming_hook()
         is_raw_sse_stream: Final = bool(request_data.get("_litellm_raw_sse_stream"))
         strip_stream_usage: Final = bool(request_data.get("_litellm_strip_stream_usage"))
+        price_stream_usage: Final = _stream_usage_pricer(request_data)
         raw_sse_buffer = ""
 
         # A stream can start on a deployment with keepalive off and fall back
@@ -10558,7 +10576,7 @@ async def async_data_generator(
                 error_state.observe_chunk(chunk)
             raw_passthrough = False
             if isinstance(chunk, BaseModel):
-                chunk = _serialize_streaming_chunk(chunk)
+                chunk = _serialize_streaming_chunk(with_client_stream_usage_cost(chunk, price_stream_usage))
             elif isinstance(chunk, bytes):
                 chunk = chunk.decode("utf-8", errors="replace")
                 if is_raw_sse_stream:

litellm/proxy/usage_cost_egress.py

@@ -0,0 +1,59 @@
+"""What a client reads in a streamed ``usage.cost``: litellm's bill, written once on the way out of the proxy, onto
+a copy, so the final usage frame and the spend row agree. Inside the pipeline ``usage.cost`` stays what the
+upstream reported."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable
+from typing import Final, TypeVar
+
+from pydantic import BaseModel
+
+from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
+from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
+
+Pricer = Callable[[ModelResponse | ResponsesAPIResponse], float | None]
+"""Prices a reply the way the request's logging object bills it (deployment rates, discount, margin)."""
+
+_RESPONSES_USAGE_EVENTS: Final = frozenset({"response.completed", "response.incomplete"})
+_ReplyT = TypeVar("_ReplyT", bound=BaseModel)
+
+
+def _billable(cost: object) -> float | None:
+    if isinstance(cost, bool) or not isinstance(cost, (int, float)) or not math.isfinite(cost) or cost < 0:
+        return None
+    return float(cost)
+
+
+def _client_usage_cost(upstream_cost: object, bill: object, *, wanted: bool) -> float | None:
+    """The bill, when the upstream put a number in usage.cost or the event carries one regardless (``wanted``).
+    None leaves the usage as the upstream sent it."""
+    reported: Final = isinstance(upstream_cost, (int, float)) and not isinstance(upstream_cost, bool)
+    return _billable(bill) if reported or wanted else None
+
+
+def _with_usage_cost(reply: _ReplyT, usage: Usage | ResponseAPIUsage, cost: float) -> _ReplyT:
+    return reply.model_copy(update={"usage": usage.model_copy(update={"cost": cost})})
+
+
+def with_client_stream_usage_cost(chunk: BaseModel, price: Pricer) -> BaseModel:
+    """A streamed event as the client gets it. Only the usage-bearing events are priced: the chat chunk carrying
+    ``usage``, rewritten where the upstream reported a cost, and the Responses ``response.completed`` /
+    ``response.incomplete`` events, which carry the bill whether or not the upstream reported one."""
+    if isinstance(chunk, ModelResponseStream):
+        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))
+        cost: Final = _client_usage_cost(getattr(usage, "cost", None), price(priced), wanted=False)
+        return chunk if cost is None else _with_usage_cost(chunk, usage, cost)
+    if getattr(chunk, "type", None) not in _RESPONSES_USAGE_EVENTS:
+        return chunk
+    response: Final = getattr(chunk, "response", None)
+    if not isinstance(response, ResponsesAPIResponse) or not isinstance(response.usage, ResponseAPIUsage):
+        return chunk
+    response_cost: Final = _client_usage_cost(getattr(response.usage, "cost", None), price(response), wanted=True)
+    if response_cost is None:
+        return chunk
+    return chunk.model_copy(update={"response": _with_usage_cost(response, response.usage, response_cost)})

tests/unit/proxy/test_usage_cost_egress.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+from typing import Final
+
+import pytest
+
+from litellm.proxy.usage_cost_egress import with_client_stream_usage_cost
+from litellm.types.llms.openai import (
+    ResponseAPIUsage,
+    ResponseCompletedEvent,
+    ResponseIncompleteEvent,
+    ResponsesAPIResponse,
+    ResponsesAPIStreamEvents,
+)
+from litellm.types.utils import ModelResponseStream, StreamingChoices, Usage
+
+UPSTREAM_COST: Final = 0.037
+OPERATOR_BILL: Final = 100 * 2e-6 + 20 * 6e-6
+
+
+def _responses_reply(upstream_cost: float | None) -> ResponsesAPIResponse:
+    return ResponsesAPIResponse(
+        id="resp_egress",
+        created_at=1,
+        status="completed",
+        model="omni-model",
+        object="response",
+        output=[],
+        usage=ResponseAPIUsage(input_tokens=100, output_tokens=20, total_tokens=120, cost=upstream_cost),
+    )
+
+
+def _usage_chunk(upstream_cost: float | None, service_tier: str | None = None) -> ModelResponseStream:
+    return ModelResponseStream(
+        id="chatcmpl-egress",
+        model="omni-model",
+        choices=[],
+        service_tier=service_tier,
+        usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=upstream_cost),
+    )
+
+
+def _bill(_reply: object) -> float | None:
+    return OPERATOR_BILL
+
+
+def _unpriceable(_reply: object) -> float | None:
+    return None
+
+
+def test_the_streamed_usage_chunk_tells_the_client_what_litellm_billed() -> None:
+    chunk: Final = _usage_chunk(UPSTREAM_COST)
+
+    billed: Final = with_client_stream_usage_cost(chunk, _bill)
+
+    assert isinstance(billed, ModelResponseStream) and billed is not chunk
+    assert billed.usage is not None and billed.usage.cost == pytest.approx(OPERATOR_BILL)
+    assert chunk.usage is not None and chunk.usage.cost == UPSTREAM_COST
+
+
+def test_a_zero_bill_replaces_the_upstream_cost() -> None:
+    billed: Final = with_client_stream_usage_cost(_usage_chunk(UPSTREAM_COST), lambda _reply: 0.0)
+
+    assert isinstance(billed, ModelResponseStream)
+    assert billed.usage is not None and billed.usage.cost == 0.0
+
+
+@pytest.mark.parametrize("bill", [None, -0.001, float("nan"), float("inf"), True])
+def test_an_unpriceable_stream_keeps_the_upstream_cost(bill: float | bool | None) -> None:
+    chunk: Final = _usage_chunk(UPSTREAM_COST)
+
+    assert with_client_stream_usage_cost(chunk, lambda _reply: bill) is chunk
+    assert chunk.usage is not None and chunk.usage.cost == UPSTREAM_COST
+
+
+def test_a_streamed_usage_chunk_without_an_upstream_cost_goes_out_as_sent() -> None:
+    chunk: Final = _usage_chunk(None)
+
+    assert with_client_stream_usage_cost(chunk, _bill) is chunk
+
+
+def test_a_content_chunk_is_not_priced() -> None:
+    chunk: Final = ModelResponseStream(
+        id="chatcmpl-egress", model="omni-model", choices=[StreamingChoices(index=0, delta={"content": "hi"})]
+    )
+    priced: list[object] = []
+
+    def price(reply: object) -> float | None:
+        priced.append(reply)
+        return OPERATOR_BILL
+
+    assert with_client_stream_usage_cost(chunk, price) is chunk
+    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 price(reply: object) -> float | None:
+        seen.append((getattr(getattr(reply, "usage", None), "cost", None), getattr(reply, "service_tier", None)))
+        return OPERATOR_BILL
+
+    with_client_stream_usage_cost(_usage_chunk(UPSTREAM_COST, service_tier="default"), price)
+
+    assert seen == [(UPSTREAM_COST, "default")]
+
+
+@pytest.mark.parametrize("upstream_cost", [UPSTREAM_COST, None])
+@pytest.mark.parametrize(
+    "event_type", [ResponsesAPIStreamEvents.RESPONSE_COMPLETED, ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE]
+)
+def test_a_terminal_responses_event_carries_the_bill_with_or_without_an_upstream_cost(
+    upstream_cost: float | None, event_type: ResponsesAPIStreamEvents
+) -> None:
+    event_cls: Final = (
+        ResponseCompletedEvent if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED else ResponseIncompleteEvent
+    )
+    event: Final = event_cls(type=event_type, response=_responses_reply(upstream_cost))
+
+    billed: Final = with_client_stream_usage_cost(event, _bill)
+
+    assert billed is not event
+    assert getattr(billed, "response").usage.cost == pytest.approx(OPERATOR_BILL)
+    assert event.response.usage is not None and event.response.usage.cost == upstream_cost
+
+
+def test_an_unpriceable_terminal_responses_event_goes_out_as_sent() -> None:
+    event: Final = ResponseCompletedEvent(
+        type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_responses_reply(None)
+    )
+
+    assert with_client_stream_usage_cost(event, _unpriceable) is event