mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
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 <noreply@anthropic.com> * chore(ui): regenerate dashboard API types after rebasing on main Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * 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 <noreply@anthropic.com> --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
joshua-berri
GitHub
parent
5bc7f6434d
commit
02b8da2f16
32 changed files with 845 additions and 5902 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
304
litellm/accounting/projection.py
Normal file
304
litellm/accounting/projection.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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."""
|
||||
...
|
||||
|
|
|
|||
97
litellm/accounting/snapshot.py
Normal file
97
litellm/accounting/snapshot.py
Normal file
|
|
@ -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)),
|
||||
}
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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
|
||||
|
|
@ -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})
|
||||
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
431
tests/unit/accounting/test_projection.py
Normal file
431
tests/unit/accounting/test_projection.py
Normal file
|
|
@ -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"
|
||||
]
|
||||
|
|
@ -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),
|
||||
],
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 [])
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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()
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
44
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
44
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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 */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue