diff --git a/litellm/proxy/usage_cost_egress.py b/litellm/proxy/usage_cost_egress.py index 827c86e1ea7..817c8372610 100644 --- a/litellm/proxy/usage_cost_egress.py +++ b/litellm/proxy/usage_cost_egress.py @@ -26,11 +26,14 @@ def _billable(cost: object) -> float | 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.""" +def _client_usage_cost(upstream_cost: object, bill: object) -> float | None: + """The bill, where the upstream put a number in usage.cost (a $0 bill included) or litellm priced the + stream above zero. None leaves the usage as the upstream sent it.""" + billable: Final = _billable(bill) + if billable is None: + return None reported: Final = isinstance(upstream_cost, (int, float)) and not isinstance(upstream_cost, bool) - return _billable(bill) if reported or wanted else None + return billable if reported or billable > 0 else None def _with_usage_cost(reply: _ReplyT, usage: Usage | ResponseAPIUsage, cost: float) -> _ReplyT: @@ -39,21 +42,20 @@ def _with_usage_cost(reply: _ReplyT, usage: Usage | ResponseAPIUsage, cost: floa 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.""" + ``usage`` and the Responses ``response.completed`` / ``response.incomplete`` events.""" 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) + 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: 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) + response_cost: Final = _client_usage_cost(getattr(response.usage, "cost", None), price(response)) 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 index bca4b94c46b..cab33035443 100644 --- a/tests/unit/proxy/test_usage_cost_egress.py +++ b/tests/unit/proxy/test_usage_cost_egress.py @@ -73,10 +73,20 @@ def test_an_unpriceable_stream_keeps_the_upstream_cost(bill: float | None) -> No 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: +def test_a_priced_stream_without_an_upstream_cost_still_tells_the_client_the_bill() -> None: chunk: Final = _usage_chunk(None) - assert with_client_stream_usage_cost(chunk, _bill) is chunk + 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) + + +@pytest.mark.parametrize("bill", [None, 0.0]) +def test_an_unpriced_stream_without_an_upstream_cost_goes_out_as_sent(bill: float | None) -> None: + chunk: Final = _usage_chunk(None) + + assert with_client_stream_usage_cost(chunk, lambda _reply: bill) is chunk def test_a_content_chunk_is_not_priced() -> None: