scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(accounting): retain callback fees and terminal pricing context

Browse source
This commit is contained in:
Joshua Valluru 2026-10-31 12:15:03 -07:00
parent 8c23bec984
commit e8b1f20be1
3 changed files with 100 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/proxy/spend_tracking/cost_compiler.py
View file

@ -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")

18
litellm/proxy/spend_tracking/usage_capture.py
View file

@ -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,

81
tests/unit/proxy/spend_tracking/test_cost_compiler.py
View file

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