diff --git a/litellm/proxy/spend_tracking/background_interaction_settlement.py b/litellm/proxy/spend_tracking/background_interaction_settlement.py index 096b3d46ab5..7585b2d2cac 100644 --- a/litellm/proxy/spend_tracking/background_interaction_settlement.py +++ b/litellm/proxy/spend_tracking/background_interaction_settlement.py @@ -1,7 +1,7 @@ import asyncio import os import socket -from collections.abc import Awaitable, Sequence +from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from typing import TYPE_CHECKING, Final, Protocol @@ -69,6 +69,7 @@ class _Claim(TypedDict): class _Outcome(TypedDict): settled_at: ReadOnly[datetime] outcome: ReadOnly[SettlementOutcome] + create_context: ReadOnly[object] class _SettlementTableActions(Protocol): @@ -85,6 +86,12 @@ def _settlement_table(prisma_client: "PrismaClient") -> _SettlementTableActions: return BackgroundInteractionSettlementRepository(prisma_client).table +def _json(data: Mapping[str, object]) -> object: + from prisma import Json # noqa: PLC0415 # local import: prisma may be ungenerated at module load in some tools + + return Json(data) + + def _pending_rows(rows: Sequence[_SettlementRow]) -> tuple[PendingBackgroundInteraction, ...]: return tuple(pending for row in rows for pending in _pending_row(row)) @@ -114,13 +121,11 @@ class PrismaBackgroundSettlementStore: claimed_by: str async def register(self, pending: PendingBackgroundInteraction) -> None: - from prisma import Json # noqa: PLC0415 # local import: prisma may be ungenerated at module load in some tools - await self.table.create( data=_NewSettlementRow( interaction_id=pending.interaction_id, custom_llm_provider=pending.custom_llm_provider, - create_context=Json(pending.create_context.model_dump(mode="json")), + create_context=_json(pending.create_context.model_dump(mode="json")), created_at=pending.created_at, ) ) @@ -144,7 +149,7 @@ class PrismaBackgroundSettlementStore: async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: await self.table.update_many( - data=_Outcome(settled_at=datetime.now(timezone.utc), outcome=outcome), + data=_Outcome(settled_at=datetime.now(timezone.utc), outcome=outcome, create_context=_json({})), where=_RowKey(interaction_id=interaction_id), ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_background_interaction_settlement.py b/tests/test_litellm/proxy/spend_tracking/test_background_interaction_settlement.py index fe59d20dec8..0d026f260d3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_background_interaction_settlement.py +++ b/tests/test_litellm/proxy/spend_tracking/test_background_interaction_settlement.py @@ -74,7 +74,7 @@ class _FakeSettlementTable: matched = self._matching(where) for row in matched: for column, value in data.items(): - setattr(row, column, value) + setattr(row, column, getattr(value, "data", value) if column == "create_context" else value) return len(matched) def _matching(self, where) -> list: @@ -184,11 +184,12 @@ async def test_unclaimed_skips_claimed_and_unreadable_rows(): @pytest.mark.asyncio -async def test_record_outcome_keeps_the_audit_trail_on_the_row(): +async def test_record_outcome_keeps_the_audit_trail_and_drops_the_stored_request_context(): table = _FakeSettlementTable() store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") await store.register(_pending("interactions/bg-1")) assert await store.claim("interactions/bg-1") + assert table.rows["interactions/bg-1"].create_context await store.record_outcome("interactions/bg-1", "billed") @@ -196,6 +197,7 @@ async def test_record_outcome_keeps_the_audit_trail_on_the_row(): assert row.outcome == "billed" assert row.settled_at is not None assert row.claimed_at <= row.settled_at + assert row.create_context == {} @pytest.mark.asyncio