scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267

feat(accounting): add durable event ledger writer (#48883)

Browse source 
* feat(accounting): add durable event ledger writer

* fix(accounting): validate delegated transaction managers locally

* docs(accounting): clarify ledger integration contracts

* fix(accounting): return ledger conflicts as values and keep cleanup on the kind index

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(accounting): import assert_never from typing_extensions for Python 3.10

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(accounting): annotate ledger integration test locals as Final

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(accounting): finish ledger test conventions

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
joshua-berri 2026-10-28 16:11:51 -07:00 • committed by GitHub
parent 3f526a5017
commit d30c702c5c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 781 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

269
litellm/proxy/db/accounting_ledger.py Normal file
View file

@ -0,0 +1,269 @@
from __future__ import annotations
import hashlib
import json
from collections.abc import AsyncGenerator, Callable
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from pydantic import AwareDatetime, ConfigDict, JsonValue, TypeAdapter
from typing_extensions import assert_never
from litellm.accounting import CostEvent, DeliveryTarget, EventKind, InboundUsageEvent, OutboundUsageEvent
from litellm.proxy.db.db_span import db_span
from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
if TYPE_CHECKING:
from prisma import Prisma
from litellm.proxy.utils import PrismaClient
AccountingEvent: TypeAlias = InboundUsageEvent | OutboundUsageEvent | CostEvent
_JSON: Final = TypeAdapter(dict[str, JsonValue])
_INSTANT: Final[TypeAdapter[datetime]] = TypeAdapter(AwareDatetime)
_STRINGS: Final = TypeAdapter(tuple[dict[str, str], ...])
_TIMEOUT: Final = timedelta(seconds=10)
class AccountingConflictError(ValueError):
pass
_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"]
@dataclass(frozen=True, slots=True)
class _AcceptedEvent:
event_id: str
kind: EventKind
source_event_id: str | None
occurred_at: datetime
payload: str
payload_hash: str
targets: tuple[str, ...]
schema_version: int
version: int
def _json(value: object) -> str:
return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False)
def _prepare(
event: AccountingEvent,
targets: tuple[DeliveryTarget, ...],
envelope: str,
schema_version: int,
version: int,
) -> _AcceptedEvent:
if type(schema_version) is not int or type(version) is not int or min(schema_version, version) < 1:
raise ValueError("Accounting versions must be positive integers")
validated: Final = type(event).model_validate_json(event.model_dump_json())
payload: Final = _JSON.validate_json(envelope)
if payload.get("event") != validated.model_dump(mode="json"):
raise ValueError("Accounting envelope must contain the matching public event")
destinations: Final = tuple(sorted(frozenset(DeliveryTarget(target).value for target in targets)))
serialized: Final = _json(payload)
identity: Final = _json((schema_version, version, destinations, payload))
kind: Final = (
EventKind.COST
if isinstance(validated, CostEvent)
else EventKind.INBOUND_USAGE
if isinstance(validated, InboundUsageEvent)
else EventKind.OUTBOUND_USAGE
)
return _AcceptedEvent(
event_id=validated.event_id,
kind=kind,
source_event_id=validated.usage_event_id if isinstance(validated, CostEvent) else None,
occurred_at=validated.occurred_at if isinstance(validated, CostEvent) else validated.ended_at,
payload=serialized,
payload_hash=hashlib.sha256(identity.encode()).hexdigest(),
targets=destinations,
schema_version=schema_version,
version=version,
)
_INSERT_EVENT: Final = """
INSERT INTO "LiteLLM_AccountingEvent"
(event_id, kind, source_event_id, schema_version, payload_hash, occurred_at, payload)
VALUES ($1, $2::"AccountingEventKind", $3, $4::int, $5, $6::timestamp, $7::jsonb)
ON CONFLICT (event_id) DO NOTHING RETURNING event_id
"""
_INSERT_DELIVERIES: Final = """
INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version)
SELECT $1, value::"AccountingDeliveryTarget", $3::int FROM jsonb_array_elements_text($2::jsonb)
"""
_LOCK_SOURCE: Final = """
SELECT event_id FROM "LiteLLM_AccountingEvent"
WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE
"""
_EXISTING_COST: Final = """
SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_event_id = $1
"""
_LOCK_EVENT: Final = """
SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE
"""
_COMPLETED_FAMILY: Final = """
root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage')
AND root.recorded_at < $1::timestamp
AND (root.kind = 'inbound_usage' OR EXISTS (
SELECT 1 FROM "LiteLLM_AccountingEvent" AS cost WHERE cost.source_event_id = root.event_id AND cost.kind = 'cost'
))
AND NOT EXISTS (
SELECT 1 FROM "LiteLLM_AccountingEvent" AS member
WHERE (member.event_id = root.event_id OR (member.kind = 'cost' AND member.source_event_id = root.event_id))
AND (member.recorded_at >= $1::timestamp
OR EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" AS delivery
WHERE delivery.event_id = member.event_id AND delivery.state != 'applied'))
)
"""
_RETENTION_CANDIDATES: Final = (
'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE '
+ _COMPLETED_FAMILY
+ " ORDER BY root.recorded_at, root.event_id LIMIT $2::int FOR UPDATE OF root SKIP LOCKED"
)
_ELIGIBLE_GROUP: Final = (
'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE root.event_id = $2 AND ' + _COMPLETED_FAMILY
)
_DELETE_DELIVERIES: Final = """
DELETE FROM "LiteLLM_AccountingDelivery" WHERE event_id IN (
SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1)
)
"""
_DELETE_EVENTS: Final = """
DELETE FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1)
"""
_Conflict: TypeAlias = Literal["missing_source", "competing_cost", "identity_conflict"]
async def _accept(db: Prisma, record: _AcceptedEvent) -> Literal["inserted", "duplicate"] | _Conflict:
if record.source_event_id is not None:
source: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_SOURCE, record.source_event_id))
if not source:
return "missing_source"
costs: Final = _STRINGS.validate_python(tuple(await db.query_raw(_EXISTING_COST, record.source_event_id)))
if any(cost["event_id"] != record.event_id for cost in costs):
return "competing_cost"
inserted: Final = _STRINGS.validate_python(
await db.query_raw(
_INSERT_EVENT,
record.event_id,
record.kind.value,
record.source_event_id,
record.schema_version,
record.payload_hash,
record.occurred_at.astimezone(timezone.utc).isoformat(),
record.payload,
)
)
if not inserted:
existing: Final = _STRINGS.validate_python(tuple(await db.query_raw(_LOCK_EVENT, record.event_id)))
if not existing or existing[0]["payload_hash"] != record.payload_hash:
return "identity_conflict"
return "duplicate"
await db.execute_raw(_INSERT_DELIVERIES, record.event_id, _json(record.targets), record.version)
return "inserted"
async def _delete_completed(db: Prisma, event_id: str, cutoff: str) -> int:
if not _STRINGS.validate_python(await db.query_raw(_ELIGIBLE_GROUP, cutoff, event_id)):
return 0
await db.execute_raw(_DELETE_DELIVERIES, event_id)
return await db.execute_raw(_DELETE_EVENTS, event_id)
class AccountingLedger:
def __init__(self, transaction: Callable[[], _Transaction]) -> None:
self._transaction: Final = transaction
@classmethod
def for_client(cls, client: PrismaClient) -> AccountingLedger:
from prisma.client import TransactionManager
managers: Final = TypeAdapter(TransactionManager, config=ConfigDict(arbitrary_types_allowed=True))
@asynccontextmanager
async def transaction() -> AsyncGenerator[Prisma]:
primary: Final = writer_wrapper(client.db)
manager: Final = managers.validate_python(primary.tx(timeout=_TIMEOUT))
async with db_span("accounting_ledger", "LiteLLM_AccountingEvent"), manager as db:
await db.execute_raw("SET LOCAL statement_timeout = 5000")
await db.execute_raw("SET LOCAL lock_timeout = 1000")
yield db
return cls(transaction)
async def insert_event(
self,
event: AccountingEvent,
targets: tuple[DeliveryTarget, ...],
*,
envelope: str,
schema_version: int = 1,
version: int = 1,
) -> Literal["inserted", "duplicate"]:
"""Accept T5/T6's serialized envelope, containing its public event under `event`.
The immutable acceptance hash includes the target set and delivery version.
Changing either for an accepted event conflicts. T7 must provide a separate
add-target or new-version workflow that preserves receipts and prevents
duplicate contributions; redelivery cannot silently create or reset work.
The Prisma adapter opens one interactive transaction per event, including
two SET LOCAL statements and a source-row lock for costs. I/O is async,
but awaiting this method inline adds commit latency per event, potentially
several times per request. T5 owns bounded asynchronous capture and an
explicit batching strategy; looping over this method does not batch commits.
Success acknowledges commit, independently of HTTP response acknowledgement.
"""
record: Final = _prepare(event, targets, envelope, schema_version, version)
async with self._transaction() as db:
result: Final = await _accept(db, record)
match result:
case "inserted" | "duplicate":
return result
case "missing_source":
raise AccountingConflictError("Cost source must be a durably accepted outbound usage event")
case "competing_cost":
raise AccountingConflictError("Usage event already has an authoritative cost")
case "identity_conflict":
raise AccountingConflictError("Event identity conflicts with its accepted content or delivery set")
case _:
assert_never(result)
async def purge_completed(
self,
*,
redelivery_before: datetime | None = None,
recovery_before: datetime | None = None,
limit: int = 100,
) -> int:
"""Delete at most `limit` completed families once both retention owners permit it.
T8 must supply a coherent recovery watermark including outstanding holds.
Omitting either watermark disables deletion. Each usage row and its cost
remain together so a retained charge never loses its source or vice versa.
Applied receipts are deleted with the family, ending its retained dedup
protection. The caller must guarantee events recorded before
redelivery_before cannot be redelivered through supported paths.
T8's recovery_before must also protect checkpoints,
outstanding holds and unfinished settlements; age alone is insufficient.
"""
if type(limit) is not int or not 1 <= limit <= 1000:
raise ValueError("Accounting retention batch size must be between 1 and 1000")
if redelivery_before is None or recovery_before is None:
return 0
cutoff: Final = min(
_INSTANT.validate_python(redelivery_before).astimezone(timezone.utc),
_INSTANT.validate_python(recovery_before).astimezone(timezone.utc),
).isoformat()
async with self._transaction() as db:
roots: Final = _STRINGS.validate_python(tuple(await db.query_raw(_RETENTION_CANDIDATES, cutoff, limit)))
deleted: Final = tuple([await _delete_completed(db, root["event_id"], cutoff) for root in roots])
return sum(deleted)

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

@ -1,16 +1,48 @@
from __future__ import annotations
import asyncio
import json
import os
import subprocess
import sys
from collections.abc import Iterator
from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from decimal import Decimal
from itertools import chain
from pathlib import Path
from typing import Final
from typing import TYPE_CHECKING, Final, TypeAlias
import psycopg
import pytest
import pytest_asyncio
from psycopg import sql
from pydantic import JsonValue
from litellm.accounting import DeliveryState, DeliveryTarget, EventKind
from litellm.accounting import (
CostEvent,
DeliveryState,
DeliveryTarget,
EventKind,
InboundUsageEvent,
Metric,
MetricName,
OutboundUsageEvent,
Outcome,
PricingSnapshot,
)
from litellm.proxy.db import accounting_ledger
from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger
from litellm.types.utils import CallTypes
from tests.integration._support.database import scratch_database
from tests.integration._support.prisma_clients import connected_proxy_prisma
if TYPE_CHECKING:
from prisma import Prisma
from litellm.proxy.utils import PrismaClient
_Writer: TypeAlias = tuple["PrismaClient", AccountingLedger]
_SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma"
_EVENT_INSERT: Final = (
@ -126,3 +158,377 @@ def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Conne
(DeliveryTarget.KEY_SPEND.value,),
).fetchone()
assert row == ("pending", 0)
_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
_CUTOFF: Final = datetime(2026, 2, 1, tzinfo=timezone.utc)
_TARGETS: Final = (DeliveryTarget.USER_SPEND, DeliveryTarget.TEAM_SPEND)
def _usage(event_id: str = "usage", *, inbound: bool = False) -> InboundUsageEvent | OutboundUsageEvent:
fields: Final = dict(
event_id=event_id,
request_id="request",
call_type=CallTypes.completion,
started_at=_AT,
ended_at=_AT,
outcome=Outcome.SUCCEEDED,
dimensions=(),
)
if inbound:
return InboundUsageEvent(**fields, metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),))
return OutboundUsageEvent(**fields, attempt_id=event_id, retry_of_attempt_id=None, metrics=None)
def _cost(source: str = "usage", event_id: str = "cost", amount: str = "0.125") -> CostEvent:
return CostEvent(
event_id=event_id,
usage_event_id=source,
occurred_at=_AT,
dimensions=(),
charges=(Metric(name=MetricName.SPEND, value=Decimal(amount), unit="USD"),),
snapshot=PricingSnapshot(
calculator_version="fixture",
model="fixture",
provider=None,
pricing_source="cost_map",
priced_model_key="fixture",
total_cost=Decimal(amount),
),
)
def _envelope(event: AccountingEvent, *, context: Mapping[str, JsonValue] | None = None) -> str:
return json.dumps({"event": event.model_dump(mode="json"), "context": context or {}})
async def _insert(
store: AccountingLedger,
event: AccountingEvent,
targets: tuple[DeliveryTarget, ...] = _TARGETS,
*,
context: Mapping[str, JsonValue] | None = None,
) -> str:
return await store.insert_event(event, targets, envelope=_envelope(event, context=context))
def _rows(database_url: str, statement: str, parameters: tuple[object, ...] = ()) -> Sequence[tuple[object, ...]]:
with psycopg.connect(database_url) as connection:
cursor: Final = connection.execute(statement, parameters)
return cursor.fetchall() if cursor.description is not None else []
def _complete_and_age(database_url: str) -> None:
with psycopg.connect(database_url) as connection:
connection.execute('UPDATE "LiteLLM_AccountingEvent" SET recorded_at = %s', (_AT,))
connection.execute("UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'applied', applied_at = %s", (_AT,))
@pytest_asyncio.fixture(loop_scope="function")
async def writer(ledger_database: str, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_Writer]:
monkeypatch.setenv("DATABASE_URL", ledger_database)
monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
_rows(ledger_database, 'TRUNCATE "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"')
async with connected_proxy_prisma(ledger_database) as client:
yield client, AccountingLedger.for_client(client)
@pytest.mark.asyncio
async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, ledger_database: str) -> None:
_, store = writer
event: Final = _usage()
assert await _insert(store, event, context={"requested_region": "eu", "quantity": "1.25"}) == "inserted"
rows: Final = _rows(ledger_database, 'SELECT kind::text, source_event_id, payload FROM "LiteLLM_AccountingEvent"')
assert rows == [
("outbound_usage", None, json.loads(_envelope(event, context={"requested_region": "eu", "quantity": "1.25"})))
]
deliveries: Final = _rows(
ledger_database, 'SELECT target::text, state::text, attempts FROM "LiteLLM_AccountingDelivery"'
)
assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS}
@pytest.mark.asyncio
async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None:
_, store = writer
event: Final = _usage(inbound=True)
results: Final = await asyncio.gather(*(_insert(store, event) for _ in range(8)))
assert results.count("inserted") == 1 and results.count("duplicate") == 7
_complete_and_age(ledger_database)
before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target')
reordered: Final = json.dumps(json.loads(_envelope(event)), sort_keys=True, indent=2)
assert await store.insert_event(event, tuple(reversed(_TARGETS)) + _TARGETS, envelope=reordered) == "duplicate"
assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') == before
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)]
@pytest.mark.asyncio
@pytest.mark.parametrize("change", ["payload", "context", "targets", "schema_version", "version"])
async def test_conflicting_acceptance_cannot_overwrite_or_add_work(
writer: _Writer, ledger_database: str, change: str
) -> None:
_, store = writer
event: Final = _usage()
await _insert(store, event)
altered: Final = event.model_copy(update={"outcome": Outcome.FAILED}) if change == "payload" else event
targets: Final = (DeliveryTarget.KEY_SPEND,) if change == "targets" else _TARGETS
context: Final = {"override": "0.5"} if change == "context" else {}
options: Final = {change: 2} if change in ("schema_version", "version") else {}
with pytest.raises(AccountingConflictError, match="conflicts"):
await store.insert_event(altered, targets, envelope=_envelope(altered, context=context), **options)
assert _rows(ledger_database, 'SELECT payload FROM "LiteLLM_AccountingEvent"') == [(json.loads(_envelope(event)),)]
assert set(_rows(ledger_database, 'SELECT target::text FROM "LiteLLM_AccountingDelivery"')) == {
(t.value,) for t in _TARGETS
}
@pytest.mark.asyncio
async def test_competing_cost_identities_accept_only_one_authoritative_charge(
writer: _Writer, ledger_database: str
) -> None:
_, store = writer
await _insert(store, _usage())
results: Final = await asyncio.gather(
_insert(store, _cost()), _insert(store, _cost(event_id="other-cost")), return_exceptions=True
)
assert sum(result == "inserted" for result in results) == 1
assert sum(isinstance(result, AccountingConflictError) for result in results) == 1
accepted: Final = _rows(
ledger_database, "SELECT event_id, payload FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'"
)
assert len(accepted) == 1
event: Final = CostEvent.model_validate_json(json.dumps(accepted[0][1]["event"]))
assert await _insert(store, event) == "duplicate"
assert event.snapshot.total_cost == Decimal("0.125")
with pytest.raises(AccountingConflictError):
await _insert(store, _cost(event_id=event.event_id, amount="0.5"))
@pytest.mark.asyncio
@pytest.mark.parametrize("source", ["missing", "inbound", "cost"])
async def test_cost_requires_an_existing_outbound_source(writer: _Writer, ledger_database: str, source: str) -> None:
_, store = writer
if source == "inbound":
await _insert(store, _usage("source", inbound=True))
if source == "cost":
await _insert(store, _usage())
await _insert(store, _cost(event_id="source"))
with pytest.raises(AccountingConflictError, match="outbound usage"):
await _insert(store, _cost(source="source", event_id="invalid"))
assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'invalid'") == [
(0,)
]
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"])
async def test_interruption_rolls_back_or_retry_recognizes_commit(
writer: _Writer, ledger_database: str, failure: str
) -> None:
client, store = writer
@asynccontextmanager
async def interrupted() -> AsyncGenerator[Prisma]:
async with client.writer_db.tx() as tx:
yield tx
if failure == "before_commit":
raise RuntimeError("injected pre-commit failure")
if failure == "cancellation":
raise asyncio.CancelledError
raise RuntimeError("injected lost acknowledgement")
error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError
with pytest.raises(error):
await _insert(AccountingLedger(interrupted), _usage())
expected: Final = 1 if failure == "after_commit" else 0
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(expected,)]
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(expected * len(_TARGETS),)]
assert await _insert(store, _usage()) == ("duplicate" if expected else "inserted")
@pytest.mark.asyncio
async def test_delivery_insert_failure_rolls_back_event(writer: _Writer, ledger_database: str) -> None:
from prisma.errors import RawQueryError
_, store = writer
with psycopg.connect(ledger_database) as connection:
connection.execute(
'ALTER TABLE "LiteLLM_AccountingDelivery" ADD CONSTRAINT reject_delivery CHECK (version < 1)'
)
try:
with pytest.raises(RawQueryError, match="reject_delivery"):
await _insert(store, _usage())
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)]
finally:
with psycopg.connect(ledger_database) as connection:
connection.execute('ALTER TABLE "LiteLLM_AccountingDelivery" DROP CONSTRAINT reject_delivery')
@pytest.mark.asyncio
@pytest.mark.parametrize("blocked", ["pending", "parked", "unpriced", "recent_cost", "recovery"])
async def test_retention_preserves_incomplete_or_recovery_needed_families(
writer: _Writer, ledger_database: str, blocked: str
) -> None:
_, store = writer
await _insert(store, _usage())
if blocked != "unpriced":
await _insert(store, _cost())
_complete_and_age(ledger_database)
with psycopg.connect(ledger_database) as connection:
if blocked in ("pending", "parked"):
connection.execute(
"UPDATE \"LiteLLM_AccountingDelivery\" SET state = %s WHERE event_id = 'cost'", (blocked,)
)
if blocked == "recent_cost":
connection.execute(
"UPDATE \"LiteLLM_AccountingEvent\" SET recorded_at = %s WHERE event_id = 'cost'", (_CUTOFF,)
)
before: Final = _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id')
assert (
await store.purge_completed(
redelivery_before=_CUTOFF, recovery_before=_AT if blocked == "recovery" else _CUTOFF
)
== 0
)
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == before
@pytest.mark.asyncio
async def test_retention_removes_completed_family_and_receipts_together(writer: _Writer, ledger_database: str) -> None:
_, store = writer
await _insert(store, _usage())
await _insert(store, _cost())
_complete_and_age(ledger_database)
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)]
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)]
@pytest.mark.asyncio
async def test_retention_respects_batch_bound_and_skips_locked_families(writer: _Writer, ledger_database: str) -> None:
client, store = writer
for name in ("a", "b", "c"):
await _insert(store, _usage(name, inbound=True))
_complete_and_age(ledger_database)
async with client.writer_db.tx() as tx:
await tx.query_raw("SELECT event_id FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'a' FOR UPDATE")
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [
("a",),
("c",),
]
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
@pytest.mark.asyncio
async def test_retention_makes_progress_past_pending_and_unpriced_events(writer: _Writer, ledger_database: str) -> None:
_, store = writer
await _insert(store, _usage("a", inbound=True))
await _insert(store, _usage("b"))
await _insert(store, _usage("c", inbound=True))
_complete_and_age(ledger_database)
_rows(ledger_database, "UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'pending' WHERE event_id = 'a'")
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [
("a",),
("b",),
]
@pytest.mark.asyncio
async def test_zero_cost_and_empty_delivery_sets_are_durable(writer: _Writer, ledger_database: str) -> None:
_, store = writer
await _insert(store, _usage(), targets=())
event: Final = _cost(amount="0")
assert await _insert(store, event, targets=()) == "inserted"
assert await _insert(store, event, targets=()) == "duplicate"
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)]
_complete_and_age(ledger_database)
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
@pytest.mark.asyncio
async def test_retention_cannot_delete_source_while_cost_acceptance_is_uncommitted(
writer: _Writer, ledger_database: str
) -> None:
client, store = writer
await _insert(store, _usage())
_complete_and_age(ledger_database)
inserted: Final = asyncio.Event()
release: Final = asyncio.Event()
@asynccontextmanager
async def paused_commit() -> AsyncGenerator[Prisma]:
async with client.writer_db.tx() as tx:
yield tx
inserted.set()
await release.wait()
task: Final = asyncio.create_task(_insert(AccountingLedger(paused_commit), _cost()))
try:
await asyncio.wait_for(inserted.wait(), timeout=5)
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 0
finally:
release.set()
await task
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [
("cost",),
("usage",),
]
@pytest.mark.asyncio
async def test_retention_failure_rolls_back_event_and_receipt_deletion(writer: _Writer, ledger_database: str) -> None:
client, store = writer
await _insert(store, _usage(inbound=True))
_complete_and_age(ledger_database)
@asynccontextmanager
async def interrupted() -> AsyncGenerator[Prisma]:
async with client.writer_db.tx() as tx:
yield tx
raise RuntimeError("injected cleanup failure")
with pytest.raises(RuntimeError, match="cleanup failure"):
await AccountingLedger(interrupted).purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF)
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)]
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)]
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 1
@pytest.mark.asyncio
async def test_ledger_reads_and_writes_use_primary_when_replica_is_configured(
writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch
) -> None:
with scratch_database("replica") as replica_url:
monkeypatch.setenv("DATABASE_URL_READ_REPLICA", replica_url)
async with connected_proxy_prisma(ledger_database) as client:
store: Final = AccountingLedger.for_client(client)
await _insert(store, _usage())
assert await _insert(store, _usage()) == "duplicate"
assert await _insert(store, _cost()) == "inserted"
_complete_and_age(ledger_database)
assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2
assert _rows(replica_url, "SELECT to_regclass('public.\"LiteLLM_AccountingEvent\"')") == [(None,)]
def _index_conditions(plan: object) -> tuple[tuple[object, object], ...]:
if isinstance(plan, dict):
own: Final = ((plan.get("Index Name"), plan.get("Index Cond")),) if "Index Name" in plan else ()
children: Final = (plan.get("Plan"), *plan.get("Plans", ()))
return own + tuple(chain.from_iterable(_index_conditions(child) for child in children if child is not None))
return tuple(chain.from_iterable(_index_conditions(item) for item in plan)) if isinstance(plan, list) else ()
@pytest.mark.parametrize("statement", [accounting_ledger._DELETE_EVENTS, accounting_ledger._DELETE_DELIVERIES])
def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, statement: str) -> None:
explained: Final = "EXPLAIN (FORMAT JSON) " + statement.replace("$1", "%s")
with psycopg.connect(ledger_database) as connection:
connection.execute("SET enable_seqscan = off")
plan: Final = connection.execute(explained, ("usage",) * statement.count("$1")).fetchone()
assert plan is not None
conditions: Final = _index_conditions(plan[0])
assert any(
name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition)
for name, condition in conditions
), conditions

103
tests/unit/proxy/db/test_accounting_ledger.py Normal file
View file

@ -0,0 +1,103 @@
import json
from datetime import datetime, timezone
from decimal import Decimal
from typing import Final, Literal, NoReturn
import pytest
from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome
from litellm.proxy.db.accounting_ledger import AccountingLedger
from litellm.types.utils import CallTypes
def _no_transaction() -> NoReturn:
raise AssertionError("Invalid accounting input must not open a database transaction")
@pytest.fixture
def event() -> InboundUsageEvent:
return InboundUsageEvent(
event_id="caller",
request_id="caller",
call_type=CallTypes.completion,
started_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
outcome=Outcome.SUCCEEDED,
dimensions=(),
metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("envelope", ["{}", "[]", '{"event":null}', '{"event":{}}', "not json"])
async def test_invalid_envelopes_fail_before_database_access(event: InboundUsageEvent, envelope: str) -> None:
with pytest.raises(ValueError, match=r"Accounting envelope|validation error"):
await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope)
@pytest.mark.asyncio
@pytest.mark.parametrize("version", [0, -1, True, 1.5])
@pytest.mark.parametrize("field", ["version", "schema_version"])
async def test_invalid_versions_fail_before_database_access(
event: InboundUsageEvent, version: int | float, field: Literal["version", "schema_version"]
) -> None:
with pytest.raises(ValueError, match="positive integers"):
await AccountingLedger(_no_transaction).insert_event(
event,
(),
envelope="{}",
**{field: version}, # pyright: ignore[reportArgumentType] # Exercise invalid runtime versions
)
@pytest.mark.asyncio
async def test_constructed_invalid_events_are_revalidated(event: InboundUsageEvent) -> None:
invalid: Final = event.model_copy(update={"event_id": ""})
with pytest.raises(ValueError, match="event_id"):
await AccountingLedger(_no_transaction).insert_event(invalid, (), envelope="{}")
@pytest.mark.asyncio
@pytest.mark.parametrize("limit", [0, -1, 1001, True, 1.5])
async def test_retention_rejects_invalid_batch_bounds(limit: int | float) -> None:
with pytest.raises(ValueError, match="batch size"):
await AccountingLedger(_no_transaction).purge_completed(
limit=limit # pyright: ignore[reportArgumentType] # Exercise invalid runtime batch sizes
)
@pytest.mark.asyncio
@pytest.mark.parametrize("missing", ["both", "redelivery", "recovery"])
async def test_retention_requires_both_safety_watermarks(missing: Literal["both", "redelivery", "recovery"]) -> None:
cutoff: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
deleted: Final = await AccountingLedger(_no_transaction).purge_completed(
redelivery_before=cutoff if missing == "recovery" else None,
recovery_before=cutoff if missing == "redelivery" else None,
)
assert deleted == 0
@pytest.mark.asyncio
async def test_retention_rejects_ambiguous_naive_watermarks() -> None:
with pytest.raises(ValueError, match="timezone"):
await AccountingLedger(_no_transaction).purge_completed(
redelivery_before=datetime(2026, 1, 1), recovery_before=datetime(2026, 1, 1, tzinfo=timezone.utc)
)
@pytest.mark.asyncio
async def test_nonfinite_context_is_rejected_before_database_access(event: InboundUsageEvent) -> None:
envelope: Final = json.dumps({"event": event.model_dump(mode="json"), "context": {"price": float("nan")}})
with pytest.raises(ValueError, match="Out of range float"):
await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope)
@pytest.mark.asyncio
async def test_unknown_delivery_target_is_rejected_before_database_access(event: InboundUsageEvent) -> None:
envelope: Final = json.dumps({"event": event.model_dump(mode="json")})
with pytest.raises(ValueError, match="not a valid DeliveryTarget"):
await AccountingLedger(_no_transaction).insert_event(
event,
("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target
envelope=envelope,
)