fix(proxy): keep the bill in a priced chat stream's usage.cost when the upstream reported none

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

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)})

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: