feat(accounting): capture request usage through the durable meter

This commit is contained in:
Joshua Valluru 2026-10-29 22:34:19 -07:00
parent 02a6a96dec
commit 85a987823c
31 changed files with 2395 additions and 55 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -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 */