diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py index 0a9c5db1619..37417244457 100644 --- a/litellm/proxy/db/accounting_ledger.py +++ b/litellm/proxy/db/accounting_ledger.py @@ -11,7 +11,17 @@ 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.accounting import ( + CostEvent, + DeliveryTarget, + EventKind, + InboundUsageEvent, + Metric, + MetricName, + OutboundUsageEvent, + PricingSnapshot, + Unpriced, +) from litellm.proxy.db.db_span import db_span from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper @@ -143,6 +153,14 @@ SELECT records.event_id, destination.value::"AccountingDeliveryTarget", 1 FROM jsonb_to_recordset($1::jsonb) AS records(event_id text, targets jsonb) CROSS JOIN LATERAL jsonb_array_elements_text(records.targets) AS destination """ +_LOAD_SOURCE: Final = """ +SELECT payload::text AS payload, schema_version::text AS schema_version +FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE +""" +_LOAD_COST: Final = """ +SELECT payload::text AS payload, schema_version::text AS schema_version, payload_hash +FROM "LiteLLM_AccountingEvent" WHERE source_event_id = $1 AND kind = 'cost' +""" _COMPLETED_FAMILY: Final = """ root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage') AND root.recorded_at < $1::timestamp @@ -320,6 +338,56 @@ class AccountingLedger: case _: assert_never(result) + async def compile_cost( + self, + event: OutboundUsageEvent, + price: Callable[[str], PricingSnapshot | Unpriced], + targets: tuple[DeliveryTarget, ...], + ) -> CostEvent | Unpriced: + """Serialize pure pricing with other compilers and commit its complete delivery set.""" + validated: Final = OutboundUsageEvent.model_validate_json(event.model_dump_json()) + async with self._transaction() as db: + sources: Final = _STRINGS.validate_python(await db.query_raw(_LOAD_SOURCE, validated.event_id)) + if not sources: + return Unpriced(usage_event_id=validated.event_id, reason="usage_not_committed") + source: Final = sources[0] + if source["schema_version"] != "1": + return Unpriced(usage_event_id=validated.event_id, reason="unsupported_usage_schema") + payload: Final = _JSON.validate_json(source["payload"]) + if payload.get("event") != validated.model_dump(mode="json"): + raise AccountingConflictError("Compilation input differs from the accepted usage event") + costs: Final = _STRINGS.validate_python(await db.query_raw(_LOAD_COST, validated.event_id)) + if costs: + if costs[0]["schema_version"] != "1": + return Unpriced(usage_event_id=validated.event_id, reason="unsupported_cost_schema") + existing: Final = CostEvent.model_validate_json( + _json(_JSON.validate_json(costs[0]["payload"])["event"]) + ) + if existing.usage_event_id != validated.event_id: + raise AccountingConflictError("Committed cost does not match its source usage") + replay: Final = _prepare(existing, targets, costs[0]["payload"], 1, 1) + if replay.payload_hash != costs[0]["payload_hash"]: + raise AccountingConflictError("Compilation delivery set differs from the accepted cost") + return existing + priced: Final = price(source["payload"]) + if isinstance(priced, Unpriced): + if priced.usage_event_id != validated.event_id: + raise AccountingConflictError("Unpriced result does not match its source usage") + return priced + cost: Final = CostEvent( + event_id=f"cost:{validated.event_id}", + usage_event_id=validated.event_id, + occurred_at=validated.ended_at, + dimensions=validated.dimensions, + charges=(Metric(name=MetricName.SPEND, value=priced.total_cost, unit="USD"),), + snapshot=priced, + ) + record: Final = _prepare(cost, targets, _json({"event": cost.model_dump(mode="json")}), 1, 1) + accepted: Final = await _accept(db, record) + if accepted not in ("inserted", "duplicate"): + raise AccountingConflictError(f"Cannot commit compiled cost: {accepted}") + return cost + async def purge_completed( self, *, diff --git a/litellm/proxy/spend_tracking/cost_compiler.py b/litellm/proxy/spend_tracking/cost_compiler.py new file mode 100644 index 00000000000..9d2438b93dc --- /dev/null +++ b/litellm/proxy/spend_tracking/cost_compiler.py @@ -0,0 +1,20 @@ +from collections.abc import Callable +from typing import Final + +from litellm.accounting import CostEvent, DeliveryTarget, OutboundUsageEvent, PricingSnapshot, Unpriced +from litellm.proxy.db.accounting_ledger import AccountingLedger + + +class LoggingCostCompiler: + def __init__( + self, + ledger: AccountingLedger, + price_usage: Callable[[str], PricingSnapshot | Unpriced], + delivery_targets: tuple[DeliveryTarget, ...], + ) -> None: + self._ledger: Final = ledger + self._price_usage: Final = price_usage + self._delivery_targets: Final = delivery_targets + + async def compile(self, event: OutboundUsageEvent) -> CostEvent | Unpriced: + return await self._ledger.compile_cost(event, self._price_usage, self._delivery_targets) diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py index 15cd5413149..e107e8e130c 100644 --- a/tests/integration/database/test_accounting_ledger.py +++ b/tests/integration/database/test_accounting_ledger.py @@ -20,6 +20,7 @@ from psycopg import sql from pydantic import JsonValue from litellm.accounting import ( + CostCompiler, CostEvent, DeliveryState, DeliveryTarget, @@ -30,9 +31,11 @@ from litellm.accounting import ( OutboundUsageEvent, Outcome, PricingSnapshot, + Unpriced, ) from litellm.proxy.db import accounting_ledger from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger, UsageWrite +from litellm.proxy.spend_tracking.cost_compiler import LoggingCostCompiler from litellm.types.utils import CallTypes from tests.integration._support.database import scratch_database from tests.integration._support.prisma_clients import connected_proxy_prisma @@ -248,6 +251,181 @@ async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS} +@pytest.mark.asyncio +async def test_compilers_share_one_committed_charge_and_preserve_delivery_receipts( + writer: _Writer, ledger_database: str +) -> None: + _, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + await _insert(store, event, context={"price": "0.125"}) + + def price(envelope: str) -> PricingSnapshot: + return _cost(amount=json.loads(envelope)["context"]["price"]).snapshot + + compiler: Final[CostCompiler] = LoggingCostCompiler(store, price, _TARGETS) + results: Final = await asyncio.gather(*(compiler.compile(event) for _ in range(8))) + assert all(result == results[0] for result in results) + assert isinstance(results[0], CostEvent) + assert results[0].snapshot.total_cost == Decimal("0.125") + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(1,)] + deliveries: Final = _rows( + ledger_database, + 'SELECT target::text, state::text FROM "LiteLLM_AccountingDelivery" WHERE event_id = %s', + (results[0].event_id,), + ) + assert set(deliveries) == {(target.value, "pending") for target in _TARGETS} + _complete_and_age(ledger_database) + before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') + + def unavailable_pricing(envelope: str) -> PricingSnapshot: + raise AssertionError("Committed cost replay must not invoke pricing") + + replay: Final = await LoggingCostCompiler(store, unavailable_pricing, _TARGETS).compile(event) + assert replay == results[0] + assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') == before + + +@pytest.mark.asyncio +async def test_compilation_rejects_changed_delivery_targets_without_resetting_receipts( + writer: _Writer, ledger_database: str +) -> None: + _, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + await _insert(store, event) + original: Final = await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS[:1]).compile(event) + _complete_and_age(ledger_database) + before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') + with pytest.raises(AccountingConflictError, match="delivery set"): + await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS).compile(event) + assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') == before + assert await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS[:1]).compile(event) == original + + +@pytest.mark.asyncio +async def test_compilation_preserves_unknown_cost_without_creating_a_zero_charge( + writer: _Writer, ledger_database: str +) -> None: + _, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + unresolved: Final = Unpriced(usage_event_id=event.event_id, reason="unknown_usage") + compiler: Final = LoggingCostCompiler(store, lambda _: unresolved, _TARGETS) + assert await compiler.compile(event) == Unpriced(usage_event_id=event.event_id, reason="usage_not_committed") + await _insert(store, event) + assert await compiler.compile(event) == unresolved + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(0,)] + + +@pytest.mark.asyncio +async def test_compilation_rejects_an_unpriced_result_for_another_source(writer: _Writer, ledger_database: str) -> None: + _, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + await _insert(store, event) + compiler: Final = LoggingCostCompiler( + store, lambda _: Unpriced(usage_event_id="other", reason="missing_price"), _TARGETS + ) + with pytest.raises(AccountingConflictError, match="Unpriced result"): + await compiler.compile(event) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(0,)] + + +@pytest.mark.asyncio +async def test_compilation_rejects_a_cost_identity_collision(writer: _Writer, ledger_database: str) -> None: + _, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + await _insert(store, event) + await _insert(store, _usage(f"cost:{event.event_id}", inbound=True)) + with pytest.raises(AccountingConflictError, match="identity_conflict"): + await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS).compile(event) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(0,)] + + +@pytest.mark.asyncio +async def test_compilation_rejects_changed_source_before_reusing_a_charge(writer: _Writer) -> None: + _, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + await _insert(store, event) + compiler: Final = LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS) + assert isinstance(await compiler.compile(event), CostEvent) + changed: Final = event.model_copy(update={"outcome": Outcome.CANCELLED}) + with pytest.raises(AccountingConflictError, match="differs from the accepted usage"): + await compiler.compile(changed) + + +@pytest.mark.asyncio +async def test_compilation_failure_rolls_back_and_can_be_retried(writer: _Writer, ledger_database: str) -> None: + _, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + await _insert(store, event) + + def failed_pricing(envelope: str) -> PricingSnapshot: + raise ValueError("malformed retained pricing") + + with pytest.raises(ValueError, match="malformed retained pricing"): + await LoggingCostCompiler(store, failed_pricing, _TARGETS).compile(event) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(0,)] + result: Final = await LoggingCostCompiler(store, lambda _: _cost(amount="0").snapshot, _TARGETS).compile(event) + assert isinstance(result, CostEvent) + assert result.charges[0].value == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["usage", "cost"]) +async def test_compilation_does_not_interpret_unknown_envelope_versions(writer: _Writer, kind: str) -> None: + _, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + await store.insert_event(event, _TARGETS, envelope=_envelope(event), schema_version=2 if kind == "usage" else 1) + if kind == "cost": + cost: Final = _cost() + await store.insert_event(cost, _TARGETS, envelope=_envelope(cost), schema_version=2) + result: Final = await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS).compile(event) + assert result == Unpriced(usage_event_id=event.event_id, reason=f"unsupported_{kind}_schema") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"]) +async def test_compiler_recovery_preserves_atomic_charge_and_delivery_acceptance( + writer: _Writer, ledger_database: str, failure: str +) -> None: + client, store = writer + event: Final = _usage() + assert isinstance(event, OutboundUsageEvent) + await _insert(store, event) + + @asynccontextmanager + async def interrupted() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + if failure == "before_commit": + raise RuntimeError("interrupted cost commit") + if failure == "cancellation": + raise asyncio.CancelledError + raise RuntimeError("lost cost acknowledgement") + + compiler: Final = LoggingCostCompiler(AccountingLedger(interrupted), lambda _: _cost().snapshot, _TARGETS) + error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError + with pytest.raises(error): + await compiler.compile(event) + count: Final = int(failure == "after_commit") + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(count,)] + assert _rows( + ledger_database, + 'SELECT count(*) FROM "LiteLLM_AccountingDelivery" WHERE event_id = %s', + (f"cost:{event.event_id}",), + ) == [(count * len(_TARGETS),)] + result: Final = await LoggingCostCompiler(store, lambda _: _cost().snapshot, _TARGETS).compile(event) + assert isinstance(result, CostEvent) + assert result.snapshot == _cost().snapshot + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent" WHERE kind = \'cost\'') == [(1,)] + + @pytest.mark.asyncio async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None: _, store = writer