fix(accounting): retain upstream usage prices during capture

litellm/litellm_core_utils/litellm_logging.py

@@ -85,6 +85,7 @@ from litellm.litellm_core_utils.classifier_logging import (
     is_classifier_call,
 )
 from litellm.litellm_core_utils.core_helpers import (
+    billed_usage_cost,
     finite_cost,
     get_provider_response_headers_from_hidden_params,
     is_expected_client_error,
@@ -3076,6 +3077,7 @@ class Logging(LiteLLMLoggingBaseClass):
             {
                 "response_ms": getattr(result, "_response_ms", None),
                 "provider_reported_cost": self._provider_reported_cost,
+                "provider_usage_cost": billed_usage_cost(result, self.custom_llm_provider),
                 "current_billing_time": self.billing_time_for_cost().isoformat(),
                 "custom_pricing": self.custom_pricing_for(result),
                 "custom_cost_per_token": self.custom_cost_per_token,

litellm/proxy/spend_tracking/cost_compiler.py

@@ -71,7 +71,11 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced:
     )
     provider: Final = next((d.value for d in event.dimensions if d.name is DimensionName.PROVIDER), None)
     logged: Final = _fields(pricing.get("logging"))
-    reported: Final = logged.get("provider_reported_cost")
+    reported: Final = (
+        logged.get("provider_reported_cost")
+        if logged.get("provider_reported_cost") is not None
+        else logged.get("provider_usage_cost")
+    )
     upstream: Final = (
         isinstance(reported, (int, float))
         and not isinstance(reported, bool)
@@ -122,6 +126,12 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced:
     )
     response_inputs: Final = _fields(pricing.get("response_inputs"))
     tool_response: Final = _fields(pricing.get("tool_response"))
+    replay_usage: Final = Usage.model_validate(
+        {
+            **{name: value for name, value in (captured.usage or {}).items() if name != "cost"},
+            **({"cost": reported} if upstream and logged.get("provider_reported_cost") is None else {}),
+        }
+    )
     token_response: Final = ModelResponse(
         model=model,
         choices=[
@@ -136,7 +146,7 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced:
                 }
             }
         ],
-        usage=Usage.model_validate({name: value for name, value in (captured.usage or {}).items() if name != "cost"}),
+        usage=replay_usage,
         service_tier=response_inputs.get("service_tier"),
         meta={"billed_units": response_inputs.get("billed_units")},
         hidden_params={
@@ -182,7 +192,7 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced:
         custom_llm_provider=provider,
         litellm_params={**deployment, **request, "metadata": {"model_info": _fields(pricing.get("model_info"))}},
         additional_response_cost=request.get("additional_response_cost"),
-        usage=Usage.model_validate({name: value for name, value in (captured.usage or {}).items() if name != "cost"}),
+        usage=replay_usage,
         end_time=event.ended_at,
         mcp_tool_call_metadata={
             "name": _fields(pricing.get("tool")).get("name"),
@@ -202,7 +212,9 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced:
         }
     )
     logger.record_provider_reported_cost(
-        TypeAdapter[float | None](float | None).validate_python(reported if upstream else None)
+        TypeAdapter[float | None](float | None).validate_python(
+            logged.get("provider_reported_cost") if upstream else None
+        )
     )
     billed_at: Final = TypeAdapter(datetime).validate_python(logged.get("current_billing_time", event.ended_at))
     if event.call_type is CallTypes.arealtime and captured.usage and captured.usage.get("metered_usage") is not None:
@@ -250,7 +262,7 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced:
             "priced_model_key": key,
             "total_cost": Decimal(str(finalized.amount)),
             "guardrail_cost": Decimal(str(guardrail)),
-            "upstream_reported_cost": Decimal(str(logged["provider_reported_cost"])) if upstream else None,
+            "upstream_reported_cost": Decimal(str(reported)) if upstream else None,
             "additional_costs": tuple(
                 {"name": name, "amount": Decimal(str(value))}
                 for name, value in {

tests/unit/proxy/spend_tracking/test_cost_compiler.py

@@ -347,25 +347,40 @@ async def test_real_agent_capture_prices_retained_per_query_fee() -> None:
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("reported_cost", (0.0, 0.123))
-async def test_real_capture_prices_upstream_cost_without_a_model_map(reported_cost: float) -> None:
+@pytest.mark.parametrize("stamped", (False, True))
+@pytest.mark.parametrize("adjusted", (False, True))
+async def test_real_capture_prices_upstream_cost_without_a_model_map(
+    monkeypatch: pytest.MonkeyPatch, reported_cost: float, stamped: bool, adjusted: bool
+) -> None:
     from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage
     from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings
     from litellm.types.utils import ModelResponse, Usage
     from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger
     from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush
 
+    import litellm
+    from litellm.litellm_core_utils.core_helpers import stamp_litellm_bill_on_usage
+
+    monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.05} if adjusted else {})
+    monkeypatch.setattr(
+        litellm, "cost_margin_config", {"global": {"percentage": 0.1, "fixed_amount": 0.001}} if adjusted else {}
+    )
+    expected: Final = reported_cost * 0.95 * 1.1 + 0.001 if adjusted else reported_cost
     logger: Final = _logger()
-    logger.record_provider_reported_cost(reported_cost)
+
+    response: Final = ModelResponse(
+        id="upstream", model="capture-model", usage=Usage(prompt_tokens=10, cost=reported_cost)
+    )
+    if stamped:
+        stamp_litellm_bill_on_usage(response, response.usage, reported_cost + 0.9)
     database: Final = StoredUsage()
     meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush)
     try:
-        capture_logging_usage(
-            meter, logger, ModelResponse(id="upstream", model="capture-model", usage=Usage(prompt_tokens=10)), _AT, _AT
-        )
+        capture_logging_usage(meter, logger, response, _AT, _AT)
         await meter.close()
         result: Final = price_usage(database.rows[-1].envelope)
         assert isinstance(result, PricingSnapshot)
-        assert float(result.total_cost) == pytest.approx(reported_cost)
+        assert float(result.total_cost) == pytest.approx(expected)
         assert result.upstream_reported_cost == Decimal(str(reported_cost))
     finally:
         await meter.stop(1)
@@ -881,13 +896,13 @@ async def test_stream_replay_applies_adjustments_once_to_legacy_stamped_usage(
     )
     expected: Final = _logger(stream=True).response_cost_calculator(response, ignore_stamped_cost=True)
     assert expected is not None
-    stamped: Final = response.model_copy(
-        update={"usage": Usage(prompt_tokens=10, completion_tokens=completion_tokens, cost=expected)}
-    )
+    from litellm.litellm_core_utils.core_helpers import stamp_litellm_bill_on_usage
+
+    stamp_litellm_bill_on_usage(response, response.usage, expected)
     database: Final = StoredUsage()
     meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush)
     try:
-        capture_logging_usage(meter, logger, stamped, _AT, _AT)
+        capture_logging_usage(meter, logger, response, _AT, _AT)
         await meter.close()
         priced: Final = price_usage(database.rows[-1].envelope)
         assert isinstance(priced, PricingSnapshot)