From 02b8da2f169b6e30a7ade9e1b7df4a9099294916 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Wed, 4 Nov 2026 10:37:07 -0800 Subject: [PATCH] refactor: replace accounting event scaffolding with SpendLogs projection (#49843) * refactor(accounting): replace event scaffolding with SpendLogs projection * test: preserve spend timestamps without removed capture callback * fix: guard accounting event removal across schema update paths * refactor(accounting): retain schema and prepare durable payloads * fix(accounting): preserve dates in serialized payloads * fix(accounting): retain schema push compatibility * fix(accounting): retain span registration for the legacy schema * test(accounting): remove obsolete capture acknowledgement test * fix(accounting): date every daily projection row the way the model rollup does A space-separated startTime kept the whole timestamp as the daily key, while the model usage rollup took the calendar date. Every daily row now uses the same date helper Co-Authored-By: Claude Opus 5.5 * chore(ui): regenerate dashboard API types after rebasing on main Co-Authored-By: Claude Opus 5.5 * refactor(accounting): project only from the stored accounting payload project() now reads entities only from the payload's accounting_entities, so the raw SpendLogs fallback branches are gone. build_accounting_payload() builds the metadata once instead of dumping it to JSON and parsing it back. Guardrail usage dates use a tolerant UTC parse, so a malformed startTime drops only the guardrail rows instead of failing the whole projection. Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .../common_utils/check_responses_cost.py | 1 - litellm/accounting/__init__.py | 5 +- litellm/accounting/projection.py | 304 +++ litellm/accounting/protocols.py | 74 +- litellm/accounting/snapshot.py | 97 + litellm/litellm_core_utils/litellm_logging.py | 165 +- .../litellm_core_utils/realtime_streaming.py | 2 - litellm/proxy/_types.py | 17 - litellm/proxy/common_request_processing.py | 4 - litellm/proxy/db/accounting_ledger.py | 352 ---- litellm/proxy/db/model_usage_rollup.py | 4 +- .../proxy/hooks/proxy_track_cost_callback.py | 125 +- litellm/proxy/proxy_server.py | 33 - .../spend_tracking/spend_tracking_utils.py | 12 +- litellm/proxy/spend_tracking/usage_capture.py | 1084 ---------- litellm/proxy/spend_tracking/usage_meter.py | 330 ---- litellm/proxy/utils.py | 38 - litellm/router_python_core.py | 5 - litellm/utils.py | 7 - .../database/test_accounting_ledger.py | 610 ------ tests/unit/accounting/test_projection.py | 431 ++++ tests/unit/accounting/test_types.py | 2 - ...responses_transformation_transformation.py | 49 - tests/unit/proxy/db/test_accounting_ledger.py | 149 -- .../hooks/test_proxy_track_cost_callback.py | 137 -- .../spend_tracking/test_cost_finalization.py | 24 - .../test_spend_tracking_utils.py | 6 +- .../spend_tracking/test_usage_capture.py | 1738 ----------------- .../proxy/spend_tracking/test_usage_meter.py | 473 ----- tests/unit/proxy/test_proxy_server.py | 335 ---- ..._utils_model_creation_and_error_logging.py | 90 - ui/litellm-dashboard/src/lib/http/schema.d.ts | 44 - 32 files changed, 845 insertions(+), 5902 deletions(-) create mode 100644 litellm/accounting/projection.py create mode 100644 litellm/accounting/snapshot.py delete mode 100644 litellm/proxy/db/accounting_ledger.py delete mode 100644 litellm/proxy/spend_tracking/usage_capture.py delete mode 100644 litellm/proxy/spend_tracking/usage_meter.py delete mode 100644 tests/integration/database/test_accounting_ledger.py create mode 100644 tests/unit/accounting/test_projection.py delete mode 100644 tests/unit/proxy/db/test_accounting_ledger.py delete mode 100644 tests/unit/proxy/spend_tracking/test_usage_capture.py delete mode 100644 tests/unit/proxy/spend_tracking/test_usage_meter.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index cfe4ca98145..69e36a0c076 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -200,7 +200,6 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi unpriced: Final = _without_read_price(billed.response) cost: Final = logging_obj.response_cost_calculator(result=unpriced) end_time: Final = datetime.now() - logging_obj.capture_usage_before_logging(unpriced, logging_obj.start_time, end_time) GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( async_coroutine=logging_obj.async_success_handler( result=unpriced, start_time=logging_obj.start_time, end_time=end_time diff --git a/litellm/accounting/__init__.py b/litellm/accounting/__init__.py index d4e56f61279..5d59a35a987 100644 --- a/litellm/accounting/__init__.py +++ b/litellm/accounting/__init__.py @@ -1,7 +1,7 @@ """Public accounting contracts. Importing these values does not enable accounting.""" from litellm.accounting.measurements import measurements_for -from litellm.accounting.protocols import BudgetEnforcer, CostAggregator, CostCompiler, UsageMeter +from litellm.accounting.protocols import CostAggregator from litellm.accounting.types import ( AdmissionDecision, AggregateQuery, @@ -33,9 +33,7 @@ __all__ = [ "AdmissionDecision", "AggregateQuery", "AggregateRow", - "BudgetEnforcer", "CostAggregator", - "CostCompiler", "CostEvent", "DeliveryState", "DeliveryTarget", @@ -56,7 +54,6 @@ __all__ = [ "PricingSnapshot", "Rate", "Unpriced", - "UsageMeter", "measurements_for", "validate_aggregate_rows", ] diff --git a/litellm/accounting/projection.py b/litellm/accounting/projection.py new file mode 100644 index 00000000000..fc3a89ae306 --- /dev/null +++ b/litellm/accounting/projection.py @@ -0,0 +1,304 @@ +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from datetime import datetime, timezone +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.constants import ( + BILLED_RETRY_ATTEMPT_METADATA_KEY, + INTERNAL_CALL_ORIGIN_METADATA_KEY, + MODEL_INSIGHTS_DEFAULT_TASK, + MODEL_INSIGHTS_TASK_TAG_PREFIX, +) +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import GuardrailCostByUnitEntry +from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, daily_spend_entity_ids +from litellm.proxy.db.model_insights_tasks import load_model_insight_tasks +from litellm.proxy.db.model_usage_rollup import date_from_start_time +from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING +from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens +from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + +_FIELDS: Final = TypeAdapter(dict[str, object]) +_STRINGS: Final = TypeAdapter(tuple[str, ...]) +_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...]) +_TASK_TYPES: Final = frozenset(load_model_insight_tasks()) + + +@dataclass(frozen=True, slots=True) +class Contribution: + table: str + key: tuple[tuple[str, str], ...] + increments: tuple[tuple[str, int | float], ...] + values: tuple[tuple[str, str | None], ...] = () + + +def _fields(value: object) -> Mapping[str, object]: + if isinstance(value, str): + return _FIELDS.validate_json(value) + return _FIELDS.validate_python(value) if isinstance(value, Mapping) else MappingProxyType({}) + + +def _text(value: object) -> str: + return "" if value is None else str(value) + + +def _number(value: object) -> int | float: + return value if isinstance(value, (int, float)) else 0 + + +def _tags(row: Mapping[str, object]) -> tuple[str, ...]: + value: Final = row.get("request_tags") + if value is None: + return () + return _STRINGS.validate_json(value) if isinstance(value, str) else _STRINGS.validate_python(value) + + +def _date(row: Mapping[str, object]) -> str | None: + start: Final = row.get("startTime") + return date_from_start_time(start) if isinstance(start, (datetime, str)) else None + + +def _spend(table: str, key: tuple[tuple[str, str], ...], spend: float) -> Contribution: + return Contribution(table, key, (("spend", spend),)) + + +def _entity_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + spend: Final = _number(row.get("spend")) + entities: Final = _fields(metadata.get("accounting_entities")) + identities: Final = ( + ("LiteLLM_VerificationToken", "token", entities.get("api_key")), + ("LiteLLM_UserTable", "user_id", entities.get("user")), + ("LiteLLM_UserTable", "user_id", entities.get("proxy_budget_name")), + ("LiteLLM_TeamTable", "team_id", entities.get("team_id")), + ("LiteLLM_OrganizationTable", "organization_id", entities.get("organization_id")), + ("LiteLLM_EndUserTable", "user_id", entities.get("end_user")), + ("LiteLLM_AgentsTable", "agent_id", entities.get("agent_id")), + ("LiteLLM_ProjectTable", "project_id", entities.get("project_id")), + ) + groups: Final = _STRINGS.validate_python(metadata.get("accounting_model_access_groups", ())) + return ( + *( + _spend(table, ((column, identity),), spend) + for table, column, identity in identities + if isinstance(identity, str) + ), + *(_spend("LiteLLM_TagTable", (("tag_name", tag),), spend) for tag in _tags(row) if tag), + *(_spend("LiteLLM_ModelAccessGroupBudgetTable", (("access_group_name", group),), spend) for group in groups), + *_membership_contributions(entities, spend), + ) + + +def _membership_contributions(row: Mapping[str, object], spend: float) -> tuple[Contribution, ...]: + user: Final = row.get("user") + if not isinstance(user, str): + return () + owners: Final = ( + ("LiteLLM_TeamMembership", "team_id", row.get("team_id")), + ("LiteLLM_OrganizationMembership", "organization_id", row.get("organization_id")), + ) + return tuple( + _spend(table, ((column, identity), ("user_id", user)), spend) + for table, column, identity in owners + if isinstance(identity, str) + ) + + +def _daily_counters(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[tuple[str, int | float], ...]: + usage: Final = _fields(metadata.get("usage_object")) + savings: Final = _fields(metadata.get("accounting_savings")) + internal: Final = bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)) + caller: Final = not internal and metadata.get(BILLED_RETRY_ATTEMPT_METADATA_KEY) is not True + success: Final = metadata.get("status") != "failure" + duration: Final = row.get("request_duration_ms") + timed: Final = not internal and success and isinstance(duration, int) and duration >= 0 + return ( + ("spend", _number(row.get("spend"))), + ("prompt_tokens", _number(row.get("prompt_tokens"))), + ("completion_tokens", _number(row.get("completion_tokens"))), + ("api_requests", int(caller)), + ("successful_requests", int(caller and success)), + ("failed_requests", int(caller and not success)), + ("cache_read_input_tokens", extract_cache_read_tokens(usage)), + ("cache_creation_input_tokens", extract_cache_creation_tokens(usage)), + ("compression_saved_tokens", extract_compression_saved_tokens(metadata)), + ("compression_savings_spend", _number(savings.get("compression"))), + ("prompt_caching_savings_spend", _number(savings.get("prompt_caching"))), + ("gateway_injected_caching_savings_spend", _number(savings.get("gateway_injected_caching"))), + ("autorouter_savings_spend", 0.0 if internal else _number(savings.get("autorouter"))), + ("total_response_time_ms", _number(duration) if timed else 0), + ("timed_requests", int(timed)), + ) + + +def _daily_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + date: Final = _date(row) + if date is None or "api_key" not in row: + return () + if "mcp_namespaced_tool_name" not in row and ( + "model" not in row or "custom_llm_provider" not in row or "model_group" not in row + ): + return () + counters: Final = _daily_counters(row, metadata) + endpoint: Final = ROUTE_ENDPOINT_MAPPING.get(_text(row.get("call_type"))) + shared_key: Final = ( + ("date", date), + ("api_key", _text(row.get("api_key"))), + ("model", _text(row.get("model"))), + ("custom_llm_provider", _text(row.get("custom_llm_provider"))), + ("mcp_namespaced_tool_name", _text(row.get("mcp_namespaced_tool_name"))), + ("endpoint", _text(endpoint)), + ) + # The legacy daily organization writer uses the charged org argument, not + # SpendLogs' normalized empty-string identity for a request with no org. + daily_row: Final = {**row, "organization_id": _fields(metadata.get("accounting_entities")).get("organization_id")} + return tuple(_daily_rows(daily_row, shared_key, counters)) + + +def _daily_rows( + row: Mapping[str, object], + shared_key: tuple[tuple[str, str], ...], + counters: tuple[tuple[str, int | float], ...], +) -> Iterator[Contribution]: + for entity, table in DAILY_SPEND_TABLES.items(): + for identity in daily_spend_entity_ids(row, entity): + yield Contribution( + table.name, + ((table.entity_id_column, _text(identity)), *shared_key), + counters, + ( + ("model_group", None if row.get("model_group") is None else _text(row.get("model_group"))), + *((("request_id", _text(row.get("request_id"))),) if table.carries_request_id else ()), + ), + ) + + +def _window_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + windows: Final = _ROWS.validate_python(metadata.get("accounting_windows", ())) + return tuple( + Contribution( + "LiteLLM_BudgetWindowSpend", + tuple((column, _text(window[column])) for column in ("entity_type", "entity_id", "window_duration")), + (("spend", _number(window.get("spend", row.get("spend")))),), + (("window_start", _text(window["window_start"])), ("started_at", _text(row.get("startTime")))), + ) + for window in windows + ) + + +def project(row: Mapping[str, object]) -> tuple[Contribution, ...]: + metadata: Final = _fields(row.get("metadata")) + return ( + *_entity_contributions(row, metadata), + *_daily_contributions(row, metadata), + *_window_contributions(row, metadata), + *_model_contributions(row, metadata), + *_guardrail_contributions(row, metadata), + ) + + +def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + date: Final = _date(row) + origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) + if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN): + return () + model: Final = _text(row.get("model")) or "unknown" + caller: Final = not origin + success: Final = row.get("status") == "success" + return ( + Contribution( + "LiteLLM_DailyModelUsage", + ( + ("date", date), + ("model_group", _text(row.get("model_group")) or model), + ("model", model), + ("custom_llm_provider", _text(row.get("custom_llm_provider")) or "unknown"), + ("task_type", _model_task(row)), + ), + ( + ("spend", _number(row.get("spend"))), + ("prompt_tokens", _number(row.get("prompt_tokens"))), + ("completion_tokens", _number(row.get("completion_tokens"))), + ("request_count", int(caller)), + ("successful_requests", int(caller and success)), + ("failed_requests", int(caller and not success)), + ), + ), + ) + + +def _utc_date(row: Mapping[str, object]) -> str | None: + start: Final = row.get("startTime") + try: + parsed: Final = datetime.fromisoformat(start.replace("Z", "+00:00")) if isinstance(start, str) else start + except ValueError: + return None + if not isinstance(parsed, datetime): + return None + return ( + (parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed) + .astimezone(timezone.utc) + .date() + .isoformat() + ) + + +def _guardrail_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]: + date: Final = _utc_date(row) + info: Final = metadata.get("guardrail_information") or metadata.get("standard_logging_guardrail_information") + if not row.get("request_id") or date is None or not isinstance(info, list): + return () + return tuple(chain.from_iterable(_guardrail_entry(row, date, entry) for entry in _ROWS.validate_python(info))) + + +def _guardrail_entry(row: Mapping[str, object], date: str, entry: Mapping[str, object]) -> tuple[Contribution, ...]: + identity: Final = _text(entry.get("guardrail_id") or entry.get("guardrail_name")) + usage: Final = entry.get("guardrail_usage") + if not identity or not isinstance(usage, Mapping): + return () + costs: Final = _guardrail_unit_costs(entry) + return tuple( + _guardrail_unit(row, (identity, date, unit), units, None if costs is None else costs.get(unit)) + for unit, units in _fields(usage).items() + if isinstance(units, int) and not isinstance(units, bool) and units > 0 + ) + + +def _guardrail_unit( + row: Mapping[str, object], identity: tuple[str, str, str], units: int, cost: float | None +) -> Contribution: + guardrail_id, date, unit = identity + return Contribution( + "LiteLLM_DailyGuardrailUsageUnits", + ( + ("guardrail_id", guardrail_id), + ("date", date), + ("team_id", _text(row.get("team_id"))), + ("api_key", _text(row.get("api_key"))), + ("usage_unit", unit), + ), + (("units", units), ("cost", cost or 0.0), ("untracked_units", units if cost is None else 0)), + ) + + +def _model_task(row: Mapping[str, object]) -> str: + return next( + ( + task + for tag in _tags(row) + if tag.startswith(MODEL_INSIGHTS_TASK_TAG_PREFIX) + and (task := tag.removeprefix(MODEL_INSIGHTS_TASK_TAG_PREFIX)) in _TASK_TYPES + ), + MODEL_INSIGHTS_DEFAULT_TASK, + ) + + +def _guardrail_unit_costs(entry: Mapping[str, object]) -> Mapping[str, float | None] | None: + try: + cost: Final = GuardrailCostByUnitEntry.model_validate(entry) + except ValidationError: + return None + return None if cost.guardrail_cost_in_spend is False else cost.guardrail_cost_by_unit diff --git a/litellm/accounting/protocols.py b/litellm/accounting/protocols.py index 0208e37fbdf..f6ec61983a0 100644 --- a/litellm/accounting/protocols.py +++ b/litellm/accounting/protocols.py @@ -1,58 +1,14 @@ -"""Contracts only. Implementations own persistence, pricing, projections and live budgets. - -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, 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. -""" +from __future__ import annotations +from collections.abc import Mapping from typing import Protocol -from litellm.accounting.types import ( - AdmissionDecision, - AggregateQuery, - AggregateRow, - CostEvent, - Dimension, - InboundUsageEvent, - Measurement, - Metric, - OutboundUsageEvent, - Unpriced, -) - - -class UsageMeter(Protocol): - async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None: - """Acknowledge commit, no-op identical IDs and reject conflicting content. - - Persist the event and its request-scoped pricing context before acknowledging. - HTTP acknowledgement is a separate policy. Unknown usage remains None. - """ - ... - - -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 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. - """ - ... +from litellm.accounting.types import AggregateQuery, AggregateRow class CostAggregator(Protocol): - async def apply(self, measurements: tuple[Measurement, ...]) -> None: - """Apply one source event's complete contribution set per target/version. - - Each target commits all contributions and its receipt atomically. A receipt - for the first metric must not suppress the rest. Usage need not wait for cost. - """ + async def apply(self, rows: tuple[Mapping[str, object], ...]) -> None: + """Apply stored request contributions atomically with their delivery receipts.""" ... async def query(self, query: AggregateQuery) -> tuple[AggregateRow, ...]: @@ -63,23 +19,3 @@ class CostAggregator(Protocol): combinations. Different groups may overlap, including non-additive tags. """ ... - - -class BudgetEnforcer(Protocol): - async def reserve( - self, reservation_id: str, dimensions: tuple[Dimension, ...], estimated: tuple[Metric, ...] - ) -> AdmissionDecision: - """Apply existing admission policies and estimates; duplicate IDs add no hold.""" - ... - - async def settle(self, reservation_id: str | None, cost_event: CostEvent) -> None: - """Settle the committed charge with atomic deduplication and hold mutation. - - Preserve scope/window attribution and retry unfinished settlement after - restart. SQL projection receipts do not provide live-counter atomicity. - """ - ... - - async def release(self, reservation_id: str) -> None: - """Idempotently release unused exposure at the request boundary, preserving charges.""" - ... diff --git a/litellm/accounting/snapshot.py b/litellm/accounting/snapshot.py new file mode 100644 index 00000000000..51c7bbad81b --- /dev/null +++ b/litellm/accounting/snapshot.py @@ -0,0 +1,97 @@ +from collections.abc import Callable, Mapping, Sequence +from datetime import datetime +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure +from litellm.proxy._types import SpendLogsPayload +from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens +from litellm.proxy.spend_tracking.savings import compute_savings_spend, marks_gateway_injection + +if TYPE_CHECKING: + from litellm.router import Router + +_FIELDS: Final = TypeAdapter(dict[str, object]) + + +def _mapping(value: object) -> Mapping[str, object] | None: + return _FIELDS.validate_python(value) if isinstance(value, Mapping) else None + + +_METADATA_FIELDS: Final = frozenset( + { + "usage_object", + "status", + "guardrail_information", + "standard_logging_guardrail_information", + "compression_savings", + "internal_call_origin", + "billed_retry_attempt", + } +) + + +_ACCOUNTING_FIELDS: Final = frozenset( + { + "request_id", + "call_type", + "api_key", + "spend", + "prompt_tokens", + "completion_tokens", + "startTime", + "model", + "model_group", + "custom_llm_provider", + "user", + "team_id", + "organization_id", + "end_user", + "agent_id", + "request_tags", + "status", + "request_duration_ms", + "mcp_namespaced_tool_name", + } +) + + +def build_accounting_payload( + payload: SpendLogsPayload, + *, + entities: Mapping[str, str | None], + model_access_groups: Sequence[str], + windows: Sequence[Mapping[str, object]], + router: "Callable[[], Router | None]", +) -> Mapping[str, object]: + metadata: Final = _FIELDS.validate_json(payload["metadata"]) + savings: Final = compute_savings_spend( + model=payload["model"], + custom_llm_provider=payload["custom_llm_provider"], + compression_saved_tokens=extract_compression_saved_tokens(metadata), + gateway_injected_cache=marks_gateway_injection(metadata, payload["model_id"]), + routing_decision=_mapping(metadata.get("routing_decision")), + model_id=payload["model_id"], + llm_router=router, + usage_object=_mapping(metadata.get("usage_object")), + cost_breakdown=_mapping(metadata.get("cost_breakdown")), + recorded_autorouter_savings=metadata.get("autorouter_savings"), + recorded_autorouter_savings_estimate=_mapping(metadata.get("autorouter_savings_estimate")), + billed_at=payload["endTime"], + ) + accounting_metadata: Final = { + **{key: value for key, value in metadata.items() if key in _METADATA_FIELDS}, + "accounting_entities": entities, + "accounting_savings": savings._asdict(), + "accounting_model_access_groups": model_access_groups, + "accounting_windows": windows, + } + return { + **{ + key: value.isoformat() if isinstance(value, datetime) else value + for key, value in payload.items() + if key in _ACCOUNTING_FIELDS + }, + "metadata": _FIELDS.validate_python(safe_json_structure(accounting_metadata, non_finite_as_null=True)), + } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f8080531961..5494d0d82b7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,7 +19,7 @@ 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, ClassVar, Final, Literal, NamedTuple, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, TypeVar, cast from httpx import Response, Timeout from openai import AsyncStream as OpenAIAsyncStream @@ -283,8 +283,6 @@ if TYPE_CHECKING: 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, @@ -662,54 +660,7 @@ class UsageObservation(Enum): 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, \ @@ -949,7 +900,6 @@ class Logging(LiteLLMLoggingBaseClass): self._stash_owed_on_success: bool = False self._stash_billed_by_continuation: bool = False self._failed_attempts_owed_on_success: tuple[ServedRound, ...] = () - self._failed_attempt_ids_owed_on_success: tuple[str, ...] = () self.failed_attempts_billed_on_success: tuple[ServedRound, ...] = () self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None @@ -2582,11 +2532,6 @@ 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, @@ -2698,20 +2643,12 @@ class Logging(LiteLLMLoggingBaseClass): self.finalized_cost = finalized self.model_call_details["response_cost"] = finalized.amount - def record_usage_capture(self, captured: "CapturedUsage") -> None: - self.model_call_details["_accounting_capture"] = captured - - def record_usage_pricing(self, pricing: object) -> None: - self.model_call_details["_accounting_pricing_snapshot"] = pricing - def record_partial_usage_for_failure( self, usage: Usage, response_cost: float, *, 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 @@ -2795,12 +2732,6 @@ class Logging(LiteLLMLoggingBaseClass): owed_on_success: Final = self._stash_owed_on_success self._stash_billed_by_continuation = False self._stash_owed_on_success = False - captured_attempt_ids: Final = ( - self.capture_failed_attempt_usage(usage) - if isinstance(usage, Usage) and type(self).usage_capture_callback is not None - else () - ) - 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) @@ -2816,60 +2747,10 @@ class Logging(LiteLLMLoggingBaseClass): if not owed_on_success or billed_by_continuation: return self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, attempt) - self._failed_attempt_ids_owed_on_success = (*self._failed_attempt_ids_owed_on_success, *captured_attempt_ids) def mark_stash_billed_by_continuation(self) -> None: self._stash_billed_by_continuation = True - def capture_failed_attempt_usage(self, usage: Usage) -> tuple[str, ...]: - 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 not isinstance(captured, CapturedUsage): - return () - self.model_call_details["_accounting_prior_captures"] = ( - *self.model_call_details.get("_accounting_prior_captures", ()), - captured, - ) - return tuple(write.event.event_id for write in captured.writes) - 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 @@ -2909,8 +2790,6 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details.pop("failed_attempts_usage", None) self.failed_attempts_billed_on_success = failed_attempts self.model_call_details["failed_attempts_billed_on_success"] = failed_attempts - if self._failed_attempt_ids_owed_on_success: - self.model_call_details["_accounting_included_attempt_ids"] = self._failed_attempt_ids_owed_on_success self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects( usage_objects=[own_usage, *(usage for usage, _ in failed_attempts)] ) @@ -3073,45 +2952,6 @@ class Logging(LiteLLMLoggingBaseClass): 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. @@ -3826,7 +3666,6 @@ class Logging(LiteLLMLoggingBaseClass): self._as_served_chat_response(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( @@ -4352,8 +4191,6 @@ 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") diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 8e6e71a31d3..b6af181980e 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -836,8 +836,6 @@ class RealTimeStreaming: if self.session_tools or self.tool_calls: self.logging_obj.model_call_details["realtime_tools"] = self.session_tools self.logging_obj.model_call_details["realtime_tool_calls"] = self.tool_calls - if LiteLLMLogging.usage_capture_callback is not None: - self.logging_obj.capture_usage_before_logging(self.messages) # Route through the bounded logging worker (per-coroutine timeout + # concurrency cap) instead of a bare create_task, so a slow callback # can't leave suspended tasks pinning each call's response in memory. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index de1096701af..9d84770b288 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -38,7 +38,6 @@ 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 ) @@ -2481,22 +2480,6 @@ 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.", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5ae6c628324..17e6d3ee411 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -140,7 +140,6 @@ 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, @@ -3397,9 +3396,6 @@ 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: diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py deleted file mode 100644 index 0a9c5db1619..00000000000 --- a/litellm/proxy/db/accounting_ledger.py +++ /dev/null @@ -1,352 +0,0 @@ -from __future__ import annotations - -import hashlib -import json -from collections.abc import AsyncGenerator, Callable -from contextlib import AbstractAsyncContextManager, asynccontextmanager -from dataclasses import dataclass -from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Final, Literal, TypeAlias - -from pydantic import AwareDatetime, ConfigDict, JsonValue, TypeAdapter -from typing_extensions import assert_never - -from litellm.accounting import CostEvent, DeliveryTarget, EventKind, InboundUsageEvent, OutboundUsageEvent -from litellm.proxy.db.db_span import db_span -from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper - -if TYPE_CHECKING: - from prisma import Prisma - - from litellm.proxy.utils import PrismaClient - -AccountingEvent: TypeAlias = InboundUsageEvent | OutboundUsageEvent | CostEvent -_JSON: Final = TypeAdapter(dict[str, JsonValue]) -_INSTANT: Final[TypeAdapter[datetime]] = TypeAdapter(AwareDatetime) -_STRINGS: Final = TypeAdapter(tuple[dict[str, str], ...]) -_TIMEOUT: Final = timedelta(seconds=10) - - -class AccountingConflictError(ValueError): - pass - - -@dataclass(frozen=True, slots=True) -class UsageWrite: - event: InboundUsageEvent | OutboundUsageEvent - envelope: str - targets: tuple[DeliveryTarget, ...] = () - - -_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"] - - -@dataclass(frozen=True, slots=True) -class _AcceptedEvent: - event_id: str - kind: EventKind - source_event_id: str | None - occurred_at: datetime - payload: str - payload_hash: str - targets: tuple[str, ...] - schema_version: int - version: int - - -def _json(value: object) -> str: - return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) - - -def _prepare( - event: AccountingEvent, - targets: tuple[DeliveryTarget, ...], - envelope: str, - schema_version: int, - version: int, -) -> _AcceptedEvent: - if type(schema_version) is not int or type(version) is not int or min(schema_version, version) < 1: - raise ValueError("Accounting versions must be positive integers") - validated: Final = type(event).model_validate_json(event.model_dump_json()) - payload: Final = _JSON.validate_json(envelope) - if payload.get("event") != validated.model_dump(mode="json"): - raise ValueError("Accounting envelope must contain the matching public event") - destinations: Final = tuple(sorted(frozenset(DeliveryTarget(target).value for target in targets))) - serialized: Final = _json(payload) - identity: Final = _json((schema_version, version, destinations, payload)) - kind: Final = ( - EventKind.COST - if isinstance(validated, CostEvent) - else EventKind.INBOUND_USAGE - if isinstance(validated, InboundUsageEvent) - else EventKind.OUTBOUND_USAGE - ) - return _AcceptedEvent( - event_id=validated.event_id, - kind=kind, - source_event_id=validated.usage_event_id if isinstance(validated, CostEvent) else None, - occurred_at=validated.occurred_at if isinstance(validated, CostEvent) else validated.ended_at, - payload=serialized, - payload_hash=hashlib.sha256(identity.encode()).hexdigest(), - targets=destinations, - schema_version=schema_version, - version=version, - ) - - -_INSERT_EVENT: Final = """ -INSERT INTO "LiteLLM_AccountingEvent" - (event_id, kind, source_event_id, schema_version, payload_hash, occurred_at, payload) -VALUES ($1, $2::"AccountingEventKind", $3, $4::int, $5, $6::timestamp, $7::jsonb) -ON CONFLICT (event_id) DO NOTHING RETURNING event_id -""" -_INSERT_DELIVERIES: Final = """ -INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) -SELECT $1, value::"AccountingDeliveryTarget", $3::int FROM jsonb_array_elements_text($2::jsonb) -""" -_LOCK_SOURCE: Final = """ -SELECT event_id FROM "LiteLLM_AccountingEvent" -WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE -""" -_EXISTING_COST: Final = """ -SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_event_id = $1 -""" -_LOCK_EVENT: Final = """ -SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE -""" -_INSERT_USAGE_BATCH: Final = """ -INSERT INTO "LiteLLM_AccountingEvent" - (event_id, kind, schema_version, payload_hash, occurred_at, payload) -SELECT event_id, kind::"AccountingEventKind", 1, payload_hash, occurred_at::timestamp, payload -FROM jsonb_to_recordset($1::jsonb) AS records( - event_id text, kind text, payload_hash text, occurred_at text, payload jsonb) -ORDER BY event_id -ON CONFLICT (event_id) DO NOTHING RETURNING event_id -""" -_LOCK_USAGE_BATCH: Final = """ -SELECT event_id, payload_hash FROM "LiteLLM_AccountingEvent" -WHERE event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb)) -ORDER BY event_id FOR UPDATE -""" -_LOAD_USAGE_BATCH: Final = """ -SELECT event.payload, ARRAY( - SELECT target::text FROM "LiteLLM_AccountingDelivery" delivery - WHERE delivery.event_id = event.event_id ORDER BY target::text -) AS targets -FROM "LiteLLM_AccountingEvent" event -WHERE event.event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb)) - AND event.kind = 'outbound_usage' AND event.schema_version = 1 -""" -_INSERT_USAGE_DELIVERIES: Final = """ -INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) -SELECT records.event_id, destination.value::"AccountingDeliveryTarget", 1 -FROM jsonb_to_recordset($1::jsonb) AS records(event_id text, targets jsonb) -CROSS JOIN LATERAL jsonb_array_elements_text(records.targets) AS destination -""" -_COMPLETED_FAMILY: Final = """ -root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage') -AND root.recorded_at < $1::timestamp -AND (root.kind = 'inbound_usage' OR EXISTS ( - SELECT 1 FROM "LiteLLM_AccountingEvent" AS cost WHERE cost.source_event_id = root.event_id AND cost.kind = 'cost' -)) -AND NOT EXISTS ( - SELECT 1 FROM "LiteLLM_AccountingEvent" AS member - WHERE (member.event_id = root.event_id OR (member.kind = 'cost' AND member.source_event_id = root.event_id)) - AND (member.recorded_at >= $1::timestamp - OR EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" AS delivery - WHERE delivery.event_id = member.event_id AND delivery.state != 'applied')) -) -""" -_RETENTION_CANDIDATES: Final = ( - 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE ' - + _COMPLETED_FAMILY - + " ORDER BY root.recorded_at, root.event_id LIMIT $2::int FOR UPDATE OF root SKIP LOCKED" -) -_ELIGIBLE_GROUP: Final = ( - 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE root.event_id = $2 AND ' + _COMPLETED_FAMILY -) -_DELETE_DELIVERIES: Final = """ -DELETE FROM "LiteLLM_AccountingDelivery" WHERE event_id IN ( - SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) -) -""" -_DELETE_EVENTS: Final = """ -DELETE FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) -""" - - -_Conflict: TypeAlias = Literal["missing_source", "competing_cost", "identity_conflict"] - - -async def _accept(db: Prisma, record: _AcceptedEvent) -> Literal["inserted", "duplicate"] | _Conflict: - if record.source_event_id is not None: - source: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_SOURCE, record.source_event_id)) - if not source: - return "missing_source" - costs: Final = _STRINGS.validate_python(tuple(await db.query_raw(_EXISTING_COST, record.source_event_id))) - if any(cost["event_id"] != record.event_id for cost in costs): - return "competing_cost" - inserted: Final = _STRINGS.validate_python( - await db.query_raw( - _INSERT_EVENT, - record.event_id, - record.kind.value, - record.source_event_id, - record.schema_version, - record.payload_hash, - record.occurred_at.astimezone(timezone.utc).isoformat(), - record.payload, - ) - ) - if not inserted: - existing: Final = _STRINGS.validate_python(tuple(await db.query_raw(_LOCK_EVENT, record.event_id))) - if not existing or existing[0]["payload_hash"] != record.payload_hash: - return "identity_conflict" - return "duplicate" - await db.execute_raw(_INSERT_DELIVERIES, record.event_id, _json(record.targets), record.version) - return "inserted" - - -async def _delete_completed(db: Prisma, event_id: str, cutoff: str) -> int: - if not _STRINGS.validate_python(await db.query_raw(_ELIGIBLE_GROUP, cutoff, event_id)): - return 0 - await db.execute_raw(_DELETE_DELIVERIES, event_id) - return await db.execute_raw(_DELETE_EVENTS, event_id) - - -class AccountingLedger: - def __init__(self, transaction: Callable[[], _Transaction]) -> None: - self._transaction: Final = transaction - - @classmethod - def for_client(cls, client: PrismaClient) -> AccountingLedger: - from prisma.client import TransactionManager - - managers: Final = TypeAdapter(TransactionManager, config=ConfigDict(arbitrary_types_allowed=True)) - - @asynccontextmanager - async def transaction() -> AsyncGenerator[Prisma]: - primary: Final = writer_wrapper(client.db) - manager: Final = managers.validate_python(primary.tx(timeout=_TIMEOUT)) - async with db_span("accounting_ledger", "LiteLLM_AccountingEvent"), manager as db: - await db.execute_raw("SET LOCAL statement_timeout = 5000") - await db.execute_raw("SET LOCAL lock_timeout = 1000") - yield db - - return cls(transaction) - - async def insert_usage_batch(self, writes: tuple[UsageWrite, ...]) -> None: - records: Final = tuple(_prepare(write.event, write.targets, write.envelope, 1, 1) for write in writes) - if not records: - return - unique: Final = {record.event_id: record for record in records} - if any(record.payload_hash != unique[record.event_id].payload_hash for record in records): - raise AccountingConflictError("Batch contains conflicting copies of a usage identity") - batch: Final = _json( - tuple( - { - "event_id": record.event_id, - "kind": record.kind.value, - "payload_hash": record.payload_hash, - "occurred_at": record.occurred_at.astimezone(timezone.utc).isoformat(), - "payload": _JSON.validate_json(record.payload), - } - for record in unique.values() - ) - ) - async with self._transaction() as db: - inserted: Final = _STRINGS.validate_python(await db.query_raw(_INSERT_USAGE_BATCH, batch)) - accepted: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_USAGE_BATCH, _json(tuple(unique)))) - if len(accepted) != len(unique) or any( - row["payload_hash"] != unique[row["event_id"]].payload_hash for row in accepted - ): - raise AccountingConflictError("Usage identity conflicts with accepted content or delivery set") - deliveries: Final = _json( - tuple({"event_id": row["event_id"], "targets": unique[row["event_id"]].targets} for row in inserted) - ) - await db.execute_raw(_INSERT_USAGE_DELIVERIES, deliveries) - - async def load_usage_batch(self, event_ids: tuple[str, ...]) -> tuple[UsageWrite, ...]: - if not event_ids: - return () - async with self._transaction() as db: - rows: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python( - await db.query_raw(_LOAD_USAGE_BATCH, _json(event_ids)) - ) - return tuple( - UsageWrite( - OutboundUsageEvent.model_validate_json(_json(_JSON.validate_python(row["payload"])["event"])), - _json(row["payload"]), - TypeAdapter(tuple[DeliveryTarget, ...]).validate_python(row["targets"]), - ) - for row in rows - ) - - async def insert_event( - self, - event: AccountingEvent, - targets: tuple[DeliveryTarget, ...], - *, - envelope: str, - schema_version: int = 1, - version: int = 1, - ) -> Literal["inserted", "duplicate"]: - """Accept T5/T6's serialized envelope, containing its public event under `event`. - - The immutable acceptance hash includes the target set and delivery version. - Changing either for an accepted event conflicts. T7 must provide a separate - add-target or new-version workflow that preserves receipts and prevents - duplicate contributions; redelivery cannot silently create or reset work. - - The Prisma adapter opens one interactive transaction per event, including - two SET LOCAL statements and a source-row lock for costs. I/O is async, - but awaiting this method inline adds commit latency per event, potentially - several times per request. T5 owns bounded asynchronous capture and an - explicit batching strategy; looping over this method does not batch commits. - Success acknowledges commit, independently of HTTP response acknowledgement. - """ - record: Final = _prepare(event, targets, envelope, schema_version, version) - async with self._transaction() as db: - result: Final = await _accept(db, record) - match result: - case "inserted" | "duplicate": - return result - case "missing_source": - raise AccountingConflictError("Cost source must be a durably accepted outbound usage event") - case "competing_cost": - raise AccountingConflictError("Usage event already has an authoritative cost") - case "identity_conflict": - raise AccountingConflictError("Event identity conflicts with its accepted content or delivery set") - case _: - assert_never(result) - - async def purge_completed( - self, - *, - redelivery_before: datetime | None = None, - recovery_before: datetime | None = None, - limit: int = 100, - ) -> int: - """Delete at most `limit` completed families once both retention owners permit it. - - T8 must supply a coherent recovery watermark including outstanding holds. - Omitting either watermark disables deletion. Each usage row and its cost - remain together so a retained charge never loses its source or vice versa. - Applied receipts are deleted with the family, ending its retained dedup - protection. The caller must guarantee events recorded before - redelivery_before cannot be redelivered through supported paths. - T8's recovery_before must also protect checkpoints, - outstanding holds and unfinished settlements; age alone is insufficient. - """ - if type(limit) is not int or not 1 <= limit <= 1000: - raise ValueError("Accounting retention batch size must be between 1 and 1000") - if redelivery_before is None or recovery_before is None: - return 0 - cutoff: Final = min( - _INSTANT.validate_python(redelivery_before).astimezone(timezone.utc), - _INSTANT.validate_python(recovery_before).astimezone(timezone.utc), - ).isoformat() - async with self._transaction() as db: - roots: Final = _STRINGS.validate_python(tuple(await db.query_raw(_RETENTION_CANDIDATES, cutoff, limit))) - deleted: Final = tuple([await _delete_completed(db, root["event_id"], cutoff) for root in roots]) - return sum(deleted) diff --git a/litellm/proxy/db/model_usage_rollup.py b/litellm/proxy/db/model_usage_rollup.py index 4566b10b260..e6fbc299504 100644 --- a/litellm/proxy/db/model_usage_rollup.py +++ b/litellm/proxy/db/model_usage_rollup.py @@ -88,14 +88,14 @@ def _internal_call_origin(metadata: str) -> object: return decoded.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) -def _date_from_start_time(start_time: datetime | str) -> str | None: +def date_from_start_time(start_time: datetime | str) -> str | None: if isinstance(start_time, datetime): return start_time.date().isoformat() return start_time[:10] if len(start_time) >= 10 else None def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransaction | None: - date: Final = _date_from_start_time(payload["startTime"]) + date: Final = date_from_start_time(payload["startTime"]) origin: Final = _internal_call_origin(payload["metadata"]) if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN): return None diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 031bdfa19cd..10b6ec2f661 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -24,7 +24,6 @@ from litellm.litellm_core_utils.core_helpers import ( ) 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 @@ -67,7 +66,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # sanitize_error_information_for_spend_logs, # pyright: ignore[reportUnusedImport] # backwards-compatible package export 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, @@ -375,26 +373,13 @@ def _attempt_spend_row_kwargs( "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 @@ -480,9 +465,6 @@ 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 @@ -519,10 +501,8 @@ class _ProxyDBLogger(CustomLogger): _error_information.pop("traceback", None) _metadata["error_information"] = _error_information - _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 = ( # rebind-ok: enrich identity before building failure spend rows + await ProxyDBLogger._enrich_failure_metadata_unless_db_stalled( metadata=_metadata, original_exception=original_exception ) ) @@ -629,106 +609,6 @@ 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, @@ -914,7 +794,6 @@ 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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c00783854a7..87bc78730b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -877,8 +877,6 @@ 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 @@ -1392,7 +1390,6 @@ 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() @@ -3309,39 +3306,9 @@ 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 diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 696f4e1ec48..eb8c28fcbd7 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -3,7 +3,7 @@ import json import os import re import secrets -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt from functools import reduce @@ -58,7 +58,6 @@ from litellm.litellm_core_utils.internal_call_metadata import ( ) 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, @@ -592,7 +591,6 @@ 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 = {} @@ -982,14 +980,6 @@ def get_logging_payload( payload.get("spend"), ) - if capture_usage is not None: - observed_usage: Final = ( - StandardLoggingPayloadSetup.get_usage_as_dict({"usage": usage}) - if usage and kwargs.get("_accounting_usage_observed") is not UsageObservation.UNKNOWN - else None - ) - capture_usage(payload, safe_dumps(observed_usage) if observed_usage is not None else None) - # Explicitly clear large intermediate objects to reduce memory pressure del response_obj_dict, usage, clean_metadata, additional_usage_values diff --git a/litellm/proxy/spend_tracking/usage_capture.py b/litellm/proxy/spend_tracking/usage_capture.py deleted file mode 100644 index 19480a66678..00000000000 --- a/litellm/proxy/spend_tracking/usage_capture.py +++ /dev/null @@ -1,1084 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import re -from collections.abc import Awaitable, Callable, Iterator, Mapping -from dataclasses import dataclass, replace -from datetime import datetime, timezone -from decimal import Decimal -from hashlib import sha256 -from itertools import chain, groupby, islice, repeat -from typing import TYPE_CHECKING, Final - -from pydantic import BaseModel, JsonValue, TypeAdapter - -from litellm._logging import verbose_proxy_logger -from litellm.accounting import ( - Dimension, - DimensionName, - InboundUsageEvent, - Metric, - MetricName, - OutboundUsageEvent, - Outcome, -) -from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY -from litellm.proxy.spend_tracking.usage_meter import ( - DurableUsageMeter, - UsageCaptureError, - UsageCaptureSettings, - usage_delivery_targets, -) -from litellm.types.utils import CallTypes, CustomPricingLiteLLMParams - -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy._types import SpendLogsPayload - - -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite - - -class UsageAcknowledgementError(UsageCaptureError): - def __init__(self, cause: Exception) -> None: - super().__init__("Usage persistence did not acknowledge commit") - self.__cause__ = cause - - -@dataclass(frozen=True, slots=True) -class CapturedUsage: - payload: SpendLogsPayload - writes: tuple[UsageWrite, ...] - receipt: asyncio.Future[None] - - -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) - - -_FIELDS: Final = TypeAdapter(dict[str, object]) - - -_INSTANT: Final = TypeAdapter(datetime) - - -_DIMENSIONS: Final = ( - ("api_key", DimensionName.KEY), - ("user", DimensionName.USER), - ("team_id", DimensionName.TEAM), - ("organization_id", DimensionName.ORGANIZATION), - ("end_user", DimensionName.END_USER), - ("model", DimensionName.MODEL), - ("model_group", DimensionName.MODEL_GROUP), - ("model_id", DimensionName.DEPLOYMENT), - ("custom_llm_provider", DimensionName.PROVIDER), - ("mcp_namespaced_tool_name", DimensionName.TOOL), - ("agent_id", DimensionName.AGENT), -) - - -_USAGE_METRICS: Final = tuple( - metric - for metric in MetricName - if metric not in (MetricName.SPEND, MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS) -) - - -def _fields(value: object) -> Mapping[str, object]: - return _FIELDS.validate_python(value) if isinstance(value, Mapping) else {} - - -def _text(value: object) -> str | None: - return value if isinstance(value, str) and value else None - - -def _utc_time(value: datetime | str) -> datetime: - parsed: Final = _INSTANT.validate_python(value) - return parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed.astimezone(timezone.utc) - - -def _json(value: object) -> str: - return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) - - -def _usage_metrics(usage: Mapping[str, JsonValue] | None) -> tuple[Metric, ...] | None: - if usage is None: - return None - return tuple( - Metric(name=metric, value=Decimal(str(value))) - for metric in _USAGE_METRICS - if isinstance(value := usage.get(metric.value), (int, float)) and not isinstance(value, bool) - ) - - -@dataclass(frozen=True, slots=True) -class UsageCaptureContext: - caller_id: str | None - attempt_id: str - parent_request_id: str | None - retry_of_attempt_id: str | None - job_id: str | None - charge_eligible: bool - request_count_eligible: bool - ownership: str - pricing: str - finalized_cost: str | None = None - allocation: str | None = None - outcome: Outcome | None = None - started_at: datetime | None = None - ended_at: datetime | None = None - phase: str = "execution" - extra_dimensions: tuple[Dimension, ...] = () - provider_timestamps: str = "{}" - usage_source: str = "measured" - included_attempt_ids: tuple[str, ...] = () - related_attempt_ids: tuple[str, ...] = () - - -def _database_usage_write(event: InboundUsageEvent | OutboundUsageEvent, envelope: Mapping[str, object]) -> UsageWrite: - from sys import maxsize - - from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure - - normalized: Final = _FIELDS.validate_python(safe_json_structure(envelope, max_depth=maxsize)) - if normalized.get("event") != event.model_dump(mode="json"): - raise UsageCaptureError("Usage event text cannot be normalized without changing identity or attribution") - serialized: Final = _json(normalized) - _JSON_OBJECT.validate_json(serialized) - return UsageWrite(event, serialized, usage_delivery_targets(event.dimensions)) - - -def usage_writes( - payload: SpendLogsPayload, usage_json: str | None, context: UsageCaptureContext -) -> tuple[UsageWrite, ...]: - fields: Final = _FIELDS.validate_python(payload) - tags: Final = TypeAdapter(tuple[str, ...]).validate_json(payload["request_tags"]) - dimensions: Final = ( - *(Dimension(name=name, value=value) for field, name in _DIMENSIONS if (value := _text(fields.get(field)))), - *(Dimension(name=DimensionName.TAG, value=tag) for tag in dict.fromkeys(tags)), - *context.extra_dimensions, - ) - outcome: Final = context.outcome or (Outcome.FAILED if payload["status"] == "failure" else Outcome.SUCCEEDED) - started: Final = context.started_at or _utc_time(payload["startTime"]) - ended: Final = context.ended_at or _utc_time(payload["endTime"]) - usage: Final = None if usage_json is None else _JSON_OBJECT.validate_json(usage_json) - outbound: Final = OutboundUsageEvent( - event_id=context.attempt_id, - attempt_id=context.attempt_id, - request_id=context.caller_id, - retry_of_attempt_id=context.retry_of_attempt_id, - call_type=CallTypes(payload["call_type"]), - started_at=started, - ended_at=ended, - outcome=outcome, - dimensions=dimensions, - metrics=_usage_metrics(None if context.usage_source == "unknown" else usage), - ) - envelope: Final = { - "event": outbound.model_dump(mode="json"), - "usage": usage, - "usage_source": "unknown" if usage is None else context.usage_source, - "included_attempt_ids": context.included_attempt_ids, - "related_attempt_ids": context.related_attempt_ids, - "pricing": _JSON_OBJECT.validate_json(context.pricing), - "charge_eligible": context.charge_eligible, - "request_count_eligible": context.request_count_eligible, - "ownership": context.ownership, - "parent_request_id": context.parent_request_id, - "job_id": context.job_id, - "phase": context.phase, - "provider_timestamps": _JSON_OBJECT.validate_json(context.provider_timestamps), - "finalized_cost": None - if context.finalized_cost is None - else _JSON_OBJECT.validate_json(context.finalized_cost), - "allocation": None if context.allocation is None else _JSON_OBJECT.validate_json(context.allocation), - } - outbound_write: Final = _database_usage_write(outbound, envelope) - if not context.request_count_eligible or context.caller_id is None: - return (outbound_write,) - inbound: Final = InboundUsageEvent( - event_id=f"{context.caller_id}:inbound", - request_id=context.caller_id, - call_type=outbound.call_type, - started_at=started, - ended_at=ended, - outcome=outcome, - dimensions=dimensions, - metrics=( - Metric(name=MetricName.REQUESTS, value=Decimal(1)), - Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(int(outcome is Outcome.SUCCEEDED))), - Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(int(outcome is not Outcome.SUCCEEDED))), - ), - ) - return ( - _database_usage_write(inbound, {"event": inbound.model_dump(mode="json")}), - outbound_write, - ) - - -_INCLUDED_IN_PARENT: Final = object() - - -_PRICING_FIELDS: Final = frozenset( - { - "service_tier", - "data_residency", - "vertex_location", - "region_name", - "base_model", - "custom_pricing", - "input_cost_per_token", - "output_cost_per_token", - "input_cost_per_second", - "output_cost_per_second", - "input_cost_per_character", - "output_cost_per_character", - "size", - "quality", - "n", - "duration", - "seconds", - "search_context_size", - "additional_response_cost", - } -) - - -def _selected_pricing(fields: Mapping[str, object]) -> Mapping[str, JsonValue]: - return _JSON_OBJECT.validate_python( - { - key: value - for key, value in fields.items() - if key in _PRICING_FIELDS or key in CustomPricingLiteLLMParams.model_fields - } - ) - - -def _response_fields(response: object) -> Mapping[str, object]: - if isinstance(response, BaseModel): - return _FIELDS.validate_python(response.model_dump()) - return _fields(response) - - -def _job_time(response: Mapping[str, object], key: str) -> datetime | None: - stamp: Final = response.get(key) - if isinstance(stamp, (int, float)) and not isinstance(stamp, bool): - return datetime.fromtimestamp(stamp, tz=timezone.utc) - return None - - -def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], response: object) -> UsageCaptureContext: - from dataclasses import asdict - - from litellm.litellm_core_utils.hidden_params import get_hidden_params - from litellm.litellm_core_utils.internal_call_metadata import is_background_response, is_unbilled_non_inference_call - from litellm.litellm_core_utils.litellm_logging import ( - Logging, - UsageCaptureObservation, - UsageObservation, - mask_api_base_credentials, - ) - from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost - from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation - from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups - - metadata: Final = _JSON_OBJECT.validate_json(payload["metadata"]) - params: Final = _fields(kwargs.get("litellm_params")) - request_metadata: Final = _fields(params.get("litellm_metadata") or params.get("metadata")) - standard: Final = _fields(kwargs.get("standard_logging_object")) - logging_obj: Final = kwargs.get("litellm_logging_obj") - raw_response: Final = _response_fields(response) - hidden: Final = _fields(get_hidden_params(response)) - tool: Final = _fields(kwargs.get("mcp_tool_call_metadata")) - tool_pricing: Final = _fields(tool.get("mcp_server_cost_info")) - caller: Final = payload["litellm_call_id"] - cached: Final = payload["cache_hit"].lower() == "true" - included_in_parent: Final = kwargs.get("_accounting_internal_aggregate") is _INCLUDED_IN_PARENT - billed_id: Final = payload["request_id"] - terminal: Final = raw_response.get("status") in { - "completed", - "complete", - "failed", - "cancelled", - "expired", - "incomplete", - } - deferred: Final = billed_id.endswith(("_batch_cost", "_background_cost")) and terminal - management: Final = ( - (billed_id.endswith(("_batch_cost", "_background_cost")) and not terminal) - or (is_background_response(response) and not terminal) - or is_unbilled_non_inference_call(payload["call_type"], request_metadata, response) - ) - attempt: Final = ( - f"{caller}_cache_hit" - if cached and caller - else f"{caller}:status" - if management - else f"{billed_id}:included" - if included_in_parent - else billed_id - ) - if not attempt or attempt == "None" or (cached and not caller): - raise UsageCaptureError("Usage capture requires a stable execution identity") - retry: Final = metadata.get(BILLED_RETRY_ATTEMPT_METADATA_KEY) is True - origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) - parent: Final = _text(kwargs.get("parent_request_id")) or _text(request_metadata.get("parent_request_id")) - finalized_value: Final = ( - logging_obj.finalized_cost if isinstance(logging_obj, Logging) else kwargs.get("_accounting_finalized_cost") - ) - finalized: Final = finalized_value if isinstance(finalized_value, FinalizedCost) else None - allocation_value: Final = ( - logging_obj.cost_allowance_allocation - if isinstance(logging_obj, Logging) - else kwargs.get("_accounting_allocation") - ) - allocation: Final = allocation_value if isinstance(allocation_value, SearchAllowanceAllocation) else None - pricing_snapshot: Final = kwargs.get("_accounting_pricing_snapshot") - observation: Final = kwargs.get("_accounting_usage_snapshot") - captured_pricing: Final = { - "model_map_information": ( - _JSON_OBJECT.validate_json(pricing_snapshot.model_map_information) - if isinstance(pricing_snapshot, PricingCapture) - else metadata.get("model_map_information") - ), - "cost_breakdown": metadata.get("cost_breakdown"), - **( - _JSON_OBJECT.validate_json(pricing_snapshot.inputs) - if isinstance(pricing_snapshot, PricingCapture) - else _pricing_inputs(kwargs, logging_obj, response) - ), - "tool": { - "name": tool.get("name"), - "cost_per_query": _fields(tool_pricing.get("tool_name_to_cost_per_query")).get( - _text(tool.get("name")) or "", tool_pricing.get("default_cost_per_query") - ), - } - if tool - else None, - "response_inputs": { - "audio_transcription_duration": hidden.get("audio_transcription_duration", raw_response.get("duration")), - "billed_units": _fields(raw_response.get("meta")).get("billed_units"), - "batch_models": hidden.get("batch_models"), - "batch_successful_requests": hidden.get("batch_successful_requests"), - "batch_failed_requests": hidden.get("batch_failed_requests"), - "service_tier": raw_response.get("service_tier"), - "region_name": hidden.get("region_name"), - }, - "reported_cost": observation.response_cost - if isinstance(observation, UsageCaptureObservation) - else standard.get( - "response_cost", kwargs.get("response_cost", _fields(get_hidden_params(response)).get("response_cost")) - ), - "adjustments": ( - _JSON_OBJECT.validate_json(pricing_snapshot.adjustments) - if isinstance(pricing_snapshot, PricingCapture) - else None - ), - } - return UsageCaptureContext( - caller_id=None if deferred else parent if retry or included_in_parent else caller, - attempt_id=attempt, - parent_request_id=parent, - retry_of_attempt_id=_text(request_metadata.get("retry_of_attempt_id")), - job_id=_text(raw_response.get("id")) if deferred or management else None, - charge_eligible=not cached and not included_in_parent and not management, - request_count_eligible=bool(caller) and not origin and not retry and not deferred and not included_in_parent, - ownership="included_in_parent" - if included_in_parent - else "aggregate" - if kwargs.get("additional_response_cost") is not None - else "attempt", - pricing=_json(captured_pricing), - finalized_cost=_json( - { - "response_cost": finalized.amount, - "cost_breakdown": finalized.logging_breakdown(), - "allowance_discount": finalized.allowance_discount, - } - ) - if finalized is not None - else ( - _json({"response_cost": standard.get("response_cost"), "cost_breakdown": standard.get("cost_breakdown")}) - if standard.get("cost_finalized") is True - else None - ), - allocation=_json(asdict(allocation)) if allocation is not None else None, - started_at=_job_time(raw_response, "created_at") if deferred else None, - ended_at=( - _job_time(raw_response, "completed_at") - or _job_time(raw_response, "failed_at") - or _job_time(raw_response, "cancelled_at") - or _job_time(raw_response, "expired_at") - ) - if deferred - else None, - outcome=( - Outcome.CANCELLED - if raw_response.get("status") == "cancelled" - or isinstance(response, asyncio.CancelledError) - or getattr(response, "status_code", None) == 499 - else Outcome.FAILED - if raw_response.get("status") in {"failed", "expired", "incomplete"} or isinstance(response, BaseException) - else None - ), - phase="terminal" if deferred else "management" if management else "execution", - usage_source="estimated" - if kwargs.get("_accounting_usage_observed") is UsageObservation.ESTIMATED - else "measured", - included_attempt_ids=TypeAdapter(tuple[str, ...]).validate_python( - kwargs.get("_accounting_included_attempt_ids", ()) - ), - related_attempt_ids=tuple(write.event.event_id for write in _prior_writes(kwargs)), - provider_timestamps=_json( - { - key: raw_response.get(key) - for key in ("created_at", "completed_at", "failed_at", "cancelled_at", "expired_at") - if raw_response.get(key) is not None - } - ) - if deferred - else "{}", - extra_dimensions=( - *( - (Dimension(name=DimensionName.TOOL, value=tool_name),) - if not payload.get("mcp_namespaced_tool_name") - and (tool_name := _text(tool.get("namespaced_tool_name"))) - else () - ), - *( - (Dimension(name=DimensionName.API_BASE, value=mask_api_base_credentials(api_base)),) - if (api_base := _text(payload["api_base"])) - else () - ), - *( - Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value=group) - for group in get_request_model_access_groups(kwargs) - ), - *( - (Dimension(name=DimensionName.PROJECT, value=project),) - if (project := _text(request_metadata.get("user_api_key_project_id"))) - else () - ), - *( - (Dimension(name=DimensionName.SERVICE_TIER, value=tier),) - if (tier := _text(raw_response.get("service_tier") or kwargs.get("service_tier"))) - else () - ), - ), - ) - - -class _CaptureBuilder: - def __init__(self, kwargs: Mapping[str, object], response: object) -> None: - self.kwargs: Final = kwargs - self.response: Final = response - self.writes: tuple[UsageWrite, ...] = () - - def observe(self, payload: SpendLogsPayload, usage: str | None) -> None: - from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation, UsageObservation - - snapshot: Final = self.kwargs.get("_accounting_usage_snapshot") - if isinstance(snapshot, UsageCaptureError): - raise snapshot - effective: Final = snapshot.usage_json if isinstance(snapshot, UsageCaptureObservation) else usage - parsed: Final = _JSON_OBJECT.validate_json(effective) if effective is not None else {} - batch: Final = _fields(parsed.get("batch_usage_by_model")) - unknown: Final = (bool(batch) or bool(parsed.get("batch_unknown_usage_requests"))) and all( - value is None for value in batch.values() - ) - observed: Final = ( - _json( - { - "batch_usage_by_model": batch, - "batch_unknown_usage_requests": parsed.get("batch_unknown_usage_requests"), - **( - {"batch_usage_by_request": parsed["batch_usage_by_request"]} - if "batch_usage_by_request" in parsed - else {} - ), - } - ) - if unknown - else effective - ) - context: Final = capture_context(payload, self.kwargs, self.response) - provenance: Final = ( - "unknown" - if unknown - else "estimated" - if isinstance(snapshot, UsageCaptureObservation) and snapshot.provenance is UsageObservation.ESTIMATED - else context.usage_source - ) - self.writes = usage_writes(payload, observed, replace(context, usage_source=provenance)) - - -def _prior_captures(kwargs: Mapping[str, object]) -> tuple[CapturedUsage, ...]: - prior: Final = kwargs.get("_accounting_prior_captures", ()) - return ( - tuple( - item for item in TypeAdapter(tuple[object, ...]).validate_python(prior) if isinstance(item, CapturedUsage) - ) - if isinstance(prior, tuple) - else () - ) - - -def _prior_writes(kwargs: Mapping[str, object]) -> tuple[UsageWrite, ...]: - return tuple(chain.from_iterable(item.writes for item in _prior_captures(kwargs))) - - -def capture_callback( - meter: DurableUsageMeter, - kwargs: Mapping[str, object], - response: object, - start_time: datetime, - end_time: datetime, -) -> CapturedUsage: - from litellm.proxy.db.db_spend_update_writer import get_llm_router - from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload - - existing: Final = kwargs.get("_accounting_capture") - if isinstance(existing, CapturedUsage): - return existing - builder: Final = _CaptureBuilder(kwargs, response) - payload: Final = get_logging_payload( - dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe - ) - prerequisites: Final = tuple(item.receipt for item in _prior_captures(kwargs)) - return CapturedUsage(payload, builder.writes, meter.submit(builder.writes, prerequisites=prerequisites)) - - -def _batch_rows(response: object) -> tuple[Mapping[str, JsonValue], ...]: - usage: Final = _response_fields(_response_fields(response).get("usage")) - rows: Final = usage.get("batch_usage_by_request") - return TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(rows) if rows else () - - -def _batch_root(write: UsageWrite, rows: tuple[Mapping[str, JsonValue], ...]) -> UsageWrite: - event: Final = write.event.model_copy(update={"metrics": None}) - envelope: Final = _JSON_OBJECT.validate_json(write.envelope) - return _database_usage_write( - event, - { - **envelope, - "event": event.model_dump(mode="json"), - "usage": { - "batch_row_count": len(rows), - "batch_rows_sha256": sha256(_json(rows).encode()).hexdigest(), - "batch_charge_eligible": envelope["charge_eligible"], - }, - "usage_source": "unknown", - "charge_eligible": False, - "ownership": "aggregate", - }, - ) - - -def _batch_row_write(root: UsageWrite, row: Mapping[str, JsonValue], index: int) -> UsageWrite: - envelope: Final = _JSON_OBJECT.validate_json(root.envelope) - pricing: Final = _JSON_OBJECT.validate_python(envelope["pricing"]) - model: Final = _text(row["model"]) - usage: Final = _JSON_OBJECT.validate_python(row["usage"]) if row["usage"] is not None else None - identity: Final = f"{root.event.event_id}:row:{index}" - event: Final = OutboundUsageEvent.model_validate( - { - **root.event.model_dump(), - "event_id": identity, - "attempt_id": identity, - "metrics": _usage_metrics(usage), - "dimensions": ( - *(dimension for dimension in root.event.dimensions if dimension.name is not DimensionName.MODEL), - *((Dimension(name=DimensionName.MODEL, value=model),) if model else ()), - ), - } - ) - return _database_usage_write( - event, - { - **envelope, - "event": event.model_dump(mode="json"), - "usage": usage, - "usage_source": "unknown" if usage is None else "measured", - "charge_eligible": _fields(envelope["usage"])["batch_charge_eligible"], - "ownership": "attempt", - "parent_request_id": root.event.event_id, - "finalized_cost": None, - "allocation": None, - "pricing": { - **pricing, - "model_map_information": _fields(pricing.get("batch_models")).get(model or ""), - "batch_models": {}, - "reported_cost": None, - "cost_breakdown": None, - "response_inputs": { - key: value - for key, value in _fields(pricing.get("response_inputs")).items() - if key not in {"batch_models", "batch_successful_requests"} - }, - }, - }, - ) - - -@dataclass(frozen=True, slots=True) -class _BatchCapture: - payload: SpendLogsPayload - root: UsageWrite - rows: bytes - prerequisites: tuple[asyncio.Future[None], ...] - - -def _prepare_batch_capture( - kwargs: Mapping[str, object], response: object, start_time: datetime, end_time: datetime -) -> _BatchCapture: - from litellm.proxy.db.db_spend_update_writer import get_llm_router - from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload - - rows: Final = tuple(sorted(_batch_rows(response), key=_json)) - builder: Final = _CaptureBuilder(kwargs, response) - payload: Final = get_logging_payload( - dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe - ) - if len(builder.writes) != 1 or _terminal_facts(builder.writes[0]) is None: - raise UsageCaptureError("Batch rows require one terminal job observation") - return _BatchCapture( - payload, - _batch_root(builder.writes[0], rows), - b"\n".join(_json((row, sum(1 for _ in group))).encode("utf-8") for row, group in groupby(map(_json, rows))), - tuple(prior.receipt for prior in _prior_captures(kwargs)), - ) - - -def _batch_groups(rows: bytes) -> Iterator[tuple[str, int]]: - return (TypeAdapter(tuple[str, int]).validate_json(match.group()) for match in re.finditer(rb"[^\n]+", rows)) - - -def _batch_writes(root: UsageWrite, rows: bytes) -> Iterator[UsageWrite]: - expanded: Final = chain.from_iterable(repeat(row, count) for row, count in _batch_groups(rows)) - return (_batch_row_write(root, _JSON_OBJECT.validate_json(row), index) for index, row in enumerate(expanded)) - - -def _batch_write_bound(root: UsageWrite, rows: bytes) -> int: - last_index: Final = sum(count for _, count in _batch_groups(rows)) - 1 - return max( - len(root.envelope.encode("utf-8")), - max( - len(_batch_row_write(root, _JSON_OBJECT.validate_json(row), last_index).envelope.encode("utf-8")) - for row, _ in _batch_groups(rows) - ), - ) - - -def _batch_capacity( - meter: DurableUsageMeter, root: UsageWrite, rows: bytes, *, extra_bytes: int = 0 -) -> tuple[int, int, int]: - retained: Final = len(rows) + len(root.envelope.encode("utf-8")) + extra_bytes - bound: Final = _batch_write_bound(root, rows) - chunk_bytes: Final = max( - bound, - min( - bound * min(meter.settings.batch_size, sum(count for _, count in _batch_groups(rows))), - meter.settings.max_pending_bytes // 4, - (meter.settings.max_pending_bytes - retained) // 2, - ), - ) - return retained + chunk_bytes, chunk_bytes, bound - - -async def _persist_batch_capture( - meter: DurableUsageMeter, - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]], - root: UsageWrite, - rows: bytes, - prerequisites: tuple[asyncio.Future[None], ...], - chunk_bytes: int, - reservation: object | None = None, -) -> None: - for receipt in prerequisites: - await asyncio.shield(receipt) - await meter.record_writes((root,)) - accepted: Final = await load((root.event.event_id,)) - if len(accepted) != 1: - raise UsageCaptureError("Committed batch snapshot could not be loaded") - frozen: Final = _reuse_terminal_snapshot(root, accepted[0]) - retained, capacity, bound = _batch_capacity(meter, frozen, rows, extra_bytes=len(root.envelope.encode("utf-8"))) - if reservation is not None: - meter.resize_producer(reservation, retained, capacity) - write_capacity: Final = capacity if reservation is not None else chunk_bytes - if bound > write_capacity: - raise UsageCaptureError("Batch observation exceeds reserved producer write capacity") - group_size: Final = min(meter.settings.batch_size, max(1, write_capacity // bound)) - remaining: Final = _batch_writes(frozen, rows) - while group := tuple(islice(remaining, group_size)): - await meter.record_writes(group) - - -async def capture_batch_callback( - meter: DurableUsageMeter, - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None, - kwargs: Mapping[str, object], - response: object, - start_time: datetime, - end_time: datetime, -) -> CapturedUsage: - if load is None: - raise UsageCaptureError("Batch capture requires the existing ledger snapshot reader") - capture: Final = _prepare_batch_capture(kwargs, response, start_time, end_time) - await _persist_batch_capture( - meter, load, capture.root, capture.rows, capture.prerequisites, meter.settings.max_pending_bytes - ) - receipt: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() - receipt.set_result(None) - return CapturedUsage(capture.payload, (capture.root,), receipt) - - -async def captured_callback_kwargs( - kwargs: Mapping[str, object], response: object, start_time: datetime | None, end_time: datetime | None -) -> Mapping[str, object]: - from litellm.proxy.proxy_server import usage_capture - - usage_meter: Final = usage_capture.meter - - if usage_meter is None: - return kwargs - if start_time is None or end_time is None: - verbose_proxy_logger.error("Usage capture requires callback start and end timestamps") - return kwargs - try: - captured: Final = ( - usage_capture.submit_batch(kwargs, response, start_time, end_time) - if _batch_rows(response) - else capture_callback(usage_meter, kwargs, response, start_time, end_time) - ) - except Exception: # noqa: BLE001 # capture failure must not suppress legacy billing - verbose_proxy_logger.exception("Usage capture failed before legacy cost tracking") - return kwargs - return {**kwargs, "_accounting_capture": captured} - - -async def capture_proxy_response(logging_obj: object, response: object, *, background: bool = False) -> None: - from fastapi import HTTPException - - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy.proxy_server import general_settings, usage_capture - - settings: Final = _FIELDS.validate_python(general_settings) - usage_meter: Final = usage_capture.meter - - if usage_meter is None or not isinstance(logging_obj, Logging): - return - sync_ack: Final = ( - settings.get("accounting_sync_ack") is True - and logging_obj.stream is not True - and not background - and logging_obj.call_type in {"acompletion", "atext_completion", "aembedding", "aresponses"} - ) - try: - capture_logging_usage(usage_meter, logging_obj, response, logging_obj.start_time, datetime.now(timezone.utc)) - captured: Final = logging_obj.model_call_details.get("_accounting_capture") - if sync_ack: - if not isinstance(captured, CapturedUsage): - raise UsageCaptureError("Usage capture was not admitted") - timeout: Final = TypeAdapter(float).validate_python( - settings.get("accounting_sync_ack_timeout_seconds", 20.0) - ) - await asyncio.wait_for(asyncio.shield(captured.receipt), timeout=timeout) - except Exception as error: # noqa: BLE001 # the configured response policy owns every capture failure - verbose_proxy_logger.exception("Usage capture failed before the response boundary") - if sync_ack and settings.get("accounting_sync_ack_failure_policy", "fail_closed") == "fail_closed": - acknowledgement: Final = UsageAcknowledgementError(error) - raise HTTPException(503, str(acknowledgement)) from acknowledgement - - -def capture_logging_usage( - meter: DurableUsageMeter, - logging_obj: object, - response: object, - start_time: datetime, - end_time: datetime, - freeze_only: bool = False, -) -> CapturedUsage | UsageCaptureError | None: - from litellm._internal_context import internal_parent_request_id, is_internal_call - from litellm.batches.batch_utils import batch_cost_is_final - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.types.utils import LiteLLMBatch - - if not isinstance(logging_obj, Logging): - return - internal: Final = is_internal_call.get() - if isinstance(response, LiteLLMBatch) and batch_cost_is_final(response) and response.usage is None: - return - try: - if not internal and isinstance(logging_obj.model_call_details.get("_accounting_capture"), CapturedUsage): - return capture_callback(meter, logging_obj.model_call_details, response, start_time, end_time) - normalized_result: Final = logging_obj.usage_capture_response(response) - prepared: Final = logging_obj.prepared_cost - normalized: Final = ( - {**_response_fields(normalized_result), "usage": prepared.usage.model_dump()} - if prepared is not None - else normalized_result - ) - if normalized is None: - return - details: Final = _FIELDS.validate_python(logging_obj.model_call_details) - params: Final = _fields(details.get("litellm_params")) - metadata: Final = _fields(params.get("litellm_metadata") or params.get("metadata")) - if metadata.get("batch_ignore_default_logging") is True: - return - existing_pricing: Final = details.get("_accounting_pricing_snapshot") - pricing: Final = ( - existing_pricing - if isinstance(existing_pricing, PricingCapture) and not internal - else _capture_pricing(logging_obj, normalized, details) - ) - if freeze_only or _batch_rows(normalized): - if not internal: - logging_obj.record_usage_pricing(pricing) - return - captured: Final = capture_callback( - meter, - { - **details, - "litellm_logging_obj": logging_obj, - "_accounting_capture": None if internal else details.get("_accounting_capture"), - "_accounting_pricing_snapshot": pricing, - "_accounting_internal_aggregate": _INCLUDED_IN_PARENT if internal else None, - "parent_request_id": internal_parent_request_id.get() if internal else details.get("parent_request_id"), - }, - normalized, - start_time, - end_time, - ) - if not internal: - logging_obj.record_usage_capture(captured) - except Exception as error: # noqa: BLE001 # failed capture must not prevent normal logging and legacy charging - verbose_proxy_logger.exception("Usage capture failed at the execution boundary") - return UsageCaptureError(str(error)) - else: - return captured - - -@dataclass(frozen=True, slots=True) -class PricingCapture: - model_map_information: str - adjustments: str - inputs: str - - -def _capture_pricing(logging_obj: Logging, normalized: object, details: Mapping[str, object]) -> PricingCapture: - import litellm - from litellm.litellm_core_utils.cost_discount import resolve_cost_discount - from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup - - params: Final = _fields(details.get("litellm_params")) - prepared: Final = logging_obj.prepared_cost - provider: Final = _text(details.get("custom_llm_provider")) - builtins: Final = logging_obj.standard_built_in_tools_params - pricing: Final = ( - prepared.payload["model_map_information"] - if prepared is not None - else StandardLoggingPayloadSetup.get_model_cost_information( - base_model=logging_obj.get_deployment_model_for_cost(), - custom_pricing=bool(params.get("custom_pricing")), - custom_llm_provider=_text(details.get("custom_llm_provider")), - init_response_obj=normalized, - api_base=_text(params.get("api_base")), - router_model_id=logging_obj.get_router_model_id(), - ) - ) - adjustments: Final = _json( - { - "discount": resolve_cost_discount( - litellm.cost_discount_config, - provider, - logging_obj.get_deployment_model_for_cost(), - _text(params.get("region_name") or params.get("vertex_location")), - ), - "margin": litellm.cost_margin_config.get(provider or "", litellm.cost_margin_config.get("global")), - "web_search_options": ( - _selected_pricing(_fields(builtins.get("web_search_options"))) - if builtins.get("web_search_options") is not None - else None - ), - "file_search": builtins.get("file_search") is not None, - } - ) - batch_usage: Final = _fields(_fields(_response_fields(normalized).get("usage")).get("batch_usage_by_model")) - return PricingCapture( - _json(pricing), - adjustments, - _json( - { - **_pricing_inputs(details, logging_obj, normalized), - "batch_models": { - model: StandardLoggingPayloadSetup.get_model_cost_information( - base_model=model, - custom_pricing=bool(params.get("custom_pricing")), - custom_llm_provider=provider, - init_response_obj={"model": model}, - api_base=_text(params.get("api_base")), - router_model_id=logging_obj.get_router_model_id(), - ) - for model in batch_usage - }, - } - ), - ) - - -def _pricing_inputs(kwargs: Mapping[str, object], logging_obj: object, response: object) -> Mapping[str, object]: - from litellm.litellm_core_utils.litellm_logging import Logging - - return { - "request": _selected_pricing(kwargs), - "deployment": _selected_pricing(_fields(kwargs.get("litellm_params"))), - "optional_params": _selected_pricing(_fields(kwargs.get("optional_params"))), - "logging": logging_obj.usage_capture_pricing(response) if isinstance(logging_obj, Logging) else None, - "model_info": _selected_pricing( - _fields( - _fields( - _fields(kwargs.get("litellm_params")).get("litellm_metadata") - or _fields(kwargs.get("litellm_params")).get("metadata") - ).get("model_info") - ) - ), - } - - -def _terminal_facts(write: UsageWrite) -> Mapping[str, JsonValue] | None: - payload: Final = _JSON_OBJECT.validate_json(write.envelope) - if payload.get("phase") != "terminal": - return None - event: Final = _JSON_OBJECT.validate_python(payload["event"]) - return { - **{ - key: value - for key, value in payload.items() - if key not in {"event", "pricing", "finalized_cost", "allocation"} - }, - "event": {key: value for key, value in event.items() if key not in {"started_at", "ended_at"}}, - } - - -def _reuse_terminal_snapshot(candidate: UsageWrite, accepted: UsageWrite) -> UsageWrite: - facts: Final = _terminal_facts(candidate) - if facts is None or facts != _terminal_facts(accepted) or set(candidate.targets) != set(accepted.targets): - raise AccountingConflictError("Repeated terminal job disagrees with its accepted usage or attribution") - return accepted - - -class UsageBatchWriter: - def __init__( - self, - insert: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]], - ) -> None: - self._insert: Final = insert - self._load: Final = load - - async def _restore(self, writes: tuple[UsageWrite, ...]) -> tuple[UsageWrite, ...]: - terminal: Final = tuple(write for write in writes if _terminal_facts(write) is not None) - if not terminal: - return writes - identities: Final = tuple(dict.fromkeys(write.event.event_id for write in terminal)) - stored: Final = await self._load(identities) - # Re-read the first accepted observation and pricing before the ledger's unchanged id/hash check. - snapshots: Final = {write.event.event_id: write for write in (*reversed(terminal), *stored)} - return tuple( - _reuse_terminal_snapshot(write, snapshots[write.event.event_id]) - if write.event.event_id in snapshots - else write - for write in writes - ) - - async def write(self, writes: tuple[UsageWrite, ...]) -> None: - restored: Final = await self._restore(writes) - try: - await self._insert(restored) - except AccountingConflictError: - if not any(_terminal_facts(write) is not None for write in writes): - raise - await self._insert(await self._restore(writes)) - - -class UsageCaptureRuntime: - def __init__(self) -> None: - self._batches: Final[set[asyncio.Task[None]]] = set() # mutable-ok: owned producer lifecycle - self._stopping = False - self.meter: DurableUsageMeter | None = None - self.load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None = None - - def start( - self, - write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], - settings: UsageCaptureSettings, - *, - load: Callable[[tuple[str, ...]], Awaitable[tuple[UsageWrite, ...]]] | None = None, - ) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - - if self.meter is not None: - if self.meter.settings != settings: - raise ValueError("Changing accounting_capture requires a worker restart") - return - meter: Final = DurableUsageMeter(write_batch, settings) - self.meter = meter - self.load = load - self._stopping = False - Logging.usage_capture_callback = lambda logger, response, start, end, freeze: capture_logging_usage( - meter, logger, response, start, end, freeze - ) - - def submit_batch( - self, kwargs: Mapping[str, object], response: object, start_time: datetime, end_time: datetime - ) -> CapturedUsage: - meter: Final = self.meter - load: Final = self.load - if meter is None or load is None or self._stopping: - raise UsageCaptureError("Batch capture requires an active ledger writer and snapshot reader") - if len(self._batches) >= meter.settings.max_pending_records: - raise UsageCaptureError("Batch capture producer capacity exhausted") - capture: Final = _prepare_batch_capture(kwargs, response, start_time, end_time) - retained, chunk_bytes, _ = _batch_capacity( - meter, capture.root, capture.rows, extra_bytes=len(capture.root.envelope.encode("utf-8")) - ) - reservation: Final = meter.reserve_producer(retained, chunk_bytes) - task: Final = asyncio.create_task( - _persist_batch_capture( - meter, load, capture.root, capture.rows, capture.prerequisites, chunk_bytes, reservation - ), - name="accounting batch capture", - ) - task.add_done_callback(lambda _: meter.release_producer(reservation)) - self._batches.add(task) - task.add_done_callback(self._batch_finished) - return CapturedUsage(capture.payload, (capture.root,), task) - - def _batch_finished(self, task: asyncio.Task[None]) -> None: - self._batches.discard(task) - if task.cancelled(): - verbose_proxy_logger.error("Batch usage capture stopped before complete persistence") - elif (error := task.exception()) is not None: - verbose_proxy_logger.error("Batch usage capture failed", exc_info=error) - - async def stop(self, timeout: float) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - - if self.meter is None: - return - self._stopping = True - deadline: Final = asyncio.get_running_loop().time() + timeout - try: - if self._batches: - await asyncio.wait(tuple(self._batches), timeout=max(0, timeout)) - finally: - remaining: Final = tuple(self._batches) - for task in remaining: - task.cancel() - await asyncio.gather(*remaining, return_exceptions=True) - try: - await self.meter.stop(max(0, deadline - asyncio.get_running_loop().time())) - finally: - Logging.usage_capture_callback = None - self.meter = None - self.load = None diff --git a/litellm/proxy/spend_tracking/usage_meter.py b/litellm/proxy/spend_tracking/usage_meter.py deleted file mode 100644 index 431e4c2c3c3..00000000000 --- a/litellm/proxy/spend_tracking/usage_meter.py +++ /dev/null @@ -1,330 +0,0 @@ -from __future__ import annotations - -import asyncio -from collections import deque -from collections.abc import Awaitable, Callable, Mapping -from concurrent.futures import Future -from dataclasses import dataclass -from itertools import accumulate, chain, count, takewhile -from threading import Lock -from typing import Final - -from pydantic import BaseModel, ConfigDict, Field - -from litellm._logging import verbose_proxy_logger -from litellm.accounting import ( - DeliveryTarget, - Dimension, - DimensionName, - InboundUsageEvent, - OutboundUsageEvent, -) -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite - -_DAILY_TARGETS: Final = ( - (DimensionName.KEY, DeliveryTarget.DAILY_USER_SPEND), - (DimensionName.USER, DeliveryTarget.DAILY_USER_SPEND), - (DimensionName.TEAM, DeliveryTarget.DAILY_TEAM_SPEND), - (DimensionName.ORGANIZATION, DeliveryTarget.DAILY_ORGANIZATION_SPEND), - (DimensionName.END_USER, DeliveryTarget.DAILY_END_USER_SPEND), - (DimensionName.AGENT, DeliveryTarget.DAILY_AGENT_SPEND), - (DimensionName.TAG, DeliveryTarget.DAILY_TAG_SPEND), - (DimensionName.MODEL, DeliveryTarget.DAILY_MODEL_USAGE), - (DimensionName.TOOL, DeliveryTarget.DAILY_TOOL_SPEND), -) - - -def usage_delivery_targets(dimensions: tuple[Dimension, ...]) -> tuple[DeliveryTarget, ...]: - names: Final = frozenset(dimension.name for dimension in dimensions) - return tuple(dict.fromkeys(target for dimension, target in _DAILY_TARGETS if dimension in names)) - - -class UsageCaptureSettings(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - batch_size: int = Field(default=1000, ge=2, le=1000) - flush_interval_seconds: float = Field(default=5.0, gt=0, le=10) - max_pending_bytes: int = Field(default=16 * 1024 * 1024, ge=1024) - max_pending_records: int = Field(default=10000, ge=2) - - -class UsageCaptureError(RuntimeError): - pass - - -@dataclass(frozen=True, slots=True) -class _Pending: - writes: tuple[UsageWrite, ...] - size: int - committed: Future[None] - finished: asyncio.Event - prerequisites: tuple[asyncio.Future[None], ...] - - -async def _wait_for_batch(wakeup: asyncio.Event, timeout: float) -> None: - try: - await asyncio.wait_for(wakeup.wait(), timeout) - except asyncio.TimeoutError: - return - - -def _report_failure(receipt: Future[None]) -> None: - if not receipt.cancelled() and receipt.exception() is not None: - verbose_proxy_logger.error("Usage capture did not acknowledge commit", exc_info=receipt.exception()) - - -def _retryable_persistence_failure(error: Exception) -> bool: - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain - - outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(error) - return (outage is not None and not PrismaDBExceptionHandler.is_permanent_database_fault(outage)) or any( - isinstance(cause, Exception) - and ( - PrismaDBExceptionHandler.is_deadlock_error(cause) - or PrismaDBExceptionHandler.is_lock_timeout_error(cause) - or PrismaDBExceptionHandler.postgres_sqlstate(cause) == "40001" - ) - for cause in exception_chain(error) - ) - - -def _isolatable_persistence_failure(error: Exception) -> bool: - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain - - return isinstance(error, AccountingConflictError) or any( - isinstance(cause, Exception) and (PrismaDBExceptionHandler.postgres_sqlstate(cause) or "").startswith("22") - for cause in exception_chain(error) - ) - - -class DurableUsageMeter: - def __init__( - self, - write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]], - settings: UsageCaptureSettings, - *, - wait_for_batch: Callable[[asyncio.Event, float], Awaitable[None]] = _wait_for_batch, - retry_delay: Callable[[float], Awaitable[None]] = asyncio.sleep, - ) -> None: - self._write_batch: Final = write_batch - self.settings: Final = settings - self._wait_for_batch: Final = wait_for_batch - self._retry_delay: Final = retry_delay - self._loop: Final = asyncio.get_running_loop() - 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 - self.pending_records = 0 - - def submit( - self, writes: tuple[UsageWrite, ...], *, prerequisites: tuple[asyncio.Future[None], ...] = () - ) -> asyncio.Future[None]: - size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes) - with self._admission: - if ( - self._closed - or not writes - or len(writes) > self.settings.batch_size - or self.pending_bytes + size > self.settings.max_pending_bytes - or self.pending_records + len(writes) > self.settings.max_pending_records - ): - raise UsageCaptureError("Usage capture admission rejected: closed, empty or capacity exhausted") - receipt: Final = Future[None]() - receipt.set_running_or_notify_cancel() - receipt.add_done_callback(_report_failure) - self._queue.append(_Pending(writes, size, receipt, asyncio.Event(), prerequisites)) - self.pending_bytes += size - self.pending_records += len(writes) - self._loop.call_soon_threadsafe(self._ensure_worker) - wrapped: Final = asyncio.wrap_future(receipt, loop=self._loop) - 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() - if self._task is None or self._task.done(): - self._task = asyncio.create_task(self._run(), name="accounting usage capture") - - async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None: - 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() - if self._task is not None: - await asyncio.shield(self._task) - - async def close(self) -> None: - with self._admission: - self._closed = True - await self.flush() - - async def stop(self, timeout: float) -> None: - with self._admission: - self._closed = True - try: - await asyncio.wait_for(self.close(), timeout) - except asyncio.TimeoutError: - if self._task is not None: - self._task.cancel() - await asyncio.gather(self._task, return_exceptions=True) - while self._queue: - self._finish( - self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement") - ) - verbose_proxy_logger.error("Usage capture shutdown deadline exhausted; commit acknowledgement incomplete") - - def _take_batch(self) -> tuple[_Pending, ...]: - with self._admission: - if self._queue[0].prerequisites: - return (self._queue.popleft(),) - independent: Final = takewhile(lambda item: not item.prerequisites, self._queue) - sizes: Final = accumulate(len(item.writes) for item in independent) - count: Final = sum(1 for _ in takewhile(lambda size: size <= self.settings.batch_size, sizes)) - return tuple(self._queue.popleft() for _ in range(count)) - - async def _persist(self, writes: tuple[UsageWrite, ...]) -> None: - for attempt in count(): - try: - await self._write_batch(writes) - except Exception as error: # noqa: BLE001 # retry only classified database failures, preserving frozen writes - if not _retryable_persistence_failure(error): - raise - verbose_proxy_logger.warning("Usage persistence unavailable; retaining admitted batch for retry") - await self._retry_delay(min(10.0, 0.2 * 2 ** min(attempt, 6))) - else: - return - - async def _commit(self, batch: tuple[_Pending, ...]) -> None: - writes: Final = tuple(chain.from_iterable(item.writes for item in batch)) - try: - if batch[0].prerequisites: - outcomes: Final = await asyncio.gather( - *(asyncio.shield(receipt) for receipt in batch[0].prerequisites), return_exceptions=True - ) - for outcome in outcomes: - if isinstance(outcome, BaseException): - raise UsageCaptureError("Usage dependency did not acknowledge commit") from outcome - await self._persist(writes) - except asyncio.CancelledError: - self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted")) - raise - except Exception as error: # noqa: BLE001 # isolate data faults, surface systemic failures to all producers - if len(batch) > 1 and _isolatable_persistence_failure(error): - try: - for item in batch: - await self._commit((item,)) - except asyncio.CancelledError: - self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted")) - raise - return - self._finish( - batch, - UsageCaptureError("Usage identity conflicts with accepted content") - if isinstance(error, AccountingConflictError) - else error, - ) - else: - self._finish(batch, None) - - def _finish(self, batch: tuple[_Pending, ...], error: Exception | None) -> None: - for item in batch: - if item.finished.is_set(): - continue - item.finished.set() - 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: - item.committed.set_result(None) - else: - item.committed.set_exception(error) - - async def _run(self) -> None: - try: - while self._queue: - await self._wait_for_batch(self._wakeup, self.settings.flush_interval_seconds) - self._wakeup.clear() - await self._commit(self._take_batch()) - if ( - self._closed - or self.pending_records >= self.settings.batch_size - or self._queue - and self._queue[0].prerequisites - ): - self._wakeup.set() - except (Exception, asyncio.CancelledError): # noqa: BLE001 # reject queued receipts if the writer exits abnormally - while self._queue: - self._finish( - self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement") - ) - raise diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e68d348c16c..145c8e38229 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -168,7 +168,6 @@ from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int 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 @@ -1805,39 +1804,6 @@ 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, - _estimate_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), @@ -1846,16 +1812,12 @@ 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}) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index e28ebf439b5..3ea900642bd 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -3842,11 +3842,6 @@ class PythonCore(DeprecatedRouterMembers): priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost" } - if LiteLLMLogging.usage_capture_callback is not None: - from litellm._internal_context import internal_sub_call - - with internal_sub_call(cut_logging.litellm_call_id): - cut_logging.capture_usage_before_logging(priced_partial) return ContinuedPartial( usage=BaseTokenUsageProcessor.combine_usage_objects( usage_objects=[partial_usage, inherited_usage] diff --git a/litellm/utils.py b/litellm/utils.py index b457da058ad..e210b2fe026 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1636,7 +1636,6 @@ 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( @@ -1675,14 +1674,9 @@ 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 @@ -2280,7 +2274,6 @@ def client(original_function): # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") - logging_obj.capture_usage_before_logging(TypeAdapter(object).validate_python(result), start_time, end_time) if not is_internal_call.get(): # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py deleted file mode 100644 index 15cd5413149..00000000000 --- a/tests/integration/database/test_accounting_ledger.py +++ /dev/null @@ -1,610 +0,0 @@ -from __future__ import annotations - -import asyncio -import json -import os -import subprocess -import sys -from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping, Sequence -from contextlib import asynccontextmanager -from datetime import datetime, timezone -from decimal import Decimal -from itertools import chain -from pathlib import Path -from typing import TYPE_CHECKING, Final, TypeAlias - -import psycopg -import pytest -import pytest_asyncio -from psycopg import sql -from pydantic import JsonValue - -from litellm.accounting import ( - CostEvent, - DeliveryState, - DeliveryTarget, - EventKind, - InboundUsageEvent, - Metric, - MetricName, - OutboundUsageEvent, - Outcome, - PricingSnapshot, -) -from litellm.proxy.db import accounting_ledger -from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger, UsageWrite -from litellm.types.utils import CallTypes -from tests.integration._support.database import scratch_database -from tests.integration._support.prisma_clients import connected_proxy_prisma - -if TYPE_CHECKING: - from prisma import Prisma - - from litellm.proxy.utils import PrismaClient - -_Writer: TypeAlias = tuple["PrismaClient", AccountingLedger] - -_SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma" -_EVENT_INSERT: Final = ( - 'INSERT INTO "LiteLLM_AccountingEvent" ' - "(event_id, kind, schema_version, payload_hash, occurred_at, payload) " - "VALUES ('event', %s, %s, 'hash', '2026-01-01', '{}') RETURNING kind::text, schema_version" -) -_DELIVERY_INSERT: Final = ( - 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state, attempts) ' - "VALUES ('event', %s, %s, %s, %s) RETURNING target::text, version, state::text, attempts" -) - - -@pytest.fixture(scope="module") -def ledger_database() -> Iterator[str]: - with scratch_database("accounting") as database_url: - result: Final = subprocess.run( - [sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(_SCHEMA)], - env={**os.environ, "DATABASE_URL": database_url}, - capture_output=True, - text=True, - timeout=300, - ) - assert result.returncode == 0, result.stdout + result.stderr - yield database_url - - -@pytest.fixture -def ledger(ledger_database: str) -> Iterator[psycopg.Connection[tuple[object, ...]]]: - with psycopg.connect(ledger_database) as connection: - try: - yield connection - finally: - connection.rollback() - - -@pytest.mark.parametrize("kind", tuple(EventKind)) -def test_ledger_accepts_every_public_event_kind( - ledger: psycopg.Connection[tuple[object, ...]], kind: EventKind -) -> None: - row: Final = ledger.execute(_EVENT_INSERT, (kind.value, 1)).fetchone() - assert row == (kind.value, 1) - - -@pytest.mark.parametrize("target", tuple(DeliveryTarget)) -@pytest.mark.parametrize("state", tuple(DeliveryState)) -def test_delivery_accepts_every_public_target_and_state( - ledger: psycopg.Connection[tuple[object, ...]], target: DeliveryTarget, state: DeliveryState -) -> None: - row: Final = ledger.execute(_DELIVERY_INSERT, (target.value, 1, state.value, 0)).fetchone() - assert row == (target.value, 1, state.value, 0) - - -@pytest.mark.parametrize( - ("enum_name", "values"), - [ - ("AccountingEventKind", tuple(value.value for value in EventKind)), - ("AccountingDeliveryState", tuple(value.value for value in DeliveryState)), - ("AccountingDeliveryTarget", tuple(value.value for value in DeliveryTarget)), - ], -) -def test_migrated_enums_match_the_public_contract( - ledger: psycopg.Connection[tuple[object, ...]], enum_name: str, values: tuple[str, ...] -) -> None: - rows: Final = ledger.execute( - "SELECT e.enumlabel FROM pg_enum e JOIN pg_type t ON t.oid = e.enumtypid WHERE t.typname = %s", - (enum_name,), - ).fetchall() - assert frozenset(row[0] for row in rows) == frozenset(values) - - -@pytest.mark.parametrize("version", [0, -1]) -def test_event_rejects_nonpositive_schema_versions( - ledger: psycopg.Connection[tuple[object, ...]], version: int -) -> None: - with pytest.raises(psycopg.errors.CheckViolation, match="AccountingEvent_schema_version_check"): - ledger.execute(_EVENT_INSERT, (EventKind.OUTBOUND_USAGE.value, version)) - - -@pytest.mark.parametrize( - ("version", "attempts", "constraint"), - [ - (0, 0, "AccountingDelivery_version_check"), - (-1, 0, "AccountingDelivery_version_check"), - (1, -1, "AccountingDelivery_attempts_check"), - ], -) -def test_delivery_rejects_invalid_versions_and_attempts( - ledger: psycopg.Connection[tuple[object, ...]], version: int, attempts: int, constraint: str -) -> None: - with pytest.raises(psycopg.errors.CheckViolation, match=constraint): - ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, version, "pending", attempts)) - - -@pytest.mark.parametrize("column", ["state", "target"]) -def test_delivery_rejects_unknown_enum_values(ledger: psycopg.Connection[tuple[object, ...]], column: str) -> None: - ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, 1, "pending", 0)) - with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): - ledger.execute( - sql.SQL('UPDATE "LiteLLM_AccountingDelivery" SET {} = %s').format(sql.Identifier(column)), ("invalid",) - ) - - -def test_event_rejects_unknown_kind(ledger: psycopg.Connection[tuple[object, ...]]) -> None: - with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): - ledger.execute(_EVENT_INSERT, ("invalid", 1)) - - -def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Connection[tuple[object, ...]]) -> None: - row: Final = ledger.execute( - 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) ' - "VALUES ('event', %s, 1) RETURNING state::text, attempts", - (DeliveryTarget.KEY_SPEND.value,), - ).fetchone() - assert row == ("pending", 0) - - -_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) -_CUTOFF: Final = datetime(2026, 2, 1, tzinfo=timezone.utc) -_TARGETS: Final = (DeliveryTarget.USER_SPEND, DeliveryTarget.TEAM_SPEND) - - -def _usage(event_id: str = "usage", *, inbound: bool = False) -> InboundUsageEvent | OutboundUsageEvent: - fields: Final = dict( - event_id=event_id, - request_id="request", - call_type=CallTypes.completion, - started_at=_AT, - ended_at=_AT, - outcome=Outcome.SUCCEEDED, - dimensions=(), - ) - if inbound: - return InboundUsageEvent(**fields, metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),)) - return OutboundUsageEvent(**fields, attempt_id=event_id, retry_of_attempt_id=None, metrics=None) - - -def _cost(source: str = "usage", event_id: str = "cost", amount: str = "0.125") -> CostEvent: - return CostEvent( - event_id=event_id, - usage_event_id=source, - occurred_at=_AT, - dimensions=(), - charges=(Metric(name=MetricName.SPEND, value=Decimal(amount), unit="USD"),), - snapshot=PricingSnapshot( - calculator_version="fixture", - model="fixture", - provider=None, - pricing_source="cost_map", - priced_model_key="fixture", - total_cost=Decimal(amount), - ), - ) - - -def _envelope(event: AccountingEvent, *, context: Mapping[str, JsonValue] | None = None) -> str: - return json.dumps({"event": event.model_dump(mode="json"), "context": context or {}}) - - -async def _insert( - store: AccountingLedger, - event: AccountingEvent, - targets: tuple[DeliveryTarget, ...] = _TARGETS, - *, - context: Mapping[str, JsonValue] | None = None, -) -> str: - return await store.insert_event(event, targets, envelope=_envelope(event, context=context)) - - -def _rows(database_url: str, statement: str, parameters: tuple[object, ...] = ()) -> Sequence[tuple[object, ...]]: - with psycopg.connect(database_url) as connection: - cursor: Final = connection.execute(statement, parameters) - return cursor.fetchall() if cursor.description is not None else [] - - -def _complete_and_age(database_url: str) -> None: - with psycopg.connect(database_url) as connection: - connection.execute('UPDATE "LiteLLM_AccountingEvent" SET recorded_at = %s', (_AT,)) - connection.execute("UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'applied', applied_at = %s", (_AT,)) - - -@pytest_asyncio.fixture(loop_scope="function") -async def writer(ledger_database: str, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_Writer]: - monkeypatch.setenv("DATABASE_URL", ledger_database) - monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) - _rows(ledger_database, 'TRUNCATE "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"') - async with connected_proxy_prisma(ledger_database) as client: - yield client, AccountingLedger.for_client(client) - - -@pytest.mark.asyncio -async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, ledger_database: str) -> None: - _, store = writer - event: Final = _usage() - assert await _insert(store, event, context={"requested_region": "eu", "quantity": "1.25"}) == "inserted" - rows: Final = _rows(ledger_database, 'SELECT kind::text, source_event_id, payload FROM "LiteLLM_AccountingEvent"') - assert rows == [ - ("outbound_usage", None, json.loads(_envelope(event, context={"requested_region": "eu", "quantity": "1.25"}))) - ] - deliveries: Final = _rows( - ledger_database, 'SELECT target::text, state::text, attempts FROM "LiteLLM_AccountingDelivery"' - ) - assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS} - - -@pytest.mark.asyncio -async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None: - _, store = writer - event: Final = _usage(inbound=True) - results: Final = await asyncio.gather(*(_insert(store, event) for _ in range(8))) - assert results.count("inserted") == 1 and results.count("duplicate") == 7 - _complete_and_age(ledger_database) - before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') - reordered: Final = json.dumps(json.loads(_envelope(event)), sort_keys=True, indent=2) - assert await store.insert_event(event, tuple(reversed(_TARGETS)) + _TARGETS, envelope=reordered) == "duplicate" - assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') == before - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("change", ["payload", "context", "targets", "schema_version", "version"]) -async def test_conflicting_acceptance_cannot_overwrite_or_add_work( - writer: _Writer, ledger_database: str, change: str -) -> None: - _, store = writer - event: Final = _usage() - await _insert(store, event) - altered: Final = event.model_copy(update={"outcome": Outcome.FAILED}) if change == "payload" else event - targets: Final = (DeliveryTarget.KEY_SPEND,) if change == "targets" else _TARGETS - context: Final = {"override": "0.5"} if change == "context" else {} - options: Final = {change: 2} if change in ("schema_version", "version") else {} - with pytest.raises(AccountingConflictError, match="conflicts"): - await store.insert_event(altered, targets, envelope=_envelope(altered, context=context), **options) - assert _rows(ledger_database, 'SELECT payload FROM "LiteLLM_AccountingEvent"') == [(json.loads(_envelope(event)),)] - assert set(_rows(ledger_database, 'SELECT target::text FROM "LiteLLM_AccountingDelivery"')) == { - (t.value,) for t in _TARGETS - } - - -@pytest.mark.asyncio -async def test_competing_cost_identities_accept_only_one_authoritative_charge( - writer: _Writer, ledger_database: str -) -> None: - _, store = writer - await _insert(store, _usage()) - results: Final = await asyncio.gather( - _insert(store, _cost()), _insert(store, _cost(event_id="other-cost")), return_exceptions=True - ) - assert sum(result == "inserted" for result in results) == 1 - assert sum(isinstance(result, AccountingConflictError) for result in results) == 1 - accepted: Final = _rows( - ledger_database, "SELECT event_id, payload FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'" - ) - assert len(accepted) == 1 - event: Final = CostEvent.model_validate_json(json.dumps(accepted[0][1]["event"])) - assert await _insert(store, event) == "duplicate" - assert event.snapshot.total_cost == Decimal("0.125") - with pytest.raises(AccountingConflictError): - await _insert(store, _cost(event_id=event.event_id, amount="0.5")) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("source", ["missing", "inbound", "cost"]) -async def test_cost_requires_an_existing_outbound_source(writer: _Writer, ledger_database: str, source: str) -> None: - _, store = writer - if source == "inbound": - await _insert(store, _usage("source", inbound=True)) - if source == "cost": - await _insert(store, _usage()) - await _insert(store, _cost(event_id="source")) - with pytest.raises(AccountingConflictError, match="outbound usage"): - await _insert(store, _cost(source="source", event_id="invalid")) - assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'invalid'") == [ - (0,) - ] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"]) -async def test_interruption_rolls_back_or_retry_recognizes_commit( - writer: _Writer, ledger_database: str, failure: str -) -> None: - client, store = writer - - @asynccontextmanager - async def interrupted() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - if failure == "before_commit": - raise RuntimeError("injected pre-commit failure") - if failure == "cancellation": - raise asyncio.CancelledError - raise RuntimeError("injected lost acknowledgement") - - error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError - with pytest.raises(error): - await _insert(AccountingLedger(interrupted), _usage()) - expected: Final = 1 if failure == "after_commit" else 0 - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(expected,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(expected * len(_TARGETS),)] - assert await _insert(store, _usage()) == ("duplicate" if expected else "inserted") - - -@pytest.mark.asyncio -async def test_delivery_insert_failure_rolls_back_event(writer: _Writer, ledger_database: str) -> None: - from prisma.errors import RawQueryError - - _, store = writer - with psycopg.connect(ledger_database) as connection: - connection.execute( - 'ALTER TABLE "LiteLLM_AccountingDelivery" ADD CONSTRAINT reject_delivery CHECK (version < 1)' - ) - try: - with pytest.raises(RawQueryError, match="reject_delivery"): - await _insert(store, _usage()) - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] - finally: - with psycopg.connect(ledger_database) as connection: - connection.execute('ALTER TABLE "LiteLLM_AccountingDelivery" DROP CONSTRAINT reject_delivery') - - -@pytest.mark.asyncio -@pytest.mark.parametrize("blocked", ["pending", "parked", "unpriced", "recent_cost", "recovery"]) -async def test_retention_preserves_incomplete_or_recovery_needed_families( - writer: _Writer, ledger_database: str, blocked: str -) -> None: - _, store = writer - await _insert(store, _usage()) - if blocked != "unpriced": - await _insert(store, _cost()) - _complete_and_age(ledger_database) - with psycopg.connect(ledger_database) as connection: - if blocked in ("pending", "parked"): - connection.execute( - "UPDATE \"LiteLLM_AccountingDelivery\" SET state = %s WHERE event_id = 'cost'", (blocked,) - ) - if blocked == "recent_cost": - connection.execute( - "UPDATE \"LiteLLM_AccountingEvent\" SET recorded_at = %s WHERE event_id = 'cost'", (_CUTOFF,) - ) - before: Final = _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') - assert ( - await store.purge_completed( - redelivery_before=_CUTOFF, recovery_before=_AT if blocked == "recovery" else _CUTOFF - ) - == 0 - ) - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == before - - -@pytest.mark.asyncio -async def test_retention_removes_completed_family_and_receipts_together(writer: _Writer, ledger_database: str) -> None: - _, store = writer - await _insert(store, _usage()) - await _insert(store, _cost()) - _complete_and_age(ledger_database) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] - - -@pytest.mark.asyncio -async def test_retention_respects_batch_bound_and_skips_locked_families(writer: _Writer, ledger_database: str) -> None: - client, store = writer - for name in ("a", "b", "c"): - await _insert(store, _usage(name, inbound=True)) - _complete_and_age(ledger_database) - async with client.writer_db.tx() as tx: - await tx.query_raw("SELECT event_id FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'a' FOR UPDATE") - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ - ("a",), - ("c",), - ] - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - - -@pytest.mark.asyncio -async def test_retention_makes_progress_past_pending_and_unpriced_events(writer: _Writer, ledger_database: str) -> None: - _, store = writer - await _insert(store, _usage("a", inbound=True)) - await _insert(store, _usage("b")) - await _insert(store, _usage("c", inbound=True)) - _complete_and_age(ledger_database) - _rows(ledger_database, "UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'pending' WHERE event_id = 'a'") - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ - ("a",), - ("b",), - ] - - -@pytest.mark.asyncio -async def test_zero_cost_and_empty_delivery_sets_are_durable(writer: _Writer, ledger_database: str) -> None: - _, store = writer - await _insert(store, _usage(), targets=()) - event: Final = _cost(amount="0") - assert await _insert(store, event, targets=()) == "inserted" - assert await _insert(store, event, targets=()) == "duplicate" - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] - _complete_and_age(ledger_database) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - - -@pytest.mark.asyncio -async def test_retention_cannot_delete_source_while_cost_acceptance_is_uncommitted( - writer: _Writer, ledger_database: str -) -> None: - client, store = writer - await _insert(store, _usage()) - _complete_and_age(ledger_database) - inserted: Final = asyncio.Event() - release: Final = asyncio.Event() - - @asynccontextmanager - async def paused_commit() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - inserted.set() - await release.wait() - - task: Final = asyncio.create_task(_insert(AccountingLedger(paused_commit), _cost())) - try: - await asyncio.wait_for(inserted.wait(), timeout=5) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 0 - finally: - release.set() - await task - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ - ("cost",), - ("usage",), - ] - - -@pytest.mark.asyncio -async def test_retention_failure_rolls_back_event_and_receipt_deletion(writer: _Writer, ledger_database: str) -> None: - client, store = writer - await _insert(store, _usage(inbound=True)) - _complete_and_age(ledger_database) - - @asynccontextmanager - async def interrupted() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - raise RuntimeError("injected cleanup failure") - - with pytest.raises(RuntimeError, match="cleanup failure"): - await AccountingLedger(interrupted).purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 1 - - -@pytest.mark.asyncio -async def test_ledger_reads_and_writes_use_primary_when_replica_is_configured( - writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch -) -> None: - with scratch_database("replica") as replica_url: - monkeypatch.setenv("DATABASE_URL_READ_REPLICA", replica_url) - async with connected_proxy_prisma(ledger_database) as client: - store: Final = AccountingLedger.for_client(client) - await _insert(store, _usage()) - assert await _insert(store, _usage()) == "duplicate" - assert await _insert(store, _cost()) == "inserted" - _complete_and_age(ledger_database) - assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 - assert _rows(replica_url, "SELECT to_regclass('public.\"LiteLLM_AccountingEvent\"')") == [(None,)] - - -def _index_conditions(plan: object) -> tuple[tuple[object, object], ...]: - if isinstance(plan, dict): - own: Final = ((plan.get("Index Name"), plan.get("Index Cond")),) if "Index Name" in plan else () - children: Final = (plan.get("Plan"), *plan.get("Plans", ())) - return own + tuple(chain.from_iterable(_index_conditions(child) for child in children if child is not None)) - return tuple(chain.from_iterable(_index_conditions(item) for item in plan)) if isinstance(plan, list) else () - - -@pytest.mark.parametrize("statement", [accounting_ledger._DELETE_EVENTS, accounting_ledger._DELETE_DELIVERIES]) -def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, statement: str) -> None: - explained: Final = "EXPLAIN (FORMAT JSON) " + statement.replace("$1", "%s") - with psycopg.connect(ledger_database) as connection: - connection.execute("SET enable_seqscan = off") - plan: Final = connection.execute(explained, ("usage",) * statement.count("$1")).fetchone() - assert plan is not None - conditions: Final = _index_conditions(plan[0]) - assert any( - name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition) - for name, condition in conditions - ), conditions - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "conflict"]) -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").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 - async def interrupted() -> AsyncGenerator[Prisma]: - async with client.writer_db.tx() as tx: - yield tx - if failure == "before_commit": - raise RuntimeError("injected pre-commit failure") - if failure == "after_commit": - raise RuntimeError("injected lost acknowledgement") - - if failure == "conflict": - await store.insert_usage_batch((writes[1],)) - _complete_and_age(ledger_database) - rejected: Final = (writes[0], UsageWrite(events[1], _envelope(events[1], context={"changed": True}), _TARGETS)) - with pytest.raises(AccountingConflictError): - await store.insert_usage_batch(rejected) - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("attempt",)] - else: - with pytest.raises(RuntimeError, match="injected"): - await AccountingLedger(interrupted).insert_usage_batch(writes) - count: Final = 2 if failure == "after_commit" else 0 - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(count,)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(count * len(_TARGETS),)] - await store.insert_usage_batch(writes) - recovered: Final = await store.load_usage_batch(tuple(event.event_id for event in events)) - assert len(recovered) == 1 - assert recovered[0].event == events[1] - assert json.loads(recovered[0].envelope) == json.loads(writes[1].envelope) - assert set(recovered[0].targets) == set(_TARGETS) - _complete_and_age(ledger_database) - before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') - await store.insert_usage_batch((*writes, *writes)) - assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') == before - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(2,)] - - -@pytest.mark.asyncio -async def test_invalid_jsonb_submission_cannot_reject_another_producers_commit( - writer: _Writer, ledger_database: str -) -> None: - from prisma.errors import DataError - - from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings - - _, store = writer - invalid: Final = _usage("invalid") - valid: Final = _usage("valid") - meter: Final = DurableUsageMeter(store.insert_usage_batch, UsageCaptureSettings(batch_size=2)) - rejected: Final = meter.submit( - (UsageWrite(invalid, _envelope(invalid, context={"parent": "bad\x00text"}), _TARGETS),) - ) - accepted: Final = meter.submit((UsageWrite(valid, _envelope(valid), _TARGETS),)) - try: - with pytest.raises(DataError): - await rejected - await accepted - finally: - await meter.close() - assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("valid",)] - assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] - assert (meter.pending_bytes, meter.pending_records) == (0, 0) diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py new file mode 100644 index 00000000000..ec0ab44e505 --- /dev/null +++ b/tests/unit/accounting/test_projection.py @@ -0,0 +1,431 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.accounting.projection import project +from litellm.accounting.snapshot import build_accounting_payload +from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.proxy._types import SpendLogsPayload +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.proxy_server import proxy_logging_obj +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.utils import PrismaClient +from litellm.types.utils import ModelResponse, Usage + +_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc) +_FIELDS: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.parametrize("status", ["success", "failure"]) +@pytest.mark.parametrize("origin", [None, "auto_router", "billed_retry"]) +@pytest.mark.parametrize("spend", [0.0, 0.25]) +@pytest.mark.asyncio +async def test_daily_projection_matches_legacy_writer(status: str, origin: str | None, spend: float) -> None: + payload: Final = get_logging_payload( + kwargs={ + "model": "projection-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "projection-request", + "litellm_params": {"metadata": {"user_api_key": "key", "user_api_key_user_id": "user", "status": status}}, + }, + response_obj=ModelResponse( + id="projection-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + ) + metadata: Final = { + **_FIELDS.validate_json(payload["metadata"]), + **({BILLED_RETRY_ATTEMPT_METADATA_KEY: True} if origin == "billed_retry" else {}), + **({INTERNAL_CALL_ORIGIN_METADATA_KEY: origin} if origin not in (None, "billed_retry") else {}), + } + row: Final = {**payload, "spend": spend, "metadata": json.dumps(metadata), "request_duration_ms": 25} + legacy: Final = await DBSpendUpdateWriter()._common_add_spend_log_transaction_to_daily_transaction( + row, PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj), "user" + ) + assert legacy is not None + accounting_input: Final[SpendLogsPayload] = {**row} + accounting: Final = build_accounting_payload( + accounting_input, + entities={"api_key": payload["api_key"], "user": payload["user"]}, + model_access_groups=(), + windows=(), + router=lambda: None, + ) + restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str)) + stored_daily: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend") + assert dict(stored_daily.increments) == {name: legacy[name] for name, _ in stored_daily.increments} + assert dict(stored_daily.key)["date"] == legacy["date"] + assert dict(stored_daily.key)["api_key"] == legacy["api_key"] + assert not {"messages", "response", "proxy_server_request", "requester_ip_address"}.intersection(accounting) + + +def test_projection_uses_frozen_savings_and_keeps_duplicate_entity_tags() -> None: + row: Final[Mapping[str, object]] = { + "request_id": "saved", + "startTime": _NOW, + "api_key": "key", + "user": "user", + "model": "unpriced-projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "request_tags": '["tag", "tag"]', + "spend": 0.5, + "metadata": json.dumps( + { + "accounting_savings": { + "compression": 1.0, + "prompt_caching": -2.0, + "gateway_injected_caching": 3.0, + "autorouter": 4.0, + }, + "autorouter_savings": 9000, + } + ), + } + contributions: Final = project(row) + tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_TagTable") + daily_tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_DailyTagSpend") + assert sum(dict(item.increments)["spend"] for item in tags) == 1.0 + assert len(daily_tags) == 1 + assert {key: value for key, value in daily_tags[0].increments if key.endswith("savings_spend")} == { + "compression_savings_spend": 1.0, + "prompt_caching_savings_spend": -2.0, + "gateway_injected_caching_savings_spend": 3.0, + "autorouter_savings_spend": 4.0, + } + + +def test_projection_carries_windows_and_does_not_reprice_guardrails() -> None: + row: Final = { + "request_id": "guarded", + "startTime": _NOW, + "api_key": "key", + "spend": 0.75, + "metadata": json.dumps( + { + "accounting_windows": [ + { + "entity_type": "key", + "entity_id": "key", + "window_duration": "1d", + "window_start": _NOW.isoformat(), + } + ], + "guardrail_information": [ + { + "guardrail_id": "guard", + "guardrail_usage": {"text": 3, "images": 2}, + "guardrail_cost_by_unit": {"text": 0.3, "images": None}, + } + ], + } + ), + } + contributions: Final = project(row) + window: Final = next(item for item in contributions if item.table == "LiteLLM_BudgetWindowSpend") + units: Final = { + dict(item.key)["usage_unit"]: dict(item.increments) + for item in contributions + if item.table == "LiteLLM_DailyGuardrailUsageUnits" + } + assert dict(window.increments) == {"spend": 0.75} + assert dict(window.values)["window_start"] == _NOW.isoformat() + assert units == { + "text": {"units": 3, "cost": 0.3, "untracked_units": 0}, + "images": {"units": 2, "cost": 0.0, "untracked_units": 2}, + } + + +@pytest.mark.asyncio +async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.responses.streaming_iterator import drain_logging_tasks + + client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj) + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr(client, "spend_log_transactions", []) + monkeypatch.setattr(client, "model_usage_transactions", []) + writer: Final = DBSpendUpdateWriter() + charged: Final = await writer.update_database( + token="sk-projection", + user_id="user", + end_user_id="end-user", + team_id="team", + org_id="org", + project_id="project", + kwargs={ + "model": "projection-test", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "entity-request", + "agent_id": "agent", + "user": "end-user", + "litellm_params": { + "metadata": { + "user_api_key": "sk-projection", + "user_api_key_user_id": "user", + "user_api_key_team_id": "team", + "user_api_key_org_id": "org", + "user_api_key_project_id": "project", + "user_api_key_end_user_id": "end-user", + "tags": ["tag", "tag"], + }, + "user": "end-user", + }, + }, + completion_response=ModelResponse( + id="entity-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + response_cost=0.25, + ) + await drain_logging_tasks() + assert charged + assert len(client.spend_log_transactions) == 1 + queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue() + logged: Final[SpendLogsPayload] = client.spend_log_transactions[0] + projected: Final = project( + _FIELDS.validate_json( + json.dumps( + build_accounting_payload( + logged, + entities={ + "api_key": logged["api_key"], + "user": "user", + "end_user": "end-user", + "team_id": "team", + "organization_id": "org", + "project_id": "project", + "agent_id": "agent", + }, + model_access_groups=(), + windows=(), + router=lambda: None, + ), + default=str, + ) + ) + ) + legacy: Final = {(item["entity_type"].value, item["entity_id"]): item["response_cost"] for item in queued} + assert legacy[("key", proxy_server.hash_token("sk-projection"))] == next( + dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_VerificationToken" + ) + for table, entity, identity in ( + ("LiteLLM_UserTable", "user", "user"), + ("LiteLLM_TeamTable", "team", "team"), + ("LiteLLM_OrganizationTable", "organization", "org"), + ("LiteLLM_EndUserTable", "end_user", "end-user"), + ("LiteLLM_AgentsTable", "agent", "agent"), + ("LiteLLM_ProjectTable", "project", "project"), + ): + assert legacy[(entity, identity)] == next( + dict(item.increments)["spend"] for item in projected if item.table == table + ) + assert sum(dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_TagTable") == sum( + item["response_cost"] or 0 for item in queued if item["entity_type"].value == "tag" + ) + + daily_queues: Final = ( + ("LiteLLM_DailyUserSpend", writer.daily_spend_update_queue), + ("LiteLLM_DailyTeamSpend", writer.daily_team_spend_update_queue), + ("LiteLLM_DailyOrganizationSpend", writer.daily_org_spend_update_queue), + ("LiteLLM_DailyEndUserSpend", writer.daily_end_user_spend_update_queue), + ("LiteLLM_DailyAgentSpend", writer.daily_agent_spend_update_queue), + ("LiteLLM_DailyTagSpend", writer.daily_tag_spend_update_queue), + ) + for table, queue in daily_queues: + transactions = await queue.flush_and_get_aggregated_daily_spend_update_transactions() + matching = tuple(item for item in projected if item.table == table) + assert matching, table + contribution = matching[0] + assert len(transactions) == 1 + transaction = next(iter(transactions.values())) + assert dict(contribution.increments) == {name: transaction[name] for name, _ in contribution.increments} + assert len(client.model_usage_transactions) == 1 + model: Final = client.model_usage_transactions[0] + model_projection: Final = next(item for item in projected if item.table == "LiteLLM_DailyModelUsage") + assert dict(model_projection.increments) == { + "spend": model.spend, + "prompt_tokens": model.prompt_tokens, + "completion_tokens": model.completion_tokens, + "request_count": int(model.counts_as_request), + "successful_requests": int(model.successful), + "failed_requests": int(not model.successful), + } + + +@pytest.mark.parametrize("served_id", ["requested", "fallback", "missing"]) +@pytest.mark.asyncio +async def test_projection_preserves_fallback_narrowed_model_access_groups(served_id: str) -> None: + from litellm import Router + from litellm.proxy.db.db_spend_update_writer import debitable_model_access_groups + + router: Final = Router( + model_list=[ + { + "model_name": "first", + "litellm_params": {"model": "openai/projection-test", "api_key": "unused"}, + "model_info": {"id": "requested", "access_groups": ["first-group"]}, + }, + { + "model_name": "second", + "litellm_params": {"model": "openai/projection-test", "api_key": "unused"}, + "model_info": {"id": "fallback", "access_groups": ["second-group", "unattributed"]}, + }, + ] + ) + attributed: Final = ("first-group", "second-group", "first-group") + writer: Final = DBSpendUpdateWriter() + client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj) + await writer._update_model_access_group_db(0.5, attributed, served_id, client, router) + queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue() + frozen_groups: Final = debitable_model_access_groups(attributed, served_id, router) + projected: Final = project( + {"spend": 0.5, "metadata": json.dumps({"accounting_model_access_groups": frozen_groups})} + ) + assert {dict(item.key)["access_group_name"]: dict(item.increments)["spend"] for item in projected} == { + item["entity_id"]: item["response_cost"] for item in queued + } + assert "unattributed" not in frozen_groups + + +def test_guardrail_projection_skips_unmeasured_counters() -> None: + row: Final = { + "request_id": "guarded", + "startTime": _NOW, + "metadata": json.dumps( + { + "guardrail_information": [ + {"guardrail_usage": {"text": 3}}, + { + "guardrail_id": "guard", + "guardrail_usage": {"zero": 0, "negative": -1, "bool": True, "unknown": None, "text": 2}, + }, + ] + } + ), + } + units: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyGuardrailUsageUnits") + assert len(units) == 1 + assert dict(units[0].key)["usage_unit"] == "text" + assert dict(units[0].increments) == {"units": 2, "cost": 0.0, "untracked_units": 2} + + +def test_lean_row_without_model_does_not_fabricate_daily_entity_usage() -> None: + projected: Final = project( + { + "request_id": "lean", + "api_key": "key", + "startTime": _NOW, + "spend": 0.5, + "metadata": {"accounting_entities": {"api_key": "key"}}, + } + ) + assert tuple(item.table for item in projected) == ("LiteLLM_VerificationToken", "LiteLLM_DailyModelUsage") + + +@pytest.mark.parametrize( + "cost_fields", + [ + {"guardrail_cost_in_spend": False, "guardrail_cost_by_unit": {"text": 1.0}}, + {"guardrail_cost_by_unit": {"text": -1.0}}, + ], +) +def test_report_only_or_invalid_guardrail_prices_remain_untracked(cost_fields: Mapping[str, object]) -> None: + row: Final = { + "request_id": "guarded", + "startTime": _NOW, + "request_tags": '["task:debugging"]', + "metadata": json.dumps( + { + "guardrail_information": [ + {"guardrail_id": "guard", "guardrail_usage": {"text": 2}, **cost_fields}, + ] + } + ), + } + contributions: Final = project(row) + guardrail: Final = next(item for item in contributions if item.table == "LiteLLM_DailyGuardrailUsageUnits") + model: Final = next(item for item in contributions if item.table == "LiteLLM_DailyModelUsage") + assert dict(guardrail.increments) == {"units": 2, "cost": 0.0, "untracked_units": 2} + assert dict(model.key)["task_type"] == "debugging" + + +@pytest.mark.parametrize("organization_id", [None, "", "charged-org"]) +def test_daily_organization_uses_frozen_charged_identity(organization_id: str | None) -> None: + row: Final[Mapping[str, object]] = { + "startTime": _NOW, + "api_key": "key", + "model": "projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "organization_id": "stored-org", + "spend": 0.25, + "metadata": {"accounting_entities": {"organization_id": organization_id}}, + } + daily: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyOrganizationSpend") + if organization_id is None: + assert daily == () + else: + assert len(daily) == 1 + assert dict(daily[0].key)["organization_id"] == organization_id + assert dict(daily[0].increments)["spend"] == 0.25 + + +@pytest.mark.parametrize("start", [_NOW, "2026-11-02T12:00:00+00:00", "2026-11-02 12:00:00+00:00"]) +def test_every_daily_row_uses_the_calendar_date_of_the_request(start: object) -> None: + row: Final[Mapping[str, object]] = { + "request_id": "dated", + "startTime": start, + "api_key": "key", + "user": "user", + "team_id": "team", + "model": "unpriced-projection-test", + "model_group": "group", + "custom_llm_provider": "openai", + "spend": 0.5, + "metadata": "{}", + } + dates: Final = { + item.table: dict(item.key)["date"] for item in project(row) if item.table.startswith("LiteLLM_Daily") + } + assert {"LiteLLM_DailyUserSpend", "LiteLLM_DailyModelUsage"} <= dates.keys() + assert set(dates.values()) == {"2026-11-02"} + + +def test_unparseable_start_time_skips_only_guardrail_usage() -> None: + projected: Final = project( + { + "request_id": "bad-time", + "startTime": "yesterday at noon", + "spend": 0.5, + "metadata": { + "accounting_entities": {"api_key": "key"}, + "guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 2}}], + }, + } + ) + tables: Final = {item.table for item in projected} + assert "LiteLLM_VerificationToken" in tables + assert "LiteLLM_DailyGuardrailUsageUnits" not in tables + + +def test_guardrail_usage_is_dated_in_utc() -> None: + projected: Final = project( + { + "request_id": "offset", + "startTime": "2026-11-02T22:30:00-05:00", + "metadata": {"guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 1}}]}, + } + ) + assert [dict(item.key)["date"] for item in projected if item.table == "LiteLLM_DailyGuardrailUsageUnits"] == [ + "2026-11-03" + ] diff --git a/tests/unit/accounting/test_types.py b/tests/unit/accounting/test_types.py index f645ee6c2cc..c1be39ab111 100644 --- a/tests/unit/accounting/test_types.py +++ b/tests/unit/accounting/test_types.py @@ -19,7 +19,6 @@ from litellm.accounting import ( Dimension, DimensionFilter, DimensionName, - EventKind, Granularity, InboundUsageEvent, Metric, @@ -442,7 +441,6 @@ def test_subset_token_counts_round_trip_without_changing_totals() -> None: @pytest.mark.parametrize( ("enum_name", "python_enum"), [ - ("AccountingEventKind", EventKind), ("AccountingDeliveryState", DeliveryState), ("AccountingDeliveryTarget", DeliveryTarget), ], diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 2d02468fd8c..04ba823d08f 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -484,55 +484,6 @@ async def test_bridged_text_completion_is_logged_as_the_text_completion_it_serve assert logged.payload["response"]["choices"][0]["text"] == "pong" -@pytest.mark.asyncio -async def test_a_bridged_chat_response_waits_on_the_usage_ack_like_any_chat_response( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier - - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - logging_obj: Final = Logging( - model="openai/gpt-5.3-codex", - messages=[], - stream=False, - call_type="acompletion", - start_time=datetime.datetime.now(), - litellm_call_id="bridged-chat", - function_id="bridged-chat", - ) - with respx.mock(assert_all_called=True) as router: - router.post(url__startswith="https://api.openai.com/v1/responses").mock( - return_value=httpx.Response(200, json=_searched_responses_body("openai/gpt-5.3-codex")) - ) - served: Final = await litellm.acompletion( - model="openai/gpt-5.3-codex", - messages=[{"role": "user", "content": "ping"}], - api_key="test-key", - litellm_logging_obj=logging_obj, - ) - response: Final = asyncio.create_task(capture_proxy_response(logging_obj, served)) - try: - await database.entered.wait() - assert not response.done() - assert database.committed == () - database.release.set() - await response - assert "bridged-chat:inbound" in database.committed - finally: - database.release.set() - await runtime.stop(1) - await response - - def _served_message_start_id(events: Sequence[bytes]) -> object: lines: Final = b"".join(events).decode().splitlines() starts: Final = tuple( diff --git a/tests/unit/proxy/db/test_accounting_ledger.py b/tests/unit/proxy/db/test_accounting_ledger.py deleted file mode 100644 index 6b301146301..00000000000 --- a/tests/unit/proxy/db/test_accounting_ledger.py +++ /dev/null @@ -1,149 +0,0 @@ -import json -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final, Literal, NoReturn - -import pytest - -from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome -from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingLedger, UsageWrite -from litellm.types.utils import CallTypes - - -def _no_transaction() -> NoReturn: - raise AssertionError("Invalid accounting input must not open a database transaction") - - -@pytest.fixture -def event() -> InboundUsageEvent: - return InboundUsageEvent( - event_id="caller", - request_id="caller", - call_type=CallTypes.completion, - started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - outcome=Outcome.SUCCEEDED, - dimensions=(), - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("envelope", ["{}", "[]", '{"event":null}', '{"event":{}}', "not json"]) -async def test_invalid_envelopes_fail_before_database_access(event: InboundUsageEvent, envelope: str) -> None: - with pytest.raises(ValueError, match=r"Accounting envelope|validation error"): - await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("version", [0, -1, True, 1.5]) -@pytest.mark.parametrize("field", ["version", "schema_version"]) -async def test_invalid_versions_fail_before_database_access( - event: InboundUsageEvent, version: int | float, field: Literal["version", "schema_version"] -) -> None: - with pytest.raises(ValueError, match="positive integers"): - await AccountingLedger(_no_transaction).insert_event( - event, - (), - envelope="{}", - **{field: version}, # pyright: ignore[reportArgumentType] # Exercise invalid runtime versions - ) - - -@pytest.mark.asyncio -async def test_constructed_invalid_events_are_revalidated(event: InboundUsageEvent) -> None: - invalid: Final = event.model_copy(update={"event_id": ""}) - with pytest.raises(ValueError, match="event_id"): - await AccountingLedger(_no_transaction).insert_event(invalid, (), envelope="{}") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("limit", [0, -1, 1001, True, 1.5]) -async def test_retention_rejects_invalid_batch_bounds(limit: int | float) -> None: - with pytest.raises(ValueError, match="batch size"): - await AccountingLedger(_no_transaction).purge_completed( - limit=limit # pyright: ignore[reportArgumentType] # Exercise invalid runtime batch sizes - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("missing", ["both", "redelivery", "recovery"]) -async def test_retention_requires_both_safety_watermarks(missing: Literal["both", "redelivery", "recovery"]) -> None: - cutoff: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - deleted: Final = await AccountingLedger(_no_transaction).purge_completed( - redelivery_before=cutoff if missing == "recovery" else None, - recovery_before=cutoff if missing == "redelivery" else None, - ) - assert deleted == 0 - - -@pytest.mark.asyncio -async def test_retention_rejects_ambiguous_naive_watermarks() -> None: - with pytest.raises(ValueError, match="timezone"): - await AccountingLedger(_no_transaction).purge_completed( - redelivery_before=datetime(2026, 1, 1), recovery_before=datetime(2026, 1, 1, tzinfo=timezone.utc) - ) - - -@pytest.mark.asyncio -async def test_nonfinite_context_is_rejected_before_database_access(event: InboundUsageEvent) -> None: - envelope: Final = json.dumps({"event": event.model_dump(mode="json"), "context": {"price": float("nan")}}) - with pytest.raises(ValueError, match="Out of range float"): - await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) - - -@pytest.mark.asyncio -async def test_unknown_delivery_target_is_rejected_before_database_access(event: InboundUsageEvent) -> None: - envelope: Final = json.dumps({"event": event.model_dump(mode="json")}) - with pytest.raises(ValueError, match="not a valid DeliveryTarget"): - await AccountingLedger(_no_transaction).insert_event( - event, - ("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target - envelope=envelope, - ) - - -@pytest.mark.asyncio -async def test_empty_usage_batches_do_not_open_a_transaction() -> None: - ledger: Final = AccountingLedger(_no_transaction) - await ledger.insert_usage_batch(()) - assert await ledger.load_usage_batch(()) == () - - -@pytest.mark.asyncio -async def test_conflicting_usage_within_one_batch_is_rejected_before_database_access(event: InboundUsageEvent) -> None: - original: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")})) - 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 diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 21b5bd300c1..c8434007a12 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -3981,140 +3981,3 @@ 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") diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index a426d8f4e14..ef70a193280 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -92,12 +92,6 @@ 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))) @@ -116,24 +110,6 @@ 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 diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 4e113ce03a4..32f6eea91ac 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -7695,13 +7695,11 @@ def test_optional_mcp_metadata_does_not_change_legacy_spend_payload(tool_metadat 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: +def test_spend_payload_preserves_unknown_completion_start(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", @@ -7713,9 +7711,7 @@ def test_spend_payload_preserves_unknown_completion_start(capture_enabled: bool, {"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 []) diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py deleted file mode 100644 index 0c59473114a..00000000000 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ /dev/null @@ -1,1738 +0,0 @@ -import asyncio -import json -from collections.abc import Awaitable, Mapping -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final - -import pytest - -from litellm.accounting import Metric, MetricName, Outcome -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite -from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings -from litellm.types.utils import CallTypes -from tests.unit.proxy.spend_tracking.test_usage_meter import CommitBarrier, _manual_flush, _write - - -@pytest.mark.asyncio -@pytest.mark.parametrize("usage", [None, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}]) -async def test_capture_keeps_unknown_usage_distinct_from_measured_zero(usage: Mapping[str, int] | None) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - captured: Final = capture_callback( - meter, - {"call_type": "acompletion", "litellm_call_id": "caller", "model": "unknown-model", "litellm_params": {}}, - {"id": "provider", "usage": usage}, - timestamp, - timestamp, - ) - await meter.close() - outbound: Final = captured.writes[1].event - assert outbound.metrics == ( - None - if usage is None - else tuple(Metric(name=MetricName(key), value=Decimal(value)) for key, value in usage.items()) - ) - envelope: Final = json.loads(captured.writes[1].envelope) - assert envelope["usage"] == usage - assert envelope["pricing"]["reported_cost"] is None - assert database.committed == ("caller:inbound", "provider") - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "call_type", - [ - call.value - for call in ( - CallTypes.acompletion, - CallTypes.arealtime, - CallTypes.call_mcp_tool, - CallTypes.asend_message, - CallTypes.pass_through, - ) - ], -) -async def test_shared_callback_capture_retains_usage_and_attribution_without_bodies(call_type: str) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - captured: Final = capture_callback( - meter, - { - "call_type": call_type, - "litellm_call_id": "caller", - "model": "test-model", - "messages": "private prompt", - "litellm_params": {"api_key": "private-secret", "metadata": {"user_api_key_team_id": "team"}}, - }, - { - "id": "provider", - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, - "output": "private answer", - }, - timestamp, - timestamp, - ) - await meter.close() - envelope: Final = captured.writes[1].envelope - assert all(secret not in envelope for secret in ("private prompt", "private-secret", "private answer")) - event: Final = captured.writes[1].event - assert event.call_type.value == call_type - assert event.metrics == ( - Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(3)), - Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), - Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(5)), - ) - assert any(dimension.value == "team" for dimension in event.dimensions) - - -@pytest.mark.asyncio -async def test_cache_duplicate_uses_caller_identity_despite_timestamped_diagnostic_ids() -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = CommitBarrier() - database.release.set() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - kwargs: Final = { - "call_type": "acompletion", - "litellm_call_id": "caller", - "model": "test-model", - "cache_hit": True, - "litellm_params": {}, - } - response: Final = {"id": "cached-provider-id", "usage": {"total_tokens": 5}} - first: Final = capture_callback(meter, kwargs, response, timestamp, timestamp) - second: Final = capture_callback(meter, kwargs, response, timestamp, timestamp) - await meter.close() - assert first.writes == second.writes - assert first.writes[1].event.event_id == "caller_cache_hit" - assert json.loads(first.writes[1].envelope)["charge_eligible"] is False - - -NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - - -class StoredUsage: - def __init__(self) -> None: - self.rows: tuple[UsageWrite, ...] = () - - async def load(self, identities: tuple[str, ...]) -> tuple[UsageWrite, ...]: - return tuple(row for row in self.rows if row.event.event_id in identities) - - async def insert(self, writes: tuple[UsageWrite, ...]) -> None: - existing: Final = {row.event.event_id: row for row in self.rows} - if any(existing.get(row.event.event_id, row) != row for row in writes): - raise AccountingConflictError("Same identity has different content") - unique: Final = {row.event.event_id: row for row in writes} - if any(unique[row.event.event_id] != row for row in writes): - raise AccountingConflictError("Submission contains conflicting content") - self.rows += tuple(row for identity, row in unique.items() if identity not in existing) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("text", ["parent\x00value", "pärent-✓-日本", r"\u0000"]) -async def test_capture_normalizes_database_text_without_poisoning_neighbors(text: str) -> None: - from litellm.proxy.spend_tracking.usage_capture import PricingCapture, capture_callback - - database: Final = StoredUsage() - batches: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - assert all("\x00" not in (json.loads(write.envelope).get("parent_request_id") or "") for write in writes) - batches.put_nowait(writes) - await database.insert(writes) - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(), wait_for_batch=_manual_flush) - nested: Final = json.loads("[" * 100 + "3" + "]" * 100) - kwargs: Final = { - "call_type": "acompletion", - "litellm_call_id": "caller", - "model": "capture-model", - "search_context_size": text, - "_accounting_pricing_snapshot": PricingCapture( - "{}", "{}", json.dumps({"request": {"search_context_size": text}, "nested": nested}) - ), - "litellm_params": {"metadata": {"parent_request_id": text}}, - } - response: Final = {"id": "provider", "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}} - captured: Final = capture_callback(meter, kwargs, response, NOW, NOW) - neighbor: Final = capture_callback( - meter, - {**kwargs, "litellm_call_id": "neighbor", "litellm_params": {}}, - {**response, "id": "neighbor-provider"}, - NOW, - NOW, - ) - try: - await meter.flush() - await asyncio.gather(captured.receipt, neighbor.receipt) - assert batches.qsize() == 1 - assert len(batches.get_nowait()) == 4 - envelope: Final = json.loads(captured.writes[-1].envelope) - assert envelope["parent_request_id"] == text.replace("\x00", "") - assert envelope["pricing"]["request"]["search_context_size"] == text.replace("\x00", "") - assert envelope["pricing"]["nested"] == nested - assert envelope["event"] == captured.writes[-1].event.model_dump(mode="json") - assert envelope["usage"] == response["usage"] - assert kwargs["litellm_params"] == {"metadata": {"parent_request_id": text}} - repeated: Final = capture_callback(meter, kwargs, response, NOW, NOW) - await meter.flush() - await repeated.receipt - assert repeated.writes == captured.writes - assert len(database.rows) == 4 - finally: - await meter.stop(0) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("invalid_field", ["identity", "context"]) -async def test_invalid_capture_text_is_rejected_before_admission(invalid_field: str) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - kwargs: Final = {"call_type": "acompletion", "litellm_call_id": "caller", "litellm_params": {}} - try: - accepted: Final = capture_callback(meter, kwargs, {"id": "provider"}, NOW, NOW) - invalid: Final = ( - {**kwargs, "litellm_call_id": "caller\x00"} - if invalid_field == "identity" - else {**kwargs, "litellm_params": {"metadata": {"parent_request_id": "parent\ud800"}}} - ) - with pytest.raises((UsageCaptureError, ValueError)): - capture_callback(meter, invalid, {"id": "other"}, NOW, NOW) - assert meter.pending_records == 2 - await meter.flush() - await accepted.receipt - assert tuple(write.event.event_id for write in database.rows) == ("caller:inbound", "provider") - finally: - await meter.stop(0) - - -def _logger(call_type: str = "acompletion", *, stream: bool = False) -> Logging: - from litellm.litellm_core_utils.litellm_logging import Logging - - logger: Final = Logging( - model="capture-model", - messages=[], - stream=stream, - call_type=call_type, - start_time=NOW, - litellm_call_id="caller", - function_id="function", - ) - logger.update_environment_variables( - model="capture-model", custom_llm_provider="openai", litellm_params={}, optional_params={} - ) - return logger - - -@pytest.mark.asyncio -async def test_terminal_retrieval_reuses_accepted_snapshot_but_rejects_changed_usage() -> None: - from datetime import timedelta - - from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY - from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, capture_callback - from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - - database: Final = StoredUsage() - writer: Final = UsageBatchWriter(database.insert, database.load) - meter: Final = DurableUsageMeter(writer.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - kwargs: Final = { - "call_type": "aget_responses", - "litellm_call_id": "poll-one", - "model": "capture-model", - "response_cost": 2.0, - "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, - } - response: Final = { - "id": "job", - "status": "completed", - "background": True, - "created_at": NOW.timestamp(), - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - } - first: Final = capture_callback(meter, kwargs, response, NOW, NOW) - await meter.flush() - later: Final = capture_callback( - meter, {**kwargs, "litellm_call_id": "poll-two", "response_cost": 9.0}, response, NOW, NOW + timedelta(hours=2) - ) - await meter.flush() - await later.receipt - assert database.rows == first.writes - assert len(database.rows) == 1 - assert database.rows[0].event.event_id == "job_background_cost" - changed: Final = capture_callback(meter, kwargs, {**response, "usage": {"total_tokens": 8}}, NOW, NOW) - await meter.close() - with pytest.raises(UsageCaptureError, match="identity"): - await changed.receipt - assert database.rows == first.writes - - -@pytest.mark.asyncio -@pytest.mark.parametrize("terminal", [False, True]) -async def test_competing_writers_reuse_only_terminal_pricing_snapshots(terminal: bool) -> None: - from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY - from litellm.proxy.spend_tracking.usage_capture import UsageBatchWriter, capture_callback - from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - - database: Final = StoredUsage() - attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - competing: Final = asyncio.Event() - - async def insert(writes: tuple[UsageWrite, ...]) -> None: - attempts.put_nowait(writes) - if attempts.qsize() >= 2: - competing.set() - await competing.wait() - await database.insert(writes) - - meters: Final = tuple( - DurableUsageMeter(UsageBatchWriter(insert, database.load).write, UsageCaptureSettings()) for _ in range(2) - ) - kwargs: Final = { - "call_type": "aget_responses" if terminal else "acompletion", - "litellm_call_id": "caller", - "model": "capture-model", - "litellm_params": { - "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} - if terminal - else {} - }, - } - response: Final = { - "id": "job", - "status": "completed", - "background": terminal, - "created_at": NOW.timestamp(), - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - } - captured: Final = tuple( - capture_callback(meter, {**kwargs, "response_cost": price}, response, NOW, NOW) - for meter, price in zip(meters, (2.0, 9.0), strict=True) - ) - try: - await asyncio.gather(*(meter.flush() for meter in meters)) - results: Final = await asyncio.gather(*(capture.receipt for capture in captured), return_exceptions=True) - assert sum(result is None for result in results) == (2 if terminal else 1) - assert sum(isinstance(result, UsageCaptureError) for result in results) == (0 if terminal else 1) - assert database.rows in tuple(capture.writes for capture in captured) - assert len(database.rows) == (1 if terminal else 2) - assert attempts.qsize() == (3 if terminal else 2) - finally: - await asyncio.gather(*(meter.stop(0) for meter in meters)) - - -@pytest.mark.asyncio -async def test_batch_poll_has_distinct_identity_and_only_terminal_usage_is_chargeable() -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - kwargs: Final = { - "call_type": "aretrieve_batch", - "litellm_call_id": "poll", - "model": "capture-model", - "litellm_params": {}, - } - pending: Final = capture_callback(meter, kwargs, {"id": "batch", "status": "in_progress"}, NOW, NOW) - completed: Final = capture_callback( - meter, kwargs, {"id": "batch", "status": "completed", "usage": {"total_tokens": 9}}, NOW, NOW - ) - await meter.close() - assert pending.writes[-1].event.event_id == "poll:status" - assert json.loads(pending.writes[-1].envelope)["charge_eligible"] is False - assert completed.writes[-1].event.event_id == "batch_batch_cost" - assert json.loads(completed.writes[-1].envelope)["charge_eligible"] is True - assert len(completed.writes) == 1 - - -@pytest.mark.asyncio -async def test_internal_capture_does_not_claim_reused_parent_logger(monkeypatch: pytest.MonkeyPatch) -> None: - import litellm - from litellm._internal_context import internal_sub_call - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.utils import ModelResponse, Usage - - monkeypatch.setitem( - litellm.model_cost, - "capture-model", - { - "input_cost_per_token": 0.1, - "output_cost_per_token": 0.2, - "litellm_provider": "openai", - "mode": "chat", - }, - ) - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger() - response: Final = ModelResponse( - id="response", model="capture-model", usage=Usage(prompt_tokens=2, completion_tokens=1, total_tokens=3) - ) - with internal_sub_call("caller"): - capture_logging_usage(meter, logger, response, NOW, NOW) - monkeypatch.setitem( - litellm.model_cost, - "capture-model", - { - "input_cost_per_token": 0.7, - "output_cost_per_token": 0.9, - "litellm_provider": "openai", - "mode": "chat", - }, - ) - litellm.register_model({"capture-model": litellm.model_cost["capture-model"]}) - capture_logging_usage(meter, logger, response, NOW, NOW) - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - assert tuple(row.event.event_id for row in database.rows) == ("response:included", "caller:inbound", "response") - child: Final = json.loads(database.rows[0].envelope) - parent: Final = json.loads(database.rows[-1].envelope) - assert (child["ownership"], child["charge_eligible"], child["parent_request_id"]) == ( - "included_in_parent", - False, - "caller", - ) - assert parent["charge_eligible"] is True - assert child["pricing"]["model_map_information"]["model_map_value"]["input_cost_per_token"] == 0.1 - assert parent["pricing"]["model_map_information"]["model_map_value"]["input_cost_per_token"] == 0.7 - - -@pytest.mark.asyncio -async def test_stream_capture_ignores_chunks_and_preserves_terminal_usage_extras() -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.utils import ModelResponse, ModelResponseStream, Usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(stream=True) - capture_logging_usage(meter, logger, ModelResponseStream(id="chunk"), NOW, NOW) - assert meter.pending_records == 0 - capture_logging_usage( - meter, logger, ModelResponse(id="stream", usage=Usage(prompt_tokens=3, provider_units=4)), NOW, NOW - ) - await meter.close() - assert json.loads(database.rows[-1].envelope)["usage"]["provider_units"] == 4 - assert database.rows[-1].event.event_id == "stream" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("call_type", ["acompletion", "atext_completion", "aembedding", "aresponses"]) -async def test_supported_sync_ack_waits_for_commit(call_type: str, monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy import proxy_server - from litellm.proxy._types import ConfigGeneralSettings - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - deadlines: Final = asyncio.Queue[float | None]() - original_wait: Final = asyncio.wait_for - - async def wait_with_deadline(awaitable: Awaitable[object], timeout: float | None) -> object: - deadlines.put_nowait(timeout) - return await original_wait(awaitable, timeout) - - monkeypatch.setattr(asyncio, "wait_for", wait_with_deadline) - response: Final = asyncio.create_task( - capture_proxy_response( - _logger(call_type), - ModelResponse(id="provider", usage=Usage(prompt_tokens=3)), - ) - ) - try: - await database.entered.wait() - deadline: Final = deadlines.get_nowait() - assert deadline == ConfigGeneralSettings().accounting_sync_ack_timeout_seconds - assert deadline is not None and deadline > UsageCaptureSettings().flush_interval_seconds + 10 - assert not response.done() - assert database.committed == () - database.release.set() - await response - assert database.committed == ("caller:inbound", "provider") - finally: - database.release.set() - await runtime.stop(1) - await response - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure_policy", ["fail_closed", "log"]) -async def test_sync_ack_commit_failure_follows_response_policy( - failure_policy: str, monkeypatch: pytest.MonkeyPatch -) -> None: - from fastapi import HTTPException - - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier(failure=ValueError("invalid envelope")) - database.release.set() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - meter: Final = runtime.meter - assert meter is not None - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr( - proxy_server, - "general_settings", - {"accounting_sync_ack": True, "accounting_sync_ack_failure_policy": failure_policy}, - ) - logger: Final = _logger() - response: Final = ModelResponse(id="provider", usage=Usage(prompt_tokens=4)) - if failure_policy == "fail_closed": - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, response) - assert error.value.status_code == 503 - else: - await capture_proxy_response(logger, response) - await runtime.stop(1.0) - assert database.committed == () - assert meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_failed_attempt_keeps_its_usage_without_claiming_the_callers_success() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - for _ in range(2): - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.4) - logger.carry_failed_attempt_usage() - logger.capture_usage_before_logging( - ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, total_tokens=3)), NOW, NOW - ) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ( - "caller-attempt-0", - "caller-attempt-1", - "caller:inbound", - "fallback", - ) - failed: Final = json.loads(database.rows[0].envelope) - assert failed["usage"]["total_tokens"] == 7 - assert failed["parent_request_id"] == "caller" - assert failed["request_count_eligible"] is False - assert failed["charge_eligible"] is True - assert failed["event"]["retry_of_attempt_id"] is None - assert failed["event"]["outcome"] == "failed" - assert failed["pricing"]["reported_cost"] == 0.4 - assert all(json.loads(row.envelope)["included_attempt_ids"] == [] for row in database.rows[:2]) - successful: Final = json.loads(database.rows[-1].envelope) - assert successful["related_attempt_ids"] == ["caller-attempt-0", "caller-attempt-1"] - assert successful["included_attempt_ids"] == [] - assert successful["usage"]["total_tokens"] == 3 - assert successful["charge_eligible"] is True - - -@pytest.mark.asyncio -async def test_success_that_bills_a_failed_follow_up_names_that_attempt_as_included() -> None: - from litellm.litellm_core_utils.hidden_params import set_hidden_param - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - fallback: Final = ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, completion_tokens=1, total_tokens=4)) - set_hidden_param(fallback, "response_cost", 0.1) - try: - logger.record_served_rounds_for_failure(((Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.4),)) - logger.carry_failed_attempt_usage() - logger.prepare_response_cost(fallback, NOW) - logger.capture_usage_before_logging(fallback, NOW, NOW) - finally: - await runtime.stop(1) - envelopes: Final = {row.event.event_id: json.loads(row.envelope) for row in database.rows} - assert envelopes["caller-attempt-0"]["usage"]["total_tokens"] == 7 - assert envelopes["caller-attempt-0"]["charge_eligible"] is True - assert envelopes["fallback"]["usage"]["total_tokens"] == 4 + 7 - assert envelopes["fallback"]["included_attempt_ids"] == ["caller-attempt-0"] - assert database.rows[2].event.outcome is Outcome.SUCCEEDED - - -@pytest.mark.asyncio -@pytest.mark.parametrize("estimated", [False, True]) -async def test_cancelled_usage_keeps_measured_or_estimated_provenance(estimated: bool) -> None: - from fastapi import HTTPException - - from litellm.litellm_core_utils.litellm_logging import UsageObservation - from litellm.proxy.spend_tracking.usage_capture import capture_callback - from litellm.types.utils import Usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - capture_callback( - meter, - { - "litellm_call_id": "cancelled", - "call_type": "acompletion", - "model": "capture-model", - "litellm_params": {}, - "combined_usage_object": Usage(prompt_tokens=4, total_tokens=4), - "_accounting_usage_observed": UsageObservation.ESTIMATED if estimated else UsageObservation.KNOWN, - }, - HTTPException(499, "Client disconnected"), - NOW, - NOW, - ) - await meter.close() - payload: Final = json.loads(database.rows[-1].envelope) - assert payload["event"]["outcome"] == "cancelled" - assert payload["usage"]["total_tokens"] == 4 - assert payload["usage_source"] == ("estimated" if estimated else "measured") - assert payload["pricing"]["reported_cost"] is None - - -@pytest.mark.asyncio -async def test_realtime_capture_keeps_metered_units_and_provider_usage_without_frames() -> None: - from litellm.constants import REALTIME_METERED_USAGE_KEY - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.llms.xai import RealtimeMeteredUsage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(CallTypes.arealtime.value) - logger.model_call_details[REALTIME_METERED_USAGE_KEY] = RealtimeMeteredUsage(True, 1.5, 2.5, 3) - capture_logging_usage( - meter, - logger, - [ - { - "type": "response.done", - "response": { - "id": "turn", - "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, - "output": "private realtime transcript", - }, - } - ], - NOW, - NOW, - ) - await meter.close() - payload: Final = json.loads(database.rows[-1].envelope) - assert payload["usage"]["total_tokens"] == 7 - assert payload["usage"]["metered_usage"] == { - "push_to_talk": True, - "audio_seconds_sent": 1.5, - "audio_seconds_received": 2.5, - "text_messages": 3, - } - assert "private realtime transcript" not in database.rows[-1].envelope - - -@pytest.mark.asyncio -async def test_default_http_capture_returns_before_database_commit(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {}) - try: - await capture_proxy_response( - _logger(), - ModelResponse(id="asynchronous", usage=Usage(prompt_tokens=1)), - ) - await database.entered.wait() - assert database.committed == () - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == ("caller:inbound", "asynchronous") - - -@pytest.mark.asyncio -async def test_deferred_capture_freezes_pricing_before_configuration_changes(monkeypatch: pytest.MonkeyPatch) -> None: - import litellm - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_logging_usage - from litellm.types.utils import ModelResponse, Usage - from litellm.utils import dispatch_success_logging - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - logger.defer_async_logging = True - logger.custom_cost_per_second = 0.25 - monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.1}) - response: Final = ModelResponse(id="deferred", usage=Usage(prompt_tokens=3)) - try: - dispatch_success_logging(logger, response, NOW, NOW, False, False) - assert database.rows == () - logger.custom_cost_per_second = 9 - monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.9}) - assert runtime.meter is not None - capture_logging_usage(runtime.meter, logger, response, NOW, NOW) - await runtime.meter.flush() - finally: - await runtime.stop(1) - pricing: Final = json.loads(database.rows[-1].envelope)["pricing"] - assert pricing["logging"]["custom_cost_per_second"] == 0.25 - assert pricing["adjustments"]["discount"] == 0.1 - - -@pytest.mark.asyncio -async def test_synchronous_internal_capture_reserves_capacity_before_scheduling() -> None: - from litellm._internal_context import internal_sub_call - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(max_pending_records=2, batch_size=2)) - logger: Final = _logger() - - def capture_child() -> None: - with internal_sub_call("caller"): - logger.capture_usage_before_logging(ModelResponse(id="child", usage=Usage(prompt_tokens=2)), NOW, NOW) - - try: - await asyncio.to_thread(capture_child) - assert runtime.meter is not None - assert runtime.meter.pending_records == 1 - receipt: Final = await asyncio.to_thread(runtime.meter.submit, (_write("second"),)) - await database.entered.wait() - with pytest.raises(UsageCaptureError, match="capacity"): - await asyncio.to_thread(runtime.meter.submit, (_write("overflow"),)) - database.release.set() - await receipt - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == ("child:included", "second") - - -@pytest.mark.asyncio -@pytest.mark.parametrize("has_metered_usage", [False, True]) -async def test_realtime_without_provider_token_usage_does_not_invent_zero(has_metered_usage: bool) -> None: - from litellm.constants import REALTIME_METERED_USAGE_KEY - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - from litellm.types.llms.xai import RealtimeMeteredUsage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(CallTypes.arealtime.value) - if has_metered_usage: - logger.model_call_details[REALTIME_METERED_USAGE_KEY] = RealtimeMeteredUsage(True, 0, 0, 0) - capture_logging_usage(meter, logger, [], NOW, NOW) - await meter.close() - usage: Final = json.loads(database.rows[-1].envelope)["usage"] - assert usage == ( - { - "metered_usage": { - "push_to_talk": True, - "audio_seconds_sent": 0, - "audio_seconds_received": 0, - "text_messages": 0, - } - } - if has_metered_usage - else None - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("committed_before_failure", [False, True]) -async def test_database_recovery_retries_frozen_usage_without_callback_or_duplicate_rows( - committed_before_failure: bool, -) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - failures: Final = asyncio.Queue[None]() - failures.put_nowait(None) - retrying: Final = asyncio.Event() - recovered: Final = asyncio.Event() - delays: Final = asyncio.Queue[float]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - if not failures.empty(): - failures.get_nowait() - if committed_before_failure: - await database.insert(writes) - raise ConnectionError("database unavailable or acknowledgement lost") - await database.insert(writes) - - async def retry_delay(delay: float) -> None: - delays.put_nowait(delay) - retrying.set() - await recovered.wait() - - meter: Final = DurableUsageMeter( - persist, UsageCaptureSettings(max_pending_records=2), wait_for_batch=_manual_flush, retry_delay=retry_delay - ) - kwargs: Final = { - "call_type": "acompletion", - "litellm_call_id": "caller", - "model": "capture-model", - "litellm_params": {}, - } - response: Final = {"id": "provider", "usage": {"total_tokens": 3}} - first: Final = capture_callback(meter, kwargs, response, NOW, NOW) - reserved: Final = meter.pending_bytes - flushing: Final = asyncio.create_task(meter.flush()) - await retrying.wait() - assert not first.receipt.done() - assert (meter.pending_records, meter.pending_bytes) == (2, reserved) - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("overflow"),)) - recovered.set() - await flushing - await first.receipt - duplicate: Final = capture_callback(meter, {**kwargs, "_accounting_capture": first}, response, NOW, NOW) - await meter.close() - assert duplicate is first - assert database.rows == first.writes - assert delays.get_nowait() == 0.2 - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -async def test_cut_stream_keeps_original_model_usage_as_part_of_the_parent_aggregate() -> None: - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.router import Router - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger(stream=True) - stream: Final = CustomStreamWrapper(completion_stream=iter(()), model="capture-model", logging_obj=logger) - partial: Final = ModelResponse( - id="cut", model="capture-model", usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5) - ) - try: - Router._price_cut_partial(stream, partial) - finally: - await runtime.stop(1) - assert len(database.rows) == 1 - payload: Final = json.loads(database.rows[0].envelope) - assert payload["event"]["event_id"] == "cut:included" - assert payload["parent_request_id"] == "caller" - assert payload["charge_eligible"] is False - assert payload["usage"]["total_tokens"] == 5 - assert payload["pricing"]["model_map_information"]["model_map_key"] == "capture-model" - - -@pytest.mark.asyncio -async def test_direct_mcp_logging_captures_tool_pricing_and_identity_without_arguments() -> None: - from mcp.types import CallToolResult, TextContent - - from litellm.accounting import DimensionName - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger(CallTypes.call_mcp_tool.value) - logger.model_call_details["mcp_tool_call_metadata"] = { - "name": "lookup", - "namespaced_tool_name": "server.lookup", - "arguments": {"query": "private query"}, - "mcp_server_cost_info": {"default_cost_per_query": 0.2, "tool_name_to_cost_per_query": {"lookup": 0.3}}, - } - try: - await logger.async_success_handler( - CallToolResult(content=[TextContent(type="text", text="private result")]), NOW, NOW - ) - finally: - await runtime.stop(1) - outbound: Final = database.rows[-1] - assert any(d.name is DimensionName.TOOL and d.value == "server.lookup" for d in outbound.event.dimensions) - payload: Final = json.loads(outbound.envelope) - assert payload["pricing"]["tool"] == {"name": "lookup", "cost_per_query": 0.3} - assert payload["usage"] is None - assert all(secret not in outbound.envelope for secret in ("private query", "private result")) - - -@pytest.mark.asyncio -async def test_callback_overload_is_visible_without_suppressing_legacy_tracking( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, captured_callback_kwargs - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(max_pending_records=2, batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - kwargs: Final = {"call_type": "acompletion", "litellm_call_id": "one", "litellm_params": {}} - try: - for start, end in ((None, NOW), (NOW, None)): - missing_timestamp: Final = await captured_callback_kwargs(kwargs, {"id": "one"}, start, end) - assert missing_timestamp == kwargs - assert "Usage capture requires callback start and end timestamps" in caplog.text - admitted: Final = await captured_callback_kwargs(kwargs, {"id": "one"}, NOW, NOW) - await database.entered.wait() - rejected: Final = await captured_callback_kwargs({**kwargs, "litellm_call_id": "two"}, {"id": "two"}, NOW, NOW) - assert "_accounting_capture" in admitted - assert "_accounting_capture" not in rejected - assert "Usage capture failed before legacy cost tracking" in caplog.text - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == ("one:inbound", "one") - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("call_type", "streaming", "background"), - [("acompletion", True, False), ("aresponses", False, True), ("aocr", False, False)], -) -@pytest.mark.parametrize("overloaded", [False, True]) -async def test_async_capture_preserves_response_without_commit_ack( - call_type: str, streaming: bool, background: bool, overloaded: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr( - proxy_server, - "general_settings", - { - "accounting_sync_ack": True, - "accounting_sync_ack_timeout_seconds": 0, - }, - ) - meter: Final = runtime.meter - assert meter is not None - if overloaded: - meter.submit((_write("occupied-one"), _write("occupied-two"))) - try: - await capture_proxy_response( - _logger(call_type, stream=streaming), - ModelResponse(id="provider", usage=Usage(prompt_tokens=3)), - background=background, - ) - assert database.committed == () - assert meter.pending_records == 2 - finally: - database.release.set() - await runtime.stop(1) - assert database.committed == (("occupied-one", "occupied-two") if overloaded else ("caller:inbound", "provider")) - - -@pytest.mark.asyncio -async def test_new_connection_turn_does_not_reuse_the_previous_turns_capture() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - logger.capture_usage_before_logging(ModelResponse(id="first", usage=Usage(prompt_tokens=2)), NOW, NOW) - next_turn: Final = logger.copy_for_new_call("second-caller", NOW) - next_turn.capture_usage_before_logging(ModelResponse(id="second", usage=Usage(prompt_tokens=4)), NOW, NOW) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ( - "caller:inbound", - "first", - "second-caller:inbound", - "second", - ) - assert json.loads(database.rows[-1].envelope)["usage"]["prompt_tokens"] == 4 - - -@pytest.mark.asyncio -async def test_billed_failure_is_captured_before_legacy_persistence_gates(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "prisma_client", None) - request: Final = { - "call_type": "acompletion", - "litellm_call_id": "failed", - "model": "capture-model", - "start_time": NOW, - "combined_usage_object": Usage(prompt_tokens=3, total_tokens=3), - "response_cost": 0.2, - "litellm_params": {}, - } - try: - await ProxyDBLogger().async_post_call_failure_hook(request, RuntimeError("provider failed"), UserAPIKeyAuth()) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("failed:inbound", "failed") - payload: Final = json.loads(database.rows[-1].envelope) - assert payload["event"]["outcome"] == "failed" - assert payload["usage"]["total_tokens"] == 3 - assert payload["charge_eligible"] is True - assert payload["pricing"]["reported_cost"] == 0.2 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("pages", [None, 0, 3]) -async def test_prepared_ocr_capture_preserves_reported_pages(pages: int | None) -> None: - from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger("aocr") - response: Final = OCRResponse(pages=[], model="capture-model", usage_info=OCRUsageInfo(pages_processed=pages)) - logger.prepare_response_cost(response, NOW) - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - envelope: Final = json.loads(database.rows[-1].envelope) - assert (None if envelope["usage"] is None else envelope["usage"]["pages_processed"]) == pages - assert envelope["usage_source"] == ("unknown" if pages is None else "measured") - assert response.usage_info is not None and response.usage_info.pages_processed == pages - - -@pytest.mark.asyncio -@pytest.mark.parametrize("measured", [None, {"total_tokens": 0}, {"total_tokens": 3}]) -async def test_batch_capture_provenance_tracks_observations_not_envelope_presence( - measured: Mapping[str, int] | None, -) -> None: - from litellm.proxy.spend_tracking.usage_capture import capture_callback - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - usage: Final = { - "total_tokens": 0 if measured is None else measured["total_tokens"], - "batch_usage_by_model": {"missing-model": None, "other-model": measured}, - "batch_unknown_usage_requests": 2 if measured is None else 1, - } - capture_callback( - meter, - {"call_type": "aretrieve_batch", "litellm_call_id": "caller", "litellm_params": {}}, - {"id": "batch", "usage": usage}, - NOW, - NOW, - ) - await meter.close() - envelope: Final = json.loads(database.rows[-1].envelope) - assert envelope["usage_source"] == ("unknown" if measured is None else "measured") - assert envelope["usage"]["batch_usage_by_model"] == usage["batch_usage_by_model"] - assert database.rows[-1].event.metrics == ( - None if measured is None else (Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(measured["total_tokens"])),) - ) - - -@pytest.mark.asyncio -async def test_ack_admission_failure_does_not_turn_provider_success_into_failed_usage( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from fastapi import HTTPException - - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - monkeypatch.setattr(proxy_server, "prisma_client", None) - meter: Final = runtime.meter - assert meter is not None - logger: Final = _logger() - try: - meter.submit((_write("occupied-one"), _write("occupied-two"))) - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, ModelResponse(id="caller", usage=Usage(prompt_tokens=3))) - assert error.value.status_code == 503 - await meter.flush() - request: Final = {**logger.model_call_details, "litellm_logging_obj": logger, "start_time": NOW} - await ProxyDBLogger().async_post_call_failure_hook(request, error.value, UserAPIKeyAuth()) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("occupied-one", "occupied-two") - - -@pytest.mark.asyncio -async def test_response_ack_timeout_preserves_success_and_later_commits(monkeypatch: pytest.MonkeyPatch) -> None: - from fastapi import HTTPException - - from litellm.proxy import proxy_server - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_proxy_response - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=2)) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": True}) - logger: Final = _logger() - - async def expired_deadline(future: object, timeout: float) -> None: - raise TimeoutError("controlled response deadline") - - try: - with monkeypatch.context() as clock: - clock.setattr(asyncio, "wait_for", expired_deadline) - with pytest.raises(HTTPException) as error: - await capture_proxy_response(logger, ModelResponse(id="provider", usage=Usage(prompt_tokens=3))) - assert error.value.status_code == 503 - await database.entered.wait() - captured: Final = logger.model_call_details["_accounting_capture"] - request: Final = {**logger.model_call_details, "litellm_logging_obj": logger, "start_time": NOW} - await ProxyDBLogger().async_post_call_failure_hook(request, error.value, UserAPIKeyAuth()) - database.release.set() - await captured.receipt - assert tuple(row.event.outcome for row in captured.writes) == (Outcome.SUCCEEDED, Outcome.SUCCEEDED) - assert database.committed == ("caller:inbound", "provider") - assert database.transactions == 1 - finally: - await runtime.stop(1) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("call_type", ["generate_content", "agenerate_content"]) -@pytest.mark.parametrize("tokens", [None, 0, 7]) -async def test_prepared_native_google_capture_preserves_reported_usage( - call_type: str, tokens: int | None, monkeypatch: pytest.MonkeyPatch -) -> None: - import httpx - - import litellm - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(call_type) - monkeypatch.setitem( - litellm.model_cost, - "capture-model", - {"input_cost_per_token": 0.1, "output_cost_per_token": 0.2, "litellm_provider": "openai", "mode": "chat"}, - ) - reported: Final = ( - {"trafficType": "ON_DEMAND"} - if tokens is None - else { - "promptTokenCount": tokens, - "candidatesTokenCount": tokens, - "totalTokenCount": 2 * tokens, - "cachedContentTokenCount": tokens, - "promptTokensDetails": [{"modality": "DOCUMENT", "tokenCount": tokens}], - "cacheTokensDetails": [{"modality": "document", "token_count": tokens}], - } - ) - response: Final = { - "responseId": "native-google", - "candidates": [{"content": {"parts": [{"text": "answer"}], "role": "model"}, "finishReason": "STOP"}], - "usageMetadata": reported, - } - native_response: Final = httpx.Response(200, json=response) - logger.model_call_details["httpx_response"] = native_response - prepared: Final = logger.prepare_response_cost(response, NOW) - assert prepared is not None - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - envelope: Final = json.loads(database.rows[-1].envelope) - assert (None if envelope["usage"] is None else envelope["usage"]["total_tokens"]) == ( - None if tokens is None else 2 * tokens - ) - assert (None if envelope["usage"] is None else envelope["usage"]["prompt_tokens_details"]["text_tokens"]) == tokens - assert ( - None - if envelope["usage"] is None - else envelope["usage"]["prompt_tokens_details"]["cached_tokens_details"]["text_tokens"] - ) == tokens - assert envelope["usage_source"] == ("unknown" if tokens is None else "measured") - assert response == native_response.json() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("rejected", [False, True]) -async def test_retry_links_only_admitted_usage_and_preserves_later_measurement(rejected: bool) -> None: - from litellm.litellm_core_utils.litellm_logging import UsageObservation - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings(max_pending_records=4)) - logger: Final = _logger() - assert runtime.meter is not None - try: - if rejected: - runtime.meter.submit(tuple(_write(f"occupied-{index}") for index in range(4))) - logger.model_call_details["_accounting_usage_observed"] = UsageObservation.UNKNOWN - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - await runtime.meter.flush() - logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - finally: - await runtime.stop(1) - successful: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "success")) - assert successful.get("related_attempt_ids", successful["included_attempt_ids"]) == ( - [] if rejected else ["caller-attempt-0"] - ) - assert successful["usage"]["prompt_tokens"] == 3 - if not rejected: - failed: Final = json.loads( - next(row.envelope for row in database.rows if row.event.event_id == "caller-attempt-0") - ) - assert failed["usage"]["prompt_tokens"] == 5 - assert failed["usage_source"] == "measured" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("current_tokens", [None, 0, 3]) -async def test_final_failure_captures_current_attempt_without_legacy_retry_aggregate( - current_tokens: int | None, -) -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.utils import _failure_fields_to_lift - from litellm.types.utils import Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - if current_tokens is not None: - logger.record_partial_usage_for_failure( - Usage(prompt_tokens=current_tokens, total_tokens=current_tokens), 0.2 - ) - lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) - assert lifted["combined_usage_object"].total_tokens == 5 + (current_tokens or 0) - assert runtime.meter is not None - capture_callback( - runtime.meter, - {**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"}, - None, - NOW, - NOW, - ) - finally: - await runtime.stop(1) - final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller")) - assert (None if final["usage"] is None else final["usage"]["total_tokens"]) == current_tokens - assert final["usage_source"] == ("unknown" if current_tokens is None else "measured") - assert final["pricing"]["reported_cost"] == (None if current_tokens is None else 0.2) - assert final["included_attempt_ids"] == [] - assert final["related_attempt_ids"] == ["caller-attempt-0"] - - -@pytest.mark.asyncio -async def test_estimated_final_attempt_keeps_its_provenance_after_measured_retry() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback - from litellm.proxy.utils import _failure_fields_to_lift - from litellm.types.utils import Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - logger.messages = [{"role": "user", "content": "hello"}] - logger.model_call_details["messages"] = logger.messages - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - logger.model_call_details["current_attempt_api_call_start_time"] = NOW - lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True) - assert runtime.meter is not None - capture_callback(runtime.meter, {**logger.model_call_details, **lifted}, None, NOW, NOW) - finally: - await runtime.stop(1) - final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller")) - assert final["usage_source"] == "estimated" - assert final["usage"]["total_tokens"] > 0 - assert lifted["combined_usage_object"].total_tokens == final["usage"]["total_tokens"] + 5 - assert final["pricing"]["reported_cost"] == 0 - - -@pytest.mark.asyncio -async def test_parent_receipt_cannot_acknowledge_when_a_linked_attempt_conflicts() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - assert runtime.meter is not None - await runtime.meter.flush() - other: Final = _logger() - other.record_partial_usage_for_failure(Usage(prompt_tokens=8), 0.8) - other.carry_failed_attempt_usage() - other.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - captured: Final = other.model_call_details["_accounting_capture"] - await runtime.meter.flush() - with pytest.raises(UsageCaptureError, match="dependency"): - await captured.receipt - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("caller-attempt-0",) - - -@pytest.mark.asyncio -async def test_unserializable_observation_does_not_interrupt_retry_or_legacy_usage() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = StoredUsage() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.insert, UsageCaptureSettings()) - logger: Final = _logger() - usage: Final = Usage(prompt_tokens=5, total_tokens=5, opaque_usage=object()) - try: - logger.record_partial_usage_for_failure(usage, 0.4) - assert logger.model_call_details["combined_usage_object"] is usage - logger.carry_failed_attempt_usage() - assert logger.model_call_details["failed_attempts_usage"] == ((usage, 0.4),) - logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - finally: - await runtime.stop(1) - assert tuple(row.event.event_id for row in database.rows) == ("caller:inbound", "success") - assert json.loads(database.rows[-1].envelope)["related_attempt_ids"] == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("prior_committed", [False, True]) -@pytest.mark.parametrize("batch_size", [2, 3]) -async def test_retry_dependencies_do_not_resubmit_usage_or_expand_final_submission( - prior_committed: bool, batch_size: int -) -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.types.utils import ModelResponse, Usage - - database: Final = CommitBarrier() - database.release.set() - runtime: Final = UsageCaptureRuntime() - runtime.start(database.write, UsageCaptureSettings(batch_size=batch_size, max_pending_records=3)) - logger: Final = _logger() - try: - logger.record_partial_usage_for_failure(Usage(prompt_tokens=5), 0.4) - logger.carry_failed_attempt_usage() - assert runtime.meter is not None - if prior_committed: - await runtime.meter.flush() - logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW) - captured: Final = logger.model_call_details.get("_accounting_capture") - assert captured is not None, "Final usage must be admitted independently of prior usage batch size" - await runtime.meter.close() - await captured.receipt - assert runtime.meter.pending_records == 0 - finally: - await runtime.stop(1) - assert database.committed == ("caller-attempt-0", "caller:inbound", "success") - final: Final = json.loads(captured.writes[-1].envelope) - assert final["related_attempt_ids"] == ["caller-attempt-0"] - assert final["usage"]["prompt_tokens"] == 3 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("prepared", [False, True]) -@pytest.mark.parametrize( - ("call_type", "response", "field", "quantity"), - [ - ("arerank", {"id": "rerank", "results": [], "meta": {"tokens": {"input_tokens": 7}}}, "prompt_tokens", 7), - ("aocr", {"id": "ocr", "usage_info": {"pages_processed": 3}}, "pages_processed", 3), - ( - "acompletion", - {"id": "google", "usageMetadata": {"promptTokenCount": 7, "totalTokenCount": 7}}, - "prompt_tokens", - 7, - ), - ], -) -async def test_capture_preserves_native_usage_before_and_after_preparation( - prepared: bool, call_type: str, response: dict[str, object], field: str, quantity: int -) -> None: - from litellm._internal_context import internal_sub_call - from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage - - database: Final = StoredUsage() - meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger(call_type) - logger.model_call_details["response_cost"] = 0.25 - if prepared: - assert logger.prepare_response_cost(response, NOW) is not None - with internal_sub_call("parent"): - capture_logging_usage(meter, logger, response, NOW, NOW) - await meter.close() - assert len(database.rows) == 1 - envelope: Final = json.loads(database.rows[0].envelope) - assert envelope["usage"][field] == quantity - assert envelope["usage_source"] == "measured" - assert envelope["ownership"] == "included_in_parent" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("tokens", [(120000, 120000), (210000, 30000)]) -@pytest.mark.parametrize("priced", [False, True]) -async def test_batch_capture_retains_request_basis_for_later_tiered_pricing( - tokens: tuple[int, int], priced: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - import litellm - from litellm.batches.batch_utils import batch_cost_and_usage - from litellm.cost_calculator import batch_cost_calculator - from litellm.proxy.spend_tracking.usage_capture import ( - UsageBatchWriter, - capture_batch_callback, - capture_logging_usage, - ) - from litellm.types.utils import Usage - - rates: Final = { - "input_cost_per_token_batches": 0.001, - "input_cost_per_token_above_200k_tokens_batches": 0.002, - "output_cost_per_token_batches": 0.0, - "litellm_provider": "openai", - "mode": "chat", - } - monkeypatch.setitem(litellm.model_cost, "capture-model", rates if priced else {"litellm_provider": "openai"}) - rows: Final = [ - { - "response": { - "status_code": 200, - "body": { - "model": "capture-model", - "usage": {"prompt_tokens": n, "completion_tokens": 0, "total_tokens": n}, - }, - } - } - for n in tokens - ] - result: Final = batch_cost_and_usage(rows, "openai", model_info=rates if priced else {}) - database: Final = StoredUsage() - writer: Final = UsageBatchWriter(database.insert, database.load) - meter: Final = DurableUsageMeter(writer.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - logger: Final = _logger("aretrieve_batch") - logger.model_call_details["response_cost"] = result.cost - capture_logging_usage(meter, logger, {"id": "batch", "status": "completed", "usage": result.usage}, NOW, NOW) - await capture_batch_callback( - meter, - database.load, - logger.model_call_details, - {"id": "batch", "status": "completed", "usage": result.usage}, - NOW, - NOW, - ) - await meter.close() - retained: Final = tuple( - {"model": "capture-model", "usage": json.loads(write.envelope)["usage"]} for write in database.rows[1:] - ) - assert sorted(row["usage"]["prompt_tokens"] for row in retained) == sorted(tokens) - reconstructed: Final = sum( - sum(batch_cost_calculator(Usage(**row["usage"]), row["model"], "openai", model_info=rates)) for row in retained - ) - expected: Final = sum(n * (0.002 if n > 200000 else 0.001) for n in tokens) - assert reconstructed == pytest.approx(expected) - assert result.cost == pytest.approx(expected if priced else 0) - assert database.rows[0].event.metrics is None - assert json.loads(database.rows[0].envelope)["charge_eligible"] is False - assert all(json.loads(write.envelope)["finalized_cost"] is None for write in database.rows[1:]) - - -@pytest.mark.asyncio -async def test_large_batch_recovers_partial_commit_with_frozen_pricing_and_rejects_changed_usage() -> None: - from litellm.proxy.spend_tracking.usage_capture import ( - PricingCapture, - UsageBatchWriter, - capture_batch_callback, - capture_callback, - ) - - database: Final = StoredUsage() - transactions: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - failure: Final = asyncio.Queue[bool]() - failure.put_nowait(True) - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - assert sum(len(write.envelope.encode()) for write in writes) <= 65536 - assert len(writes) <= 32 - if len(database.rows) > 1 and not failure.empty(): - failure.get_nowait() - raise RuntimeError("Interrupted batch persistence") - await database.insert(writes) - transactions.put_nowait(writes) - - writer: Final = UsageBatchWriter(persist, database.load) - meter: Final = DurableUsageMeter( - writer.write, - UsageCaptureSettings(batch_size=32, max_pending_records=32, max_pending_bytes=65536), - wait_for_batch=_manual_flush, - ) - rows: Final = tuple( - { - "model": "capture-model" if index % 7 else None, - "usage": {"prompt_tokens": index % 5 + 1} if index % 7 else None, - } - for index in range(1000) - ) - response: Final = {"id": "large-batch", "status": "completed", "usage": {"batch_usage_by_request": rows}} - pricing: Final = PricingCapture("{}", "{}", json.dumps({"batch_models": {"capture-model": {"rate": 1}}})) - kwargs: Final = { - "call_type": "aretrieve_batch", - "litellm_call_id": "poll", - "litellm_params": {}, - "_accounting_pricing_snapshot": pricing, - } - failed_database: Final = CommitBarrier(failure=RuntimeError("Earlier capture failed")) - failed_database.release.set() - prior_meter: Final = DurableUsageMeter(failed_database.write, UsageCaptureSettings(), wait_for_batch=_manual_flush) - prior: Final = capture_callback( - prior_meter, - {"call_type": "acompletion", "litellm_call_id": "earlier", "litellm_params": {}}, - {"id": "earlier"}, - NOW, - NOW, - ) - await prior_meter.close() - with pytest.raises(RuntimeError, match="Earlier capture failed"): - await capture_batch_callback( - meter, - database.load, - {**kwargs, "_accounting_prior_captures": (prior,)}, - response, - NOW, - NOW, - ) - assert database.rows == () - with pytest.raises(RuntimeError, match="Interrupted batch persistence"): - await capture_batch_callback(meter, database.load, kwargs, response, NOW, NOW) - partial: Final = database.rows - assert 1 < len(partial) < len(rows) + 1 - changed_pricing: Final = PricingCapture("{}", "{}", json.dumps({"batch_models": {"capture-model": {"rate": 2}}})) - retried: Final = await capture_batch_callback( - meter, - database.load, - {**kwargs, "_accounting_pricing_snapshot": changed_pricing}, - {**response, "usage": {"batch_usage_by_request": tuple(reversed(rows))}}, - NOW, - NOW, - ) - assert retried.receipt.done() and retried.receipt.exception() is None - assert len(database.rows) == len(rows) + 1 - assert database.rows[: len(partial)] == partial - envelopes: Final = tuple(json.loads(write.envelope) for write in database.rows[1:]) - assert all( - envelope["pricing"]["model_map_information"] == ({"rate": 1} if envelope["usage"] is not None else None) - for envelope in envelopes - ) - assert sum(envelope["usage"] is None for envelope in envelopes) == sum(row["usage"] is None for row in rows) - assert all(envelope["request_count_eligible"] is False for envelope in envelopes) - assert sum(envelope["usage"]["prompt_tokens"] for envelope in envelopes if envelope["usage"] is not None) == sum( - row["usage"]["prompt_tokens"] for row in rows if row["usage"] is not None - ) - before_conflict: Final = database.rows - with pytest.raises(UsageCaptureError, match="identity conflicts"): - await capture_batch_callback( - meter, - database.load, - kwargs, - {**response, "usage": {"batch_usage_by_request": rows[:-1]}}, - NOW, - NOW, - ) - await meter.close() - assert database.rows == before_conflict - assert transactions.qsize() > 2 - assert meter.pending_bytes == meter.pending_records == 0 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("outcome", ["drain", "cancel", "failure"]) -async def test_batch_callback_handoff_is_bounded_and_runtime_owns_persistence( - outcome: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - from litellm.proxy import proxy_server - from litellm.proxy.spend_tracking.usage_capture import CapturedUsage, UsageCaptureRuntime, captured_callback_kwargs - - database: Final = StoredUsage() - entered: Final = asyncio.Event() - release: Final = asyncio.Event() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - entered.set() - await release.wait() - if outcome == "failure": - raise RuntimeError("Persistence unavailable") - await database.insert(writes) - - runtime: Final = UsageCaptureRuntime() - runtime.start(persist, UsageCaptureSettings(batch_size=2, max_pending_records=2), load=database.load) - monkeypatch.setattr(proxy_server, "usage_capture", runtime) - kwargs: Final = {"call_type": "aretrieve_batch", "litellm_call_id": "poll", "litellm_params": {}} - response: Final = { - "id": "batch", - "status": "completed", - "usage": {"batch_usage_by_request": ({"model": "model", "usage": {"prompt_tokens": 2}},) * 3}, - } - callback: Final = asyncio.create_task(captured_callback_kwargs(kwargs, response, NOW, NOW)) - try: - await entered.wait() - assert callback.done(), "Batch persistence must not consume the legacy callback deadline" - captured: Final = (await callback)["_accounting_capture"] - assert isinstance(captured, CapturedUsage) and not captured.receipt.done() - second: Final = await captured_callback_kwargs(kwargs, response, NOW, NOW) - assert isinstance(second.get("_accounting_capture"), CapturedUsage) - overflow: Final = await captured_callback_kwargs(kwargs, response, NOW, NOW) - assert "_accounting_capture" not in overflow - assert "Batch capture producer capacity exhausted" in caplog.text - if outcome != "cancel": - release.set() - await runtime.stop(0 if outcome == "cancel" else 1) - assert captured.receipt.done() - assert captured.receipt.cancelled() is (outcome == "cancel") - if outcome == "failure": - assert isinstance(captured.receipt.exception(), RuntimeError) - assert "Batch usage capture failed" in caplog.text - assert len(database.rows) == (4 if outcome == "drain" else 0) - with pytest.raises(UsageCaptureError, match="active ledger writer"): - runtime.submit_batch(kwargs, response, NOW, NOW) - finally: - release.set() - await asyncio.gather(callback, return_exceptions=True) - await runtime.stop(1) - - -@pytest.mark.asyncio -async def test_batch_producers_share_byte_capacity_and_release_it_after_drain() -> None: - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - - database: Final = StoredUsage() - release: Final = asyncio.Event() - runtime: Final = UsageCaptureRuntime() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - await release.wait() - assert runtime.meter is not None - assert runtime.meter.pending_bytes <= 65536 - await database.insert(writes) - - runtime.start(persist, UsageCaptureSettings(batch_size=32, max_pending_bytes=65536), load=database.load) - from litellm.batches.batch_utils import batch_cost_and_usage - - parsed: Final = batch_cost_and_usage( - [ - { - "response": { - "status_code": 200, - "body": { - "model": "capture-model", - "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, - }, - } - } - for _ in range(1000) - ], - "openai", - model_info={}, - ) - rows: Final = ({"model": "capture-model", "usage": {"prompt_tokens": 2}},) - response: Final = {"id": "batch", "status": "completed", "usage": parsed.usage} - kwargs: Final = {"call_type": "aretrieve_batch", "litellm_call_id": "poll", "litellm_params": {}} - try: - large: Final = asyncio.Queue[asyncio.Future[None]]() - for index in range(100): - try: - captured: Final = runtime.submit_batch(kwargs, {**response, "id": f"batch-{index}"}, NOW, NOW) - except UsageCaptureError: - break - large.put_nowait(captured.receipt) - large_count: Final = large.qsize() - assert 0 < large_count < 100 - assert runtime.meter is not None and 0 < runtime.meter.pending_bytes <= 65536 - release.set() - await asyncio.gather(*(large.get_nowait() for _ in range(large_count))) - assert len(database.rows) == 1001 * large_count - observations: Final = tuple(json.loads(write.envelope) for write in database.rows) - assert sum(row["usage"].get("prompt_tokens", 0) for row in observations) == 2000 * large_count - assert sum(row["usage"].get("completion_tokens", 0) for row in observations) == 1000 * large_count - assert runtime.meter.pending_bytes == 0 - admitted: Final = asyncio.Queue[asyncio.Future[None]]() - for index in range(100): - try: - small: Final = runtime.submit_batch( - kwargs, - {**response, "id": f"small-{index}", "usage": {"batch_usage_by_request": rows}}, - NOW, - NOW, - ) - except UsageCaptureError: - break - admitted.put_nowait(small.receipt) - count: Final = admitted.qsize() - assert 1 < count < 100 - await asyncio.gather(*(admitted.get_nowait() for _ in range(count))) - assert len(database.rows) == 1001 * large_count + 2 * count - assert runtime.meter.pending_bytes == 0 - finally: - release.set() - await runtime.stop(1) - - -@pytest.mark.asyncio -async def test_batch_replay_reserves_the_accepted_pricing_snapshot_size() -> None: - from litellm.proxy.spend_tracking.usage_capture import PricingCapture, UsageBatchWriter, UsageCaptureRuntime - - database: Final = StoredUsage() - writer: Final = UsageBatchWriter(database.insert, database.load) - runtime: Final = UsageCaptureRuntime() - runtime.start(writer.write, UsageCaptureSettings(max_pending_bytes=16384), load=database.load) - response: Final = { - "id": "batch", - "status": "completed", - "usage": {"batch_usage_by_request": ({"model": "model", "usage": {"prompt_tokens": 2}},)}, - } - kwargs: Final = {"call_type": "aretrieve_batch", "litellm_call_id": "poll", "litellm_params": {}} - pricing: Final = PricingCapture("{}", "{}", json.dumps({"batch_models": {"model": {"source": "x" * 1200}}})) - try: - initial: Final = runtime.submit_batch({**kwargs, "_accounting_pricing_snapshot": pricing}, response, NOW, NOW) - await asyncio.shield(initial.receipt) - accepted: Final = database.rows - replay: Final = runtime.submit_batch(kwargs, response, NOW, NOW) - await asyncio.shield(replay.receipt) - assert database.rows == accepted and len(accepted) == 2 - assert runtime.meter is not None and runtime.meter.pending_bytes == 0 - finally: - await runtime.stop(1) diff --git a/tests/unit/proxy/spend_tracking/test_usage_meter.py b/tests/unit/proxy/spend_tracking/test_usage_meter.py deleted file mode 100644 index bd3a0b9ed70..00000000000 --- a/tests/unit/proxy/spend_tracking/test_usage_meter.py +++ /dev/null @@ -1,473 +0,0 @@ -import asyncio -import json -from datetime import datetime, timezone -from decimal import Decimal -from typing import Final - -import pytest - -from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome -from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite -from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings -from litellm.types.utils import CallTypes - - -def _write(identity: str) -> UsageWrite: - event: Final = InboundUsageEvent( - event_id=identity, - request_id=identity, - call_type=CallTypes.acompletion, - started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), - outcome=Outcome.SUCCEEDED, - dimensions=(), - metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), - ) - return UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")})) - - -class CommitBarrier: - def __init__(self, *, conflict: str | None = None, failure: Exception | None = None) -> None: - self.entered: Final = asyncio.Event() - self.release: Final = asyncio.Event() - self.committed: tuple[str, ...] = () - self.transactions = 0 - self.conflict: Final = conflict - self.failure: Final = failure - - async def write(self, writes: tuple[UsageWrite, ...]) -> None: - self.transactions += 1 - self.entered.set() - await self.release.wait() - if self.failure is not None: - raise self.failure - identities: Final = tuple(write.event.event_id for write in writes) - if self.conflict in identities: - raise AccountingConflictError("Conflicting test event") - self.committed += identities - - -async def _manual_flush(wakeup: asyncio.Event, timeout: float) -> None: - await wakeup.wait() - - -@pytest.mark.asyncio -async def test_receipts_wait_for_one_batched_commit_and_keep_inflight_capacity_reserved() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2)) - first: Final = meter.submit((_write("one"),)) - second: Final = meter.submit((_write("two"),)) - await database.entered.wait() - assert not first.done() and not second.done() - assert meter.pending_records == 2 - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("three"),)) - database.release.set() - await asyncio.gather(first, second) - await meter.close() - assert database.committed == ("one", "two") - assert database.transactions == 1 - assert (meter.pending_bytes, meter.pending_records) == (0, 0) - - -@pytest.mark.asyncio -async def test_byte_limit_rejects_new_work_without_evicting_admitted_work() -> None: - database: Final = CommitBarrier() - write: Final = _write("large" * 300) - size: Final = len(write.envelope.encode()) - 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"),)) - assert meter.pending_bytes == size - database.release.set() - await meter.close() - await receipt - 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 -async def test_conflicting_group_rolls_back_and_valid_neighbor_still_commits() -> None: - database: Final = CommitBarrier(conflict="conflict") - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3)) - rejected: Final = meter.submit((_write("caller"), _write("conflict"))) - accepted: Final = meter.submit((_write("valid"),)) - await database.entered.wait() - database.release.set() - with pytest.raises(UsageCaptureError, match="identity"): - await rejected - await accepted - await meter.close() - assert database.committed == ("valid",) - assert meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_permanent_persistence_failure_reaches_each_receipt_and_releases_capacity() -> None: - failure: Final = ValueError("invalid envelope") - database: Final = CommitBarrier(failure=failure) - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - first: Final = meter.submit((_write("one"),)) - second: Final = meter.submit((_write("two"),)) - await database.entered.wait() - database.release.set() - results: Final = await asyncio.gather(first, second, return_exceptions=True) - await meter.close() - assert results == [failure, failure] - assert meter.pending_bytes == 0 - - -@pytest.mark.asyncio -async def test_cancelling_record_waiter_does_not_cancel_persistence() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - waiting: Final = asyncio.create_task(meter.record(_write("caller").event)) - await database.entered.wait() - waiting.cancel() - with pytest.raises(asyncio.CancelledError): - await waiting - database.release.set() - await meter.close() - assert database.committed == ("caller",) - assert meter.pending_records == 0 - - -@pytest.mark.asyncio -async def test_shutdown_deadline_rejects_unacknowledged_work_and_releases_capacity() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - first: Final = meter.submit((_write("in-flight"), _write("in-flight-two"))) - await database.entered.wait() - queued: Final = meter.submit((_write("queued"),)) - await meter.stop(0) - results: Final = await asyncio.gather(first, queued, return_exceptions=True) - assert all(isinstance(result, UsageCaptureError) for result in results) - assert meter.pending_records == 0 - assert meter.pending_bytes == 0 - assert database.committed == () - - -@pytest.mark.asyncio -async def test_shutdown_during_database_backoff_releases_reserved_capacity() -> None: - retrying: Final = asyncio.Event() - blocked: Final = asyncio.Event() - - async def unavailable(writes: tuple[UsageWrite, ...]) -> None: - raise ConnectionError("database unavailable") - - async def backoff(delay: float) -> None: - retrying.set() - await blocked.wait() - - meter: Final = DurableUsageMeter(unavailable, UsageCaptureSettings(batch_size=2), retry_delay=backoff) - receipt: Final = meter.submit((_write("one"), _write("two"))) - await retrying.wait() - await meter.stop(0) - with pytest.raises(UsageCaptureError, match="interrupted"): - await receipt - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("wrapped", [False, True]) -async def test_closed_pooled_connection_retains_batch_until_automatic_recovery(wrapped: bool) -> None: - import httpx - from prisma.engine.errors import EngineRequestError - from prisma.http import Response - - database: Final = CommitBarrier() - attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]() - retrying: Final = asyncio.Event() - recovered: Final = asyncio.Event() - closure: Final = EngineRequestError( - Response(httpx.Response(500)), '{"error_code":"P1017","message":"Server has closed the connection."}' - ) - failure: Final = RuntimeError("ledger transaction failed") if wrapped else closure - if wrapped: - failure.__cause__ = closure - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - attempts.put_nowait(writes) - if attempts.qsize() == 1: - raise failure - await database.write(writes) - - async def backoff(delay: float) -> None: - retrying.set() - await recovered.wait() - - meter: Final = DurableUsageMeter( - persist, UsageCaptureSettings(batch_size=2, max_pending_records=2), retry_delay=backoff - ) - writes: Final = (_write("caller"), _write("attempt")) - receipt: Final = meter.submit(writes) - retry_signal: Final = asyncio.create_task(retrying.wait()) - reserved: Final = meter.pending_bytes - try: - await asyncio.wait((receipt, retry_signal), return_when=asyncio.FIRST_COMPLETED) - assert not receipt.done(), "A dropped pooled connection must retain admitted usage for retry" - assert (meter.pending_records, meter.pending_bytes) == (2, reserved) - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("overflow"),)) - database.release.set() - recovered.set() - await receipt - assert attempts.get_nowait() == attempts.get_nowait() == writes - assert database.committed == ("caller", "attempt") - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - finally: - retry_signal.cancel() - await asyncio.gather(retry_signal, return_exceptions=True) - await meter.stop(0) - await asyncio.gather(receipt, return_exceptions=True) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("first_commits", [False, True]) -async def test_shutdown_during_conflict_isolation_completes_every_receipt(first_commits: bool) -> None: - isolating: Final = asyncio.Event() - blocked: Final = asyncio.Event() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - if len(writes) > 1: - raise AccountingConflictError("one producer conflicts") - if first_commits and writes[0].event.event_id == "one": - return - isolating.set() - await blocked.wait() - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=2)) - first: Final = meter.submit((_write("one"),)) - second: Final = meter.submit((_write("two"),)) - await isolating.wait() - await meter.stop(0) - results: Final = await asyncio.gather(first, second, return_exceptions=True) - assert results[0] is None if first_commits else isinstance(results[0], UsageCaptureError) - assert isinstance(results[1], UsageCaptureError) - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -async def test_elapsed_flush_deadline_commits_a_partial_batch() -> None: - from litellm.proxy.spend_tracking.usage_meter import _wait_for_batch - - database: Final = CommitBarrier() - database.release.set() - - async def expired_flush(wakeup: asyncio.Event, timeout: float) -> None: - assert timeout == 5.0 - await _wait_for_batch(wakeup, 0) - - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=10), wait_for_batch=expired_flush) - receipt: Final = meter.submit((_write("partial-batch"),)) - try: - await receipt - assert database.committed == ("partial-batch",) - assert database.transactions == 1 - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - finally: - await meter.stop(0) - - -@pytest.mark.asyncio -async def test_burst_drains_full_batches_without_waiting_for_partial_batch_deadline() -> None: - database: Final = CommitBarrier() - database.release.set() - waits: Final = asyncio.Queue[bool]() - - async def observe_wait(wakeup: asyncio.Event, timeout: float) -> None: - waits.put_nowait(wakeup.is_set()) - await wakeup.wait() - - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2), wait_for_batch=observe_wait) - receipts: Final = tuple(meter.submit((_write(str(index)),)) for index in range(5)) - try: - assert await waits.get() is True - await asyncio.gather(*receipts[:2]) - assert await waits.get() is True, "A full queued batch must not wait for the flush deadline" - await asyncio.gather(*receipts[2:4]) - assert await waits.get() is False - assert database.committed == ("0", "1", "2", "3") - assert meter.pending_records == 1 - assert not receipts[-1].done() - await meter.close() - await receipts[-1] - assert database.committed == ("0", "1", "2", "3", "4") - assert database.transactions == 3 - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - finally: - await meter.stop(0) - await asyncio.gather(*receipts, return_exceptions=True) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("sqlstate", ["22P05", "22003", "42P01"]) -async def test_invalid_submission_isolated_without_retrying_systemic_faults(sqlstate: str) -> None: - from prisma.errors import DataError - - database: Final = CommitBarrier() - database.release.set() - attempts: Final = asyncio.Queue[tuple[str, ...]]() - - async def persist(writes: tuple[UsageWrite, ...]) -> None: - identities: Final = tuple(write.event.event_id for write in writes) - attempts.put_nowait(identities) - if "invalid" in identities or sqlstate == "42P01": - raise DataError( - data={"user_facing_error": {"message": "database rejected input", "meta": {"code": sqlstate}}} - ) - await database.write(writes) - - meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=3)) - rejected: Final = meter.submit((_write("caller"), _write("invalid"))) - neighbor: Final = meter.submit((_write("neighbor"),)) - try: - with pytest.raises(DataError): - await rejected - if sqlstate == "42P01": - with pytest.raises(DataError): - await neighbor - assert attempts.qsize() == 1 - else: - await neighbor - assert database.committed == ("neighbor",) - assert attempts.qsize() == 3 - finally: - await meter.close() - assert (meter.pending_records, meter.pending_bytes) == (0, 0) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("dependency_outcome", ["commit", "failure", "cancel"]) -async def test_commit_dependencies_preserve_capacity_and_isolate_failed_dependents(dependency_outcome: str) -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3, max_pending_records=4)) - prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() - dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,)) - neighbor: Final = meter.submit((_write("neighbor"),)) - assert meter.pending_records == 3 - assert not dependent.done() - with pytest.raises(UsageCaptureError, match="capacity"): - meter.submit((_write("excess-one"), _write("excess-two"))) - if dependency_outcome == "commit": - prerequisite.set_result(None) - elif dependency_outcome == "failure": - prerequisite.set_exception(ValueError("dependency failed")) - else: - prerequisite.cancel() - database.release.set() - await meter.close() - await neighbor - if dependency_outcome == "commit": - await dependent - assert database.committed == ("parent", "final", "neighbor") - else: - with pytest.raises(UsageCaptureError, match="dependency"): - await dependent - assert database.committed == ("neighbor",) - assert meter.pending_records == 0 - assert meter.pending_bytes == 0 - - -@pytest.mark.asyncio -async def test_shutdown_does_not_cancel_dependency_receipt_or_leak_reserved_capacity() -> None: - database: Final = CommitBarrier() - meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2)) - prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() - dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,)) - await meter.stop(0) - with pytest.raises(UsageCaptureError, match="acknowledgement"): - await dependent - assert not prerequisite.done() - assert database.committed == () - 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() diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index a50db09128f..dae9ffcf7fd 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -268,76 +268,6 @@ async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_comm token = MASTER_KEY -@pytest.mark.asyncio -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() - started: Final = asyncio.Event() - release: Final = asyncio.Event() - first_flush: Final = asyncio.Event() - 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: - started.set() - await release.wait() - spawn_logging_task(asyncio.get_running_loop(), enqueue_spend(), task_name="accounting") - - async def flush_spend() -> None: - while not pending.empty(): - persisted.put_nowait(pending.get_nowait()) - first_flush.set() - - worker.ensure_initialized_and_enqueue(callback()) - await started.wait() - shutdown: Final = asyncio.create_task( - litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker) - ) - try: - await first_flush.wait() - waited_for_callback: Final = not shutdown.done() - release.set() - await shutdown - 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 @@ -4377,276 +4307,11 @@ async def test_shutdown_waits_for_enqueue_producers_before_stopping_worker() -> 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 @pytest.mark.parametrize( diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index d28619c0938..3b000ba741f 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -3735,42 +3735,6 @@ def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) -@pytest.mark.asyncio -async def test_a_billed_responses_route_failure_captures_usage_from_the_newest_compaction_item() -> None: - from litellm.litellm_core_utils.get_litellm_params import get_litellm_params - from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation - from litellm.proxy.db.accounting_ledger import UsageWrite - from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime - from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings - from litellm.proxy.utils import _failure_fields_to_lift - - async def persist(_: tuple[UsageWrite, ...]) -> None: - return None - - def captured_prompt_tokens(request_input: list[dict[str, str]]) -> int: - logging_obj: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params()) - logging_obj.model_call_details["call_type"] = "aresponses" - logging_obj.model_call_details["messages"] = request_input - lifted: Final = _failure_fields_to_lift( - {"litellm_logging_obj": logging_obj}, provider_billed=True, route_call_type="aresponses" - ) - snapshot: Final = lifted["_accounting_usage_snapshot"] - assert isinstance(snapshot, UsageCaptureObservation) and snapshot.usage_json is not None - return json.loads(snapshot.usage_json)["prompt_tokens"] - - runtime: Final = UsageCaptureRuntime() - runtime.start(persist, UsageCaptureSettings()) - try: - sent: Final = captured_prompt_tokens( - [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] - ) - since_compaction: Final = captured_prompt_tokens( - [COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION] - ) - finally: - await runtime.stop(1) - - assert sent == since_compaction @pytest.mark.parametrize( @@ -4405,57 +4369,3 @@ 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") diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 1963a0302c5..2191454a2f7 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30261,27 +30261,6 @@ 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 @@ -49983,29 +49962,6 @@ export interface components { */ version: string | 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 */