scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 233

feat(accounting): add opt-in SpendLogs ledger (#49856)

Browse source 
* feat(accounting): add opt-in SpendLogs ledger delivery

* fix(accounting): preserve missing organization daily spend parity

* fix(accounting): preserve ledger receipts when replaying migration

* fix(accounting): preserve callback windows and tune ledger batches

* fix(accounting): preserve window charges and bound ledger statements

* test(accounting): verify budget callback registration order

* test(accounting): verify captured charges after repeated startup

* feat(accounting): store durable payloads independently of spend logs

* fix(accounting): preserve first delivery and trace ledger writes

* fix(accounting): keep SDK imports independent of proxy tracing

* fix(accounting): honor log retention during batch takeover

* test(accounting): cover retained batch alerts and invalid settings

* test(accounting): verify downgrade command against pending work

* fix(accounting): drain retries without starving fresh deliveries

* test(accounting): cover delivery projection, metrics and window key reads

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(accounting): read delivery metrics without the shared registry

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): charge a leftover $0 batch row when its claim fails

A failed batch claim buffers the retrieve's row and reports it charged. At flush the
row conflicted with the pre-upgrade $0 SpendLogs row, so its delivery was rolled back
and the charge never reached the spend tables. The flush now takes that row over in the
same transaction and keeps the delivery.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): trace the batch takeover at flush as a SpendLogs update

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): keep local spend logs when SPEND_LOGS_URL has no export client

The export client is never created, so main always writes SpendLogs to the proxy's own
database even with SPEND_LOGS_URL set. The ledger skipped local rows whenever the URL was
set, which dropped those request logs. It now treats a flush as an export only when the
client exists, and batch claims write local rows as before.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): keep every request's charge when its log row is shared, evicted or rejected

- Key non-batch deliveries on a per-request delivery_id so a reused client call id is still
  charged once per request, while buffer replays stay deduplicated
- Persist the redacted api_key instead of an opaque bearer token
- Keep the accounting payload when the spend-log queue evicts a row or Postgres rejects it
- Strip NUL from payloads and skip non-text tags so projection never parks a paid request

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): isolate a raw-SQL spend log row Postgres rejects instead of requeueing the batch

A rejected row on the ledger insert path raises RawQueryError, which the exact DataError check missed, so the whole batch was requeued every flush. Data and integrity SQLSTATE classes now bisect like before, while missing tables and deadlocks keep retrying

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): keep charging a request whose tag is too long for the daily tag index

A tag past Postgres's btree entry limit failed the DailyTagSpend upsert, and on the ledger path that rolled back the request's key, team and user charges with it. Such tags are now left out of the daily tag rollup only

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): store only guardrail billing fields in the accounting payload

The ledger payload copied whole guardrail records, including responses and match details, into LiteLLM_SpendLogDelivery even when spend logs were disabled. It now keeps the id, name, usage and cost fields the projection reads. NUL stripping reuses safe_json_structure instead of a recursive walk

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(proxy): move spend payload stamping out of update_database

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(proxy): restore the modules the import-layering guard evicts

The guard popped litellm.proxy.utils from sys.modules for good, so later tests on the same xdist worker patched a fresh module while calling functions bound to the old one, and spend queue tests failed depending on scheduling

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(proxy): keep update_database's payload stamping on its own local

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(accounting): type guardrail entries and use a public key redaction helper

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): keep charging requests whose end-user id is too long for an index

An oversized end-user id failed the SpendLogs insert with SQLSTATE 54000, which was requeued forever, and failed the end-user upsert in the delivery, which parked the key and team charge with it. Class 54 rows are now isolated like other rejected rows, and such ids are left out of the end-user rollups only

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(proxy): type the spend log queue snapshot

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): leave out any projected row Postgres cannot index or store

Caller-supplied ids and counters, such as a long tag or end-user id or an out-of-range compression count, failed the delivery transaction and rolled back the request's key, team and user charges. One check in project() now drops only the row that cannot be stored, replacing the per-field cutoffs

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): keep the guardrail fields compression savings are counted from

Slimming the guardrail records dropped the provider and the tokens_saved figure that Headroom compression savings are read from, and dropped a bare-dict record entirely. The payload now keeps both while still leaving out guardrail content

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): cap projected counters so daily totals cannot overflow

Counters such as a caller-supplied compression count fit one by one but overflowed the bigint daily total once summed across requests, rolling back the key, team and user charges in the same transaction. Each counter is now capped at 2^40, far below any total that could overflow

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): let a buffered batch charge take over its zero-cost receipt

A batch charge queued after a failed lookup hit the $0 receipt an earlier poll left behind, and ON CONFLICT dropped it. The buffered insert now runs the same receipt takeover as a live retrieve

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): read a rescued batch row's spend from its accounting payload

Rows rescued from an overflowing queue keep only the accounting payload, so the buffered takeover never saw their positive spend and left the $0 receipt uncharged

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): read budget windows through the delivery transaction

Window seed and existing-row reads used a separate connection while the delivery transaction held one, so with a one-connection pool they waited on their own transaction until it timed out. Inside an enclosing transaction they now share its connection

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): retry a transient failure while keeping a dropped row's charge

The rescue swallowed every error, so a database blip right after a row was rejected lost the unpaid charge. Only a rejected row is still logged and skipped, anything else propagates so the flush retries

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(accounting): give delivery projection payloads their charged entities

project() now reads entities only from accounting_entities, so the hand-built
delivery payloads carry them like real stored payloads do.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(accounting): own ledger rows by their payload and key batches on the batch id

Drop the accounting_path stamp and its SpendLogs column. A buffered row that
carries an accounting payload is ledger owned, one without it is legacy, and
spend counters get a plain windows_from_ledger flag. Window seeding tells
applied ledger rows apart by their delivery row.

A ledger batch is plain transaction A keyed by the batch id, so a repeated
completed retrieve conflicts and charges nothing. A $0 retrieve records no
delivery, a batch a legacy retrieve already charged is not charged again, and
the legacy batch path is unchanged from main apart from refusing a batch the
ledger already owns. Batch takeover and its delta bookkeeping are removed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(accounting): drop a projection test duplicated while rebasing

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(accounting): drop previous_payload now that deliveries are never replaced

Batch takeover was the only writer of previous_payload, so the worker applies
each delivery's payload as stored and the column leaves the migration.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(accounting): drop the ledger downgrade checklist script

The use_spend_logs_ledger flag is the rollback path, so a separate downgrade
gate is no longer needed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): name the window seed reads so their Prisma span resolves

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(accounting): remove the unused accounting event types and protocols

types.py, protocols.py and measurements.py were left from the old event
scaffolding. Nothing outside the package imports them, so they go with their
tests.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): park only deliveries whose own data fails

A delivery that keeps failing on a database outage, lock timeout or deadlock
stays pending and retries with capped backoff, so a recovered database applies
it. Only a data error or an invalid payload parks the isolated row.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): isolate a spend log row the insert statement rejects

The raw SpendLogs insert validates columns and JSON before Postgres sees the
row, so a row that fails there is now bisected out and dropped alone while its
charge is kept, instead of failing and replaying the whole flush.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): carry a job's model window on every completed poll in ledger mode

The ledger keeps the poll whose delivery claims the job, which need not be the
poll that won the limiter's in-memory charge. Every completed poll now carries
the job's window, so the delivery that lands always has it, while the live
model budget is still charged once.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(accounting): leave model-budget windows on the existing window writer

Model-budget windows are written by the model budget limiter rather than
derived from a request's spend, so they stay on main's window queue like live
budget counters do. This drops window capture from the accounting payload,
the deferred limiter registration and the job-window capture, and the ledger
keeps owning key and team budget windows.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): keep a delivery pending when the database is unreachable

Prisma reports a can't reach database server outage as a DataError, so the
worker now checks for an unavailable database before treating a data error
as the delivery's own fault.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): isolate only rows the raw insert cannot encode

Catching every ValueError also swallowed unexpected insert errors that main
re-raises. The raw insert now raises RejectedSpendLogRow for an unknown column
or non-JSON value, and only that is isolated as a rejected row.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(accounting): store zero savings when caller data makes them uncomputable

A huge caller-supplied tokens_saved overflowed the savings math inside
build_accounting_payload, which dropped the request's charge in ledger mode.
Savings that cannot be computed are now recorded as zero and the charge lands.

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 2026-11-04 17:49:15 -08:00 • committed by GitHub
parent b080c59e78
commit 906fa7ed28
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
48 changed files with 4041 additions and 1482 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm-proxy-extras/litellm_proxy_extras/migrations/20261103010000_spendlogs_accounting_delivery/migration.sql Normal file
View file

@ -0,0 +1,19 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogDelivery" (
"delivery_id" TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text,
"request_id" TEXT NOT NULL,
"state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending',
"attempts" INTEGER NOT NULL DEFAULT 0,
"next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"applied_at" TIMESTAMP(3),
"last_error" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"accounting_payload" JSONB NOT NULL,
"is_batch" BOOLEAN NOT NULL
);
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_next_attempt_at_idx"
ON "LiteLLM_SpendLogDelivery" ("state", "next_attempt_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_applied_at_idx"
ON "LiteLLM_SpendLogDelivery" ("state", "applied_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_request_id_idx"
ON "LiteLLM_SpendLogDelivery" ("request_id");

17
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -2141,6 +2141,23 @@ model LiteLLM_AccountingDelivery {
@@index([target, state, next_attempt_at])
}
model LiteLLM_SpendLogDelivery {
delivery_id String @id @default(dbgenerated("gen_random_uuid()::text"))
request_id String
state AccountingDeliveryState @default(pending)
accounting_payload Json
is_batch Boolean
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
created_at DateTime @default(now())
@@index([state, next_attempt_at])
@@index([state, applied_at])
@@index([request_id])
}
model LiteLLM_SubscriptionAccountTable {
subscription_account_id String @id @default(uuid())
custom_llm_provider String

59
litellm/accounting/__init__.py
View file

@ -1,59 +0,0 @@
"""Public accounting contracts. Importing these values does not enable accounting."""
from litellm.accounting.measurements import measurements_for
from litellm.accounting.protocols import CostAggregator
from litellm.accounting.types import (
AdmissionDecision,
AggregateQuery,
AggregateRow,
CostEvent,
DeliveryState,
DeliveryTarget,
Dimension,
DimensionFilter,
DimensionName,
EventKind,
Granularity,
InboundUsageEvent,
Measurement,
Metric,
MetricName,
Multiplier,
NamedCost,
OutboundUsageEvent,
Outcome,
PricedQuantity,
PricingSnapshot,
Rate,
Unpriced,
validate_aggregate_rows,
)
__all__ = [
"AdmissionDecision",
"AggregateQuery",
"AggregateRow",
"CostAggregator",
"CostEvent",
"DeliveryState",
"DeliveryTarget",
"Dimension",
"DimensionFilter",
"DimensionName",
"EventKind",
"Granularity",
"InboundUsageEvent",
"Measurement",
"Metric",
"MetricName",
"Multiplier",
"NamedCost",
"OutboundUsageEvent",
"Outcome",
"PricedQuantity",
"PricingSnapshot",
"Rate",
"Unpriced",
"measurements_for",
"validate_aggregate_rows",
]

71
litellm/accounting/application.py Normal file
View file

@ -0,0 +1,71 @@
from typing import TYPE_CHECKING, Final
from litellm.accounting.delivery_projection import entity_transactions, project_deliveries
from litellm.accounting.projection import Contribution
from litellm.proxy.db.budget_window_spend_writer import commit_window_spend_updates
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key
from litellm.proxy.db.db_span import db_span
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction
from litellm.proxy.db.model_usage_rollup import queue_model_usage_upsert
from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction
if TYPE_CHECKING:
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.utils import PrismaClient, ProxyLogging
async def apply_deliveries(
transaction: LedgerTransaction,
deliveries: tuple[DeliveryProjection, ...],
*,
prisma_client: "PrismaClient",
writer: "DBSpendUpdateWriter",
proxy_logging_obj: "ProxyLogging",
) -> tuple[Contribution, ...]:
from litellm.proxy.guardrails.usage_tracking import UsageUnitIncrement, UsageUnitKey, upsert_usage_unit_row
contributions: Final = project_deliveries(deliveries)
await writer.apply_accounting_entities(
prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions)
)
for table in sorted(DAILY_SPEND_TABLES.values(), key=lambda item: item.name):
rows = tuple(
{**dict(item.key), **dict(item.values), **dict(item.increments)}
for item in contributions
if item.table == table.name
)
merged = merge_by_conflict_key(table, rows)
for offset in range(0, len(merged), 100):
sql, params = build_bulk_upsert(table, merged[offset : offset + 100])
async with db_span("upsert_daily_spend", table.name):
await transaction.execute_raw(sql, *params)
async with transaction.batch_() as batcher:
for item in contributions:
if item.table == "LiteLLM_DailyModelUsage":
queue_model_usage_upsert(batcher, dict(item.key), dict(item.increments))
for item in contributions:
if item.table == "LiteLLM_DailyGuardrailUsageUnits":
key = dict(item.key)
counters = dict(item.increments)
await upsert_usage_unit_row(
prisma_client,
UsageUnitKey(key["guardrail_id"], key["date"], key["team_id"], key["api_key"], key["usage_unit"]),
UsageUnitIncrement(int(counters["units"]), float(counters["cost"]), int(counters["untracked_units"])),
outer_transaction=transaction,
)
windows: Final = tuple(_window(item) for item in contributions if item.table == "LiteLLM_BudgetWindowSpend")
await commit_window_spend_updates(prisma_client, windows, outer_transaction=transaction)
return contributions
def _window(item: Contribution) -> WindowSpendTransaction:
key: Final = dict(item.key)
values: Final = dict(item.values)
return WindowSpendTransaction(
entity_type=key["entity_type"],
entity_id=key["entity_id"],
window_duration=key["window_duration"],
window_start=values["window_start"] or "",
spend=float(dict(item.increments)["spend"]),
started_at=values["started_at"],
)

78
litellm/accounting/delivery_projection.py Normal file
View file

@ -0,0 +1,78 @@
from collections.abc import Iterable
from itertools import chain, groupby
from typing import TYPE_CHECKING, Final
from urllib.parse import quote
from litellm.accounting.projection import Contribution, project
from litellm.repositories.accounting_delivery_repository import DeliveryProjection
def _destination(contribution: Contribution) -> tuple[str, tuple[tuple[str, str], ...], str]:
return (
contribution.table,
contribution.key,
(dict(contribution.values).get("window_start") or "")
if contribution.table == "LiteLLM_BudgetWindowSpend"
else "",
)
def _sum(contributions: tuple[Contribution, ...]) -> Contribution:
first: Final = contributions[0]
increments: Final = tuple(
(column, sum(dict(item.increments).get(column, 0) for item in contributions))
for column in dict.fromkeys(
column for column, _ in chain.from_iterable(item.increments for item in contributions)
)
)
starts: Final = tuple(
value
for column, value in chain.from_iterable(item.values for item in contributions)
if column == "started_at" and value
)
values: Final = tuple(
(column, min(starts) if column == "started_at" and starts else value) for column, value in first.values
)
return Contribution(first.table, first.key, increments, values)
def project_deliveries(deliveries: Iterable[DeliveryProjection]) -> tuple[Contribution, ...]:
ordered: Final = sorted(chain.from_iterable(project(delivery.row) for delivery in deliveries), key=_destination)
return tuple(_sum(tuple(group)) for _, group in groupby(ordered, key=_destination))
def entity_transactions(contributions: tuple[Contribution, ...]) -> "DBSpendUpdateTransactions":
return {
"user_list_transactions": dict(_entity_totals(contributions, "LiteLLM_UserTable")),
"end_user_list_transactions": dict(_entity_totals(contributions, "LiteLLM_EndUserTable")),
"key_list_transactions": dict(_entity_totals(contributions, "LiteLLM_VerificationToken")),
"team_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TeamTable")),
"team_member_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TeamMembership")),
"org_list_transactions": dict(_entity_totals(contributions, "LiteLLM_OrganizationTable")),
"org_member_list_transactions": dict(_entity_totals(contributions, "LiteLLM_OrganizationMembership")),
"project_list_transactions": dict(_entity_totals(contributions, "LiteLLM_ProjectTable")),
"tag_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TagTable")),
"agent_list_transactions": dict(_entity_totals(contributions, "LiteLLM_AgentsTable")),
"model_access_group_list_transactions": dict(
_entity_totals(contributions, "LiteLLM_ModelAccessGroupBudgetTable")
),
}
def _entity_totals(contributions: tuple[Contribution, ...], table: str) -> tuple[tuple[str, float], ...]:
return tuple(
(_entity_key(item), float(dict(item.increments)["spend"])) for item in contributions if item.table == table
)
def _entity_key(item: Contribution) -> str:
fields: Final = dict(item.key)
if item.table == "LiteLLM_TeamMembership":
return f"team_id::{fields['team_id']}::user_id::{fields['user_id']}"
if item.table == "LiteLLM_OrganizationMembership":
return f"org_id::{quote(fields['organization_id'], safe='')}::user_id::{quote(fields['user_id'], safe='')}"
return item.key[0][1]
if TYPE_CHECKING:
from litellm.proxy._types import DBSpendUpdateTransactions

13
litellm/accounting/measurements.py
View file

@ -1,13 +0,0 @@
from typing import Final
from litellm.accounting.types import CostEvent, InboundUsageEvent, Measurement, OutboundUsageEvent
def measurements_for(event: InboundUsageEvent | OutboundUsageEvent | CostEvent) -> tuple[Measurement, ...]:
"""Return the complete contribution set for one event, ready for one apply call."""
metrics: Final = event.charges if isinstance(event, CostEvent) else event.metrics
occurred_at: Final = event.occurred_at if isinstance(event, CostEvent) else event.started_at
return tuple(
Measurement(source_event_id=event.event_id, occurred_at=occurred_at, dimensions=event.dimensions, metric=metric)
for metric in metrics or ()
)

30
litellm/accounting/metrics.py Normal file
View file

@ -0,0 +1,30 @@
from typing import Final
from prometheus_client import Gauge, Histogram
from litellm.repositories.accounting_delivery_repository import DeliveryMetrics
_PENDING: Final = Gauge("litellm_accounting_pending", "Pending SpendLogs deliveries", multiprocess_mode="livemax")
_PARKED: Final = Gauge("litellm_accounting_parked", "Parked SpendLogs deliveries", multiprocess_mode="livemax")
_OLDEST: Final = Gauge(
"litellm_accounting_oldest_pending_seconds",
"Age of the oldest pending SpendLogs delivery",
multiprocess_mode="livemax",
)
_BATCH_TIME: Final = Histogram("litellm_accounting_batch_seconds", "SpendLogs delivery worker batch time")
_BATCH_COUNT: Final = Gauge(
"litellm_accounting_batch_deliveries", "Retained batch deliveries", multiprocess_mode="livemax"
)
_TABLE_BYTES: Final = Gauge(
"litellm_accounting_delivery_table_bytes", "Delivery table size in bytes", multiprocess_mode="livemax"
)
def record_deliveries(metrics: DeliveryMetrics, seconds: float) -> None:
_PENDING.set(metrics.pending_count)
_PARKED.set(metrics.parked_count)
_OLDEST.set(max(0.0, metrics.oldest_pending_seconds))
_BATCH_TIME.observe(seconds)
_BATCH_COUNT.set(metrics.batch_count)
_TABLE_BYTES.set(metrics.table_bytes)

28
litellm/accounting/projection.py
View file

@ -1,5 +1,5 @@
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from dataclasses import dataclass, replace
from datetime import datetime, timezone
from itertools import chain
from types import MappingProxyType
@ -24,8 +24,11 @@ from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
_FIELDS: Final = TypeAdapter(dict[str, object])
_STRINGS: Final = TypeAdapter(tuple[str, ...])
_VALUES: Final = TypeAdapter(tuple[object, ...])
_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...])
_TASK_TYPES: Final = frozenset(load_model_insight_tasks())
_MAX_INDEXED_TEXT_BYTES: Final = 2000
_MAX_COUNTER: Final = 2**40
@dataclass(frozen=True, slots=True)
@ -54,7 +57,8 @@ 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)
parsed: Final = _VALUES.validate_json(value) if isinstance(value, str) else _VALUES.validate_python(value)
return tuple(tag for tag in parsed if isinstance(tag, str))
def _date(row: Mapping[str, object]) -> str | None:
@ -191,13 +195,31 @@ def _window_contributions(row: Mapping[str, object], metadata: Mapping[str, obje
def project(row: Mapping[str, object]) -> tuple[Contribution, ...]:
metadata: Final = _fields(row.get("metadata"))
return (
contributions: Final = (
*_entity_contributions(row, metadata),
*_daily_contributions(row, metadata),
*_window_contributions(row, metadata),
*_model_contributions(row, metadata),
*_guardrail_contributions(row, metadata),
)
return tuple(_capped(contribution) for contribution in contributions if _indexable(contribution))
def _indexable(contribution: Contribution) -> bool:
texts: Final = (value for _, value in (*contribution.key, *contribution.values) if value is not None)
return all(len(text.encode("utf-8", "surrogatepass")) <= _MAX_INDEXED_TEXT_BYTES for text in texts)
def _capped(contribution: Contribution) -> Contribution:
"""Counters are summed into bigint columns in the charge's transaction, so one capped far below that range
can never overflow a day's total."""
return replace(
contribution,
increments=tuple(
(name, max(-_MAX_COUNTER, min(_MAX_COUNTER, value)) if isinstance(value, int) else value)
for name, value in contribution.increments
),
)
def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:

21
litellm/accounting/protocols.py
View file

@ -1,21 +0,0 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import Protocol
from litellm.accounting.types import AggregateQuery, AggregateRow
class CostAggregator(Protocol):
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, ...]:
"""Return disjoint [start, end) rows per group using existing daily projections.
Filters combine with AND, values within a filter with IN. DAY emits daily
buckets; None emits whole-interval totals. Reject unsupported metric/scope
combinations. Different groups may overlap, including non-additive tags.
"""
...

94
litellm/accounting/snapshot.py
View file

@ -2,12 +2,13 @@ from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
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
from litellm.proxy.spend_tracking.savings import SavingsSpend, compute_savings_spend, marks_gateway_injection
if TYPE_CHECKING:
from litellm.router import Router
@ -66,7 +67,47 @@ def build_accounting_payload(
router: "Callable[[], Router | None]",
) -> Mapping[str, object]:
metadata: Final = _FIELDS.validate_json(payload["metadata"])
savings: Final = compute_savings_spend(
savings: Final = _savings_or_zero(payload, metadata, router)
accounting_metadata: Final = {
**{key: _accounting_metadata_field(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 _FIELDS.validate_python(
safe_json_structure(
{
**{
key: value.isoformat() if isinstance(value, datetime) else value
for key, value in payload.items()
if key in _ACCOUNTING_FIELDS
},
"metadata": accounting_metadata,
},
non_finite_as_null=True,
)
)
def _savings_or_zero(
payload: SpendLogsPayload, metadata: Mapping[str, object], router: "Callable[[], Router | None]"
) -> SavingsSpend:
"""Savings are analytics computed from caller-supplied counts, so one that cannot be computed is stored as zero
rather than losing the charge it rides with."""
try:
return _savings(payload, metadata, router)
except (ArithmeticError, ValueError, TypeError) as error:
verbose_proxy_logger.warning(
"Recording zero savings for request_id=%s, savings could not be computed: %s", payload["request_id"], error
)
return SavingsSpend(compression=0.0, prompt_caching=0.0)
def _savings(
payload: SpendLogsPayload, metadata: Mapping[str, object], router: "Callable[[], Router | None]"
) -> SavingsSpend:
return compute_savings_spend(
model=payload["model"],
custom_llm_provider=payload["custom_llm_provider"],
compression_saved_tokens=extract_compression_saved_tokens(metadata),
@ -80,18 +121,41 @@ def build_accounting_payload(
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,
_GUARDRAIL_ENTRIES: Final = TypeAdapter(tuple[dict[str, object], ...])
_GUARDRAIL_METADATA_FIELDS: Final = frozenset({"guardrail_information", "standard_logging_guardrail_information"})
_GUARDRAIL_ACCOUNTING_FIELDS: Final = frozenset(
{
"guardrail_id",
"guardrail_name",
"guardrail_provider",
"guardrail_usage",
"guardrail_cost_by_unit",
"guardrail_cost_in_spend",
}
)
_GUARDRAIL_RESPONSE: Final = TypeAdapter(dict[str, object])
def _accounting_guardrail_entries(value: object) -> object:
try:
entries: Final = _GUARDRAIL_ENTRIES.validate_python((value,) if isinstance(value, dict) else value)
except ValidationError:
return None
return [_accounting_guardrail_entry(entry) for entry in entries]
def _accounting_guardrail_entry(entry: Mapping[str, object]) -> Mapping[str, object]:
kept: Final = {key: item for key, item in entry.items() if key in _GUARDRAIL_ACCOUNTING_FIELDS}
response: Final = entry.get("guardrail_response")
if not isinstance(response, dict) or "tokens_saved" not in response:
return kept
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)),
**kept,
"guardrail_response": {"tokens_saved": _GUARDRAIL_RESPONSE.validate_python(response)["tokens_saved"]},
}
def _accounting_metadata_field(key: str, value: object) -> object:
return _accounting_guardrail_entries(value) if key in _GUARDRAIL_METADATA_FIELDS else value

428
litellm/accounting/types.py
View file

@ -1,428 +0,0 @@
"""Accounting values use UTC instants, decimal strings on the wire and USD charges."""
from collections.abc import Mapping
from datetime import datetime, timezone
from decimal import Decimal
from enum import Enum
from typing import Annotated, Final, Literal, TypeAlias
from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field, model_validator
from typing_extensions import Self
from litellm.types.utils import CallTypes
_CONFIG: Final = ConfigDict(frozen=True, extra="forbid", validate_default=True)
Identifier: TypeAlias = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")]
Quantity: TypeAlias = Annotated[Decimal, Field(strict=True, ge=0, allow_inf_nan=False)]
def _utc(value: datetime) -> datetime:
return value.astimezone(timezone.utc)
Instant: TypeAlias = Annotated[AwareDatetime, AfterValidator(_utc)]
class MetricName(str, Enum):
"""Reasoning tokens are included in completion tokens; above-1hr writes are included in cache creation tokens."""
REQUESTS = "requests"
SUCCESSFUL_REQUESTS = "successful_requests"
FAILED_REQUESTS = "failed_requests"
PROMPT_TOKENS = "prompt_tokens"
COMPLETION_TOKENS = "completion_tokens"
REASONING_TOKENS = "reasoning_tokens"
TOTAL_TOKENS = "total_tokens"
CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens"
CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens"
CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR = "cache_creation_input_tokens_above_1hr"
SPEND = "spend"
class DimensionName(str, Enum):
KEY = "key"
USER = "user"
TEAM = "team"
ORGANIZATION = "organization"
END_USER = "end_user"
MODEL = "model"
MODEL_GROUP = "model_group"
DEPLOYMENT = "deployment"
PROVIDER = "provider"
TOOL = "tool"
AGENT = "agent"
PROJECT = "project"
MODEL_ACCESS_GROUP = "model_access_group"
API_BASE = "api_base"
SERVICE_TIER = "service_tier"
TAG = "tag"
class Outcome(str, Enum):
SUCCEEDED = "succeeded"
FAILED = "failed"
CANCELLED = "cancelled"
class Granularity(str, Enum):
DAY = "day"
class EventKind(str, Enum):
INBOUND_USAGE = "inbound_usage"
OUTBOUND_USAGE = "outbound_usage"
COST = "cost"
class DeliveryState(str, Enum):
PENDING = "pending"
APPLIED = "applied"
PARKED = "parked"
class DeliveryTarget(str, Enum):
"""One member per existing table the projector writes; each is applied and receipted independently."""
KEY_SPEND = "LiteLLM_VerificationToken"
USER_SPEND = "LiteLLM_UserTable"
TEAM_SPEND = "LiteLLM_TeamTable"
TEAM_MEMBER_SPEND = "LiteLLM_TeamMembership"
ORGANIZATION_SPEND = "LiteLLM_OrganizationTable"
ORGANIZATION_MEMBER_SPEND = "LiteLLM_OrganizationMembership"
END_USER_SPEND = "LiteLLM_EndUserTable"
AGENT_SPEND = "LiteLLM_AgentsTable"
TAG_SPEND = "LiteLLM_TagTable"
PROJECT_SPEND = "LiteLLM_ProjectTable"
MODEL_ACCESS_GROUP_SPEND = "LiteLLM_ModelAccessGroupBudgetTable"
BUDGET_SPEND = "LiteLLM_BudgetTable"
BUDGET_WINDOW_SPEND = "LiteLLM_BudgetWindowSpend"
DAILY_USER_SPEND = "LiteLLM_DailyUserSpend"
DAILY_TEAM_SPEND = "LiteLLM_DailyTeamSpend"
DAILY_ORGANIZATION_SPEND = "LiteLLM_DailyOrganizationSpend"
DAILY_END_USER_SPEND = "LiteLLM_DailyEndUserSpend"
DAILY_AGENT_SPEND = "LiteLLM_DailyAgentSpend"
DAILY_TAG_SPEND = "LiteLLM_DailyTagSpend"
DAILY_MODEL_USAGE = "LiteLLM_DailyModelUsage"
DAILY_TOOL_SPEND = "LiteLLM_DailyToolSpend"
AUTOROUTER_DAILY_SPEND = "LiteLLM_AutoRouterDailySpend"
_SUBSET_PAIRS: Final = (
(MetricName.REASONING_TOKENS, MetricName.COMPLETION_TOKENS),
(MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, MetricName.CACHE_CREATION_INPUT_TOKENS),
)
def _check_subsets(values: Mapping[str, Decimal]) -> None:
for part, whole in _SUBSET_PAIRS:
if part.value in values and whole.value in values and values[part.value] > values[whole.value]:
raise ValueError(f"{part.value} cannot exceed {whole.value}")
class Metric(BaseModel):
model_config = _CONFIG
name: MetricName
value: Quantity
unit: Literal["count", "USD"] = "count"
@model_validator(mode="after")
def validate_unit(self) -> Self:
if (self.name is MetricName.SPEND) != (self.unit == "USD"):
raise ValueError("SPEND requires USD; usage metrics require count")
if self.unit == "count" and self.value != self.value.to_integral_value():
raise ValueError("Count metrics require integral values")
return self
class Dimension(BaseModel):
model_config = _CONFIG
name: DimensionName
value: Identifier
_REPEATABLE_DIMENSIONS: Final = frozenset({DimensionName.TAG, DimensionName.MODEL_ACCESS_GROUP})
def _dimensions(values: tuple[Dimension, ...]) -> tuple[Dimension, ...]:
if len(frozenset(values)) != len(values):
raise ValueError("Duplicate dimension values")
single_names: Final = tuple(value.name for value in values if value.name not in _REPEATABLE_DIMENSIONS)
if len(frozenset(single_names)) != len(single_names):
raise ValueError("Only tag and model access group dimensions may have multiple values")
return values
Dimensions: TypeAlias = Annotated[tuple[Dimension, ...], AfterValidator(_dimensions)]
_CALLER_METRICS: Final = frozenset((MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS))
def _metrics(values: tuple[Metric, ...]) -> tuple[Metric, ...]:
if len(frozenset(value.name for value in values)) != len(values):
raise ValueError("Duplicate metric names")
return values
Metrics: TypeAlias = Annotated[tuple[Metric, ...], AfterValidator(_metrics)]
def _interval(start: datetime, end: datetime, *, allow_empty: bool = False) -> None:
if end < start or (end == start and not allow_empty):
raise ValueError("End must follow start")
class InboundUsageEvent(BaseModel):
model_config = _CONFIG
event_id: Identifier
request_id: Identifier
call_type: CallTypes
started_at: Instant
ended_at: Instant
outcome: Outcome
dimensions: Dimensions
metrics: Metrics
@model_validator(mode="after")
def validate_caller_activity(self) -> Self:
_interval(self.started_at, self.ended_at, allow_empty=True)
if any(metric.name not in _CALLER_METRICS for metric in self.metrics):
raise ValueError("Inbound events own only caller request metrics")
if not any(metric.name is MetricName.REQUESTS and metric.value == 1 for metric in self.metrics):
raise ValueError("Inbound events count exactly one caller request")
for metric in self.metrics:
if metric.name is MetricName.SUCCESSFUL_REQUESTS and metric.value != int(self.outcome is Outcome.SUCCEEDED):
raise ValueError("Successful request count must match outcome")
if metric.name is MetricName.FAILED_REQUESTS and metric.value != int(self.outcome is not Outcome.SUCCEEDED):
raise ValueError("Failed request count must match outcome")
return self
class OutboundUsageEvent(BaseModel):
model_config = _CONFIG
event_id: Identifier
attempt_id: Identifier
request_id: Identifier | None
retry_of_attempt_id: Identifier | None
call_type: CallTypes
started_at: Instant
ended_at: Instant
outcome: Outcome
dimensions: Dimensions
metrics: Metrics | None
@model_validator(mode="after")
def validate_attempt_usage(self) -> Self:
_interval(self.started_at, self.ended_at, allow_empty=True)
if self.retry_of_attempt_id == self.attempt_id:
raise ValueError("An attempt cannot retry itself")
if self.metrics is not None and any(
metric.name in _CALLER_METRICS or metric.name is MetricName.SPEND for metric in self.metrics
):
raise ValueError("Outbound events own only attempt usage metrics")
if self.metrics is not None:
_check_subsets({metric.name.value: metric.value for metric in self.metrics})
return self
PricingUnit: TypeAlias = Literal["token", "second", "character", "image", "request"]
class Rate(BaseModel):
model_config = _CONFIG
measure: Identifier
unit: PricingUnit
usd_per_unit: Quantity
class PricedQuantity(BaseModel):
"""A pricing input in its billing unit. Only seconds may be fractional."""
model_config = _CONFIG
measure: Identifier
unit: PricingUnit
value: Quantity
@model_validator(mode="after")
def validate_discrete_quantity(self) -> Self:
if self.unit != "second" and self.value != self.value.to_integral_value():
raise ValueError("Discrete priced quantities require integral values")
return self
class Multiplier(BaseModel):
"""A regional or tier factor applied on top of the priced key, e.g. data_residency=eu or vertex_location=us-east5."""
model_config = _CONFIG
name: Identifier
value: Identifier
factor: Annotated[Decimal, Field(strict=True, gt=0, allow_inf_nan=False)]
class NamedCost(BaseModel):
"""A free-form extra charge by name, e.g. azure_model_router_flat_cost."""
model_config = _CONFIG
name: Identifier
amount: Quantity
class PricingSnapshot(BaseModel):
"""Retained inputs and result of authoritative compilation, reproducible after cost-map rates change.
Request-scoped inputs, overrides and effective rates come from the usage row's versioned payload.
Retries reuse the committed charge and snapshot.
Component names match ``CostBreakdown``: reasoning_cost is part of output_cost, not an additional charge.
Component sums and float tolerances belong to reconciliation, not this structural contract.
"""
model_config = _CONFIG
schema_version: Literal[1] = 1
calculator_version: Identifier
model: Identifier
provider: Identifier | None
pricing_source: Literal["cost_map", "deployment", "per_request", "upstream_header"]
priced_model_key: Identifier
multipliers: tuple[Multiplier, ...] = ()
rates: tuple[Rate, ...] = ()
usage: tuple[PricedQuantity, ...] = ()
input_cost: Quantity | None = None
output_cost: Quantity | None = None
cache_read_cost: Quantity | None = None
cache_creation_cost: Quantity | None = None
reasoning_cost: Quantity | None = None
tool_usage_cost: Quantity | None = None
guardrail_cost: Quantity | None = None
additional_costs: tuple[NamedCost, ...] = ()
original_cost: Quantity | None = None
discount_percent: Quantity | None = None
discount_amount: Quantity | None = None
margin_percent: Quantity | None = None
margin_fixed_amount: Quantity | None = None
margin_total_amount: Quantity | None = None
upstream_reported_cost: Quantity | None = None
total_cost: Quantity
@model_validator(mode="after")
def validate_snapshot(self) -> Self:
rate_keys: Final = frozenset((rate.measure, rate.unit) for rate in self.rates)
quantity_keys: Final = frozenset((quantity.measure, quantity.unit) for quantity in self.usage)
if len(rate_keys) != len(self.rates):
raise ValueError("Duplicate rate measures and units")
if len(quantity_keys) != len(self.usage):
raise ValueError("Duplicate priced quantity measures and units")
if not rate_keys.issubset(quantity_keys):
raise ValueError("Every rate requires one priced quantity with the same measure and unit")
if len(frozenset(cost.name for cost in self.additional_costs)) != len(self.additional_costs):
raise ValueError("Duplicate additional cost names")
if len(frozenset(multiplier.name for multiplier in self.multipliers)) != len(self.multipliers):
raise ValueError("Duplicate multiplier names")
_check_subsets({quantity.measure: quantity.value for quantity in self.usage if quantity.unit == "token"})
return self
class CostEvent(BaseModel):
model_config = _CONFIG
event_id: Identifier
usage_event_id: Identifier
occurred_at: Instant
dimensions: Dimensions
charges: Metrics
snapshot: PricingSnapshot
@model_validator(mode="after")
def validate_charge(self) -> Self:
if len(self.charges) != 1 or self.charges[0].name is not MetricName.SPEND:
raise ValueError("A cost event requires one authoritative USD spend metric")
if self.event_id == self.usage_event_id:
raise ValueError("Cost and usage events require distinct identities")
if self.snapshot.total_cost != self.charges[0].value:
raise ValueError("Snapshot total must equal the authoritative spend")
return self
class Unpriced(BaseModel):
model_config = _CONFIG
usage_event_id: Identifier
reason: Identifier
class Measurement(BaseModel):
model_config = _CONFIG
source_event_id: Identifier
occurred_at: Instant
dimensions: Dimensions
metric: Metric
class DimensionFilter(BaseModel):
model_config = _CONFIG
name: DimensionName
values: Annotated[tuple[Identifier, ...], Field(min_length=1)]
class AggregateQuery(BaseModel):
model_config = _CONFIG
metric: MetricName
filters: tuple[DimensionFilter, ...]
group_by: tuple[DimensionName, ...]
start: Instant
end: Instant
granularity: Granularity | None
@model_validator(mode="after")
def validate_query(self) -> Self:
_interval(self.start, self.end)
if len(frozenset(item.name for item in self.filters)) != len(self.filters):
raise ValueError("Use one IN filter per dimension; filters combine with AND")
if len(frozenset(self.group_by)) != len(self.group_by):
raise ValueError("Duplicate group-by dimensions")
if any(
(instant.hour, instant.minute, instant.second, instant.microsecond) != (0, 0, 0, 0)
for instant in (self.start, self.end)
):
raise ValueError("Initial reporting queries require UTC day boundaries")
return self
class AggregateRow(BaseModel):
model_config = _CONFIG
start: Instant
end: Instant
dimensions: Dimensions
value: Quantity
@model_validator(mode="after")
def validate_bounds(self) -> Self:
_interval(self.start, self.end)
return self
def validate_aggregate_rows(rows: tuple[AggregateRow, ...]) -> tuple[AggregateRow, ...]:
"""Validate [start, end) buckets per group; overlapping tag groups are not additive."""
groups: Final = frozenset(frozenset(row.dimensions) for row in rows)
for group in groups:
ordered = sorted((row for row in rows if frozenset(row.dimensions) == group), key=lambda row: row.start)
if any(left.end > right.start for left, right in zip(ordered, ordered[1:])):
raise ValueError("Aggregate buckets overlap within a dimension group")
return rows
class AdmissionDecision(BaseModel):
model_config = _CONFIG
allowed: bool
reason: Identifier | None = None

102
litellm/accounting/worker.py Normal file
View file

@ -0,0 +1,102 @@
import os
import time
from typing import (
TYPE_CHECKING,
Final,
Literal,
Protocol,
cast, # noqa: TID251 # generated Prisma transaction signatures cannot express the borrowed transaction protocol
)
from litellm._logging import verbose_proxy_logger
from litellm.accounting.application import apply_deliveries
from litellm.accounting.metrics import record_deliveries
from litellm.accounting.projection import Contribution
from litellm.repositories.accounting_delivery_repository import (
AccountingDeliveryRepository,
DeliveryProjection,
LedgerDatabase,
LedgerTransaction,
)
from litellm.types.integrations.slack_alerting import AlertType
if TYPE_CHECKING:
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.utils import PrismaClient, ProxyLogging
class _AccountingAlerts(Protocol):
async def alerting_handler(
self, message: str, level: Literal["Low", "Medium", "High"], alert_type: AlertType
) -> None: ...
_JOB: Final = "accounting_delivery"
async def flush_accounting_deliveries(
writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging"
) -> None:
acquired: Final = await writer.pod_lock_manager.acquire_lock(cronjob_id=_JOB, allow_reentrant=False)
if acquired is False:
return
applied: tuple[Contribution, ...] = () # rebind-ok: keep committed contributions for cache invalidation
async def apply(transaction: LedgerTransaction, deliveries: tuple[DeliveryProjection, ...]) -> None:
nonlocal applied
applied = await apply_deliveries( # rebind-ok: invalidate caches only after the outer transaction commits
transaction, deliveries, prisma_client=prisma_client, writer=writer, proxy_logging_obj=proxy_logging_obj
)
try:
database: Final = cast( # cast-ok: PrismaClient.tx owns writer routing and yields the generated transaction
LedgerDatabase, prisma_client
)
repository: Final = AccountingDeliveryRepository(database)
started: Final = time.monotonic()
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "10000"))
async def run_batch(*, retry_only: bool) -> bool:
result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only)
if result.applied:
await writer.invalidate_accounting_entities(applied, proxy_logging_obj)
if result.parked:
verbose_proxy_logger.error(
"Accounting delivery parked %d requests after isolated failures", len(result.parked)
)
await _alert(
proxy_logging_obj,
message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.",
)
return bool(result.applied or result.failed)
await run_batch(retry_only=False)
retry_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True):
break
await repository.cleanup_applied(
retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")),
batch_size=batch_size * 2,
)
metrics: Final = await repository.metrics()
record_deliveries(metrics, time.monotonic() - started)
if metrics.batch_count >= int(
os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000")
) or metrics.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")):
await _alert(
proxy_logging_obj,
message=f"Accounting delivery storage requires review. Retained batches={metrics.batch_count}, table bytes={metrics.table_bytes}.",
)
if metrics.oldest_pending_seconds > 300:
await _alert(
proxy_logging_obj,
message=f"Accounting deliveries are delayed. Pending={metrics.pending_count}, oldest={metrics.oldest_pending_seconds:.0f}s.",
)
finally:
if acquired:
await writer.pod_lock_manager.release_lock(cronjob_id=_JOB)
async def _alert(alerts: _AccountingAlerts, message: str) -> None:
await alerts.alerting_handler(message=message, level="High", alert_type=AlertType.failed_tracking_spend)

1
litellm/integrations/otel/model/spans.py
View file

@ -340,6 +340,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_AccessGroupTable",
"LiteLLM_AccountingDelivery",
"LiteLLM_AccountingEvent",
"LiteLLM_SpendLogDelivery",
"LiteLLM_ClaudeCodePluginTable",
"LiteLLM_ClaudeCodePluginUpload",
"LiteLLM_MemoryTable",

8
litellm/proxy/_types.py
View file

@ -2480,6 +2480,11 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
Documents all the fields supported by `general_settings` in config.yaml
"""
use_spend_logs_ledger: bool = Field(
False,
description="Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.",
)
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.",
@ -3800,6 +3805,9 @@ class SpendLogsMetadata(TypedDict):
class SpendLogsPayload(TypedDict):
accounting_payload: NotRequired[Mapping[str, object]] # writable-ok: capture before buffering
accounting_log_disabled: NotRequired[bool] # writable-ok: capture logging mode before buffering
accounting_delivery_id: NotRequired[str] # writable-ok: assigned once before buffering
request_id: str
call_type: str
api_key: str

98
litellm/proxy/db/budget_window_spend_writer.py
View file

@ -14,11 +14,13 @@ interval of increments queued elsewhere: the same lag the SpendLogs aggregate
it replaces (and every other spend column) already has.
"""
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from contextlib import nullcontext
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import TypeAdapter
from typing_extensions import LiteralString, ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
@ -74,30 +76,44 @@ _OPEN_WINDOW_SPEND_SQL: Final = (
"WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')"
)
_HAS_DELIVERY: Final = 'EXISTS (SELECT 1 FROM "LiteLLM_SpendLogDelivery" delivery WHERE delivery.request_id = "LiteLLM_SpendLogs".request_id)'
_SETTLED_SPEND_LOG_FILTER: Final = (
' AND NOT EXISTS (SELECT 1 FROM "LiteLLM_SpendLogDelivery" delivery '
'WHERE delivery.request_id = "LiteLLM_SpendLogs".request_id '
"AND delivery.state IN ('pending', 'parked'))"
)
_SEED_FROM_SPEND_LOGS_KEY_SQL: Final = (
"SELECT COALESCE(SUM(spend), 0.0) AS total, "
"COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch "
"COALESCE(SUM(spend) FILTER (WHERE "
+ _HAS_DELIVERY
+ " OR \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) "
"AS before_batch "
'FROM "LiteLLM_SpendLogs" '
"WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')"
"WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER
)
_SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = (
"SELECT COALESCE(SUM(spend), 0.0) AS total, "
"COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch "
"COALESCE(SUM(spend) FILTER (WHERE "
+ _HAS_DELIVERY
+ " OR \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) "
"AS before_batch "
'FROM "LiteLLM_SpendLogs" '
"WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')"
"WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER
)
_SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = (
"SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch "
'FROM "LiteLLM_SpendLogs" '
"WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')"
"WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER
)
_SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL: Final = (
"SELECT COALESCE(SUM(spend), 0.0) AS total, COALESCE(SUM(spend), 0.0) AS before_batch "
'FROM "LiteLLM_SpendLogs" '
"WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')"
"WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC')" + _SETTLED_SPEND_LOG_FILTER
)
_UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60)
@ -114,6 +130,10 @@ class _ExistingKeyRow(TypedDict):
window_duration: ReadOnly[str]
_SEED_TOTALS_ROWS: Final = TypeAdapter(tuple[_SeedTotalsRow, ...])
_EXISTING_KEY_ROWS: Final = TypeAdapter(tuple[_ExistingKeyRow, ...])
class _SeedTotalsQuery(Protocol):
async def __call__(self, query: LiteralString, *args: object) -> Sequence[_SeedTotalsRow]: ...
@ -135,6 +155,8 @@ class _UpsertBatchManager(Protocol):
class _UpsertTransaction(Protocol):
def batch_(self) -> _UpsertBatchManager: ...
async def query_raw(self, query: LiteralString, *args: object) -> Sequence[Mapping[str, object]]: ...
class _UpsertTransactionManager(Protocol):
async def __aenter__(self) -> _UpsertTransaction: ...
@ -178,6 +200,7 @@ class WindowSpendLogsAggregate(Protocol):
entity_id: str,
window_start: datetime,
batch_started_at: datetime | None,
query: _SeedTotalsQuery | None = None,
) -> WindowSeedTotals | None: ...
@ -187,6 +210,7 @@ async def spend_logs_seed_totals(
entity_id: str,
window_start: datetime,
batch_started_at: datetime | None,
query: _SeedTotalsQuery | None = None,
) -> WindowSeedTotals | None:
"""LiteLLM_SpendLogs spend for one entity since window_start, both in full
and up to the start of the batch being flushed, in one scan.
@ -210,11 +234,12 @@ async def spend_logs_seed_totals(
bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_TEAM_SQL, _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL
else:
return None
read: Final = query or _seed_totals_query(prisma_client)
async with db_span("seed_window_spend", "LiteLLM_SpendLogs"):
rows: Final = (
await _seed_totals_query(prisma_client)(unbounded_sql, entity_id, window_start)
await read(unbounded_sql, entity_id, window_start)
if batch_started_at is None
else await _seed_totals_query(prisma_client)(
else await read(
bounded_sql,
entity_id,
window_start,
@ -246,9 +271,10 @@ def _primary_key(transaction: WindowSpendTransaction) -> tuple[str, str, str]:
async def _existing_primary_keys(
prisma_client: "PrismaClient",
transactions: tuple[WindowSpendTransaction, ...],
query: _ExistingKeysQuery | None = None,
) -> frozenset[tuple[str, str, str]]:
async with db_span("select_window_spend_rows", "LiteLLM_BudgetWindowSpend"):
rows: Final = await _existing_keys_query(prisma_client)(
rows: Final = await (query or _existing_keys_query(prisma_client))(
_SELECT_EXISTING_ROWS_SQL,
tuple(transaction["entity_type"] for transaction in transactions),
tuple(transaction["entity_id"] for transaction in transactions),
@ -262,6 +288,7 @@ async def _seed_base_for_missing_row(
transaction: WindowSpendTransaction,
existing_primary_keys: frozenset[tuple[str, str, str]],
spend_logs_aggregate: WindowSpendLogsAggregate,
query: _SeedTotalsQuery | None = None,
) -> float:
"""Spend already recorded for a window that has no row yet.
@ -272,12 +299,24 @@ async def _seed_base_for_missing_row(
"""
if _primary_key(transaction) in existing_primary_keys:
return 0.0
totals: Final = await spend_logs_aggregate(
prisma_client=prisma_client,
entity_type=transaction["entity_type"],
entity_id=transaction["entity_id"],
window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc),
batch_started_at=_transaction_started_at(transaction),
window_start: Final = datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc)
totals: Final = (
await spend_logs_aggregate(
prisma_client=prisma_client,
entity_type=transaction["entity_type"],
entity_id=transaction["entity_id"],
window_start=window_start,
batch_started_at=_transaction_started_at(transaction),
)
if query is None
else await spend_logs_aggregate(
prisma_client=prisma_client,
entity_type=transaction["entity_type"],
entity_id=transaction["entity_id"],
window_start=window_start,
batch_started_at=_transaction_started_at(transaction),
query=query,
)
)
if totals is None:
return 0.0
@ -324,10 +363,30 @@ def _upsert_params(
)
def _existing_keys_in(transaction: _UpsertTransaction) -> _ExistingKeysQuery:
async def query_raw(query: LiteralString, *args: object) -> Sequence[_ExistingKeyRow]:
return _EXISTING_KEY_ROWS.validate_python(await transaction.query_raw(query, *args))
return query_raw
def _seed_totals_in(transaction: _UpsertTransaction) -> _SeedTotalsQuery:
async def query_raw(query: LiteralString, *args: object) -> Sequence[_SeedTotalsRow]:
return _SEED_TOTALS_ROWS.validate_python(await transaction.query_raw(query, *args))
return query_raw
def _transaction_reads(transaction: _UpsertTransaction) -> tuple[_ExistingKeysQuery, _SeedTotalsQuery]:
"""Reads made inside an enclosing transaction share its connection, so they cannot wait on it."""
return _existing_keys_in(transaction), _seed_totals_in(transaction)
async def commit_window_spend_updates(
prisma_client: "PrismaClient",
transactions: Sequence[WindowSpendTransaction],
spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_seed_totals,
outer_transaction: _UpsertTransaction | None = None,
) -> None:
"""Apply aggregated window increments to LiteLLM_BudgetWindowSpend.
@ -343,9 +402,11 @@ async def commit_window_spend_updates(
return
ordered: Final = tuple(sorted(transactions, key=window_spend_group_key))
reads: Final = _transaction_reads(outer_transaction) if outer_transaction is not None else (None, None)
existing_primary_keys: Final = await _existing_primary_keys(
prisma_client=prisma_client,
transactions=ordered,
query=reads[0],
)
seed_bases: Final = tuple(
[
@ -354,6 +415,7 @@ async def commit_window_spend_updates(
transaction=transaction,
existing_primary_keys=existing_primary_keys,
spend_logs_aggregate=spend_logs_aggregate,
query=reads[1],
)
for transaction in ordered
]
@ -367,7 +429,9 @@ async def commit_window_spend_updates(
)
async with (
db_span("commit_window_spend_updates", "LiteLLM_BudgetWindowSpend"),
_upsert_transaction(prisma_client) as db_transaction,
(
nullcontext(outer_transaction) if outer_transaction is not None else _upsert_transaction(prisma_client)
) as db_transaction,
db_transaction.batch_() as batcher,
):
for transaction, seed_base in zip(ordered, seed_bases):

4
litellm/proxy/db/daily_spend_bulk_upsert.py
View file

@ -55,7 +55,7 @@ _ENTITY_INPUT_KEYS: Final[Mapping[DailySpendEntity, str]] = MappingProxyType(
"tag": "request_tags",
}
)
_TAGS: Final = TypeAdapter(tuple[str, ...])
_TAGS: Final = TypeAdapter(tuple[object, ...])
def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEntity) -> tuple[str | None, ...]:
@ -67,7 +67,7 @@ def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEnti
if value is None:
return ()
tags: Final = _TAGS.validate_json(value) if isinstance(value, str) else _TAGS.validate_python(value)
return tuple(dict.fromkeys(tags))
return tuple(dict.fromkeys(tag for tag in tags if isinstance(tag, str)))
if value is None:
return (None,) if entity == "user" else ()
if not isinstance(value, str) or (entity == "end_user" and not value):

306
litellm/proxy/db/db_spend_update_writer.py
View file

@ -12,6 +12,7 @@ import json
import os
import random
import time
import uuid
from collections.abc import AsyncGenerator, Callable, Coroutine, Iterable, Mapping, Sequence
from contextlib import asynccontextmanager
from contextvars import ContextVar
@ -96,6 +97,7 @@ from litellm.types.utils import CallTypes
from litellm.types.videos.utils import VIDEO_COST_REQUEST_ID_SUFFIX, VIDEO_STATUS_CALL_TYPES
if TYPE_CHECKING:
from litellm.accounting.projection import Contribution
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -160,6 +162,18 @@ def _org_member_transaction_key(org_id: str, user_id: str) -> str:
return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}"
ACCOUNTING_LEDGER_KWARG: Final = "accounting_ledger"
def _persistable_api_key(token: str | None, logged_api_key: str) -> str | None:
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
from litellm.proxy.spend_tracking.spend_tracking_utils import redact_logged_api_key
if token is None or token == logged_api_key or is_valid_sha256_hash(token):
return token
return redact_logged_api_key(token)
async def _track_guardrail_usage_of_cost_row(
payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool
) -> None:
@ -200,12 +214,20 @@ _COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "star
def _cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]:
"""Reduce a job's cost row to what tells its polls apart when logging is off.
A proxy run with spend logs disabled still needs one row per job to charge it once,
A legacy proxy run with spend logs disabled still needs one row per job to charge it once,
so the row is written either way, but it carries no request of its own: no metadata,
no requester IP, no key, model, or token counts (LIT-7048).
no requester IP, no key, model, or token counts (LIT-7048). A ledger batch is deduplicated
by its delivery row instead, so it writes no local row when logging is off.
"""
if disable_spend_logs is False:
return payload
if "accounting_payload" in payload:
return {
"request_id": payload["request_id"],
"call_type": payload["call_type"],
"accounting_payload": payload["accounting_payload"],
"accounting_log_disabled": True,
}
return MappingProxyType({field: value for field, value in payload.items() if field in _COST_CLAIM_FIELDS})
@ -219,8 +241,17 @@ class _MemberSpendRow(TypedDict):
cost: ReadOnly[float]
class SpendUpsertTable(Protocol):
def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
class _SpendBatch(Protocol):
litellm_usertable: BatchTable
litellm_endusertable: SpendUpsertTable
litellm_dailymodelusage: SpendUpsertTable
def execute_raw(self, query: LiteralString, *args: object) -> None: ...
litellm_verificationtoken: BatchTable
litellm_teamtable: BatchTable
litellm_teammembership: BatchTable
@ -251,14 +282,14 @@ def _entity_spend_table(batcher: _SpendBatch, table_accessor: _EntitySpendTable)
return _ENTITY_SPEND_TABLES[table_accessor](batcher)
class _SpendBatchManager(Protocol):
class SpendBatchManager(Protocol):
async def __aenter__(self) -> _SpendBatch: ...
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
class _SpendTransaction(Protocol):
def batch_(self) -> _SpendBatchManager: ...
def batch_(self) -> SpendBatchManager: ...
async def execute_raw(self, query: LiteralString, *args: object) -> int: ...
@ -413,8 +444,14 @@ _ENTITY_SPEND_MODELS: Final[Mapping[_EntitySpendTable, str]] = MappingProxyType(
@asynccontextmanager
async def _spend_update_tx(
prisma_client: PrismaClient, table: str, call_type: str = "commit_spend_updates"
prisma_client: PrismaClient,
table: str,
call_type: str = "commit_spend_updates",
outer_transaction: _SpendTransaction | None = None,
) -> AsyncGenerator[_SpendTransaction]:
if outer_transaction is not None:
yield outer_transaction
return
tx: Final[_SpendTransactionManager] = prisma_client.db.tx(timeout=timedelta(seconds=60))
async with db_span(call_type, table), tx as transaction:
await apply_rollup_lock_timeout(transaction)
@ -601,6 +638,21 @@ def debitable_model_access_groups(
return tuple(group for group in ordered if group in served)
def _billable_cost(response_cost: float | None, payload: SpendLogsPayload) -> float | None:
billable_cost: Final = finite_cost(response_cost)
if response_cost is not None and billable_cost is None:
verbose_proxy_logger.warning(
"Spend tracking - recording a non-finite response_cost=%s as 0 for request_id=%s",
response_cost,
single_log_line(payload["request_id"]),
)
return billable_cost
def _iso_timestamp(value: datetime | str) -> str:
return value.isoformat() if isinstance(value, datetime) else value
class DBSpendUpdateWriter:
"""
Module responsible for
@ -649,6 +701,7 @@ class DBSpendUpdateWriter:
so the caller leaves the key, team, and user counters alone (LIT-7048).
"""
from litellm.proxy.proxy_server import (
accounting_window_keys,
disable_spend_logs,
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy proxy configuration is narrowed below
litellm_proxy_budget_name,
@ -656,6 +709,13 @@ class DBSpendUpdateWriter:
)
from litellm.proxy.utils import ProxyUpdateSpend, hash_token
proxy_settings: Final = cast( # cast-ok: proxy settings use string keys
Mapping[str, object], general_settings
)
use_ledger: Final = proxy_settings.get("use_spend_logs_ledger") is True
request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {})
if kwargs is not None:
kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker
try:
verbose_proxy_logger.debug(
"Enters prisma db call, response_cost: %s, token: %s; user_id: %s; team_id: %s",
@ -666,10 +726,9 @@ class DBSpendUpdateWriter:
)
if ProxyUpdateSpend.disable_spend_updates() is True:
return True
if token is not None and isinstance(token, str) and token.startswith("sk-"):
hashed_token = hash_token(token=token)
else:
hashed_token = token
hashed_token: Final = (
hash_token(token=token) if isinstance(token, str) and token.startswith("sk-") else token
)
## CREATE SPEND LOG PAYLOAD ##
from litellm.proxy.spend_tracking.spend_tracking_utils import (
@ -685,32 +744,42 @@ class DBSpendUpdateWriter:
end_time=end_time,
llm_router=llm_router,
)
billable_cost: Final = finite_cost(response_cost)
if response_cost is not None and billable_cost is None:
verbose_proxy_logger.warning(
"Spend tracking - recording a non-finite response_cost=%s as 0 for request_id=%s",
response_cost,
single_log_line(payload["request_id"]),
)
billable_cost: Final = _billable_cost(response_cost, payload)
payload["spend"] = billable_cost or 0.0
if isinstance(payload["startTime"], datetime):
payload["startTime"] = payload["startTime"].isoformat()
if isinstance(payload["endTime"], datetime):
payload["endTime"] = payload["endTime"].isoformat()
if org_id is not None and org_id != "":
payload["startTime"] = _iso_timestamp(payload["startTime"])
payload["endTime"] = _iso_timestamp(payload["endTime"])
if org_id:
payload["organization_id"] = org_id
if team_id is not None and team_id != "":
if team_id:
payload["team_id"] = team_id
error_log_settings: Final = cast( # cast-ok: proxy settings use string keys
Mapping[str, object], general_settings
)
if use_ledger:
from litellm.accounting.snapshot import build_accounting_payload
if not _is_settle_once_cost_row(payload):
payload["accounting_delivery_id"] = uuid.uuid4().hex
payload["accounting_payload"] = build_accounting_payload(
payload,
entities={
"api_key": _persistable_api_key(hashed_token, payload.get("api_key") or ""),
"user": user_id,
"end_user": end_user_id,
"team_id": team_id,
"organization_id": org_id,
"project_id": project_id,
"agent_id": payload["agent_id"],
"proxy_budget_name": litellm_proxy_budget_name if litellm.max_budget > 0 else None,
},
model_access_groups=debitable_model_access_groups(
get_request_model_access_groups(request_fields), payload["model_id"], llm_router
),
windows=await accounting_window_keys(hashed_token, team_id) if billable_cost else (),
router=get_llm_router,
)
keep_spend_logs: Final = disable_spend_logs is False and not (
"status" in payload
and payload["status"] == "failure"
and error_log_settings.get("disable_error_logs") is True
and proxy_settings.get("disable_error_logs") is True
)
if not await self._record_spend_log(
payload=payload,
@ -739,24 +808,26 @@ class DBSpendUpdateWriter:
"disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur."
)
spawn_logging_task(
asyncio.get_running_loop(),
self._batch_database_updates(
response_cost=billable_cost,
user_id=user_id,
hashed_token=hashed_token,
team_id=team_id,
org_id=org_id,
project_id=project_id,
end_user_id=end_user_id,
prisma_client=prisma_client,
litellm_proxy_budget_name=litellm_proxy_budget_name,
payload=payload,
request_model_access_groups=get_request_model_access_groups(kwargs),
),
task_name="spend rollups",
)
if not use_ledger:
spawn_logging_task(
asyncio.get_running_loop(),
self._batch_database_updates(
response_cost=billable_cost,
user_id=user_id,
hashed_token=hashed_token,
team_id=team_id,
org_id=org_id,
project_id=project_id,
end_user_id=end_user_id,
prisma_client=prisma_client,
litellm_proxy_budget_name=litellm_proxy_budget_name,
payload=payload,
request_model_access_groups=get_request_model_access_groups(request_fields),
),
task_name="spend rollups",
)
else:
self._record_request_error_safely(payload=payload, prisma_client=prisma_client)
self._enqueue_tool_registry_upsert(
kwargs=kwargs,
completion_response=completion_response,
@ -791,7 +862,13 @@ class DBSpendUpdateWriter:
return await self._claim_settle_once_cost_row(
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
)
if disable_spend_logs is False:
if "accounting_payload" in payload:
await self._insert_spend_log_to_db(
payload={**payload, "accounting_log_disabled": disable_spend_logs},
prisma_client=prisma_client,
llm_router=llm_router,
)
elif disable_spend_logs is False:
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client, llm_router=llm_router)
return True
@ -804,7 +881,9 @@ class DBSpendUpdateWriter:
the charge and every later one finds the row and charges nothing (LIT-7048). Only
a row that recorded a charge counts: a failed poll, a request whose client
picked the job id as its call id, and the $0 row a poll left behind
while the job was still running all leave the charge to be made.
while the job was still running all leave the charge to be made. A ledger job
is owned by its delivery row, keyed by its cost row id, in every logging mode, and a
$0 poll records no delivery so it cannot block the poll that carries the cost.
"""
from litellm.repositories.table_repositories import SpendLogsRepository
@ -817,13 +896,24 @@ class DBSpendUpdateWriter:
row: Final = spend_log_row_with_retained_metadata(
_cost_row_to_write(payload, disable_spend_logs), configured_spend_logs_metadata_fields()
)
spend_logs: Final = SpendLogsRepository(prisma_client).table
repository: Final = SpendLogsRepository(prisma_client)
try:
claimed: Final = await spend_logs.create_many(
if "accounting_payload" in payload:
if payload["spend"] <= 0:
return True
accepted: Final = await repository.insert_with_deliveries([prisma_client.jsonify_object(row)]) == 1
if accepted:
await _track_guardrail_usage_of_cost_row(payload, prisma_client, disable_spend_logs)
return accepted
if await repository.has_delivery(request_id):
return False
claimed: Final = await repository.table.create_many(
data=[prisma_client.jsonify_object(row)],
skip_duplicates=True,
)
existing: Final = None if claimed == 1 else await spend_logs.find_unique(where={"request_id": request_id})
existing: Final = (
None if claimed == 1 else await repository.table.find_unique(where={"request_id": request_id})
)
except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log
verbose_proxy_logger.warning(
"Could not claim spend row %s for a job's cost, queueing it: %s", request_id, e
@ -1268,10 +1358,7 @@ class DBSpendUpdateWriter:
self.add_spend_log_transaction_to_daily_team_transaction(payload=payload_copy, prisma_client=prisma_client),
payload=payload_copy,
)
try:
self._record_request_error(payload=payload_copy, prisma_client=prisma_client)
except Exception as e: # noqa: BLE001 # the rollup must never skip the sibling spend writes
_report_spend_update_step_failure("_record_request_error", payload=payload_copy, exc=e)
self._record_request_error_safely(payload=payload_copy, prisma_client=prisma_client)
await self._run_spend_update_step(
"add_spend_log_transaction_to_daily_org_transaction",
self.add_spend_log_transaction_to_daily_org_transaction(
@ -1601,7 +1688,7 @@ class DBSpendUpdateWriter:
async def _insert_spend_log_to_db(
self,
payload: dict | SpendLogsPayload,
payload: Mapping[str, object],
prisma_client: PrismaClient | None = None,
spend_logs_url: str | None = os.getenv("SPEND_LOGS_URL"),
llm_router: "Router | None" = None,
@ -1645,6 +1732,13 @@ class DBSpendUpdateWriter:
else:
- Regular flow of this method
"""
from litellm.accounting.worker import flush_accounting_deliveries
try:
await flush_accounting_deliveries(self, prisma_client, proxy_logging_obj)
except Exception: # noqa: BLE001 # pending deliveries survive outages while legacy queues must still drain
verbose_proxy_logger.exception("Accounting delivery flush failed; pending work will retry")
if RedisUpdateBuffer.should_commit_spend_updates_to_redis():
await self._commit_spend_updates_to_db_with_redis(
prisma_client=prisma_client,
@ -2171,6 +2265,39 @@ class DBSpendUpdateWriter:
_handle_spend_update_failure = handle_spend_update_failure
@staticmethod
async def invalidate_accounting_entities(
contributions: tuple["Contribution", ...], proxy_logging_obj: ProxyLogging
) -> None:
await DBSpendUpdateWriter._invalidate_project_caches(
tuple(dict(item.key)["project_id"] for item in contributions if item.table == "LiteLLM_ProjectTable"),
proxy_logging_obj,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
details: Final = _SPEND_METADATA_ADAPTER.validate_python(
proxy_logging_obj.call_details # pyright: ignore[reportUnknownMemberType] # legacy logging metadata is validated at this boundary
)
user_api_key_cache: Final = details.get("user_api_key_cache")
if isinstance(user_api_key_cache, UserApiKeyCache):
with service_target(AUTH_OBJECTS_TARGET):
for item in contributions:
if item.table == "LiteLLM_TeamMembership":
await user_api_key_cache.async_delete_cache(
key=f"team_membership:{dict(item.key)['user_id']}:{dict(item.key)['team_id']}"
)
async def apply_accounting_entities(
self,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
transaction: _SpendTransaction,
updates: DBSpendUpdateTransactions,
) -> None:
await self._commit_spend_updates_to_db(
prisma_client, 0, proxy_logging_obj, updates, outer_transaction=transaction
)
async def _commit_spend_updates_to_db(
self,
prisma_client: PrismaClient,
@ -2178,6 +2305,7 @@ class DBSpendUpdateWriter:
proxy_logging_obj: ProxyLogging,
db_spend_update_transactions: DBSpendUpdateTransactions,
on_table_committed: Callable[[_SpendTableName], None] | None = None,
outer_transaction: _SpendTransaction | None = None,
):
"""
Commits all the spend `UPDATE` transactions to the Database
@ -2192,7 +2320,9 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(prisma_client, "LiteLLM_UserTable") as transaction:
async with _spend_update_tx(
prisma_client, "LiteLLM_UserTable", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by ID for consistent lock ordering across pods to prevent deadlocks.
# batch_() issues statements sequentially within the tx, so iteration
@ -2204,6 +2334,8 @@ class DBSpendUpdateWriter:
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
@ -2223,6 +2355,7 @@ class DBSpendUpdateWriter:
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
end_user_list_transactions=end_user_list_transactions,
outer_transaction=outer_transaction,
)
if on_table_committed is not None:
on_table_committed("end_user_list_transactions")
@ -2233,7 +2366,9 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(prisma_client, "LiteLLM_VerificationToken") as transaction:
async with _spend_update_tx(
prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by token for consistent lock ordering across pods to prevent deadlocks.
for token, response_cost in sorted(key_list_transactions.items()):
@ -2248,6 +2383,8 @@ class DBSpendUpdateWriter:
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
@ -2265,7 +2402,9 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(prisma_client, "LiteLLM_TeamTable") as transaction:
async with _spend_update_tx(
prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by team_id for consistent lock ordering across pods to prevent deadlocks.
for team_id, response_cost in sorted(team_list_transactions.items()):
@ -2278,6 +2417,8 @@ class DBSpendUpdateWriter:
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
@ -2303,11 +2444,15 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(prisma_client, "LiteLLM_TeamMembership") as transaction:
async with _spend_update_tx(
prisma_client, "LiteLLM_TeamMembership", outer_transaction=outer_transaction
) as transaction:
await _write_team_member_spend(transaction, team_member_list_transactions)
# Transaction succeeded, break out of retry loop
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
@ -2320,7 +2465,7 @@ class DBSpendUpdateWriter:
# Invalidate cache for updated team memberships
# This ensures budget checks read fresh spend data from the database
if team_memberships_to_invalidate and proxy_logging_obj is not None:
if outer_transaction is None and team_memberships_to_invalidate and proxy_logging_obj is not None:
user_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache")
if user_api_key_cache is not None:
with service_target(AUTH_OBJECTS_TARGET):
@ -2340,7 +2485,9 @@ class DBSpendUpdateWriter:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(prisma_client, "LiteLLM_OrganizationTable") as transaction:
async with _spend_update_tx(
prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by org_id for consistent lock ordering across pods to prevent deadlocks.
for org_id, response_cost in sorted(org_list_transactions.items()):
@ -2350,6 +2497,8 @@ class DBSpendUpdateWriter:
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
@ -2367,7 +2516,9 @@ class DBSpendUpdateWriter:
start_time = time.time()
try:
async with (
_spend_update_tx(prisma_client, "LiteLLM_OrganizationMembership") as transaction,
_spend_update_tx(
prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction
) as transaction,
transaction.batch_() as batcher,
):
for key, response_cost in sorted(org_member_list_transactions.items()):
@ -2378,6 +2529,8 @@ class DBSpendUpdateWriter:
)
break
except Exception as e:
if outer_transaction is not None:
raise
await self.handle_spend_update_failure(
e=e,
attempt=i,
@ -2393,6 +2546,7 @@ class DBSpendUpdateWriter:
await DBSpendUpdateWriter._update_entity_spend_in_db(
entity_name="Project",
transactions=project_list_transactions,
outer_transaction=outer_transaction,
table_accessor="litellm_projecttable",
where_field="project_id",
n_retry_times=n_retry_times,
@ -2401,16 +2555,18 @@ class DBSpendUpdateWriter:
)
if on_table_committed is not None:
on_table_committed("project_list_transactions")
await DBSpendUpdateWriter._invalidate_project_caches(
project_ids=tuple(project_list_transactions or ()),
proxy_logging_obj=proxy_logging_obj,
)
if outer_transaction is None:
await DBSpendUpdateWriter._invalidate_project_caches(
project_ids=tuple(project_list_transactions or ()),
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE TAG TABLE ###
tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"]
await DBSpendUpdateWriter._update_entity_spend_in_db(
entity_name="Tag",
transactions=tag_list_transactions,
outer_transaction=outer_transaction,
table_accessor="litellm_tagtable",
where_field="tag_name",
n_retry_times=n_retry_times,
@ -2427,6 +2583,7 @@ class DBSpendUpdateWriter:
await DBSpendUpdateWriter._update_entity_spend_in_db(
entity_name="Model access group",
transactions=model_access_group_list_transactions,
outer_transaction=outer_transaction,
table_accessor="litellm_modelaccessgroupbudgettable",
where_field="access_group_name",
n_retry_times=n_retry_times,
@ -2441,6 +2598,7 @@ class DBSpendUpdateWriter:
await DBSpendUpdateWriter._update_entity_spend_in_db(
entity_name="Agent",
transactions=agent_list_transactions,
outer_transaction=outer_transaction,
table_accessor="litellm_agentstable",
where_field="agent_id",
n_retry_times=n_retry_times,
@ -2470,6 +2628,7 @@ class DBSpendUpdateWriter:
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
outer_transaction: _SpendTransaction | None = None,
):
"""
Helper function to update spend for any entity type (team, org, tag, etc).
@ -2482,13 +2641,16 @@ class DBSpendUpdateWriter:
n_retry_times: Number of retries on failure
prisma_client: Prisma client instance
proxy_logging_obj: Proxy logging object
outer_transaction: Existing ledger transaction owning the complete delivery update
"""
verbose_proxy_logger.debug("%s Spend transactions: %s", entity_name, transactions)
if transactions is not None and len(transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(prisma_client, _ENTITY_SPEND_MODELS[table_accessor]) as transaction:
async with _spend_update_tx(
prisma_client, _ENTITY_SPEND_MODELS[table_accessor], outer_transaction=outer_transaction
) as transaction:
async with transaction.batch_() as batcher:
# Sort by entity_id for consistent lock ordering across pods to prevent deadlocks.
for entity_id, response_cost in sorted(transactions.items()):
@ -2505,6 +2667,8 @@ class DBSpendUpdateWriter:
)
break
except Exception as e:
if outer_transaction is not None:
raise
await DBSpendUpdateWriter.handle_spend_update_failure(
e=e,
attempt=i,
@ -2838,6 +3002,12 @@ class DBSpendUpdateWriter:
entity_id_field="tag",
)
def _record_request_error_safely(self, *, payload: SpendLogsPayload, prisma_client: PrismaClient | None) -> None:
try:
self._record_request_error(payload=payload, prisma_client=prisma_client)
except Exception as error: # noqa: BLE001 # error rollup failure must not skip sibling writes
_report_spend_update_step_failure("_record_request_error", payload=payload, exc=error)
@staticmethod
def _record_request_error(*, payload: SpendLogsPayload, prisma_client: PrismaClient | None) -> None:
if prisma_client is None:

15
litellm/proxy/db/exception_handler.py
View file

@ -53,6 +53,7 @@ _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024"
_LOCK_TIMEOUT_SQLSTATE: Final = "55P03"
_LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout"
_STATEMENT_TIMEOUT_SQLSTATE: Final = "57014"
_REJECTED_ROW_SQLSTATE_CLASSES: Final = frozenset({"22", "23", "54"})
_STATEMENT_TIMEOUT_PHRASE: Final = "canceling statement due to statement timeout"
@ -221,6 +222,20 @@ class PrismaDBExceptionHandler:
return type(e) is prisma.errors.DataError
@staticmethod
def is_rejected_row_error(e: Exception) -> bool:
"""True iff Postgres refused the rows themselves: a base ``DataError``, or a raw query failing with a data
exception (SQLSTATE class 22), integrity violation (class 23) or a value too large for an index (class 54). Systemic raw failures such as a missing
table (42P01) or a deadlock (40P01) stay False so the batch is retried instead of dropped."""
import prisma
if PrismaDBExceptionHandler.is_prisma_data_error(e):
return True
if not isinstance(e, _exception_types(prisma.errors.RawQueryError)):
return False
sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e)
return sqlstate is not None and sqlstate[:2] in _REJECTED_ROW_SQLSTATE_CLASSES
@staticmethod
def is_database_transport_error(e: Exception) -> bool:
"""

55
litellm/proxy/db/model_usage_rollup.py
View file

@ -32,14 +32,15 @@ class _UpsertTable(Protocol):
def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
class _ModelUsageBatch(Protocol):
litellm_dailymodelusage: _UpsertTable
class ModelUsageBatch(Protocol):
@property
def litellm_dailymodelusage(self) -> _UpsertTable: ...
def execute_raw(self, query: LiteralString, *args: object) -> None: ...
class _ModelUsageBatchManager(Protocol):
async def __aenter__(self) -> _ModelUsageBatch: ...
class ModelUsageBatchManager(Protocol):
async def __aenter__(self) -> ModelUsageBatch: ...
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
@ -116,8 +117,8 @@ def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransa
)
def _model_usage_batch(prisma_client: PrismaClient) -> _ModelUsageBatchManager:
batch: Final[_ModelUsageBatchManager] = prisma_client.db.batch_()
def _model_usage_batch(prisma_client: PrismaClient) -> ModelUsageBatchManager:
batch: Final[ModelUsageBatchManager] = prisma_client.db.batch_()
return batch
@ -156,26 +157,16 @@ async def flush_model_usage_transactions(
"custom_llm_provider": key.custom_llm_provider,
"task_type": key.task_type,
}
batcher.litellm_dailymodelusage.upsert(
where={"date_model_group_model_custom_llm_provider_task_type": key_fields},
data={
"create": {
**key_fields,
"spend": spend,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"request_count": len(requests),
"successful_requests": successful,
"failed_requests": failed,
},
"update": {
"spend": {"increment": spend},
"prompt_tokens": {"increment": prompt_tokens},
"completion_tokens": {"increment": completion_tokens},
"request_count": {"increment": len(requests)},
"successful_requests": {"increment": successful},
"failed_requests": {"increment": failed},
},
queue_model_usage_upsert(
batcher,
key_fields,
{
"spend": spend,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"request_count": len(requests),
"successful_requests": successful,
"failed_requests": failed,
},
)
return
@ -183,3 +174,15 @@ async def flush_model_usage_transactions(
if attempt >= n_retry_times:
raise
await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
def queue_model_usage_upsert(
batcher: ModelUsageBatch, key_fields: Mapping[str, str], increments: Mapping[str, int | float]
) -> None:
batcher.litellm_dailymodelusage.upsert(
where={"date_model_group_model_custom_llm_provider_task_type": key_fields},
data={
"create": {**key_fields, **increments},
"update": {name: {"increment": value} for name, value in increments.items()},
},
)

61
litellm/proxy/guardrails/usage_tracking.py
View file

@ -11,7 +11,7 @@ from datetime import datetime, timezone
from functools import partial
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeVar
from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict
@ -22,6 +22,7 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import billed_guard
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
from litellm.proxy.db.spend_log_batching import spend_log_write_batches
from litellm.proxy.utils import PrismaClient
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
DailyGuardrailMetricsRepository,
DailyGuardrailUsageUnitsRepository,
@ -40,7 +41,7 @@ _RowKey = TypeVar("_RowKey")
_RowValue = TypeVar("_RowValue")
class _UsageUnitKey(NamedTuple):
class UsageUnitKey(NamedTuple):
guardrail_id: str
date: str
team_id: str
@ -48,7 +49,7 @@ class _UsageUnitKey(NamedTuple):
usage_unit: str
class _UsageUnitIncrement(NamedTuple):
class UsageUnitIncrement(NamedTuple):
units: int
cost: float
"""USD for the priced share of units."""
@ -56,14 +57,14 @@ class _UsageUnitIncrement(NamedTuple):
"""Units recorded with no known price, the share cost leaves out."""
def _usage_unit_key(entry: tuple[_UsageUnitKey, _UsageUnitIncrement]) -> _UsageUnitKey:
def _usage_unit_key(entry: tuple[UsageUnitKey, UsageUnitIncrement]) -> UsageUnitKey:
return entry[0]
def _usage_unit_increment(units: int, cost: float | None) -> _UsageUnitIncrement:
def _usage_unit_increment(units: int, cost: float | None) -> UsageUnitIncrement:
if cost is None:
return _UsageUnitIncrement(units=units, cost=0.0, untracked_units=units)
return _UsageUnitIncrement(units=units, cost=cost, untracked_units=0)
return UsageUnitIncrement(units=units, cost=0.0, untracked_units=units)
return UsageUnitIncrement(units=units, cost=cost, untracked_units=0)
class _MetricsKey(NamedTuple):
@ -89,22 +90,22 @@ class PendingRollups:
def __init__(self) -> None:
self.lock: Final = asyncio.Lock()
self.metrics: Mapping[_MetricsKey, Mapping[str, int]] = MappingProxyType({})
self.units: Mapping[_UsageUnitKey, _UsageUnitIncrement] = MappingProxyType({})
self.units: Mapping[UsageUnitKey, UsageUnitIncrement] = MappingProxyType({})
_PENDING_ROLLUPS: Final = PendingRollups()
_NO_COUNTERS: Final[Mapping[str, int]] = MappingProxyType({})
_NO_INCREMENT: Final = _UsageUnitIncrement(units=0, cost=0.0, untracked_units=0)
_NO_INCREMENT: Final = UsageUnitIncrement(units=0, cost=0.0, untracked_units=0)
def _merged_keys(base: Mapping[_RowKey, object], extra: Mapping[_RowKey, object]) -> tuple[_RowKey, ...]:
return (*base, *(key for key in extra if key not in base))
def _summed_increments(increments: Iterable[_UsageUnitIncrement]) -> _UsageUnitIncrement:
def _summed_increments(increments: Iterable[UsageUnitIncrement]) -> UsageUnitIncrement:
materialized: Final = tuple(increments)
return _UsageUnitIncrement(
return UsageUnitIncrement(
units=sum(i.units for i in materialized),
cost=sum(i.cost for i in materialized),
untracked_units=sum(i.untracked_units for i in materialized),
@ -112,8 +113,8 @@ def _summed_increments(increments: Iterable[_UsageUnitIncrement]) -> _UsageUnitI
def _merged_unit_rows(
base: Mapping[_UsageUnitKey, _UsageUnitIncrement], extra: Mapping[_UsageUnitKey, _UsageUnitIncrement]
) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]:
base: Mapping[UsageUnitKey, UsageUnitIncrement], extra: Mapping[UsageUnitKey, UsageUnitIncrement]
) -> Mapping[UsageUnitKey, UsageUnitIncrement]:
return MappingProxyType(
{
key: _summed_increments((base.get(key, _NO_INCREMENT), extra.get(key, _NO_INCREMENT)))
@ -250,7 +251,7 @@ def _parse_payload_start_time(payload: Mapping[str, object]) -> datetime | None:
def _iter_usage_unit_increments(
logs_to_process: Sequence[Mapping[str, object]],
) -> Iterator[tuple[_UsageUnitKey, _UsageUnitIncrement]]:
) -> Iterator[tuple[UsageUnitKey, UsageUnitIncrement]]:
for payload in logs_to_process:
start_time = _parse_payload_start_time(payload)
if not payload.get("request_id") or start_time is None:
@ -266,14 +267,14 @@ def _iter_usage_unit_increments(
cost_by_unit = billed_guardrail_cost_by_unit(entry)
for unit_name, units in usage.items():
if isinstance(units, int) and not isinstance(units, bool) and units > 0:
key = _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name))
key = UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name))
cost = cost_by_unit.get(str(unit_name)) if cost_by_unit is not None else None
yield key, _usage_unit_increment(units=units, cost=cost)
def _sum_usage_unit_increments(
logs_to_process: Sequence[Mapping[str, object]],
) -> Mapping[_UsageUnitKey, _UsageUnitIncrement]:
) -> Mapping[UsageUnitKey, UsageUnitIncrement]:
ordered: Final = sorted(_iter_usage_unit_increments(logs_to_process), key=_usage_unit_key)
return MappingProxyType(
{
@ -283,8 +284,16 @@ def _sum_usage_unit_increments(
)
async def _upsert_usage_unit_row(
prisma_client: PrismaClient, key: _UsageUnitKey, increment: _UsageUnitIncrement
class GuardrailUsageDatabase(Protocol):
@property
def litellm_dailyguardrailusageunits(self) -> TableActions[object]: ...
async def upsert_usage_unit_row(
prisma_client: PrismaClient,
key: UsageUnitKey,
increment: UsageUnitIncrement,
outer_transaction: GuardrailUsageDatabase | None = None,
) -> None:
row: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsCreateInput] = {
"guardrail_id": key.guardrail_id,
@ -314,7 +323,12 @@ async def _upsert_usage_unit_row(
"untracked_units": {"increment": increment.untracked_units},
},
}
await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data)
table: Final = (
outer_transaction.litellm_dailyguardrailusageunits
if outer_transaction is not None
else DailyGuardrailUsageUnitsRepository(prisma_client).table
)
await table.upsert(where=where, data=data)
async def _upsert_metrics_row(prisma_client: PrismaClient, key: _MetricsKey, agg: Mapping[str, int]) -> None:
@ -363,7 +377,7 @@ async def process_spend_logs_guardrail_usage(
)
index_rows_by_key: Final[dict[tuple[str, str], dict[str, object]]] = {}
for payload in logs_to_process:
for payload in (row for row in logs_to_process if row.get("accounting_log_disabled") is not True):
request_id = payload.get("request_id")
start_time = _parse_payload_start_time(payload)
if not isinstance(request_id, str) or not request_id or start_time is None:
@ -416,7 +430,10 @@ async def process_spend_logs_guardrail_usage(
{key: agg for key, agg in daily_guardrail.items() if int(agg["requests_evaluated"]) > 0}
)
metrics_rows: Final = _merged_metric_rows(pending_metrics, evaluated_metrics)
unit_rows: Final = _merged_unit_rows(pending_units, _sum_usage_unit_increments(logs_to_process))
unit_rows: Final = _merged_unit_rows(
pending_units,
_sum_usage_unit_increments(tuple(row for row in logs_to_process if "accounting_payload" not in row)),
)
if not metrics_rows and not index_rows and not unit_rows:
return
@ -435,7 +452,7 @@ async def process_spend_logs_guardrail_usage(
metrics_rows, partial(_upsert_metrics_row, prisma_client), "daily metrics", sleep
)
failed_units: Final = await _upsert_rows_with_retry(
unit_rows, partial(_upsert_usage_unit_row, prisma_client), "usage unit", sleep
unit_rows, partial(upsert_usage_unit_row, prisma_client), "usage unit", sleep
)
if failed_metrics or failed_units:
async with pending.lock:

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

@ -9,6 +9,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast
from pydantic import TypeAdapter, ValidationError
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
@ -38,6 +39,7 @@ from litellm.proxy.auth.auth_checks import (
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded
from litellm.proxy.db.db_spend_update_writer import (
ACCOUNTING_LEDGER_KWARG,
DBSpendUpdateWriter,
debitable_model_access_groups,
get_llm_router,
@ -657,6 +659,7 @@ class _ProxyDBLogger(CustomLogger):
token: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
end_user_id: Final = None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id
accounting_rows: Final = tuple((cost, dict(row)) for cost, row in spend_rows)
charged: Final = [
await self._spend_writer().update_database(
token=token,
@ -664,17 +667,17 @@ class _ProxyDBLogger(CustomLogger):
user_id=user_api_key_dict.user_id,
end_user_id=end_user_id,
team_id=user_api_key_dict.team_id,
kwargs=dict(row_kwargs),
kwargs=row_kwargs,
completion_response=original_exception,
start_time=start_time,
end_time=datetime.now(),
org_id=user_api_key_dict.org_id,
project_id=user_api_key_dict.project_id,
)
for row_cost, row_kwargs in spend_rows
for row_cost, row_kwargs in accounting_rows
]
counted_rows: Final = tuple(
row for row, row_charged in zip(spend_rows, charged, strict=True) if row_charged and row[0] > 0
row for row, row_charged in zip(accounting_rows, charged, strict=True) if row_charged and row[0] > 0
)
if counted_rows and not ProxyUpdateSpend.disable_spend_updates():
await self._increment_failure_spend_counters(
@ -753,6 +756,7 @@ class _ProxyDBLogger(CustomLogger):
router=get_llm_router(),
),
project_id=user_api_key_dict.project_id,
**_ledger_counter_options(row_kwargs.get(ACCOUNTING_LEDGER_KWARG)),
)
except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing
verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row")
@ -1196,6 +1200,14 @@ def _get_request_tags_for_cost_tracking(
return None
class _LedgerCounterOptions(TypedDict):
windows_from_ledger: ReadOnly[NotRequired[bool]]
def _ledger_counter_options(ledger: object) -> _LedgerCounterOptions:
return {"windows_from_ledger": True} if ledger is True else {}
class _ChargeSpend(Protocol):
async def __call__(
self,
@ -1256,6 +1268,7 @@ class _IncrementSpendCounters(Protocol):
request_started_at: datetime | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
windows_from_ledger: bool = False,
) -> None: ...
@ -1393,6 +1406,7 @@ async def _update_database_and_spend_counters_in_batch(
request_started_at=start_time,
model_access_groups=model_access_groups,
project_id=project_id,
**_ledger_counter_options(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get(ACCOUNTING_LEDGER_KWARG)),
)
except Exception:
if budget_reservation is not None:

66
litellm/proxy/proxy_server.py
View file

@ -3617,6 +3617,7 @@ async def increment_spend_counters(
request_started_at: datetime | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
windows_from_ledger: bool = False,
):
"""
Atomically increment spend counters for budget enforcement.
@ -3653,6 +3654,7 @@ async def increment_spend_counters(
request_started_at=request_started_at,
model_access_groups=model_access_groups,
project_id=project_id,
windows_from_ledger=windows_from_ledger,
)
@ -3682,6 +3684,50 @@ def _finish_pending_read(absent_key: str, pending_read: str, remember_absent: bo
user_api_key_cache.budget_windows_absent_cache.delete_cache(absent_key)
async def accounting_window_keys(token: str | None, team_id: str | None) -> tuple[Mapping[str, object], ...]:
scopes: Final = tuple(
(entity_type, entity_id, cache_key)
for entity_type, entity_id, cache_key in (
(Litellm_EntityType.KEY, token, token),
(Litellm_EntityType.TEAM, team_id, f"team_id:{team_id}"),
)
if entity_id is not None and cache_key is not None
)
snapshots: Final = tuple(
[
await _accounting_scope_windows(entity_type, entity_id, cache_key)
for entity_type, entity_id, cache_key in scopes
]
)
from itertools import chain
return tuple(chain.from_iterable(snapshots))
async def _accounting_scope_windows(
entity_type: Litellm_EntityType, entity_id: str, cache_key: str
) -> tuple[Mapping[str, object], ...]:
from litellm.proxy.spend_tracking.budget_reservation import fixed_budget_window
cached: Final[object] = _UNTYPED_VALUE.validate_python(await user_api_key_cache.async_get_cache(key=cache_key))
limits: Final = (
await _stored_budget_limits(entity_type, entity_id) if cached is None else _cached_budget_limits(cached)
)
if not isinstance(limits, list):
return ()
windows: Final = TypeAdapter(tuple[object, ...]).validate_python(limits)
return tuple(
{
"entity_type": entity_type.value,
"entity_id": entity_id,
"window_duration": fixed[0],
"window_start": fixed[1].isoformat(),
}
for window in windows
if (fixed := fixed_budget_window(window)) is not None
)
async def _stored_budget_limits(entity_type: Litellm_EntityType, entity_id: str) -> object:
"""Read only: caching the row here could race a delete and put a just-deleted key back into the auth cache."""
if prisma_client is None:
@ -3723,6 +3769,7 @@ async def _increment_spend_counters_batched(
request_started_at: datetime | None,
model_access_groups: Sequence[str] | None,
project_id: str | None = None,
windows_from_ledger: bool = False,
):
"""Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET, and
the reconcile adjustments go out in the same INCRBYFLOAT pipeline as the counter increments."""
@ -3764,15 +3811,16 @@ async def _increment_spend_counters_batched(
if window_counter not in reserved_counter_keys
else None
)
await _enqueue_window_spend_row_update(
entity_type=entity_type,
entity_id=entity_id,
reset_at=window_reset_at,
window_duration=duration,
window_start=window_start,
increment=cost,
request_started_at=request_started_at,
)
if not windows_from_ledger:
await _enqueue_window_spend_row_update(
entity_type=entity_type,
entity_id=entity_id,
reset_at=window_reset_at,
window_duration=duration,
window_start=window_start,
increment=cost,
request_started_at=request_started_at,
)
return pending_window
async def _window_increments(

17
litellm/proxy/schema.prisma
View file

@ -2141,6 +2141,23 @@ model LiteLLM_AccountingDelivery {
@@index([target, state, next_attempt_at])
}
model LiteLLM_SpendLogDelivery {
delivery_id String @id @default(dbgenerated("gen_random_uuid()::text"))
request_id String
state AccountingDeliveryState @default(pending)
accounting_payload Json
is_batch Boolean
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
created_at DateTime @default(now())
@@index([state, next_attempt_at])
@@index([state, applied_at])
@@index([request_id])
}
model LiteLLM_SubscriptionAccountTable {
subscription_account_id String @id @default(uuid())
custom_llm_provider String

9
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -1264,6 +1264,15 @@ def _get_entry_reserved_cost(entry: dict, default_reserved_cost: float) -> float
return default_reserved_cost
def fixed_budget_window(window: object) -> tuple[str, datetime] | None:
fields: Final = _coerce_window(window)
start: Final = get_budget_window_start(window)
duration: Final = fields.get("budget_duration")
if not fields.get("reset_at") or not isinstance(duration, str) or start is None:
return None
return duration, start
def get_budget_window_start(window: object) -> datetime | None:
window_dict: Final = _coerce_window(window)
budget_duration: Final = window_dict.get("budget_duration")

10
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -161,9 +161,7 @@ def _is_non_secret_key_value(value: str, *, key_alias: object = None) -> bool:
)
def _redact_logged_api_key(
value: str | None, *, already_redacted: bool = False, key_alias: object = None
) -> str | None:
def redact_logged_api_key(value: str | None, *, already_redacted: bool = False, key_alias: object = None) -> str | None:
if not isinstance(value, str) or not value:
return None
stripped: Final = re.sub(r"(?i)^bearer ", "", value)
@ -327,7 +325,7 @@ def _get_spend_logs_metadata(
and _is_non_secret_key_value(_trusted_hash, key_alias=_key_alias)
and _trusted_hash == _raw_key
)
clean_metadata["user_api_key"] = _redact_logged_api_key(
clean_metadata["user_api_key"] = redact_logged_api_key(
_raw_key, already_redacted=_already_redacted, key_alias=_key_alias
)
clean_metadata["applied_guardrails"] = applied_guardrails
@ -692,14 +690,14 @@ def get_logging_payload(
and _is_non_secret_key_value(_trusted_hash, key_alias=_key_alias)
and _trusted_hash == api_key
)
api_key = _redact_logged_api_key(api_key, already_redacted=_key_already_redacted, key_alias=_key_alias) or ""
api_key = redact_logged_api_key(api_key, already_redacted=_key_already_redacted, key_alias=_key_alias) or ""
if (
standard_logging_payload is not None
): # [TODO] migrate completely to sl payload. currently missing pass-through endpoint data
api_key = (
api_key
or _redact_logged_api_key(
or redact_logged_api_key(
standard_logging_payload["metadata"].get("user_api_key_hash"),
already_redacted=True,
key_alias=standard_logging_payload["metadata"].get("user_api_key_alias"),

96
litellm/proxy/utils.py
View file

@ -278,6 +278,7 @@ from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.spend_log_insert import RejectedSpendLogRow
from litellm.repositories.table_repositories import (
EndUserRepository,
HealthCheckRepository,
@ -8418,14 +8419,17 @@ async def enqueue_spend_logs(
"""
added: Final = sum(spend_log_row_bytes(row) for row in logs)
async with prisma_client._spend_log_transactions_lock:
queued: Final = (
queued: Final[tuple[Mapping[str, object], ...]] = (
tuple(logs) + tuple(prisma_client.spend_log_transactions)
if at_head
else tuple(prisma_client.spend_log_transactions) + tuple(logs)
)
kept, kept_bytes = spend_log_queue_within_budget(queued, PrismaClient.spend_log_queue_bytes + added, max_bytes)
prisma_client.spend_log_transactions[:] = kept
PrismaClient.spend_log_queue_bytes = kept_bytes
rescued: Final = tuple(
accounting_only_spend_log(row) for row in queued[: len(queued) - len(kept)] if _carries_a_charge(row)
)
prisma_client.spend_log_transactions[:] = (*rescued, *kept)
PrismaClient.spend_log_queue_bytes = kept_bytes + sum(spend_log_row_bytes(row) for row in rescued)
if len(kept) < len(queued):
verbose_proxy_logger.error(
"Spend tracking - spend log queue is at its %d byte budget; dropped the %d oldest spend logs",
@ -8480,13 +8484,18 @@ class ProxyUpdateSpend:
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
end_user_list_transactions: dict[str, float],
outer_transaction: _EndUserSpendTransaction | None = None,
):
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with (
db_span("update_end_user_spend", "LiteLLM_EndUserTable"),
_end_user_spend_transaction(prisma_client.db, timedelta(seconds=60)) as transaction,
(
contextlib.nullcontext(outer_transaction)
if outer_transaction is not None
else _end_user_spend_transaction(prisma_client.db, timedelta(seconds=60))
) as transaction,
):
batcher: _EndUserSpendBatch
async with transaction.batch_() as batcher:
@ -8508,6 +8517,8 @@ class ProxyUpdateSpend:
break
except Exception as e:
if outer_transaction is not None:
raise
await DBSpendUpdateWriter.handle_spend_update_failure(
e=e,
attempt=i,
@ -8547,11 +8558,18 @@ class ProxyUpdateSpend:
for i in range(n_retry_times + 1):
try:
base_url = os.getenv("SPEND_LOGS_URL", None)
if len(logs_to_process) > 0 and base_url is not None and db_writer_client is not None:
await _commit_export_accounting(
prisma_client, rows_to_write, base_url if db_writer_client is not None else None
)
if (
base_url is not None
and db_writer_client is not None
and any(row.get("accounting_log_disabled") is not True for row in rows_to_write)
):
if not base_url.endswith("/"):
base_url += "/"
verbose_proxy_logger.debug("base_url: %s", base_url)
json_data = json.dumps(rows_to_write)
json_data = json.dumps(_spend_log_export_rows(rows_to_write))
response = await db_writer_client.post(
url=base_url + "spend/update",
data=json_data,
@ -8629,6 +8647,28 @@ class ProxyUpdateSpend:
return False
def _spend_log_export_rows(rows: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]:
return tuple(
{
key: value
for key, value in row.items()
if key not in {"accounting_payload", "accounting_log_disabled", "accounting_delivery_id"}
}
for row in rows
if row.get("accounting_log_disabled") is not True
)
async def _commit_export_accounting(
prisma_client: PrismaClient, rows: Sequence[Mapping[str, object]], base_url: str | None
) -> None:
if base_url is None:
return
await SpendLogsRepository(prisma_client).insert_with_deliveries(
tuple(row for row in rows if "accounting_payload" in row), write_spend_logs=False
)
async def commit_tool_discoveries(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None:
"""Scheduler job writing discovered tools to LiteLLM_ToolTable, on its own lock, next to update_spend."""
writer: Final = proxy_logging_obj.db_spend_update_writer
@ -9127,6 +9167,40 @@ def _is_transient_spend_log_write_error(e: Exception) -> bool:
)
_ACCOUNTING_ONLY_KEYS: Final = frozenset({"request_id", "call_type", "accounting_payload", "accounting_delivery_id"})
def _carries_a_charge(row: Mapping[str, object]) -> bool:
return "accounting_payload" in row
def accounting_only_spend_log(row: Mapping[str, object]) -> Mapping[str, object]:
return {
**{key: value for key, value in row.items() if key in _ACCOUNTING_ONLY_KEYS},
"accounting_log_disabled": True,
}
async def _keep_charges_of_dropped_logs(repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]]) -> None:
charged: Final = tuple(accounting_only_spend_log(row) for row in rows if _carries_a_charge(row))
if not charged:
return
try:
await repo.insert_with_deliveries(charged)
except Exception as e: # noqa: BLE001 # the request log is already lost; keep the flush going for other rows
if not (
PrismaDBExceptionHandler.is_rejected_row_error(e)
or isinstance(e, (UnicodeEncodeError, RejectedSpendLogRow))
):
raise
spend_log_error(
"Spend tracking - could not keep the charges of %d dropped spend log rows. error=%s",
len(charged),
str(e),
exc=e,
)
async def _create_spend_logs_with_poison_isolation(
repo: SpendLogsRepository,
rows: Sequence[Mapping[str, object]],
@ -9155,10 +9229,13 @@ async def _create_spend_logs_with_poison_isolation(
persist. Returns the budget left after this subtree.
"""
try:
await repo.table.create_many(data=rows, skip_duplicates=True)
await repo.insert_with_deliveries(rows)
return failure_budget
except Exception as e:
if not (PrismaDBExceptionHandler.is_prisma_data_error(e) or isinstance(e, UnicodeEncodeError)):
if not (
PrismaDBExceptionHandler.is_rejected_row_error(e)
or isinstance(e, (UnicodeEncodeError, RejectedSpendLogRow))
):
raise
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
raise
@ -9173,6 +9250,7 @@ async def _create_spend_logs_with_poison_isolation(
str(e),
exc=e,
)
await _keep_charges_of_dropped_logs(repo, rows)
return budget_left
if budget_left <= 0:
spend_log_error(
@ -9180,6 +9258,7 @@ async def _create_spend_logs_with_poison_isolation(
"isolation failure budget exhausted for this flush",
len(rows),
)
await _keep_charges_of_dropped_logs(repo, rows)
return 0
mid: Final = len(rows) // 2
remaining: Final = await _create_spend_logs_with_poison_isolation(repo, rows[:mid], budget_left)
@ -9189,6 +9268,7 @@ async def _create_spend_logs_with_poison_isolation(
"isolation failure budget exhausted for this flush",
len(rows) - mid,
)
await _keep_charges_of_dropped_logs(repo, rows[mid:])
return 0
return await _create_spend_logs_with_poison_isolation(repo, rows[mid:], remaining)

193
litellm/repositories/accounting_delivery_repository.py Normal file
View file

@ -0,0 +1,193 @@
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import timedelta
from itertools import chain
from types import TracebackType
from typing import TYPE_CHECKING, Final, Literal, Protocol
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
from litellm.repositories.prisma_protocols import TableActions
if TYPE_CHECKING:
from litellm.proxy.db.db_spend_update_writer import SpendBatchManager
class LedgerTransaction(Protocol):
@property
def litellm_dailyguardrailusageunits(self) -> TableActions[object]: ...
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ...
async def execute_raw(self, query: str, *args: object) -> int: ...
def batch_(self) -> "SpendBatchManager": ...
class _TransactionManager(Protocol):
async def __aenter__(self) -> LedgerTransaction: ...
async def __aexit__(
self, exc_type: type[BaseException] | None, exc: BaseException | None, exc_tb: TracebackType | None
) -> None: ...
class LedgerDatabase(Protocol):
def tx(self, *, timeout: timedelta) -> _TransactionManager: ...
class _Delivery(BaseModel):
model_config = ConfigDict(frozen=True)
delivery_id: str
attempts: int
accounting_payload: JsonValue
@dataclass(frozen=True, slots=True)
class DeliveryProjection:
row: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class DeliveryBatchResult:
applied: int = 0
failed: int = 0
parked: tuple[str, ...] = ()
error: str | None = None
class DeliveryMetrics(BaseModel):
model_config = ConfigDict(frozen=True)
pending_count: int
parked_count: int
oldest_pending_seconds: float
batch_count: int
table_bytes: int
_DELIVERIES: Final = TypeAdapter(tuple[_Delivery, ...])
_PAYLOAD: Final = TypeAdapter(dict[str, object])
_TIMEOUT: Final = timedelta(seconds=60)
_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload FROM "LiteLLM_SpendLogDelivery"
WHERE state = 'pending' AND next_attempt_at <= CURRENT_TIMESTAMP
AND (attempts > 0) = $2::boolean
ORDER BY next_attempt_at, delivery_id LIMIT $1::integer FOR UPDATE SKIP LOCKED"""
_APPLIED: Final = """UPDATE "LiteLLM_SpendLogDelivery"
SET state = 'applied', applied_at = CURRENT_TIMESTAMP, last_error = NULL
WHERE delivery_id = ANY($1::text[]) AND state = 'pending' """
_FAILED: Final = """UPDATE "LiteLLM_SpendLogDelivery"
SET attempts = attempts + 1, last_error = $3,
next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)))::integer),
state = $4::"AccountingDeliveryState"
WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer
RETURNING delivery_id, attempts, accounting_payload"""
def _is_permanent_failure(error: Exception) -> bool:
"""A delivery whose own data failed fails the same way on every retry, while any other error is the database's
and clears once it recovers."""
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
if PrismaDBExceptionHandler.is_database_service_unavailable_error(error):
return False
return isinstance(error, (ValueError, TypeError, KeyError)) or PrismaDBExceptionHandler.is_rejected_row_error(error)
class AccountingDeliveryRepository:
def __init__(self, database: LedgerDatabase) -> None:
self._database = database
async def metrics(self) -> DeliveryMetrics:
async with self._database.tx(timeout=_TIMEOUT) as transaction:
rows: Final = await transaction.query_raw(
"SELECT COUNT(*) FILTER (WHERE state = 'pending') AS pending_count, "
"COUNT(*) FILTER (WHERE state = 'parked') AS parked_count, "
"COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at) FILTER (WHERE state = 'pending')), 0) "
"AS oldest_pending_seconds, COUNT(*) FILTER (WHERE is_batch) AS batch_count, "
"pg_total_relation_size('\"LiteLLM_SpendLogDelivery\"') AS table_bytes "
'FROM "LiteLLM_SpendLogDelivery"'
)
return DeliveryMetrics.model_validate(rows[0])
async def cleanup_applied(self, *, retention_days: int = 1, batch_size: int = 1000) -> int:
if retention_days < 1 or batch_size < 1:
raise ValueError("Delivery retention and cleanup batch size must be positive")
async with self._database.tx(timeout=_TIMEOUT) as transaction:
return await transaction.execute_raw(
'DELETE FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id IN ('
'SELECT delivery_id FROM "LiteLLM_SpendLogDelivery" '
"WHERE state = 'applied' AND NOT is_batch AND applied_at < CURRENT_TIMESTAMP - make_interval(days => $1::integer) "
"ORDER BY applied_at, delivery_id LIMIT $2::integer FOR UPDATE SKIP LOCKED)",
retention_days,
batch_size,
)
async def apply_pending(
self,
apply: Callable[[LedgerTransaction, tuple[DeliveryProjection, ...]], Awaitable[None]],
*,
batch_size: int = 100,
max_attempts: int = 5,
retry_only: bool | None = None,
) -> DeliveryBatchResult:
if batch_size < 1 or max_attempts < 1:
raise ValueError("Delivery batch size and maximum attempts must be positive")
claimed: tuple[_Delivery, ...] = () # rebind-ok: preserve claims for retry accounting after rollback
try:
async with self._database.tx(timeout=_TIMEOUT) as transaction:
retries: Final = (
_DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, 1, True))
if retry_only is not False
else ()
)
claimed = retries or ( # rebind-ok: retain claimed identities after rollback
_DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, batch_size, False))
if retry_only is not True
else ()
)
if not claimed:
return DeliveryBatchResult()
await apply(
transaction,
tuple(
DeliveryProjection(_PAYLOAD.validate_python(delivery.accounting_payload))
for delivery in claimed
),
)
await transaction.execute_raw(_APPLIED, tuple(delivery.delivery_id for delivery in claimed))
return DeliveryBatchResult(applied=len(claimed))
except Exception as error: # noqa: BLE001 # roll back all destinations before recording a retry
if not claimed:
raise
return await self._record_failure(claimed, error, max_attempts)
async def _record_failure(
self, claimed: tuple[_Delivery, ...], error: Exception, max_attempts: int
) -> DeliveryBatchResult:
parks: Final = len(claimed) == 1 and _is_permanent_failure(error)
async with self._database.tx(timeout=_TIMEOUT) as transaction:
updates: Final = tuple(
[
await self._record_row_failure(
transaction,
delivery,
type(error).__name__,
"parked" if parks and delivery.attempts + 1 >= max_attempts else "pending",
)
for delivery in claimed
]
)
return DeliveryBatchResult(
failed=len(claimed), parked=tuple(chain.from_iterable(updates)), error=type(error).__name__
)
@staticmethod
async def _record_row_failure(
transaction: LedgerTransaction, delivery: _Delivery, error: str, state: Literal["pending", "parked"]
) -> tuple[str, ...]:
updated: Final = _DELIVERIES.validate_python(
await transaction.query_raw(_FAILED, delivery.delivery_id, delivery.attempts, error, state)
)
return tuple(row.delivery_id for row in updated) if state == "parked" else ()

88
litellm/repositories/spend_log_insert.py Normal file
View file

@ -0,0 +1,88 @@
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import Final
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import LiteralString
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_COLUMNS: Final[tuple[LiteralString, ...]] = (
"request_id",
"call_type",
"api_key",
"spend",
"total_tokens",
"prompt_tokens",
"completion_tokens",
"startTime",
"endTime",
"request_duration_ms",
"completionStartTime",
"model",
"model_id",
"model_group",
"custom_llm_provider",
"api_base",
"user",
"metadata",
"cache_hit",
"cache_key",
"request_tags",
"team_id",
"organization_id",
"end_user",
"requester_ip_address",
"messages",
"response",
"session_id",
"status",
"mcp_namespaced_tool_name",
"billing_agent_id",
"agent_id",
"proxy_server_request",
"litellm_call_id",
"created_at",
"updated_at",
)
_JSON_COLUMNS: Final = frozenset({"metadata", "request_tags", "messages", "response", "proxy_server_request"})
class RejectedSpendLogRow(ValueError):
"""A row the raw insert cannot encode, so it is isolated like a row Postgres refuses."""
def _json_value(column: str, value: object) -> object:
if isinstance(value, datetime):
return value.isoformat()
if column not in _JSON_COLUMNS or not isinstance(value, str):
return value
try:
return _JSON.validate_json(value)
except ValidationError as error:
raise RejectedSpendLogRow(f"Spend log column {column} is not valid JSON") from error
def spend_log_insert_statement(rows: Sequence[Mapping[str, object]]) -> tuple[LiteralString, str]:
if not rows:
raise ValueError("Spend log insert requires at least one row")
fields: Final = frozenset(rows[0]) | {"created_at", "updated_at"}
if not fields.issubset(_COLUMNS) or any(frozenset(row) | {"created_at", "updated_at"} != fields for row in rows):
raise RejectedSpendLogRow("Spend log insert requires matching known columns")
columns: Final = ", ".join('"' + column + '"' for column in _COLUMNS if column in fields)
now: Final = datetime.now(timezone.utc).isoformat()
values: Final = tuple(
{"created_at": now, "updated_at": now, **{column: _json_value(column, value) for column, value in row.items()}}
for row in rows
)
sql: Final = (
'INSERT INTO "LiteLLM_SpendLogs" ('
+ columns
+ ") SELECT "
+ columns
+ ' FROM jsonb_populate_recordset(NULL::"LiteLLM_SpendLogs", $1::jsonb) '
"ON CONFLICT DO NOTHING RETURNING request_id"
)
return sql, safe_dumps(values, non_finite_as_null=True)

115
litellm/repositories/table_repositories.py
View file

@ -7,8 +7,15 @@ These are thin wrappers for tables that do not (yet) need domain-specific query
methods; richer repositories live in their own modules.
"""
from collections.abc import Mapping, Sequence
from datetime import timedelta
from itertools import groupby
from typing import TYPE_CHECKING, Any, Final, Generic
from pydantic import TypeAdapter
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.repositories.accounting_delivery_repository import LedgerDatabase, LedgerTransaction
from litellm.repositories.prisma_protocols import RowT_co, TableActions
if TYPE_CHECKING:
@ -80,9 +87,117 @@ class OrganizationMembershipRepository(PrismaTableRepository["prisma_models.Lite
table_name = "litellm_organizationmembership"
_SPEND_LOG_FIELDS: Final = TypeAdapter(dict[str, object])
def _accounting_payload(value: object) -> Mapping[str, object]:
return (
_SPEND_LOG_FIELDS.validate_json(value) if isinstance(value, str) else _SPEND_LOG_FIELDS.validate_python(value)
)
_ACCOUNTING_ONLY_FIELDS: Final = frozenset({"accounting_payload", "accounting_log_disabled", "accounting_delivery_id"})
def _settles_once(row: Mapping[str, object]) -> bool:
"""A batch or video cost row: update_database gives every other ledger row its own delivery id."""
return "accounting_delivery_id" not in row
def _delivery_id(row: Mapping[str, object]) -> str:
return str(row["accounting_delivery_id"]) if not _settles_once(row) else str(row["request_id"])
class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs"]):
table_name = "litellm_spendlogs"
async def insert_with_deliveries(
self, rows: Sequence[Mapping[str, object]], *, write_spend_logs: bool = True
) -> int:
legacy: Final = tuple(row for row in rows if "accounting_payload" not in row)
legacy_count: Final = await self.table.create_many(data=legacy, skip_duplicates=True) if legacy else 0
ledger: Final = tuple(
sorted(
{_delivery_id(row): row for row in reversed(rows) if "accounting_payload" in row}.values(),
key=_delivery_id,
)
)
if not ledger:
return legacy_count
database: Final[LedgerDatabase] = self.prisma_client.writer_db
async with database.tx(timeout=timedelta(seconds=60)) as transaction:
accepted: Final = await self._insert_ledger(transaction, ledger, write_spend_logs=write_spend_logs)
return legacy_count + accepted
@staticmethod
async def _insert_ledger(
transaction: LedgerTransaction, rows: Sequence[Mapping[str, object]], *, write_spend_logs: bool
) -> int:
from litellm.proxy.db.db_span import db_span
from litellm.repositories.spend_log_insert import spend_log_insert_statement
payloads: Final = tuple(
{
"delivery_id": _delivery_id(row),
"request_id": row["request_id"],
"accounting_payload": _accounting_payload(row["accounting_payload"]),
"is_batch": _settles_once(row),
}
for row in rows
)
inserted_deliveries: Final = await transaction.query_raw(
'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch) '
"SELECT delivery_id, request_id, accounting_payload, is_batch FROM jsonb_to_recordset($1::jsonb) "
"AS inputs(delivery_id text, request_id text, accounting_payload jsonb, is_batch boolean) "
"ON CONFLICT (delivery_id) DO NOTHING RETURNING delivery_id",
safe_dumps(payloads, non_finite_as_null=True),
)
accepted_ids: Final = frozenset(str(row["delivery_id"]) for row in inserted_deliveries)
charged_by_legacy: Final = await SpendLogsRepository._jobs_charged_by_legacy(
transaction,
tuple(_delivery_id(row) for row in rows if _settles_once(row) and _delivery_id(row) in accepted_ids),
)
if charged_by_legacy:
await transaction.execute_raw(
'DELETE FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id = ANY($1::text[])', list(charged_by_legacy)
)
accepted: Final = accepted_ids - charged_by_legacy
local_rows: Final = tuple(
{key: value for key, value in row.items() if key not in _ACCOUNTING_ONLY_FIELDS}
for row in rows
if _delivery_id(row) in accepted and row.get("accounting_log_disabled") is not True
)
if write_spend_logs and local_rows:
groups: Final = tuple(
tuple(group)
for _, group in groupby(
sorted(local_rows, key=lambda row: tuple(sorted(row))), key=lambda row: tuple(sorted(row))
)
)
async with db_span("insert_spend_logs", "LiteLLM_SpendLogs"):
for group in groups:
await transaction.query_raw(*spend_log_insert_statement(group))
return len(accepted)
@staticmethod
async def _jobs_charged_by_legacy(transaction: LedgerTransaction, cost_row_ids: Sequence[str]) -> frozenset[str]:
"""Batches and videos a pre-ledger poll already charged through their SpendLogs cost row."""
if not cost_row_ids:
return frozenset()
rows: Final = await transaction.query_raw(
'SELECT DISTINCT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY($1::text[]) '
"AND status = 'success' AND spend > 0",
list(cost_row_ids),
)
return frozenset(str(row["request_id"]) for row in rows)
async def has_delivery(self, delivery_id: str) -> bool:
database: Final[LedgerTransaction] = self.prisma_client.writer_db
rows: Final = await database.query_raw(
'SELECT 1 FROM "LiteLLM_SpendLogDelivery" WHERE delivery_id = $1', delivery_id
)
return bool(rows)
class BudgetWindowSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_BudgetWindowSpend"]):
table_name = "litellm_budgetwindowspend"

17
schema.prisma
View file

@ -2141,6 +2141,23 @@ model LiteLLM_AccountingDelivery {
@@index([target, state, next_attempt_at])
}
model LiteLLM_SpendLogDelivery {
delivery_id String @id @default(dbgenerated("gen_random_uuid()::text"))
request_id String
state AccountingDeliveryState @default(pending)
accounting_payload Json
is_batch Boolean
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
created_at DateTime @default(now())
@@index([state, next_attempt_at])
@@index([state, applied_at])
@@index([request_id])
}
model LiteLLM_SubscriptionAccountTable {
subscription_account_id String @id @default(uuid())
custom_llm_provider String

512
tests/integration/database/test_spendlogs_delivery.py Normal file
View file

@ -0,0 +1,512 @@
import asyncio
import os
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import psycopg
import pytest
import pytest_asyncio
from prisma import Prisma
from litellm.repositories.accounting_delivery_repository import (
AccountingDeliveryRepository,
DeliveryProjection,
LedgerTransaction,
)
from litellm.repositories.spend_log_insert import spend_log_insert_statement
from tests.integration._support.database import scratch_database
_MIGRATIONS: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/migrations"
@pytest.fixture
def delivery_database(request: pytest.FixtureRequest) -> Iterator[psycopg.Connection[tuple[object, ...]]]:
with scratch_database("spendlogs_delivery") as url, psycopg.connect(url, autocommit=True) as connection:
if getattr(request, "param", "public") == "ledger_tenant":
connection.execute("CREATE SCHEMA ledger_tenant")
connection.execute("SET search_path TO ledger_tenant")
connection.execute((_MIGRATIONS / "20261027000000_add_accounting_ledger_tables/migration.sql").read_text())
connection.execute("""CREATE TABLE "LiteLLM_SpendLogs" (
request_id text PRIMARY KEY, call_type text NOT NULL, spend float8 NOT NULL DEFAULT 0,
status text,
metadata jsonb DEFAULT '{}', "startTime" timestamp(3) NOT NULL,
created_at timestamp(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at timestamp(3) NOT NULL DEFAULT CURRENT_TIMESTAMP
)""")
connection.execute((_MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql").read_text())
yield connection
def _insert(
connection: psycopg.Connection[tuple[object, ...]], rows: Sequence[Mapping[str, object]], *, ledger: bool = True
) -> None:
from psycopg.types.json import Jsonb
statement: Final = spend_log_insert_statement(rows)
with connection.transaction():
inserted: Final = {
row[0] for row in connection.execute(statement[0].replace("$1", "%s"), (statement[1],)).fetchall()
}
for row in rows:
if ledger and row["request_id"] in inserted:
connection.execute(
'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch) '
"VALUES (%s, %s, %s, %s)",
(row["request_id"], row["request_id"], Jsonb(dict(row)), row.get("call_type") == "aretrieve_batch"),
)
def _with_payload(row: Mapping[str, object]) -> Mapping[str, object]:
return {**row, "accounting_payload": dict(row)}
def _row(request_id: str) -> Mapping[str, object]:
return {
"request_id": request_id,
"call_type": "acompletion",
"spend": 0.5,
"startTime": "2026-11-02T12:00:00+00:00",
"metadata": '{"accounting_savings":{"compression":0.2}}',
}
@pytest.mark.parametrize("call_type", ["acompletion", "aretrieve_batch"])
@pytest.mark.parametrize("mode", ["local", "url", "disabled"])
@pytest.mark.parametrize("partitioned", [False, True])
@pytest.mark.asyncio
async def test_delivery_first_acceptance_deduplicates_all_storage_modes(
delivery_database, delivery_client, mode, partitioned, call_type
):
from types import SimpleNamespace
from litellm.repositories.table_repositories import SpendLogsRepository
if partitioned:
delivery_database.execute('ALTER TABLE "LiteLLM_SpendLogs" RENAME TO spend_template')
delivery_database.execute(
'CREATE TABLE "LiteLLM_SpendLogs" (LIKE spend_template INCLUDING DEFAULTS INCLUDING CONSTRAINTS, PRIMARY KEY(request_id, "startTime")) PARTITION BY RANGE ("startTime")'
)
delivery_database.execute('CREATE TABLE "LiteLLM_SpendLogs_p20261102" PARTITION OF "LiteLLM_SpendLogs" DEFAULT')
delivery_database.execute("DROP TABLE spend_template")
repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client))
row = {**_with_payload({**_row("shared"), "call_type": call_type}), "accounting_log_disabled": mode == "disabled"}
results = await asyncio.gather(
repository.insert_with_deliveries((row,), write_spend_logs=mode != "url"),
repository.insert_with_deliveries(
({**row, "startTime": "2026-11-03T12:00:00"},), write_spend_logs=mode != "url"
),
)
assert sorted(results) == [0, 1]
assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (int(mode == "local"),)
assert delivery_database.execute(
"SELECT accounting_payload->>'spend', state::text FROM \"LiteLLM_SpendLogDelivery\""
).fetchall() == [("0.5", "pending")]
if partitioned:
delivery_database.execute('DROP TABLE "LiteLLM_SpendLogs_p20261102"')
else:
delivery_database.execute('DELETE FROM "LiteLLM_SpendLogs"')
async def apply(transaction, deliveries):
await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(deliveries[0].row["spend"]))
worker = AccountingDeliveryRepository(delivery_client)
assert (await worker.apply_pending(apply)).applied == 1
assert (await worker.apply_pending(apply)).applied == 0
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,)
@pytest.mark.parametrize("mode", ["local", "url", "disabled"])
@pytest.mark.asyncio
async def test_duplicate_rows_in_one_flush_keep_first_accounting_payload(delivery_database, delivery_client, mode):
from types import SimpleNamespace
from litellm.repositories.table_repositories import SpendLogsRepository
repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client))
first = {**_with_payload(_row("repeated")), "accounting_log_disabled": mode == "disabled"}
later = {
**_with_payload({**_row("repeated"), "spend": 9.0, "startTime": "2026-11-03T12:00:00+00:00"}),
"accounting_log_disabled": mode == "disabled",
}
assert await repository.insert_with_deliveries((first, later), write_spend_logs=mode != "url") == 1
assert delivery_database.execute(
'SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"'
).fetchone() == (first["accounting_payload"],)
assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == (
[(first["spend"],)] if mode == "local" else []
)
@pytest.mark.parametrize(("legacy_spend", "accepted"), [(0.5, 0), (0.0, 1)])
@pytest.mark.asyncio
async def test_batch_charged_before_the_ledger_is_not_charged_again(
delivery_database, delivery_client, legacy_spend, accepted
):
from types import SimpleNamespace
from litellm.repositories.table_repositories import SpendLogsRepository
batch: Final = {**_row("batch_legacy"), "call_type": "aretrieve_batch", "status": "success"}
_insert(delivery_database, ({**batch, "spend": legacy_spend},), ledger=False)
repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client))
assert await repository.insert_with_deliveries((_with_payload({**batch, "spend": 1.0}),)) == accepted
assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (accepted,)
assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == [(legacy_spend,)]
@pytest.mark.parametrize("table", ["LiteLLM_SpendLogs", "LiteLLM_SpendLogDelivery"])
@pytest.mark.asyncio
async def test_transaction_a_rolls_back_payload_and_log_on_failure(delivery_database, delivery_client, table):
from types import SimpleNamespace
from prisma.errors import PrismaError
from psycopg import sql
from litellm.repositories.table_repositories import SpendLogsRepository
delivery_database.execute(
sql.SQL("ALTER TABLE {} ADD CONSTRAINT reject_poison CHECK (request_id <> 'poison')").format(
sql.Identifier(table)
)
)
repository = SpendLogsRepository(SimpleNamespace(db=delivery_client, writer_db=delivery_client))
with pytest.raises(PrismaError):
await repository.insert_with_deliveries(tuple(_with_payload(_row(name)) for name in ("good", "poison")))
assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,)
assert delivery_database.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogs"').fetchone() == (0,)
@pytest.mark.asyncio
async def test_worker_rolls_back_destinations_and_parks_only_isolated_failure(
delivery_database: psycopg.Connection[tuple[object, ...]],
) -> None:
url: Final = urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{delivery_database.info.dbname}"))
client: Final = Prisma(datasource={"url": url})
delivery_database.execute("CREATE TABLE totals (spend float8 NOT NULL)")
delivery_database.execute("INSERT INTO totals VALUES (0)")
_insert(delivery_database, (_row("bad"), _row("good")))
async def apply(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None:
await transaction.execute_raw(
"UPDATE totals SET spend = spend + $1::float8", sum(float(str(row.row["spend"])) for row in rows)
)
if any(row.row["request_id"] == "bad" for row in rows):
raise ValueError("Invalid projected destination")
await client.connect()
try:
repository: Final = AccountingDeliveryRepository(client)
failed: Final = await repository.apply_pending(apply, max_attempts=2)
assert failed.failed == 2
assert failed.parked == ()
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,)
assert delivery_database.execute(
'SELECT DISTINCT state::text, attempts FROM "LiteLLM_SpendLogDelivery"'
).fetchall() == [("pending", 1)]
delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at = CURRENT_TIMESTAMP')
isolated: Final = await repository.apply_pending(apply, max_attempts=2)
assert isolated.parked == ("bad",)
applied: Final = await repository.apply_pending(apply, max_attempts=2)
assert applied.applied == 1
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,)
assert delivery_database.execute(
'SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id'
).fetchall() == [("bad", "parked"), ("good", "applied")]
assert (await repository.apply_pending(apply)).applied == 0
finally:
await client.disconnect()
@pytest.mark.asyncio
async def test_a_locked_database_never_parks_a_valid_delivery(delivery_database, delivery_client):
_insert(delivery_database, (_row("valid"),))
async def apply(transaction, deliveries):
await transaction.execute_raw("SET LOCAL lock_timeout = '50ms'")
await transaction.execute_raw(
"UPDATE totals SET spend = spend + $1::float8", sum(float(row.row["spend"]) for row in deliveries)
)
repository = AccountingDeliveryRepository(delivery_client)
with psycopg.connect(
urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{delivery_database.info.dbname}"))
) as locker:
locker.execute("LOCK TABLE totals IN ACCESS EXCLUSIVE MODE")
for _ in range(4):
result = await repository.apply_pending(apply, max_attempts=2)
assert result.failed == 1 and result.parked == ()
delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at = CURRENT_TIMESTAMP')
assert delivery_database.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [
("pending", 4)
]
assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,)
@pytest.mark.asyncio
async def test_an_unreachable_database_reported_as_a_data_error_never_parks_a_delivery(
delivery_database, delivery_client
):
from prisma.errors import DataError
_insert(delivery_database, (_row("valid"),))
async def unreachable(transaction, deliveries):
raise DataError(data={"user_facing_error": {"message": "Can't reach database server at `db`:`5432`"}})
repository = AccountingDeliveryRepository(delivery_client)
for _ in range(3):
assert (await repository.apply_pending(unreachable, max_attempts=2)).parked == ()
delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at = CURRENT_TIMESTAMP')
assert delivery_database.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [
("pending", 3)
]
@pytest_asyncio.fixture(loop_scope="function")
async def delivery_client(delivery_database: psycopg.Connection[tuple[object, ...]]) -> AsyncIterator[Prisma]:
url: Final = urlunsplit(urlsplit(os.environ["DATABASE_URL"])._replace(path=f"/{delivery_database.info.dbname}"))
client: Final = Prisma(datasource={"url": url})
delivery_database.execute("CREATE TABLE totals (spend float8 NOT NULL)")
delivery_database.execute("INSERT INTO totals VALUES (0)")
await client.connect()
try:
yield client
finally:
await client.disconnect()
@pytest.mark.asyncio
async def test_two_workers_claim_disjoint_rows_and_apply_once(
delivery_database: psycopg.Connection[tuple[object, ...]],
delivery_client: Prisma,
) -> None:
_insert(delivery_database, (_row("first"), _row("second")))
entered: Final[asyncio.Queue[str]] = asyncio.Queue()
release: Final = asyncio.Event()
async def apply(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None:
assert len(rows) == 1
await entered.put(str(rows[0].row["request_id"]))
await release.wait()
await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"])))
repository: Final = AccountingDeliveryRepository(delivery_client)
first: Final = asyncio.create_task(repository.apply_pending(apply, batch_size=1))
first_id: Final = await asyncio.wait_for(entered.get(), timeout=10)
second: Final = asyncio.create_task(repository.apply_pending(apply, batch_size=1))
try:
second_id: Final = await asyncio.wait_for(entered.get(), timeout=10)
assert first_id != second_id
finally:
release.set()
results: Final = await asyncio.gather(first, second)
assert sum(result.applied for result in results) == 2
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (1.0,)
assert delivery_database.execute(
"SELECT COUNT(*) FROM \"LiteLLM_SpendLogDelivery\" WHERE state = 'applied'"
).fetchone() == (2,)
@pytest.mark.asyncio
async def test_cancelled_worker_rolls_back_and_restart_applies_once(
delivery_database: psycopg.Connection[tuple[object, ...]],
delivery_client: Prisma,
) -> None:
_insert(delivery_database, (_row("restart"),))
wrote: Final = asyncio.Event()
hold: Final = asyncio.Event()
async def interrupted(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None:
await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"])))
wrote.set()
await hold.wait()
repository: Final = AccountingDeliveryRepository(delivery_client)
task: Final = asyncio.create_task(repository.apply_pending(interrupted))
await asyncio.wait_for(wrote.wait(), timeout=10)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,)
assert delivery_database.execute('SELECT state::text, attempts FROM "LiteLLM_SpendLogDelivery"').fetchone() == (
"pending",
0,
)
async def complete(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None:
await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"])))
assert (await repository.apply_pending(complete)).applied == 1
assert (await repository.apply_pending(complete)).applied == 0
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,)
@pytest.mark.asyncio
async def test_worker_applies_accounting_after_spend_logs_are_deleted(
delivery_database: psycopg.Connection[tuple[object, ...]], delivery_client: Prisma
) -> None:
_insert(delivery_database, (_row("removed"),))
delivery_database.execute('DELETE FROM "LiteLLM_SpendLogs"')
async def apply(transaction: LedgerTransaction, rows: tuple[DeliveryProjection, ...]) -> None:
await transaction.execute_raw("UPDATE totals SET spend = spend + $1::float8", float(str(rows[0].row["spend"])))
repository: Final = AccountingDeliveryRepository(delivery_client)
assert (await repository.apply_pending(apply)).applied == 1
assert (await repository.apply_pending(apply)).applied == 0
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,)
@pytest.mark.asyncio
async def test_receipt_cleanup_keeps_pending_and_parked_rows_and_reports_lag(
delivery_database: psycopg.Connection[tuple[object, ...]], delivery_client: Prisma
) -> None:
_insert(delivery_database, tuple(_row(state) for state in ("pending", "parked", "applied")))
delivery_database.execute(
'UPDATE "LiteLLM_SpendLogDelivery" '
'SET state = request_id::"AccountingDeliveryState", '
"created_at = CURRENT_TIMESTAMP - INTERVAL '40 days', "
"applied_at = CURRENT_TIMESTAMP - INTERVAL '40 days'"
)
repository = AccountingDeliveryRepository(delivery_client)
assert await repository.cleanup_applied() == 1
metrics = await repository.metrics()
assert metrics.pending_count == 1
assert metrics.parked_count == 1
assert metrics.oldest_pending_seconds >= 40 * 86400
assert delivery_database.execute(
'SELECT request_id FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id'
).fetchall() == [("parked",), ("pending",)]
@pytest.mark.parametrize("delivery_database", ["public", "ledger_tenant"], indirect=True)
@pytest.mark.parametrize("state", [None, "pending", "parked", "applied"])
def test_delivery_migration_replay_preserves_ledger_receipts(delivery_database, state):
if state is not None:
_insert(delivery_database, (_row("replay"),))
delivery_database.execute(
'UPDATE "LiteLLM_SpendLogDelivery" SET state=%s, attempts=3, last_error=%s',
(state, "retained failure"),
)
before = delivery_database.execute('SELECT * FROM "LiteLLM_SpendLogDelivery"').fetchall()
migration = (_MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql").read_text()
delivery_database.execute(migration)
delivery_database.execute(migration)
assert delivery_database.execute('SELECT * FROM "LiteLLM_SpendLogDelivery"').fetchall() == before
assert delivery_database.execute('SELECT count(*) FROM "LiteLLM_SpendLogs"').fetchone() == (int(state is not None),)
@pytest.mark.parametrize("state", ["pending", "parked"])
def test_delivery_migration_preserves_unused_old_receipts(state):
with scratch_database("spendlogs_old_delivery") as url, psycopg.connect(url, autocommit=True) as connection:
connection.execute((_MIGRATIONS / "20261027000000_add_accounting_ledger_tables/migration.sql").read_text())
connection.execute('CREATE TABLE "LiteLLM_SpendLogs" (request_id text PRIMARY KEY)')
connection.execute(
'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state) VALUES (%s, %s, 1, %s)',
("old-event", "LiteLLM_UserTable", state),
)
connection.execute((_MIGRATIONS / "20261103010000_spendlogs_accounting_delivery/migration.sql").read_text())
assert connection.execute('SELECT event_id, state::text FROM "LiteLLM_AccountingDelivery"').fetchall() == [
("old-event", state)
]
assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,)
@pytest.mark.asyncio
async def test_malformed_payload_is_parked_without_losing_other_deliveries(delivery_database, delivery_client):
_insert(delivery_database, (_row("bad"), _row("good")))
delivery_database.execute(
'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload=$$ [] $$::jsonb WHERE request_id=$$bad$$'
)
async def apply(transaction, deliveries):
await transaction.execute_raw(
"UPDATE totals SET spend = spend + $1::float8", sum(float(row.row["spend"]) for row in deliveries)
)
repository = AccountingDeliveryRepository(delivery_client)
assert (await repository.apply_pending(apply, max_attempts=2)).failed == 2
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,)
delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP')
assert (await repository.apply_pending(apply, max_attempts=2)).parked == ("bad",)
assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1
assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,)
@pytest.mark.asyncio
async def test_retention_keeps_batch_receipts_indefinitely(delivery_database, delivery_client):
_insert(delivery_database, (_row("normal"), {**_row("batch"), "call_type": "aretrieve_batch"}))
delivery_database.execute(
"UPDATE \"LiteLLM_SpendLogDelivery\" SET state='applied', applied_at=CURRENT_TIMESTAMP - INTERVAL '400 days'"
)
repository = AccountingDeliveryRepository(delivery_client)
assert await repository.cleanup_applied() == 1
assert await repository.cleanup_applied() == 0
assert delivery_database.execute('SELECT request_id FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("batch",)]
metrics = await repository.metrics()
assert metrics.batch_count == 1
assert metrics.table_bytes > 0
@pytest.mark.parametrize("disabled", [False, True])
@pytest.mark.asyncio
async def test_ledger_batch_retrieves_charge_once_in_every_logging_mode(
delivery_database, delivery_client, monkeypatch, disabled
):
from types import SimpleNamespace
from litellm.proxy import proxy_server
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.utils import jsonify_object
completed = {
**_row("batch_completed"), "call_type": "aretrieve_batch", "status": "success", "spend": 0.5,
"metadata": '{"status":"success","private":"request detail"}',
"accounting_payload": {
"request_id": "batch_completed", "call_type": "aretrieve_batch", "status": "success", "spend": 0.5,
"metadata": {"accounting_savings": {"compression": 0.2}},
},
}
monkeypatch.setattr(proxy_server, "general_settings", {"spend_logs_metadata_fields": {"include": ["status"]}})
client = SimpleNamespace(db=delivery_client, writer_db=delivery_client, jsonify_object=jsonify_object)
writer = DBSpendUpdateWriter()
first = await writer._claim_settle_once_cost_row(payload=completed, prisma_client=client, disable_spend_logs=disabled)
again = await writer._claim_settle_once_cost_row(
payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, prisma_client=client, disable_spend_logs=disabled
)
assert (first, again) == (True, False)
assert delivery_database.execute('SELECT spend, metadata FROM "LiteLLM_SpendLogs"').fetchall() == (
[] if disabled else [(0.5, {"status": "success"})]
)
assert delivery_database.execute(
'SELECT delivery_id, accounting_payload FROM "LiteLLM_SpendLogDelivery"'
).fetchall() == [("batch_completed", completed["accounting_payload"])]
@pytest.mark.parametrize("setting", ["retention_days", "batch_size"])
@pytest.mark.asyncio
async def test_invalid_cleanup_settings_preserve_receipts(delivery_database, delivery_client, setting):
_insert(delivery_database, (_row("protected"),))
with pytest.raises(ValueError, match="positive"):
await AccountingDeliveryRepository(delivery_client).cleanup_applied(**{setting: 0})
assert delivery_database.execute('SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery"').fetchall() == [
("protected", "pending")
]
@pytest.mark.parametrize("setting", ["max_attempts", "batch_size"])
@pytest.mark.asyncio
async def test_invalid_worker_settings_preserve_pending_work(delivery_database, delivery_client, setting):
from unittest.mock import AsyncMock
_insert(delivery_database, (_row("protected"),))
apply = AsyncMock()
with pytest.raises(ValueError, match="positive"):
await AccountingDeliveryRepository(delivery_client).apply_pending(apply, **{setting: 0})
apply.assert_not_awaited()
assert delivery_database.execute(
'SELECT request_id, state::text, attempts FROM "LiteLLM_SpendLogDelivery"'
).fetchall() == [("protected", "pending", 0)]

1437
tests/integration/database/test_spendlogs_ledger_parity.py Normal file
View file

File diff suppressed because it is too large Load diff

85
tests/unit/accounting/test_delivery_projection.py Normal file
View file

@ -0,0 +1,85 @@
import json
from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
from typing import Final
from urllib.parse import unquote
from litellm.accounting.delivery_projection import entity_transactions, project_deliveries
from litellm.accounting.projection import Contribution
from litellm.repositories.accounting_delivery_repository import DeliveryProjection
_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc)
def _payload(spend: float, started: datetime = _NOW, window_start: datetime = _NOW) -> Mapping[str, object]:
return {
"request_id": f"request-{spend}-{started.isoformat()}",
"startTime": started,
"api_key": "key",
"user": "user/1",
"team_id": "team",
"organization_id": "org::1",
"spend": spend,
"metadata": json.dumps(
{
"accounting_entities": {
"api_key": "key",
"user": "user/1",
"team_id": "team",
"organization_id": "org::1",
},
"accounting_windows": [
{
"entity_type": "key",
"entity_id": "key",
"window_duration": "1d",
"window_start": window_start.isoformat(),
}
],
}
),
}
def _spend(contributions: tuple[Contribution, ...], table: str) -> tuple[float, ...]:
return tuple(float(dict(item.increments)["spend"]) for item in contributions if item.table == table)
def test_deliveries_for_one_destination_sum_and_keep_the_earliest_request_start() -> None:
later: Final = _NOW + timedelta(hours=1)
contributions: Final = project_deliveries(
(DeliveryProjection(_payload(0.5, started=later)), DeliveryProjection(_payload(0.25)))
)
assert _spend(contributions, "LiteLLM_VerificationToken") == (0.75,)
windows: Final = tuple(item for item in contributions if item.table == "LiteLLM_BudgetWindowSpend")
assert len(windows) == 1
assert dict(windows[0].increments) == {"spend": 0.75}
assert dict(windows[0].values)["started_at"] == str(_NOW)
def test_the_same_window_in_a_new_period_is_charged_to_its_own_row() -> None:
next_day: Final = _NOW + timedelta(days=1)
contributions: Final = project_deliveries(
(DeliveryProjection(_payload(0.25)), DeliveryProjection(_payload(0.5, window_start=next_day)))
)
windows: Final = sorted(
(dict(item.values)["window_start"], dict(item.increments)["spend"])
for item in contributions
if item.table == "LiteLLM_BudgetWindowSpend"
)
assert windows == [(_NOW.isoformat(), 0.25), (next_day.isoformat(), 0.5)]
def test_entity_transactions_key_memberships_the_way_the_writer_applies_them() -> None:
transactions: Final = entity_transactions(project_deliveries((DeliveryProjection(_payload(0.25)),)))
assert transactions["key_list_transactions"] == {"key": 0.25}
assert transactions["org_list_transactions"] == {"org::1": 0.25}
assert transactions["team_member_list_transactions"] == {"team_id::team::user_id::user/1": 0.25}
((org_member_key, org_member_spend),) = transactions["org_member_list_transactions"].items()
_, quoted_org_id, _, quoted_user_id = org_member_key.split("::")
assert (unquote(quoted_org_id), unquote(quoted_user_id), org_member_spend) == ("org::1", "user/1", 0.25)

235
tests/unit/accounting/test_measurements.py
View file

@ -1,235 +0,0 @@
from datetime import datetime, timezone
from decimal import Decimal
from typing import Final
import pytest
from litellm.accounting import (
CostEvent,
Dimension,
DimensionName,
InboundUsageEvent,
Measurement,
Metric,
MetricName,
Outcome,
OutboundUsageEvent,
PricingSnapshot,
measurements_for,
)
from litellm.types.utils import CallTypes
_START: Final = datetime(2026, 1, 1, 23, 59, 59, tzinfo=timezone.utc)
_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc)
_DIMENSIONS: Final = (
Dimension(name=DimensionName.TEAM, value="team"),
Dimension(name=DimensionName.TAG, value="alpha"),
Dimension(name=DimensionName.TAG, value="beta"),
)
def _snapshot(amount: Decimal) -> PricingSnapshot:
return PricingSnapshot(
calculator_version="test",
model="test-model",
provider=None,
pricing_source="upstream_header",
priced_model_key="test-model",
upstream_reported_cost=amount,
total_cost=amount,
)
@pytest.mark.parametrize(
("outcome", "outcome_metric"),
[
(Outcome.SUCCEEDED, MetricName.SUCCESSFUL_REQUESTS),
(Outcome.FAILED, MetricName.FAILED_REQUESTS),
(Outcome.CANCELLED, MetricName.FAILED_REQUESTS),
],
)
def test_caller_measurements_keep_both_request_and_outcome_counts(outcome, outcome_metric):
event = InboundUsageEvent(
event_id="caller-event",
request_id="caller",
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=outcome,
dimensions=_DIMENSIONS,
metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)), Metric(name=outcome_metric, value=Decimal(1))),
)
assert measurements_for(event) == (
Measurement(
source_event_id="caller-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.REQUESTS, value=Decimal(1)),
),
Measurement(
source_event_id="caller-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=outcome_metric, value=Decimal(1)),
),
)
def test_usage_measurements_keep_every_metric_in_the_start_time_bucket():
event = OutboundUsageEvent(
event_id="usage-event",
attempt_id="attempt",
request_id="caller",
retry_of_attempt_id=None,
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.SUCCEEDED,
dimensions=_DIMENSIONS,
metrics=(
Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)),
Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)),
Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),
Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)),
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)),
),
)
assert measurements_for(event) == (
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)),
),
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)),
),
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),
),
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)),
),
Measurement(
source_event_id="usage-event",
occurred_at=_START,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)),
),
)
@pytest.mark.parametrize("metrics", [None, ()])
def test_unknown_or_empty_usage_does_not_fabricate_zero_measurements(metrics):
event = OutboundUsageEvent(
event_id="usage-event",
attempt_id="attempt",
request_id=None,
retry_of_attempt_id=None,
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.SUCCEEDED,
dimensions=(),
metrics=metrics,
)
assert measurements_for(event) == ()
@pytest.mark.parametrize("amount", [Decimal(0), Decimal("0.1234567890123456789")])
def test_cost_measurements_preserve_authoritative_amount_identity_and_occurrence(amount):
event = CostEvent(
event_id="cost-event",
usage_event_id="usage-event",
occurred_at=_END,
dimensions=_DIMENSIONS,
charges=(Metric(name=MetricName.SPEND, value=amount, unit="USD"),),
snapshot=_snapshot(amount),
)
assert measurements_for(event) == (
Measurement(
source_event_id="cost-event",
occurred_at=_END,
dimensions=_DIMENSIONS,
metric=Metric(name=MetricName.SPEND, value=amount, unit="USD"),
),
)
def test_billed_retry_and_success_contribute_one_caller_request_and_both_charges():
caller = InboundUsageEvent(
event_id="caller-event",
request_id="caller",
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.SUCCEEDED,
dimensions=_DIMENSIONS,
metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),),
)
retry = OutboundUsageEvent(
event_id="retry-usage",
attempt_id="retry",
request_id="caller",
retry_of_attempt_id=None,
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.FAILED,
dimensions=_DIMENSIONS,
metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(4)),),
)
success = OutboundUsageEvent(
event_id="success-usage",
attempt_id="success",
request_id="caller",
retry_of_attempt_id="retry",
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_END,
outcome=Outcome.SUCCEEDED,
dimensions=_DIMENSIONS,
metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),),
)
retry_cost = CostEvent(
event_id="retry-cost",
usage_event_id="retry-usage",
occurred_at=_START,
dimensions=_DIMENSIONS,
charges=(Metric(name=MetricName.SPEND, value=Decimal("0.025"), unit="USD"),),
snapshot=_snapshot(Decimal("0.025")),
)
success_cost = CostEvent(
event_id="success-cost",
usage_event_id="success-usage",
occurred_at=_START,
dimensions=_DIMENSIONS,
charges=(Metric(name=MetricName.SPEND, value=Decimal("0.125"), unit="USD"),),
snapshot=_snapshot(Decimal("0.125")),
)
measurements = (
*measurements_for(caller),
*measurements_for(retry),
*measurements_for(success),
*measurements_for(retry_cost),
*measurements_for(success_cost),
)
assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.REQUESTS) == Decimal(1)
assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.TOTAL_TOKENS) == Decimal(16)
assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.SPEND) == Decimal("0.150")
assert {m.source_event_id for m in measurements} == {
"caller-event",
"retry-usage",
"success-usage",
"retry-cost",
"success-cost",
}

28
tests/unit/accounting/test_metrics.py Normal file
View file

@ -0,0 +1,28 @@
from prometheus_client.metrics import MetricWrapperBase
from litellm.accounting import metrics
from litellm.accounting.metrics import record_deliveries
from litellm.repositories.accounting_delivery_repository import DeliveryMetrics
def _sample(metric: MetricWrapperBase, suffix: str = "") -> float:
family = metric.collect()[0]
return next(sample.value for sample in family.samples if sample.name == family.name + suffix)
def test_delivery_metrics_publish_backlog_and_clamp_clock_skew() -> None:
batches_before = _sample(metrics._BATCH_TIME, "_count")
record_deliveries(
DeliveryMetrics(pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192),
0.5,
)
assert (
_sample(metrics._PENDING),
_sample(metrics._PARKED),
_sample(metrics._OLDEST),
_sample(metrics._BATCH_COUNT),
_sample(metrics._TABLE_BYTES),
) == (4.0, 1.0, 0.0, 7.0, 8192.0)
assert _sample(metrics._BATCH_TIME, "_count") == batches_before + 1

126
tests/unit/accounting/test_projection.py
View file

@ -429,3 +429,129 @@ def test_guardrail_usage_is_dated_in_utc() -> None:
assert [dict(item.key)["date"] for item in projected if item.table == "LiteLLM_DailyGuardrailUsageUnits"] == [
"2026-11-03"
]
@pytest.mark.parametrize(
("untrusted", "metadata"),
[
({"request_tags": json.dumps(["x" * 16384, "normal"])}, {}),
({"end_user": "x" * 3072}, {}),
({}, {"compression_savings": {"tokens_saved": 2**63}}),
],
)
def test_untrusted_analytics_values_never_block_the_charge(
untrusted: Mapping[str, str], metadata: Mapping[str, object]
) -> None:
row: Final[Mapping[str, object]] = {
"request_id": "untrusted",
"startTime": _NOW,
"api_key": "key",
"user": "user",
"team_id": "team",
"model": "unpriced-projection-test",
"model_group": "group",
"custom_llm_provider": "openai",
"spend": 0.5,
**untrusted,
"metadata": {
**metadata,
"accounting_entities": {
"api_key": "key",
"user": "user",
"team_id": "team",
"end_user": untrusted.get("end_user"),
},
},
}
contributions: Final = project(row)
texts: Final = tuple(value for item in contributions for _, value in (*item.key, *item.values) if value)
numbers: Final = tuple(value for item in contributions for _, value in item.increments)
assert {"LiteLLM_VerificationToken", "LiteLLM_TeamTable", "LiteLLM_UserTable"} <= {
item.table for item in contributions
}
assert "LiteLLM_DailyUserSpend" in {item.table for item in contributions}
assert all(len(text.encode()) <= 2000 for text in texts)
assert all(abs(number) <= 2**40 for number in numbers)
def test_stored_payload_keeps_guardrail_billing_but_not_guardrail_content() -> None:
payload: Final = get_logging_payload(
kwargs={
"model": "projection-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "guardrail-request",
"litellm_params": {"metadata": {"user_api_key": "key", "user_api_key_user_id": "user"}},
},
response_obj=ModelResponse(
id="guardrail-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
start_time=_NOW,
end_time=_NOW,
)
guardrail: Final = {
"guardrail_name": "pii",
"guardrail_usage": {"requests": 1},
"guardrail_cost_by_unit": {"requests": 0.25},
"guardrail_response": "card 4111 1111 1111 1111",
"match_details": [{"matched": "4111 1111 1111 1111"}],
}
headroom: Final = {
"guardrail_name": "compress",
"guardrail_provider": "headroom",
"guardrail_response": {"tokens_saved": 7, "compressed_prompt": "card 4111 1111 1111 1111"},
}
metadata: Final = {**_FIELDS.validate_json(payload["metadata"]), "guardrail_information": [guardrail, headroom]}
accounting_input: Final[SpendLogsPayload] = {
**payload,
"spend": 0.5,
"metadata": json.dumps(metadata),
"request_tags": '["nul\x00tag"]',
}
accounting: Final = build_accounting_payload(
accounting_input, entities={"api_key": "key"}, model_access_groups=(), windows=(), router=lambda: None
)
stored: Final = json.dumps(accounting, default=str)
restored: Final = _FIELDS.validate_json(stored)
usage_rows: Final = tuple(item for item in project(restored) if item.table == "LiteLLM_DailyGuardrailUsageUnits")
assert "4111" not in stored
assert "\\u0000" not in stored
assert [dict(item.increments) for item in usage_rows] == [{"units": 1, "cost": 0.25, "untracked_units": 0}]
daily_user: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend")
assert dict(daily_user.increments)["compression_saved_tokens"] == 7
def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(monkeypatch: pytest.MonkeyPatch) -> None:
import litellm
monkeypatch.setitem(
litellm.model_cost,
"savings-overflow-test",
{"input_cost_per_token": 0.002, "output_cost_per_token": 0.004, "litellm_provider": "openai", "mode": "chat"},
)
payload: Final[SpendLogsPayload] = {
**get_logging_payload(
kwargs={
"model": "savings-overflow-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "overflow-request",
"litellm_params": {"metadata": {"user_api_key": "key"}},
},
response_obj=ModelResponse(
id="overflow-request", model="savings-overflow-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
start_time=_NOW,
end_time=_NOW,
),
"spend": 0.5,
"metadata": json.dumps({"compression_savings": {"tokens_saved": 10**400}}),
}
accounting: Final = build_accounting_payload(
payload, entities={"api_key": "key"}, model_access_groups=(), windows=(), router=lambda: None
)
restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str))
assert dict(restored["metadata"])["accounting_savings"]["compression"] == 0.0
assert [dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"] == [
0.5
]

482
tests/unit/accounting/test_types.py
View file

@ -1,482 +0,0 @@
import re
from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from enum import Enum
from pathlib import Path
from typing import Final
import pytest
from pydantic import ValidationError
from litellm.accounting import (
AdmissionDecision,
AggregateQuery,
AggregateRow,
CostEvent,
DeliveryState,
DeliveryTarget,
Dimension,
DimensionFilter,
DimensionName,
Granularity,
InboundUsageEvent,
Metric,
MetricName,
Multiplier,
Outcome,
OutboundUsageEvent,
PricedQuantity,
PricingSnapshot,
Rate,
validate_aggregate_rows,
)
from litellm.accounting.types import PricingUnit
from litellm.types.utils import CallTypes
_START: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc)
_REQUEST: Final = Metric(name=MetricName.REQUESTS, value=Decimal(1))
_TOKEN: Final = Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10))
_SPEND: Final = Metric(name=MetricName.SPEND, value=Decimal("0.1234567890123456789"), unit="USD")
_TEAM: Final = Dimension(name=DimensionName.TEAM, value="team")
def _inbound(metrics=(_REQUEST,), outcome=Outcome.SUCCEEDED):
return InboundUsageEvent(
event_id="inbound:request",
request_id="request",
call_type=CallTypes.acompletion,
started_at=_START,
ended_at=_START,
outcome=outcome,
dimensions=(_TEAM,),
metrics=metrics,
)
def _outbound(metrics=(_TOKEN,), **changes):
return OutboundUsageEvent.model_validate(
{
"event_id": "outbound:attempt",
"attempt_id": "attempt",
"request_id": "request",
"retry_of_attempt_id": None,
"call_type": CallTypes.acompletion,
"started_at": _START,
"ended_at": _END,
"outcome": Outcome.SUCCEEDED,
"dimensions": (_TEAM,),
"metrics": metrics,
**changes,
}
)
def _query(**changes):
return AggregateQuery.model_validate(
{
"metric": MetricName.SPEND,
"filters": (DimensionFilter(name=DimensionName.TEAM, values=("one", "two")),),
"group_by": (DimensionName.MODEL,),
"start": _START,
"end": _END,
"granularity": Granularity.DAY,
**changes,
}
)
def test_wire_values_preserve_decimal_precision_and_normalize_aware_times():
offset = timezone(timedelta(hours=-8))
event = CostEvent(
event_id="cost:attempt",
usage_event_id="outbound:attempt",
occurred_at=datetime(2025, 12, 31, 16, tzinfo=offset),
dimensions=(_TEAM,),
charges=(_SPEND,),
snapshot=_snapshot(),
)
wire = event.model_dump(mode="json")
assert wire["occurred_at"] == "2026-01-01T00:00:00Z"
assert wire["charges"] == [{"name": "spend", "value": "0.1234567890123456789", "unit": "USD"}]
assert event.occurred_at == _START
assert CostEvent.model_validate_json(event.model_dump_json()).charges[0].value == _SPEND.value
@pytest.mark.parametrize(
("name", "value", "unit", "message"),
[
(MetricName.SPEND, Decimal(1), "count", "SPEND requires USD"),
(MetricName.REQUESTS, Decimal(1), "USD", "SPEND requires USD"),
(MetricName.PROMPT_TOKENS, Decimal("1.5"), "count", "integral"),
(MetricName.SPEND, Decimal("NaN"), "USD", "finite"),
(MetricName.SPEND, Decimal("Infinity"), "USD", "finite"),
(MetricName.SPEND, Decimal(-1), "USD", "greater than or equal"),
(MetricName.SPEND, 0.1, "USD", "instance of Decimal"),
],
)
def test_metrics_reject_wrong_units_nonfinite_values_and_binary_float_inputs(name, value, unit, message):
with pytest.raises(ValidationError, match=message):
Metric(name=name, value=value, unit=unit)
@pytest.mark.parametrize("metrics", [(_TOKEN,), (_SPEND,), ()])
def test_inbound_activity_cannot_own_usage_or_cost_or_omit_request_count(metrics):
with pytest.raises(ValidationError, match="Inbound events"):
_inbound(metrics)
@pytest.mark.parametrize(
("outcome", "metric"),
[
(Outcome.FAILED, Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(1))),
(Outcome.SUCCEEDED, Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(1))),
],
)
def test_inbound_outcome_counts_cannot_contradict_outcome(outcome, metric):
with pytest.raises(ValidationError, match="count must match outcome"):
_inbound((_REQUEST, metric), outcome)
@pytest.mark.parametrize("metrics", [(_REQUEST,), (_SPEND,)])
def test_outbound_attempts_cannot_own_caller_counts_or_charges(metrics):
with pytest.raises(ValidationError, match="Outbound events own only"):
_outbound(metrics)
@pytest.mark.parametrize(
("changes", "message"),
[
({"started_at": datetime(2026, 1, 1)}, "timezone"),
({"ended_at": _START - timedelta(seconds=1)}, "End must follow start"),
({"retry_of_attempt_id": "attempt"}, "cannot retry itself"),
({"dimensions": (_TEAM, _TEAM)}, "Duplicate dimension values"),
({"dimensions": (_TEAM, Dimension(name=DimensionName.TEAM, value="other"))}, "Only tag and model access group"),
(
{
"metrics": (
Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(5)),
Metric(name=MetricName.REASONING_TOKENS, value=Decimal(6)),
)
},
"reasoning_tokens cannot exceed completion_tokens",
),
(
{
"metrics": (
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(5)),
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(6)),
)
},
"cache_creation_input_tokens_above_1hr cannot exceed cache_creation_input_tokens",
),
({"event_id": " "}, "pattern"),
({"metrics": (_TOKEN, _TOKEN)}, "Duplicate metric names"),
],
)
def test_outbound_events_reject_invalid_identity_time_and_duplicate_attribution(changes, message):
with pytest.raises(ValidationError, match=message):
_outbound(**changes)
def test_unknown_and_known_zero_usage_are_distinct_and_values_are_immutable():
unknown = _outbound(None)
zero = _outbound((Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(0)),))
assert unknown.metrics is None
assert zero.metrics is not None and zero.metrics[0].value == 0
assert _inbound().started_at == _inbound().ended_at
with pytest.raises(ValidationError, match="frozen"):
zero.attempt_id = "different"
with pytest.raises(ValidationError, match="frozen"):
zero.metrics[0].value = Decimal(10)
@pytest.mark.parametrize("charges", [(), (_TOKEN,), (_SPEND, _SPEND)])
def test_cost_events_require_one_authoritative_spend(charges):
with pytest.raises(ValidationError, match=r"authoritative USD spend|Duplicate metric"):
CostEvent(
event_id="cost",
usage_event_id="usage",
occurred_at=_START,
dimensions=(),
charges=charges,
snapshot=_snapshot(),
)
def test_cost_event_cannot_reuse_its_usage_identity():
with pytest.raises(ValidationError, match="distinct identities"):
CostEvent(
event_id="usage",
usage_event_id="usage",
occurred_at=_START,
dimensions=(),
charges=(_SPEND,),
snapshot=_snapshot(),
)
@pytest.mark.parametrize("granularity", [Granularity.DAY, None])
def test_daily_and_whole_interval_queries_preserve_filter_and_group_semantics(granularity):
query = _query(granularity=granularity)
assert query.granularity is granularity
assert query.filters[0].values == ("one", "two")
assert query.group_by == (DimensionName.MODEL,)
assert query.start == _START and query.end == _END
@pytest.mark.parametrize(
("changes", "message"),
[
({"end": _START}, "End must follow start"),
({"start": _START + timedelta(hours=1)}, "UTC day boundaries"),
({"end": _END + timedelta(microseconds=1)}, "UTC day boundaries"),
({"granularity": "hour"}, "day"),
({"group_by": (DimensionName.TEAM, DimensionName.TEAM)}, "Duplicate group-by"),
({"filters": (DimensionFilter(name=DimensionName.TEAM, values=("one",)),) * 2}, "one IN filter"),
({"filters": ({"name": "team", "values": ()},)}, "at least 1"),
],
)
def test_queries_reject_unsupported_or_ambiguous_ranges_and_filters(changes, message):
with pytest.raises(ValidationError, match=message):
_query(**changes)
def test_rows_allow_adjacent_buckets_and_overlapping_distinct_groups():
first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM,), value=Decimal(1))
second = AggregateRow(start=_END, end=_END + timedelta(days=1), dimensions=(_TEAM,), value=Decimal(2))
other = AggregateRow(start=_START, end=_END, dimensions=(), value=Decimal(3))
rows = (second, other, first)
assert validate_aggregate_rows(rows) is rows
assert validate_aggregate_rows(()) == ()
def test_rows_reject_overlap_within_a_group_regardless_of_dimension_order():
model = Dimension(name=DimensionName.MODEL, value="model")
first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM, model), value=Decimal(1))
overlapping = AggregateRow(start=_START, end=_END, dimensions=(model, _TEAM), value=Decimal(2))
with pytest.raises(ValueError, match="overlap within a dimension group"):
validate_aggregate_rows((first, overlapping))
def test_rows_reject_empty_intervals():
with pytest.raises(ValidationError, match="End must follow start"):
AggregateRow(start=_START, end=_START, dimensions=(), value=Decimal(0))
def test_admission_decision_preserves_a_denial_reason():
decision = AdmissionDecision(allowed=False, reason="Team budget exhausted")
assert not decision.allowed
assert decision.reason == "Team budget exhausted"
def _snapshot(**changes: object) -> PricingSnapshot:
return PricingSnapshot.model_validate(
{
"calculator_version": "1.80.0",
"model": "gpt-4o",
"provider": "openai",
"pricing_source": "cost_map",
"priced_model_key": "openai/gpt-4o",
"multipliers": (Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")),),
"rates": (Rate(measure="input", unit="token", usd_per_unit=Decimal("0.0000025")),),
"usage": (PricedQuantity(measure="input", unit="token", value=_TOKEN.value),),
"input_cost": Decimal("0.000025"),
"total_cost": _SPEND.value,
**changes,
}
)
def test_cost_event_carries_a_snapshot_whose_total_matches_the_spend():
event = CostEvent(
event_id="cost:attempt",
usage_event_id="outbound:attempt",
occurred_at=_START,
dimensions=(_TEAM,),
charges=(_SPEND,),
snapshot=_snapshot(),
)
assert event.snapshot is not None
assert event.snapshot.total_cost == _SPEND.value
assert event.snapshot.multipliers[0].factor == Decimal("1.1")
with pytest.raises(ValidationError, match="Snapshot total must equal the authoritative spend"):
CostEvent(
event_id="cost:attempt",
usage_event_id="outbound:attempt",
occurred_at=_START,
dimensions=(_TEAM,),
charges=(_SPEND,),
snapshot=_snapshot(total_cost=Decimal("0.5")),
)
@pytest.mark.parametrize(
("changes", "message"),
(
(
{
"rates": (
Rate(measure="input", unit="token", usd_per_unit=Decimal(1)),
Rate(measure="input", unit="token", usd_per_unit=Decimal(2)),
)
},
"Duplicate rate measures and units",
),
({"usage": (_SPEND,)}, "PricedQuantity"),
({"usage": (_REQUEST,)}, "PricedQuantity"),
({"pricing_source": "guess"}, "pricing_source"),
({"multipliers": ({"name": "x", "value": "y", "factor": Decimal(0)},)}, "factor"),
({"schema_version": 2}, "schema_version"),
(
{"additional_costs": ({"name": "fee", "amount": Decimal(1)}, {"name": "fee", "amount": Decimal(2)})},
"Duplicate additional cost names",
),
),
)
def test_snapshots_reject_ambiguous_rates_wrong_usage_and_unknown_sources(changes: Mapping[str, object], message: str):
with pytest.raises(ValidationError, match=message):
_snapshot(**changes)
@pytest.mark.parametrize(
("unit", "value"),
[("token", "10"), ("second", "1.25"), ("character", "25"), ("image", "2"), ("request", "1")],
)
def test_snapshot_preserves_matching_rates_and_quantities_for_each_billing_unit(unit: PricingUnit, value: str) -> None:
quantity: Final = PricedQuantity(measure="input", unit=unit, value=Decimal(value))
rate: Final = Rate(measure="input", unit=unit, usd_per_unit=Decimal("0.003"))
amount: Final = quantity.value * rate.usd_per_unit
snapshot: Final = _snapshot(rates=(rate,), usage=(quantity,), total_cost=amount)
restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json())
assert restored == snapshot
assert restored.usage[0].value * restored.rates[0].usd_per_unit == amount
@pytest.mark.parametrize(
("usage", "message"),
[
((), "Every rate requires"),
(({"measure": "output", "unit": "token", "value": Decimal(10)},), "Every rate requires"),
(({"measure": "input", "unit": "second", "value": Decimal(10)},), "Every rate requires"),
(
({"measure": "input", "unit": "token", "value": Decimal(10)},) * 2,
"Duplicate priced quantity",
),
],
)
def test_snapshot_rejects_missing_mismatched_or_duplicate_priced_quantities(
usage: tuple[Mapping[str, object], ...], message: str
) -> None:
with pytest.raises(ValidationError, match=message):
_snapshot(usage=usage)
@pytest.mark.parametrize("unit", ["token", "character", "image", "request"])
def test_priced_counts_remain_integral(unit: PricingUnit) -> None:
with pytest.raises(ValidationError, match="integral"):
PricedQuantity(measure="input", unit=unit, value=Decimal("1.25"))
@pytest.mark.parametrize(
("value", "message"),
[
(Decimal(-1), "greater than or equal"),
(Decimal("NaN"), "finite"),
(Decimal("Infinity"), "finite"),
(1.25, "instance of Decimal"),
],
)
def test_pricing_values_reject_invalid_amounts(value: Decimal | float, message: str) -> None:
with pytest.raises(ValidationError, match=message):
PricedQuantity.model_validate({"measure": "input", "unit": "second", "value": value})
with pytest.raises(ValidationError, match=message):
Rate.model_validate({"measure": "input", "unit": "second", "usd_per_unit": value})
@pytest.mark.parametrize("snapshot_fields", [{}, {"snapshot": None}])
def test_cost_event_cannot_omit_its_derivation(snapshot_fields: Mapping[str, object]) -> None:
with pytest.raises(ValidationError, match="snapshot"):
CostEvent.model_validate(
{
"event_id": "cost",
"usage_event_id": "usage",
"occurred_at": _START,
"dimensions": (),
"charges": (_SPEND,),
**snapshot_fields,
}
)
def test_upstream_charge_snapshot_preserves_reported_cost_without_inventing_rates() -> None:
snapshot: Final = _snapshot(
pricing_source="upstream_header", rates=(), usage=(), upstream_reported_cost=_SPEND.value
)
restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json())
assert restored.upstream_reported_cost == restored.total_cost == _SPEND.value
assert restored.rates == restored.usage == ()
def test_subset_token_counts_round_trip_without_changing_totals() -> None:
metrics: Final = (
Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(10)),
Metric(name=MetricName.REASONING_TOKENS, value=Decimal(3)),
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(8)),
Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(2)),
)
event: Final = _outbound(metrics=metrics)
assert OutboundUsageEvent.model_validate_json(event.model_dump_json()).metrics == metrics
@pytest.mark.parametrize(
"schema_path",
[
"schema.prisma",
"litellm/proxy/schema.prisma",
"litellm-proxy-extras/litellm_proxy_extras/schema.prisma",
],
)
@pytest.mark.parametrize(
("enum_name", "python_enum"),
[
("AccountingDeliveryState", DeliveryState),
("AccountingDeliveryTarget", DeliveryTarget),
],
)
def test_prisma_accepts_exactly_the_public_accounting_enum_values(
schema_path: str, enum_name: str, python_enum: type[Enum]
) -> None:
schema: Final = (Path(__file__).resolve().parents[3] / schema_path).read_text()
declaration: Final = re.search(r"enum " + enum_name + r"\s*\{([^}]+)\}", schema)
assert declaration is not None, f"Missing persisted enum {enum_name}"
assert frozenset(declaration[1].split()) == frozenset(member.value for member in python_enum)
def test_one_event_carries_every_matched_model_access_group():
groups = (
Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="gpt-group"),
Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="claude-group"),
)
assert _outbound(dimensions=(_TEAM, *groups)).dimensions == (_TEAM, *groups)
with pytest.raises(ValidationError, match="Duplicate dimension values"):
_outbound(dimensions=(groups[0], groups[0]))
def test_snapshots_reject_duplicate_multipliers_and_contradictory_token_counts():
multiplier = Multiplier(name="data_residency", value="eu", factor=Decimal("1.1"))
with pytest.raises(ValidationError, match="Duplicate multiplier names"):
_snapshot(multipliers=(multiplier, multiplier))
with pytest.raises(ValidationError, match="reasoning_tokens cannot exceed completion_tokens"):
_snapshot(
usage=(
PricedQuantity(measure="input", unit="token", value=_TOKEN.value),
PricedQuantity(measure="completion_tokens", unit="token", value=Decimal(5)),
PricedQuantity(measure="reasoning_tokens", unit="token", value=Decimal(6)),
)
)
def test_agent_live_spend_is_a_delivery_target():
assert DeliveryTarget.AGENT_SPEND.value == "LiteLLM_AgentsTable"

50
tests/unit/proxy/common_utils/test_check_batch_cost.py
View file

@ -24,12 +24,12 @@ import httpx
import pytest
import respx
from fastapi import HTTPException
from litellm_enterprise.proxy.common_utils.check_batch_cost import GIVE_UP_POLLING_AFTER
import litellm
from litellm.constants import MANAGED_OBJECT_STALENESS_CUTOFF_DAYS
from litellm.types.router import LiteLLM_Params
from litellm.types.utils import LiteLLMBatch
from litellm_enterprise.proxy.common_utils.check_batch_cost import GIVE_UP_POLLING_AFTER
if TYPE_CHECKING:
from litellm.batches.batch_utils import BatchCostUsageResult
@ -254,7 +254,6 @@ class TestCheckBatchCost:
@pytest.mark.asyncio
async def test_column_absence_cached_across_cycles(self, check_batch_cost_instance, mock_prisma_client):
"""After column absence is discovered, subsequent cycles skip the primary query entirely."""
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
# Simulate column already known absent from a previous cycle
@ -2387,8 +2386,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
@pytest.mark.asyncio
async def test_target_model_names_comes_from_input_file_not_provider_model(self):
from litellm.proxy.openai_files_endpoints.common_utils import (
is_base64_encoded_unified_file_id,
get_models_from_unified_file_id,
is_base64_encoded_unified_file_id,
)
output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)))
@ -2422,8 +2421,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
@pytest.mark.asyncio
async def test_falls_back_to_deployment_model_group_without_managed_input_file(self):
from litellm.proxy.openai_files_endpoints.common_utils import (
is_base64_encoded_unified_file_id,
get_models_from_unified_file_id,
is_base64_encoded_unified_file_id,
)
output_file_id = await self._run(self._job("file-raw-provider-input"))
@ -3309,7 +3308,9 @@ class TestMultiPodBatchCostClaim:
patch(
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
new_callable=AsyncMock,
return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, ["gpt-4"]),
return_value=_batch_cost_result(
0.01, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, ["gpt-4"]
),
),
patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
@ -3329,15 +3330,15 @@ class TestMultiPodBatchCostClaim:
from datetime import datetime
import httpx
import litellm
import respx
import litellm
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.utils import PrismaClient
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.types.utils import SpecialEnums
@ -3383,6 +3384,14 @@ class TestMultiPodBatchCostClaim:
database_url="postgresql://unused:unused@localhost:5432/unused", proxy_logging_obj=dependency
)
create_spend_rows: Final = AsyncMock(return_value=1)
from litellm.proxy.db.prisma_client import PrismaWrapper
spend_prisma.db = PrismaWrapper(
original_prisma=MagicMock(
query_raw=AsyncMock(return_value=[]),
litellm_spendlogs=MagicMock(create_many=create_spend_rows),
)
)
input_rate: Final = 0.000002
litellm.register_model(
{
@ -3434,11 +3443,6 @@ class TestMultiPodBatchCostClaim:
respx.mock,
patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency),
patch("litellm.proxy.proxy_server.prisma_client", spend_prisma),
patch.object(
spend_prisma.writer_db._original_prisma,
"litellm_spendlogs",
MagicMock(create_many=create_spend_rows),
),
):
respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content").mock(
return_value=httpx.Response(200, content=content)
@ -3483,9 +3487,9 @@ class TestMultiPodBatchCostClaim:
from datetime import datetime
import httpx
import litellm
import respx
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.types.utils import LiteLLMBatch, SpecialEnums
@ -3587,9 +3591,9 @@ class TestMultiPodBatchCostClaim:
import base64
import httpx
import litellm
import respx
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import SpecialEnums
@ -4423,9 +4427,10 @@ class TestBatchLineItemsFromThePoller:
@pytest.mark.asyncio
async def test_a_batch_bigger_than_one_part_reaches_the_callbacks_as_numbered_parts(self, monkeypatch):
import litellm
from litellm_enterprise.proxy.common_utils import check_batch_cost
import litellm
journal: list = []
_, instance = self._setup(journal)
logger = self._recording_logger(journal)
@ -4933,9 +4938,10 @@ async def test_failed_retirement_keeps_batch_available() -> None:
@pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl"))
@pytest.mark.parametrize("other_id", ("{id}4", "prefix-{id}", "{id}/child", "{id}.backup", "{id}?generation=2"))
def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other_id: str) -> None:
from litellm.exceptions import NotFoundError
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
from litellm.exceptions import NotFoundError
error: Final = NotFoundError(
message=f"No object found with id '{other_id.format(id=resource_id)}'.", model="test", llm_provider="openai"
)
@ -4958,9 +4964,10 @@ def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other
),
)
def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, message: str) -> None:
from litellm.exceptions import NotFoundError
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
from litellm.exceptions import NotFoundError
error: Final = NotFoundError(message=message.format(id=resource_id), model="test", llm_provider="openai")
assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is True
assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is True
@ -4970,9 +4977,10 @@ def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, mes
("missing_key", "gone"), (("out/job-123/input.jsonl.out", True), ("out/job-123/other.jsonl.out", False))
)
def test_an_s3_no_such_key_for_the_output_object_is_permanent_loss(missing_key: str, gone: bool) -> None:
from litellm.llms.bedrock.common_utils import BedrockError
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
from litellm.llms.bedrock.common_utils import BedrockError
# S3 GetObject 404 body shape, checked live against a missing key on 2026-10-20
error: Final = BedrockError(
status_code=404,
@ -4985,9 +4993,10 @@ def test_an_s3_no_such_key_for_the_output_object_is_permanent_loss(missing_key:
@pytest.mark.parametrize("output_file_id", (None, ""))
def test_a_404_without_an_output_identity_remains_retryable(output_file_id: str | None) -> None:
from litellm.exceptions import NotFoundError
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
from litellm.exceptions import NotFoundError
error: Final = NotFoundError(message="Output file not found", model="test", llm_provider="openai")
assert CheckBatchCost._is_output_file_gone_at_provider(error, output_file_id) is False
@ -5003,11 +5012,12 @@ async def test_batch_cost_reaches_configured_success_callbacks(
initialized: bool,
claimed: bool,
) -> None:
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.types.utils import LiteLLMBatch
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
events: Final[SimpleQueue[tuple[str, object, int, int]]] = SimpleQueue()

33
tests/unit/proxy/db/test_budget_window_spend_writer.py
View file

@ -407,9 +407,7 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o
@pytest.mark.asyncio
@pytest.mark.parametrize(
("replaces_window_start", "replaced_param"), [(None, None), (WINDOW_A, datetime(2026, 8, 1))]
)
@pytest.mark.parametrize(("replaces_window_start", "replaced_param"), [(None, None), (WINDOW_A, datetime(2026, 8, 1))])
async def test_open_window_spend_row_moves_only_the_window_the_caller_saw_end(replaces_window_start, replaced_param):
db = _FakeDB()
@ -425,7 +423,9 @@ async def test_open_window_spend_row_moves_only_the_window_the_caller_saw_end(re
((query, params),) = db.execute_raw_calls
normalized = " ".join(query.split())
assert "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE" in normalized
assert normalized.endswith("WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')")
assert normalized.endswith(
"WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')"
)
assert params == ("model_max_budget", "k1", "1d", datetime(2026, 8, 31), replaced_param)
@ -561,7 +561,12 @@ async def test_seed_aggregate_sql_splits_the_window_at_the_batch_start(entity_ty
((query, params),) = db.query_raw_calls
normalized = " ".join(query.split())
assert expected_column in normalized
assert "FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC'))" in normalized
assert (
'FILTER (WHERE EXISTS (SELECT 1 FROM "LiteLLM_SpendLogDelivery" delivery '
'WHERE delivery.request_id = "LiteLLM_SpendLogs".request_id) '
"OR \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC'))"
in normalized
)
assert 'FROM "LiteLLM_SpendLogs"' in normalized
# startTime is TIMESTAMP(3): the bound is floored to the second so the
# batch's own earliest row cannot round under it.
@ -635,3 +640,21 @@ async def test_rolling_a_window_row_renders_a_postgres_update_span(
new_window_start=WINDOW_B,
)
assert await postgres_span_names() == ("postgres.update LiteLLM_BudgetWindowSpend",)
@pytest.mark.asyncio
async def test_window_reads_inside_an_enclosing_transaction_use_its_connection():
client_db = _FakeDB()
transaction = _FakeDB(existing_rows=[])
await commit_window_spend_updates(
prisma_client=_FakePrismaClient(client_db),
transactions=(
build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0, started_at=BATCH_STARTED_AT),
),
outer_transaction=transaction,
)
assert client_db.query_raw_calls == []
assert len(transaction.query_raw_calls) == 2
assert len(transaction.batcher.calls) == 1

38
tests/unit/proxy/db/test_db_spend_update_writer.py
View file

@ -21,6 +21,7 @@ from redis.exceptions import DataError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._service_logger import ServiceTypes
from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS
from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem
@ -4382,7 +4383,6 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_only_on_rows_the_bri
assert prisma.spend_log_flush_requested.is_set() is expects_flush
@pytest.mark.asyncio
@pytest.mark.timeout(30)
async def test_a_chatgpt_responses_row_queues_without_a_default_chatgpt_login(
@ -4424,6 +4424,8 @@ def _batch_cost_payload() -> dict:
def _spend_logs_prisma(inserted: int, existing: object, taken_over: int = 1) -> MagicMock:
prisma = _tool_usage_prisma()
prisma.jsonify_object = lambda data: dict(data)
prisma.writer_db = prisma.db
prisma.db.query_raw = AsyncMock(return_value=[])
prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted)
prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing)
prisma.db.litellm_spendlogs.update_many = AsyncMock(return_value=taken_over)
@ -4745,6 +4747,7 @@ def _spend_logs_table_prisma() -> MagicMock:
prisma = _tool_usage_prisma()
prisma.jsonify_object = lambda data: dict(data)
prisma.db.litellm_spendlogs = _SpendLogsTable()
prisma.writer_db.query_raw = AsyncMock(return_value=[])
return prisma
@ -5839,3 +5842,36 @@ async def test_completed_spend_survives_error_log_suppression(
assert len(rows) == 1
assert rows[0]["spend"] == pytest.approx(0.23)
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (10, 5)
@pytest.mark.parametrize("ledger_enabled", [False, True])
@pytest.mark.asyncio
async def test_request_error_counts_survive_the_ledger_path_switch(monkeypatch: pytest.MonkeyPatch, ledger_enabled: bool):
from litellm.proxy import proxy_server
from litellm.proxy.db import db_spend_update_writer as writer_module
from litellm.proxy.db.request_error_tracking import RequestErrorAccumulator
from litellm.types.proxy.request_errors import RequestErrorKey
accumulator = RequestErrorAccumulator()
prisma = _tool_usage_prisma()
prisma.get_request_status.return_value = "failure"
monkeypatch.setattr(writer_module, "request_error_accumulator", accumulator)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "general_settings", {"use_spend_logs_ledger": ledger_enabled})
monkeypatch.setattr(proxy_server, "disable_spend_logs", False)
now = datetime(2026, 11, 2, tzinfo=timezone.utc)
await DBSpendUpdateWriter().update_database(
token="key", user_id="user", team_id=None, org_id=None, end_user_id=None,
kwargs={
"model": "ledger-test", "call_type": "acompletion", "litellm_call_id": "failed-request",
"litellm_params": {"metadata": {
"user_api_key": "key", "user_api_key_user_id": "user", "status": "failure",
"error_information": {"error_code": "503"},
}},
},
completion_response=None, start_time=now, end_time=now, response_cost=0,
)
await drain_logging_tasks()
assert accumulator.drain() == {RequestErrorKey(
date="2026-11-02", api_key=prisma.spend_log_transactions[0]["api_key"], team_id="", user_id="user", model_group="ledger-test", status_code=503,
): 1}

12
tests/unit/proxy/guardrails/test_usage_tracking.py
View file

@ -654,3 +654,15 @@ async def test_one_failing_index_statement_does_not_drop_the_others_or_the_rollu
assert prisma.db.litellm_spendlogguardrailindex.create_many.await_count == 3
assert prisma.db.litellm_dailyguardrailmetrics.upsert.await_count == len(guardrail_ids)
@pytest.mark.asyncio
async def test_disabled_log_deliveries_do_not_enable_guardrail_analytics():
prisma = _prisma()
row = _payload("lean", usage={"topicPolicyUnits": 2})
row["accounting_payload"] = {"request_id": "lean"}
row["accounting_log_disabled"] = True
await process_spend_logs_guardrail_usage(prisma, [row], pending=PendingRollups())
prisma.db.litellm_dailyguardrailmetrics.upsert.assert_not_awaited()
prisma.db.litellm_spendlogguardrailindex.create_many.assert_not_awaited()
prisma.db.litellm_dailyguardrailusageunits.upsert.assert_not_awaited()

4
tests/unit/proxy/hooks/test_proxy_hooks_init.py
View file

@ -33,7 +33,7 @@ def test_managed_vector_stores_hook_registered():
assert hook_cls.__name__ == "_PROXY_LiteLLMManagedVectorStores"
def test_isolation_module_does_not_pull_in_proxy_utils():
def test_isolation_module_does_not_pull_in_proxy_utils(monkeypatch):
"""Layering guard: litellm.llms.* must not transitively import
litellm.proxy.utils, which would reintroduce the import cycle."""
import importlib
@ -44,7 +44,7 @@ def test_isolation_module_does_not_pull_in_proxy_utils():
"litellm.proxy.management_endpoints.common_utils",
"litellm.llms.base_llm.managed_resources.isolation",
]:
sys.modules.pop(mod, None)
monkeypatch.delitem(sys.modules, mod, raising=False)
importlib.import_module("litellm.llms.base_llm.managed_resources.isolation")
assert "litellm.proxy.utils" not in sys.modules

7
tests/unit/proxy/spend_tracking/test_budget_reservation.py
View file

@ -588,3 +588,10 @@ async def test_group_reserves_on_its_current_deployments_after_another_vendors_d
by_own_name: Final = litellm.token_counter(model="prod", messages=_MULTILINGUAL_MESSAGES)
assert openai_only != by_own_name
assert (await _input_tokens(router, "prod"))["prod"] == by_own_name
@pytest.mark.parametrize("window", [{}, {"budget_duration": "1d"}, {"reset_at": "2026-11-03T00:00:00Z"}])
def test_fixed_budget_window_requires_duration_and_reset_boundary(window: object) -> None:
from litellm.proxy.spend_tracking.budget_reservation import fixed_budget_window
assert fixed_budget_window(window) is None

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

@ -47,7 +47,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
_get_spend_logs_metadata,
_get_vector_store_request_for_spend_logs_payload,
is_master_key,
_redact_logged_api_key,
redact_logged_api_key,
_redact_prompt_leaks_in_error_string,
sanitize_error_information_for_spend_logs,
_sanitize_guardrail_information_for_spend_logs,
@ -4038,7 +4038,7 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form(
assert REDACTED_BY_LITELM_STRING in sanitized["error_message"]
# ── _redact_logged_api_key unit tests ──────────────────────────────────────
# ── redact_logged_api_key unit tests ──────────────────────────────────────
@pytest.mark.parametrize(
@ -4067,16 +4067,16 @@ def test_sanitize_error_information_persists_no_raw_model_for_an_unknown_model_r
def test_redact_logged_api_key_none_returns_none():
assert _redact_logged_api_key(None) is None
assert redact_logged_api_key(None) is None
def test_redact_logged_api_key_empty_string_returns_none():
assert _redact_logged_api_key("") is None
assert redact_logged_api_key("") is None
def test_redact_logged_api_key_sk_key_is_hashed():
raw = "sk-9876secret"
result = _redact_logged_api_key(raw)
result = redact_logged_api_key(raw)
assert result == hash_token(raw)
assert result is not None
assert not result.startswith("sk-")
@ -4085,15 +4085,15 @@ def test_redact_logged_api_key_sk_key_is_hashed():
def test_redact_logged_api_key_bearer_sk_equals_sk_hash():
raw = "sk-9876secret"
result_plain = _redact_logged_api_key(raw)
result_bearer = _redact_logged_api_key(f"Bearer {raw}")
result_plain = redact_logged_api_key(raw)
result_bearer = redact_logged_api_key(f"Bearer {raw}")
assert result_bearer == result_plain
def test_redact_logged_api_key_bearer_case_insensitive():
raw = "sk-9876secret"
result_lower = _redact_logged_api_key(f"bearer {raw}")
result_upper = _redact_logged_api_key(f"BEARER {raw}")
result_lower = redact_logged_api_key(f"bearer {raw}")
result_upper = redact_logged_api_key(f"BEARER {raw}")
expected = hash_token(raw)
assert result_lower == expected
assert result_upper == expected
@ -4101,7 +4101,7 @@ def test_redact_logged_api_key_bearer_case_insensitive():
def test_redact_logged_api_key_non_sk_raw_key_is_hashed():
raw = "anthropic-raw-key-xyz"
result = _redact_logged_api_key(raw)
result = redact_logged_api_key(raw)
assert result is not None
assert result != raw
assert len(result) == 64
@ -4111,7 +4111,7 @@ def test_redact_logged_api_key_non_sk_raw_key_is_hashed():
def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag():
already_hashed = hash_token("sk-some-key")
assert len(already_hashed) == 64
result = _redact_logged_api_key(already_hashed, already_redacted=True)
result = redact_logged_api_key(already_hashed, already_redacted=True)
assert result == already_hashed
assert hash_token(already_hashed) != result # no double-hash
@ -4119,7 +4119,7 @@ def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag():
def test_redact_logged_api_key_sha256_without_flag_is_hashed():
already_hashed = hash_token("sk-some-key")
assert len(already_hashed) == 64
result = _redact_logged_api_key(already_hashed)
result = redact_logged_api_key(already_hashed)
assert result is not None
assert result != already_hashed
assert len(result) == 64
@ -4129,7 +4129,7 @@ def test_redact_logged_api_key_sha256_without_flag_is_hashed():
def test_redact_logged_api_key_long_opaque_token_is_hashed():
raw = "x1" * 450
assert len(raw) == 900
result = _redact_logged_api_key(raw)
result = redact_logged_api_key(raw)
assert result is not None
assert result != raw
assert raw not in result
@ -4139,27 +4139,27 @@ def test_redact_logged_api_key_long_opaque_token_is_hashed():
def test_redact_logged_api_key_hashed_jwt_passes_through():
jwt_hash = "hashed-jwt-" + "a" * 64
result = _redact_logged_api_key(jwt_hash, already_redacted=True)
result = redact_logged_api_key(jwt_hash, already_redacted=True)
assert result == jwt_hash
def test_redact_logged_api_key_hashed_jwt_shape_without_provenance_is_hashed():
lookalike = "hashed-jwt-" + "a" * 64
result = _redact_logged_api_key(lookalike)
result = redact_logged_api_key(lookalike)
assert result == hash_token(lookalike)
assert result != lookalike
def test_redact_logged_api_key_hashed_jwt_trailing_newline_is_hashed():
trailing = "hashed-jwt-" + "a" * 64 + "\n"
result = _redact_logged_api_key(trailing, already_redacted=True)
result = redact_logged_api_key(trailing, already_redacted=True)
assert result == hash_token(trailing)
assert result != trailing
def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed():
short_jwt = "hashed-jwt-tooshort"
result = _redact_logged_api_key(short_jwt)
result = redact_logged_api_key(short_jwt)
assert result is not None
assert result != short_jwt
assert len(result) == 64
@ -4169,14 +4169,14 @@ def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed():
def test_redact_logged_api_key_master_key_alias_passes_through():
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS, already_redacted=True)
result = redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS, already_redacted=True)
assert result == LITELLM_PROXY_MASTER_KEY_ALIAS
def test_redact_logged_api_key_master_key_alias_without_provenance_is_hashed():
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS)
result = redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS)
assert result == hash_token(LITELLM_PROXY_MASTER_KEY_ALIAS)
assert result != LITELLM_PROXY_MASTER_KEY_ALIAS
@ -4282,9 +4282,9 @@ def test_get_logging_payload_reads_used_client_oauth_token_from_the_bucket_the_p
def test_redact_logged_api_key_bearer_only_returns_none():
# "bearer " with nothing after stripping is equivalent to no key
assert _redact_logged_api_key("bearer ") is None
assert _redact_logged_api_key("Bearer ") is None
assert _redact_logged_api_key("BEARER ") is None
assert redact_logged_api_key("bearer ") is None
assert redact_logged_api_key("Bearer ") is None
assert redact_logged_api_key("BEARER ") is None
# ── _get_spend_logs_metadata key-hash invariant tests ─────────────────────
@ -4645,7 +4645,7 @@ class TestSpendLogKeyRedaction:
def test_bearer_prefixed_sk_key_is_hashed(self):
raw = "Bearer sk-WLi4iRn4JmbVlTaYw12IOA"
result = _redact_logged_api_key(raw)
result = redact_logged_api_key(raw)
assert result is not None
assert not result.startswith("Bearer")
assert not result.startswith("sk-")
@ -4653,14 +4653,14 @@ class TestSpendLogKeyRedaction:
def test_bare_sk_key_is_hashed(self):
raw = "sk-WLi4iRn4JmbVlTaYw12IOA"
result = _redact_logged_api_key(raw)
result = redact_logged_api_key(raw)
assert result is not None
assert not result.startswith("sk-")
assert len(result) == 64
def test_bearer_lowercase_is_handled(self):
raw = "bearer sk-WLi4iRn4JmbVlTaYw12IOA"
result = _redact_logged_api_key(raw)
result = redact_logged_api_key(raw)
assert result is not None
assert not result.startswith("bearer")
assert not result.startswith("sk-")
@ -4668,11 +4668,11 @@ class TestSpendLogKeyRedaction:
def test_already_hashed_key_unchanged(self):
hashed = "bcfe8173f5447f10be0e7fb37aaa8b97829d5c9e0498232152f9d123456789ab"
assert _redact_logged_api_key(hashed, already_redacted=True) == hashed
assert redact_logged_api_key(hashed, already_redacted=True) == hashed
def test_bearer_prefixed_non_sk_key_is_hashed(self):
raw = "Bearer some-other-token-format"
result = _redact_logged_api_key(raw)
result = redact_logged_api_key(raw)
assert result == hash_token("some-other-token-format")
assert result is not None
assert not result.startswith("Bearer")
@ -4680,7 +4680,7 @@ class TestSpendLogKeyRedaction:
def test_bearer_and_bare_produce_same_hash(self):
bare = "sk-WLi4iRn4JmbVlTaYw12IOA"
bearer = "Bearer sk-WLi4iRn4JmbVlTaYw12IOA"
assert _redact_logged_api_key(bare) == _redact_logged_api_key(bearer)
assert redact_logged_api_key(bare) == redact_logged_api_key(bearer)
@patch("litellm.proxy.proxy_server.master_key", None)
@ -4777,7 +4777,7 @@ def test_get_logging_payload_keeps_internal_service_account_key_readable(service
def test_redact_logged_api_key_service_account_name_without_provenance_is_hashed():
result = _redact_logged_api_key(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME)
result = redact_logged_api_key(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME)
assert result == hash_token(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME)
@ -5889,7 +5889,7 @@ def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal
assert payload["api_key"] == "", (
f"Expected empty string but got {payload['api_key']!r}; "
"dropping _redact_logged_api_key's 'or \"\"' guard would yield 'None' here"
"dropping redact_logged_api_key's 'or \"\"' guard would yield 'None' here"
)
@ -5940,7 +5940,7 @@ def test_get_logging_payload_stores_user_email_in_spend_log_metadata(user_email,
def test_redact_logged_api_key_partial_sha256_is_hashed():
partial_hex = "a" * 63
result = _redact_logged_api_key(partial_hex)
result = redact_logged_api_key(partial_hex)
assert result is not None
assert result != partial_hex
assert len(result) == 64
@ -5950,7 +5950,7 @@ def test_redact_logged_api_key_partial_sha256_is_hashed():
def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag():
already_hashed = hash_token("sk-some-key")
assert len(already_hashed) == 64
result = _redact_logged_api_key(f"Bearer {already_hashed}", already_redacted=True)
result = redact_logged_api_key(f"Bearer {already_hashed}", already_redacted=True)
assert result == already_hashed
assert hash_token(already_hashed) != result
@ -5958,7 +5958,7 @@ def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag():
def test_redact_logged_api_key_bearer_sha256_without_flag_is_hashed():
already_hashed = hash_token("sk-some-key")
assert len(already_hashed) == 64
result = _redact_logged_api_key(f"Bearer {already_hashed}")
result = redact_logged_api_key(f"Bearer {already_hashed}")
assert result is not None
assert result != already_hashed
assert result == hash_token(already_hashed)
@ -6570,12 +6570,12 @@ def test_get_logging_payload_never_lands_a_raw_cli_session_token(call_type: str)
def test_redact_logged_api_key_cli_session_alias_needs_alias_provenance():
assert _redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True, key_alias=_CLI_SESSION_ALIAS) == (
assert redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True, key_alias=_CLI_SESSION_ALIAS) == (
_CLI_SESSION_ALIAS
)
assert _redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True) == hash_token(_CLI_SESSION_ALIAS)
assert _redact_logged_api_key(_CLI_SESSION_ALIAS, key_alias=_CLI_SESSION_ALIAS) == hash_token(_CLI_SESSION_ALIAS)
assert _redact_logged_api_key("alice", already_redacted=True, key_alias="alice") == hash_token("alice")
assert redact_logged_api_key(_CLI_SESSION_ALIAS, already_redacted=True) == hash_token(_CLI_SESSION_ALIAS)
assert redact_logged_api_key(_CLI_SESSION_ALIAS, key_alias=_CLI_SESSION_ALIAS) == hash_token(_CLI_SESSION_ALIAS)
assert redact_logged_api_key("alice", already_redacted=True, key_alias="alice") == hash_token("alice")
def test_azure_spillover_stamped_from_response_headers():
@ -7814,3 +7814,45 @@ def test_spend_payload_preserves_unknown_completion_start(completion_start: str
assert payload["completionStartTime"] == (
end if completion_start == "absent" else reported if completion_start else None
)
def test_spend_log_metadata_retention_drops_accounting_inputs_it_was_not_told_to_keep() -> None:
metadata: Final = {
"status": "success",
"accounting_savings": {
"compression": 0.4,
"prompt_caching": 0.1,
"gateway_injected_caching": 0.2,
"autorouter": 0.3,
},
"accounting_windows": [
{"entity_type": "key", "entity_id": "key", "window_duration": "1d", "window_start": "2026-11-02T00:00:00"}
],
"accounting_model_access_groups": ["group"],
"user_api_key_project_id": "project",
"usage_object": {"prompt_tokens_details": {"cached_tokens": 2}},
"compression_savings": {"tokens_saved": 3},
"guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 2}}],
"internal_call_origin": "classifier",
"billed_retry_attempt": True,
"user_api_key_alias": "discard",
}
row: Final = {
"request_id": "retained-accounting",
"spend": 0.5,
"startTime": "2026-11-02T12:00:00",
"endTime": "2026-11-02T12:00:01",
"model": "model",
"model_group": "group",
"api_key": "key",
"custom_llm_provider": "openai",
"user": "user",
"team_id": "team",
"agent_id": "agent",
"request_tags": '["tag"]',
"metadata": json.dumps(metadata),
}
stored: Final = spend_log_row_with_retained_metadata(row, SpendLogsMetadataFields(include=("usage_object",)))
parsed: Final = json.loads(str(stored["metadata"]))
assert "user_api_key_alias" not in parsed
assert parsed == {"usage_object": metadata["usage_object"], "status": metadata["status"]}

60
tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
View file

@ -16589,3 +16589,63 @@ async def test_initialize_prints_the_banner_as_one_json_record_in_json_mode(
banner_records: Final = [record for record in records if record["message"] == LITELLM_BANNER]
assert len(banner_records) == 1, records
assert banner_records[0]["level"] == "INFO"
@contextlib.contextmanager
def _accounting_window_env(cached: Mapping[str, object], client: MagicMock) -> Iterator[None]:
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
cache: Final = UserApiKeyCache()
for key, value in cached.items():
cache.set_cache(key=key, value=value)
originals: Final = (proxy_server_module.user_api_key_cache, proxy_server_module.prisma_client)
proxy_server_module.user_api_key_cache = cache
proxy_server_module.prisma_client = client
try:
yield
finally:
proxy_server_module.user_api_key_cache, proxy_server_module.prisma_client = originals
@pytest.mark.asyncio
async def test_accounting_window_keys_cover_the_stored_key_and_the_cached_team():
key_reset: Final = datetime(2026, 11, 13, tzinfo=timezone.utc)
team_reset: Final = datetime(2026, 11, 6, tzinfo=timezone.utc)
key_row: Final = _StoredRow(
budget_limits=[{"budget_duration": "30d", "max_budget": 100.0, "reset_at": key_reset.isoformat()}]
)
team_obj: Final = MagicMock()
team_obj.budget_limits = [{"budget_duration": "7d", "max_budget": 50.0, "reset_at": team_reset.isoformat()}]
with _accounting_window_env({"team_id:team-1": team_obj}, _prisma_with_rows(keys={"hashed-key": key_row})):
windows = await proxy_server_module.accounting_window_keys("hashed-key", "team-1")
assert windows == (
{
"entity_type": "key",
"entity_id": "hashed-key",
"window_duration": "30d",
"window_start": (key_reset - timedelta(days=30)).isoformat(),
},
{
"entity_type": "team",
"entity_id": "team-1",
"window_duration": "7d",
"window_start": (team_reset - timedelta(days=7)).isoformat(),
},
)
@pytest.mark.asyncio
async def test_accounting_window_keys_remember_a_key_without_windows_but_retry_a_failed_read():
client: Final = _prisma_with_rows(keys={"no-windows": _StoredRow(budget_limits=[])})
with _accounting_window_env({}, client):
assert await proxy_server_module.accounting_window_keys("no-windows", None) == ()
assert await proxy_server_module.accounting_window_keys("no-windows", None) == ()
assert client.db.litellm_verificationtoken.find_unique.await_count == 1
client.db.litellm_verificationtoken.find_unique.side_effect = RuntimeError("database unavailable")
assert await proxy_server_module.accounting_window_keys("unreadable", None) == ()
assert await proxy_server_module.accounting_window_keys("unreadable", None) == ()
assert client.db.litellm_verificationtoken.find_unique.await_count == 3

122
tests/unit/proxy/test_update_spend.py
View file

@ -351,3 +351,125 @@ async def test_update_spend_logs_drops_row_prisma_cannot_encode_and_keeps_the_re
assert sorted(written) == ["good-1", "good-2"]
assert prisma_client.spend_log_transactions == []
proxy_logging_obj.failure_handler.assert_not_called()
def _raw_query_error(sqlstate: str):
from prisma.errors import RawQueryError
return RawQueryError(
{
"user_facing_error": {
"error_code": "P2010",
"meta": {"code": sqlstate, "message": f"ERROR: sqlstate {sqlstate}"},
}
}
)
class _RawInsertRepository:
def __init__(self, sqlstate: str):
self.sqlstate = sqlstate
self.written: list[str] = []
async def insert_with_deliveries(self, rows):
if any(row["request_id"] == "rejected" for row in rows):
raise _raw_query_error(self.sqlstate)
self.written.extend(str(row["request_id"]) for row in rows)
@pytest.mark.asyncio
@pytest.mark.parametrize("sqlstate", ["22P02", "23502", "54000"])
async def test_raw_insert_rejecting_one_row_drops_only_that_row(sqlstate):
from litellm.proxy.utils import _create_spend_logs_with_poison_isolation
repo = _RawInsertRepository(sqlstate)
rows = [{"request_id": "good-1"}, {"request_id": "rejected"}, {"request_id": "good-2"}]
await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10)
assert sorted(repo.written) == ["good-1", "good-2"]
@pytest.mark.asyncio
@pytest.mark.parametrize("sqlstate", ["42P01", "40P01"])
async def test_systemic_raw_insert_failure_is_retried_not_dropped(sqlstate):
from prisma.errors import RawQueryError
from litellm.proxy.utils import _create_spend_logs_with_poison_isolation
repo = _RawInsertRepository(sqlstate)
rows = [{"request_id": "good-1"}, {"request_id": "rejected"}]
with pytest.raises(RawQueryError):
await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10)
assert repo.written == []
class _RescueFailsRepository:
def __init__(self, rescue_error: Exception):
self.rescue_error = rescue_error
async def insert_with_deliveries(self, rows):
if any(row.get("accounting_log_disabled") for row in rows):
raise self.rescue_error
if any(row["request_id"] == "rejected" for row in rows):
raise _raw_query_error("22P02")
@pytest.mark.asyncio
async def test_a_transient_failure_while_keeping_a_dropped_rows_charge_is_retried_not_lost():
from litellm.proxy.utils import _create_spend_logs_with_poison_isolation
outage = _raw_query_error("40P01")
repo = _RescueFailsRepository(outage)
rows = [{"request_id": "rejected", "accounting_payload": {"spend": 1.0}}]
with pytest.raises(type(outage)):
await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10)
@pytest.mark.asyncio
async def test_a_rejected_charge_rescue_is_logged_and_the_flush_continues():
from litellm.proxy.utils import _create_spend_logs_with_poison_isolation
repo = _RescueFailsRepository(_raw_query_error("22P02"))
rows = [{"request_id": "rejected", "accounting_payload": {"spend": 1.0}}]
assert await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) == 9
class _StatementValidatingRepository:
def __init__(self):
self.written = []
self.charges_kept = []
async def insert_with_deliveries(self, rows):
from litellm.repositories.spend_log_insert import spend_log_insert_statement
if all(row.get("accounting_log_disabled") for row in rows):
self.charges_kept.extend(row["request_id"] for row in rows)
return len(rows)
spend_log_insert_statement(rows)
self.written.extend(row["request_id"] for row in rows)
return len(rows)
@pytest.mark.asyncio
async def test_a_row_the_insert_statement_rejects_is_dropped_alone_and_keeps_its_charge():
from litellm.proxy.utils import _create_spend_logs_with_poison_isolation
repo = _StatementValidatingRepository()
rows = [
{"request_id": "good", "call_type": "acompletion", "messages": "[]"},
{
"request_id": "unparseable",
"call_type": "acompletion",
"messages": "not json",
"accounting_payload": {"spend": 1.0},
},
]
assert await _create_spend_logs_with_poison_isolation(repo, rows, failure_budget=10) < 10
assert repo.written == ["good"]
assert repo.charges_kept == ["unparseable"]

31
tests/unit/repositories/test_spend_log_insert.py Normal file
View file

@ -0,0 +1,31 @@
from datetime import datetime, timezone
import json
import pytest
from litellm.repositories.spend_log_insert import spend_log_insert_statement
@pytest.mark.parametrize(
"rows",
[(), ({"request_id": "request", "unknown_column": 1},), ({"request_id": "first"}, {"spend": 1})],
)
def test_insert_rejects_empty_unknown_or_inconsistent_columns(rows):
with pytest.raises(ValueError, match="Spend log insert requires"):
spend_log_insert_statement(rows)
def test_insert_keeps_json_and_timestamps_in_bound_data():
timestamp = datetime(2026, 11, 2, tzinfo=timezone.utc)
request_id = "request'); DROP TABLE unrelated; --"
statement, encoded = spend_log_insert_statement(({
"request_id": request_id, "startTime": timestamp, "metadata": '{"project":"project"}',
"request_tags": '["tag"]', "messages": '[{"role":"user","content":"hello"}]',
},))
row = json.loads(encoded)[0]
assert request_id not in statement
assert row["request_id"] == request_id
assert row["startTime"] == timestamp.isoformat()
assert row["metadata"] == {"project": "project"}
assert row["request_tags"] == ["tag"]
assert row["messages"] == [{"role": "user", "content": "hello"}]

6
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -30862,6 +30862,12 @@ export interface components {
* @description decrypt keys with google kms
*/
use_google_kms?: boolean | null;
/**
* Use Spend Logs Ledger
* @description Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.
* @default false
*/
use_spend_logs_ledger: boolean;
/**
* Use Spend Logs Partitioning
* @description If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False.