From e8b1f20be1a315588918a0df2ec7560207a759e9 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 12:15:03 -0700 Subject: [PATCH] fix(accounting): retain callback fees and terminal pricing context --- litellm/proxy/spend_tracking/cost_compiler.py | 11 ++- litellm/proxy/spend_tracking/usage_capture.py | 18 ++++- .../spend_tracking/test_cost_compiler.py | 81 +++++++++++++++++-- 3 files changed, 100 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/spend_tracking/cost_compiler.py b/litellm/proxy/spend_tracking/cost_compiler.py index d85513da70a..e40ed19afcc 100644 --- a/litellm/proxy/spend_tracking/cost_compiler.py +++ b/litellm/proxy/spend_tracking/cost_compiler.py @@ -129,10 +129,17 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: for name, snapshot in _fields(pricing.get("realtime_models")).items() if _fields(snapshot).get("model_map_value") is not None } + base_pricing: Final = _fields(pricing.get("base_model_information")) + base_key: Final = base_pricing.get("model_map_key") context: Final = RetainedPricing( _json( { **secondary, + **( + {base_key: _fields(base_pricing.get("model_map_value"))} + if isinstance(base_key, str) and base_pricing.get("model_map_value") is not None + else {} + ), **{ alias: info for alias in dict.fromkeys((model, key, f"{provider}/{model}", model.removeprefix(f"{provider}/"))) @@ -267,7 +274,9 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: return Unpriced(usage_event_id=event.event_id, reason="calculation_failed") retained_final: Final = captured.finalized_cost or {} breakdown: Final = _fields(retained_final.get("cost_breakdown")) - guardrail: Final = TypeAdapter(float).validate_python(breakdown.get("guardrail_cost", 0.0)) + guardrail: Final = TypeAdapter(float).validate_python( + breakdown.get("guardrail_cost", pricing.get("guardrail_cost", 0.0)) + ) allowance: Final = TypeAdapter(float).validate_python(retained_final.get("allowance_discount", 0.0)) if allowance and captured.allocation is None: return Unpriced(usage_event_id=event.event_id, reason="allowance_allocation_missing") diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py index 674862713dc..5c735eaa6d9 100644 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ b/litellm/proxy/spend_tracking/usage_capture.py @@ -281,6 +281,7 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res mask_api_base_credentials, ) from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost + from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups @@ -336,7 +337,11 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res else kwargs.get("_accounting_allocation") ) allocation: Final = allocation_value if isinstance(allocation_value, SearchAllowanceAllocation) else None - pricing_snapshot: Final = kwargs.get("_accounting_pricing_snapshot") + pricing_snapshot: Final = ( + _capture_pricing(logging_obj, response, kwargs) + if isinstance(logging_obj, Logging) and _batch_rows(response) + else kwargs.get("_accounting_pricing_snapshot") + ) observation: Final = kwargs.get("_accounting_usage_snapshot") captured_pricing: Final = { "model_map_information": ( @@ -345,6 +350,7 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res else metadata.get("model_map_information") ), "cost_breakdown": metadata.get("cost_breakdown"), + "guardrail_cost": guardrail_information_cost(metadata.get("guardrail_information")), **( _JSON_OBJECT.validate_json(pricing_snapshot.inputs) if isinstance(pricing_snapshot, PricingCapture) @@ -889,6 +895,7 @@ def _capture_effective_pricing( import litellm from litellm.litellm_core_utils.cost_discount import resolve_cost_discount from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + from litellm.utils import get_base_model_from_metadata params: Final = _fields(details.get("litellm_params")) prepared: Final = logging_obj.prepared_cost @@ -929,6 +936,7 @@ def _capture_effective_pricing( else None, } ) + base_model: Final = get_base_model_from_metadata(model_call_details=logging_obj.model_call_details) realtime_models: Final = tuple( dict.fromkeys( chain.from_iterable(_realtime_pricing_models(event) for event in _realtime_cost_events(normalized)) @@ -943,6 +951,14 @@ def _capture_effective_pricing( _json( { **_pricing_inputs(details, logging_obj, normalized), + "base_model_information": StandardLoggingPayloadSetup.get_model_cost_information( + base_model=base_model, + custom_pricing=False, + custom_llm_provider=None, + init_response_obj={"model": base_model}, + ) + if base_model is not None + else None, "realtime_models": { model: StandardLoggingPayloadSetup.get_model_cost_information( base_model=model, diff --git a/tests/unit/proxy/spend_tracking/test_cost_compiler.py b/tests/unit/proxy/spend_tracking/test_cost_compiler.py index c54f5bc7039..212b5c556ce 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_compiler.py +++ b/tests/unit/proxy/spend_tracking/test_cost_compiler.py @@ -352,15 +352,14 @@ async def test_real_agent_capture_prices_retained_per_query_fee() -> None: async def test_real_capture_prices_upstream_cost_without_a_model_map( monkeypatch: pytest.MonkeyPatch, reported_cost: float, stamped: bool, adjusted: bool ) -> None: + import litellm + from litellm.litellm_core_utils.core_helpers import stamp_litellm_bill_on_usage 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 {} @@ -914,7 +913,10 @@ async def test_stream_replay_applies_adjustments_once_to_legacy_stamped_usage( @pytest.mark.asyncio @pytest.mark.parametrize("amount", [0.0, 0.123]) -async def test_capture_retains_provider_header_price_without_catalog(amount: float) -> None: +@pytest.mark.parametrize("mapped", [False, True]) +async def test_capture_retains_provider_header_price_without_catalog( + amount: float, mapped: bool, monkeypatch: pytest.MonkeyPatch +) -> 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 @@ -928,14 +930,37 @@ async def test_capture_retains_provider_header_price_without_catalog(amount: flo from litellm.litellm_core_utils.core_helpers import set_response_cost_in_hidden_params set_response_cost_in_hidden_params(response, amount) + import litellm + + logger: Final = _logger() + if mapped: + monkeypatch.setitem( + litellm.model_cost, + "mapped-cost-model", + { + "mode": "chat", + "litellm_provider": "openai", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + ) + logger.model_call_details["litellm_params"]["metadata"] = {"model_info": {"base_model": "mapped-cost-model"}} + expected: Final = logger.response_cost_calculator(response, ignore_stamped_cost=True) + expected_breakdown: Final = logger.cost_breakdown database: Final = StoredUsage() meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) try: - capture_logging_usage(meter, _logger(), response, _AT, _AT) + capture_logging_usage(meter, logger, response, _AT, _AT) await meter.close() + if mapped: + monkeypatch.setitem(litellm.model_cost, "mapped-cost-model", {"input_cost_per_token": 99.0}) result: Final = price_usage(database.rows[-1].envelope) assert isinstance(result, PricingSnapshot) assert result.total_cost == Decimal(str(amount)) + assert float(result.total_cost) == expected + assert expected_breakdown is not None + assert float(result.input_cost or 0) == pytest.approx(expected_breakdown["input_cost"]) + assert float(result.output_cost or 0) == pytest.approx(expected_breakdown["output_cost"]) assert (result.input_cost or Decimal(0)) + (result.output_cost or Decimal(0)) == result.total_cost assert result.pricing_source == "upstream_header" finally: @@ -961,7 +986,10 @@ def test_terminal_response_replay_prices_usage_without_changing_event_identity(c @pytest.mark.asyncio -async def test_batch_terminal_capture_refreshes_row_prices_after_creation(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize("freeze_pricing", [False, True]) +async def test_batch_terminal_capture_refreshes_row_prices_after_creation( + monkeypatch: pytest.MonkeyPatch, freeze_pricing: bool +) -> None: import litellm from litellm.proxy.spend_tracking.usage_capture import capture_batch_callback, capture_logging_usage from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings @@ -998,8 +1026,11 @@ async def test_batch_terminal_capture_refreshes_row_prices_after_creation(monkey } try: capture_logging_usage(meter, logger, {"id": "batch", "status": "in_progress"}, _AT, _AT, freeze_only=True) - capture_logging_usage(meter, logger, response, _AT, _AT) - await capture_batch_callback(meter, database.load, logger.model_call_details, response, _AT, _AT) + if freeze_pricing: + capture_logging_usage(meter, logger, response, _AT, _AT) + await capture_batch_callback( + meter, database.load, {**logger.model_call_details, "litellm_logging_obj": logger}, response, _AT, _AT + ) first_rows: Final = database.rows monkeypatch.setitem(litellm.model_cost, "capture-model", {**rates, "input_cost_per_token_batches": 99.0}) capture_logging_usage(meter, logger, response, _AT, _AT) @@ -1087,3 +1118,37 @@ async def test_failed_fixed_call_without_usage_does_not_invent_a_free_charge(cal ) finally: await meter.stop(1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("billable", [False, True]) +async def test_capture_preserves_completed_guardrail_fee_without_header_finalization( + stream: bool, billable: 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 + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + logger: Final = _capture_logger(stream=stream) + logger.model_call_details["litellm_params"]["metadata"] = { + "standard_logging_guardrail_information": [ + {"guardrail_name": "late-fee", "guardrail_cost": 0.03, "guardrail_cost_in_spend": billable} + ] + } + response: Final = ModelResponse( + id="guarded", model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=4, total_tokens=14) + ) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, logger, response, _AT, _AT) + await meter.close() + result: Final = price_usage(database.rows[-1].envelope) + assert isinstance(result, PricingSnapshot) + assert result.guardrail_cost == Decimal("0.03" if billable else "0") + assert float(result.total_cost) == pytest.approx(0.18 + (0.03 if billable else 0)) + finally: + await meter.stop(1)