From 26c6c23c48362d996ea7beeed1895757215fe9f2 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 08:31:00 -0700 Subject: [PATCH] feat(accounting): compile retained usage into authoritative costs --- .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 3 + litellm/cost_calculator.py | 97 +- litellm/litellm_core_utils/litellm_logging.py | 126 ++- .../llm_cost_calc/pricing_context.py | 65 ++ litellm/main.py | 55 +- litellm/proxy/db/accounting_ledger.py | 44 +- litellm/proxy/proxy_server.py | 22 +- litellm/proxy/schema.prisma | 3 + litellm/proxy/spend_tracking/cost_compiler.py | 324 ++++++- litellm/proxy/spend_tracking/usage_capture.py | 112 ++- litellm/utils.py | 5 + schema.prisma | 3 + .../database/test_accounting_ledger.py | 210 ++++- .../spend_tracking/test_cost_compiler.py | 849 ++++++++++++++++++ .../spend_tracking/test_usage_capture.py | 74 +- tests/unit/test_cost_calculator.py | 47 + .../test_register_model_custom_pricing.py | 370 ++++---- 18 files changed, 2132 insertions(+), 282 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_accounting_compilation_state/migration.sql create mode 100644 litellm/litellm_core_utils/llm_cost_calc/pricing_context.py create mode 100644 tests/unit/proxy/spend_tracking/test_cost_compiler.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_accounting_compilation_state/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_accounting_compilation_state/migration.sql new file mode 100644 index 00000000000..04fbdfdc237 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261031000000_accounting_compilation_state/migration.sql @@ -0,0 +1,5 @@ +ALTER TABLE "LiteLLM_AccountingEvent" +ADD COLUMN "compilation_state" TEXT NOT NULL DEFAULT 'pending', +ADD COLUMN "compilation_error" TEXT; +CREATE INDEX "LiteLLM_AccountingEvent_kind_compilation_state_recorded_at_idx" +ON "LiteLLM_AccountingEvent"("kind", "compilation_state", "recorded_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 179b741fe7c..048e0e1e17a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -2111,9 +2111,12 @@ model LiteLLM_AccountingEvent { occurred_at DateTime recorded_at DateTime @default(now()) payload Json + compilation_state String @default("pending") + compilation_error String? @@index([recorded_at]) @@index([kind, source_event_id]) + @@index([kind, compilation_state, recorded_at]) } // One row per (event, target table, version): pending work for the projector and, diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index e66be233be4..76cced0c1c9 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -26,6 +26,12 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping, billed_usage from litellm.litellm_core_utils.cost_discount import resolve_cost_discount from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, is_decisions_call_type from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR +from litellm.litellm_core_utils.llm_cost_calc.pricing_context import ( + pricing_discounts, + pricing_margins, + retained_model_info, + retained_models, +) from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) @@ -379,7 +385,9 @@ def _per_second_pricing_cost( audio_seconds: float = 0.0, ) -> tuple[float, float] | None: try: - model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + model_info: Final = retained_model_info(model, custom_llm_provider) or cached_get_model_info_helper( + model=model, custom_llm_provider=custom_llm_provider + ) except Exception: # noqa: BLE001 # the lookup raises plain Exception for an unmapped model return None if _has_token_or_tiered_pricing(model_info) or not _bills_wall_clock_seconds(model_info): @@ -538,7 +546,7 @@ def cost_per_token( if response_cost is not None: return response_cost[0], response_cost[1] - model_cost_ref: Final = litellm.model_cost + model_cost_ref: Final = retained_models() or litellm.model_cost # Only callers that explicitly pass `custom_llm_provider` get the # dedup/prefix-join treatment. When provider is omitted, preserve legacy # behavior: `model_with_provider` stays equal to the raw `model` string @@ -600,9 +608,11 @@ def cost_per_token( model = model_without_prefix if is_decisions_call_type(call_type): - decisions_model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + decisions_model_info: Final = retained_model_info(model, custom_llm_provider) or cached_get_model_info_helper( + model=model, custom_llm_provider=custom_llm_provider + ) decisions_overrides: Final = decisions_pricing_fields( - litellm.model_cost.get(decisions_model_info["key"], {}) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # model cost map is dynamically loaded + (retained_models() or litellm.model_cost).get(decisions_model_info["key"], {}) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # model cost map is dynamically loaded ) if decisions_overrides: return generic_cost_per_token( @@ -623,7 +633,9 @@ def cost_per_token( ) if lyria_generation_cost is not None: return 0.0, lyria_generation_cost - speech_model_info: Final = cached_get_model_info_helper( + speech_model_info: Final = retained_model_info( + model_without_prefix, custom_llm_provider + ) or cached_get_model_info_helper( model=model_without_prefix, custom_llm_provider=custom_llm_provider, preserve_unset_costs=True, @@ -808,7 +820,9 @@ def cost_per_token( service_tier=service_tier, ) else: - model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + model_info: Final = retained_model_info(model, custom_llm_provider) or cached_get_model_info_helper( + model=model, custom_llm_provider=custom_llm_provider + ) if _has_token_or_tiered_pricing(model_info): return generic_cost_per_token( model=model, @@ -995,8 +1009,8 @@ def select_model_name_for_cost_calc( ) if custom_pricing is True: - if router_model_id is not None and router_model_id in litellm.model_cost: - entry: Final = litellm.model_cost[router_model_id] + if router_model_id is not None and router_model_id in (retained_models() or litellm.model_cost): + entry: Final = (retained_models() or litellm.model_cost)[router_model_id] if _cost_map_entry_prices_anything(entry): return_model = router_model_id else: @@ -1075,14 +1089,16 @@ def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> registered cost key ("vertex_ai/claude-opus-5"), keeping the model unchanged when it already resolves downstream (custom-priced router ids) or no stripped candidate is registered (#38069).""" segments: Final = model.split("/") - if "/".join(segments[1:]) in litellm.model_cost: + if "/".join(segments[1:]) in (retained_models() or litellm.model_cost): return model head_len: Final = 2 if region_name is not None and len(segments) > 2 and segments[1] == region_name else 1 head: Final = "/".join(segments[:head_len]) tail: Final = segments[head_len:] strippable: Final = next((index for index, segment in enumerate(tail) if segment in LlmProvidersSet), len(tail)) candidates: Final = (f"{head}/{'/'.join(tail[start:])}" for start in range(min(strippable, len(tail) - 1) + 1)) - return next((candidate for candidate in candidates if candidate in litellm.model_cost), model) + return next( + (candidate for candidate in candidates if candidate in (retained_models() or litellm.model_cost)), model + ) @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) @@ -1348,9 +1364,7 @@ def _apply_cost_discount( Tuple of (final_cost, discount_percent, discount_amount) """ original_cost: Final = base_cost - resolved_discount: Final = resolve_cost_discount( - litellm.cost_discount_config, custom_llm_provider, model, region_name - ) + resolved_discount: Final = resolve_cost_discount(pricing_discounts(), custom_llm_provider, model, region_name) discount_percent: Final = 0.0 if resolved_discount is None else resolved_discount discount_amount: Final = original_cost * discount_percent @@ -1389,12 +1403,12 @@ def _apply_cost_margin( # Get margin config - check provider-specific first, then global margin_config = None - if custom_llm_provider and custom_llm_provider in litellm.cost_margin_config: - margin_config = litellm.cost_margin_config[custom_llm_provider] + if custom_llm_provider and custom_llm_provider in pricing_margins(): + margin_config = pricing_margins()[custom_llm_provider] if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug("Found provider-specific margin config for %s: %s", custom_llm_provider, margin_config) - elif "global" in litellm.cost_margin_config: - margin_config = litellm.cost_margin_config["global"] + elif "global" in pricing_margins(): + margin_config = pricing_margins()["global"] if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug("Using global margin config: %s", margin_config) else: @@ -1402,7 +1416,7 @@ def _apply_cost_margin( verbose_logger.debug( "No margin config found. Provider: %s, Available configs: %s", custom_llm_provider, - list(litellm.cost_margin_config.keys()), + list(pricing_margins().keys()), ) if margin_config is not None: @@ -2159,12 +2173,14 @@ def completion_cost( # Calculate cost based on prompt_tokens, completion_tokens if ( "togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai" - ) and not has_together_registry_pricing(model, litellm.model_cost): + ) and not has_together_registry_pricing(model, (retained_models() or litellm.model_cost)): model = get_model_params_and_category(model, call_type=CallTypes(call_type)) # replicate llms are calculate based on time for request running # see https://replicate.com/pricing - elif (model in litellm.replicate_models or "replicate" in model) and model not in litellm.model_cost: + elif (model in litellm.replicate_models or "replicate" in model) and model not in ( + retained_models() or litellm.model_cost + ): # for unmapped replicate model, default to replicate's time tracking logic return get_replicate_completion_pricing(completion_response, response_time_ms) @@ -2261,7 +2277,7 @@ def completion_cost( _final_cost += sum(additional_costs.values()) original_cost = _final_cost - if litellm.cost_discount_config: + if pricing_discounts(): ( _final_cost, discount_percent, @@ -2277,7 +2293,7 @@ def completion_cost( discount_amount = 0.0 # Apply margin from module-level config if configured - if litellm.cost_margin_config: + if pricing_margins(): ( _final_cost, margin_percent, @@ -2629,7 +2645,7 @@ def _priced_candidates_before_unpriced_rules(candidates: tuple[_CostCandidate, . def _raw_cost_map_entry(key: str) -> Mapping[str, object] | None: - raw_entry: Final = litellm.model_cost.get(key) + raw_entry: Final = (retained_models() or litellm.model_cost).get(key) return raw_entry if isinstance(raw_entry, Mapping) else None @@ -2973,12 +2989,19 @@ def default_image_cost_calculator( model_name_without_custom_llm_provider, ) matched_model: Final = next( - (_model for _model in models_to_check if _model is not None and _model in litellm.model_cost), None + ( + _model + for _model in models_to_check + if _model is not None and _model in (retained_models() or litellm.model_cost) + ), + None, ) if matched_model is None and model_info is None: raise Exception(f"Model not found in cost map. Tried checking {models_to_check}") - shared_cost_info: Final = litellm.model_cost[matched_model] if matched_model is not None else None + shared_cost_info: Final = ( + (retained_models() or litellm.model_cost)[matched_model] if matched_model is not None else None + ) price_tables: Final = tuple(table for table in (model_info, shared_cost_info) if table is not None) image_count: Final = n if n is not None else 1 unit_counts: Final = ( @@ -3016,7 +3039,7 @@ def default_video_cost_calculator( custom_llm_provider (Optional[str]): Custom LLM provider model_info (Optional[ModelInfo]): Deployment-level model info containing custom video pricing. When provided, used before falling back to - the global litellm.model_cost lookup. + the effective model cost lookup. video_resolution (Optional[str]): From usage (e.g. ``720p``, ``1080p``) for tiered per-second pricing. Returns: @@ -3026,9 +3049,9 @@ def default_video_cost_calculator( Exception: If model pricing not found in cost map """ # Use custom model_info pricing if provided (deployment-specific pricing) - cost_info: dict | None = None + cost_info: ModelInfo | None = None if model_info is not None: - cost_info = dict(model_info) + cost_info = model_info else: # Build model names for cost lookup base_model_name = model @@ -3049,15 +3072,15 @@ def default_video_cost_calculator( model_name_without_custom_llm_provider, ] for _model in models_to_check: - if _model is not None and _model in litellm.model_cost: - cost_info = litellm.model_cost[_model] + if _model is not None and _model in (retained_models() or litellm.model_cost): + cost_info = (retained_models() or litellm.model_cost)[_model] break # If still not found, try with custom_llm_provider prefix if cost_info is None and custom_llm_provider: prefixed_model: Final = f"{custom_llm_provider}/{model}" - if prefixed_model in litellm.model_cost: - cost_info = litellm.model_cost[prefixed_model] + if prefixed_model in (retained_models() or litellm.model_cost): + cost_info = (retained_models() or litellm.model_cost)[prefixed_model] if cost_info is None: raise Exception(f"Model not found in cost map for model={model}") @@ -3663,9 +3686,9 @@ def _cost_map_entry_declares_pricing(model_name: str, custom_llm_provider: str) """ resolved: Final = _cost_map_model_info(model_name, custom_llm_provider) entries: Final = ( - litellm.model_cost.get(resolved.get("key")) if resolved is not None else None, - litellm.model_cost.get(model_name), - litellm.model_cost.get(f"{custom_llm_provider}/{model_name}"), + (retained_models() or litellm.model_cost).get(resolved.get("key")) if resolved is not None else None, + (retained_models() or litellm.model_cost).get(model_name), + (retained_models() or litellm.model_cost).get(f"{custom_llm_provider}/{model_name}"), ) return any(entry is not None and _cost_map_entry_prices_anything(entry) for entry in entries) @@ -3910,11 +3933,11 @@ def _model_info_rate(info: ModelInfo, key: _TranscriptionRateKey) -> float | Non def _declared_transcription_rate(info: ModelInfo | None, keys: tuple[_TranscriptionRateKey, ...]) -> float | None: - """First of ``keys`` this entry prices, read off the raw ``litellm.model_cost`` entry + """First of ``keys`` this entry prices, read off the raw model cost entry because ``get_model_info`` synthesizes zero token rates for entries that omit them.""" if info is None: return None - declared: Final = litellm.model_cost.get(info.get("key")) + declared: Final = (retained_models() or litellm.model_cost).get(info.get("key")) if declared is None: return None return next( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 639003c9a6a..6d9df8a7ad5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -14,6 +14,7 @@ import sys import time import traceback from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence +from contextlib import nullcontext from dataclasses import dataclass from datetime import datetime as dt_object from enum import Enum @@ -25,7 +26,13 @@ from httpx import Response, Timeout from pydantic import BaseModel, JsonValue, TypeAdapter import litellm -from litellm._internal_context import in_mirrored_call, outside_mirrored_call, post_response_phase +from litellm._internal_context import ( + current_billing_time, + in_mirrored_call, + outside_mirrored_call, + pinned_billing_time, + post_response_phase, +) from litellm._logging import ( is_debugging_on, redact_string, @@ -112,6 +119,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( object_mapping, with_guardrail_entry, ) +from litellm.litellm_core_utils.llm_cost_calc.pricing_context import RetainedPricing, retained_pricing from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) @@ -666,6 +674,10 @@ class UsageCaptureObservation: return cls(serialized, provenance, response_cost) +_REQUEST_MODEL_INFO: Final = TypeAdapter(ModelInfo) +_REQUEST_PRICING_MODELS: Final = TypeAdapter(dict[str, ModelInfo]) + + class Logging(LiteLLMLoggingBaseClass): usage_capture_callback: ClassVar[ Callable[ @@ -703,6 +715,8 @@ class Logging(LiteLLMLoggingBaseClass): logfireLogger, \ prometheusLogger, \ slack_app + cost_billing_time: datetime.datetime | None = None + request_pricing: RetainedPricing | None = None custom_pricing: bool = False custom_cost_per_token: CostPerToken | None = None custom_cost_per_second: float | None = None @@ -2107,6 +2121,57 @@ class Logging(LiteLLMLoggingBaseClass): if margin_total_amount is not None: self.cost_breakdown["margin_total_amount"] = margin_total_amount + def _request_pricing_model_info( + self, model: str, provider: str, overrides: Mapping[str, object], deployment_id: str | None + ) -> ModelInfo: + for candidate in (deployment_id, f"{provider}/{model}"): + if candidate is None: + continue + try: + return litellm.get_model_info(model=candidate, custom_llm_provider=provider) + except Exception: # noqa: BLE001 # unmapped SDK models can supply their complete pricing + verbose_logger.debug("No base pricing for request override candidate %s", candidate) + return _REQUEST_MODEL_INFO.validate_python( + { + "key": deployment_id or f"{provider}/{model}", + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + "input_cost_per_token": None, + "output_cost_per_token": None, + "litellm_provider": provider, + "mode": overrides.get( + "mode", "embedding" if "embedding" in str(self.model_call_details.get("call_type", "")) else "chat" + ), + "supported_openai_params": None, + } + ) + + def set_request_pricing( + self, + model: str, + provider: str, + overrides: Mapping[str, object], + deployment_id: str | None = None, + ) -> None: + from litellm.utils import model_info_with_pricing_overrides + + base: Final = self._request_pricing_model_info(model, provider, overrides, deployment_id) + normalized: Final = _REQUEST_MODEL_INFO.validate_python( + { + **base, + "mode": base.get("mode") + or ("embedding" if "embedding" in str(self.model_call_details.get("call_type", "")) else "chat"), + } + ) + effective: Final = model_info_with_pricing_overrides(normalized, overrides) + keys: Final = (model, f"{provider}/{model}", deployment_id or model, effective["key"]) + self.request_pricing = RetainedPricing( + models_json=_REQUEST_PRICING_MODELS.dump_json({name: effective for name in keys}).decode(), + discounts_json=json.dumps(litellm.cost_discount_config), + margins_json=json.dumps(litellm.cost_margin_config), + ) + def response_cost_calculator( self, result: object, @@ -2114,6 +2179,24 @@ class Logging(LiteLLMLoggingBaseClass): litellm_model_name: str | None = None, router_model_id: str | None = None, ignore_stamped_cost: bool = False, + ) -> float | None: + context: Final = ( + retained_pricing(self.request_pricing) + if self.request_pricing is not None and not self.spilled_over_ptu_for(result) + else nullcontext() + ) + with context, pinned_billing_time(self.billing_time_for_cost()): + return self._calculate_response_cost( + result, cache_hit, litellm_model_name, router_model_id, ignore_stamped_cost + ) + + def _calculate_response_cost( + self, + result: object, + cache_hit: bool | None = None, + litellm_model_name: str | None = None, + router_model_id: str | None = None, + ignore_stamped_cost: bool = False, ) -> float | None: """ Calculate response cost using result + logging object variables. @@ -2404,7 +2487,9 @@ class Logging(LiteLLMLoggingBaseClass): def custom_pricing_for(self, result: object) -> bool: if self.spilled_over_ptu_for(result): return False - return use_custom_pricing_for_model(litellm_params=getattr(self, "litellm_params", None)) + return self.request_pricing is not None or use_custom_pricing_for_model( + litellm_params=getattr(self, "litellm_params", None) + ) def _prompt_for_cost_calculation(self) -> str: """ @@ -2965,6 +3050,8 @@ class Logging(LiteLLMLoggingBaseClass): metered: Final = self.model_call_details.get(REALTIME_METERED_USAGE_KEY) return { "id": self.litellm_call_id, + "model": self.model, + "results": result, "usage": {"metered_usage": TypeAdapter(object).dump_python(metered, mode="json")} if metered is not None else None, @@ -2974,14 +3061,29 @@ class Logging(LiteLLMLoggingBaseClass): ) return normalized if assembled is None else assembled + def record_provider_reported_cost(self, cost: float | None) -> None: + self._provider_reported_cost = cost + + def billing_time_for_cost(self) -> datetime.datetime: + if self.cost_billing_time is None: + self.cost_billing_time = current_billing_time() + return self.cost_billing_time + def usage_capture_pricing(self, result: object) -> Mapping[str, JsonValue]: from litellm.utils import count_characters return TypeAdapter(dict[str, JsonValue]).validate_python( { + "response_ms": getattr(result, "_response_ms", None), + "provider_reported_cost": self._provider_reported_cost, + "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, "custom_cost_per_second": self.custom_cost_per_second, + "tool_response_cost": self.model_call_details.get("response_cost") + if self.model_call_details.get("call_type") + in (CallTypes.call_mcp_tool.value, CallTypes.send_message.value, CallTypes.asend_message.value) + else None, "prompt_characters": count_characters(self._prompt_for_cost_calculation()) if self.model_call_details.get("call_type") in (CallTypes.speech.value, CallTypes.aspeech.value) else None, @@ -7838,14 +7940,20 @@ def get_standard_logging_object_payload( timing_metrics: Final = getattr(logging_obj, "response_timing_metrics", None) or {} clean_hidden_params["litellm_overhead_time_ms"] = timing_metrics.get("litellm_overhead_time_ms") - model_cost_information: Final = StandardLoggingPayloadSetup.get_model_cost_information( - base_model=base_model, - custom_pricing=custom_pricing, - custom_llm_provider=kwargs.get("custom_llm_provider"), - init_response_obj=init_response_obj, - api_base=litellm_params.get("api_base"), - router_model_id=logging_obj.get_router_model_id(), + pricing_context: Final = ( + retained_pricing(logging_obj.request_pricing) + if logging_obj.request_pricing is not None and not spilled_over + else nullcontext() ) + with pricing_context: + model_cost_information: Final = StandardLoggingPayloadSetup.get_model_cost_information( + base_model=base_model, + custom_pricing=custom_pricing, + custom_llm_provider=kwargs.get("custom_llm_provider"), + init_response_obj=init_response_obj, + api_base=litellm_params.get("api_base"), + router_model_id=logging_obj.get_router_model_id(), + ) error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload( metadata=metadata, diff --git a/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py b/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py new file mode 100644 index 00000000000..f9509cac41e --- /dev/null +++ b/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py @@ -0,0 +1,65 @@ +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter + +from litellm.types.utils import ModelInfo + +_MODELS: Final = TypeAdapter(dict[str, ModelInfo]) +_DISCOUNTS: Final = TypeAdapter(dict[str, float]) +_MARGINS: Final = TypeAdapter(dict[str, float | dict[str, float]]) + + +@dataclass(frozen=True, slots=True) +class RetainedPricing: + models_json: str + discounts_json: str + margins_json: str + + def models(self) -> Mapping[str, ModelInfo]: + return _MODELS.validate_json(self.models_json) + + +_CURRENT: Final[ContextVar[RetainedPricing | None]] = ContextVar("retained_pricing", default=None) + + +@contextmanager +def retained_pricing(pricing: RetainedPricing) -> Generator[None]: + token: Final = _CURRENT.set(pricing) + try: + yield + finally: + _CURRENT.reset(token) + + +def retained_model_info(model: str, provider: str | None) -> ModelInfo | None: + context: Final = _CURRENT.get() + if context is None: + return None + models: Final = context.models() + result: Final = models.get(f"{provider}/{model}") or models.get(model) + if result is None: + raise ValueError(f"Pricing snapshot has no retained model {model}") + return result + + +def retained_models() -> Mapping[str, ModelInfo] | None: + context: Final = _CURRENT.get() + return None if context is None else context.models() + + +def pricing_discounts() -> Mapping[str, float]: + import litellm + + context: Final = _CURRENT.get() + return litellm.cost_discount_config if context is None else _DISCOUNTS.validate_json(context.discounts_json) + + +def pricing_margins() -> Mapping[str, float | Mapping[str, float]]: + import litellm + + context: Final = _CURRENT.get() + return litellm.cost_margin_config if context is None else _MARGINS.validate_json(context.margins_json) diff --git a/litellm/main.py b/litellm/main.py index 78767b03eb7..170a4cd3551 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1316,7 +1316,6 @@ from litellm.responses.bridge_check import ( # noqa: E402 # re-export at the o ) from litellm.responses.bridge_check import ( bridge_reasoning_effort, - keep_registered_mode, str_keyed_mapping, ) from litellm.responses.bridge_check import ( # noqa: E402 # re-export at the old definition site @@ -1472,50 +1471,20 @@ def _record_agentic_loop_params( ) -def _register_custom_pricing_for_request( +def _set_request_pricing( + logging_obj: object, model: str, custom_llm_provider: str, kwargs: dict, model_info: dict | None, ) -> None: - """Register per-request custom pricing in litellm.model_cost. - - Router-originated requests (identified by the deployment id the router puts - in metadata) get their full pricing registered under that unique id only; - the shared ``{provider}/{model}`` key receives the entry with pricing fields - stripped, mirroring Router._create_deployment. This keeps one deployment's - pricing overrides (e.g. a zero-cost wildcard) from clobbering built-in - pricing used by sibling deployments of the same backend model. Direct SDK - calls keep the legacy behavior of registering the shared key with pricing. - Either way the shared key's ``mode`` follows ``keep_registered_mode``, as on - deployment registration, since a call routes by its own ``model_info``. - """ - entry: Final = _build_custom_pricing_entry( - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, - model_info=model_info, - ) - shared_key: Final = f"{custom_llm_provider}/{model}" - deployment_id: Final = _get_router_deployment_id(kwargs) - if deployment_id is None: - litellm.register_model({shared_key: keep_registered_mode(shared_key, entry)}, persist_across_reloads=False) + if not isinstance(logging_obj, LiteLLMLoggingObj): return - shared_entry: Final = keep_registered_mode( - shared_key, CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry) - ) - litellm.register_model( - {deployment_id: entry}, - persist_across_reloads=False, - warning_display_name=shared_key, - custom_llm_provider=custom_llm_provider, - add_to_provider_model_sets=False, - ) - litellm.register_model( - {shared_key: shared_entry}, - persist_across_reloads=False, - warning_display_name=shared_key, - custom_llm_provider=custom_llm_provider, - add_to_provider_model_sets=False, + logging_obj.set_request_pricing( + model, + custom_llm_provider, + _build_custom_pricing_entry(custom_llm_provider, kwargs, model_info), + _get_router_deployment_id(kwargs), ) @@ -5973,14 +5942,14 @@ def completion( "aws_region_name", None ) # support region-based pricing for bedrock - ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### if ( (input_cost_per_token is not None and output_cost_per_token is not None) or input_cost_per_second is not None or output_cost_per_second is not None or cost_per_second is not None ): - _register_custom_pricing_for_request( + _set_request_pricing( + logging_obj=litellm_logging_obj, model=model, custom_llm_provider=custom_llm_provider, kwargs=kwargs, @@ -7044,14 +7013,14 @@ def embedding( **non_default_params, ) - ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### if ( (input_cost_per_token is not None and output_cost_per_token is not None) or input_cost_per_second is not None or output_cost_per_second is not None or cost_per_second is not None ): - _register_custom_pricing_for_request( + _set_request_pricing( + logging_obj=litellm_logging_obj, model=model, custom_llm_provider=custom_llm_provider, kwargs=kwargs, diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py index 37417244457..b479cea6a0d 100644 --- a/litellm/proxy/db/accounting_ledger.py +++ b/litellm/proxy/db/accounting_ledger.py @@ -161,6 +161,24 @@ _LOAD_COST: Final = """ SELECT payload::text AS payload, schema_version::text AS schema_version, payload_hash FROM "LiteLLM_AccountingEvent" WHERE source_event_id = $1 AND kind = 'cost' """ +_PENDING_COMPILATION: Final = """ +SELECT payload::text AS payload FROM "LiteLLM_AccountingEvent" +WHERE kind = 'outbound_usage' AND compilation_state = 'pending' +ORDER BY recorded_at, event_id LIMIT $1::int +""" +_COMPILATION_RESULT: Final = """ +UPDATE "LiteLLM_AccountingEvent" SET compilation_state = $2, compilation_error = $3 +WHERE event_id = $1 AND kind = 'outbound_usage' +""" + + +async def _unpriced(db: Prisma, event_id: str, reason: str) -> Unpriced: + await db.execute_raw( + _COMPILATION_RESULT, event_id, "excluded" if reason == "not_charge_owner" else "unpriced", reason + ) + return Unpriced(usage_event_id=event_id, reason=reason) + + _COMPLETED_FAMILY: Final = """ root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage') AND root.recorded_at < $1::timestamp @@ -338,6 +356,15 @@ class AccountingLedger: case _: assert_never(result) + async def pending_compilation(self, limit: int = 100) -> tuple[OutboundUsageEvent, ...]: + if not 1 <= limit <= 1000: + raise ValueError("Compilation batch size must be between 1 and 1000") + async with self._transaction() as db: + rows: Final = _STRINGS.validate_python(await db.query_raw(_PENDING_COMPILATION, limit)) + return tuple( + OutboundUsageEvent.model_validate_json(_json(_JSON.validate_json(row["payload"])["event"])) for row in rows + ) + async def compile_cost( self, event: OutboundUsageEvent, @@ -352,14 +379,14 @@ class AccountingLedger: return Unpriced(usage_event_id=validated.event_id, reason="usage_not_committed") source: Final = sources[0] if source["schema_version"] != "1": - return Unpriced(usage_event_id=validated.event_id, reason="unsupported_usage_schema") + return await _unpriced(db, validated.event_id, "unsupported_usage_schema") payload: Final = _JSON.validate_json(source["payload"]) if payload.get("event") != validated.model_dump(mode="json"): raise AccountingConflictError("Compilation input differs from the accepted usage event") costs: Final = _STRINGS.validate_python(await db.query_raw(_LOAD_COST, validated.event_id)) if costs: if costs[0]["schema_version"] != "1": - return Unpriced(usage_event_id=validated.event_id, reason="unsupported_cost_schema") + return await _unpriced(db, validated.event_id, "unsupported_cost_schema") existing: Final = CostEvent.model_validate_json( _json(_JSON.validate_json(costs[0]["payload"])["event"]) ) @@ -368,12 +395,16 @@ class AccountingLedger: replay: Final = _prepare(existing, targets, costs[0]["payload"], 1, 1) if replay.payload_hash != costs[0]["payload_hash"]: raise AccountingConflictError("Compilation delivery set differs from the accepted cost") + await db.execute_raw(_COMPILATION_RESULT, validated.event_id, "compiled", None) return existing - priced: Final = price(source["payload"]) + try: + priced: Final = price(source["payload"]) + except ValueError: + return await _unpriced(db, validated.event_id, "invalid_pricing_inputs") if isinstance(priced, Unpriced): if priced.usage_event_id != validated.event_id: raise AccountingConflictError("Unpriced result does not match its source usage") - return priced + return await _unpriced(db, validated.event_id, priced.reason) cost: Final = CostEvent( event_id=f"cost:{validated.event_id}", usage_event_id=validated.event_id, @@ -382,10 +413,13 @@ class AccountingLedger: charges=(Metric(name=MetricName.SPEND, value=priced.total_cost, unit="USD"),), snapshot=priced, ) - record: Final = _prepare(cost, targets, _json({"event": cost.model_dump(mode="json")}), 1, 1) + record: Final = _prepare( + cost, targets, _json({"event": cost.model_dump(mode="json"), "pricing_inputs": payload}), 1, 1 + ) accepted: Final = await _accept(db, record) if accepted not in ("inserted", "duplicate"): raise AccountingConflictError(f"Cannot commit compiled cost: {accepted}") + await db.execute_raw(_COMPILATION_RESULT, validated.event_id, "compiled", None) return cost async def purge_completed( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b5f4a36f138..7980069b0ff 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3228,11 +3228,15 @@ def _configure_usage_capture() -> None: "Usage capture requires local Prisma writes and disabled collector, remote writer and search free tier" ) from litellm.proxy.db.accounting_ledger import AccountingLedger + from litellm.proxy.spend_tracking.cost_compiler import LoggingCostCompiler ledger: Final = AccountingLedger.for_client(prisma_client) writer: Final = UsageBatchWriter(ledger.insert_usage_batch, ledger.load_usage_batch) usage_capture.start( - writer.write, UsageCaptureSettings.model_validate(capture_settings), load=ledger.load_usage_batch + writer.write, + UsageCaptureSettings.model_validate(capture_settings), + load=ledger.load_usage_batch, + compiler=LoggingCostCompiler(ledger), ) @@ -11591,6 +11595,22 @@ class ProxyStartupEvent: await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) + async def compile_accounting_costs() -> None: + from litellm.proxy.db.accounting_ledger import AccountingLedger + from litellm.proxy.spend_tracking.cost_compiler import LoggingCostCompiler + + if usage_capture.meter is not None: + await LoggingCostCompiler(AccountingLedger.for_client(prisma_client)).compile_pending() + + scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # APScheduler has no typed add_job interface + compile_accounting_costs, + "interval", + seconds=5, + id="accounting_cost_compilation", + replace_existing=True, + max_instances=1, + ) + cls._initialize_spend_capture_rate_check_job( scheduler=scheduler, proxy_logging_obj=proxy_logging_obj, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 179b741fe7c..048e0e1e17a 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -2111,9 +2111,12 @@ model LiteLLM_AccountingEvent { occurred_at DateTime recorded_at DateTime @default(now()) payload Json + compilation_state String @default("pending") + compilation_error String? @@index([recorded_at]) @@index([kind, source_event_id]) + @@index([kind, compilation_state, recorded_at]) } // One row per (event, target table, version): pending work for the projector and, diff --git a/litellm/proxy/spend_tracking/cost_compiler.py b/litellm/proxy/spend_tracking/cost_compiler.py index 9d2438b93dc..3580d2854e2 100644 --- a/litellm/proxy/spend_tracking/cost_compiler.py +++ b/litellm/proxy/spend_tracking/cost_compiler.py @@ -1,20 +1,334 @@ -from collections.abc import Callable +import json +import math +from collections.abc import Callable, Mapping +from datetime import datetime +from decimal import Decimal from typing import Final -from litellm.accounting import CostEvent, DeliveryTarget, OutboundUsageEvent, PricingSnapshot, Unpriced +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +from litellm._internal_context import pinned_billing_time +from litellm._logging import verbose_proxy_logger +from litellm.accounting import CostEvent, DeliveryTarget, DimensionName, OutboundUsageEvent, PricingSnapshot, Unpriced +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_cost_calc.finalization import finalize_cost +from litellm.litellm_core_utils.llm_cost_calc.pricing_context import RetainedPricing, retained_pricing from litellm.proxy.db.accounting_ledger import AccountingLedger +from litellm.proxy.spend_tracking.usage_meter import usage_delivery_targets +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ( + CallTypes, + CostPerToken, + LiteLLMRealtimeStreamLoggingObject, + ModelResponse, + StandardBuiltInToolsParams, + Usage, + is_server_derived_pricing_key, +) + +_FIELDS: Final = TypeAdapter(dict[str, JsonValue]) + + +def _fields(value: object) -> Mapping[str, JsonValue]: + return _FIELDS.validate_python(value) if isinstance(value, dict) else {} + + +def _json(value: object) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +class _UsageEnvelope(BaseModel): + model_config = ConfigDict(frozen=True) + + event: OutboundUsageEvent + usage: Mapping[str, JsonValue] | None = None + pricing: Mapping[str, JsonValue] + charge_eligible: bool + finalized_cost: Mapping[str, JsonValue] | None = None + allocation: Mapping[str, JsonValue] | None = None + + +def price_usage(envelope: str) -> PricingSnapshot | Unpriced: + captured: Final = _UsageEnvelope.model_validate_json(envelope) + event: Final = captured.event + if not captured.charge_eligible: + return Unpriced(usage_event_id=event.event_id, reason="not_charge_owner") + fixed_call: Final = event.call_type in {CallTypes.call_mcp_tool, CallTypes.send_message, CallTypes.asend_message} + realtime_events: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python( + captured.pricing.get("realtime_events") or () + ) + realtime_usage: Final = event.call_type is CallTypes.arealtime and any( + bool(item.get("usage")) for item in realtime_events + ) + if captured.usage is None and not fixed_call and not realtime_usage: + return Unpriced(usage_event_id=event.event_id, reason="usage_unknown") + pricing: Final = captured.pricing + mapped: Final = _fields(pricing.get("model_map_information")) + rates: Final = _fields(mapped.get("model_map_value")) + model: Final = next( + (d.value for d in event.dimensions if d.name is DimensionName.MODEL), + event.call_type.value if fixed_call else None, + ) + 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") + upstream: Final = ( + isinstance(reported, (int, float)) + and not isinstance(reported, bool) + and math.isfinite(reported) + and reported >= 0 + and logged.get("custom_pricing") is not True + ) + key: Final = mapped.get("model_map_key") or (model if fixed_call or upstream else None) + has_rates: Final = any(is_server_derived_pricing_key(name) and value is not None for name, value in rates.items()) + if model is None or not isinstance(key, str) or (not has_rates and not fixed_call and not upstream): + return Unpriced(usage_event_id=event.event_id, reason="pricing_inputs_missing") + info: Final = { + "key": key, + "litellm_provider": provider, + "mode": "chat", + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + "input_cost_per_token": None, + "output_cost_per_token": None, + "supported_openai_params": None, + **rates, + } + adjustments: Final = _fields(pricing.get("adjustments")) + discount: Final = adjustments.get("discount") + margin: Final = adjustments.get("margin") + secondary: Final = { + name: { + **info, + **_fields(_fields(snapshot).get("model_map_value")), + "key": _fields(snapshot).get("model_map_key") or name, + } + for name, snapshot in _fields(pricing.get("realtime_models")).items() + if _fields(snapshot).get("model_map_value") is not None + } + context: Final = RetainedPricing( + _json( + { + **secondary, + **{ + alias: info + for alias in dict.fromkeys((model, key, f"{provider}/{model}", model.removeprefix(f"{provider}/"))) + }, + } + ), + _json({provider: discount} if provider and discount is not None else {}), + _json({"global": margin} if margin is not None else {}), + ) + response_inputs: Final = _fields(pricing.get("response_inputs")) + tool_response: Final = _fields(pricing.get("tool_response")) + token_response: Final = ModelResponse( + model=model, + choices=[ + { + "message": { + "role": "assistant", + "annotations": tuple( + {"type": annotation} + for flag, annotation in (("web_search", "url_citation"),) + if tool_response.get(flag) + ), + } + } + ], + usage=Usage.model_validate(captured.usage or {}), + service_tier=response_inputs.get("service_tier"), + meta={"billed_units": response_inputs.get("billed_units")}, + hidden_params={ + "region_name": response_inputs.get("region_name"), + "audio_transcription_duration": response_inputs.get("audio_transcription_duration"), + }, + response_ms=logged.get("response_ms"), + ) + response: Final = ( + LiteLLMRealtimeStreamLoggingObject.model_validate( + { + "results": pricing.get("realtime_events", []), + "usage": captured.usage or {}, + "service_tier": response_inputs.get("service_tier"), + } + ) + if event.call_type in {CallTypes.arealtime, CallTypes.aresponses_websocket} + else ResponsesAPIResponse.model_validate( + { + "id": event.event_id, + "created_at": int(event.started_at.timestamp()), + "model": model, + "output": tool_response.get("output", []), + "usage": tool_response.get("usage"), + } + ) + if tool_response.get("responses_api") + else token_response + ) + logger: Final = Logging( + model=model, + messages=[], + stream=False, + call_type=event.call_type.value, + start_time=event.started_at, + litellm_call_id=event.event_id, + function_id=event.attempt_id, + ) + request: Final = _fields(pricing.get("request")) + deployment: Final = _fields(pricing.get("deployment")) + logger.update_environment_variables( # pyright: ignore[reportUnknownMemberType] # legacy Logging environment method accepts validated JSON fields + model=model, + 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(captured.usage or {}), + end_time=event.ended_at, + mcp_tool_call_metadata={ + "name": _fields(pricing.get("tool")).get("name"), + "mcp_server_cost_info": {"default_cost_per_query": _fields(pricing.get("tool")).get("cost_per_query")}, + }, + optional_params=dict(_fields(pricing.get("optional_params"))), + ) + if logged.get("tool_response_cost") is not None: + logger.model_call_details["response_cost"] = logged["tool_response_cost"] + logger.request_pricing = context if logged.get("custom_pricing") is True else None + logger.custom_cost_per_token = TypeAdapter(CostPerToken | None).validate_python(logged.get("custom_cost_per_token")) + logger.custom_cost_per_second = TypeAdapter(float | None).validate_python(logged.get("custom_cost_per_second")) + logger.standard_built_in_tools_params = TypeAdapter(StandardBuiltInToolsParams).validate_python( + { + "web_search_options": adjustments.get("web_search_options"), + "file_search": adjustments.get("file_search") or None, + } + ) + logger.record_provider_reported_cost( + TypeAdapter[float | None](float | None).validate_python(reported 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: + from litellm.constants import REALTIME_METERED_USAGE_KEY + from litellm.types.llms.xai import RealtimeMeteredUsage + + logger.model_call_details[REALTIME_METERED_USAGE_KEY] = TypeAdapter(RealtimeMeteredUsage).validate_python( + captured.usage["metered_usage"] + ) + logger.cost_billing_time = billed_at + with retained_pricing(context), pinned_billing_time(billed_at): + amount: Final = logger.response_cost_calculator(response, ignore_stamped_cost=True) + if amount is None: + 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)) + 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") + finalized: Final = finalize_cost(amount + guardrail, logger.cost_breakdown, allowance) + if retained_final.get("response_cost") is not None and Decimal(str(finalized.amount)) != Decimal( + str(retained_final["response_cost"]) + ): + return Unpriced(usage_event_id=event.event_id, reason="finalized_cost_mismatch") + final_breakdown: Final = dict(finalized.breakdown or {}) + components: Final = { + name: Decimal(str(value)) + for name, value in final_breakdown.items() + if name in _COST_COMPONENTS and isinstance(value, (int, float)) + } + return PricingSnapshot.model_validate( + { + **components, + "calculator_version": "logging-v1", + "model": model, + "provider": provider, + "pricing_source": "upstream_header" + if upstream + else "per_request" + if request + else "deployment" + if deployment + else "cost_map", + "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, + "additional_costs": tuple( + {"name": name, "amount": Decimal(str(value))} + for name, value in { + **_fields(final_breakdown.get("additional_costs")), + **( + {"request": request["additional_response_cost"]} + if request.get("additional_response_cost") + else {} + ), + }.items() + ), + } + ) + + +_COST_COMPONENTS: Final = frozenset( + { + "input_cost", + "output_cost", + "cache_read_cost", + "cache_creation_cost", + "reasoning_cost", + "tool_usage_cost", + "original_cost", + "discount_percent", + "discount_amount", + "margin_percent", + "margin_fixed_amount", + "margin_total_amount", + } +) class LoggingCostCompiler: def __init__( self, ledger: AccountingLedger, - price_usage: Callable[[str], PricingSnapshot | Unpriced], - delivery_targets: tuple[DeliveryTarget, ...], + price_usage: Callable[[str], PricingSnapshot | Unpriced] = price_usage, + delivery_targets: tuple[DeliveryTarget, ...] | None = None, ) -> None: self._ledger: Final = ledger self._price_usage: Final = price_usage self._delivery_targets: Final = delivery_targets async def compile(self, event: OutboundUsageEvent) -> CostEvent | Unpriced: - return await self._ledger.compile_cost(event, self._price_usage, self._delivery_targets) + targets: Final = self._delivery_targets if self._delivery_targets is not None else cost_delivery_targets(event) + return await self._ledger.compile_cost(event, self._price_usage, targets) + + async def compile_pending(self, limit: int = 100) -> None: + for event in await self._ledger.pending_compilation(limit): + try: + await self.compile(event) + except Exception: # noqa: BLE001 # retain pending work for retry without starving other events + verbose_proxy_logger.exception("Cost compilation failed for %s", event.event_id) + + +_SPEND_TARGETS: Final = ( + (DimensionName.KEY, DeliveryTarget.KEY_SPEND), + (DimensionName.USER, DeliveryTarget.USER_SPEND), + (DimensionName.TEAM, DeliveryTarget.TEAM_SPEND), + (DimensionName.ORGANIZATION, DeliveryTarget.ORGANIZATION_SPEND), + (DimensionName.END_USER, DeliveryTarget.END_USER_SPEND), + (DimensionName.AGENT, DeliveryTarget.AGENT_SPEND), + (DimensionName.TAG, DeliveryTarget.TAG_SPEND), + (DimensionName.PROJECT, DeliveryTarget.PROJECT_SPEND), + (DimensionName.MODEL_ACCESS_GROUP, DeliveryTarget.MODEL_ACCESS_GROUP_SPEND), +) + + +def cost_delivery_targets(event: OutboundUsageEvent) -> tuple[DeliveryTarget, ...]: + names: Final = frozenset(d.name for d in event.dimensions) + return ( + *usage_delivery_targets(event.dimensions), + *(target for dimension, target in _SPEND_TARGETS if dimension in names), + *((DeliveryTarget.TEAM_MEMBER_SPEND,) if {DimensionName.TEAM, DimensionName.USER} <= names else ()), + *( + (DeliveryTarget.ORGANIZATION_MEMBER_SPEND,) + if {DimensionName.ORGANIZATION, DimensionName.USER} <= names + else () + ), + ) diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py index 19480a66678..c1d87d752d8 100644 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ b/litellm/proxy/spend_tracking/usage_capture.py @@ -15,6 +15,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.accounting import ( + CostCompiler, Dimension, DimensionName, InboundUsageEvent, @@ -240,6 +241,7 @@ _PRICING_FIELDS: Final = frozenset( "seconds", "search_context_size", "additional_response_cost", + "cost_per_query", } ) @@ -782,7 +784,13 @@ async def capture_proxy_response(logging_obj: object, response: object, *, backg timeout: Final = TypeAdapter(float).validate_python( settings.get("accounting_sync_ack_timeout_seconds", 20.0) ) - await asyncio.wait_for(asyncio.shield(captured.receipt), timeout=timeout) + async with asyncio.timeout(timeout): + await asyncio.shield(captured.receipt) + compiler: Final = usage_capture.compiler + if compiler is not None: + for write in captured.writes: + if isinstance(write.event, OutboundUsageEvent): + await compiler.compile(write.event) except Exception as error: # noqa: BLE001 # the configured response policy owns every capture failure verbose_proxy_logger.exception("Usage capture failed before the response boundary") if sync_ack and settings.get("accounting_sync_ack_failure_policy", "fail_closed") == "fail_closed": @@ -866,6 +874,17 @@ class PricingCapture: def _capture_pricing(logging_obj: Logging, normalized: object, details: Mapping[str, object]) -> PricingCapture: + from litellm.litellm_core_utils.llm_cost_calc.pricing_context import retained_pricing + + if logging_obj.request_pricing is not None: + with retained_pricing(logging_obj.request_pricing): + return _capture_effective_pricing(logging_obj, normalized, details) + return _capture_effective_pricing(logging_obj, normalized, details) + + +def _capture_effective_pricing( + logging_obj: Logging, normalized: object, details: Mapping[str, object] +) -> PricingCapture: import litellm from litellm.litellm_core_utils.cost_discount import resolve_cost_discount from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -900,16 +919,38 @@ def _capture_pricing(logging_obj: Logging, normalized: object, details: Mapping[ if builtins.get("web_search_options") is not None else None ), - "file_search": builtins.get("file_search") is not None, + "file_search": { + key: value + for key, value in _fields(builtins.get("file_search")).items() + if key in {"type", "storage_gb", "days"} + } + if builtins.get("file_search") is not None + else None, } ) - batch_usage: Final = _fields(_fields(_response_fields(normalized).get("usage")).get("batch_usage_by_model")) + realtime_models: Final = tuple( + dict.fromkeys( + chain.from_iterable(_realtime_pricing_models(event) for event in _realtime_cost_events(normalized)) + ) + ) + batch_usage: Final = _fields( + _response_fields(_response_fields(normalized).get("usage")).get("batch_usage_by_model") + ) return PricingCapture( _json(pricing), adjustments, _json( { **_pricing_inputs(details, logging_obj, normalized), + "realtime_models": { + model: StandardLoggingPayloadSetup.get_model_cost_information( + base_model=model, + custom_pricing=False, + custom_llm_provider=provider, + init_response_obj={"model": model}, + ) + for model in realtime_models + }, "batch_models": { model: StandardLoggingPayloadSetup.get_model_cost_information( base_model=model, @@ -926,10 +967,71 @@ def _capture_pricing(logging_obj: Logging, normalized: object, details: Mapping[ ) +def _tool_cost_inputs(response: object) -> Mapping[str, object]: + from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import StandardBuiltInToolCostTracking + from litellm.types.llms.openai import ResponsesAPIResponse + + raw: Final = _response_fields(response) + flags: Final = { + "web_search": StandardBuiltInToolCostTracking.response_object_includes_web_search_call(response), + "file_search": StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response), + } + if not isinstance(response, ResponsesAPIResponse) and raw.get("object") != "response": + return flags if any(flags.values()) else {} + return { + **flags, + "responses_api": True, + "usage": _response_fields(raw.get("usage")), + "output": tuple( + {key: value for key, value in _fields(item).items() if key in {"type", "id", "status", "usage"}} + for item in TypeAdapter(tuple[object, ...]).validate_python(raw.get("output") or ()) + ), + } + + +def _realtime_cost_fields(value: object) -> Mapping[str, JsonValue]: + fields: Final = _fields(value) + return { + **{ + key: _JSON_OBJECT.validate_python(fields[key]) for key in ("usage",) if isinstance(fields.get(key), Mapping) + }, + **{ + key: value + for key, value in fields.items() + if key in {"type", "id", "model", "service_tier"} and isinstance(value, str) + }, + **{ + key: dict(_realtime_cost_fields(fields[key])) + for key in ("response", "event", "session", "audio", "input", "transcription", "input_audio_transcription") + if isinstance(fields.get(key), Mapping) + }, + } + + +def _realtime_cost_events(response: object) -> tuple[Mapping[str, JsonValue], ...]: + results: Final = _response_fields(response).get("results") + return ( + tuple(_realtime_cost_fields(result) for result in TypeAdapter(tuple[object, ...]).validate_python(results)) + if isinstance(results, list) + else () + ) + + +def _realtime_pricing_models(value: Mapping[str, JsonValue]) -> Iterator[str]: + model: Final = value.get("model") + if isinstance(model, str): + yield model + for key in ("response", "event", "session", "audio", "input", "transcription", "input_audio_transcription"): + if isinstance(nested := value.get(key), dict): + yield from _realtime_pricing_models(nested) + + def _pricing_inputs(kwargs: Mapping[str, object], logging_obj: object, response: object) -> Mapping[str, object]: from litellm.litellm_core_utils.litellm_logging import Logging return { + "tool_response": _tool_cost_inputs(response), + "realtime_events": _realtime_cost_events(response), "request": _selected_pricing(kwargs), "deployment": _selected_pricing(_fields(kwargs.get("litellm_params"))), "optional_params": _selected_pricing(_fields(kwargs.get("optional_params"))), @@ -1006,6 +1108,7 @@ class UsageCaptureRuntime: self._batches: Final[set[asyncio.Task[None]]] = set() # mutable-ok: owned producer lifecycle self._stopping = False self.meter: DurableUsageMeter | None = None + self.compiler: CostCompiler | None = None self.load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None = None def start( @@ -1014,6 +1117,7 @@ class UsageCaptureRuntime: settings: UsageCaptureSettings, *, load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None = None, + compiler: CostCompiler | None = None, ) -> None: from litellm.litellm_core_utils.litellm_logging import Logging @@ -1024,6 +1128,7 @@ class UsageCaptureRuntime: meter: Final = DurableUsageMeter(write_batch, settings) self.meter = meter self.load = load + self.compiler = compiler self._stopping = False Logging.usage_capture_callback = lambda logger, response, start, end, freeze: capture_logging_usage( meter, logger, response, start, end, freeze @@ -1082,3 +1187,4 @@ class UsageCaptureRuntime: Logging.usage_capture_callback = None self.meter = None self.load = None + self.compiler = None diff --git a/litellm/utils.py b/litellm/utils.py index b0826f84006..5fb27ad6ab6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7531,6 +7531,11 @@ def get_model_info( "supported_openai_params": ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"] } """ + from litellm.litellm_core_utils.llm_cost_calc.pricing_context import retained_model_info + + retained: Final = retained_model_info(model, custom_llm_provider) + if retained is not None: + return retained # api_key is a per-caller credential, not part of the model identity, so it is # kept out of the cache key; explicit keys are resolved without the cache. if api_key is not None: diff --git a/schema.prisma b/schema.prisma index 179b741fe7c..048e0e1e17a 100644 --- a/schema.prisma +++ b/schema.prisma @@ -2111,9 +2111,12 @@ model LiteLLM_AccountingEvent { occurred_at DateTime recorded_at DateTime @default(now()) payload Json + compilation_state String @default("pending") + compilation_error String? @@index([recorded_at]) @@index([kind, source_event_id]) + @@index([kind, compilation_state, recorded_at]) } // One row per (event, target table, version): pending work for the projector and, diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py index e107e8e130c..f783a128bc7 100644 --- a/tests/integration/database/test_accounting_ledger.py +++ b/tests/integration/database/test_accounting_ledger.py @@ -35,7 +35,7 @@ from litellm.accounting import ( ) from litellm.proxy.db import accounting_ledger from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger, UsageWrite -from litellm.proxy.spend_tracking.cost_compiler import LoggingCostCompiler +from litellm.proxy.spend_tracking.cost_compiler import LoggingCostCompiler, price_usage from litellm.types.utils import CallTypes from tests.integration._support.database import scratch_database from tests.integration._support.prisma_clients import connected_proxy_prisma @@ -268,7 +268,7 @@ async def test_compilers_share_one_committed_charge_and_preserve_delivery_receip assert all(result == results[0] for result in results) assert isinstance(results[0], CostEvent) assert results[0].snapshot.total_cost == Decimal("0.125") - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(1,)] + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(1,)] deliveries: Final = _rows( ledger_database, 'SELECT target::text, state::text FROM "LiteLLM_AccountingDelivery" WHERE event_id = %s', @@ -315,7 +315,7 @@ async def test_compilation_preserves_unknown_cost_without_creating_a_zero_charge assert await compiler.compile(event) == Unpriced(usage_event_id=event.event_id, reason="usage_not_committed") await _insert(store, event) assert await compiler.compile(event) == unresolved - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(0,)] + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(0,)] @pytest.mark.asyncio @@ -329,7 +329,7 @@ async def test_compilation_rejects_an_unpriced_result_for_another_source(writer: ) with pytest.raises(AccountingConflictError, match="Unpriced result"): await compiler.compile(event) - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(0,)] + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(0,)] @pytest.mark.asyncio @@ -341,7 +341,7 @@ async def test_compilation_rejects_a_cost_identity_collision(writer: _Writer, le await _insert(store, _usage(f"cost:{event.event_id}", inbound=True)) with pytest.raises(AccountingConflictError, match="identity_conflict"): await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS).compile(event) - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(0,)] + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(0,)] @pytest.mark.asyncio @@ -365,11 +365,11 @@ async def test_compilation_failure_rolls_back_and_can_be_retried(writer: _Writer await _insert(store, event) def failed_pricing(envelope: str) -> PricingSnapshot: - raise ValueError("malformed retained pricing") + raise RuntimeError("pricing temporarily unavailable") - with pytest.raises(ValueError, match="malformed retained pricing"): + with pytest.raises(RuntimeError, match="pricing temporarily unavailable"): await LoggingCostCompiler(store, failed_pricing, _TARGETS).compile(event) - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(0,)] + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(0,)] result: Final = await LoggingCostCompiler(store, lambda _: _cost(amount="0").snapshot, _TARGETS).compile(event) assert isinstance(result, CostEvent) assert result.charges[0].value == 0 @@ -414,7 +414,7 @@ async def test_compiler_recovery_preserves_atomic_charge_and_delivery_acceptance with pytest.raises(error): await compiler.compile(event) count: Final = int(failure == "after_commit") - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(count,)] + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(count,)] assert _rows( ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery" WHERE event_id = %s', @@ -423,7 +423,7 @@ async def test_compiler_recovery_preserves_atomic_charge_and_delivery_acceptance result: Final = await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS).compile(event) assert isinstance(result, CostEvent) assert result.snapshot == _cost().snapshot - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(1,)] + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(1,)] @pytest.mark.asyncio @@ -786,3 +786,193 @@ async def test_invalid_jsonb_submission_cannot_reject_another_producers_commit( assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("valid",)] assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] assert (meter.pending_bytes, meter.pending_records) == (0, 0) + + +@pytest.mark.asyncio +async def test_compilation_persists_invalid_pricing_reason_without_charging( + writer: _Writer, ledger_database: str +) -> None: + _, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + await _insert(store, event) + accepted: Final = _rows( + ledger_database, + 'SELECT payload, payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = %s', + (event.event_id,), + ) + + def malformed(envelope: str) -> PricingSnapshot: + raise ValueError("invalid private input must not appear in durable diagnostics") + + result: Final = await LoggingCostCompiler(store, malformed, _TARGETS).compile(event) + assert result == Unpriced(usage_event_id=event.event_id, reason="invalid_pricing_inputs") + assert _rows( + ledger_database, + 'SELECT compilation_state, compilation_error FROM "LiteLLM_AccountingEvent" WHERE event_id = %s', + (event.event_id,), + ) == [("unpriced", "invalid_pricing_inputs")] + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(0,)] + assert ( + _rows( + ledger_database, + 'SELECT payload, payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = %s', + (event.event_id,), + ) + == accepted + ) + async with connected_proxy_prisma(ledger_database) as restarted: + assert await AccountingLedger.for_client(restarted).pending_compilation() == () + + +@pytest.mark.asyncio +async def test_pending_compilation_only_returns_unresolved_billable_work(writer: _Writer, ledger_database: str) -> None: + _, store = writer + events: Final = tuple(_usage(name) for name in ("compiled", "excluded", "unknown", "pending")) + assert all(isinstance(event, OutboundUsageEvent) for event in events) + for event in events: + await _insert(store, event) + await _insert(store, _usage("inbound", inbound=True)) + compiled, excluded, unknown, pending = events + assert isinstance(compiled, OutboundUsageEvent) + assert isinstance(excluded, OutboundUsageEvent) + assert isinstance(unknown, OutboundUsageEvent) + assert isinstance(pending, OutboundUsageEvent) + assert isinstance( + await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS).compile(compiled), CostEvent + ) + await LoggingCostCompiler( + store, lambda _: Unpriced(usage_event_id=excluded.event_id, reason="not_charge_owner"), _TARGETS + ).compile(excluded) + await LoggingCostCompiler( + store, lambda _: Unpriced(usage_event_id=unknown.event_id, reason="usage_unknown"), _TARGETS + ).compile(unknown) + assert await store.pending_compilation(limit=1) == (pending,) + assert _rows( + ledger_database, + "SELECT event_id, compilation_state FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'outbound_usage' ORDER BY event_id", + ) == [("compiled", "compiled"), ("excluded", "excluded"), ("pending", "pending"), ("unknown", "unpriced")] + + +@pytest.mark.asyncio +async def test_restarted_compiler_recovers_transient_failures_without_revisiting_terminal_work( + writer: _Writer, ledger_database: str +) -> None: + _, store = writer + delayed: Final = _usage("delayed") + healthy: Final = _usage("healthy") + assert isinstance(delayed, OutboundUsageEvent) + assert isinstance(healthy, OutboundUsageEvent) + await _insert(store, delayed) + await _insert(store, healthy) + + def transient(envelope: str) -> PricingSnapshot: + if OutboundUsageEvent.model_validate(json.loads(envelope)["event"]).event_id == delayed.event_id: + raise RuntimeError("temporary pricing dependency failure") + return _cost().snapshot + + await LoggingCostCompiler(store, transient, _TARGETS).compile_pending() + assert await store.pending_compilation() == (delayed,) + assert _rows(ledger_database, "SELECT source_event_id FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [ + (healthy.event_id,) + ] + async with connected_proxy_prisma(ledger_database) as restarted: + recovered: Final = AccountingLedger.for_client(restarted) + await LoggingCostCompiler(recovered, lambda _: _cost().snapshot, _TARGETS).compile_pending() + assert await recovered.pending_compilation() == () + await LoggingCostCompiler(recovered, transient, _TARGETS).compile_pending() + assert _rows( + ledger_database, + "SELECT source_event_id FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost' ORDER BY source_event_id", + ) == [(delayed.event_id,), (healthy.event_id,)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_ack", (False, True)) +async def test_proxy_capture_and_background_compiler_converge_on_final_header_charge( + writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch, sync_ack: bool +) -> None: + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.spend_tracking.cost_finalization import finalize_response_cost + from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, UsageCaptureRuntime, capture_proxy_response + from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings + from litellm.types.utils import ModelResponse, Usage + + _, store = writer + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + batches: Final = UsageBatchWriter(store.insert_usage_batch, store.load_usage_batch) + + async def delayed_write(writes: tuple[UsageWrite, ...]) -> None: + entered.set() + await release.wait() + await batches.write(writes) + + compiler: Final = LoggingCostCompiler(store) + runtime: Final = UsageCaptureRuntime() + runtime.start(delayed_write, UsageCaptureSettings(batch_size=2), load=store.load_usage_batch, compiler=compiler) + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": sync_ack}) + logger: Final = Logging( + model="capture-priced-model", + messages=[], + stream=False, + call_type="acompletion", + start_time=_AT, + litellm_call_id="gateway-request", + function_id="gateway-attempt", + ) + logger.update_environment_variables( + model="capture-priced-model", custom_llm_provider="openai", litellm_params={}, optional_params={} + ) + logger.set_request_pricing( + "capture-priced-model", "openai", {"input_cost_per_token": 0.01, "output_cost_per_token": 0.02} + ) + response: Final = ModelResponse( + id="gateway-provider", model="capture-priced-model", usage=Usage(prompt_tokens=10, completion_tokens=4) + ) + try: + finalized: Final = await finalize_response_cost(logger, response, None, _AT, guardrail_cost=0.03) + assert finalized is not None + headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=UserAPIKeyAuth(), response_cost=finalized.amount, litellm_logging_obj=logger + ) + capture: Final = asyncio.create_task(capture_proxy_response(logger, response)) + await asyncio.wait_for(entered.wait(), timeout=5) + if sync_ack: + assert not capture.done() + else: + await asyncio.wait_for(asyncio.shield(capture), timeout=5) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] + release.set() + await asyncio.wait_for(capture, timeout=10) + assert runtime.meter is not None + await runtime.meter.flush() + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [ + (1 if sync_ack else 0,) + ] + await asyncio.gather(*(compiler.compile_pending() for _ in range(3))) + costs: Final = _rows( + ledger_database, "SELECT payload::text FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'" + ) + assert len(costs) == 1 + cost_envelope: Final = json.loads(str(costs[0][0])) + committed: Final = CostEvent.model_validate_json(json.dumps(cost_envelope["event"])) + assert committed.snapshot.total_cost == Decimal(headers["x-litellm-response-cost"]) + assert float(committed.snapshot.total_cost) == pytest.approx(0.21) + monkeypatch.setattr(litellm, "model_cost", {"openai/capture-priced-model": {"input_cost_per_token": 99.0}}) + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.99}) + monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"fixed_amount": 999.0}}) + replayed: Final = price_usage(json.dumps(cost_envelope["pricing_inputs"])) + assert replayed == committed.snapshot + assert litellm.model_cost["openai/capture-priced-model"]["input_cost_per_token"] == 99.0 + await capture_proxy_response(logger, response) + await compiler.compile_pending() + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'") == [(1,)] + finally: + release.set() + await runtime.stop(5) diff --git a/tests/unit/proxy/spend_tracking/test_cost_compiler.py b/tests/unit/proxy/spend_tracking/test_cost_compiler.py new file mode 100644 index 00000000000..bd093d6e012 --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_cost_compiler.py @@ -0,0 +1,849 @@ +import json +from datetime import datetime, timezone +from decimal import Decimal +from typing import Final + +import pytest + +from litellm.accounting import Dimension, DimensionName, OutboundUsageEvent, Outcome, PricingSnapshot, Unpriced +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy.spend_tracking.cost_compiler import price_usage +from litellm.types.utils import CallTypes + +_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +def _envelope() -> dict[str, object]: + event: Final = OutboundUsageEvent( + event_id="attempt", + attempt_id="attempt", + request_id="request", + retry_of_attempt_id=None, + call_type=CallTypes.acompletion, + started_at=_AT, + ended_at=_AT, + outcome=Outcome.SUCCEEDED, + dimensions=( + Dimension(name=DimensionName.MODEL, value="openai/retained-model"), + Dimension(name=DimensionName.PROVIDER, value="openai"), + ), + metrics=None, + ) + return { + "event": event.model_dump(mode="json"), + "charge_eligible": True, + "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, + "pricing": { + "model_map_information": { + "model_map_key": "openai/retained-model", + "model_map_value": { + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + "litellm_provider": "openai", + "mode": "chat", + }, + }, + "adjustments": {"discount": 0.1, "margin": {"fixed_amount": 0.008}}, + }, + } + + +def test_compiler_prices_retained_rates_without_registering_a_model() -> None: + import litellm + + result: Final = price_usage(json.dumps(_envelope())) + assert isinstance(result, PricingSnapshot) + assert result.total_cost == Decimal(str((10 * 0.01 + 4 * 0.02) * 0.9 + 0.008)) + assert "openai/retained-model" not in litellm.model_cost + + +@pytest.mark.parametrize( + ("field", "value", "reason"), + [ + ("usage", None, "usage_unknown"), + ("charge_eligible", False, "not_charge_owner"), + ("pricing", {}, "pricing_inputs_missing"), + ("finalized_cost", {"response_cost": 999}, "finalized_cost_mismatch"), + ], +) +def test_compiler_does_not_invent_a_charge(field: str, value: object, reason: str) -> None: + result: Final = price_usage(json.dumps({**_envelope(), field: value})) + assert result == Unpriced(usage_event_id="attempt", reason=reason) + + +def _retained_envelope( + *, + rates: dict[str, object] | None = None, + usage: dict[str, object] | None = None, + logging: dict[str, object] | None = None, + response_inputs: dict[str, object] | None = None, + deployment: dict[str, object] | None = None, + request: dict[str, object] | None = None, +) -> dict[str, object]: + from pydantic import TypeAdapter + + original: Final = _envelope() + pricing: Final = TypeAdapter(dict[str, object]).validate_python(original["pricing"]) + mapped: Final = TypeAdapter(dict[str, object]).validate_python(pricing["model_map_information"]) + base_rates: Final = TypeAdapter(dict[str, object]).validate_python(mapped["model_map_value"]) + return { + **original, + "usage": usage if usage is not None else original["usage"], + "pricing": { + **pricing, + "model_map_information": {**mapped, "model_map_value": {**base_rates, **(rates or {})}}, + "logging": logging or {}, + "response_inputs": response_inputs or {}, + "deployment": deployment or {}, + "request": request or {}, + "adjustments": {"discount": None, "margin": None}, + }, + } + + +def test_retained_pricing_ignores_later_global_price_and_adjustment_changes(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + + captured: Final = json.dumps(_envelope()) + original: Final = price_usage(captured) + monkeypatch.setattr(litellm, "model_cost", {"openai/retained-model": {"input_cost_per_token": 99.0}}) + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.99}) + monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"fixed_amount": 999.0}}) + recomputed: Final = price_usage(captured) + assert isinstance(original, PricingSnapshot) + assert recomputed == original + assert litellm.model_cost["openai/retained-model"]["input_cost_per_token"] == 99.0 + + +def test_repricing_recomputes_retained_usage_instead_of_returning_a_reported_total() -> None: + from pydantic import TypeAdapter + + captured: Final = _retained_envelope() + pricing: Final = TypeAdapter(dict[str, object]).validate_python(captured["pricing"]) + original: Final = price_usage(json.dumps({**captured, "pricing": {**pricing, "reported_cost": 999.0}})) + changed: Final = price_usage(json.dumps({**captured, "usage": {"prompt_tokens": 20, "completion_tokens": 8}})) + assert isinstance(original, PricingSnapshot) + assert isinstance(changed, PricingSnapshot) + assert float(original.total_cost) == pytest.approx(0.18) + assert float(changed.total_cost) == pytest.approx(float(original.total_cost) * 2) + + +def test_retained_token_pricing_preserves_cache_reasoning_and_service_tier() -> None: + captured: Final = _retained_envelope( + rates={ + "input_cost_per_token_priority": 0.03, + "output_cost_per_token_priority": 0.05, + "cache_read_input_token_cost_priority": 0.002, + "output_cost_per_reasoning_token_priority": 0.07, + }, + usage={ + "prompt_tokens": 100, + "completion_tokens": 20, + "prompt_tokens_details": {"cached_tokens": 40}, + "completion_tokens_details": {"reasoning_tokens": 5}, + }, + response_inputs={"service_tier": "priority"}, + ) + result: Final = price_usage(json.dumps(captured)) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(60 * 0.03 + 40 * 0.002 + 15 * 0.05 + 5 * 0.07) + + +def test_upstream_reported_price_takes_precedence_without_custom_pricing() -> None: + result: Final = price_usage(json.dumps(_retained_envelope(logging={"provider_reported_cost": 0.037}))) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(0.037) + + +def test_per_request_pricing_overrides_upstream_reported_price() -> None: + result: Final = price_usage( + json.dumps( + _retained_envelope( + logging={ + "provider_reported_cost": 0.037, + "custom_pricing": True, + "custom_cost_per_token": {"input_cost_per_token": 0.04, "output_cost_per_token": 0.06}, + } + ) + ) + ) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(10 * 0.04 + 4 * 0.06) + + +def test_retained_regional_pricing_preserves_data_residency_uplift() -> None: + result: Final = price_usage( + json.dumps( + _retained_envelope( + rates={"regional_processing_uplift_multiplier_eu": 1.1}, + deployment={"data_residency": "eu"}, + ) + ) + ) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx((10 * 0.01 + 4 * 0.02) * 1.1) + + +def test_retained_request_additional_cost_is_applied_once() -> None: + result: Final = price_usage(json.dumps(_retained_envelope(request={"additional_response_cost": 0.3}))) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(10 * 0.01 + 4 * 0.02 + 0.3) + + +def _capture_logger(call_type: str = "acompletion", *, stream: bool = False) -> Logging: + from tests.unit.proxy.spend_tracking.test_usage_capture import _logger + + logger: Final = _logger(call_type, stream=stream) + logger.set_request_pricing("capture-model", "openai", {"input_cost_per_token": 0.01, "output_cost_per_token": 0.02}) + return logger + + +@pytest.mark.asyncio +async def test_real_capture_preserves_synchronous_finalization_and_guardrail_parity() -> None: + from litellm.proxy.spend_tracking.cost_finalization import finalize_response_cost + 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() + response: Final = ModelResponse( + id="finalized", model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=4, total_tokens=14) + ) + finalized: Final = await finalize_response_cost(logger, response, None, _AT, guardrail_cost=0.03) + assert finalized is not None + 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 float(result.total_cost) == pytest.approx(finalized.amount) + assert float(result.guardrail_cost or 0) == pytest.approx(0.03) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +async def test_real_stream_capture_only_prices_final_usage() -> 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, ModelResponseStream, 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=True) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, logger, ModelResponseStream(id="chunk"), _AT, _AT) + assert meter.pending_records == 0 + capture_logging_usage( + meter, + logger, + ModelResponse(id="terminal", model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=4)), + _AT, + _AT, + ) + await meter.close() + result: Final = price_usage(database.rows[-1].envelope) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(0.18) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +async def test_real_internal_capture_excludes_usage_owned_by_parent() -> None: + from litellm._internal_context import internal_sub_call + 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 + + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + with internal_sub_call("parent"): + capture_logging_usage( + meter, + _capture_logger(), + ModelResponse(id="child", model="capture-model", usage=Usage(prompt_tokens=10)), + _AT, + _AT, + ) + await meter.close() + assert price_usage(database.rows[-1].envelope) == Unpriced( + usage_event_id=database.rows[-1].event.event_id, reason="not_charge_owner" + ) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("already_calculated", (False, True)) +async def test_real_mcp_capture_prices_retained_tool_fee_without_token_usage(already_calculated: bool) -> None: + from mcp.types import CallToolResult, TextContent + + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + logger: Final = _logger(CallTypes.call_mcp_tool.value) + logger.model_call_details["mcp_tool_call_metadata"] = { + "name": "lookup", + "namespaced_tool_name": "server.lookup", + "mcp_server_cost_info": {"default_cost_per_query": 0.2, "tool_name_to_cost_per_query": {"lookup": 0.3}}, + } + if already_calculated: + logger.model_call_details["response_cost"] = 0.3 + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, logger, CallToolResult(content=[TextContent(type="text", text="ok")]), _AT, _AT) + await meter.close() + result: Final = price_usage(database.rows[-1].envelope) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(0.3) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +async def test_real_agent_capture_prices_retained_per_query_fee() -> 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.agents import LiteLLMSendMessageResponse + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + logger: Final = _logger(CallTypes.asend_message.value) + logger.update_environment_variables( + model="capture-model", custom_llm_provider="a2a", litellm_params={"cost_per_query": 0.4}, optional_params={} + ) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, logger, LiteLLMSendMessageResponse(id="agent"), _AT, _AT) + await meter.close() + result: Final = price_usage(database.rows[-1].envelope) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(0.4) + finally: + await meter.stop(1) + + +@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: + 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 + + logger: Final = _logger() + logger.record_provider_reported_cost(reported_cost) + 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 + ) + 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 result.upstream_reported_cost == Decimal(str(reported_cost)) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancelled", (False, True)) +async def test_real_failure_capture_preserves_incurred_charge(cancelled: bool) -> None: + from fastapi import HTTPException + + from litellm.litellm_core_utils.litellm_logging import UsageObservation + from litellm.proxy.spend_tracking.usage_capture import capture_callback, 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() + usage: Final = Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage( + meter, logger, ModelResponse(model="capture-model", usage=usage), _AT, _AT, freeze_only=True + ) + capture_callback( + meter, + { + **logger.model_call_details, + "litellm_logging_obj": logger, + "combined_usage_object": usage, + "_accounting_usage_observed": UsageObservation.KNOWN, + "log_event_type": "failed_api_call", + }, + HTTPException(499, "Client disconnected") if cancelled else RuntimeError("provider failed after billing"), + _AT, + _AT, + ) + await meter.close() + outbound: Final = database.rows[-1] + assert outbound.event.outcome is (Outcome.CANCELLED if cancelled else Outcome.FAILED) + result: Final = price_usage(outbound.envelope) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(0.04) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +async def test_captured_price_replay_uses_original_off_peak_clock(monkeypatch: pytest.MonkeyPatch) -> None: + from datetime import timedelta + + import litellm + from litellm._internal_context import pinned_billing_time + 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 + + logger: Final = _logger() + monkeypatch.setitem( + litellm.model_cost, + "capture-model", + { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + "off_peak_pricing": { + "hours_utc": "00:00-01:00", + "input_cost_per_token": 0.005, + "output_cost_per_token": 0.01, + }, + }, + ) + response: Final = ModelResponse(model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=4)) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + with pinned_billing_time(_AT): + source_cost: Final = logger.response_cost_calculator(response, ignore_stamped_cost=True) + capture_logging_usage(meter, logger, response, _AT, _AT) + await meter.close() + with pinned_billing_time(_AT + timedelta(hours=3)): + replay: Final = price_usage(database.rows[-1].envelope) + assert source_cost == pytest.approx(0.09) + assert isinstance(replay, PricingSnapshot) + assert float(replay.total_cost) == pytest.approx(source_cost) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +async def test_real_realtime_capture_preserves_audio_and_text_token_prices() -> None: + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + 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(CallTypes.arealtime.value) + logger.set_request_pricing( + "capture-model", + "openai", + { + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + "input_cost_per_audio_token": 0.03, + "output_cost_per_audio_token": 0.04, + }, + ) + events: Final = [ + { + "type": "response.done", + "response": { + "id": "turn", + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "input_token_details": {"text_tokens": 4, "audio_tokens": 6, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 2, "audio_tokens": 3}, + }, + "output": "private transcript excluded from accounting", + }, + } + ] + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, logger, events, _AT, _AT) + await meter.close() + captured: Final = database.rows[-1].envelope + assert "private transcript" not in captured + result: Final = price_usage(captured) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(4 * 0.01 + 6 * 0.03 + 2 * 0.02 + 3 * 0.04) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +async def test_real_realtime_capture_preserves_metered_audio_and_message_fees() -> None: + from litellm.constants import REALTIME_METERED_USAGE_KEY + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from litellm.types.llms.xai import RealtimeMeteredUsage + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + logger: Final = _logger(CallTypes.arealtime.value) + logger.update_environment_variables( + model="capture-model", custom_llm_provider="xai", litellm_params={}, optional_params={} + ) + logger.set_request_pricing( + "capture-model", + "xai", + { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "input_cost_per_second": 0.02, + "input_cost_per_text_message": 0.1, + }, + ) + logger.model_call_details[REALTIME_METERED_USAGE_KEY] = RealtimeMeteredUsage(True, 1.5, 2.5, 3) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage( + meter, + logger, + [ + { + "type": "response.done", + "response": { + "id": "metered", + "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, + }, + } + ], + _AT, + _AT, + ) + await meter.close() + result: Final = price_usage(database.rows[-1].envelope) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx((1.5 + 2.5) * 0.02 + 3 * 0.1) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +async def test_real_batch_rows_keep_per_request_tiers_across_repeated_terminal_retrieval( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm.batches.batch_utils import batch_cost_and_usage + from litellm.proxy.spend_tracking.usage_capture import ( + UsageBatchWriter, + capture_batch_callback, + capture_logging_usage, + ) + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + rates: Final = { + "input_cost_per_token_batches": 0.001, + "input_cost_per_token_above_200k_tokens_batches": 0.002, + "output_cost_per_token_batches": 0.0, + "litellm_provider": "openai", + "mode": "chat", + } + monkeypatch.setitem(litellm.model_cost, "capture-model", rates) + rows: Final = [ + { + "response": { + "status_code": 200, + "body": { + "model": "capture-model", + "usage": {"prompt_tokens": tokens, "completion_tokens": 0, "total_tokens": tokens}, + }, + } + } + for tokens in (10, 200001) + ] + billed: Final = batch_cost_and_usage(rows, "openai") + response: Final = {"id": "batch", "status": "completed", "usage": billed.usage} + database: Final = StoredUsage() + writer: Final = UsageBatchWriter(database.insert, database.load) + meter: Final = DurableUsageMeter(writer.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) + logger: Final = _logger("aretrieve_batch") + logger.model_call_details["response_cost"] = billed.cost + try: + capture_logging_usage(meter, logger, response, _AT, _AT) + await capture_batch_callback(meter, database.load, logger.model_call_details, response, _AT, _AT) + await meter.flush() + first: Final = tuple(database.rows) + await capture_batch_callback(meter, database.load, logger.model_call_details, response, _AT, _AT) + await meter.close() + assert tuple(database.rows) == first + assert price_usage(database.rows[0].envelope) == Unpriced( + usage_event_id=database.rows[0].event.event_id, reason="not_charge_owner" + ) + costs: Final = tuple(price_usage(row.envelope) for row in database.rows[1:]) + assert all(isinstance(cost, PricingSnapshot) for cost in costs), costs + assert sum(float(cost.total_cost) for cost in costs if isinstance(cost, PricingSnapshot)) == pytest.approx( + 10 * 0.001 + 200001 * 0.002 + ) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transcription", (False, True)) +async def test_realtime_capture_replays_sanitized_lifecycle_and_transcription_charges(transcription: bool) -> None: + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + 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(CallTypes.arealtime.value) + logger.set_request_pricing( + "capture-model", + "openai", + { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "input_cost_per_second": 0.02, + }, + ) + events: Final = ( + [ + { + "type": "transcription_session.created", + "session": {"input_audio_transcription": {"model": "capture-model"}}, + }, + { + "type": "conversation.item.input_audio_transcription.completed", + "item_id": "audio", + "transcript": "private transcript", + "usage": {"type": "duration", "seconds": 2.5}, + }, + ] + if transcription + else [ + {"type": "session.started", "session": {"model": "capture-model"}}, + {"type": "session.closed", "usage": {"seconds": 2.5}, "transcript": "private transcript"}, + ] + ) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, logger, events, _AT, _AT) + await meter.close() + captured: Final = database.rows[-1].envelope + assert "private transcript" not in captured + result: Final = price_usage(captured) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(2.5 * 0.02) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("file_search", (False, True)) +async def test_capture_replays_builtin_search_charge_without_retaining_results(file_search: 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.llms.openai import ResponsesAPIResponse + 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("aresponses") + logger.set_request_pricing( + "capture-model", + "openai", + { + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + "search_context_cost_per_query": {"search_context_size_medium": 0.1}, + "file_search_cost_per_1k_calls": 0.025, + }, + ) + logger.standard_built_in_tools_params = ( + {"file_search": {"type": "file_search", "vector_store_ids": ["store"]}} + if file_search + else {"web_search_options": {"search_context_size": "medium"}} + ) + response: Final = ResponsesAPIResponse.model_validate( + { + "id": "search", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "capture-model", + "output": [ + {"type": "file_search_call", "id": "search", "queries": ["private query"], "status": "completed"} + ] + if file_search + else [ + { + "type": "web_search_call", + "id": "search", + "status": "completed", + "action": {"type": "search", "query": "private query"}, + } + ], + "usage": {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}, + } + ) + source_cost: Final = logger.response_cost_calculator(response, ignore_stamped_cost=True) + assert source_cost is not None and source_cost > 0.18 + 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() + captured: Final = database.rows[-1].envelope + assert "private query" not in captured + result: Final = price_usage(captured) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(source_cost) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +async def test_capture_replays_multi_round_aggregate_with_one_fixed_request_margin( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + from litellm.responses.mcp.round_billing import billed_for_every_round, summed_cost_breakdown + 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 + + monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"fixed_amount": 0.1}}) + first_logger: Final = _capture_logger() + final_logger: Final = _capture_logger() + first: Final = ModelResponse(id="first", model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=4)) + final: Final = ModelResponse(id="final", model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=4)) + first._hidden_params["response_cost"] = first_logger.response_cost_calculator(first, ignore_stamped_cost=True) + final._hidden_params["response_cost"] = final_logger.response_cost_calculator(final, ignore_stamped_cost=True) + aggregate: Final = billed_for_every_round(first, first_logger.cost_breakdown, final, final_logger.cost_breakdown) + breakdown: Final = summed_cost_breakdown(first, first_logger.cost_breakdown, final, final_logger.cost_breakdown) + final_logger.cost_breakdown = breakdown + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + try: + capture_logging_usage(meter, final_logger, aggregate, _AT, _AT) + await meter.close() + result: Final = price_usage(database.rows[-1].envelope) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(2 * (10 * 0.01 + 4 * 0.02) + 0.1) + assert float(result.total_cost) == pytest.approx(aggregate._hidden_params["response_cost"]) + finally: + await meter.stop(1) + + +@pytest.mark.parametrize("reported", (-0.1, float("inf"), float("-inf"), float("nan"), True, False)) +@pytest.mark.parametrize("has_rates", (False, True)) +def test_invalid_upstream_price_cannot_become_an_authoritative_charge(reported: float, has_rates: bool) -> None: + from pydantic import TypeAdapter + + envelope: Final = _retained_envelope(logging={"provider_reported_cost": reported}) + pricing: Final = TypeAdapter(dict[str, object]).validate_python(envelope["pricing"]) + if not has_rates: + pricing["model_map_information"] = {} + envelope["pricing"] = pricing + result: Final = price_usage(json.dumps(envelope)) + if has_rates: + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(0.18) + assert result.upstream_reported_cost is None + else: + assert result == Unpriced(usage_event_id="attempt", reason="pricing_inputs_missing") + + +def test_repricing_rejects_unretained_allowance_allocation() -> None: + envelope: Final = _retained_envelope() + envelope["finalized_cost"] = {"allowance_discount": 0.1} + assert price_usage(json.dumps(envelope)) == Unpriced( + usage_event_id="attempt", reason="allowance_allocation_missing" + ) + + +@pytest.mark.asyncio +async def test_retained_pricing_is_task_local_and_resets_on_cancellation() -> None: + import asyncio + + from litellm.litellm_core_utils.llm_cost_calc.pricing_context import ( + RetainedPricing, + pricing_discounts, + retained_model_info, + retained_models, + retained_pricing, + ) + + outer: Final = RetainedPricing("{}", '{"openai":0.1}', "{}") + inner: Final = RetainedPricing("{}", '{"openai":0.9}', "{}") + ready: Final = asyncio.Event() + hold: Final = asyncio.Event() + restored: Final = asyncio.Event() + + async def child() -> None: + try: + with retained_pricing(inner): + assert pricing_discounts()["openai"] == 0.9 + with pytest.raises(ValueError, match="no retained model"): + retained_model_info("gpt-4o-mini", "openai") + ready.set() + await hold.wait() + finally: + assert pricing_discounts()["openai"] == 0.1 + restored.set() + + assert retained_models() is None + with retained_pricing(outer): + task: Final = asyncio.create_task(child()) + try: + await asyncio.wait_for(ready.wait(), timeout=5) + assert pricing_discounts()["openai"] == 0.1 + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert restored.is_set() + assert pricing_discounts()["openai"] == 0.1 + assert retained_models() is None + assert retained_model_info("gpt-4o-mini", "openai") is None + + +def test_incomplete_retained_rates_remain_unpriced() -> None: + result: Final = price_usage( + json.dumps(_retained_envelope(rates={"input_cost_per_token": None, "output_cost_per_token": None})) + ) + assert result == Unpriced(usage_event_id="attempt", reason="pricing_inputs_missing") + + +def test_explicit_zero_retained_token_rates_are_known_free() -> None: + result: Final = price_usage( + json.dumps(_retained_envelope(rates={"input_cost_per_token": 0.0, "output_cost_per_token": 0.0})) + ) + assert isinstance(result, PricingSnapshot) + assert result.total_cost == Decimal("0") + assert result.pricing_source == "cost_map" diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py index 32710e594ee..24bb3480560 100644 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ b/tests/unit/proxy/spend_tracking/test_usage_capture.py @@ -1,13 +1,13 @@ import asyncio import json -from collections.abc import Awaitable, Mapping +from collections.abc import Mapping from datetime import datetime, timezone from decimal import Decimal from typing import Final import pytest -from litellm.accounting import Metric, MetricName, Outcome +from litellm.accounting import Metric, MetricName, OutboundUsageEvent, Outcome, Unpriced from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings @@ -441,13 +441,13 @@ async def test_supported_sync_ack_waits_for_commit(call_type: str, monkeypatch: monkeypatch.setattr(proxy_server, "usage_capture", runtime) monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) deadlines: Final = asyncio.Queue[float | None]() - original_wait: Final = asyncio.wait_for + original_timeout: Final = asyncio.timeout - async def wait_with_deadline(awaitable: Awaitable[object], timeout: float | None) -> object: - deadlines.put_nowait(timeout) - return await original_wait(awaitable, timeout) + def timeout_with_deadline(delay: float | None) -> asyncio.Timeout: + deadlines.put_nowait(delay) + return original_timeout(delay) - monkeypatch.setattr(asyncio, "wait_for", wait_with_deadline) + monkeypatch.setattr(asyncio, "timeout", timeout_with_deadline) response: Final = asyncio.create_task( capture_proxy_response( _logger(call_type), @@ -1075,17 +1075,14 @@ async def test_response_ack_timeout_preserves_success_and_later_commits(monkeypa runtime.start(database.write, UsageCaptureSettings(batch_size=2)) monkeypatch.setattr(proxy_server, "usage_capture", runtime) monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) + monkeypatch.setattr( + proxy_server, "general_settings", {"accounting_sync_ack": True, "accounting_sync_ack_timeout_seconds": 0.01} + ) logger: Final = _logger() - async def expired_deadline(future: object, timeout: float) -> None: - raise TimeoutError("controlled response deadline") - try: - with monkeypatch.context() as clock: - clock.setattr(asyncio, "wait_for", expired_deadline) - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, ModelResponse(id="provider", usage=Usage(prompt_tokens=3))) + with pytest.raises(HTTPException) as error: + await capture_proxy_response(logger, ModelResponse(id="provider", usage=Usage(prompt_tokens=3))) assert error.value.status_code == 503 await database.entered.wait() captured: Final = logger.model_call_details["_accounting_capture"] @@ -1710,3 +1707,50 @@ async def test_batch_replay_reserves_the_accepted_pricing_snapshot_size() -> Non assert runtime.meter is not None and runtime.meter.pending_bytes == 0 finally: await runtime.stop(1) + + +class DelayedCostCompiler: + def __init__(self) -> None: + self.entered: Final = asyncio.Event() + self.cancelled: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def compile(self, event: OutboundUsageEvent) -> Unpriced: + self.entered.set() + try: + await self.release.wait() + except asyncio.CancelledError: + self.cancelled.set() + raise + return Unpriced(usage_event_id=event.event_id, reason="usage_unknown") + + +@pytest.mark.asyncio +async def test_sync_ack_deadline_bounds_compilation_after_usage_commit(monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import HTTPException + + from litellm.proxy import proxy_server + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response + from litellm.types.utils import ModelResponse, Usage + + database: Final = CommitBarrier() + database.release.set() + compiler: Final = DelayedCostCompiler() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.write, UsageCaptureSettings(batch_size=2), compiler=compiler) + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + monkeypatch.setattr( + proxy_server, "general_settings", {"accounting_sync_ack": True, "accounting_sync_ack_timeout_seconds": 0.1} + ) + try: + with pytest.raises(HTTPException) as error: + await asyncio.wait_for( + capture_proxy_response(_logger(), ModelResponse(id="provider", usage=Usage(prompt_tokens=3))), 2 + ) + assert error.value.status_code == 503 + assert compiler.entered.is_set() + assert compiler.cancelled.is_set() + assert database.committed == ("caller:inbound", "provider") + finally: + compiler.release.set() + await runtime.stop(1) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index cca138e0bb5..bb4035e8725 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -11108,3 +11108,50 @@ def test_upstream_cost_header_total_on_an_unmapped_alias_is_split_by_the_served_ input_share: Final = (10 * 0.001) / (10 * 0.001 + 20 * 0.002) assert breakdown["input_cost"] == pytest.approx(0.25 * input_share) assert breakdown["output_cost"] == pytest.approx(0.25 * (1 - input_share)) + + +@pytest.mark.parametrize( + ("model", "retained_key", "deployment_rate"), + ( + ("video-model", "openai/video-model", None), + ("openai/video-model", "video-model", None), + ("openai/video-model", "openai/video-model", 0.7), + ), +) +def test_video_pricing_uses_retained_catalog_and_deployment_precedence( + monkeypatch: pytest.MonkeyPatch, model: str, retained_key: str, deployment_rate: float | None +) -> None: + from pydantic import TypeAdapter + + from litellm.cost_calculator import default_video_cost_calculator + from litellm.litellm_core_utils.llm_cost_calc.pricing_context import RetainedPricing, retained_pricing + from litellm.types.utils import ModelInfo + + info: Final = TypeAdapter(ModelInfo).validate_python( + { + "key": retained_key, + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + "input_cost_per_token": None, + "output_cost_per_token": None, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_openai_params": None, + "output_cost_per_video_per_second": 0.2, + } + ) + retained: Final = RetainedPricing(json.dumps({retained_key: info}), "{}", "{}") + changed: Final = {**info, "output_cost_per_video_per_second": 9.0} + monkeypatch.setattr(litellm, "model_cost", {"video-model": changed, "openai/video-model": changed}) + deployment: Final = ( + TypeAdapter(ModelInfo).validate_python({**info, "output_cost_per_video_per_second": deployment_rate}) + if deployment_rate is not None + else None + ) + with retained_pricing(retained): + price: Final = default_video_cost_calculator( + model, duration_seconds=4, custom_llm_provider="openai", model_info=deployment + ) + assert price == pytest.approx(4 * (deployment_rate if deployment_rate is not None else 0.2)) + assert litellm.model_cost["video-model"]["output_cost_per_video_per_second"] == 9.0 diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 35fbd28eae7..9ee21032343 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -1,13 +1,4 @@ -""" -Test that register_model() in completion() and embedding() passes all -custom pricing fields from kwargs and model_info, not just the base -input/output costs. - -Previously, only input_cost_per_token, output_cost_per_token, and -litellm_provider were forwarded. Fields like cache_read_input_token_cost, -mode, and supports_prompt_caching were dropped, causing incorrect cost -calculations for DB-sourced models with prompt caching pricing. -""" +"""Verify explicit model registration and isolated request pricing preserve effective rates.""" import copy import itertools @@ -17,9 +8,9 @@ from typing import Final, NoReturn import pytest - import litellm from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases, adopt_model_cost_map +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.main import _build_custom_pricing_entry from litellm.utils import _invalidate_model_cost_lowercase_map @@ -236,10 +227,7 @@ def test_register_model_strips_none_litellm_provider(): registered = litellm.model_cost.get(model_key) assert registered is not None, f"{model_key} should be in model_cost" # The key may be absent entirely, but if present it must not be None. - assert ( - "litellm_provider" not in registered - or registered["litellm_provider"] is not None - ) + assert "litellm_provider" not in registered or registered["litellm_provider"] is not None # Downstream consumers must accept this entry for any provider, # mirroring what the cost calculator does. assert _check_provider_match(registered, "openai") is True @@ -297,9 +285,7 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp # The strip must have removed the None-valued provider that # ``get_model_info`` returned. The key may be absent entirely, but # it must never be present with value ``None``. - assert "litellm_provider" not in registered or ( - registered["litellm_provider"] is not None - ), ( + assert "litellm_provider" not in registered or (registered["litellm_provider"] is not None), ( "register_model failed to strip litellm_provider=None returned " f"by get_model_info, got {registered.get('litellm_provider')!r}" ) @@ -360,14 +346,8 @@ def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(monkeypa ) registered = litellm.model_cost[registered_key] - assert ( - registered.get("cache_creation_input_token_cost") - == builtin["cache_creation_input_token_cost"] - ) - assert ( - registered.get("cache_read_input_token_cost") - == builtin["cache_read_input_token_cost"] - ) + assert registered.get("cache_creation_input_token_cost") == builtin["cache_creation_input_token_cost"] + assert registered.get("cache_read_input_token_cost") == builtin["cache_read_input_token_cost"] assert registered["litellm_provider"] == "bedrock" usage = Usage( @@ -431,8 +411,7 @@ def test_register_model_warns_when_no_builtin_match_for_cache_pricing(caplog): ) assert any( - registered_key in record.message - and "cache_creation_input_token_cost" in record.message + registered_key in record.message and "cache_creation_input_token_cost" in record.message for record in caplog.records ), "expected a warning naming the unmapped key and the cache cost fields" finally: @@ -619,17 +598,14 @@ def test_register_model_router_add_deployment_custom_pricing_applies(): # the deployment's provider. from litellm.utils import _check_provider_match - registered_keys = [ - k for k in (deployment_model, model_key) if k in litellm.model_cost - ] + registered_keys = [k for k in (deployment_model, model_key) if k in litellm.model_cost] assert registered_keys, ( - "Router.add_deployment did not register custom pricing for " - f"{model_key} / {deployment_model}" + f"Router.add_deployment did not register custom pricing for {model_key} / {deployment_model}" ) for k in registered_keys: - assert ( - _check_provider_match(litellm.model_cost[k], "openai") is True - ), f"custom pricing for {k} was dropped by _check_provider_match" + assert _check_provider_match(litellm.model_cost[k], "openai") is True, ( + f"custom pricing for {k} was dropped by _check_provider_match" + ) finally: litellm.model_cost.pop(model_key, None) litellm.model_cost.pop(deployment_model, None) @@ -646,17 +622,15 @@ def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing(): """ shared_key = "openai/text-embedding-3-small" deployment_id = "lit3991-wildcard-embed-zero" - snapshot = _snapshot_model_cost_entries( - [shared_key, "text-embedding-3-small", deployment_id] - ) - builtin_input_cost = litellm.get_model_info(model=shared_key)[ - "input_cost_per_token" - ] + snapshot = _snapshot_model_cost_entries([shared_key, "text-embedding-3-small", deployment_id]) + builtin_input_cost = litellm.get_model_info(model=shared_key)["input_cost_per_token"] assert builtin_input_cost > 0 + logger: Final = _request_logger(shared_key, "embedding") try: - litellm.embedding( + response = litellm.embedding( model=shared_key, + litellm_logging_obj=logger, input=["hello"], api_key="fake-key", input_cost_per_token=0.0, @@ -666,12 +640,10 @@ def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing(): mock_response=[0.1, 0.2], ) - assert ( - litellm.get_model_info(model=shared_key)["input_cost_per_token"] - == builtin_input_cost - ), "wildcard deployment's zero pricing leaked into the shared model_cost key" - assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.0 - assert litellm.model_cost[deployment_id]["output_cost_per_token"] == 0.0 + assert litellm.get_model_info(model=shared_key)["input_cost_per_token"] == builtin_input_cost, ( + "wildcard deployment's zero pricing leaked into the shared model_cost key" + ) + assert logger.response_cost_calculator(response, ignore_stamped_cost=True) == 0.0 sibling_response = litellm.embedding( model=shared_key, @@ -679,9 +651,7 @@ def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing(): api_key="fake-key", mock_response=[0.1, 0.2], ) - sibling_cost = litellm.completion_cost( - completion_response=sibling_response, call_type="embedding" - ) + sibling_cost = litellm.completion_cost(completion_response=sibling_response, call_type="embedding") assert sibling_cost == pytest.approx(10 * builtin_input_cost) finally: _restore_model_cost_entries(snapshot) @@ -695,17 +665,15 @@ def test_embedding_router_custom_pricing_costs_request_via_deployment_id(): shared_key = "openai/text-embedding-3-small" deployment_id = "lit3991-wildcard-embed-custom" override_input_cost = 5e-05 - snapshot = _snapshot_model_cost_entries( - [shared_key, "text-embedding-3-small", deployment_id] - ) - builtin_input_cost = litellm.get_model_info(model=shared_key)[ - "input_cost_per_token" - ] + snapshot = _snapshot_model_cost_entries([shared_key, "text-embedding-3-small", deployment_id]) + builtin_input_cost = litellm.get_model_info(model=shared_key)["input_cost_per_token"] assert builtin_input_cost != override_input_cost + logger: Final = _request_logger(shared_key, "embedding") try: response = litellm.embedding( model=shared_key, + litellm_logging_obj=logger, input=["hello"], api_key="fake-key", input_cost_per_token=override_input_cost, @@ -715,19 +683,9 @@ def test_embedding_router_custom_pricing_costs_request_via_deployment_id(): mock_response=[0.1, 0.2], ) - request_cost = litellm.completion_cost( - completion_response=response, - model=shared_key, - custom_llm_provider="openai", - call_type="embedding", - custom_pricing=True, - router_model_id=deployment_id, - ) + request_cost = logger.response_cost_calculator(response, ignore_stamped_cost=True) assert request_cost == pytest.approx(10 * override_input_cost) - assert ( - litellm.get_model_info(model=shared_key)["input_cost_per_token"] - == builtin_input_cost - ) + assert litellm.get_model_info(model=shared_key)["input_cost_per_token"] == builtin_input_cost finally: _restore_model_cost_entries(snapshot) @@ -736,17 +694,15 @@ def test_completion_router_zero_pricing_does_not_clobber_builtin_pricing(): """Same isolation as the embedding path, exercised through completion().""" shared_key = "openai/gpt-4o-mini" deployment_id = "lit3991-wildcard-chat-zero" - snapshot = _snapshot_model_cost_entries( - [shared_key, "gpt-4o-mini", deployment_id] - ) - builtin_input_cost = litellm.get_model_info(model=shared_key)[ - "input_cost_per_token" - ] + snapshot = _snapshot_model_cost_entries([shared_key, "gpt-4o-mini", deployment_id]) + builtin_input_cost = litellm.get_model_info(model=shared_key)["input_cost_per_token"] assert builtin_input_cost > 0 + logger: Final = _request_logger(shared_key, "completion") try: - litellm.completion( + response = litellm.completion( model=shared_key, + litellm_logging_obj=logger, messages=[{"role": "user", "content": "hello"}], api_key="fake-key", input_cost_per_token=0.0, @@ -756,47 +712,43 @@ def test_completion_router_zero_pricing_does_not_clobber_builtin_pricing(): mock_response="hello back", ) - assert ( - litellm.get_model_info(model=shared_key)["input_cost_per_token"] - == builtin_input_cost - ), "wildcard deployment's zero pricing leaked into the shared model_cost key" - assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.0 + assert litellm.get_model_info(model=shared_key)["input_cost_per_token"] == builtin_input_cost, ( + "wildcard deployment's zero pricing leaked into the shared model_cost key" + ) + assert logger.response_cost_calculator(response, ignore_stamped_cost=True) == 0.0 finally: _restore_model_cost_entries(snapshot) -def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key(): - """Direct SDK calls (no router deployment id in metadata) keep the legacy - behavior: custom pricing is registered under ``{provider}/{model}`` and the - request is costed with it. - """ - model_key = "openai/lit3991-direct-sdk-embed-model" - override_input_cost = 3e-05 - try: - response = litellm.embedding( - model=model_key, - input=["hello"], - api_key="fake-key", - input_cost_per_token=override_input_cost, - output_cost_per_token=override_input_cost * 2, - mock_response=[0.1, 0.2], - ) +def _request_logger(model: str, call_type: str) -> Logging: + from datetime import datetime, timezone - assert ( - litellm.model_cost[model_key]["input_cost_per_token"] - == override_input_cost - ) - cost = litellm.completion_cost( - completion_response=response, - model=model_key, - custom_llm_provider="openai", - call_type="embedding", - custom_pricing=True, - ) - assert cost == pytest.approx(10 * override_input_cost) - finally: - litellm.model_cost.pop(model_key, None) - _invalidate_model_cost_lowercase_map() + return Logging( + model=model, + messages=[], + stream=False, + call_type=call_type, + start_time=datetime.now(timezone.utc), + litellm_call_id="override", + function_id="override", + ) + + +def test_embedding_direct_sdk_custom_pricing_is_request_scoped() -> None: + model: Final = "openai/request-priced-embedding" + before: Final = _snapshot_model_cost_entries((model,)) + logger: Final = _request_logger(model, "embedding") + response: Final = litellm.embedding( + model=model, + input=["hello"], + api_key="fake-key", + input_cost_per_token=0.003, + output_cost_per_token=0.006, + mock_response=[0.1, 0.2], + litellm_logging_obj=logger, + ) + assert logger.response_cost_calculator(response, ignore_stamped_cost=True) == pytest.approx(10 * 0.003) + assert _snapshot_model_cost_entries((model,)) == before @pytest.mark.parametrize("route", ["", "converse/", "invoke/"]) @@ -1407,10 +1359,10 @@ def test_update_dictionary_merges_nested_dicts_without_aliasing(): object stays untouched, and the caller's incoming nested dict is never inserted by reference into the merged result. """ - from litellm.utils import update_dictionary - from typing import cast + from litellm.utils import update_dictionary + existing_nested = cast(dict[str, object], {"hours_utc": "01:00-02:00", 1: "existing"}) existing = {"off_peak_pricing": existing_nested} incoming_nested = cast( @@ -1609,22 +1561,24 @@ def test_completion_cost_applies_off_peak_only_deployment_pricing(): del router -def test_completion_registers_cost_per_second_pricing(): - model_key: Final = "openai/test-cost-per-second-registration" - original_entries: Final = _snapshot_model_cost_entries([model_key]) +def test_completion_per_second_pricing_is_request_scoped() -> None: + from litellm.types.utils import ModelResponse, Usage - try: - litellm.completion( - model=model_key, - messages=[{"role": "user", "content": "hello"}], - api_key="fake-key", - cost_per_second=0.02, - mock_response="hello back", - ) - - assert litellm.model_cost[model_key]["cost_per_second"] == 0.02 - finally: - _restore_model_cost_entries(original_entries) + model: Final = "openai/request-priced-seconds" + logger: Final = _request_logger(model, "completion") + before: Final = _snapshot_model_cost_entries((model,)) + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + api_key="fake-key", + cost_per_second=0.02, + mock_response="hello back", + litellm_logging_obj=logger, + ) + response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=0, completion_tokens=0)) + response._response_ms = 2000 + assert logger.response_cost_calculator(response, ignore_stamped_cost=True) == pytest.approx(0.04) + assert _snapshot_model_cost_entries((model,)) == before def test_update_model_cost(): @@ -1676,31 +1630,139 @@ def test_register_model_scopes_builtin_match_to_the_given_provider(): _restore_model_cost_entries(model_cost_entries) -def test_per_request_custom_pricing_scopes_a_colliding_deployment_id_to_its_provider(): - from litellm.main import _register_custom_pricing_for_request +def test_per_request_custom_pricing_preserves_a_colliding_deployment() -> None: + from litellm.main import _set_request_pricing + from litellm.types.utils import ModelResponse, Usage - deployment_id: Final = "baseten/zai-org/glm-5.2" - builtin_key: Final = "baseten/zai-org/GLM-5.2" - shared_key: Final = "openai/zai-org/GLM-5.2" - model_cost_entries: Final = _snapshot_model_cost_entries((deployment_id, builtin_key, shared_key)) - builtin_row_before: Final = copy.deepcopy(litellm.model_cost[builtin_key]) - openai_models_before: Final = frozenset(litellm.open_ai_chat_completion_models) - try: - _register_custom_pricing_for_request( - model="zai-org/GLM-5.2", - custom_llm_provider="openai", - kwargs={ - "input_cost_per_token": 0.00000096, - "output_cost_per_token": 0.00000302, - "metadata": {"model_info": {"id": deployment_id}}, - }, - model_info={"mode": "chat"}, + deployment: Final = "baseten/zai-org/glm-5.2" + model: Final = "zai-org/GLM-5.2" + keys: Final = (deployment, "baseten/zai-org/GLM-5.2", f"openai/{model}") + before: Final = _snapshot_model_cost_entries(keys) + logger: Final = _request_logger(model, "completion") + logger.update_environment_variables( + model=model, user=None, optional_params={}, litellm_params={}, custom_llm_provider="openai" + ) + _set_request_pricing( + logger, + model, + "openai", + {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002, "metadata": {"model_info": {"id": deployment}}}, + {"mode": "chat"}, + ) + response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5)) + assert logger.response_cost_calculator(response, ignore_stamped_cost=True) == pytest.approx(0.02) + assert _snapshot_model_cost_entries(keys) == before + + +@pytest.mark.parametrize("call_type", ("completion", "embedding")) +def test_concurrent_requests_keep_distinct_prices_and_leave_default_rates(call_type: str) -> None: + from concurrent.futures import ThreadPoolExecutor + from threading import Barrier + + model: Final = "openai/gpt-4o-mini" if call_type == "completion" else "openai/text-embedding-3-small" + before: Final = _snapshot_model_cost_entries((model, model.removeprefix("openai/"))) + ready: Final = Barrier(2) + + def execute(logger: Logging, rate: float) -> object: + if call_type == "completion": + return litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + api_key="fake-key", + input_cost_per_token=rate, + output_cost_per_token=rate, + mock_response="hello", + litellm_logging_obj=logger, + ) + return litellm.embedding( + model=model, + input=["hello"], + api_key="fake-key", + input_cost_per_token=rate, + output_cost_per_token=rate, + mock_response=[0.1], + litellm_logging_obj=logger, ) - entry: Final = litellm.model_cost[deployment_id] - assert entry["input_cost_per_token"] == 0.00000096 - assert entry["output_cost_per_token"] == 0.00000302 - assert litellm.model_cost[builtin_key] == builtin_row_before + def price(rate: float) -> float | None: + logger: Final = _request_logger(model, call_type) + response: Final = execute(logger, rate) + ready.wait(timeout=10) + return logger.response_cost_calculator(response, ignore_stamped_cost=True) + + with ThreadPoolExecutor(max_workers=2) as pool: + amounts: Final = tuple(pool.map(price, (0.001, 0.003))) + assert amounts[0] is not None and amounts[0] > 0 + assert amounts[1] == pytest.approx(amounts[0] * 3) + assert _snapshot_model_cost_entries((model, model.removeprefix("openai/"))) == before + + +def test_request_pricing_preserves_reasoning_and_cache_rates() -> None: + from litellm.types.utils import ModelResponse, Usage + + model: Final = "openai/gpt-4o-mini" + logger: Final = _request_logger(model, "completion") + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hello"}], + api_key="fake-key", + input_cost_per_token=0.01, + output_cost_per_token=0.02, + cache_read_input_token_cost=0.001, + output_cost_per_reasoning_token=0.04, + mock_response="hello", + litellm_logging_obj=logger, + ) + response: Final = ModelResponse( + model="gpt-4o-mini", + usage=Usage( + prompt_tokens=100, + completion_tokens=20, + prompt_tokens_details={"cached_tokens": 40}, + completion_tokens_details={"reasoning_tokens": 5}, + ), + ) + assert logger.response_cost_calculator(response, ignore_stamped_cost=True) == pytest.approx( + 60 * 0.01 + 40 * 0.001 + 15 * 0.02 + 5 * 0.04 + ) + + +@pytest.mark.parametrize("explicit_null_mode", (False, True)) +def test_router_custom_token_prices_accept_missing_deployment_mode(explicit_null_mode: bool) -> None: + from litellm import Router + from litellm.types.utils import ModelResponse + + deployment_id: Final = "custom-pricing-without-mode" + keys: Final = ("gpt-4o-mini", "openai/gpt-4o-mini", deployment_id) + original: Final = _snapshot_model_cost_entries(keys) + model_info: Final = {"id": deployment_id, **({"mode": None} if explicit_null_mode else {})} + try: + router: Final = Router( + model_list=[ + { + "model_name": "priced-chat", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake-key", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "mock_response": "hello back", + }, + "model_info": model_info, + } + ], + num_retries=0, + ) + registered: Final = _snapshot_model_cost_entries(keys) + logger: Final = _request_logger("openai/gpt-4o-mini", "completion") + response: Final = router.completion( + model="priced-chat", messages=[{"role": "user", "content": "hello"}], litellm_logging_obj=logger + ) + assert isinstance(response, ModelResponse) + assert response.usage.prompt_tokens > 0 + assert logger.response_cost_calculator(response, ignore_stamped_cost=True) == pytest.approx( + response.usage.prompt_tokens * 0.001 + response.usage.completion_tokens * 0.002 + ) + assert _snapshot_model_cost_entries(keys) == registered finally: - _restore_model_cost_entries(model_cost_entries) - litellm.open_ai_chat_completion_models.intersection_update(openai_models_before) + _restore_model_cost_entries(original)