scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166

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

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

* fix(accounting): preserve native Google usage compatibility

* fix(accounting): preserve prepared usage observation on failure

* fix(accounting): exclude sync acknowledgement policy rejections from capture

* test(accounting): verify capture races and failure isolation

* fix(accounting): preserve request support with sync acknowledgements

* fix(accounting): normalize usage metadata before queue admission

* fix(accounting): preserve prior attempt links on successful fallback

* fix(accounting): preserve attempt evidence through retries and failures

* fix(accounting): depend on prior receipts without resubmitting usage

* fix(accounting): preserve pricing inputs and resolved capture attribution

* fix(accounting): persist batch usage within capture bounds

* test(accounting): cover batch dependency failure and backpressure

* fix(accounting): isolate batch persistence and preserve trusted attribution

* fix(accounting): share byte admission across batch producers and writes

* fix(accounting): reserve batch snapshot reads before admission

* fix(accounting): compact repeated batch observations without losing requests

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri 2026-10-31 07:32:25 -07:00 • committed by GitHub
parent a713ed6bab
commit be4f3d36ce
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
38 changed files with 4291 additions and 102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
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

5
litellm/_internal_context.py
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

6
litellm/accounting/protocols.py
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.
"""

4
litellm/accounting/types.py
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.
"""

95
litellm/batches/batch_utils.py
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,54 @@ 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_request(usages: tuple[tuple[str | None, Usage | None], ...]) -> tuple[Mapping[str, object], ...]:
rows: Final = tuple(
{"model": model, "usage": usage.model_dump() if usage is not None else None} for model, usage in usages
)
return tuple(sorted(rows, key=lambda row: json.dumps(row, sort_keys=True)))
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 +420,9 @@ 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_usage_by_request=_batch_usage_by_request(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 +463,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 +501,13 @@ 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_usage_by_request=_batch_usage_by_request(
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]

2
litellm/endpoints/speech/speech_to_completion_bridge/handler.py
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,
)

286
litellm/litellm_core_utils/litellm_logging.py
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 (
@ -163,6 +165,7 @@ from litellm.types.llms.openai import (
ResponseIncompleteEvent,
ResponsesAPIResponse,
)
from litellm.types.llms.vertex_ai import GenerateContentResponseBody, UsageMetadata
from litellm.types.mcp import MCPPostCallResponseObject
from litellm.types.prompts.init_prompts import PromptSpec
from litellm.types.rerank import RerankResponse, get_rerank_usage
@ -266,6 +269,9 @@ 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
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError
try:
from litellm_enterprise.enterprise_callbacks.callback_controls import (
EnterpriseCallbackControls,
@ -571,6 +577,7 @@ def _resolve_mantle_region_for_cost(
_PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"})
_LITELLM_PARAMS: Final = TypeAdapter(dict[str, object])
_SERVED_ANTHROPIC_BODY: Final = TypeAdapter(dict[str, object])
_GOOGLE_USAGE_FIELDS: Final = frozenset(UsageMetadata.__annotations__)
_REDACTED_RESULT: Final = TypeAdapter(object)
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
@ -630,7 +637,60 @@ 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"
@dataclass(frozen=True, slots=True)
class UsageCaptureObservation:
usage_json: str | None
provenance: UsageObservation
response_cost: float | None = None
@classmethod
def from_usage(
cls, usage: Usage | None, provenance: UsageObservation, response_cost: float | None = None
) -> "UsageCaptureObservation | UsageCaptureError":
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError
try:
serialized: Final = (
None if usage is None or provenance is UsageObservation.UNKNOWN else usage.model_dump_json()
)
except Exception as error: # noqa: BLE001 # observation failures must not interrupt legacy retry handling
return UsageCaptureError(str(error))
else:
return cls(serialized, provenance, response_cost)
class Logging(LiteLLMLoggingBaseClass):
usage_capture_callback: ClassVar[
Callable[
["Logging", object, datetime.datetime, datetime.datetime, bool], "CapturedUsage | UsageCaptureError | 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 +888,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 +945,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 +955,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 +2486,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 +2521,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 +2539,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 +2570,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,
@ -2540,17 +2618,29 @@ class Logging(LiteLLMLoggingBaseClass):
},
usage,
)
self.record_partial_usage_for_failure(usage, raw_cost)
self.record_partial_usage_for_failure(
usage, raw_cost, observation=self.model_call_details["_accounting_usage_observed"]
)
return self.prepared_cost
def record_finalized_cost(self, finalized: FinalizedCost) -> None:
self.finalized_cost = finalized
self.model_call_details["response_cost"] = finalized.amount
def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None:
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, *, observation: UsageObservation = UsageObservation.KNOWN
) -> 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
self.model_call_details["response_cost"] = response_cost
if type(self).usage_capture_callback is not None:
self.model_call_details["_accounting_usage_observed"] = observation
def record_failure_after_earlier_rounds(
self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float
@ -2630,6 +2720,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 +2734,53 @@ 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
from litellm.proxy.spend_tracking.usage_capture import CapturedUsage
capture: Final = type(self).usage_capture_callback
if capture is None:
return
attempt: Final = self.attempt_scoped_copy()
attempt.prepared_cost = None
attempt.finalized_cost = None
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)
observed: Final[object] = self.model_call_details.get("_accounting_usage_observed")
cost: Final[object] = self.model_call_details.get("response_cost")
attempt.model_call_details = {
**attempt.model_call_details,
"litellm_call_id": identity,
"parent_request_id": self.litellm_call_id,
"_accounting_capture": None,
"_accounting_prior_captures": (),
"_accounting_usage_snapshot": UsageCaptureObservation.from_usage(
usage,
observed if isinstance(observed, UsageObservation) else UsageObservation.KNOWN,
float(cost) if isinstance(cost, (int, float)) else 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"},
},
}
captured: Final = capture(
attempt,
{"id": identity, "usage": usage.model_dump(), "status": "failed"},
self.start_time,
datetime.datetime.now(datetime.timezone.utc),
False,
)
if isinstance(captured, CapturedUsage):
self.model_call_details["_accounting_prior_captures"] = (
*self.model_call_details.get("_accounting_prior_captures", ()),
captured,
)
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 +2938,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 +2997,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 +3283,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 +3636,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 +4160,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")
@ -6808,45 +7006,19 @@ class StandardLoggingPayloadSetup:
@staticmethod
def get_usage_from_response_obj(response_obj: dict | None, combined_usage_object: Usage | None = None) -> Usage:
## BASE CASE ##
if combined_usage_object is not None:
return combined_usage_object
if response_obj is None:
return Usage(
prompt_tokens=0,
completion_tokens=0,
total_tokens=0,
)
usage: Final = response_obj.get("usage", None)
if usage is None:
rerank_usage: Final = get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj))
if rerank_usage is not None:
return Usage(**rerank_usage)
if usage is None or (not isinstance(usage, dict) and not isinstance(usage, Usage)):
return Usage(
prompt_tokens=0,
completion_tokens=0,
total_tokens=0,
)
elif isinstance(usage, Usage):
return usage
elif isinstance(usage, ResponseAPIUsage):
return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage)
elif isinstance(usage, dict):
if ResponseAPILoggingUtils.is_response_api_usage(usage):
return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage)
if InteractionsUsageObjectTransformation.is_interactions_usage_object(usage):
return InteractionsUsageObjectTransformation.transform_interactions_usage_object(usage)
return Usage(**usage)
raise ValueError(f"usage is required, got={usage} of type {type(usage)}")
reported: Final = response_obj.get("usage") if response_obj else None
if isinstance(reported, Usage):
return reported
normalized: Final = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj)
return Usage(**normalized) if normalized else Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
@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 +7030,21 @@ 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)
google_usage: Final[object] = response_obj.get("usageMetadata")
if isinstance(google_usage, dict) and _GOOGLE_USAGE_FIELDS.intersection(google_usage):
return litellm.VertexGeminiConfig.calculate_usage(
cast( # cast-ok: retain converter-supported modalities and token_count aliases
GenerateContentResponseBody, response_obj
)
).model_dump()
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):

2
litellm/litellm_core_utils/realtime_streaming.py
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.

17
litellm/llms/vertex_ai/batches/transformation.py
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 "

17
litellm/proxy/_types.py
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,22 @@ 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 on foreground non-stream chat, text, embedding and Responses requests. "
"Other requests remain supported with asynchronous usage capture.",
)
accounting_sync_ack_timeout_seconds: float = Field(
20.0, gt=0, description="Maximum response wait for usage commit, including the batch flush interval"
)
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.",

4
litellm/proxy/common_request_processing.py
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
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,
@ -3366,6 +3367,9 @@ class ProxyBaseLLMRequestProcessing:
)
),
)
await capture_proxy_response(
logging_obj, response, background=_REQUEST_FIELDS.validate_python(self.data).get("background") is True
)
except (Exception, asyncio.CancelledError) as exc:
_exception_raised = True
if logging_obj.prepared_cost is not None:

2
litellm/proxy/db/accounting_ledger.py
View file

@ -275,7 +275,7 @@ class AccountingLedger:
)
return tuple(
UsageWrite(
OutboundUsageEvent.model_validate(_JSON.validate_python(row["payload"])["event"]),
OutboundUsageEvent.model_validate_json(_json(_JSON.validate_python(row["payload"])["event"])),
_json(row["payload"]),
TypeAdapter(tuple[DeliveryTarget, ...]).validate_python(row["targets"]),
)

145
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -22,7 +22,11 @@ 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,
UsageCaptureObservation,
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 +67,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 +349,27 @@ def _attempt_spend_row_kwargs(
"litellm_params": litellm_params,
"combined_usage_object": attempt.usage,
"response_cost": attempt.response_cost,
"_accounting_usage_observed": UsageObservation.ESTIMATED,
"_accounting_usage_snapshot": UsageCaptureObservation.from_usage(
attempt.usage, UsageObservation.ESTIMATED, attempt.response_cost
),
"_accounting_prior_captures": (),
}
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
@ -385,7 +406,7 @@ class _ProxyDBLogger(CustomLogger):
self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime
) -> None:
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(dict(kwargs), response_obj, start_time, end_time)
return
event: Final = build_spend_event(
kwargs,
@ -396,7 +417,7 @@ class _ProxyDBLogger(CustomLogger):
)
if isinstance(event, SpendEventBuildError):
verbose_proxy_logger.warning("collector: tracking cost in-process, event not buildable: %s", event.reason)
await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time)
await self._PROXY_track_cost_callback(dict(kwargs), response_obj, start_time, end_time)
return
await self.spend_event_producer.publish(event)
@ -434,6 +455,9 @@ class _ProxyDBLogger(CustomLogger):
"Failed to invalidate budget reservation counters after failure release failed"
)
capture_metadata: Final = await self._try_capture_failure_usage(
request_data, original_exception, user_api_key_dict
)
if completed_cost and tracking_disabled:
return
@ -477,8 +501,12 @@ class _ProxyDBLogger(CustomLogger):
)
_metadata["error_information"] = _error_information
_metadata = await ProxyDBLogger._enrich_failure_metadata_unless_db_stalled( # rebind-ok: pre-existing rebinding on a rename-only line
metadata=_metadata, original_exception=original_exception
_metadata = ( # rebind-ok: enriched identity is shared with the earlier immutable capture
{**_metadata, **capture_metadata}
if capture_metadata is not None
else await ProxyDBLogger._enrich_failure_metadata_unless_db_stalled(
metadata=_metadata, original_exception=original_exception
)
)
request_fields: Final = _str_keyed_fields(request_data)
@ -595,6 +623,106 @@ class _ProxyDBLogger(CustomLogger):
),
)
async def _failure_capture_metadata(
self, request_data: Mapping[str, object], exception: Exception, auth: UserAPIKeyAuth
) -> Mapping[str, object] | None:
from litellm.proxy.proxy_server import usage_capture
if usage_capture.meter is None:
return None
captured_identity: Final = (
_batch_accounting_for(request_data.get("litellm_call_id")) is not None
and request_data.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is not REJECTED_BEFORE_REQUEST_SETUP
)
current_identity: Final = {
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=auth),
"user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(auth),
}
metadata: Final = (
{
**dict.fromkeys(current_identity),
**{
key: value
for key, value in _REQUEST_DATA.validate_python(
get_litellm_metadata_from_kwargs(dict(request_data))
).items()
if key in current_identity or key == "tags"
},
}
if captured_identity
else current_identity
)
return await self._enrich_failure_metadata_unless_db_stalled(
metadata=metadata, original_exception=exception, resolve_missing_key_identity=not captured_identity
)
async 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,
) -> Mapping[str, object] | None:
try:
metadata: Final = await self._failure_capture_metadata(request_data, original_exception, auth)
self._capture_failure_usage(request_data, original_exception, auth, metadata)
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")
return None
return metadata
def _capture_failure_usage(
self,
request_data: dict[str, object], # mutable-ok: receipt handoff shared with the legacy writer
original_exception: Exception,
auth: UserAPIKeyAuth,
metadata: Mapping[str, object] | None = None,
) -> 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 {}
)
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")),
**(
metadata
if metadata is not None
else {
**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,
@ -780,6 +908,7 @@ class _ProxyDBLogger(CustomLogger):
resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES,
)
_write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
kwargs.update(await captured_callback_kwargs(kwargs, completion_response, start_time, end_time))
budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
if accounting is not None and (budget_reservation is not None or litellm.apply_search_free_tier):
accounting.write_started = True
@ -967,11 +1096,13 @@ class _ProxyDBLogger(CustomLogger):
@staticmethod
async def _enrich_failure_metadata_unless_db_stalled(
metadata: dict[str, object], original_exception: Exception
metadata: dict[str, object], original_exception: Exception, resolve_missing_key_identity: bool = True
) -> dict[str, object]:
if isinstance(original_exception, DBLookupDeadlineExceeded):
return metadata
return await ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata)
return await ProxyDBLogger._enrich_failure_metadata_with_key_info(
metadata=metadata, resolve_missing_key_identity=resolve_missing_key_identity
)
@staticmethod
async def _enrich_failure_metadata_with_key_info(metadata: dict, resolve_missing_key_identity: bool = True) -> dict:

44
litellm/proxy/proxy_server.py
View file

@ -473,7 +473,6 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l
log_llm_api_exception,
open_sse_before_first_byte,
proxy_exception_from_route_error,
request_litellm_call_id,
resolve_litellm_call_id,
should_return_raw_model_name,
sse_error_payload,
@ -840,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
@ -1329,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()
@ -3202,9 +3204,39 @@ 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), load=ledger.load_usage_batch
)
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
@ -13006,8 +13038,10 @@ async def completion(
"""
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
data = {}
base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
try:
data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line
base_llm_response_processor.data = data
if user_api_key_dict is not None:
if data.get("metadata") is None:
data["metadata"] = {}
@ -13021,7 +13055,6 @@ async def completion(
data["metadata"]["user_api_key_org_alias"] = user_api_key_dict.organization_alias
if hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None:
data["metadata"]["agent_id"] = user_api_key_dict.agent_id
base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
return await base_llm_response_processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
@ -13123,9 +13156,12 @@ async def completion(
return _response
except Exception as e:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=base_llm_response_processor.data,
dispatched_call_type=base_llm_response_processor.route_type,
)
litellm_call_id: Final = request_litellm_call_id(data)
litellm_call_id: Final = base_llm_response_processor.litellm_call_id
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, ProxyException):
raise with_litellm_call_id(e, litellm_call_id)

1
litellm/proxy/spend_tracking/cost_finalization.py
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
)

23
litellm/proxy/spend_tracking/spend_tracking_utils.py
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 | None](datetime | None).validate_python(
kwargs.get("completion_start_time", end_time)
)
call_type: Final = kwargs.get("call_type")
cache_hit: Final = kwargs.get("cache_hit", False)
@ -644,7 +649,7 @@ def get_logging_payload(
_usage: Final = (
chat_usage_of_batch.model_dump()
if chat_usage_of_batch is not None
else response_obj_dict.get("usage", None) or {}
else StandardLoggingPayloadSetup.get_usage_as_dict(response_obj_dict)
)
if isinstance(_usage, litellm.Usage):
usage = dict(_usage)
@ -928,7 +933,9 @@ def get_logging_payload(
cache_hit=str(cache_hit),
startTime=_ensure_datetime_utc(start_time),
endTime=_ensure_datetime_utc(end_time),
completionStartTime=_ensure_datetime_utc(completion_start_time),
completionStartTime=(
_ensure_datetime_utc(completion_start_time) if completion_start_time is not None else None
),
model=model_name,
user=metadata.get("user_api_key_user_id", "") or "",
team_id=metadata.get("user_api_key_team_id", "") or "",
@ -984,6 +991,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

1084
litellm/proxy/spend_tracking/usage_capture.py Normal file
View file

File diff suppressed because it is too large Load diff

64
litellm/proxy/spend_tracking/usage_meter.py
View file

@ -2,7 +2,7 @@ from __future__ import annotations
import asyncio
from collections import deque
from collections.abc import Awaitable, Callable
from collections.abc import Awaitable, Callable, Mapping
from concurrent.futures import Future
from dataclasses import dataclass
from itertools import accumulate, chain, count, takewhile
@ -114,6 +114,9 @@ class DurableUsageMeter:
self._admission: Final = Lock()
self._queue: Final[deque[_Pending]] = deque() # mutable-ok: one bounded producer/consumer queue
self._wakeup: Final = asyncio.Event()
self._capacity_changed: Final = asyncio.Event()
self._reservations: Mapping[object, tuple[int, int]] = {}
self._reserved_bytes = 0
self._task: asyncio.Task[None] | None = None
self._closed = False
self.pending_bytes = 0
@ -143,6 +146,36 @@ class DurableUsageMeter:
wrapped.add_done_callback(lambda done: None if done.cancelled() else done.exception())
return wrapped
def reserve_producer(self, size: int, headroom: int) -> object:
token: Final = object()
self.resize_producer(token, size, headroom)
return token
def resize_producer(self, token: object, size: int, headroom: int) -> None:
with self._admission:
previous: Final = self._reservations.get(token, (0, 0))[0]
delta: Final = size - previous
others: Final = (value[1] for key, value in self._reservations.items() if key is not token)
required: Final = max((headroom, *others))
if (
self._closed
or self.pending_bytes + delta > self.settings.max_pending_bytes
or self._reserved_bytes + delta + required > self.settings.max_pending_bytes
):
raise UsageCaptureError("Batch capture producer byte capacity exhausted")
self._reservations = {**self._reservations, token: (size, headroom)}
self._reserved_bytes += delta
self.pending_bytes += delta
self._capacity_changed.set()
def release_producer(self, token: object) -> None:
with self._admission:
size, _ = self._reservations[token]
self._reservations = {key: value for key, value in self._reservations.items() if key is not token}
self._reserved_bytes -= size
self.pending_bytes -= size
self._capacity_changed.set()
def _ensure_worker(self) -> None:
if self.pending_records >= self.settings.batch_size or self._queue and self._queue[0].prerequisites:
self._wakeup.set()
@ -153,6 +186,34 @@ class DurableUsageMeter:
envelope: Final = '{"event":' + event.model_dump_json() + "}"
await asyncio.shield(self.submit((UsageWrite(event, envelope, usage_delivery_targets(event.dimensions)),)))
async def record_writes(self, writes: tuple[UsageWrite, ...]) -> None:
size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes)
capacity: Final = self.settings.max_pending_bytes - self._reserved_bytes
if len(writes) > min(self.settings.batch_size, self.settings.max_pending_records) or (
size > capacity and len(writes) > 1
):
middle: Final = len(writes) // 2
await self.record_writes(writes[:middle])
await self.record_writes(writes[middle:])
return
if size > capacity:
raise UsageCaptureError("One usage observation exceeds capture capacity")
while (receipt := self._submit_available(writes)) is None:
self._wakeup.set()
await self._capacity_changed.wait()
self._wakeup.set()
await asyncio.shield(receipt)
def _submit_available(self, writes: tuple[UsageWrite, ...]) -> asyncio.Future[None] | None:
self._capacity_changed.clear()
try:
receipt: Final = self.submit(writes)
except UsageCaptureError:
if self._closed or not self.pending_records or not writes:
raise
return None
return receipt
async def flush(self) -> None:
self._ensure_worker()
self._wakeup.set()
@ -240,6 +301,7 @@ class DurableUsageMeter:
with self._admission:
self.pending_bytes -= item.size
self.pending_records -= len(item.writes)
self._capacity_changed.set()
if item.committed.done():
continue
if error is None:

50
litellm/proxy/utils.py
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
@ -1739,6 +1744,39 @@ def _failure_fields_to_lift(
if _prompt_estimate is not None
else ()
)
from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation
_recorded_observation: Final[object] = _model_call_details.get("_accounting_usage_observed")
_current_usage: Final = (
_current_failure_usage_to_lift(
_model_call_details,
request_data,
_model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None,
_billed,
_deployment_id,
_base_model,
dispatched_call_type,
)
if Logging.usage_capture_callback is not None
and (_billed or _model_call_details.get("combined_usage_object") is not None)
else None
)
_current_observation: Final = (
_recorded_observation
if isinstance(_recorded_observation, UsageObservation)
else UsageObservation.UNKNOWN
if _current_usage is None
else UsageObservation.KNOWN
if _model_call_details.get("combined_usage_object") is not None
else UsageObservation.ESTIMATED
)
_capture_observation: Final = UsageCaptureObservation.from_usage(
_current_usage[0] if _current_usage is not None and isinstance(_current_usage[0], Usage) else None,
_current_observation,
float(_current_usage[1])
if _current_usage is not None and isinstance(_current_usage[1], (int, float))
else None,
)
_entries: Final = (
("start_time", _model_call_details.get("start_time")),
("first_api_call_start_time", _first_handoff),
@ -1747,6 +1785,16 @@ 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", _current_observation if Logging.usage_capture_callback is not None else None),
("_accounting_usage_snapshot", _capture_observation if Logging.usage_capture_callback is not None else None),
("_accounting_prior_captures", _model_call_details.get("_accounting_prior_captures")),
)
return MappingProxyType({key: value for key, value in _entries if value is not None})

10
litellm/rag/main.py
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,

6
litellm/responses/main.py
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,

4
litellm/responses/mcp/chat_completions_handler.py
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):

5
litellm/router.py
View file

@ -3767,6 +3767,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]

6
litellm/types/utils.py
View file

@ -677,6 +677,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"

7
litellm/utils.py
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

7
tests/integration/database/test_accounting_ledger.py
View file

@ -540,7 +540,12 @@ async def test_usage_batch_is_atomic_and_replay_preserves_receipts(
writer: _Writer, ledger_database: str, failure: str
) -> None:
client, store = writer
events: Final = (_usage("caller", inbound=True), _usage("attempt"))
events: Final = (
_usage("caller", inbound=True),
_usage("attempt").model_copy(
update={"metrics": (Metric(name=MetricName.PROMPT_TOKENS, value=Decimal("120000")),)}
),
)
writes: Final = tuple(UsageWrite(event, _envelope(event), _TARGETS) for event in events)
@asynccontextmanager

54
tests/unit/batches/test_batch_utils.py
View file

@ -3436,3 +3436,57 @@ 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
per_request: Final = captured["batch_usage_by_request"]
assert len(per_request) == len(rows)
assert per_request[-1] == {"model": "unregistered-batch-model-c", "usage": None}
reordered: Final = bu.batch_cost_and_usage(list(reversed(rows)), "openai")
assert reordered.usage.model_dump()["batch_usage_by_request"] == per_request
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)
assert sorted(row["usage"]["pages_processed"] for row in result.usage.model_dump()["batch_usage_by_request"]) == [2, 3]
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_usage_by_request"] == ({"model": "missing-usage-model", "usage": None},)
assert captured["batch_unknown_usage_requests"] == 1
assert (result.cost, result.successful_requests, result.failed_requests) == (0, 0, 1)

3
tests/unit/llms/azure/test_azure_common_utils.py
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",

31
tests/unit/proxy/db/test_accounting_ledger.py
View file

@ -116,3 +116,34 @@ async def test_conflicting_usage_within_one_batch_is_rejected_before_database_ac
changed: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json"), "pricing": {"rate": 2}}))
with pytest.raises(AccountingConflictError, match="conflicting copies"):
await AccountingLedger(_no_transaction).insert_usage_batch((original, changed))
@pytest.mark.asyncio
async def test_loaded_usage_preserves_decimal_metrics_from_jsonb(event: InboundUsageEvent) -> None:
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import cast
from unittest.mock import AsyncMock
from prisma import Prisma
from litellm.accounting import OutboundUsageEvent
outbound: Final = OutboundUsageEvent(
**event.model_dump(exclude={"metrics"}),
attempt_id="attempt",
retry_of_attempt_id=None,
metrics=(Metric(name=MetricName.PROMPT_TOKENS, value=Decimal("120000")),),
)
payload: Final = {"event": outbound.model_dump(mode="json"), "usage": {"prompt_tokens": 120000}}
database: Final = AsyncMock(spec=Prisma)
database.query_raw.return_value = [{"payload": payload, "targets": []}]
@asynccontextmanager
async def transaction() -> AsyncIterator[Prisma]:
yield cast(Prisma, database)
recovered: Final = await AccountingLedger(transaction).load_usage_batch((outbound.event_id,))
assert len(recovered) == 1
assert recovered[0].event == outbound
assert json.loads(recovered[0].envelope) == payload

137
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -3949,3 +3949,140 @@ 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
async def test_failure_capture_overload_does_not_skip_budget_release(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier, _write
database: Final = CommitBarrier()
runtime: Final = UsageCaptureRuntime()
runtime.start(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2))
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
monkeypatch.setattr(proxy_server, "prisma_client", None)
meter: Final = runtime.meter
assert meter is not None
receipt: Final = meter.submit((_write("occupied-one"), _write("occupied-two")))
auth: Final = UserAPIKeyAuth(budget_reservation={"reserved_cost": 0.5, "entries": [], "finalized": False})
request: Final = {
"call_type": "acompletion",
"litellm_call_id": "failed",
"model": "capture-model",
"start_time": _CALL_START,
"litellm_params": {},
}
try:
await database.entered.wait()
await ProxyDBLogger().async_post_call_failure_hook(request, RuntimeError("provider failed"), auth)
assert auth.budget_reservation is not None and auth.budget_reservation["finalized"] is True
assert "Usage capture failed on the request failure path" in caplog.text
assert "_accounting_capture" not in request
finally:
database.release.set()
await runtime.stop(1)
await receipt
assert database.committed == ("occupied-one", "occupied-two")
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", [False, True])
@pytest.mark.parametrize("complete", [False, True])
@pytest.mark.parametrize("deferred", [False, True])
async def test_capture_and_legacy_share_resolved_key_attribution(
failure: bool, complete: bool, deferred: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
from pydantic import TypeAdapter
from litellm.accounting import DeliveryTarget, DimensionName
from litellm.proxy import proxy_server
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
from tests.unit.proxy.spend_tracking.test_usage_capture import StoredUsage
key: Final = UserAPIKeyAuth(
api_key="a" * 64,
user_id="owner",
team_id="team",
org_id="org",
project_id="project",
key_alias="key",
team_alias="team-alias",
request_route="/v1/chat/completions",
)
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key="a" * 64, value=key)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "prisma_client", object())
writer: Final = _RecordingSpendWriter()
monkeypatch.setattr(proxy_server, "proxy_logging_obj", SimpleNamespace(db_spend_update_writer=writer))
database: Final = StoredUsage()
runtime: Final = UsageCaptureRuntime()
runtime.start(database.insert, UsageCaptureSettings())
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
rejected: Final = failure and deferred and not complete
owner: Final = "original-owner" if deferred and not rejected else "owner"
metadata: Final = {
"user_api_key": "b" * 64 if rejected else "a" * 64,
**(
{
"user_api_key_user_id": "forged-owner" if rejected else owner,
"user_api_key_team_id": "forged-team" if rejected else "team",
"user_api_key_org_id": "forged-org" if rejected else "org",
"user_api_key_project_id": None if deferred and not complete else "project",
"user_api_key_alias": "key",
"user_api_key_team_alias": "team-alias",
}
if complete or deferred
else {}
),
}
kwargs: Final = {
"call_type": "aretrieve_batch" if deferred else "call_mcp_tool",
"litellm_call_id": "attributed",
"model": "MCP: echo",
"response_cost": 0.25,
"start_time": _CALL_START,
"litellm_params": {"metadata": metadata},
}
if rejected:
from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY
kwargs[REJECTED_BEFORE_REQUEST_SETUP_KEY] = REJECTED_BEFORE_REQUEST_SETUP
logger: Final = ProxyDBLogger(spend_writer=lambda: cast(DBSpendUpdateWriter, writer))
try:
with track_batch_accounting("attributed" if deferred and not rejected else "other"):
if failure:
await logger.async_post_call_failure_hook(
kwargs,
HTTPException(401, "model not allowed"),
key if complete else UserAPIKeyAuth(api_key="a" * 64, request_route="/v1/chat/completions"),
)
else:
await logger.async_log_success_event(kwargs, {"id": "tool"}, _CALL_START, _CALL_START)
finally:
await runtime.stop(1)
assert len(writer.rows) == 1
assert len(database.rows) == 2
dimensions: Final = {item.name: item.value for item in database.rows[-1].event.dimensions}
if rejected:
assert dimensions.get(DimensionName.KEY) == key.api_key
assert {
name: dimensions.get(name)
for name in (DimensionName.USER, DimensionName.TEAM, DimensionName.ORGANIZATION, DimensionName.PROJECT)
} == {
DimensionName.USER: owner,
DimensionName.TEAM: "team",
DimensionName.ORGANIZATION: "org",
DimensionName.PROJECT: None if deferred and not complete and not rejected else "project",
}
assert set(database.rows[-1].targets) >= {DeliveryTarget.DAILY_TEAM_SPEND, DeliveryTarget.DAILY_ORGANIZATION_SPEND}
legacy: Final = TypeAdapter(dict[str, object]).validate_python(writer.rows[0]["kwargs"])
payload: Final = get_logging_payload(legacy, {"id": "tool"}, _CALL_START, _CALL_START)
assert (payload["user"], payload["team_id"], payload["organization_id"]) == (owner, "team", "org")

27
tests/unit/proxy/spend_tracking/test_cost_finalization.py
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:

44
tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -7543,3 +7543,47 @@ 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
@pytest.mark.parametrize("capture_enabled", [False, True])
@pytest.mark.parametrize("completion_start", [None, "absent", "reported"])
def test_spend_payload_preserves_unknown_completion_start(capture_enabled: bool, completion_start: str | None) -> None:
start: Final = datetime.datetime(2026, 1, 1, tzinfo=timezone.utc)
end: Final = start + datetime.timedelta(seconds=2)
reported: Final = start + datetime.timedelta(seconds=1)
observations: Final[list[object]] = []
kwargs: Final = {
"call_type": "acompletion",
"litellm_call_id": "timestamp-call",
"litellm_params": {},
**({} if completion_start == "absent" else {"completion_start_time": reported if completion_start else None}),
}
payload: Final = spend_tracking_utils.get_logging_payload(
kwargs,
{"id": "timestamp-response"},
start,
end,
capture_usage=(lambda row, usage: observations.append(row)) if capture_enabled else None,
)
assert payload["completionStartTime"] == (
end if completion_start == "absent" else reported if completion_start else None
)
assert observations == ([payload] if capture_enabled else [])

1712
tests/unit/proxy/spend_tracking/test_usage_capture.py Normal file
View file

File diff suppressed because it is too large Load diff

84
tests/unit/proxy/spend_tracking/test_usage_meter.py
View file

@ -78,6 +78,13 @@ async def test_byte_limit_rejects_new_work_without_evicting_admitted_work() -> N
meter: Final = DurableUsageMeter(
database.write, UsageCaptureSettings(max_pending_bytes=size), wait_for_batch=_manual_flush
)
reservation: Final = meter.reserve_producer(size // 2, size // 2)
with pytest.raises(UsageCaptureError, match="producer byte capacity"):
meter.reserve_producer(1, size)
with pytest.raises(UsageCaptureError, match="producer byte capacity"):
meter.resize_producer(reservation, size, 1)
assert meter.pending_bytes == size // 2
meter.release_producer(reservation)
receipt: Final = meter.submit((write,))
with pytest.raises(UsageCaptureError, match="capacity"):
meter.submit((_write("extra"),))
@ -88,6 +95,8 @@ async def test_byte_limit_rejects_new_work_without_evicting_admitted_work() -> N
assert database.committed == (write.event.event_id,)
with pytest.raises(UsageCaptureError, match="closed"):
meter.submit((_write("after-close"),))
with pytest.raises(UsageCaptureError, match="producer byte capacity"):
meter.reserve_producer(1, 1)
@pytest.mark.asyncio
@ -387,3 +396,78 @@ async def test_shutdown_does_not_cancel_dependency_receipt_or_leak_reserved_capa
assert meter.pending_records == 0
assert meter.pending_bytes == 0
prerequisite.set_result(None)
@pytest.mark.asyncio
async def test_large_write_sequence_drains_within_byte_and_record_bounds() -> None:
database: Final = CommitBarrier()
database.release.set()
meter: Final = DurableUsageMeter(
database.write,
UsageCaptureSettings(batch_size=8, max_pending_records=2, max_pending_bytes=1024),
wait_for_batch=_manual_flush,
)
writes: Final = tuple(_write(str(index)) for index in range(20))
pending: Final = meter.submit((_write("pending-a"), _write("pending-b")))
await meter.record_writes(writes)
await pending
with pytest.raises(UsageCaptureError, match="One usage observation"):
await meter.record_writes((_write("oversized" * 300),))
await meter.close()
assert database.committed == ("pending-a", "pending-b", *(write.event.event_id for write in writes))
assert database.transactions > 1
assert (meter.pending_bytes, meter.pending_records) == (0, 0)
with pytest.raises(UsageCaptureError):
await meter.record_writes((_write("closed"),))
@pytest.mark.asyncio
async def test_cancelling_bounded_capture_preserves_admitted_writes_without_admitting_remaining_rows() -> None:
database: Final = CommitBarrier()
meter: Final = DurableUsageMeter(
database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2), wait_for_batch=_manual_flush
)
recording: Final = asyncio.create_task(meter.record_writes(tuple(_write(str(index)) for index in range(8))))
await database.entered.wait()
recording.cancel()
with pytest.raises(asyncio.CancelledError):
await recording
database.release.set()
await meter.close()
assert database.committed == ("0", "1")
assert meter.pending_bytes == meter.pending_records == 0
@pytest.mark.asyncio
async def test_capacity_wait_retries_admission_before_other_producers_finish() -> None:
entered: Final = asyncio.Queue[tuple[str, ...]]()
release: Final = asyncio.Queue[None]()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
entered.put_nowait(tuple(write.event.event_id for write in writes))
await release.get()
meter: Final = DurableUsageMeter(
persist, UsageCaptureSettings(batch_size=2, max_pending_records=3), wait_for_batch=_manual_flush
)
initial: Final = meter.submit((_write("first"), _write("second")))
remaining: Final = meter.submit((_write("remaining"),))
assert await entered.get() == ("first", "second")
recording: Final = asyncio.create_task(meter.record_writes((_write("target"),)))
try:
await asyncio.sleep(0)
release.put_nowait(None)
assert await entered.get() == ("remaining",)
await asyncio.sleep(0)
assert meter.pending_records == 2
release.put_nowait(None)
assert await entered.get() == ("target",)
release.put_nowait(None)
await recording
await asyncio.gather(initial, remaining)
finally:
recording.cancel()
await asyncio.gather(recording, return_exceptions=True)
release.put_nowait(None)
release.put_nowait(None)
await meter.close()

2
tests/unit/proxy/test_chat_completion_metadata.py
View file

@ -119,7 +119,7 @@ async def test_completion_metadata_population():
# Verify
call_args = MockProcessor.call_args
assert call_args is not None
data_arg = call_args.kwargs.get("data")
data_arg = mock_instance.data
assert data_arg is not None
assert "metadata" in data_arg

302
tests/unit/proxy/test_proxy_server.py
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,275 @@ 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
@pytest.mark.asyncio
async def test_capture_configuration_requires_restart_without_replacing_the_active_writer(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.utils import PrismaClient
runtime: Final = UsageCaptureRuntime()
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True})
with pytest.raises(ValueError, match="requires accounting_capture"):
proxy_server._configure_usage_capture()
client: Final = PrismaClient("postgresql://localhost/unused", ProxyLogging(UserApiKeyCache()))
monkeypatch.setattr(proxy_server, "prisma_client", client)
monkeypatch.setattr(litellm, "apply_search_free_tier", False)
monkeypatch.setenv("LITELLM_COLLECTOR_ENABLED", "false")
monkeypatch.delenv("SPEND_LOGS_URL", raising=False)
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 4}})
try:
proxy_server._configure_usage_capture()
active: Final = runtime.meter
assert active is not None
assert active.settings.batch_size == 4
proxy_server._configure_usage_capture()
assert runtime.meter is active
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_capture": {"batch_size": 6}})
with pytest.raises(ValueError, match="Changing accounting_capture requires a worker restart"):
proxy_server._configure_usage_capture()
monkeypatch.setattr(proxy_server, "general_settings", {})
with pytest.raises(ValueError, match="Disabling accounting_capture requires a worker restart"):
proxy_server._configure_usage_capture()
assert runtime.meter is active
assert active.settings.batch_size == 4
finally:
await runtime.stop(1)
assert runtime.meter is None
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_ack", [False, True])
@pytest.mark.parametrize(
("path", "payload"),
[
(
"/v1/chat/completions",
{"model": "openai/gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True},
),
("/v1/responses", {"model": "openai/gpt-5-mini", "input": "hi", "background": True}),
(
"/v1/messages",
{"model": "openai/gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1},
),
],
)
async def test_sync_ack_preserves_async_request_support(
sync_ack: bool,
path: str,
payload: dict[str, object],
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
from fastapi import Request
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.db.accounting_ledger import UsageWrite
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
for write in writes:
persisted.put_nowait(write)
async def authorize(request: Request) -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key")
upstream: Final = {
"id": "provider-result",
"object": "chat.completion",
"created": 1,
"model": "gpt-5-mini",
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}],
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
}
if payload.get("stream"):
chunk: Final = {
**upstream,
"object": "chat.completion.chunk",
"choices": [{"index": 0, "finish_reason": "stop", "delta": {"role": "assistant", "content": "hi"}}],
}
respx_mock.post().respond(
200, text=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n", headers={"content-type": "text/event-stream"}
)
elif path == "/v1/responses":
respx_mock.post().respond(
200,
json={
"id": "resp_background",
"object": "response",
"created_at": 1,
"model": "gpt-5-mini",
"status": "in_progress",
"background": True,
"output": [],
"usage": None,
},
)
else:
respx_mock.post(url__regex=r".*/responses$").respond(
200,
json={
"id": "resp_message",
"object": "response",
"created_at": 1,
"model": "gpt-5-mini",
"status": "completed",
"output": [
{
"id": "msg_reply",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "hi", "annotations": []}],
}
],
"usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
},
)
respx_mock.post(url__regex=r".*/chat/completions$").respond(200, json=upstream)
runtime: Final = UsageCaptureRuntime()
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
monkeypatch.setattr(
proxy_server, "general_settings", {"accounting_sync_ack": sync_ack, "accounting_sync_ack_timeout_seconds": 0}
)
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(
proxy_server,
"llm_router",
litellm.Router(
model_list=[
{
"model_name": payload["model"],
"litellm_params": {"model": payload["model"], "api_key": "test-provider-key"},
}
],
num_retries=0,
),
)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache()))
monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()])
monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize)
runtime.start(persist, UsageCaptureSettings())
try:
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test"
) as client:
response: Final = await client.post(path, json=payload)
finally:
await runtime.stop(1)
assert response.status_code == 200, response.text
assert len(respx_mock.calls) == 1
if payload.get("background"):
assert response.json()["status"] == "in_progress"
else:
assert "hi" in response.text
assert not persisted.empty(), "Terminal usage must still be captured"
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_ack", [False, True])
async def test_text_completion_provider_rejection_captures_unknown_usage(
sync_ack: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
from fastapi import Request
from litellm.accounting import Outcome
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.db.accounting_ledger import UsageWrite
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
for write in writes:
persisted.put_nowait(write)
async def authorize(request: Request) -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key")
respx_mock.post().respond(
404,
json={
"error": {"message": "The model is unavailable", "type": "invalid_request_error", "code": "model_not_found"}
},
)
runtime: Final = UsageCaptureRuntime()
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": sync_ack})
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(
proxy_server,
"llm_router",
litellm.Router(
model_list=[
{
"model_name": "instruct",
"litellm_params": {"model": "openai/gpt-3.5-turbo-instruct", "api_key": "test-provider-key"},
}
],
num_retries=0,
),
)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache()))
monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()])
monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize)
runtime.start(persist, UsageCaptureSettings())
try:
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test"
) as client:
response: Final = await client.post(
"/v1/completions", json={"model": "instruct", "prompt": "hi", "max_tokens": 1}
)
finally:
await runtime.stop(1)
assert response.status_code == 404, response.text
assert len(respx_mock.calls) == 1
assert persisted.qsize() == 2
inbound: Final = persisted.get_nowait()
outbound: Final = persisted.get_nowait()
assert outbound.event.outcome is Outcome.FAILED
assert outbound.event.call_type.value == "atext_completion"
assert outbound.event.metrics is None
assert inbound.event.request_id == outbound.event.request_id
envelope: Final = json.loads(outbound.envelope)
assert envelope["usage"] is None
assert envelope["usage_source"] == "unknown"
assert envelope["pricing"]["reported_cost"] is None

54
tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py
View file

@ -4158,3 +4158,57 @@ def test_is_client_error_exception_existing_behaviour_unchanged() -> None:
ProxyException(message="x", type="x", param=None, code="429")
) is True
assert _is_client_error_exception(ValueError("x")) is False
@pytest.mark.asyncio
@pytest.mark.parametrize("tokens", [None, 0, 3])
async def test_failure_capture_preserves_prepared_usage_observation(tokens: int | None) -> None:
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
timestamp: Final = real_datetime.datetime(2026, 1, 1)
committed: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue()
async def persist(writes: tuple[UsageWrite, ...]) -> None:
for write in writes:
committed.put_nowait(write)
logger: Final = Logging(
model=FAILURE_USAGE_MODEL,
messages=[],
stream=False,
call_type="acompletion",
start_time=timestamp,
litellm_call_id="post-response-failure",
function_id="post-response-failure",
)
logger.update_environment_variables(
model=FAILURE_USAGE_MODEL, custom_llm_provider="openai", litellm_params={}, optional_params={}
)
logger.model_call_details["response_cost"] = 0.2
runtime: Final = UsageCaptureRuntime()
runtime.start(persist, UsageCaptureSettings())
try:
prepared: Final = logger.prepare_response_cost(
{"id": "provider", "usage": None if tokens is None else {"prompt_tokens": tokens, "total_tokens": tokens}},
timestamp,
)
assert prepared is not None
lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True)
assert lifted["combined_usage_object"] is logger.model_call_details["combined_usage_object"]
assert lifted["response_cost"] == 0.2
assert runtime.meter is not None
capture_callback(
runtime.meter,
{**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"},
None,
timestamp,
timestamp,
)
finally:
await runtime.stop(1)
committed.get_nowait()
envelope: Final = json.loads(committed.get_nowait().envelope)
assert (None if envelope["usage"] is None else envelope["usage"]["total_tokens"]) == tokens
assert envelope["usage_source"] == ("unknown" if tokens is None else "measured")

46
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
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 on foreground non-stream chat, text, embedding and Responses requests. Other requests remain supported with asynchronous usage capture.
* @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, including the batch flush interval
* @default 20
*/
accounting_sync_ack_timeout_seconds: number;
/**
* Admission Queue Timeout Seconds
* @description maximum time a request waits for a worker slot
@ -49503,6 +49524,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 5
*/
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 */