From 0a0281f1c11905e07e357a124c4671cd59c9e88f Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 27 Oct 2026 21:43:18 +0000 Subject: [PATCH] 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 | 20 +++- litellm/proxy/usage_cost_egress.py | 59 +++++++++ tests/unit/proxy/test_usage_cost_egress.py | 132 +++++++++++++++++++++ 3 files changed, 210 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/usage_cost_egress.py create mode 100644 tests/unit/proxy/test_usage_cost_egress.py diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bc79becf7bb..5c7388f7861 100644 --- a/litellm/proxy/proxy_server.py +++ b/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: diff --git a/litellm/proxy/usage_cost_egress.py b/litellm/proxy/usage_cost_egress.py new file mode 100644 index 00000000000..b76d6579a6a --- /dev/null +++ b/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)}) diff --git a/tests/unit/proxy/test_usage_cost_egress.py b/tests/unit/proxy/test_usage_cost_egress.py new file mode 100644 index 00000000000..0be4a6b745a --- /dev/null +++ b/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