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 index ef18da183ed..5bdedb25c25 100644 --- 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 @@ -1,6 +1,6 @@ ALTER TABLE "LiteLLM_AccountingEvent" -ADD COLUMN "compilation_state" TEXT NOT NULL DEFAULT 'pending', -ADD COLUMN "compilation_error" TEXT, -ADD COLUMN "compilation_retry_at" TIMESTAMP(3); -CREATE INDEX "LiteLLM_AccountingEvent_kind_compilation_state_recorded_at_idx" +ADD COLUMN IF NOT EXISTS "compilation_state" TEXT NOT NULL DEFAULT 'pending', +ADD COLUMN IF NOT EXISTS "compilation_error" TEXT, +ADD COLUMN IF NOT EXISTS "compilation_retry_at" TIMESTAMP(3); +CREATE INDEX IF NOT EXISTS "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 45ea5cd20d0..c2f34e3b17e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -2116,8 +2116,8 @@ model LiteLLM_AccountingEvent { compilation_retry_at DateTime? @@index([recorded_at]) - @@index([kind, source_event_id]) @@index([kind, compilation_state, recorded_at]) + @@index([kind, source_event_id]) } // 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 d50627cceb0..9f567a6c707 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2924,6 +2924,46 @@ def transcription_cost(model: str, custom_llm_provider: str | None, duration: fl return openai_cost_per_second(model=model, custom_llm_provider=custom_llm_provider, duration=duration) +def image_cost_model_candidates( + model: str, custom_llm_provider: str | None, quality: str | None, size: str | None +) -> tuple[str | None, ...]: + # Standardize size format to use "-x-" + size_str: str = size or "1024-x-1024" + size_str = size_str.replace("x", "-x-") if "x" in size_str and "-x-" not in size_str else size_str + + # Build model names for cost lookup + base_model_name = f"{size_str}/{model}" + model_name_without_custom_llm_provider: str | None = None + if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): + model_name_without_custom_llm_provider = model.replace(f"{custom_llm_provider}/", "") + base_model_name = f"{custom_llm_provider}/{size_str}/{model_name_without_custom_llm_provider}" + model_name_with_quality: Final = f"{quality}/{base_model_name}" if quality else base_model_name + provider_first_model_name_with_quality: Final = ( + f"{custom_llm_provider}/{quality}/{size_str}/{model_name_without_custom_llm_provider or model}" + if quality and custom_llm_provider + else None + ) + + # gpt-image-1 models use low, medium, high quality. If user did not specify quality, use medium fot gpt-image-1 model family + model_name_with_v2_quality: Final = f"{ImageGenerationRequestQuality.HIGH.value}/{base_model_name}" + + verbose_logger.debug("Looking up cost for models: %s, %s", model_name_with_quality, base_model_name) + + model_without_provider: Final = f"{size_str}/{model.split('/')[-1]}" + model_with_quality_without_provider = f"{quality}/{model_without_provider}" if quality else model_without_provider + + return ( + model_name_with_quality, + provider_first_model_name_with_quality, + base_model_name, + model_name_with_v2_quality, + model_with_quality_without_provider, + model_without_provider, + model, + model_name_without_custom_llm_provider, + ) + + def default_image_cost_calculator( model: str, custom_llm_provider: str | None = None, @@ -2950,44 +2990,9 @@ def default_image_cost_calculator( Raises: Exception: If model pricing not found in cost map """ - # Standardize size format to use "-x-" - size_str: str = size or "1024-x-1024" - size_str = size_str.replace("x", "-x-") if "x" in size_str and "-x-" not in size_str else size_str - - # Parse dimensions - height, width = map(int, size_str.split("-x-")) - - # Build model names for cost lookup - base_model_name = f"{size_str}/{model}" - model_name_without_custom_llm_provider: str | None = None - if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): - model_name_without_custom_llm_provider = model.replace(f"{custom_llm_provider}/", "") - base_model_name = f"{custom_llm_provider}/{size_str}/{model_name_without_custom_llm_provider}" - model_name_with_quality: Final = f"{quality}/{base_model_name}" if quality else base_model_name - provider_first_model_name_with_quality: Final = ( - f"{custom_llm_provider}/{quality}/{size_str}/{model_name_without_custom_llm_provider or model}" - if quality and custom_llm_provider - else None - ) - - # gpt-image-1 models use low, medium, high quality. If user did not specify quality, use medium fot gpt-image-1 model family - model_name_with_v2_quality: Final = f"{ImageGenerationRequestQuality.HIGH.value}/{base_model_name}" - - verbose_logger.debug("Looking up cost for models: %s, %s", model_name_with_quality, base_model_name) - - model_without_provider: Final = f"{size_str}/{model.split('/')[-1]}" - model_with_quality_without_provider = f"{quality}/{model_without_provider}" if quality else model_without_provider - - models_to_check: Final = ( - model_name_with_quality, - provider_first_model_name_with_quality, - base_model_name, - model_name_with_v2_quality, - model_with_quality_without_provider, - model_without_provider, - model, - model_name_without_custom_llm_provider, - ) + size_str: Final = (size or "1024x1024").replace("-x-", "x") + height, width = map(int, size_str.split("x")) + models_to_check: Final = image_cost_model_candidates(model, custom_llm_provider, quality, size) matched_model: Final = next( ( _model diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c87aaabc4b7..1dbcaf8b647 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2167,12 +2167,20 @@ class Logging(LiteLLMLoggingBaseClass): or ("embedding" if "embedding" in str(self.model_call_details.get("call_type", "")) else "chat"), } ) - effective: Final = model_info_with_pricing_overrides(normalized, overrides) + effective: Final = model_info_with_pricing_overrides( + normalized, + { + name: value + for name, value in overrides.items() + if deployment_id is None or normalized.get(name) != value + }, + ) 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), + allow_live_models=True, ) def response_cost_calculator( @@ -2182,6 +2190,7 @@ class Logging(LiteLLMLoggingBaseClass): litellm_model_name: str | None = None, router_model_id: str | None = None, ignore_stamped_cost: bool = False, + custom_pricing_override: bool | None = None, ) -> float | None: context: Final = ( retained_pricing(self.request_pricing) @@ -2190,7 +2199,7 @@ class Logging(LiteLLMLoggingBaseClass): ) 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 + result, cache_hit, litellm_model_name, router_model_id, ignore_stamped_cost, custom_pricing_override ) def _calculate_response_cost( @@ -2200,6 +2209,7 @@ class Logging(LiteLLMLoggingBaseClass): litellm_model_name: str | None = None, router_model_id: str | None = None, ignore_stamped_cost: bool = False, + custom_pricing_override: bool | None = None, ) -> float | None: """ Calculate response cost using result + logging object variables. @@ -2262,7 +2272,9 @@ class Logging(LiteLLMLoggingBaseClass): router_model_id = self.get_router_model_id() ## RESPONSE COST ## - custom_pricing: Final = self.custom_pricing_for(priced_result) + custom_pricing: Final = ( + self.custom_pricing_for(priced_result) if custom_pricing_override is None else custom_pricing_override + ) prompt: Final = self._prompt_for_cost_calculation() diff --git a/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py b/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py index f9509cac41e..bb575da9bff 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py +++ b/litellm/litellm_core_utils/llm_cost_calc/pricing_context.py @@ -1,8 +1,9 @@ +from collections import ChainMap from collections.abc import Generator, Mapping from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass -from typing import Final +from typing import Final, cast # noqa: TID251 # live catalog has no typed boundary yet from pydantic import TypeAdapter @@ -18,6 +19,7 @@ class RetainedPricing: models_json: str discounts_json: str margins_json: str + allow_live_models: bool = False def models(self) -> Mapping[str, ModelInfo]: return _MODELS.validate_json(self.models_json) @@ -41,14 +43,23 @@ def retained_model_info(model: str, provider: str | None) -> ModelInfo | None: return None models: Final = context.models() result: Final = models.get(f"{provider}/{model}") or models.get(model) - if result is None: + if result is None and not context.allow_live_models: raise ValueError(f"Pricing snapshot has no retained model {model}") return result def retained_models() -> Mapping[str, ModelInfo] | None: + import litellm + context: Final = _CURRENT.get() - return None if context is None else context.models() + if context is None: + return None + if not context.allow_live_models: + return context.models() + catalog: Final = cast( # cast-ok: legacy catalog getter has an untyped dict signature + dict[str, ModelInfo], vars(litellm)["model_cost"] + ) + return ChainMap(dict(context.models()), catalog) def pricing_discounts() -> Mapping[str, float]: diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py index 0dbc3d8c86a..c321c21c06b 100644 --- a/litellm/proxy/db/accounting_ledger.py +++ b/litellm/proxy/db/accounting_ledger.py @@ -175,7 +175,7 @@ WHERE event_id = $1 AND kind = 'outbound_usage' _RETRY_COMPILATION: Final = """ UPDATE "LiteLLM_AccountingEvent" -SET compilation_error = 'calculation_retryable', compilation_retry_at = CURRENT_TIMESTAMP + INTERVAL '30 seconds' +SET compilation_error = $2, compilation_retry_at = CURRENT_TIMESTAMP + INTERVAL '30 seconds' WHERE event_id = $1 AND kind = 'outbound_usage' AND compilation_state = 'pending' """ @@ -373,6 +373,10 @@ class AccountingLedger: OutboundUsageEvent.model_validate_json(_json(_JSON.validate_json(row["payload"])["event"])) for row in rows ) + async def defer_compilation(self, event: OutboundUsageEvent) -> None: + async with self._transaction() as db: + await db.execute_raw(_RETRY_COMPILATION, event.event_id, "compilation_retryable") + async def compile_cost( self, event: OutboundUsageEvent, @@ -410,7 +414,7 @@ class AccountingLedger: except ValueError: return await _unpriced(db, validated.event_id, "invalid_pricing_inputs") except Exception: # noqa: BLE001 # retain failed pricing for bounded retries without starving newer events - await db.execute_raw(_RETRY_COMPILATION, validated.event_id) + await db.execute_raw(_RETRY_COMPILATION, validated.event_id, "calculation_retryable") return Unpriced(usage_event_id=validated.event_id, reason="calculation_retryable") if isinstance(priced, Unpriced): if priced.usage_event_id != validated.event_id: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 45ea5cd20d0..c2f34e3b17e 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -2116,8 +2116,8 @@ model LiteLLM_AccountingEvent { compilation_retry_at DateTime? @@index([recorded_at]) - @@index([kind, source_event_id]) @@index([kind, compilation_state, recorded_at]) + @@index([kind, source_event_id]) } // 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 52452a79106..2c3bd8754d7 100644 --- a/litellm/proxy/spend_tracking/cost_compiler.py +++ b/litellm/proxy/spend_tracking/cost_compiler.py @@ -28,6 +28,7 @@ from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo 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 HttpxBinaryResponseContent, ResponsesAPIResponse +from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, CostBreakdown, @@ -71,6 +72,12 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: ) if not captured.charge_eligible or batch_summary: return Unpriced(usage_event_id=event.event_id, reason="not_charge_owner") + response_inputs: Final = _fields(captured.pricing.get("response_inputs")) + rerank_usage: Final = event.call_type in {CallTypes.rerank, CallTypes.arerank} and any( + isinstance(value, int) and not isinstance(value, bool) and value >= 0 + for name, value in _fields(response_inputs.get("billed_units")).items() + if name in {"search_units", "total_tokens"} + ) logged: Final = _fields(captured.pricing.get("logging")) speech_characters: Final = logged.get("prompt_characters") speech_usage: Final = ( @@ -86,7 +93,7 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: 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 and not speech_usage: + if captured.usage is None and not fixed_call and not realtime_usage and not speech_usage and not rerank_usage: return Unpriced(usage_event_id=event.event_id, reason="usage_unknown") pricing: Final = captured.pricing mapped: Final = _fields(pricing.get("model_map_information")) @@ -146,7 +153,7 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: **_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() + for name, snapshot in {**_fields(pricing.get("realtime_models")), **_fields(pricing.get("tool_models"))}.items() if _fields(snapshot).get("model_map_value") is not None } base_pricing: Final = _fields(pricing.get("base_model_information")) @@ -169,7 +176,6 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: _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")) replay_usage: Final = Usage.model_validate( { @@ -215,6 +221,10 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: } ) if event.call_type in {CallTypes.arealtime, CallTypes.aresponses_websocket} + else RerankResponse.model_validate( + {"id": event.event_id, "results": [], "meta": {"billed_units": response_inputs.get("billed_units")}} + ) + if event.call_type in {CallTypes.rerank, CallTypes.arerank} else ResponsesAPIResponse.model_validate( { "id": event.event_id, @@ -222,6 +232,9 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: "model": model, "output": tool_response.get("output", []), "usage": tool_response.get("usage"), + "service_tier": response_inputs.get("service_tier"), + "tools": tool_response.get("tools"), + "tool_usage": tool_response.get("tool_usage"), } ) if tool_response.get("responses_api") @@ -298,7 +311,13 @@ def price_usage(envelope: str) -> PricingSnapshot | Unpriced: amount: Final = ( sum(batch_costs) if batch_costs is not None - else logger.response_cost_calculator(response, ignore_stamped_cost=True) + else logger.response_cost_calculator( + response, + ignore_stamped_cost=True, + custom_pricing_override=TypeAdapter[bool | None](bool | None).validate_python( + logged.get("custom_pricing") + ), + ) ) if batch_costs is not None: logger.cost_breakdown = CostBreakdown( @@ -396,6 +415,12 @@ class LoggingCostCompiler: 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) + try: + await self._ledger.defer_compilation(event) + except Exception: # noqa: BLE001 # database outages must not prevent attempts on later events + verbose_proxy_logger.exception( + "Cost compilation retry could not be scheduled for %s", event.event_id + ) _SPEND_TARGETS: Final = ( diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py index 4fede891018..edb7015f9a4 100644 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ b/litellm/proxy/spend_tracking/usage_capture.py @@ -8,7 +8,7 @@ from dataclasses import dataclass, replace from datetime import datetime, timezone from decimal import Decimal from hashlib import sha256 -from itertools import chain, groupby, islice, repeat +from itertools import chain, groupby, islice, product, repeat from typing import TYPE_CHECKING, Final from pydantic import BaseModel, JsonValue, TypeAdapter @@ -828,7 +828,18 @@ async def capture_proxy_response(logging_obj: object, response: object, *, backg 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": + accepted: Final = logging_obj.model_call_details.get("_accounting_capture") + committed: Final = ( + isinstance(accepted, CapturedUsage) + and accepted.receipt.done() + and not accepted.receipt.cancelled() + and accepted.receipt.exception() is None + ) + if ( + sync_ack + and not committed + and settings.get("accounting_sync_ack_failure_policy", "fail_closed") == "fail_closed" + ): acknowledgement: Final = UsageAcknowledgementError(error) raise HTTPException(503, str(acknowledgement)) from acknowledgement @@ -911,7 +922,7 @@ 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: + if logging_obj.request_pricing is not None and not logging_obj.spilled_over_ptu_for(normalized): with retained_pricing(logging_obj.request_pricing): return _capture_effective_pricing(logging_obj, normalized, details) return _capture_effective_pricing(logging_obj, normalized, details) @@ -934,7 +945,7 @@ def _capture_effective_pricing( if prepared is not None else StandardLoggingPayloadSetup.get_model_cost_information( base_model=logging_obj.get_deployment_model_for_cost(), - custom_pricing=bool(params.get("custom_pricing")), + custom_pricing=logging_obj.custom_pricing_for(normalized), custom_llm_provider=_text(details.get("custom_llm_provider")), init_response_obj=normalized, api_base=_text(params.get("api_base")), @@ -996,6 +1007,16 @@ def _capture_effective_pricing( ) for model in realtime_models }, + "tool_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 _image_tool_pricing_models(normalized, provider) + if model in _fields(vars(litellm)["model_cost"]) + }, "batch_models": { model: StandardLoggingPayloadSetup.get_model_cost_information( base_model=model, @@ -1026,16 +1047,64 @@ def _tool_cost_inputs(response: object) -> Mapping[str, object]: return { **flags, "responses_api": True, + "tools": tuple( + {key: value for key, value in _fields(tool).items() if key in {"type", "model"}} + for tool in TypeAdapter(tuple[object, ...]).validate_python(raw.get("tools") or ()) + if _fields(tool).get("type") == "image_generation" + ), + "tool_usage": _response_fields(raw.get("tool_usage")) or None, "usage": _response_fields(raw.get("usage")), "output": tuple( - {key: value for key, value in _fields(item).items() if key in {"type", "id", "status", "usage"}} + { + key: value + for key, value in _fields(item).items() + if key in {"type", "id", "status", "usage", "quality", "size"} + } 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) +def _image_tool_pricing_models(response: object, provider: str | None) -> tuple[str, ...]: + from litellm.cost_calculator import image_cost_model_candidates + + raw: Final = _tool_cost_inputs(response) + models: Final = tuple( + model + for tool in TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(raw.get("tools") or ()) + if isinstance(model := tool.get("model"), str) + ) + items: Final = tuple( + item + for item in TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(raw.get("output") or ()) + if item.get("type") == "image_generation_call" and item.get("status") == "completed" + ) + candidates: Final = chain.from_iterable( + image_cost_model_candidates( + model, + provider, + _text(item.get("quality")) if item.get("quality") != "auto" else None, + _text(item.get("size")) if item.get("size") != "auto" else None, + ) + for model, item in product(models, items) + ) + return tuple(dict.fromkeys(model for model in candidates if model is not None)) + + +_REALTIME_CHILDREN: Final = ( + "response", + "event", + "session", + "audio", + "input", + "transcription", + "input_audio_transcription", +) + + +def _realtime_cost_node( + fields: Mapping[str, object], +) -> dict[str, JsonValue]: # mutable-ok: tree builder attaches sanitized children return { **{ key: _JSON_OBJECT.validate_python(fields[key]) for key in ("usage",) if isinstance(fields.get(key), Mapping) @@ -1045,14 +1114,26 @@ def _realtime_cost_fields(value: object) -> Mapping[str, JsonValue]: 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_fields(value: object) -> Mapping[str, JsonValue]: + root: Final = _realtime_cost_node(_fields(value)) + pending: Final[list[tuple[Mapping[str, object], dict[str, JsonValue]]]] = [ # mutable-ok: traversal stack + (_fields(value), root) + ] + while pending: + source, target = pending.pop() + for key in _REALTIME_CHILDREN: + if not isinstance(source.get(key), Mapping): + continue + fields: Final = _fields(source[key]) # pyright: ignore[reportGeneralTypeIssues] # loop-local traversal value + child: Final = _realtime_cost_node(fields) # pyright: ignore[reportGeneralTypeIssues] # loop-local traversal value + target[key] = child + pending.append((fields, child)) + return root + + def _realtime_cost_events(response: object) -> tuple[Mapping[str, JsonValue], ...]: results: Final = _response_fields(response).get("results") return ( @@ -1063,12 +1144,15 @@ def _realtime_cost_events(response: object) -> tuple[Mapping[str, JsonValue], .. 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) + pending: Final[list[Mapping[str, JsonValue]]] = [ # mutable-ok: iterative traversal avoids recursion + value + ] + while pending: + fields: Final = pending.pop() # pyright: ignore[reportGeneralTypeIssues] # loop-local traversal value + model: Final = fields.get("model") # pyright: ignore[reportGeneralTypeIssues] # loop-local traversal value + if isinstance(model, str): + yield model + pending.extend(nested for key in _REALTIME_CHILDREN if isinstance(nested := fields.get(key), dict)) def _pricing_inputs(kwargs: Mapping[str, object], logging_obj: object, response: object) -> Mapping[str, object]: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f4763acaeda..a5e2bb1eb68 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -488,6 +488,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): "extraction", "ocr", "realtime", + "rerank", ] ] supported_endpoints: list[str] | None diff --git a/schema.prisma b/schema.prisma index 45ea5cd20d0..c2f34e3b17e 100644 --- a/schema.prisma +++ b/schema.prisma @@ -2116,8 +2116,8 @@ model LiteLLM_AccountingEvent { compilation_retry_at DateTime? @@index([recorded_at]) - @@index([kind, source_event_id]) @@index([kind, compilation_state, recorded_at]) + @@index([kind, source_event_id]) } // 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 0171d3683ca..0dc27b09991 100644 --- a/tests/integration/database/test_accounting_ledger.py +++ b/tests/integration/database/test_accounting_ledger.py @@ -995,7 +995,10 @@ async def test_proxy_capture_and_background_compiler_converge_on_final_header_ch @pytest.mark.asyncio -async def test_failing_compilation_page_does_not_starve_later_usage(writer: _Writer, ledger_database: str) -> None: +@pytest.mark.parametrize("conflicting_result", (False, True)) +async def test_failing_compilation_page_does_not_starve_later_usage( + writer: _Writer, ledger_database: str, conflicting_result: bool +) -> None: _, store = writer broken: Final = _usage("broken-price") healthy: Final = _usage("healthy-price") @@ -1004,12 +1007,15 @@ async def test_failing_compilation_page_does_not_starve_later_usage(writer: _Wri await _insert(store, broken) await _insert(store, healthy) - def price(envelope: str) -> PricingSnapshot: + def price(envelope: str) -> PricingSnapshot | Unpriced: if OutboundUsageEvent.model_validate(json.loads(envelope)["event"]).event_id == broken.event_id: + if conflicting_result: + return Unpriced(usage_event_id="different-usage", reason="usage_unknown") raise RuntimeError("temporary calculator failure") return _cost().snapshot await LoggingCostCompiler(store, price, _TARGETS).compile_pending(limit=1) + assert await store.pending_compilation(limit=1) == (healthy,) _rows( ledger_database, 'UPDATE "LiteLLM_AccountingEvent" SET compilation_retry_at = CURRENT_TIMESTAMP WHERE event_id = %s', diff --git a/tests/unit/proxy/spend_tracking/test_cost_compiler.py b/tests/unit/proxy/spend_tracking/test_cost_compiler.py index 1b89f472248..0e9c8f48609 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_compiler.py +++ b/tests/unit/proxy/spend_tracking/test_cost_compiler.py @@ -1348,3 +1348,230 @@ def test_character_pricing_does_not_invent_zero_for_missing_or_invalid_quantity( assert price_usage(json.dumps({**original, "event": event.model_dump(mode="json")})) == Unpriced( usage_event_id=event.event_id, reason="usage_unknown" ) + + +@pytest.mark.parametrize("measured_tokens", (False, True)) +def test_rerank_replay_preserves_billable_search_units(measured_tokens: bool) -> None: + from pydantic import TypeAdapter + + original: Final = _retained_envelope( + rates={"mode": "rerank", "input_cost_per_query": 0.03}, + response_inputs={"billed_units": {"search_units": 3, "total_tokens": 14}}, + ) + event: Final = OutboundUsageEvent.model_validate(original["event"]).model_copy( + update={ + "call_type": CallTypes.arerank, + "dimensions": ( + Dimension(name=DimensionName.MODEL, value="cohere/retained-reranker"), + Dimension(name=DimensionName.PROVIDER, value="cohere"), + ), + } + ) + pricing: Final = TypeAdapter(dict[str, object]).validate_python(original["pricing"]) + result: Final = price_usage( + json.dumps( + { + **original, + "usage": original["usage"] if measured_tokens else None, + "event": event.model_dump(mode="json"), + "pricing": pricing, + } + ) + ) + assert isinstance(result, PricingSnapshot) + assert float(result.total_cost) == pytest.approx(3 * 0.03) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("served_tier", ("default", "priority")) +async def test_responses_replay_uses_served_tier(served_tier: str) -> 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, + "input_cost_per_token_priority": 0.03, + "output_cost_per_token_priority": 0.04, + }, + ) + logger.update_environment_variables( + model="capture-model", + custom_llm_provider="openai", + litellm_params={}, + optional_params={"service_tier": "priority" if served_tier == "default" else "default"}, + ) + response: Final = ResponsesAPIResponse.model_validate( + { + "id": "tier", + "created_at": 1, + "model": "capture-model", + "output": [], + "service_tier": served_tier, + "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 == pytest.approx(0.18 if served_tier == "default" else 0.46) + 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(source_cost) + finally: + await meter.stop(1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("token_pricing", "automatic_size"), ((False, False), (False, True), (True, False))) +@pytest.mark.parametrize("custom_primary", (False, True)) +async def test_responses_image_tool_replays_retained_prices_without_image_content( + monkeypatch: pytest.MonkeyPatch, token_pricing: bool, custom_primary: bool, automatic_size: bool +) -> 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.types.llms.openai import ResponsesAPIResponse + from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage, _logger + from tests.unit.proxy.spend_tracking.test_usage_meter import _manual_flush + + monkeypatch.setitem( + litellm.model_cost, + "capture-model", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, + ) + monkeypatch.setitem( + litellm.model_cost, + "retained-image", + { + "litellm_provider": "openai", + "mode": "image_generation", + "input_cost_per_token": 0.03, + "output_cost_per_image_token": 0.04, + }, + ) + image_price_key: Final = f"{'high' if automatic_size else 'low'}/1024-x-1024/retained-image" + monkeypatch.setitem( + litellm.model_cost, + image_price_key, + {"litellm_provider": "openai", "mode": "image_generation", "output_cost_per_image": 0.25}, + ) + logger: Final = _logger("aresponses") + if custom_primary: + logger.set_request_pricing( + "capture-model", "openai", {"input_cost_per_token": 0.01, "output_cost_per_token": 0.02} + ) + response: Final = ResponsesAPIResponse.model_validate( + { + "id": "image-tool", + "created_at": 1, + "model": "capture-model", + "tools": [{"type": "image_generation", "model": "retained-image"}], + "output": [ + { + "id": "image", + "type": "image_generation_call", + "status": "completed", + "quality": "auto" if automatic_size else "low", + "size": "auto" if automatic_size else "1024x1024", + "result": "private-image-content", + } + ], + "usage": {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}, + "tool_usage": { + "image_gen": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "input_tokens_details": {"text_tokens": 10, "image_tokens": 0}, + "output_tokens_details": {"text_tokens": 0, "image_tokens": 5}, + } + } + if token_pricing + else None, + } + ) + source_cost: Final = logger.response_cost_calculator(response, ignore_stamped_cost=True) + assert source_cost == pytest.approx(0.18 + (0.5 if token_pricing else 0.25)) + 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-image-content" not in captured + monkeypatch.delitem(litellm.model_cost, "retained-image") + monkeypatch.delitem(litellm.model_cost, image_price_key) + 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 +@pytest.mark.parametrize("spilled_over", (False, True)) +async def test_capture_replays_effective_ptu_spillover_prices( + monkeypatch: pytest.MonkeyPatch, spilled_over: bool +) -> None: + import litellm + from litellm.litellm_core_utils.hidden_params import set_hidden_param + 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 + + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "true") + monkeypatch.setitem( + litellm.model_cost, + "azure/capture-model", + {"litellm_provider": "azure", "mode": "chat", "input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, + ) + logger: Final = _logger() + logger.update_environment_variables( + model="capture-model", + custom_llm_provider="azure", + optional_params={}, + litellm_params={ + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "metadata": { + "model_info": { + "team_id": "team", + "ptu_count": 10, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01", + } + }, + }, + ) + logger.set_request_pricing("capture-model", "azure", {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}) + response: Final = ModelResponse( + id="spill", model="capture-model", usage=Usage(prompt_tokens=10, completion_tokens=4, total_tokens=14) + ) + set_hidden_param( + response, "additional_headers", {"llm_provider-x-ms-is-spilled-over": "true"} if spilled_over else {} + ) + source_cost: Final = logger.response_cost_calculator(response, ignore_stamped_cost=True) + assert source_cost == pytest.approx(0.18 if spilled_over else 0) + 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(source_cost) + finally: + await meter.stop(1) diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py index c83d11250c9..a1327e0a177 100644 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ b/tests/unit/proxy/spend_tracking/test_usage_capture.py @@ -1725,8 +1725,9 @@ class DelayedCostCompiler: @pytest.mark.asyncio -async def test_sync_ack_deadline_bounds_compilation_after_usage_commit(monkeypatch: pytest.MonkeyPatch) -> None: - from fastapi import HTTPException +async def test_sync_ack_preserves_committed_usage_ack_when_compilation_times_out( + monkeypatch: pytest.MonkeyPatch, +) -> None: from litellm.proxy import proxy_server from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response @@ -1742,11 +1743,9 @@ async def test_sync_ack_deadline_bounds_compilation_after_usage_commit(monkeypat 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 + await asyncio.wait_for( + capture_proxy_response(_logger(), ModelResponse(id="provider", usage=Usage(prompt_tokens=3))), 2 + ) assert compiler.entered.is_set() assert compiler.cancelled.is_set() assert database.committed == ("caller:inbound", "provider") diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 5cf872f4ea7..add903ad561 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -2333,7 +2333,7 @@ def test_register_model_price_is_used_for_completion_cost() -> None: @respx.mock @pytest.mark.usefixtures("isolated_openai_model_sets", "local_model_cost_map") -def test_completion_registers_test_owned_model_prices() -> None: +def test_completion_keeps_custom_prices_scoped_to_the_request() -> None: model: Final = "migration-owned-pricing-model" messages: Final = [{"role": "user", "content": "price control"}] litellm.register_model( @@ -2376,8 +2376,8 @@ def test_completion_registers_test_owned_model_prices() -> None: ) assert json.loads(route.calls[0].request.content) == {"model": model, "messages": messages} - assert litellm.model_cost[model]["input_cost_per_token"] == 0.001 - assert litellm.model_cost[model]["output_cost_per_token"] == 0.002 + assert litellm.model_cost[model]["input_cost_per_token"] == 0.0001 + assert litellm.model_cost[model]["output_cost_per_token"] == 0.0002 assert response._hidden_params["response_cost"] == pytest.approx(0.018)