From 85a987823c5f63e069f1cbdac89ace35908495b4 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:34:19 -0700 Subject: [PATCH] feat(accounting): capture request usage through the durable meter --- .../common_utils/check_responses_cost.py | 1 + litellm/_internal_context.py | 5 +- litellm/accounting/protocols.py | 6 +- litellm/accounting/types.py | 4 +- litellm/batches/batch_utils.py | 84 +- .../speech_to_completion_bridge/handler.py | 2 +- litellm/litellm_core_utils/litellm_logging.py | 193 +++- .../litellm_core_utils/realtime_streaming.py | 2 + .../llms/vertex_ai/batches/transformation.py | 17 +- litellm/proxy/_types.py | 11 + litellm/proxy/common_request_processing.py | 14 + .../proxy/hooks/proxy_track_cost_callback.py | 80 +- litellm/proxy/proxy_server.py | 31 + .../proxy/spend_tracking/cost_finalization.py | 1 + .../spend_tracking/spend_tracking_utils.py | 17 +- litellm/proxy/spend_tracking/usage_capture.py | 791 ++++++++++++++++ litellm/proxy/utils.py | 28 +- litellm/rag/main.py | 10 +- litellm/responses/main.py | 6 +- .../responses/mcp/chat_completions_handler.py | 4 +- litellm/router.py | 5 + litellm/types/utils.py | 6 + litellm/utils.py | 7 + tests/unit/batches/test_batch_utils.py | 47 + .../llms/azure/test_azure_common_utils.py | 3 + .../hooks/test_proxy_track_cost_callback.py | 50 + .../spend_tracking/test_cost_finalization.py | 27 + .../test_spend_tracking_utils.py | 18 + .../spend_tracking/test_usage_capture.py | 882 ++++++++++++++++++ tests/unit/proxy/test_proxy_server.py | 52 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 46 +- 31 files changed, 2395 insertions(+), 55 deletions(-) create mode 100644 litellm/proxy/spend_tracking/usage_capture.py create mode 100644 tests/unit/proxy/spend_tracking/test_usage_capture.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 69e36a0c076..cfe4ca98145 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -200,6 +200,7 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi unpriced: Final = _without_read_price(billed.response) cost: Final = logging_obj.response_cost_calculator(result=unpriced) end_time: Final = datetime.now() + logging_obj.capture_usage_before_logging(unpriced, logging_obj.start_time, end_time) GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( async_coroutine=logging_obj.async_success_handler( result=unpriced, start_time=logging_obj.start_time, end_time=end_time diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index c2243b90fbf..40d3377bf66 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -21,6 +21,7 @@ _T = TypeVar("_T") # When True, suppresses async logging and billing for internal sub-calls # (e.g., emulated file-search steps that make nested LLM calls). is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", default=False) +internal_parent_request_id: Final[ContextVar[str | None]] = ContextVar("internal_parent_request_id", default=None) # One request prices its totals, its per-token-type lines and the rates it reports on # separate code paths. Each reads the clock for off-peak pricing, so without a pinned @@ -40,13 +41,15 @@ _mirrored_call: Final[ContextVar[bool]] = ContextVar("mirrored_call", default=Fa @contextmanager -def internal_sub_call() -> Generator[None]: +def internal_sub_call(parent_request_id: str | None = None) -> Generator[None]: """A nested call the parent call logs and bills, so the parent must fold this call's cost into its own.""" token: Final = is_internal_call.set(True) + parent_token: Final = internal_parent_request_id.set(parent_request_id or internal_parent_request_id.get()) try: yield finally: is_internal_call.reset(token) + internal_parent_request_id.reset(parent_token) @contextmanager diff --git a/litellm/accounting/protocols.py b/litellm/accounting/protocols.py index 8955f56e7a0..0208e37fbdf 100644 --- a/litellm/accounting/protocols.py +++ b/litellm/accounting/protocols.py @@ -3,7 +3,7 @@ The event row's ``payload`` is a private envelope versioned by the row's ``schema_version``. Usage envelopes retain the public event and request-scoped pricing context: quantities with units, model/provider, requested and billed tier/region when known, and per-request/deployment overrides. -These inputs must survive request teardown. Cost-map rates are resolved at compilation, not capture. +These inputs, including the effective cost-map rates, must survive request teardown. Cost envelopes retain the public CostEvent, including its required finalized pricing snapshot. T5/T6 own envelope serialization; this public API does not introduce a second persisted schema. """ @@ -38,8 +38,8 @@ class CostCompiler(Protocol): async def compile(self, event: OutboundUsageEvent) -> CostEvent | Unpriced: """Reuse the committed charge or finalize one charge and its delivery intent. - Resolve captured request-scoped context by event identity and cost-map rates - at compilation. Persist the resolved inputs in the required snapshot. Replay + Resolve captured request-scoped context and rates by event identity. + Persist the resolved inputs in the required snapshot. Replay reuses the committed result and never consumes free allowance again. Missing pricing is Unpriced, never an inferred zero. """ diff --git a/litellm/accounting/types.py b/litellm/accounting/types.py index d0be29bf5b9..2ee6f60f756 100644 --- a/litellm/accounting/types.py +++ b/litellm/accounting/types.py @@ -276,8 +276,8 @@ class NamedCost(BaseModel): class PricingSnapshot(BaseModel): """Retained inputs and result of authoritative compilation, reproducible after cost-map rates change. - Request-scoped inputs and overrides come from the usage row's versioned payload. Cost-map rates are - resolved at compilation, not frozen at capture. Retries reuse the committed charge and snapshot. + Request-scoped inputs, overrides and effective rates come from the usage row's versioned payload. + Retries reuse the committed charge and snapshot. Component names match ``CostBreakdown``: reasoning_cost is part of output_cost, not an additional charge. Component sums and float tolerances belong to reconciliation, not this structural contract. """ diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 6a2fa41830f..42f9d5f74ce 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -3,9 +3,10 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from dataclasses import replace as dataclasses_replace from enum import Enum +from itertools import groupby from typing import Any, Final, Literal, cast -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import verbose_logger @@ -15,6 +16,7 @@ from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.llms.bedrock.batches.transformation import titan_embedding_usage_from_batch_output from litellm.llms.vertex_ai.batches.transformation import ( is_native_vertex_batch_output_row, + native_vertex_batch_row_model, native_vertex_batch_row_stats, ) from litellm.types.llms.openai import Batch @@ -210,6 +212,7 @@ class _BatchOutputLineStats: cache_creation_tokens: int reasoning_tokens: int model: str | None + usage: Usage | None def _classify_output_line_stats( @@ -268,7 +271,7 @@ def _compute_output_line_stats( raw_model: Final = response_body.get("model") response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None completion_details: Final = usage.completion_tokens_details - line_prompt_cost, line_completion_cost = _output_line_cost( + line_prompt_cost, line_completion_cost = _output_line_cost_or_pending( response_body=response_body, usage=usage, custom_llm_provider=custom_llm_provider, @@ -286,7 +289,15 @@ def _compute_output_line_stats( cache_read_tokens=prompt_details["cache_hit_tokens"], cache_creation_tokens=prompt_details["cache_creation_tokens"], reasoning_tokens=(completion_details.reasoning_tokens if completion_details else None) or 0, - model=response_model, + model=response_model or model_name, + usage=( + usage.model_copy(update=ocr_usage.model_dump(exclude_none=True)) + if (ocr_usage := _ocr_usage_info_from_response_body(response_body)) is not None + and ocr_usage.model_dump(exclude_none=True) + else usage + if response_body.get("usage") is not None or usage.total_tokens > 0 + else None + ), ) @@ -298,6 +309,47 @@ def _ocr_usage_info_from_response_body(response_body: Mapping[str, object]) -> O return OCRUsageInfo.model_validate(raw_usage_info) +def _output_line_cost_or_pending( + response_body: Mapping[str, object], + usage: Usage, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + model_name: str | None, + response_model: str | None, + model_info: ModelInfo | None, + vertex_location: str | None, +) -> tuple[float, float]: + try: + return _output_line_cost( + response_body, usage, custom_llm_provider, model_name, response_model, model_info, vertex_location + ) + except Exception: # noqa: BLE001 # retain measured usage when existing batch pricing fails + verbose_logger.warning("Batch line pricing unavailable; retaining usage with the existing zero-cost fallback") + return 0.0, 0.0 + + +def _combined_reported_batch_usage(usages: tuple[Usage | None, ...]) -> Mapping[str, object] | None: + from litellm.cost_calculator import BaseTokenUsageProcessor + + reported: Final = [usage for usage in usages if usage is not None] + if not reported: + return None + fields: Final = tuple(TypeAdapter(dict[str, JsonValue]).validate_python(usage.model_dump()) for usage in reported) + quantities: Final = { + key: sum(value for usage in fields if isinstance(value := usage.get(key), (int, float))) + for key in OCRUsageInfo.model_fields + if any(isinstance(usage.get(key), (int, float)) for usage in fields) + } + return {**BaseTokenUsageProcessor.combine_usage_objects(usage_objects=reported).model_dump(), **quantities} + + +def _batch_usage_by_model(usages: tuple[tuple[str | None, Usage | None], ...]) -> Mapping[str, object]: + ordered: Final = sorted(usages, key=lambda item: item[0] or "") + return { + model: _combined_reported_batch_usage(tuple(item[1] for item in group)) + for model, group in groupby(ordered, key=lambda item: item[0] or "") + } + + def _output_line_cost( response_body: Mapping[str, object], usage: Usage, @@ -361,6 +413,8 @@ def _aggregate_batch_cost_usage_models( completion_tokens=sum(stats.completion_tokens for stats in line_stats), reasoning_tokens=sum(stats.reasoning_tokens for stats in line_stats), **cache_token_params, + batch_usage_by_model=_batch_usage_by_model(tuple((stats.model, stats.usage) for stats in line_stats)), + batch_unknown_usage_requests=sum(stats.usage is None for stats in line_stats), ) batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model] total_prompt_cost: Final = sum((stats.prompt_cost for stats in line_stats), 0.0) @@ -401,17 +455,21 @@ def calculate_vertex_ai_batch_cost_and_usage( from litellm.cost_calculator import batch_cost_calculator from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig - row_stats: Final = tuple( - native_vertex_batch_row_stats( - row, - model_name, - model_info=model_info, - calculate_usage=VertexGeminiConfig.calculate_usage, - cost_calculator=batch_cost_calculator, - vertex_location=vertex_location, + observations: Final = tuple( + ( + native_vertex_batch_row_model(row, model_name), + native_vertex_batch_row_stats( + row, + model_name, + model_info=model_info, + calculate_usage=VertexGeminiConfig.calculate_usage, + cost_calculator=batch_cost_calculator, + vertex_location=vertex_location, + ), ) for row in vertex_ai_batch_responses ) + row_stats: Final = tuple(stats for _, stats in observations) priced: Final = tuple(stats for stats in row_stats if stats is not None) total_prompt_cost: Final = sum(stats.prompt_cost for stats in priced) total_completion_cost: Final = sum(stats.completion_cost for stats in priced) @@ -435,6 +493,10 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens=total_tokens, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, + batch_usage_by_model=_batch_usage_by_model( + tuple((model, stats.usage if stats is not None else None) for model, stats in observations) + ), + batch_unknown_usage_requests=sum(stats is None for stats in row_stats), ), models=( [model_name] diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index faf2daecf5f..e4d270834ff 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -113,7 +113,7 @@ class SpeechToCompletionBridgeHandler: voice=voice, ) - with internal_sub_call(): + with internal_sub_call(logging_obj.litellm_call_id): result: Final = completion( **request_data, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 924603dc7e8..22ea72fa56e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -16,9 +16,10 @@ import traceback from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime as dt_object +from enum import Enum from functools import lru_cache from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, cast from httpx import Response, Timeout from pydantic import BaseModel, JsonValue, TypeAdapter @@ -47,6 +48,7 @@ from litellm.constants import ( KUBERNETES_POD_ROUTING_KEY, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, PROVIDER_REQUEST_ID_HEADERS, + REALTIME_METERED_USAGE_KEY, REDACTED_BY_LITELLM, ) from litellm.cost_calculator import ( @@ -266,6 +268,8 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation + from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation + from litellm.proxy.spend_tracking.usage_capture import CapturedUsage try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -630,7 +634,35 @@ def same_attempt_marker(first: object, second: object) -> bool: return first is second or first == second +class UsageObservation(Enum): + KNOWN = "known" + UNKNOWN = "unknown" + ESTIMATED = "estimated" + + class Logging(LiteLLMLoggingBaseClass): + usage_capture_callback: ClassVar[ + Callable[["Logging", object, datetime.datetime, datetime.datetime, bool], None] | None + ] = None + + def capture_usage_before_logging( + self, + response: object, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + *, + freeze_only: bool = False, + ) -> None: + capture: Final = type(self).usage_capture_callback + if capture is not None and (not self.stream or self._is_assembled_stream_success(response)): + capture( + self, + response, + start_time or self.start_time, + end_time or datetime.datetime.now(datetime.timezone.utc), + freeze_only, + ) + global \ supabaseClient, \ promptLayerLogger, \ @@ -828,6 +860,7 @@ class Logging(LiteLLMLoggingBaseClass): self.finalized_cost: FinalizedCost | None = None self.prepared_cost: PreparedCost | None = None self.cost_allowance_discount: float | None = None + self.cost_allowance_allocation: SearchAllowanceAllocation | None = None self.cost_finalization_lock = asyncio.Lock() self.deferred_success_result: tuple[object, datetime.datetime, datetime.datetime] | None = None self.billed_token_rates: BilledTokenRates | None = None @@ -884,7 +917,7 @@ class Logging(LiteLLMLoggingBaseClass): new_call.start_time = start_time new_call.litellm_params = litellm_params new_call.model_call_details = { - **self.model_call_details, + **{key: value for key, value in self.model_call_details.items() if not key.startswith("_accounting_")}, "litellm_call_id": litellm_call_id, "start_time": start_time, "first_api_call_start_time": None, @@ -894,6 +927,7 @@ class Logging(LiteLLMLoggingBaseClass): new_call.finalized_cost = None new_call.prepared_cost = None new_call.cost_allowance_discount = None + new_call.cost_allowance_allocation = None new_call.cost_finalization_lock = asyncio.Lock() new_call.deferred_success_result = None new_call.streaming_chunks = [] @@ -2424,7 +2458,7 @@ class Logging(LiteLLMLoggingBaseClass): _is_sync_litellm_request = is_sync_litellm_request - def _is_assembled_stream_success(self, result=None) -> bool: + def _is_assembled_stream_success(self, result: object = None) -> bool: """Final assembled stream export (not a per-chunk success call). Per-chunk callers pass a ``ModelResponseStream`` (or ``None``); the @@ -2459,7 +2493,7 @@ class Logging(LiteLLMLoggingBaseClass): """ from litellm.litellm_core_utils.thread_pool_executor import executor - if self._is_assembled_stream_success(result): + if self._is_assembled_stream_success(TypeAdapter(object).validate_python(result)): if self.model_call_details.get("has_dispatched_final_stream_success"): return self.model_call_details["has_dispatched_final_stream_success"] = True @@ -2477,6 +2511,11 @@ class Logging(LiteLLMLoggingBaseClass): ) return + self.capture_usage_before_logging( + TypeAdapter(object).validate_python(result), + TypeAdapter[datetime.datetime | None](datetime.datetime | None).validate_python(start_time), + TypeAdapter[datetime.datetime | None](datetime.datetime | None).validate_python(end_time), + ) await self.async_success_handler( result, start_time=start_time, @@ -2503,6 +2542,17 @@ class Logging(LiteLLMLoggingBaseClass): return self.prepared_cost recorded_cache_hit: Final[object] = self.model_call_details.get("cache_hit") deferred: Final = self.deferred_success_result + observed_response: Final = _LITELLM_PARAMS.validate_python( + _extract_response_obj_and_hidden_params(deferred[0] if deferred else response, None)[0] + ) + self.model_call_details["_accounting_usage_observed"] = ( + UsageObservation.KNOWN + if ( + self.model_call_details.get("combined_usage_object") + or StandardLoggingPayloadSetup.get_usage_as_dict(observed_response) + ) + else UsageObservation.UNKNOWN + ) prepared: Final = self.prepare_success_logging( result=deferred[0] if deferred is not None else response, start_time=deferred[1] if deferred is not None else self.start_time, @@ -2547,6 +2597,12 @@ class Logging(LiteLLMLoggingBaseClass): self.finalized_cost = finalized self.model_call_details["response_cost"] = finalized.amount + def record_usage_capture(self, captured: "CapturedUsage") -> None: + self.model_call_details["_accounting_capture"] = captured + + def record_usage_pricing(self, pricing: object) -> None: + self.model_call_details["_accounting_pricing_snapshot"] = pricing + def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None: """Stash what an interrupted stream already consumed so the failure log bills it instead of zero.""" self.model_call_details["combined_usage_object"] = usage @@ -2630,6 +2686,11 @@ class Logging(LiteLLMLoggingBaseClass): own usage while the request's failure log still bills every attempt.""" usage: Final = self._model_call_details_view.get("combined_usage_object") response_cost: Final = self._model_call_details_view.get("response_cost") + if isinstance(usage, Usage) and type(self).usage_capture_callback is not None: + self.capture_failed_attempt_usage(usage) + self.model_call_details.pop("_accounting_capture", None) + self.model_call_details.pop("_accounting_pricing_snapshot", None) + self.model_call_details.pop("_accounting_usage_observed", None) self.model_call_details.pop("combined_usage_object", None) self.model_call_details.pop("response_cost", None) if isinstance(usage, Usage): @@ -2639,6 +2700,41 @@ class Logging(LiteLLMLoggingBaseClass): (usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0), ) + def capture_failed_attempt_usage(self, usage: Usage) -> None: + from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY + + capture: Final = type(self).usage_capture_callback + if capture is None: + return + attempt: Final = self.attempt_scoped_copy() + identity: Final = ( + f"{self.litellm_call_id}-attempt-{len(self.model_call_details.get('failed_attempts_usage', ()))}" + ) + params: Final = _LITELLM_PARAMS.validate_python(self.model_call_details.get("litellm_params") or {}) + metadata: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(params) + attempt.model_call_details = { + **attempt.model_call_details, + "litellm_call_id": identity, + "parent_request_id": self.litellm_call_id, + "_accounting_capture": None, + "litellm_params": { + **params, + "metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"}, + "litellm_metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"}, + }, + } + capture( + attempt, + {"id": identity, "usage": usage.model_dump(), "status": "failed"}, + self.start_time, + datetime.datetime.now(datetime.timezone.utc), + False, + ) + self.model_call_details["_accounting_prior_attempt_ids"] = ( + *self.model_call_details.get("_accounting_prior_attempt_ids", ()), + identity, + ) + def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None: """Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw.""" self.model_call_details["continued_partial_usage"] = usage @@ -2796,6 +2892,55 @@ class Logging(LiteLLMLoggingBaseClass): _update_completion_start_time = update_completion_start_time + def normalize_endpoint_logging_result(self, result: object) -> object: + call_type: Final[object] = self.call_type + if call_type == CallTypes.anthropic_messages.value: + return self._anthropic_messages_logged_response(result=result) + if call_type in (CallTypes.generate_content.value, CallTypes.agenerate_content.value): + return self._handle_non_streaming_google_genai_generate_content_response_logging(result=result) + if call_type in (CallTypes.asend_message.value, CallTypes.send_message.value): + return self._handle_a2a_response_logging(result=result) + return result + + def usage_capture_response(self, result: object) -> object: + if self.stream and not self._is_assembled_stream_success(result): + return None + copied: Final = result.model_copy(deep=True) if isinstance(result, BaseModel) else result + normalized: Final = self.normalize_logging_result(self.normalize_endpoint_logging_result(copied)) + if self.model_call_details.get("call_type") == CallTypes.arealtime.value and isinstance(result, list): + frames: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(result) + measured: Final = any( + isinstance(_LITELLM_PARAMS.validate_python(frame.get("response") or {}).get("usage"), Mapping) + for frame in frames + if frame.get("type") == "response.done" + ) + if not measured: + metered: Final = self.model_call_details.get(REALTIME_METERED_USAGE_KEY) + return { + "id": self.litellm_call_id, + "usage": {"metered_usage": TypeAdapter(object).dump_python(metered, mode="json")} + if metered is not None + else None, + } + assembled: Final = self._get_assembled_streaming_response( + normalized, self.start_time, self.start_time, True, [] + ) + return normalized if assembled is None else assembled + + def usage_capture_pricing(self, result: object) -> Mapping[str, JsonValue]: + from litellm.utils import count_characters + + return TypeAdapter(dict[str, JsonValue]).validate_python( + { + "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, + "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, + } + ) + def normalize_logging_result(self, result: object) -> object: """ Some endpoints return a different type of result than what is expected by the logging system. @@ -2806,8 +2951,20 @@ class Logging(LiteLLMLoggingBaseClass): combined_usage_object: Final = ( RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=result) ) + metered: Final[object] = self.model_call_details.get(REALTIME_METERED_USAGE_KEY) + captured_usage: Final = ( + combined_usage_object.model_copy( + update={ + "metered_usage": TypeAdapter(dict[str, JsonValue]).validate_python( + TypeAdapter(object).dump_python(metered, mode="json") + ) + } + ) + if metered is not None + else combined_usage_object + ) logging_result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object( - usage=combined_usage_object, + usage=captured_usage, results=result, ) @@ -3080,15 +3237,7 @@ class Logging(LiteLLMLoggingBaseClass): else cache_hit ) - if self.call_type == CallTypes.anthropic_messages.value: - result = self._anthropic_messages_logged_response(result=result) - elif ( - self.call_type == CallTypes.generate_content.value - or self.call_type == CallTypes.agenerate_content.value - ): - result = self._handle_non_streaming_google_genai_generate_content_response_logging(result=result) - elif self.call_type == CallTypes.asend_message.value or self.call_type == CallTypes.send_message.value: - result = self._handle_a2a_response_logging(result=result) + result = self.normalize_endpoint_logging_result(result) logging_result: Final = self.normalize_logging_result(result=result) if cache_hit is None and _served_from_cache_at_no_charge(logging_result): @@ -3441,6 +3590,7 @@ class Logging(LiteLLMLoggingBaseClass): self._with_client_facing_stream_model(complete_streaming_response) ) self.model_call_details["complete_streaming_response"] = logged_streaming_response + self.capture_usage_before_logging(logged_streaming_response, start_time, end_time) ## STANDARDIZED LOGGING PAYLOAD if not for_async_request: self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( @@ -3964,6 +4114,8 @@ class Logging(LiteLLMLoggingBaseClass): cost_for_built_in_tools_cost_usd_dollar=0.0, ) + self.capture_usage_before_logging(result, start_time, end_time) + self.truncated_messages_for_logging = await truncate_base64_in_messages_async( StandardLoggingPayloadSetup.append_system_prompt_messages( kwargs=self.model_call_details, messages=self.model_call_details.get("messages") @@ -6844,9 +6996,9 @@ class StandardLoggingPayloadSetup: @staticmethod def get_usage_as_dict( - response_obj: dict | None, + response_obj: dict[str, Any] | None, combined_usage_object: Usage | None = None, - ) -> dict: + ) -> dict[str, Any]: """ Like get_usage_from_response_obj but returns a plain dict, skipping the Pydantic Usage construction on the hot path. @@ -6858,7 +7010,14 @@ class StandardLoggingPayloadSetup: return _empty _raw: Final = response_obj.get("usage", None) if _raw is None: - return dict(get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) or _empty) + reported: Final = response_obj.get("usage_info") + if isinstance(reported, BaseModel): + return reported.model_dump(exclude_none=True) + return ( + {key: value for key, value in reported.items() if value is not None} + if isinstance(reported, dict) + else dict(get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) or _empty) + ) if isinstance(_raw, ResponseAPIUsage): return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_raw).model_dump() if isinstance(_raw, dict): diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 4db22e9d5bf..46f5ca87a93 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -631,6 +631,8 @@ class RealTimeStreaming: if self.session_tools or self.tool_calls: self.logging_obj.model_call_details["realtime_tools"] = self.session_tools self.logging_obj.model_call_details["realtime_tool_calls"] = self.tool_calls + if LiteLLMLogging.usage_capture_callback is not None: + self.logging_obj.capture_usage_before_logging(self.messages) # Route through the bounded logging worker (per-coroutine timeout + # concurrency cap) instead of a bare create_task, so a slow callback # can't leave suspended tasks pinning each call's response in memory. diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 99d78bcdca6..a1843924c43 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -129,6 +129,17 @@ def _native_vertex_row_usage( return calculate_usage(completion_response) +def native_vertex_batch_row_model(row: Mapping[str, object], model_name: str | None) -> str | None: + response: Final = as_json_object(row.get("response")) + model_version: Final = ( + response.get("model" if is_vertex_anthropic_batch_output_response(response) else "modelVersion") + if response is not None + else None + ) + deployment_model: Final = model_name if model_name and "*" not in model_name else None + return deployment_model or (model_version if isinstance(model_version, str) else model_name) + + def native_vertex_batch_row_stats( row: Mapping[str, object], model_name: str | None, @@ -156,11 +167,7 @@ def native_vertex_batch_row_stats( if usage is None: return None total_tokens: Final = usage.total_tokens or (usage.prompt_tokens + usage.completion_tokens) - model_version: Final = response_body.get( - "model" if is_vertex_anthropic_batch_output_response(response_body) else "modelVersion" - ) - deployment_model: Final = model_name if model_name and "*" not in model_name else None - model: Final = deployment_model or (model_version if isinstance(model_version, str) else model_name) + model: Final = native_vertex_batch_row_model(row, model_name) if model is None: verbose_logger.warning( "vertex_ai batch output row could not be costed, so it is billed at $0 and the rest of the batch " diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 07588053959..1f12273bfde 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -38,6 +38,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_no_callback_env_reference, ) from litellm.proxy.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled +from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings from litellm.types.agents import ( AgentCaller as AgentCaller, # noqa: PLC0414 # public re-export ) @@ -2470,6 +2471,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ + accounting_capture: UsageCaptureSettings | None = Field( + None, + description="Opt-in durable usage capture. Legacy spend updates remain active. Requires local Prisma writes.", + ) + accounting_sync_ack: bool = Field(False, description="Wait for usage commit before a supported non-stream response") + accounting_sync_ack_timeout_seconds: float = Field(5.0, gt=0, description="Maximum response wait for usage commit") + accounting_sync_ack_failure_policy: Literal["fail_closed", "log"] = Field( + "fail_closed", description="Return 503 or log the persistence failure after the response wait expires" + ) + request_tag_mappings: Sequence[RequestTagMapping] | None = Field( None, description="Map W3C baggage values to request tags. Exact value/tag rules take precedence over tag_prefix fallbacks.", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c9b6fff3a00..ebe0923ebcc 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -137,6 +137,7 @@ from litellm.proxy.route_llm_request import ( route_request, ) from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost, prepare_response_cost +from litellm.proxy.spend_tracking.usage_capture import capture_proxy_response, validate_sync_ack_request from litellm.proxy.telemetry.request_context import bind_call_id as bind_telemetry_call_id from litellm.proxy.utils import ( # noqa: F401 # legacy module exports ProxyLogging, @@ -2459,6 +2460,12 @@ class ProxyBaseLLMRequestProcessing: skip_guardrails: bool = False, logging_auth: UserAPIKeyAuth | None = None, ) -> tuple[dict, LiteLLMLoggingObj]: + validate_sync_ack_request( + _REQUEST_FIELDS.validate_python(general_settings), + route_type, + _REQUEST_FIELDS.validate_python(self.data).get("stream") is True, + _REQUEST_FIELDS.validate_python(self.data).get("background") is True, + ) start_time: Final = datetime.now() # start before calling guardrail hooks requested_model: Final = self.data.get("model") @@ -3038,6 +3045,12 @@ class ProxyBaseLLMRequestProcessing: self._debug_log_request_payload() if skip_pre_call_logic: + validate_sync_ack_request( + general_settings, + route_type, + _REQUEST_FIELDS.validate_python(self.data).get("stream") is True, + _REQUEST_FIELDS.validate_python(self.data).get("background") is True, + ) logging_obj = ( cast( # cast-ok: pre-call processing stores Logging; preserve the existing duck-typed boundary LiteLLMLoggingObj | None, _REQUEST_FIELDS.validate_python(self.data).get("litellm_logging_obj") @@ -3366,6 +3379,7 @@ class ProxyBaseLLMRequestProcessing: ) ), ) + await capture_proxy_response(logging_obj, response) except (Exception, asyncio.CancelledError) as exc: _exception_raised = True if logging_obj.prepared_cost is not None: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 2951bc02d8c..203eaf5cc1c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -22,7 +22,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_parent_otel_span_from_kwargs, ) -from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup, UsageObservation from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.litellm_core_utils.redact_messages import MESSAGE_REDACTION_ENABLE_HEADERS from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping @@ -63,6 +63,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # sanitize_error_information_for_spend_logs, should_store_prompts_and_responses_in_spend_logs, ) +from litellm.proxy.spend_tracking.usage_capture import CapturedUsage, capture_callback, captured_callback_kwargs from litellm.proxy.utils import ( REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, @@ -344,11 +345,23 @@ def _attempt_spend_row_kwargs( "litellm_params": litellm_params, "combined_usage_object": attempt.usage, "response_cost": attempt.response_cost, + "_accounting_usage_observed": UsageObservation.ESTIMATED, } row.pop("billed_failed_attempts", None) + row.pop("_accounting_capture", None) call_id: Final = request_data.get("litellm_call_id") or litellm_params.get("litellm_call_id") if isinstance(call_id, str): row["litellm_call_id"] = f"{call_id}{suffix}" + row["parent_request_id"] = call_id + metadata["parent_request_id"] = call_id + captured_attempts: Final = request_data.get("_accounting_billed_attempts") + if isinstance(captured_attempts, tuple): + for captured in TypeAdapter(tuple[object, ...]).validate_python(captured_attempts): + if ( + isinstance(captured, CapturedUsage) + and captured.payload["litellm_call_id"] == row["litellm_call_id"] + ): + row["_accounting_capture"] = captured standard_logging_object: Final = request_data.get("standard_logging_object") if isinstance(standard_logging_object, dict): payload: Final = cast("Mapping[str, object]", standard_logging_object) # cast-ok: StandardLoggingPayload dict @@ -384,11 +397,12 @@ class _ProxyDBLogger(CustomLogger): async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime ) -> None: + captured_kwargs: Final = captured_callback_kwargs(kwargs, response_obj, start_time, end_time) if self.spend_event_producer is None or not is_offloadable_success(response_obj): - await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) + await self._PROXY_track_cost_callback(captured_kwargs, response_obj, start_time, end_time) return event: Final = build_spend_event( - kwargs, + captured_kwargs, response_obj, start_time, end_time, @@ -416,6 +430,7 @@ class _ProxyDBLogger(CustomLogger): user_api_key_dict: UserAPIKeyAuth, traceback_str: str | None = None, ) -> None: + self._try_capture_failure_usage(request_data, original_exception, user_api_key_dict) completed_cost: Final = ( _str_keyed_fields(request_data.get("standard_logging_object")).get("cost_finalized") is True ) @@ -595,6 +610,65 @@ class _ProxyDBLogger(CustomLogger): ), ) + def _try_capture_failure_usage( + self, + request_data: dict[str, object], # mutable-ok: receipt handoff shared with the legacy writer + original_exception: Exception, + auth: UserAPIKeyAuth, + ) -> None: + try: + self._capture_failure_usage(request_data, original_exception, auth) + except Exception: # noqa: BLE001 # capture failure must not suppress legacy billing or reservation release + verbose_proxy_logger.exception("Usage capture failed on the request failure path") + + def _capture_failure_usage( + self, + request_data: dict[str, object], # mutable-ok: receipt handoff shared with the legacy writer + original_exception: Exception, + auth: UserAPIKeyAuth, + ) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.db.exception_handler import exception_chain + from litellm.proxy.proxy_server import usage_capture + from litellm.proxy.spend_tracking.usage_capture import UsageAcknowledgementError + + if any(isinstance(error, UsageAcknowledgementError) for error in exception_chain(original_exception)): + return + usage_meter: Final = usage_capture.meter + + if usage_meter is None: + return + logging_obj: Final = request_data.get("litellm_logging_obj") + details: Final[Mapping[str, object]] = ( + _str_keyed_fields(logging_obj.model_call_details) if isinstance(logging_obj, Logging) else {} + ) + if details.get("call_type") == "IGNORE_THIS" and is_pre_routing_rejection(original_exception, request_data): + return + params: Final = _str_keyed_fields(request_data.get("litellm_params") or details.get("litellm_params")) + failure_data: Final = { + **details, + **request_data, + "litellm_params": { + **params, + "metadata": { + **_str_keyed_fields(params.get("metadata")), + **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=auth), + "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(auth), + "status": "failure", + }, + }, + } + start: Final = logging_obj.start_time if isinstance(logging_obj, Logging) else request_data.get("start_time") + if not isinstance(start, datetime): + return + end: Final = datetime.now(timezone.utc) + captured: Final = tuple( + capture_callback(usage_meter, row, original_exception, start, end) + for _, row in _failure_spend_rows(failure_data, 0.0, None, self._deployment_lookup) + ) + request_data["_accounting_capture"] = captured[0] # rebind-ok: share one receipt with the legacy writer + request_data["_accounting_billed_attempts"] = captured[1:] # rebind-ok: share attempt receipts + async def _write_failure_spend_rows( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b9b15d332e9..1b736f52330 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -839,6 +839,8 @@ from litellm.proxy.spend_tracking.spend_event_producer import ( SpendEventProducer, build_spend_event_producer, ) +from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, UsageCaptureRuntime +from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings from litellm.proxy.telemetry.endpoints import router as telemetry_router from litellm.proxy.telemetry.middleware import TelemetryMiddleware from litellm.proxy.telemetry.runtime import TelemetryRuntime @@ -1328,6 +1330,7 @@ async def _flush_spend_around_logging_drain_on_shutdown( finally: await worker.stop(discard_pending=True, timeout=GracefulShutdownManager.remaining(2.0, clock=clock)) await drain_tasks(GracefulShutdownManager.remaining(5.0, clock=clock)) + await usage_capture.stop(GracefulShutdownManager.remaining(clock=clock)) await flush_spend() @@ -3201,9 +3204,37 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): spend_event_producer: SpendEventProducer | None = None +usage_capture: Final = UsageCaptureRuntime() + + +def _configure_usage_capture() -> None: + settings: Final = _GENERAL_SETTINGS_VIEW.validate_python(general_settings) + capture_settings: Final = settings.get("accounting_capture") + if settings.get("accounting_sync_ack") and capture_settings is None: + raise ValueError("accounting_sync_ack requires accounting_capture") + if capture_settings is None: + if usage_capture.meter is not None: + raise ValueError("Disabling accounting_capture requires a worker restart") + return + if ( + prisma_client is None + or CollectorSettings().enabled + or os.getenv("SPEND_LOGS_URL") + or litellm.apply_search_free_tier + ): + raise ValueError( + "Usage capture requires local Prisma writes and disabled collector, remote writer and search free tier" + ) + from litellm.proxy.db.accounting_ledger import AccountingLedger + + 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)) + def cost_tracking(): global prisma_client, spend_event_producer + _configure_usage_capture() if prisma_client is not None: from litellm.integrations.shadow_eval_logger import ShadowEvalLogger diff --git a/litellm/proxy/spend_tracking/cost_finalization.py b/litellm/proxy/spend_tracking/cost_finalization.py index b260412331e..c42579714e8 100644 --- a/litellm/proxy/spend_tracking/cost_finalization.py +++ b/litellm/proxy/spend_tracking/cost_finalization.py @@ -89,6 +89,7 @@ async def _finish_finalization( if litellm.apply_search_free_tier else None ) + logging_obj.cost_allowance_allocation = allocation # rebind-ok: retain the allocation for durable capture logging_obj.cost_allowance_discount = ( # rebind-ok: retain allocation across snapshot retries allocation.discount if allocation is not None else 0.0 ) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 306e70f7029..86d3711754d 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -3,7 +3,7 @@ import json import os import re import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt from functools import reduce @@ -55,6 +55,8 @@ from litellm.litellm_core_utils.get_llm_provider_logic import declared_authentic from litellm.litellm_core_utils.hidden_params import served_from_cache from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + UsageObservation, coerce_model_access_groups, is_valid_sha256_hash, request_model_access_groups_from_litellm_params, @@ -599,6 +601,7 @@ def get_logging_payload( start_time: datetime, end_time: datetime, llm_router: "Router | None" = None, + capture_usage: Callable[[SpendLogsPayload, str | None], None] | None = None, ) -> SpendLogsPayload: if kwargs is None: kwargs = {} @@ -622,7 +625,9 @@ def get_logging_payload( routing_metadata, routing_litellm_params, ) - completion_start_time: Final = kwargs.get("completion_start_time", end_time) + completion_start_time: Final = TypeAdapter(datetime).validate_python( + kwargs.get("completion_start_time") or end_time + ) call_type: Final = kwargs.get("call_type") cache_hit: Final = kwargs.get("cache_hit", False) @@ -984,6 +989,14 @@ def get_logging_payload( payload.get("spend"), ) + if capture_usage is not None: + observed_usage: Final = ( + StandardLoggingPayloadSetup.get_usage_as_dict({"usage": usage}) + if usage and kwargs.get("_accounting_usage_observed") is not UsageObservation.UNKNOWN + else None + ) + capture_usage(payload, safe_dumps(observed_usage) if observed_usage is not None else None) + # Explicitly clear large intermediate objects to reduce memory pressure del response_obj_dict, usage, clean_metadata, additional_usage_values diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py new file mode 100644 index 00000000000..e0354a5d98f --- /dev/null +++ b/litellm/proxy/spend_tracking/usage_capture.py @@ -0,0 +1,791 @@ +from __future__ import annotations + +import asyncio +import json +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from decimal import Decimal +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, JsonValue, TypeAdapter + +from litellm._logging import verbose_proxy_logger +from litellm.accounting import ( + Dimension, + DimensionName, + InboundUsageEvent, + Metric, + MetricName, + OutboundUsageEvent, + Outcome, +) +from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.proxy.spend_tracking.usage_meter import ( + DurableUsageMeter, + UsageCaptureError, + UsageCaptureSettings, + usage_delivery_targets, +) +from litellm.types.utils import CallTypes, CustomPricingLiteLLMParams + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import SpendLogsPayload + + +from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite + + +class UsageAcknowledgementError(UsageCaptureError): + def __init__(self, cause: Exception) -> None: + super().__init__("Usage persistence did not acknowledge commit") + self.__cause__ = cause + + +@dataclass(frozen=True, slots=True) +class CapturedUsage: + payload: SpendLogsPayload + writes: tuple[UsageWrite, ...] + receipt: asyncio.Future[None] + + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +_FIELDS: Final = TypeAdapter(dict[str, object]) + + +_INSTANT: Final = TypeAdapter(datetime) + + +_DIMENSIONS: Final = ( + ("api_key", DimensionName.KEY), + ("user", DimensionName.USER), + ("team_id", DimensionName.TEAM), + ("organization_id", DimensionName.ORGANIZATION), + ("end_user", DimensionName.END_USER), + ("model", DimensionName.MODEL), + ("model_group", DimensionName.MODEL_GROUP), + ("model_id", DimensionName.DEPLOYMENT), + ("custom_llm_provider", DimensionName.PROVIDER), + ("mcp_namespaced_tool_name", DimensionName.TOOL), + ("agent_id", DimensionName.AGENT), +) + + +_USAGE_METRICS: Final = tuple( + metric + for metric in MetricName + if metric not in (MetricName.SPEND, MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS) +) + + +def _fields(value: object) -> Mapping[str, object]: + return _FIELDS.validate_python(value) if isinstance(value, Mapping) else {} + + +def _text(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _utc_time(value: datetime | str) -> datetime: + parsed: Final = _INSTANT.validate_python(value) + return parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed.astimezone(timezone.utc) + + +def _json(value: object) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def _usage_metrics(usage: Mapping[str, JsonValue] | None) -> tuple[Metric, ...] | None: + if usage is None: + return None + return tuple( + Metric(name=metric, value=Decimal(str(value))) + for metric in _USAGE_METRICS + if isinstance(value := usage.get(metric.value), (int, float)) and not isinstance(value, bool) + ) + + +@dataclass(frozen=True, slots=True) +class UsageCaptureContext: + caller_id: str | None + attempt_id: str + parent_request_id: str | None + retry_of_attempt_id: str | None + job_id: str | None + charge_eligible: bool + request_count_eligible: bool + ownership: str + pricing: str + finalized_cost: str | None = None + allocation: str | None = None + outcome: Outcome | None = None + started_at: datetime | None = None + ended_at: datetime | None = None + phase: str = "execution" + extra_dimensions: tuple[Dimension, ...] = () + provider_timestamps: str = "{}" + usage_source: str = "measured" + included_attempt_ids: tuple[str, ...] = () + + +def usage_writes( + payload: SpendLogsPayload, usage_json: str | None, context: UsageCaptureContext +) -> tuple[UsageWrite, ...]: + fields: Final = _FIELDS.validate_python(payload) + tags: Final = TypeAdapter(tuple[str, ...]).validate_json(payload["request_tags"]) + dimensions: Final = ( + *(Dimension(name=name, value=value) for field, name in _DIMENSIONS if (value := _text(fields.get(field)))), + *(Dimension(name=DimensionName.TAG, value=tag) for tag in dict.fromkeys(tags)), + *context.extra_dimensions, + ) + outcome: Final = context.outcome or (Outcome.FAILED if payload["status"] == "failure" else Outcome.SUCCEEDED) + started: Final = context.started_at or _utc_time(payload["startTime"]) + ended: Final = context.ended_at or _utc_time(payload["endTime"]) + usage: Final = None if usage_json is None else _JSON_OBJECT.validate_json(usage_json) + outbound: Final = OutboundUsageEvent( + event_id=context.attempt_id, + attempt_id=context.attempt_id, + request_id=context.caller_id, + retry_of_attempt_id=context.retry_of_attempt_id, + call_type=CallTypes(payload["call_type"]), + started_at=started, + ended_at=ended, + outcome=outcome, + dimensions=dimensions, + metrics=_usage_metrics(None if context.usage_source == "unknown" else usage), + ) + envelope: Final = _json( + { + "event": outbound.model_dump(mode="json"), + "usage": usage, + "usage_source": "unknown" if usage is None else context.usage_source, + "included_attempt_ids": context.included_attempt_ids, + "pricing": _JSON_OBJECT.validate_json(context.pricing), + "charge_eligible": context.charge_eligible, + "request_count_eligible": context.request_count_eligible, + "ownership": context.ownership, + "parent_request_id": context.parent_request_id, + "job_id": context.job_id, + "phase": context.phase, + "provider_timestamps": _JSON_OBJECT.validate_json(context.provider_timestamps), + "finalized_cost": None + if context.finalized_cost is None + else _JSON_OBJECT.validate_json(context.finalized_cost), + "allocation": None if context.allocation is None else _JSON_OBJECT.validate_json(context.allocation), + } + ) + outbound_write: Final = UsageWrite(outbound, envelope, usage_delivery_targets(dimensions)) + if not context.request_count_eligible or context.caller_id is None: + return (outbound_write,) + inbound: Final = InboundUsageEvent( + event_id=f"{context.caller_id}:inbound", + request_id=context.caller_id, + call_type=outbound.call_type, + started_at=started, + ended_at=ended, + outcome=outcome, + dimensions=dimensions, + metrics=( + Metric(name=MetricName.REQUESTS, value=Decimal(1)), + Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(int(outcome is Outcome.SUCCEEDED))), + Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(int(outcome is not Outcome.SUCCEEDED))), + ), + ) + return ( + UsageWrite(inbound, _json({"event": inbound.model_dump(mode="json")}), usage_delivery_targets(dimensions)), + outbound_write, + ) + + +_INCLUDED_IN_PARENT: Final = object() + + +_PRICING_FIELDS: Final = frozenset( + { + "service_tier", + "data_residency", + "vertex_location", + "region_name", + "base_model", + "custom_pricing", + "input_cost_per_token", + "output_cost_per_token", + "input_cost_per_second", + "output_cost_per_second", + "input_cost_per_character", + "output_cost_per_character", + "size", + "quality", + "n", + "duration", + "seconds", + "search_context_size", + "additional_response_cost", + } +) + + +def _selected_pricing(fields: Mapping[str, object]) -> Mapping[str, JsonValue]: + return _JSON_OBJECT.validate_python( + { + key: value + for key, value in fields.items() + if key in _PRICING_FIELDS or key in CustomPricingLiteLLMParams.model_fields + } + ) + + +def _response_fields(response: object) -> Mapping[str, object]: + if isinstance(response, BaseModel): + return _FIELDS.validate_python(response.model_dump()) + return _fields(response) + + +def _job_time(response: Mapping[str, object], key: str) -> datetime | None: + stamp: Final = response.get(key) + if isinstance(stamp, (int, float)) and not isinstance(stamp, bool): + return datetime.fromtimestamp(stamp, tz=timezone.utc) + return None + + +def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], response: object) -> UsageCaptureContext: + from dataclasses import asdict + + from litellm.litellm_core_utils.hidden_params import get_hidden_params + from litellm.litellm_core_utils.internal_call_metadata import is_background_response, is_unbilled_non_inference_call + from litellm.litellm_core_utils.litellm_logging import Logging, UsageObservation, mask_api_base_credentials + from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost + from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation + from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups + + metadata: Final = _JSON_OBJECT.validate_json(payload["metadata"]) + params: Final = _fields(kwargs.get("litellm_params")) + request_metadata: Final = _fields(params.get("litellm_metadata") or params.get("metadata")) + standard: Final = _fields(kwargs.get("standard_logging_object")) + logging_obj: Final = kwargs.get("litellm_logging_obj") + raw_response: Final = _response_fields(response) + hidden: Final = _fields(get_hidden_params(response)) + tool: Final = _fields(kwargs.get("mcp_tool_call_metadata")) + tool_pricing: Final = _fields(tool.get("mcp_server_cost_info")) + caller: Final = payload["litellm_call_id"] + cached: Final = payload["cache_hit"].lower() == "true" + included_in_parent: Final = kwargs.get("_accounting_internal_aggregate") is _INCLUDED_IN_PARENT + billed_id: Final = payload["request_id"] + terminal: Final = raw_response.get("status") in { + "completed", + "complete", + "failed", + "cancelled", + "expired", + "incomplete", + } + deferred: Final = billed_id.endswith(("_batch_cost", "_background_cost")) and terminal + management: Final = ( + (billed_id.endswith(("_batch_cost", "_background_cost")) and not terminal) + or (is_background_response(response) and not terminal) + or is_unbilled_non_inference_call(payload["call_type"], request_metadata, response) + ) + attempt: Final = ( + f"{caller}_cache_hit" + if cached and caller + else f"{caller}:status" + if management + else f"{billed_id}:included" + if included_in_parent + else billed_id + ) + if not attempt or attempt == "None" or (cached and not caller): + raise UsageCaptureError("Usage capture requires a stable execution identity") + retry: Final = metadata.get(BILLED_RETRY_ATTEMPT_METADATA_KEY) is True + origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) + parent: Final = _text(kwargs.get("parent_request_id")) or _text(request_metadata.get("parent_request_id")) + finalized_value: Final = ( + logging_obj.finalized_cost if isinstance(logging_obj, Logging) else kwargs.get("_accounting_finalized_cost") + ) + finalized: Final = finalized_value if isinstance(finalized_value, FinalizedCost) else None + allocation_value: Final = ( + logging_obj.cost_allowance_allocation + if isinstance(logging_obj, Logging) + else kwargs.get("_accounting_allocation") + ) + allocation: Final = allocation_value if isinstance(allocation_value, SearchAllowanceAllocation) else None + pricing_snapshot: Final = kwargs.get("_accounting_pricing_snapshot") + captured_pricing: Final = { + "model_map_information": ( + _JSON_OBJECT.validate_json(pricing_snapshot.model_map_information) + if isinstance(pricing_snapshot, PricingCapture) + else metadata.get("model_map_information") + ), + "cost_breakdown": metadata.get("cost_breakdown"), + **( + _JSON_OBJECT.validate_json(pricing_snapshot.inputs) + if isinstance(pricing_snapshot, PricingCapture) + else _pricing_inputs(kwargs, logging_obj, response) + ), + "tool": { + "name": tool.get("name"), + "cost_per_query": _fields(tool_pricing.get("tool_name_to_cost_per_query")).get( + _text(tool.get("name")) or "", tool_pricing.get("default_cost_per_query") + ), + } + if tool + else None, + "response_inputs": { + "audio_transcription_duration": hidden.get("audio_transcription_duration", raw_response.get("duration")), + "billed_units": _fields(raw_response.get("meta")).get("billed_units"), + "batch_models": hidden.get("batch_models"), + "batch_successful_requests": hidden.get("batch_successful_requests"), + "batch_failed_requests": hidden.get("batch_failed_requests"), + "service_tier": raw_response.get("service_tier"), + "region_name": hidden.get("region_name"), + }, + "reported_cost": standard.get( + "response_cost", kwargs.get("response_cost", _fields(get_hidden_params(response)).get("response_cost")) + ), + "adjustments": ( + _JSON_OBJECT.validate_json(pricing_snapshot.adjustments) + if isinstance(pricing_snapshot, PricingCapture) + else None + ), + } + return UsageCaptureContext( + caller_id=None if deferred else parent if retry or included_in_parent else caller, + attempt_id=attempt, + parent_request_id=parent, + retry_of_attempt_id=_text(request_metadata.get("retry_of_attempt_id")), + job_id=_text(raw_response.get("id")) if deferred or management else None, + charge_eligible=not cached and not included_in_parent and not management, + request_count_eligible=bool(caller) and not origin and not retry and not deferred and not included_in_parent, + ownership="included_in_parent" + if included_in_parent + else "aggregate" + if kwargs.get("additional_response_cost") is not None + else "attempt", + pricing=_json(captured_pricing), + finalized_cost=_json( + { + "response_cost": finalized.amount, + "cost_breakdown": finalized.logging_breakdown(), + "allowance_discount": finalized.allowance_discount, + } + ) + if finalized is not None + else ( + _json({"response_cost": standard.get("response_cost"), "cost_breakdown": standard.get("cost_breakdown")}) + if standard.get("cost_finalized") is True + else None + ), + allocation=_json(asdict(allocation)) if allocation is not None else None, + started_at=_job_time(raw_response, "created_at") if deferred else None, + ended_at=( + _job_time(raw_response, "completed_at") + or _job_time(raw_response, "failed_at") + or _job_time(raw_response, "cancelled_at") + or _job_time(raw_response, "expired_at") + ) + if deferred + else None, + outcome=( + Outcome.CANCELLED + if raw_response.get("status") == "cancelled" + or isinstance(response, asyncio.CancelledError) + or getattr(response, "status_code", None) == 499 + else Outcome.FAILED + if raw_response.get("status") in {"failed", "expired", "incomplete"} or isinstance(response, BaseException) + else None + ), + phase="terminal" if deferred else "management" if management else "execution", + usage_source="estimated" + if kwargs.get("_accounting_usage_observed") is UsageObservation.ESTIMATED + else "measured", + included_attempt_ids=TypeAdapter(tuple[str, ...]).validate_python( + kwargs.get("_accounting_included_attempt_ids", ()) + ), + provider_timestamps=_json( + { + key: raw_response.get(key) + for key in ("created_at", "completed_at", "failed_at", "cancelled_at", "expired_at") + if raw_response.get(key) is not None + } + ) + if deferred + else "{}", + extra_dimensions=( + *( + (Dimension(name=DimensionName.TOOL, value=tool_name),) + if not payload.get("mcp_namespaced_tool_name") + and (tool_name := _text(tool.get("namespaced_tool_name"))) + else () + ), + *( + (Dimension(name=DimensionName.API_BASE, value=mask_api_base_credentials(api_base)),) + if (api_base := _text(payload["api_base"])) + else () + ), + *( + Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value=group) + for group in get_request_model_access_groups(kwargs) + ), + *( + (Dimension(name=DimensionName.PROJECT, value=project),) + if (project := _text(request_metadata.get("user_api_key_project_id"))) + else () + ), + *( + (Dimension(name=DimensionName.SERVICE_TIER, value=tier),) + if (tier := _text(raw_response.get("service_tier") or kwargs.get("service_tier"))) + else () + ), + ), + ) + + +class _CaptureBuilder: + def __init__(self, kwargs: Mapping[str, object], response: object) -> None: + self.kwargs: Final = kwargs + self.response: Final = response + self.writes: tuple[UsageWrite, ...] = () + + def observe(self, payload: SpendLogsPayload, usage: str | None) -> None: + parsed: Final = _JSON_OBJECT.validate_json(usage) if usage is not None else {} + batch: Final = _fields(parsed.get("batch_usage_by_model")) + unknown: Final = (bool(batch) or bool(parsed.get("batch_unknown_usage_requests"))) and all( + value is None for value in batch.values() + ) + observed: Final = ( + _json( + { + "batch_usage_by_model": batch, + "batch_unknown_usage_requests": parsed.get("batch_unknown_usage_requests"), + } + ) + if unknown + else usage + ) + context: Final = capture_context(payload, self.kwargs, self.response) + self.writes = usage_writes(payload, observed, replace(context, usage_source="unknown") if unknown else context) + + +def capture_callback( + meter: DurableUsageMeter, + kwargs: Mapping[str, object], + response: object, + start_time: datetime, + end_time: datetime, +) -> CapturedUsage: + from litellm.proxy.db.db_spend_update_writer import get_llm_router + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + existing: Final = kwargs.get("_accounting_capture") + if isinstance(existing, CapturedUsage): + return existing + builder: Final = _CaptureBuilder(kwargs, response) + payload: Final = get_logging_payload( + dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe + ) + return CapturedUsage(payload, builder.writes, meter.submit(builder.writes)) + + +def captured_callback_kwargs( + kwargs: Mapping[str, object], response: object, start_time: datetime, end_time: datetime +) -> Mapping[str, object]: + from litellm.proxy.proxy_server import usage_capture + + usage_meter: Final = usage_capture.meter + + if usage_meter is None: + return kwargs + try: + return {**kwargs, "_accounting_capture": capture_callback(usage_meter, kwargs, response, start_time, end_time)} + except Exception: # noqa: BLE001 # capture failure must not suppress legacy billing + verbose_proxy_logger.exception("Usage capture failed before legacy cost tracking") + return kwargs + + +def validate_sync_ack_request( + settings: Mapping[str, object], call_type: str, streaming: bool, background: bool = False +) -> None: + if settings.get("accounting_sync_ack") is not True: + return + if streaming or background or call_type not in {"acompletion", "atext_completion", "aembedding", "aresponses"}: + from fastapi import HTTPException + + raise HTTPException( + 400, "accounting_sync_ack supports foreground non-stream chat, text, embedding and Responses requests" + ) + + +async def capture_proxy_response(logging_obj: object, response: object) -> None: + from fastapi import HTTPException + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.proxy_server import general_settings, usage_capture + + settings: Final = _FIELDS.validate_python(general_settings) + usage_meter: Final = usage_capture.meter + + if usage_meter is None or not isinstance(logging_obj, Logging): + return + try: + capture_logging_usage(usage_meter, logging_obj, response, logging_obj.start_time, datetime.now(timezone.utc)) + captured: Final = logging_obj.model_call_details.get("_accounting_capture") + if settings.get("accounting_sync_ack") is True: + if not isinstance(captured, CapturedUsage): + raise UsageCaptureError("Usage capture was not admitted") + timeout: Final = TypeAdapter(float).validate_python( + settings.get("accounting_sync_ack_timeout_seconds", 5.0) + ) + await asyncio.wait_for(asyncio.shield(captured.receipt), timeout=timeout) + 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 ( + settings.get("accounting_sync_ack") is True + and settings.get("accounting_sync_ack_failure_policy", "fail_closed") == "fail_closed" + ): + acknowledgement: Final = UsageAcknowledgementError(error) + raise HTTPException(503, str(acknowledgement)) from acknowledgement + + +def capture_logging_usage( + meter: DurableUsageMeter, + logging_obj: object, + response: object, + start_time: datetime, + end_time: datetime, + freeze_only: bool = False, +) -> None: + from litellm._internal_context import internal_parent_request_id, is_internal_call + from litellm.batches.batch_utils import batch_cost_is_final + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import LiteLLMBatch + + if not isinstance(logging_obj, Logging): + return + internal: Final = is_internal_call.get() + if isinstance(response, LiteLLMBatch) and batch_cost_is_final(response) and response.usage is None: + return + try: + if not internal and isinstance(logging_obj.model_call_details.get("_accounting_capture"), CapturedUsage): + logging_obj.record_usage_capture( + capture_callback(meter, logging_obj.model_call_details, response, start_time, end_time) + ) + return + normalized_result: Final = logging_obj.usage_capture_response(response) + prepared: Final = logging_obj.prepared_cost + normalized: Final = ( + {**_response_fields(normalized_result), "usage": prepared.usage.model_dump()} + if prepared is not None + else normalized_result + ) + if normalized is None: + return + details: Final = _FIELDS.validate_python(logging_obj.model_call_details) + params: Final = _fields(details.get("litellm_params")) + metadata: Final = _fields(params.get("litellm_metadata") or params.get("metadata")) + if metadata.get("batch_ignore_default_logging") is True: + return + existing_pricing: Final = details.get("_accounting_pricing_snapshot") + pricing: Final = ( + existing_pricing + if isinstance(existing_pricing, PricingCapture) and not internal + else _capture_pricing(logging_obj, normalized, details) + ) + if freeze_only: + if not internal: + logging_obj.record_usage_pricing(pricing) + return + captured: Final = capture_callback( + meter, + { + **details, + "litellm_logging_obj": logging_obj, + "_accounting_capture": None if internal else details.get("_accounting_capture"), + "_accounting_pricing_snapshot": pricing, + "_accounting_internal_aggregate": _INCLUDED_IN_PARENT if internal else None, + "parent_request_id": internal_parent_request_id.get() if internal else details.get("parent_request_id"), + }, + normalized, + start_time, + end_time, + ) + if not internal: + logging_obj.record_usage_capture(captured) + except Exception: # noqa: BLE001 # failed capture must not prevent normal logging and legacy charging + verbose_proxy_logger.exception("Usage capture failed at the execution boundary") + + +@dataclass(frozen=True, slots=True) +class PricingCapture: + model_map_information: str + adjustments: str + inputs: str + + +def _capture_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 + + params: Final = _fields(details.get("litellm_params")) + prepared: Final = logging_obj.prepared_cost + provider: Final = _text(details.get("custom_llm_provider")) + builtins: Final = logging_obj.standard_built_in_tools_params + pricing: Final = ( + prepared.payload["model_map_information"] + 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_llm_provider=_text(details.get("custom_llm_provider")), + init_response_obj=normalized, + api_base=_text(params.get("api_base")), + router_model_id=logging_obj.get_router_model_id(), + ) + ) + adjustments: Final = _json( + { + "discount": resolve_cost_discount( + litellm.cost_discount_config, + provider, + logging_obj.get_deployment_model_for_cost(), + _text(params.get("region_name") or params.get("vertex_location")), + ), + "margin": litellm.cost_margin_config.get(provider or "", litellm.cost_margin_config.get("global")), + "web_search_options": ( + _selected_pricing(_fields(builtins.get("web_search_options"))) + if builtins.get("web_search_options") is not None + else None + ), + "file_search": builtins.get("file_search") is not None, + } + ) + batch_usage: Final = _fields(_fields(_response_fields(normalized).get("usage")).get("batch_usage_by_model")) + return PricingCapture( + _json(pricing), + adjustments, + _json( + { + **_pricing_inputs(details, logging_obj, normalized), + "batch_models": { + model: StandardLoggingPayloadSetup.get_model_cost_information( + base_model=model, + custom_pricing=bool(params.get("custom_pricing")), + custom_llm_provider=provider, + init_response_obj={"model": model}, + api_base=_text(params.get("api_base")), + router_model_id=logging_obj.get_router_model_id(), + ) + for model in batch_usage + }, + } + ), + ) + + +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 { + "request": _selected_pricing(kwargs), + "deployment": _selected_pricing(_fields(kwargs.get("litellm_params"))), + "optional_params": _selected_pricing(_fields(kwargs.get("optional_params"))), + "logging": logging_obj.usage_capture_pricing(response) if isinstance(logging_obj, Logging) else None, + "model_info": _selected_pricing( + _fields( + _fields( + _fields(kwargs.get("litellm_params")).get("litellm_metadata") + or _fields(kwargs.get("litellm_params")).get("metadata") + ).get("model_info") + ) + ), + } + + +def _terminal_facts(write: UsageWrite) -> Mapping[str, JsonValue] | None: + payload: Final = _JSON_OBJECT.validate_json(write.envelope) + if payload.get("phase") != "terminal": + return None + event: Final = _JSON_OBJECT.validate_python(payload["event"]) + return { + **{ + key: value + for key, value in payload.items() + if key not in {"event", "pricing", "finalized_cost", "allocation"} + }, + "event": {key: value for key, value in event.items() if key not in {"started_at", "ended_at"}}, + } + + +def _reuse_terminal_snapshot(candidate: UsageWrite, accepted: UsageWrite) -> UsageWrite: + facts: Final = _terminal_facts(candidate) + if facts is None or facts != _terminal_facts(accepted) or set(candidate.targets) != set(accepted.targets): + raise AccountingConflictError("Repeated terminal job disagrees with its accepted usage or attribution") + return accepted + + +class UsageBatchWriter: + def __init__( + self, + insert: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], + load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]], + ) -> None: + self._insert: Final = insert + self._load: Final = load + + async def _restore(self, writes: tuple[UsageWrite, ...]) -> tuple[UsageWrite, ...]: + terminal: Final = tuple(write for write in writes if _terminal_facts(write) is not None) + if not terminal: + return writes + identities: Final = tuple(dict.fromkeys(write.event.event_id for write in terminal)) + stored: Final = await self._load(identities) + # Re-read the first accepted observation and pricing before the ledger's unchanged id/hash check. + snapshots: Final = {write.event.event_id: write for write in (*reversed(terminal), *stored)} + return tuple( + _reuse_terminal_snapshot(write, snapshots[write.event.event_id]) + if write.event.event_id in snapshots + else write + for write in writes + ) + + async def write(self, writes: tuple[UsageWrite, ...]) -> None: + restored: Final = await self._restore(writes) + try: + await self._insert(restored) + except AccountingConflictError: + if not any(_terminal_facts(write) is not None for write in writes): + raise + await self._insert(await self._restore(writes)) + + +class UsageCaptureRuntime: + def __init__(self) -> None: + self.meter: DurableUsageMeter | None = None + + def start( + self, write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], settings: UsageCaptureSettings + ) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + + if self.meter is not None: + if self.meter.settings != settings: + raise ValueError("Changing accounting_capture requires a worker restart") + return + meter: Final = DurableUsageMeter(write_batch, settings) + self.meter = meter + Logging.usage_capture_callback = lambda logger, response, start, end, freeze: capture_logging_usage( + meter, logger, response, start, end, freeze + ) + + async def stop(self, timeout: float) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + + if self.meter is None: + return + try: + await self.meter.stop(timeout) + finally: + Logging.usage_capture_callback = None + self.meter = None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9e237ffcb57..e2f32ef4f96 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -163,7 +163,12 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.deployment_pricing import pricing_models from litellm.litellm_core_utils.duration_parser import get_budget_window_start from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int -from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info +from litellm.litellm_core_utils.litellm_logging import ( + BilledRoundsDispatch, + Logging, + UsageObservation, + deployment_pricing_model_info, +) from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -1747,6 +1752,27 @@ def _failure_fields_to_lift( ("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)), ("standard_logging_object", _model_call_details.get("standard_logging_object")), ("billed_failed_attempts", _earlier_attempts or None), + ("_accounting_capture", _model_call_details.get("_accounting_capture")), + ("_accounting_pricing_snapshot", _model_call_details.get("_accounting_pricing_snapshot")), + ("_accounting_finalized_cost", _logging_obj.finalized_cost if isinstance(_logging_obj, Logging) else None), + ( + "_accounting_allocation", + _logging_obj.cost_allowance_allocation if isinstance(_logging_obj, Logging) else None, + ), + ( + "_accounting_usage_observed", + ( + None + if Logging.usage_capture_callback is None + else UsageObservation.UNKNOWN + if _usage_to_lift is None + else UsageObservation.KNOWN + if _model_call_details.get("combined_usage_object") is not None + or _model_call_details.get("failed_attempts_usage") + else UsageObservation.ESTIMATED + ), + ), + ("_accounting_included_attempt_ids", _model_call_details.get("_accounting_prior_attempt_ids")), ) return MappingProxyType({key: value for key, value in _entries if value is not None}) diff --git a/litellm/rag/main.py b/litellm/rag/main.py index be566a0da8b..77bf6868bf1 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -223,6 +223,10 @@ async def _execute_query_pipeline( """ Execute the RAG query pipeline. """ + parent_request_id: Final[str | None] = TypeAdapter[str | None](str | None).validate_python( + TypeAdapter(dict[str, object]).validate_python(kwargs).get("litellm_call_id") + ) + # Extract router from kwargs - use it for completion if available # to properly resolve virtual model names router: Final[Router | None] = kwargs.pop("router", None) @@ -252,7 +256,7 @@ async def _execute_query_pipeline( forwarded_search_params: Final = MappingProxyType( {**provider_search_params, **kwargs, **filter_search_params, **store_search_params} ) - with internal_sub_call(): + with internal_sub_call(parent_request_id): search_response: Final = await litellm.vector_stores.asearch( vector_store_id=retrieval_config["vector_store_id"], query=query_text, @@ -286,7 +290,7 @@ async def _execute_query_pipeline( if rerank and rerank.get("enabled"): documents: Final = RAGQuery.extract_documents_from_search(search_response) if documents: - with internal_sub_call(): + with internal_sub_call(parent_request_id): rerank_response = await litellm.arerank( model=rerank["model"], query=query_text, @@ -304,7 +308,7 @@ async def _execute_query_pipeline( modified_messages: Final = messages[:-1] + [context_message] + [messages[-1]] # Use router if available to properly resolve virtual model names - with internal_sub_call(): + with internal_sub_call(parent_request_id): if router is not None: response = await router.acompletion( model=model, diff --git a/litellm/responses/main.py b/litellm/responses/main.py index b7b93a20b4e..ad1d35bbe17 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -437,7 +437,7 @@ async def aresponses_api_with_mcp( ######################################################### if isinstance(logging_obj, LiteLLMLoggingObj): logging_obj.carry_failed_attempt_usage() - with internal_sub_call(): + with internal_sub_call(logging_obj.litellm_call_id if isinstance(logging_obj, LiteLLMLoggingObj) else None): response: Final = await aresponses( input=input, model=model, @@ -514,7 +514,9 @@ async def aresponses_api_with_mcp( tool_calls=tool_calls, tool_results=tool_results ) - with internal_sub_call(): + with internal_sub_call( + logging_obj.litellm_call_id if isinstance(logging_obj, LiteLLMLoggingObj) else None + ): final_response = await LiteLLM_Proxy_MCP_Handler.make_follow_up_call( follow_up_input=follow_up_input, model=model, diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 1e07b768f62..f8a52227d1c 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -663,7 +663,7 @@ async def acompletion_with_mcp( logging_obj: Final = kwargs.get("litellm_logging_obj") if isinstance(logging_obj, LiteLLMLoggingObj): logging_obj.carry_failed_attempt_usage() - with internal_sub_call(): + with internal_sub_call(logging_obj.litellm_call_id if isinstance(logging_obj, LiteLLMLoggingObj) else None): initial_response: Final = await litellm_acompletion(**initial_call_args) first_round_breakdown: Final = logging_obj.cost_breakdown if isinstance(logging_obj, LiteLLMLoggingObj) else None @@ -713,7 +713,7 @@ async def acompletion_with_mcp( follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream - with internal_sub_call(): + with internal_sub_call(logging_obj.litellm_call_id if isinstance(logging_obj, LiteLLMLoggingObj) else None): final_response: Final = await litellm_acompletion(**follow_up_call_args) if not isinstance(final_response, ModelResponse): diff --git a/litellm/router.py b/litellm/router.py index 535072d5fcf..dd369096c74 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3735,6 +3735,11 @@ class Router(DeprecatedRouterMembers): priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost" } + if LiteLLMLogging.usage_capture_callback is not None: + from litellm._internal_context import internal_sub_call + + with internal_sub_call(cut_logging.litellm_call_id): + cut_logging.capture_usage_before_logging(priced_partial) return ContinuedPartial( usage=BaseTokenUsageProcessor.combine_usage_objects( usage_objects=[partial_usage, inherited_usage] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 0c3440e88f7..75d3f682fc9 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -676,6 +676,12 @@ class CallTypes(str, Enum): retrieve_fine_tuning_job = "retrieve_fine_tuning_job" responses = "responses" aresponses = "aresponses" + get_responses = "get_responses" + aget_responses = "aget_responses" + delete_responses = "delete_responses" + adelete_responses = "adelete_responses" + cancel_responses = "cancel_responses" + acancel_responses = "acancel_responses" alist_input_items = "alist_input_items" llm_passthrough_route = "llm_passthrough_route" allm_passthrough_route = "allm_passthrough_route" diff --git a/litellm/utils.py b/litellm/utils.py index 8fe3fd84fbd..fcb2d5f6232 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1595,6 +1595,7 @@ def dispatch_success_logging( is_litellm_internal_call: bool, ) -> None: if is_litellm_internal_call: + logging_obj.capture_usage_before_logging(result, start_time, end_time) return _schedule_async_success_logging( @@ -1633,9 +1634,14 @@ def _schedule_async_success_logging( if not logging_obj.claim_async_success_log(): return + if logging_obj.defer_async_logging and not is_completion_with_fallbacks: + logging_obj.capture_usage_before_logging(result, start_time, end_time, freeze_only=True) + def _enqueue_async_logging() -> None: from litellm.responses.streaming_iterator import spawn_logging_task + if not is_completion_with_fallbacks: + logging_obj.capture_usage_before_logging(result, start_time, end_time) if logging_obj.defer_async_logging: logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, start_time=start_time, end_time=end_time @@ -2203,6 +2209,7 @@ def client(original_function): # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") + logging_obj.capture_usage_before_logging(TypeAdapter(object).validate_python(result), start_time, end_time) if not is_internal_call.get(): # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index d154086aec9..a9e1e0c2a35 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -3436,3 +3436,50 @@ def test_invalid_batch_envelopes_preserve_recoverable_cost( assert result.usage.prompt_tokens == (200 if include_valid else 0) assert result.usage.completion_tokens == (40 if include_valid else 0) assert (result.successful_requests, result.failed_requests) == (2 if include_valid else 0, 1) + + +def test_unpriced_batch_retains_reported_usage_by_model_and_marks_missing_usage(): + rows: Final = [ + _success_row(model="unregistered-batch-model-a", usage=_usage(4, 2)), + _success_row(model="unregistered-batch-model-a", usage=_usage(3, 1)), + _success_row(model="unregistered-batch-model-b", usage=_usage(5, 2)), + {"response": {"status_code": 200, "body": {"model": "unregistered-batch-model-c"}}}, + ] + result: Final = bu.batch_cost_and_usage(rows, "openai") + captured: Final = result.usage.model_dump() + assert (result.usage.prompt_tokens, result.usage.completion_tokens) == (12, 5) + assert result.cost == 0 + by_model: Final = captured["batch_usage_by_model"] + assert by_model["unregistered-batch-model-a"]["total_tokens"] == 10 + assert by_model["unregistered-batch-model-b"]["total_tokens"] == 7 + assert by_model["unregistered-batch-model-c"] is None + assert captured["batch_unknown_usage_requests"] == 1 + + +def test_batch_ocr_preserves_additive_page_quantities_without_changing_cost(): + rows: Final = [ + { + "response": { + "status_code": 200, + "body": { + "model": "ocr-quantity-control", + "usage_info": {"pages_processed": pages, "doc_size_bytes": size}, + }, + } + } + for pages, size in ((2, 100), (3, 200)) + ] + result: Final = bu.batch_cost_and_usage(rows, "mistral", model_info={"ocr_cost_per_page_batches": 0.01}) + captured: Final = result.usage.model_dump()["batch_usage_by_model"]["ocr-quantity-control"] + assert captured["pages_processed"] == 5 + assert captured["doc_size_bytes"] == 300 + assert result.cost == pytest.approx(5 * 0.01) + + +def test_native_vertex_capture_retains_unknown_models_without_changing_legacy_counts(): + rows: Final = [{"request": {}, "response": {"modelVersion": "missing-usage-model", "candidates": []}}] + result: Final = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + captured: Final = result.usage.model_dump() + assert captured["batch_usage_by_model"] == {"missing-usage-model": None} + assert captured["batch_unknown_usage_requests"] == 1 + assert (result.cost, result.successful_requests, result.failed_requests) == (0, 0, 1) diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 8f083298555..3f9da8be2d4 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -467,6 +467,9 @@ def test_default_max_retries_env_var_reaches_azure_sdk_client(): "add_message", "arun_thread_stream", "aresponses", + "aget_responses", + "adelete_responses", + "acancel_responses", "aresponses_websocket", "alist_input_items", "acreate_fine_tuning_job", diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index e482b6552a4..8bc55053a08 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -3949,3 +3949,53 @@ async def test_async_post_call_failure_hook_disable_pre_routing_error_logs( user_api_key_dict=user_api_key_dict, ) assert len(writer.rows) == rows_written + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "streaming", "background"), + [("acompletion", True, False), ("aresponses", False, True), ("aocr", False, False)], +) +async def test_sync_ack_rejection_before_setup_does_not_capture_an_ignored_route( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + call_type: str, + streaming: bool, + background: bool, +) -> None: + import asyncio + from typing import Final + + from fastapi import HTTPException + + from litellm.proxy import proxy_server + from litellm.proxy.db.accounting_ledger import UsageWrite + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, validate_sync_ack_request + from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings + from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY + from tests.unit.proxy.spend_tracking.test_usage_capture import _logger + + persisted: Final = asyncio.Queue[tuple[UsageWrite, ...]]() + + async def persist(writes: tuple[UsageWrite, ...]) -> None: + persisted.put_nowait(writes) + + runtime: Final = UsageCaptureRuntime() + runtime.start(persist, UsageCaptureSettings()) + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + monkeypatch.setattr(proxy_server, "prisma_client", None) + logger: Final = _logger("IGNORE_THIS") + request: Final = { + "litellm_logging_obj": logger, + REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP, + } + try: + with pytest.raises(HTTPException) as rejected: + validate_sync_ack_request({"accounting_sync_ack": True}, call_type, streaming, background) + assert rejected.value.status_code == 400 + assert "supports foreground non-stream" in rejected.value.detail + await ProxyDBLogger().async_post_call_failure_hook(request, rejected.value, UserAPIKeyAuth()) + finally: + await runtime.stop(1) + assert persisted.empty() + assert "Usage capture failed on the request failure path" not in caplog.text diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index aefc4a96a28..c9515c1232e 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -91,12 +91,20 @@ def response(cost: float = 0.5) -> ModelResponse: async def test_header_and_logging_share_one_allowance_allocation( logging_obj: Logging, used: int, expected: float ) -> None: + import json + + from litellm.proxy.db.accounting_ledger import UsageWrite + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings + counter: Final = Counter(used) result: Final = response() costs: Final = await asyncio.gather(*(finalize_response_cost(logging_obj, result, counter, NOW) for _ in range(3))) assert costs[0] is costs[1] is costs[2] assert costs[0] is not None assert costs[0].amount == pytest.approx(expected) + assert logging_obj.cost_allowance_allocation is not None + assert logging_obj.cost_allowance_allocation.discount == pytest.approx(costs[0].allowance_discount) assert counter.calls == 1 logging_obj._success_handler_helper_fn(result=result, start_time=NOW, end_time=NOW) payload: Final = logging_obj.model_call_details["standard_logging_object"] @@ -107,6 +115,25 @@ async def test_header_and_logging_share_one_allowance_allocation( assert payload["cost_finalized"] is True assert payload["cost_breakdown"] == costs[0].logging_breakdown() + committed: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue() + + async def persist(writes: tuple[UsageWrite, ...]) -> None: + for write in writes: + committed.put_nowait(write) + + meter: Final = DurableUsageMeter(persist, UsageCaptureSettings()) + assert logging_obj.prepared_cost is not None + prepared_usage: Final = logging_obj.prepared_cost.usage.model_dump() + result.usage = Usage(prompt_tokens=999) + capture_logging_usage(meter, logging_obj, result, NOW, NOW) + await meter.close() + committed.get_nowait() + event: Final = json.loads(committed.get_nowait().envelope) + assert event["finalized_cost"]["response_cost"] == pytest.approx(costs[0].amount) + assert event["allocation"]["discount"] == pytest.approx(costs[0].allowance_discount) + assert event["usage"] == prepared_usage + assert counter.calls == 1 + @pytest.mark.asyncio async def test_failed_counter_freezes_full_price_without_a_retry(logging_obj: Logging) -> None: diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index a91cdd512fe..e2fac424b7d 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -7543,3 +7543,21 @@ def test_get_logging_payload_stores_end_user_only_while_end_user_cost_tracking_i end_time=datetime.datetime.now(timezone.utc), ) assert payload["end_user"] == expected_end_user + + +@pytest.mark.parametrize("tool_metadata", [{"name": "tool"}, {"name": "tool", "arguments": "invalid"}]) +def test_optional_mcp_metadata_does_not_change_legacy_spend_payload(tool_metadata: Mapping[str, object]) -> None: + now: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc) + kwargs: Final = { + "call_type": "call_mcp_tool", + "litellm_call_id": "tool-call", + "litellm_params": {}, + "litellm_trace_id": "fixed-trace", + "response_cost": 0.5, + } + response: Final = {"id": "tool-call"} + expected: Final = spend_tracking_utils.get_logging_payload(kwargs, response, now, now) + actual: Final = spend_tracking_utils.get_logging_payload( + {**kwargs, "mcp_tool_call_metadata": tool_metadata}, response, now, now + ) + assert actual == expected diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py new file mode 100644 index 00000000000..1f4f3d6da35 --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_usage_capture.py @@ -0,0 +1,882 @@ +import asyncio +import json +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.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 +from litellm.types.utils import CallTypes +from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier, _manual_flush, _write + + +@pytest.mark.asyncio +@pytest.mark.parametrize("usage", [None, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}]) +async def test_capture_keeps_unknown_usage_distinct_from_measured_zero(usage: Mapping[str, int] | None) -> None: + from litellm.proxy.spend_tracking.usage_capture import capture_callback + + database: Final = CommitBarrier() + database.release.set() + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + captured: Final = capture_callback( + meter, + {"call_type": "acompletion", "litellm_call_id": "caller", "model": "unknown-model", "litellm_params": {}}, + {"id": "provider", "usage": usage}, + timestamp, + timestamp, + ) + await meter.close() + outbound: Final = captured.writes[1].event + assert outbound.metrics == ( + None + if usage is None + else tuple(Metric(name=MetricName(key), value=Decimal(value)) for key, value in usage.items()) + ) + envelope: Final = json.loads(captured.writes[1].envelope) + assert envelope["usage"] == usage + assert envelope["pricing"]["reported_cost"] is None + assert database.committed == ("caller:inbound", "provider") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type", + [ + call.value + for call in ( + CallTypes.acompletion, + CallTypes.arealtime, + CallTypes.call_mcp_tool, + CallTypes.asend_message, + CallTypes.pass_through, + ) + ], +) +async def test_shared_callback_capture_retains_usage_and_attribution_without_bodies(call_type: str) -> None: + from litellm.proxy.spend_tracking.usage_capture import capture_callback + + database: Final = CommitBarrier() + database.release.set() + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + captured: Final = capture_callback( + meter, + { + "call_type": call_type, + "litellm_call_id": "caller", + "model": "test-model", + "messages": "private prompt", + "litellm_params": {"api_key": "private-secret", "metadata": {"user_api_key_team_id": "team"}}, + }, + { + "id": "provider", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "output": "private answer", + }, + timestamp, + timestamp, + ) + await meter.close() + envelope: Final = captured.writes[1].envelope + assert all(secret not in envelope for secret in ("private prompt", "private-secret", "private answer")) + event: Final = captured.writes[1].event + assert event.call_type.value == call_type + assert event.metrics == ( + Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(3)), + Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), + Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(5)), + ) + assert any(dimension.value == "team" for dimension in event.dimensions) + + +@pytest.mark.asyncio +async def test_cache_duplicate_uses_caller_identity_despite_timestamped_diagnostic_ids() -> None: + from litellm.proxy.spend_tracking.usage_capture import capture_callback + + database: Final = CommitBarrier() + database.release.set() + meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + kwargs: Final = { + "call_type": "acompletion", + "litellm_call_id": "caller", + "model": "test-model", + "cache_hit": True, + "litellm_params": {}, + } + response: Final = {"id": "cached-provider-id", "usage": {"total_tokens": 5}} + first: Final = capture_callback(meter, kwargs, response, timestamp, timestamp) + second: Final = capture_callback(meter, kwargs, response, timestamp, timestamp) + await meter.close() + assert first.writes == second.writes + assert first.writes[1].event.event_id == "caller_cache_hit" + assert json.loads(first.writes[1].envelope)["charge_eligible"] is False + + +NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +class StoredUsage: + def __init__(self) -> None: + self.rows: tuple[UsageWrite, ...] = () + + async def load(self, identities: tuple[str, ...]) -> tuple[UsageWrite, ...]: + return tuple(row for row in self.rows if row.event.event_id in identities) + + async def insert(self, writes: tuple[UsageWrite, ...]) -> None: + existing: Final = {row.event.event_id: row for row in self.rows} + if any(existing.get(row.event.event_id, row) != row for row in writes): + raise AccountingConflictError("Same identity has different content") + self.rows += tuple(row for row in writes if row.event.event_id not in existing) + + +def _logger(call_type: str = "acompletion", *, stream: bool = False) -> Logging: + from litellm.litellm_core_utils.litellm_logging import Logging + + logger: Final = Logging( + model="capture-model", + messages=[], + stream=stream, + call_type=call_type, + start_time=NOW, + litellm_call_id="caller", + function_id="function", + ) + logger.update_environment_variables( + model="capture-model", custom_llm_provider="openai", litellm_params={}, optional_params={} + ) + return logger + + +@pytest.mark.asyncio +async def test_terminal_retrieval_reuses_accepted_snapshot_but_rejects_changed_usage() -> None: + from datetime import timedelta + + from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY + from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, capture_callback + from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + + database: Final = StoredUsage() + writer: Final = UsageBatchWriter(database.insert, database.load) + meter: Final = DurableUsageMeter(writer.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) + kwargs: Final = { + "call_type": "aget_responses", + "litellm_call_id": "poll-one", + "model": "capture-model", + "response_cost": 2.0, + "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, + } + response: Final = { + "id": "job", + "status": "completed", + "background": True, + "created_at": NOW.timestamp(), + "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, + } + first: Final = capture_callback(meter, kwargs, response, NOW, NOW) + await meter.flush() + later: Final = capture_callback( + meter, {**kwargs, "litellm_call_id": "poll-two", "response_cost": 9.0}, response, NOW, NOW + timedelta(hours=2) + ) + await meter.flush() + await later.receipt + assert database.rows == first.writes + assert len(database.rows) == 1 + assert database.rows[0].event.event_id == "job_background_cost" + changed: Final = capture_callback(meter, kwargs, {**response, "usage": {"total_tokens": 8}}, NOW, NOW) + await meter.close() + with pytest.raises(UsageCaptureError, match="identity"): + await changed.receipt + assert database.rows == first.writes + + +@pytest.mark.asyncio +async def test_batch_poll_has_distinct_identity_and_only_terminal_usage_is_chargeable() -> None: + from litellm.proxy.spend_tracking.usage_capture import capture_callback + + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + kwargs: Final = { + "call_type": "aretrieve_batch", + "litellm_call_id": "poll", + "model": "capture-model", + "litellm_params": {}, + } + pending: Final = capture_callback(meter, kwargs, {"id": "batch", "status": "in_progress"}, NOW, NOW) + completed: Final = capture_callback( + meter, kwargs, {"id": "batch", "status": "completed", "usage": {"total_tokens": 9}}, NOW, NOW + ) + await meter.close() + assert pending.writes[-1].event.event_id == "poll:status" + assert json.loads(pending.writes[-1].envelope)["charge_eligible"] is False + assert completed.writes[-1].event.event_id == "batch_batch_cost" + assert json.loads(completed.writes[-1].envelope)["charge_eligible"] is True + assert len(completed.writes) == 1 + + +@pytest.mark.asyncio +async def test_internal_capture_does_not_claim_reused_parent_logger(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm._internal_context import internal_sub_call + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.types.utils import ModelResponse, Usage + + monkeypatch.setitem( + litellm.model_cost, + "capture-model", + { + "input_cost_per_token": 0.1, + "output_cost_per_token": 0.2, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + logger: Final = _logger() + response: Final = ModelResponse( + id="response", model="capture-model", usage=Usage(prompt_tokens=2, completion_tokens=1, total_tokens=3) + ) + with internal_sub_call("caller"): + capture_logging_usage(meter, logger, response, NOW, NOW) + monkeypatch.setitem( + litellm.model_cost, + "capture-model", + { + "input_cost_per_token": 0.7, + "output_cost_per_token": 0.9, + "litellm_provider": "openai", + "mode": "chat", + }, + ) + litellm.register_model({"capture-model": litellm.model_cost["capture-model"]}) + capture_logging_usage(meter, logger, response, NOW, NOW) + capture_logging_usage(meter, logger, response, NOW, NOW) + await meter.close() + assert tuple(row.event.event_id for row in database.rows) == ("response:included", "caller:inbound", "response") + child: Final = json.loads(database.rows[0].envelope) + parent: Final = json.loads(database.rows[-1].envelope) + assert (child["ownership"], child["charge_eligible"], child["parent_request_id"]) == ( + "included_in_parent", + False, + "caller", + ) + assert parent["charge_eligible"] is True + assert child["pricing"]["model_map_information"]["model_map_value"]["input_cost_per_token"] == 0.1 + assert parent["pricing"]["model_map_information"]["model_map_value"]["input_cost_per_token"] == 0.7 + + +@pytest.mark.asyncio +async def test_stream_capture_ignores_chunks_and_preserves_terminal_usage_extras() -> None: + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.types.utils import ModelResponse, ModelResponseStream, Usage + + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + logger: Final = _logger(stream=True) + capture_logging_usage(meter, logger, ModelResponseStream(id="chunk"), NOW, NOW) + assert meter.pending_records == 0 + capture_logging_usage( + meter, logger, ModelResponse(id="stream", usage=Usage(prompt_tokens=3, provider_units=4)), NOW, NOW + ) + await meter.close() + assert json.loads(database.rows[-1].envelope)["usage"]["provider_units"] == 4 + assert database.rows[-1].event.event_id == "stream" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_policy", ["fail_closed", "log"]) +async def test_sync_ack_commit_failure_follows_response_policy( + failure_policy: str, 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(failure=ValueError("invalid envelope")) + database.release.set() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.write, UsageCaptureSettings(batch_size=2)) + meter: Final = runtime.meter + assert meter is not None + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + monkeypatch.setattr( + proxy_server, + "general_settings", + {"accounting_sync_ack": True, "accounting_sync_ack_failure_policy": failure_policy}, + ) + logger: Final = _logger() + response: Final = ModelResponse(id="provider", usage=Usage(prompt_tokens=4)) + if failure_policy == "fail_closed": + with pytest.raises(HTTPException) as error: + await capture_proxy_response(logger, response) + assert error.value.status_code == 503 + else: + await capture_proxy_response(logger, response) + await runtime.stop(1.0) + assert database.committed == () + assert meter.pending_records == 0 + + +@pytest.mark.asyncio +async def test_failed_attempt_keeps_its_usage_without_claiming_the_callers_success() -> None: + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.types.utils import ModelResponse, Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger() + try: + logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.4) + logger.carry_failed_attempt_usage() + logger.capture_usage_before_logging( + ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, total_tokens=3)), NOW, NOW + ) + finally: + await runtime.stop(1) + assert tuple(row.event.event_id for row in database.rows) == ("caller-attempt-0", "caller:inbound", "fallback") + failed: Final = json.loads(database.rows[0].envelope) + assert failed["usage"]["total_tokens"] == 7 + assert failed["parent_request_id"] == "caller" + assert failed["request_count_eligible"] is False + assert failed["charge_eligible"] is True + assert failed["event"]["retry_of_attempt_id"] is None + assert failed["event"]["outcome"] == "failed" + assert failed["pricing"]["reported_cost"] == 0.4 + assert database.rows[1].event.outcome is Outcome.SUCCEEDED + + +@pytest.mark.asyncio +@pytest.mark.parametrize("estimated", [False, True]) +async def test_cancelled_usage_keeps_measured_or_estimated_provenance(estimated: 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 + from litellm.types.utils import Usage + + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + capture_callback( + meter, + { + "litellm_call_id": "cancelled", + "call_type": "acompletion", + "model": "capture-model", + "litellm_params": {}, + "combined_usage_object": Usage(prompt_tokens=4, total_tokens=4), + "_accounting_usage_observed": UsageObservation.ESTIMATED if estimated else UsageObservation.KNOWN, + }, + HTTPException(499, "Client disconnected"), + NOW, + NOW, + ) + await meter.close() + payload: Final = json.loads(database.rows[-1].envelope) + assert payload["event"]["outcome"] == "cancelled" + assert payload["usage"]["total_tokens"] == 4 + assert payload["usage_source"] == ("estimated" if estimated else "measured") + assert payload["pricing"]["reported_cost"] is None + + +@pytest.mark.asyncio +async def test_realtime_capture_keeps_metered_units_and_provider_usage_without_frames() -> None: + from litellm.constants import REALTIME_METERED_USAGE_KEY + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.types.llms.xai import RealtimeMeteredUsage + + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + logger: Final = _logger(CallTypes.arealtime.value) + logger.model_call_details[REALTIME_METERED_USAGE_KEY] = RealtimeMeteredUsage(True, 1.5, 2.5, 3) + capture_logging_usage( + meter, + logger, + [ + { + "type": "response.done", + "response": { + "id": "turn", + "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, + "output": "private realtime transcript", + }, + } + ], + NOW, + NOW, + ) + await meter.close() + payload: Final = json.loads(database.rows[-1].envelope) + assert payload["usage"]["total_tokens"] == 7 + assert payload["usage"]["metered_usage"] == { + "push_to_talk": True, + "audio_seconds_sent": 1.5, + "audio_seconds_received": 2.5, + "text_messages": 3, + } + assert "private realtime transcript" not in database.rows[-1].envelope + + +@pytest.mark.asyncio +async def test_default_http_capture_returns_before_database_commit(monkeypatch: pytest.MonkeyPatch) -> None: + 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() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.write, UsageCaptureSettings(batch_size=2)) + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + monkeypatch.setattr(proxy_server, "general_settings", {}) + try: + await capture_proxy_response( + _logger(), + ModelResponse(id="asynchronous", usage=Usage(prompt_tokens=1)), + ) + await database.entered.wait() + assert database.committed == () + finally: + database.release.set() + await runtime.stop(1) + assert database.committed == ("caller:inbound", "asynchronous") + + +@pytest.mark.asyncio +async def test_deferred_capture_freezes_pricing_before_configuration_changes(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_logging_usage + from litellm.types.utils import ModelResponse, Usage + from litellm.utils import dispatch_success_logging + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger() + logger.defer_async_logging = True + logger.custom_cost_per_second = 0.25 + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.1}) + response: Final = ModelResponse(id="deferred", usage=Usage(prompt_tokens=3)) + try: + dispatch_success_logging(logger, response, NOW, NOW, False, False) + assert database.rows == () + logger.custom_cost_per_second = 9 + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.9}) + assert runtime.meter is not None + capture_logging_usage(runtime.meter, logger, response, NOW, NOW) + await runtime.meter.flush() + finally: + await runtime.stop(1) + pricing: Final = json.loads(database.rows[-1].envelope)["pricing"] + assert pricing["logging"]["custom_cost_per_second"] == 0.25 + assert pricing["adjustments"]["discount"] == 0.1 + + +@pytest.mark.asyncio +async def test_synchronous_internal_capture_reserves_capacity_before_scheduling() -> None: + from litellm._internal_context import internal_sub_call + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.types.utils import ModelResponse, Usage + + database: Final = CommitBarrier() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.write, UsageCaptureSettings(max_pending_records=2, batch_size=2)) + logger: Final = _logger() + + def capture_child() -> None: + with internal_sub_call("caller"): + logger.capture_usage_before_logging(ModelResponse(id="child", usage=Usage(prompt_tokens=2)), NOW, NOW) + + try: + await asyncio.to_thread(capture_child) + assert runtime.meter is not None + assert runtime.meter.pending_records == 1 + receipt: Final = await asyncio.to_thread(runtime.meter.submit, (_write("second"),)) + await database.entered.wait() + with pytest.raises(UsageCaptureError, match="capacity"): + await asyncio.to_thread(runtime.meter.submit, (_write("overflow"),)) + database.release.set() + await receipt + finally: + database.release.set() + await runtime.stop(1) + assert database.committed == ("child:included", "second") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("has_metered_usage", [False, True]) +async def test_realtime_without_provider_token_usage_does_not_invent_zero(has_metered_usage: bool) -> None: + from litellm.constants import REALTIME_METERED_USAGE_KEY + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + from litellm.types.llms.xai import RealtimeMeteredUsage + + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + logger: Final = _logger(CallTypes.arealtime.value) + if has_metered_usage: + logger.model_call_details[REALTIME_METERED_USAGE_KEY] = RealtimeMeteredUsage(True, 0, 0, 0) + capture_logging_usage(meter, logger, [], NOW, NOW) + await meter.close() + usage: Final = json.loads(database.rows[-1].envelope)["usage"] + assert usage == ( + { + "metered_usage": { + "push_to_talk": True, + "audio_seconds_sent": 0, + "audio_seconds_received": 0, + "text_messages": 0, + } + } + if has_metered_usage + else None + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("committed_before_failure", [False, True]) +async def test_database_recovery_retries_frozen_usage_without_callback_or_duplicate_rows( + committed_before_failure: bool, +) -> None: + from litellm.proxy.spend_tracking.usage_capture import capture_callback + + database: Final = StoredUsage() + failures: Final = asyncio.Queue[None]() + failures.put_nowait(None) + retrying: Final = asyncio.Event() + recovered: Final = asyncio.Event() + delays: Final = asyncio.Queue[float]() + + async def persist(writes: tuple[UsageWrite, ...]) -> None: + if not failures.empty(): + failures.get_nowait() + if committed_before_failure: + await database.insert(writes) + raise ConnectionError("database unavailable or acknowledgement lost") + await database.insert(writes) + + async def retry_delay(delay: float) -> None: + delays.put_nowait(delay) + retrying.set() + await recovered.wait() + + meter: Final = DurableUsageMeter( + persist, UsageCaptureSettings(max_pending_records=2), wait_for_batch=_manual_flush, retry_delay=retry_delay + ) + kwargs: Final = { + "call_type": "acompletion", + "litellm_call_id": "caller", + "model": "capture-model", + "litellm_params": {}, + } + response: Final = {"id": "provider", "usage": {"total_tokens": 3}} + first: Final = capture_callback(meter, kwargs, response, NOW, NOW) + reserved: Final = meter.pending_bytes + flushing: Final = asyncio.create_task(meter.flush()) + await retrying.wait() + assert not first.receipt.done() + assert (meter.pending_records, meter.pending_bytes) == (2, reserved) + with pytest.raises(UsageCaptureError, match="capacity"): + meter.submit((_write("overflow"),)) + recovered.set() + await flushing + await first.receipt + duplicate: Final = capture_callback(meter, {**kwargs, "_accounting_capture": first}, response, NOW, NOW) + await meter.close() + assert duplicate is first + assert database.rows == first.writes + assert delays.get_nowait() == 0.2 + assert (meter.pending_records, meter.pending_bytes) == (0, 0) + + +@pytest.mark.asyncio +async def test_cut_stream_keeps_original_model_usage_as_part_of_the_parent_aggregate() -> None: + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.router import Router + from litellm.types.utils import ModelResponse, Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger(stream=True) + stream: Final = CustomStreamWrapper(completion_stream=iter(()), model="capture-model", logging_obj=logger) + partial: Final = ModelResponse( + id="cut", model="capture-model", usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5) + ) + try: + Router._price_cut_partial(stream, partial) + finally: + await runtime.stop(1) + assert len(database.rows) == 1 + payload: Final = json.loads(database.rows[0].envelope) + assert payload["event"]["event_id"] == "cut:included" + assert payload["parent_request_id"] == "caller" + assert payload["charge_eligible"] is False + assert payload["usage"]["total_tokens"] == 5 + assert payload["pricing"]["model_map_information"]["model_map_key"] == "capture-model" + + +@pytest.mark.asyncio +async def test_direct_mcp_logging_captures_tool_pricing_and_identity_without_arguments() -> None: + from mcp.types import CallToolResult, TextContent + + from litellm.accounting import DimensionName + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger(CallTypes.call_mcp_tool.value) + logger.model_call_details["mcp_tool_call_metadata"] = { + "name": "lookup", + "namespaced_tool_name": "server.lookup", + "arguments": {"query": "private query"}, + "mcp_server_cost_info": {"default_cost_per_query": 0.2, "tool_name_to_cost_per_query": {"lookup": 0.3}}, + } + try: + await logger.async_success_handler( + CallToolResult(content=[TextContent(type="text", text="private result")]), NOW, NOW + ) + finally: + await runtime.stop(1) + outbound: Final = database.rows[-1] + assert any(d.name is DimensionName.TOOL and d.value == "server.lookup" for d in outbound.event.dimensions) + payload: Final = json.loads(outbound.envelope) + assert payload["pricing"]["tool"] == {"name": "lookup", "cost_per_query": 0.3} + assert payload["usage"] is None + assert all(secret not in outbound.envelope for secret in ("private query", "private result")) + + +@pytest.mark.asyncio +async def test_callback_overload_is_visible_without_suppressing_legacy_tracking( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, captured_callback_kwargs + + database: Final = CommitBarrier() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.write, UsageCaptureSettings(max_pending_records=2, batch_size=2)) + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + kwargs: Final = {"call_type": "acompletion", "litellm_call_id": "one", "litellm_params": {}} + try: + admitted: Final = captured_callback_kwargs(kwargs, {"id": "one"}, NOW, NOW) + await database.entered.wait() + rejected: Final = captured_callback_kwargs({**kwargs, "litellm_call_id": "two"}, {"id": "two"}, NOW, NOW) + assert "_accounting_capture" in admitted + assert "_accounting_capture" not in rejected + assert "Usage capture failed before legacy cost tracking" in caplog.text + finally: + database.release.set() + await runtime.stop(1) + assert database.committed == ("one:inbound", "one") + + +@pytest.mark.parametrize( + ("call_type", "streaming", "background"), + [ + ("acompletion", True, False), + ("aresponses", False, True), + ("arealtime", False, False), + ], +) +def test_strict_ack_rejects_requests_without_a_terminal_response_boundary( + call_type: str, streaming: bool, background: bool +) -> None: + from fastapi import HTTPException + + from litellm.proxy.spend_tracking.usage_capture import validate_sync_ack_request + + validate_sync_ack_request({}, call_type, streaming, background) + with pytest.raises(HTTPException) as error: + validate_sync_ack_request({"accounting_sync_ack": True}, call_type, streaming, background) + assert error.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_new_connection_turn_does_not_reuse_the_previous_turns_capture() -> None: + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.types.utils import ModelResponse, Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger() + try: + logger.capture_usage_before_logging(ModelResponse(id="first", usage=Usage(prompt_tokens=2)), NOW, NOW) + next_turn: Final = logger.copy_for_new_call("second-caller", NOW) + next_turn.capture_usage_before_logging(ModelResponse(id="second", usage=Usage(prompt_tokens=4)), NOW, NOW) + finally: + await runtime.stop(1) + assert tuple(row.event.event_id for row in database.rows) == ( + "caller:inbound", + "first", + "second-caller:inbound", + "second", + ) + assert json.loads(database.rows[-1].envelope)["usage"]["prompt_tokens"] == 4 + + +@pytest.mark.asyncio +async def test_billed_failure_is_captured_before_legacy_persistence_gates(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.types.utils import Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + monkeypatch.setattr(proxy_server, "prisma_client", None) + request: Final = { + "call_type": "acompletion", + "litellm_call_id": "failed", + "model": "capture-model", + "start_time": NOW, + "combined_usage_object": Usage(prompt_tokens=3, total_tokens=3), + "response_cost": 0.2, + "litellm_params": {}, + } + try: + await ProxyDBLogger().async_post_call_failure_hook(request, RuntimeError("provider failed"), UserAPIKeyAuth()) + finally: + await runtime.stop(1) + assert tuple(row.event.event_id for row in database.rows) == ("failed:inbound", "failed") + payload: Final = json.loads(database.rows[-1].envelope) + assert payload["event"]["outcome"] == "failed" + assert payload["usage"]["total_tokens"] == 3 + assert payload["charge_eligible"] is True + assert payload["pricing"]["reported_cost"] == 0.2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pages", [None, 0, 3]) +async def test_prepared_ocr_capture_preserves_reported_pages(pages: int | None) -> None: + from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo + from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage + + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + logger: Final = _logger("aocr") + response: Final = OCRResponse(pages=[], model="capture-model", usage_info=OCRUsageInfo(pages_processed=pages)) + logger.prepare_response_cost(response, NOW) + capture_logging_usage(meter, logger, response, NOW, NOW) + await meter.close() + envelope: Final = json.loads(database.rows[-1].envelope) + assert (None if envelope["usage"] is None else envelope["usage"]["pages_processed"]) == pages + assert envelope["usage_source"] == ("unknown" if pages is None else "measured") + assert response.usage_info is not None and response.usage_info.pages_processed == pages + + +@pytest.mark.asyncio +@pytest.mark.parametrize("measured", [None, {"total_tokens": 0}, {"total_tokens": 3}]) +async def test_batch_capture_provenance_tracks_observations_not_envelope_presence( + measured: Mapping[str, int] | None, +) -> None: + from litellm.proxy.spend_tracking.usage_capture import capture_callback + + database: Final = StoredUsage() + meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) + usage: Final = { + "total_tokens": 0 if measured is None else measured["total_tokens"], + "batch_usage_by_model": {"missing-model": None, "other-model": measured}, + "batch_unknown_usage_requests": 2 if measured is None else 1, + } + capture_callback( + meter, + {"call_type": "aretrieve_batch", "litellm_call_id": "caller", "litellm_params": {}}, + {"id": "batch", "usage": usage}, + NOW, + NOW, + ) + await meter.close() + envelope: Final = json.loads(database.rows[-1].envelope) + assert envelope["usage_source"] == ("unknown" if measured is None else "measured") + assert envelope["usage"]["batch_usage_by_model"] == usage["batch_usage_by_model"] + assert database.rows[-1].event.metrics == ( + None if measured is None else (Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(measured["total_tokens"])),) + ) + + +@pytest.mark.asyncio +async def test_ack_admission_failure_does_not_turn_provider_success_into_failed_usage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import HTTPException + + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response + from litellm.types.utils import ModelResponse, Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings(batch_size=2, max_pending_records=2)) + monkeypatch.setattr(proxy_server, "usage_capture", runtime) + monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) + monkeypatch.setattr(proxy_server, "prisma_client", None) + meter: Final = runtime.meter + assert meter is not None + logger: Final = _logger() + try: + meter.submit((_write("occupied-one"), _write("occupied-two"))) + with pytest.raises(HTTPException) as error: + await capture_proxy_response(logger, ModelResponse(id="caller", usage=Usage(prompt_tokens=3))) + assert error.value.status_code == 503 + await meter.flush() + request: Final = {**logger.model_call_details, "litellm_logging_obj": logger, "start_time": NOW} + await ProxyDBLogger().async_post_call_failure_hook(request, error.value, UserAPIKeyAuth()) + finally: + await runtime.stop(1) + assert tuple(row.event.event_id for row in database.rows) == ("occupied-one", "occupied-two") + + +@pytest.mark.asyncio +async def test_response_ack_timeout_preserves_success_and_later_commits(monkeypatch: pytest.MonkeyPatch) -> None: + from fastapi import HTTPException + + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response + from litellm.types.utils import ModelResponse, Usage + + database: Final = CommitBarrier() + runtime: Final = UsageCaptureRuntime() + 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}) + 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))) + assert error.value.status_code == 503 + await database.entered.wait() + captured: Final = logger.model_call_details["_accounting_capture"] + request: Final = {**logger.model_call_details, "litellm_logging_obj": logger, "start_time": NOW} + await ProxyDBLogger().async_post_call_failure_hook(request, error.value, UserAPIKeyAuth()) + database.release.set() + await captured.receipt + assert tuple(row.event.outcome for row in captured.writes) == (Outcome.SUCCEEDED, Outcome.SUCCEEDED) + assert database.committed == ("caller:inbound", "provider") + assert database.transactions == 1 + finally: + await runtime.stop(1) diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index b3725f01512..5543113af60 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -268,8 +268,15 @@ token = MASTER_KEY @pytest.mark.asyncio -async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning() -> None: +async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from datetime import datetime + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.proxy.db.accounting_ledger import UsageWrite + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback + from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings from litellm.responses.streaming_iterator import spawn_logging_task worker: Final = LoggingWorker() @@ -279,7 +286,25 @@ async def test_shutdown_persists_accounting_from_worker_callbacks_before_returni pending: Final[asyncio.Queue[str]] = asyncio.Queue() persisted: Final[asyncio.Queue[str]] = asyncio.Queue() + captured: Final[asyncio.Queue[str]] = asyncio.Queue() + + async def persist_usage(writes: tuple[UsageWrite, ...]) -> None: + for write in writes: + captured.put_nowait(write.event.event_id) + + capture: Final = UsageCaptureRuntime() + capture.start(persist_usage, UsageCaptureSettings(flush_interval_seconds=10)) + monkeypatch.setattr(litellm.proxy.proxy_server, "usage_capture", capture) + async def enqueue_spend() -> None: + assert capture.meter is not None + capture_callback( + capture.meter, + {"call_type": "acompletion", "litellm_call_id": "shutdown-usage", "litellm_params": {}}, + {"id": "shutdown-attempt", "usage": {"total_tokens": 3}}, + datetime(2026, 1, 1), + datetime(2026, 1, 1), + ) pending.put_nowait("accounted request") async def callback() -> None: @@ -305,10 +330,13 @@ async def test_shutdown_persists_accounting_from_worker_callbacks_before_returni assert waited_for_callback, "Shutdown returned while accounting callbacks were still running" assert persisted.get_nowait() == "accounted request" assert pending.empty() + assert (captured.get_nowait(), captured.get_nowait()) == ("shutdown-usage:inbound", "shutdown-attempt") + assert capture.meter is None finally: release.set() await worker.flush() await worker.stop() + await capture.stop(0) @pytest.mark.asyncio @@ -4261,3 +4289,25 @@ async def test_shutdown_waits_for_enqueue_producers_before_stopping_worker() -> assert worker.pending_count == 0 finally: await worker.stop(discard_pending=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unsupported", ["database", "collector", "remote", "free_tier"]) +async def test_capture_enablement_rejects_unrecoverable_writer_modes( + unsupported: str, monkeypatch: pytest.MonkeyPatch +) -> None: + import litellm + from litellm.proxy import proxy_server + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + + monkeypatch.setattr(proxy_server, "usage_capture", UsageCaptureRuntime()) + monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {}}) + monkeypatch.setattr(proxy_server, "prisma_client", None if unsupported == "database" else object()) + monkeypatch.setattr(litellm, "apply_search_free_tier", unsupported == "free_tier") + monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "true" if unsupported == "collector" else "false") + monkeypatch.delenv("SPEND_LOGS_URL", raising=False) + if unsupported == "remote": + monkeypatch.setenv("SPEND_LOGS_URL", "https://writer.invalid") + with pytest.raises(ValueError, match="requires local Prisma"): + proxy_server._configure_usage_capture() + assert proxy_server.usage_capture.meter is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 30d61f181cb..8e4fe573130 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28644,7 +28644,7 @@ export interface components { * CallTypes * @enum {string} */ - CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "decisions" | "adecisions" | "systemone" | "asystemone" | "extraction" | "aextraction" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; + CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "decisions" | "adecisions" | "systemone" | "asystemone" | "extraction" | "aextraction" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "get_responses" | "aget_responses" | "delete_responses" | "adelete_responses" | "cancel_responses" | "acancel_responses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; /** CallbackDelete */ CallbackDelete: { /** Callback Name */ @@ -29942,6 +29942,27 @@ export interface components { * @description Documents all the fields supported by `general_settings` in config.yaml */ ConfigGeneralSettings: { + /** @description Opt-in durable usage capture. Legacy spend updates remain active. Requires local Prisma writes. */ + accounting_capture?: components["schemas"]["UsageCaptureSettings"] | null; + /** + * Accounting Sync Ack + * @description Wait for usage commit before a supported non-stream response + * @default false + */ + accounting_sync_ack: boolean; + /** + * Accounting Sync Ack Failure Policy + * @description Return 503 or log the persistence failure after the response wait expires + * @default fail_closed + * @enum {string} + */ + accounting_sync_ack_failure_policy: "fail_closed" | "log"; + /** + * Accounting Sync Ack Timeout Seconds + * @description Maximum response wait for usage commit + * @default 5 + */ + accounting_sync_ack_timeout_seconds: number; /** * Admission Queue Timeout Seconds * @description maximum time a request waits for a worker slot @@ -49501,6 +49522,29 @@ export interface components { /** User Role */ user_role?: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; + /** UsageCaptureSettings */ + UsageCaptureSettings: { + /** + * Batch Size + * @default 1000 + */ + batch_size: number; + /** + * Flush Interval Seconds + * @default 0.2 + */ + flush_interval_seconds: number; + /** + * Max Pending Bytes + * @default 16777216 + */ + max_pending_bytes: number; + /** + * Max Pending Records + * @default 10000 + */ + max_pending_records: number; + }; /** UsageChartPoint */ UsageChartPoint: { /** Blocked */