From a2b911d97b526b39d4bcb04bf7027015d496cb7b Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 10:17:00 -0700 Subject: [PATCH] fix(accounting): retain upstream usage prices during capture --- litellm/litellm_core_utils/litellm_logging.py | 2 ++ litellm/proxy/spend_tracking/cost_compiler.py | 22 +++++++++--- .../spend_tracking/test_cost_compiler.py | 35 +++++++++++++------ 3 files changed, 44 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 6d9df8a7ad5..e22ca540864 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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, diff --git a/litellm/proxy/spend_tracking/cost_compiler.py b/litellm/proxy/spend_tracking/cost_compiler.py index 5a801039b50..37516add6d4 100644 --- a/litellm/proxy/spend_tracking/cost_compiler.py +++ b/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 { diff --git a/tests/unit/proxy/spend_tracking/test_cost_compiler.py b/tests/unit/proxy/spend_tracking/test_cost_compiler.py index f5e35ec7efc..ea769816b59 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_compiler.py +++ b/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)