mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(interactions): durable cross-pod settlement for background interaction billing
Background interaction billing previously lived only in the creating pod's memory: a DELETE routed to another pod or a proxy restart meant the completed provider work was never billed and the budget reservation was refunded at the poll timeout. The poll task now persists a settlement context row (virtual key attribution, model, original request id, budget reservation) in a new LiteLLM_BackgroundInteractionSettlementTable so any pod can settle. A conditional update flipping status pending to settled is the cross-pod exactly-once claim; every outcome (bill, release, timeout give-up) must win it before touching spend counters, and a losing pod finalizes its local reservation copy without touching shared counters. The delete route settles registry misses through the store by rebuilding a billable logging context, resolving provider credentials through the router by model group. A per-pod sweep job resumes pending settlements after restarts and surfaces timed-out entries as abandoned rows instead of silently reconciling to zero. Without a registered store (SDK usage) behavior is unchanged.
This commit is contained in:
mateo-berri
parent
e458aa1230
commit
8b63b7045c
10 changed files with 1206 additions and 14 deletions
|
|
@ -0,0 +1,16 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_BackgroundInteractionSettlementTable" (
|
||||
"interaction_id" TEXT NOT NULL,
|
||||
"status" TEXT NOT NULL DEFAULT 'pending',
|
||||
"outcome" TEXT,
|
||||
"context" JSONB NOT NULL,
|
||||
"claimed_by" TEXT,
|
||||
"timeout_at" TIMESTAMP(3) NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL,
|
||||
|
||||
CONSTRAINT "LiteLLM_BackgroundInteractionSettlementTable_pkey" PRIMARY KEY ("interaction_id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_BackgroundInteractionSettlement_status_created_idx" ON "LiteLLM_BackgroundInteractionSettlementTable"("status", "created_at");
|
||||
|
|
@ -950,6 +950,19 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
|
|||
@@index([team_id, created_at(sort: Desc)])
|
||||
}
|
||||
|
||||
model LiteLLM_BackgroundInteractionSettlementTable {
|
||||
interaction_id String @id
|
||||
status String @default("pending")
|
||||
outcome String?
|
||||
context Json
|
||||
claimed_by String?
|
||||
timeout_at DateTime
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@index([status, created_at], map: "LiteLLM_BackgroundInteractionSettlement_status_created_idx")
|
||||
}
|
||||
|
||||
model LiteLLM_ManagedVectorStoreTable {
|
||||
id String @id @default(uuid())
|
||||
unified_resource_id String @unique // The base64 encoded unified vector store ID
|
||||
|
|
|
|||
|
|
@ -1468,6 +1468,9 @@ BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS = float(
|
|||
)
|
||||
_background_interaction_cost_polling_env = os.getenv("BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true").lower()
|
||||
BACKGROUND_INTERACTION_COST_POLLING_ENABLED = _background_interaction_cost_polling_env == "true"
|
||||
BACKGROUND_SETTLEMENT_SWEEP_INTERVAL_SECONDS = int(os.getenv("BACKGROUND_SETTLEMENT_SWEEP_INTERVAL_SECONDS", 300))
|
||||
BACKGROUND_SETTLEMENT_SWEEP_MIN_AGE_SECONDS = float(os.getenv("BACKGROUND_SETTLEMENT_SWEEP_MIN_AGE_SECONDS", 60))
|
||||
MAX_BACKGROUND_SETTLEMENTS_PER_SWEEP = int(os.getenv("MAX_BACKGROUND_SETTLEMENTS_PER_SWEEP", 100))
|
||||
PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605))
|
||||
PROXY_BATCH_WRITE_AT = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10
|
||||
|
||||
|
|
|
|||
|
|
@ -19,11 +19,23 @@ state with the create's credentials, bills it if it is terminal with usage,
|
|||
and releases the reservation otherwise. A settlement gate on the create's
|
||||
logging object makes the poll task and the delete path mutually exclusive, so
|
||||
the interaction is billed exactly once no matter who settles first.
|
||||
|
||||
The poll registry and the create's logging object only live in the pod that
|
||||
served the create, so a delete routed to another pod or a proxy restart would
|
||||
leave the completed work unbilled. When a durable settlement store is
|
||||
registered (the proxy registers a Prisma-backed one at startup), the poll
|
||||
task persists enough create context for any pod to settle, and every outcome
|
||||
must win the store's cross-pod compare-and-swap claim before touching spend
|
||||
counters; the in-memory gate remains the same-pod fast path and the only gate
|
||||
when no store is registered (SDK usage).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterator, Optional
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterator, Literal, Optional, Protocol
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
|
|
@ -40,6 +52,70 @@ if TYPE_CHECKING:
|
|||
|
||||
_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete", "budget_exceeded"})
|
||||
|
||||
SettlementOutcome = Literal["billed", "released", "abandoned", "error"]
|
||||
|
||||
|
||||
class SettlementReservationEntry(BaseModel):
|
||||
counter_key: str
|
||||
entity_type: str
|
||||
entity_id: str
|
||||
reserved_cost: float
|
||||
applied_adjustment: float = 0.0
|
||||
|
||||
|
||||
class SettlementReservation(BaseModel):
|
||||
reserved_cost: float
|
||||
entries: list[SettlementReservationEntry]
|
||||
finalized: bool
|
||||
input_cost: float
|
||||
|
||||
|
||||
class SettlementKeyAttribution(BaseModel):
|
||||
user_api_key: Optional[str] = None
|
||||
user_api_key_user_id: Optional[str] = None
|
||||
user_api_key_team_id: Optional[str] = None
|
||||
user_api_key_org_id: Optional[str] = None
|
||||
user_api_key_alias: Optional[str] = None
|
||||
user_api_key_team_alias: Optional[str] = None
|
||||
user_api_key_user_email: Optional[str] = None
|
||||
user_api_key_end_user_id: Optional[str] = None
|
||||
|
||||
|
||||
class BackgroundSettlementContext(BaseModel):
|
||||
interaction_id: str
|
||||
custom_llm_provider: str
|
||||
model: str
|
||||
model_group: Optional[str] = None
|
||||
litellm_call_id: str
|
||||
litellm_trace_id: Optional[str] = None
|
||||
call_type: str
|
||||
attribution: SettlementKeyAttribution
|
||||
budget_reservation: Optional[SettlementReservation] = None
|
||||
|
||||
|
||||
class BackgroundSettlementStore(Protocol):
|
||||
async def persist_pending(self, context: BackgroundSettlementContext, timeout_at: datetime) -> None: ...
|
||||
|
||||
async def claim(self, interaction_id: str) -> bool: ...
|
||||
|
||||
async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: ...
|
||||
|
||||
async def is_pending(self, interaction_id: str) -> bool: ...
|
||||
|
||||
async def settle_pending_before_delete(self, interaction_id: str) -> None: ...
|
||||
|
||||
|
||||
_SETTLEMENT_STORE: Optional[BackgroundSettlementStore] = None
|
||||
|
||||
|
||||
def get_settlement_store() -> Optional[BackgroundSettlementStore]:
|
||||
return _SETTLEMENT_STORE
|
||||
|
||||
|
||||
def set_settlement_store(store: Optional[BackgroundSettlementStore]) -> None:
|
||||
global _SETTLEMENT_STORE
|
||||
_SETTLEMENT_STORE = store
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BackgroundInteractionPollContext:
|
||||
|
|
@ -94,10 +170,148 @@ def _claim_settlement(logging_obj: "LiteLLMLoggingObj") -> bool:
|
|||
return True
|
||||
|
||||
|
||||
def _get_reservation_dict(logging_obj: "LiteLLMLoggingObj") -> Optional[dict]:
|
||||
metadata = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)
|
||||
budget_reservation = metadata.get("user_api_key_budget_reservation")
|
||||
return budget_reservation if isinstance(budget_reservation, dict) else None
|
||||
|
||||
|
||||
def _finalize_reservation_locally(logging_obj: "LiteLLMLoggingObj") -> None:
|
||||
"""
|
||||
Losing the durable store's claim means another pod owns (or already
|
||||
performed) the billing or release for this interaction. The local copy of
|
||||
the reservation dict must be marked finalized without touching the shared
|
||||
counters, or a later reconcile of this copy would re-apply its full
|
||||
adjustment on top of the winner's.
|
||||
"""
|
||||
budget_reservation = _get_reservation_dict(logging_obj)
|
||||
if budget_reservation is not None:
|
||||
budget_reservation["finalized"] = True
|
||||
|
||||
|
||||
def build_settlement_context(
|
||||
context: BackgroundInteractionPollContext,
|
||||
) -> BackgroundSettlementContext:
|
||||
metadata = get_litellm_metadata_from_kwargs(kwargs=context.logging_obj.model_call_details)
|
||||
reservation_dict = metadata.get("user_api_key_budget_reservation")
|
||||
model_group = metadata.get("deployment_model_name")
|
||||
return BackgroundSettlementContext(
|
||||
interaction_id=context.interaction_id,
|
||||
custom_llm_provider=context.custom_llm_provider,
|
||||
model=str(context.logging_obj.model_call_details.get("model") or context.logging_obj.model),
|
||||
model_group=model_group if isinstance(model_group, str) else None,
|
||||
litellm_call_id=str(context.logging_obj.litellm_call_id),
|
||||
litellm_trace_id=context.logging_obj.model_call_details.get("litellm_trace_id"),
|
||||
call_type=str(context.logging_obj.call_type),
|
||||
attribution=SettlementKeyAttribution.model_validate(metadata),
|
||||
budget_reservation=(
|
||||
SettlementReservation.model_validate(reservation_dict) if isinstance(reservation_dict, dict) else None
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _persist_pending_settlement(
|
||||
context: BackgroundInteractionPollContext,
|
||||
store: BackgroundSettlementStore,
|
||||
) -> bool:
|
||||
try:
|
||||
settlement_context = build_settlement_context(context)
|
||||
timeout_at = datetime.now(timezone.utc) + timedelta(seconds=context.timeout_seconds)
|
||||
await store.persist_pending(context=settlement_context, timeout_at=timeout_at)
|
||||
return True
|
||||
except Exception: # noqa: BLE001 # a failed persist degrades to in-memory-only settlement, never kills the poll
|
||||
verbose_logger.warning(
|
||||
"Failed to persist pending settlement for background interaction %s; settlement will be in-memory only",
|
||||
context.interaction_id,
|
||||
exc_info=True,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
async def _claim_in_store(store: BackgroundSettlementStore, interaction_id: str) -> bool:
|
||||
try:
|
||||
return await store.claim(interaction_id)
|
||||
except Exception: # noqa: BLE001 # an unreachable store defers to the sweep instead of risking a double bill
|
||||
verbose_logger.exception(
|
||||
"Failed to claim settlement for background interaction %s; deferring to the settlement sweep",
|
||||
interaction_id,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
async def _record_outcome_best_effort(
|
||||
store: BackgroundSettlementStore,
|
||||
interaction_id: str,
|
||||
outcome: SettlementOutcome,
|
||||
) -> None:
|
||||
try:
|
||||
await store.record_outcome(interaction_id, outcome)
|
||||
except Exception: # noqa: BLE001 # the outcome column is observability, never worth failing settlement over
|
||||
verbose_logger.warning(
|
||||
"Failed to record settlement outcome %s for background interaction %s",
|
||||
outcome,
|
||||
interaction_id,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
|
||||
async def _claim_across_gates(
|
||||
context: BackgroundInteractionPollContext,
|
||||
store: Optional[BackgroundSettlementStore],
|
||||
) -> bool:
|
||||
if not _claim_settlement(context.logging_obj):
|
||||
return False
|
||||
if store is not None and not await _claim_in_store(store, context.interaction_id):
|
||||
_finalize_reservation_locally(context.logging_obj)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def _settle_claimed(
|
||||
context: BackgroundInteractionPollContext,
|
||||
response: InteractionsAPIResponse,
|
||||
store: Optional[BackgroundSettlementStore],
|
||||
) -> None:
|
||||
outcome: SettlementOutcome
|
||||
try:
|
||||
if response.usage is not None:
|
||||
await context.logging_obj.async_log_background_interaction_completion(result=response)
|
||||
outcome = "billed"
|
||||
else:
|
||||
await _release_open_budget_reservation(logging_obj=context.logging_obj)
|
||||
outcome = "released"
|
||||
except Exception: # noqa: BLE001 # a billing failure after winning the claim must be surfaced, not retried
|
||||
verbose_logger.exception(
|
||||
"Billing failed after claiming settlement for background interaction %s; its spend may be under-tracked",
|
||||
context.interaction_id,
|
||||
)
|
||||
await _release_open_budget_reservation(logging_obj=context.logging_obj)
|
||||
outcome = "error"
|
||||
if store is not None:
|
||||
await _record_outcome_best_effort(store, context.interaction_id, outcome)
|
||||
|
||||
|
||||
async def _no_longer_pending_in_store(
|
||||
store: BackgroundSettlementStore,
|
||||
interaction_id: str,
|
||||
) -> bool:
|
||||
try:
|
||||
return not await store.is_pending(interaction_id)
|
||||
except Exception: # noqa: BLE001 # an unreachable store must not change the poll loop's behavior
|
||||
return False
|
||||
|
||||
|
||||
async def poll_and_log_background_interaction_cost(
|
||||
context: BackgroundInteractionPollContext,
|
||||
fetch_interaction: FetchInteraction = _fetch_interaction,
|
||||
store: Optional[BackgroundSettlementStore] = None,
|
||||
) -> None:
|
||||
configured_store = store if store is not None else get_settlement_store()
|
||||
active_store = (
|
||||
configured_store
|
||||
if configured_store is not None and await _persist_pending_settlement(context, configured_store)
|
||||
else None
|
||||
)
|
||||
for interval in _poll_intervals(
|
||||
initial=context.initial_interval_seconds,
|
||||
maximum=context.max_interval_seconds,
|
||||
|
|
@ -109,6 +323,10 @@ async def poll_and_log_background_interaction_cost(
|
|||
try:
|
||||
response = await fetch_interaction(context)
|
||||
except Exception as e: # noqa: BLE001 # any fetch error must not kill the billing poll loop
|
||||
if active_store is not None and await _no_longer_pending_in_store(active_store, context.interaction_id):
|
||||
if _claim_settlement(context.logging_obj):
|
||||
_finalize_reservation_locally(context.logging_obj)
|
||||
return
|
||||
verbose_logger.debug(
|
||||
"Background interaction cost poll for %s failed, will retry: %s",
|
||||
context.interaction_id,
|
||||
|
|
@ -117,14 +335,11 @@ async def poll_and_log_background_interaction_cost(
|
|||
continue
|
||||
if response.status not in _TERMINAL_STATUSES:
|
||||
continue
|
||||
if not _claim_settlement(context.logging_obj):
|
||||
if not await _claim_across_gates(context, active_store):
|
||||
return
|
||||
if response.usage is not None:
|
||||
await context.logging_obj.async_log_background_interaction_completion(result=response)
|
||||
else:
|
||||
await _release_open_budget_reservation(logging_obj=context.logging_obj)
|
||||
await _settle_claimed(context=context, response=response, store=active_store)
|
||||
return
|
||||
if not _claim_settlement(context.logging_obj):
|
||||
if not await _claim_across_gates(context, active_store):
|
||||
return
|
||||
verbose_logger.warning(
|
||||
"Gave up cost polling for background interaction %s after %ss; its usage will not be tracked",
|
||||
|
|
@ -132,6 +347,8 @@ async def poll_and_log_background_interaction_cost(
|
|||
context.timeout_seconds,
|
||||
)
|
||||
await _release_open_budget_reservation(logging_obj=context.logging_obj)
|
||||
if active_store is not None:
|
||||
await _record_outcome_best_effort(active_store, context.interaction_id, "abandoned")
|
||||
|
||||
|
||||
async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> None:
|
||||
|
|
@ -144,9 +361,8 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") ->
|
|||
settlement must release the reservation here or the spend counters stay
|
||||
pinned at the estimated cost.
|
||||
"""
|
||||
metadata = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)
|
||||
budget_reservation = metadata.get("user_api_key_budget_reservation")
|
||||
if not isinstance(budget_reservation, dict):
|
||||
budget_reservation = _get_reservation_dict(logging_obj)
|
||||
if budget_reservation is None:
|
||||
return
|
||||
|
||||
from litellm.proxy.spend_tracking.budget_reservation import release_budget_reservation
|
||||
|
|
@ -212,9 +428,13 @@ def maybe_schedule_background_interaction_cost_polling(
|
|||
async def maybe_settle_background_interaction_before_delete(
|
||||
interaction_id: str,
|
||||
fetch_interaction: FetchInteraction = _fetch_interaction,
|
||||
store: Optional[BackgroundSettlementStore] = None,
|
||||
) -> None:
|
||||
active_store = store if store is not None else get_settlement_store()
|
||||
entry = _ACTIVE_POLLS.get(interaction_id)
|
||||
if entry is None:
|
||||
if active_store is not None:
|
||||
await active_store.settle_pending_before_delete(interaction_id)
|
||||
return
|
||||
context = entry.context
|
||||
try:
|
||||
|
|
@ -225,12 +445,17 @@ async def maybe_settle_background_interaction_before_delete(
|
|||
interaction_id,
|
||||
e,
|
||||
)
|
||||
if _claim_settlement(context.logging_obj):
|
||||
await _release_open_budget_reservation(logging_obj=context.logging_obj)
|
||||
if not await _claim_across_gates(context, active_store):
|
||||
return
|
||||
await _release_open_budget_reservation(logging_obj=context.logging_obj)
|
||||
if active_store is not None:
|
||||
await _record_outcome_best_effort(active_store, interaction_id, "released")
|
||||
return
|
||||
if not _claim_settlement(context.logging_obj):
|
||||
if not await _claim_across_gates(context, active_store):
|
||||
return
|
||||
if response.status in _TERMINAL_STATUSES and response.usage is not None:
|
||||
await context.logging_obj.async_log_background_interaction_completion(result=response)
|
||||
await _settle_claimed(context=context, response=response, store=active_store)
|
||||
return
|
||||
await _release_open_budget_reservation(logging_obj=context.logging_obj)
|
||||
if active_store is not None:
|
||||
await _record_outcome_best_effort(active_store, interaction_id, "released")
|
||||
|
|
|
|||
|
|
@ -224,6 +224,8 @@ from litellm.constants import (
|
|||
APSCHEDULER_MAX_INSTANCES,
|
||||
APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
APSCHEDULER_REPLACE_EXISTING,
|
||||
BACKGROUND_INTERACTION_COST_POLLING_ENABLED,
|
||||
BACKGROUND_SETTLEMENT_SWEEP_INTERVAL_SECONDS,
|
||||
DAYS_IN_A_MONTH,
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
|
|
@ -7902,6 +7904,27 @@ class ProxyStartupEvent:
|
|||
)
|
||||
pass
|
||||
|
||||
### BACKGROUND INTERACTION SETTLEMENT ###
|
||||
if BACKGROUND_INTERACTION_COST_POLLING_ENABLED:
|
||||
from litellm.interactions.background_cost_polling import set_settlement_store
|
||||
from litellm.proxy.spend_tracking.background_settlement import (
|
||||
PrismaBackgroundSettlementStore,
|
||||
sweep_pending_settlements,
|
||||
)
|
||||
|
||||
settlement_store = PrismaBackgroundSettlementStore(prisma_client=prisma_client)
|
||||
set_settlement_store(settlement_store)
|
||||
scheduler.add_job(
|
||||
sweep_pending_settlements,
|
||||
"interval",
|
||||
seconds=BACKGROUND_SETTLEMENT_SWEEP_INTERVAL_SECONDS + random.randint(0, 30),
|
||||
kwargs={"store": settlement_store},
|
||||
id="background_interaction_settlement_sweep_job",
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
verbose_proxy_logger.info("Background interaction settlement sweep job scheduled successfully")
|
||||
|
||||
### CHECK RESPONSES COST ###
|
||||
if llm_router is not None and PROXY_BATCH_POLLING_ENABLED:
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -950,6 +950,19 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
|
|||
@@index([team_id, created_at(sort: Desc)])
|
||||
}
|
||||
|
||||
model LiteLLM_BackgroundInteractionSettlementTable {
|
||||
interaction_id String @id
|
||||
status String @default("pending")
|
||||
outcome String?
|
||||
context Json
|
||||
claimed_by String?
|
||||
timeout_at DateTime
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@index([status, created_at], map: "LiteLLM_BackgroundInteractionSettlement_status_created_idx")
|
||||
}
|
||||
|
||||
model LiteLLM_ManagedVectorStoreTable {
|
||||
id String @id @default(uuid())
|
||||
unified_resource_id String @unique // The base64 encoded unified vector store ID
|
||||
|
|
|
|||
337
litellm/proxy/spend_tracking/background_settlement.py
Normal file
337
litellm/proxy/spend_tracking/background_settlement.py
Normal file
|
|
@ -0,0 +1,337 @@
|
|||
"""
|
||||
Prisma-backed durable settlement for background interaction billing.
|
||||
|
||||
The creating pod's poll task persists a settlement context row here so that a
|
||||
delete routed to any pod, or a sweep after a pod restart, can rebuild a
|
||||
billable logging object and settle the interaction exactly once. The row's
|
||||
``status`` column is the cross-pod claim: whoever flips ``pending`` to
|
||||
``settled`` via a conditional update owns billing or release, and the
|
||||
``outcome`` column records what the winner did (``billed``, ``released``,
|
||||
``abandoned``, or ``error``) for audit of anything that could not be billed.
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Awaitable, Callable, Optional, Protocol
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
BACKGROUND_SETTLEMENT_SWEEP_MIN_AGE_SECONDS,
|
||||
MAX_BACKGROUND_SETTLEMENTS_PER_SWEEP,
|
||||
)
|
||||
from litellm.interactions.background_cost_polling import (
|
||||
_TERMINAL_STATUSES,
|
||||
BackgroundSettlementContext,
|
||||
SettlementOutcome,
|
||||
)
|
||||
from litellm.types.interactions import InteractionsAPIResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
class PendingSettlementRow(BaseModel):
|
||||
interaction_id: str
|
||||
context: BackgroundSettlementContext
|
||||
timeout_at: datetime
|
||||
|
||||
|
||||
class SettlementRowStore(Protocol):
|
||||
async def claim(self, interaction_id: str) -> bool: ...
|
||||
|
||||
async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: ...
|
||||
|
||||
async def list_due(self, older_than: datetime, limit: int) -> tuple[PendingSettlementRow, ...]: ...
|
||||
|
||||
|
||||
SettlementFetch = Callable[[BackgroundSettlementContext], Awaitable[InteractionsAPIResponse]]
|
||||
|
||||
|
||||
class DeploymentResolver(Protocol):
|
||||
async def async_get_available_deployment(self, model: str, request_kwargs: dict) -> dict: ...
|
||||
|
||||
|
||||
def _get_proxy_router() -> Optional[DeploymentResolver]:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
return llm_router
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SettlementCredentials:
|
||||
api_key: Optional[str]
|
||||
api_base: Optional[str]
|
||||
|
||||
|
||||
async def _resolve_settlement_credentials(
|
||||
context: BackgroundSettlementContext,
|
||||
router: Optional[DeploymentResolver] = None,
|
||||
) -> _SettlementCredentials:
|
||||
if context.model_group is None:
|
||||
return _SettlementCredentials(api_key=None, api_base=None)
|
||||
active_router = router if router is not None else _get_proxy_router()
|
||||
if active_router is None:
|
||||
return _SettlementCredentials(api_key=None, api_base=None)
|
||||
try:
|
||||
deployment = await active_router.async_get_available_deployment(
|
||||
model=context.model_group,
|
||||
request_kwargs={},
|
||||
)
|
||||
except Exception: # noqa: BLE001 # a failed deployment lookup falls back to provider env credentials
|
||||
verbose_proxy_logger.warning(
|
||||
"Could not resolve a deployment for %s to settle background interaction %s; "
|
||||
"falling back to provider environment credentials",
|
||||
context.model_group,
|
||||
context.interaction_id,
|
||||
exc_info=True,
|
||||
)
|
||||
return _SettlementCredentials(api_key=None, api_base=None)
|
||||
litellm_params = deployment.get("litellm_params") or {}
|
||||
return _SettlementCredentials(
|
||||
api_key=litellm_params.get("api_key"),
|
||||
api_base=litellm_params.get("api_base"),
|
||||
)
|
||||
|
||||
|
||||
async def fetch_interaction_for_settlement(
|
||||
context: BackgroundSettlementContext,
|
||||
router: Optional[DeploymentResolver] = None,
|
||||
) -> InteractionsAPIResponse:
|
||||
from litellm.interactions import aget
|
||||
|
||||
credentials = await _resolve_settlement_credentials(context, router=router)
|
||||
return await aget(
|
||||
interaction_id=context.interaction_id,
|
||||
custom_llm_provider=context.custom_llm_provider,
|
||||
**{"api_key": credentials.api_key, "api_base": credentials.api_base, "no-log": True},
|
||||
)
|
||||
|
||||
|
||||
def rebuild_logging_for_settlement(context: BackgroundSettlementContext) -> "LiteLLMLoggingObj":
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj = Logging(
|
||||
model=context.model,
|
||||
messages=[{"role": "user", "content": f"<background_interaction_settlement/{context.interaction_id}>"}],
|
||||
stream=False,
|
||||
call_type=context.call_type,
|
||||
start_time=datetime.now(timezone.utc),
|
||||
litellm_call_id=context.litellm_call_id,
|
||||
function_id=str(uuid.uuid4()),
|
||||
litellm_trace_id=context.litellm_trace_id,
|
||||
)
|
||||
attribution = {key: value for key, value in context.attribution.model_dump().items() if value is not None}
|
||||
reservation = context.budget_reservation.model_dump() if context.budget_reservation is not None else None
|
||||
metadata = attribution if reservation is None else {**attribution, "user_api_key_budget_reservation": reservation}
|
||||
logging_obj.update_environment_variables(
|
||||
litellm_params={
|
||||
"litellm_call_id": context.litellm_call_id,
|
||||
"proxy_server_request": {"headers": {"user-agent": "litellm-background-settlement"}},
|
||||
"metadata": metadata,
|
||||
},
|
||||
optional_params={},
|
||||
model=context.model,
|
||||
custom_llm_provider=context.custom_llm_provider,
|
||||
)
|
||||
return logging_obj
|
||||
|
||||
|
||||
async def _release_persisted_reservation(context: BackgroundSettlementContext) -> None:
|
||||
if context.budget_reservation is None:
|
||||
return
|
||||
from litellm.proxy.spend_tracking.budget_reservation import release_budget_reservation
|
||||
|
||||
try:
|
||||
await release_budget_reservation(budget_reservation=context.budget_reservation.model_dump())
|
||||
except Exception: # noqa: BLE001 # a failed release must not fail settlement; counters expire via TTL
|
||||
verbose_proxy_logger.exception(
|
||||
"Failed to release persisted budget reservation for background interaction %s",
|
||||
context.interaction_id,
|
||||
)
|
||||
|
||||
|
||||
async def settle_claimed_row(
|
||||
row: PendingSettlementRow,
|
||||
response: InteractionsAPIResponse,
|
||||
store: SettlementRowStore,
|
||||
) -> None:
|
||||
if response.status in _TERMINAL_STATUSES and response.usage is not None:
|
||||
try:
|
||||
logging_obj = rebuild_logging_for_settlement(row.context)
|
||||
await logging_obj.async_log_background_interaction_completion(result=response)
|
||||
await store.record_outcome(row.interaction_id, "billed")
|
||||
return
|
||||
except Exception: # noqa: BLE001 # a billing failure after winning the claim must be surfaced, not retried
|
||||
verbose_proxy_logger.exception(
|
||||
"Billing failed after claiming settlement for background interaction %s; "
|
||||
"its spend may be under-tracked",
|
||||
row.interaction_id,
|
||||
)
|
||||
await _release_persisted_reservation(row.context)
|
||||
await store.record_outcome(row.interaction_id, "error")
|
||||
return
|
||||
await _release_persisted_reservation(row.context)
|
||||
await store.record_outcome(row.interaction_id, "released")
|
||||
|
||||
|
||||
async def settle_row_before_delete(
|
||||
row: PendingSettlementRow,
|
||||
store: SettlementRowStore,
|
||||
fetch: SettlementFetch,
|
||||
) -> None:
|
||||
try:
|
||||
response = await fetch(row.context)
|
||||
except Exception as e: # noqa: BLE001 # unfetchable pre-delete state settles by releasing the reservation
|
||||
verbose_proxy_logger.debug(
|
||||
"Could not fetch background interaction %s before delete, releasing its reservation: %s",
|
||||
row.interaction_id,
|
||||
e,
|
||||
)
|
||||
if not await store.claim(row.interaction_id):
|
||||
return
|
||||
await _release_persisted_reservation(row.context)
|
||||
await store.record_outcome(row.interaction_id, "released")
|
||||
return
|
||||
if not await store.claim(row.interaction_id):
|
||||
return
|
||||
await settle_claimed_row(row=row, response=response, store=store)
|
||||
|
||||
|
||||
async def sweep_pending_settlements(
|
||||
store: SettlementRowStore,
|
||||
fetch: SettlementFetch = fetch_interaction_for_settlement,
|
||||
min_age_seconds: float = BACKGROUND_SETTLEMENT_SWEEP_MIN_AGE_SECONDS,
|
||||
limit: int = MAX_BACKGROUND_SETTLEMENTS_PER_SWEEP,
|
||||
) -> None:
|
||||
now = datetime.now(timezone.utc)
|
||||
rows = await store.list_due(older_than=now - timedelta(seconds=min_age_seconds), limit=limit)
|
||||
for row in rows:
|
||||
await _sweep_row(row=row, store=store, fetch=fetch, now=now)
|
||||
|
||||
|
||||
async def _sweep_row(
|
||||
row: PendingSettlementRow,
|
||||
store: SettlementRowStore,
|
||||
fetch: SettlementFetch,
|
||||
now: datetime,
|
||||
) -> None:
|
||||
if row.timeout_at <= now:
|
||||
if not await store.claim(row.interaction_id):
|
||||
return
|
||||
verbose_proxy_logger.warning(
|
||||
"Abandoning settlement for background interaction %s past its %s timeout; its usage will not be tracked",
|
||||
row.interaction_id,
|
||||
row.timeout_at,
|
||||
)
|
||||
await _release_persisted_reservation(row.context)
|
||||
await store.record_outcome(row.interaction_id, "abandoned")
|
||||
return
|
||||
try:
|
||||
response = await fetch(row.context)
|
||||
except Exception as e: # noqa: BLE001 # a sweep fetch error must leave the row pending for a later cycle
|
||||
verbose_proxy_logger.debug(
|
||||
"Settlement sweep could not fetch background interaction %s, leaving it pending: %s",
|
||||
row.interaction_id,
|
||||
e,
|
||||
)
|
||||
return
|
||||
if response.status not in _TERMINAL_STATUSES:
|
||||
return
|
||||
if not await store.claim(row.interaction_id):
|
||||
return
|
||||
await settle_claimed_row(row=row, response=response, store=store)
|
||||
|
||||
|
||||
def _parse_row(interaction_id: str, context_value: object, timeout_at: datetime) -> Optional[PendingSettlementRow]:
|
||||
try:
|
||||
raw = json.loads(context_value) if isinstance(context_value, str) else context_value
|
||||
return PendingSettlementRow(
|
||||
interaction_id=interaction_id,
|
||||
context=BackgroundSettlementContext.model_validate(raw),
|
||||
timeout_at=timeout_at,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # an unparseable row is surfaced and skipped rather than crashing the sweep
|
||||
verbose_proxy_logger.exception(
|
||||
"Could not parse persisted settlement context for background interaction %s",
|
||||
interaction_id,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class PrismaBackgroundSettlementStore:
|
||||
def __init__(
|
||||
self,
|
||||
prisma_client: "PrismaClient",
|
||||
fetch: SettlementFetch = fetch_interaction_for_settlement,
|
||||
claimed_by: Optional[str] = None,
|
||||
) -> None:
|
||||
self.prisma_client = prisma_client
|
||||
self.fetch = fetch
|
||||
self.claimed_by = claimed_by or str(uuid.uuid4())
|
||||
|
||||
@property
|
||||
def _table(self):
|
||||
return self.prisma_client.db.litellm_backgroundinteractionsettlementtable
|
||||
|
||||
async def persist_pending(self, context: BackgroundSettlementContext, timeout_at: datetime) -> None:
|
||||
context_json = context.model_dump_json()
|
||||
await self._table.upsert(
|
||||
where={"interaction_id": context.interaction_id},
|
||||
data={
|
||||
"create": {
|
||||
"interaction_id": context.interaction_id,
|
||||
"context": context_json,
|
||||
"timeout_at": timeout_at,
|
||||
},
|
||||
"update": {"context": context_json, "timeout_at": timeout_at},
|
||||
},
|
||||
)
|
||||
|
||||
async def claim(self, interaction_id: str) -> bool:
|
||||
updated = await self._table.update_many(
|
||||
where={"interaction_id": interaction_id, "status": "pending"},
|
||||
data={"status": "settled", "claimed_by": self.claimed_by},
|
||||
)
|
||||
if updated > 0:
|
||||
return True
|
||||
row = await self._table.find_unique(where={"interaction_id": interaction_id})
|
||||
return row is None
|
||||
|
||||
async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None:
|
||||
await self._table.update_many(
|
||||
where={"interaction_id": interaction_id},
|
||||
data={"outcome": outcome},
|
||||
)
|
||||
|
||||
async def is_pending(self, interaction_id: str) -> bool:
|
||||
row = await self._table.find_unique(where={"interaction_id": interaction_id})
|
||||
return row is not None and row.status == "pending"
|
||||
|
||||
async def get_pending(self, interaction_id: str) -> Optional[PendingSettlementRow]:
|
||||
row = await self._table.find_unique(where={"interaction_id": interaction_id})
|
||||
if row is None or row.status != "pending":
|
||||
return None
|
||||
return _parse_row(interaction_id=row.interaction_id, context_value=row.context, timeout_at=row.timeout_at)
|
||||
|
||||
async def list_due(self, older_than: datetime, limit: int) -> tuple[PendingSettlementRow, ...]:
|
||||
rows = await self._table.find_many(
|
||||
where={"status": "pending", "created_at": {"lt": older_than}},
|
||||
take=limit,
|
||||
order={"created_at": "asc"},
|
||||
)
|
||||
parsed = (
|
||||
_parse_row(interaction_id=row.interaction_id, context_value=row.context, timeout_at=row.timeout_at)
|
||||
for row in rows
|
||||
)
|
||||
return tuple(row for row in parsed if row is not None)
|
||||
|
||||
async def settle_pending_before_delete(self, interaction_id: str) -> None:
|
||||
row = await self.get_pending(interaction_id)
|
||||
if row is None:
|
||||
return
|
||||
await settle_row_before_delete(row=row, store=self, fetch=self.fetch)
|
||||
|
|
@ -950,6 +950,19 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
|
|||
@@index([team_id, created_at(sort: Desc)])
|
||||
}
|
||||
|
||||
model LiteLLM_BackgroundInteractionSettlementTable {
|
||||
interaction_id String @id
|
||||
status String @default("pending")
|
||||
outcome String?
|
||||
context Json
|
||||
claimed_by String?
|
||||
timeout_at DateTime
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@index([status, created_at], map: "LiteLLM_BackgroundInteractionSettlement_status_created_idx")
|
||||
}
|
||||
|
||||
model LiteLLM_ManagedVectorStoreTable {
|
||||
id String @id @default(uuid())
|
||||
unified_resource_id String @unique // The base64 encoded unified vector store ID
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
|
|
@ -7,6 +8,9 @@ import pytest
|
|||
from litellm.interactions.background_cost_polling import (
|
||||
_SETTLED_KEY,
|
||||
BackgroundInteractionPollContext,
|
||||
BackgroundSettlementContext,
|
||||
SettlementOutcome,
|
||||
get_settlement_store,
|
||||
maybe_schedule_background_interaction_cost_polling,
|
||||
maybe_settle_background_interaction_before_delete,
|
||||
poll_and_log_background_interaction_cost,
|
||||
|
|
@ -344,6 +348,231 @@ async def test_poller_exits_without_billing_once_settled_elsewhere():
|
|||
assert logging_obj.model_call_details.get("response_cost") is None
|
||||
|
||||
|
||||
class _InMemorySettlementStore:
|
||||
def __init__(self, persist_error: Optional[Exception] = None) -> None:
|
||||
self.rows: dict[str, dict] = {}
|
||||
self.contexts: dict[str, BackgroundSettlementContext] = {}
|
||||
self.outcomes: dict[str, SettlementOutcome] = {}
|
||||
self.settle_before_delete_calls: list[str] = []
|
||||
self.persist_error = persist_error
|
||||
|
||||
async def persist_pending(self, context: BackgroundSettlementContext, timeout_at: datetime) -> None:
|
||||
if self.persist_error is not None:
|
||||
raise self.persist_error
|
||||
existing = self.rows.get(context.interaction_id)
|
||||
status = existing["status"] if existing is not None else "pending"
|
||||
self.rows[context.interaction_id] = {"status": status, "timeout_at": timeout_at}
|
||||
self.contexts[context.interaction_id] = context
|
||||
|
||||
async def claim(self, interaction_id: str) -> bool:
|
||||
row = self.rows.get(interaction_id)
|
||||
if row is None:
|
||||
return True
|
||||
if row["status"] != "pending":
|
||||
return False
|
||||
row["status"] = "settled"
|
||||
return True
|
||||
|
||||
async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None:
|
||||
self.outcomes[interaction_id] = outcome
|
||||
|
||||
async def is_pending(self, interaction_id: str) -> bool:
|
||||
row = self.rows.get(interaction_id)
|
||||
return row is not None and row["status"] == "pending"
|
||||
|
||||
async def settle_pending_before_delete(self, interaction_id: str) -> None:
|
||||
self.settle_before_delete_calls.append(interaction_id)
|
||||
|
||||
|
||||
def _attributed_litellm_params(reservation: Optional[dict] = None) -> dict:
|
||||
metadata = {
|
||||
"user_api_key": "hashed-virtual-key",
|
||||
"user_api_key_user_id": "user-123",
|
||||
"user_api_key_team_id": "team-456",
|
||||
"deployment_model_name": "gemini-3-flash-preview",
|
||||
}
|
||||
if reservation is not None:
|
||||
metadata["user_api_key_budget_reservation"] = reservation
|
||||
return {"metadata": metadata}
|
||||
|
||||
|
||||
def test_no_settlement_store_registered_by_default():
|
||||
assert get_settlement_store() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poller_persists_settlement_context_with_attribution_and_reservation():
|
||||
store = _InMemorySettlementStore()
|
||||
reservation = _reservation()
|
||||
logging_obj = _logging_obj(litellm_params=_attributed_litellm_params(reservation))
|
||||
fetch, _ = _fetch_sequence(_response("completed", with_usage=True))
|
||||
|
||||
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch, store=store)
|
||||
|
||||
context = store.contexts["interactions/bg-abc"]
|
||||
assert context.attribution.user_api_key == "hashed-virtual-key"
|
||||
assert context.attribution.user_api_key_user_id == "user-123"
|
||||
assert context.attribution.user_api_key_team_id == "team-456"
|
||||
assert context.model == "gemini-2.5-flash"
|
||||
assert context.model_group == "gemini-3-flash-preview"
|
||||
assert context.custom_llm_provider == "gemini"
|
||||
assert context.call_type == "acreate_interaction"
|
||||
assert context.litellm_call_id == "bg-interactions-call-id"
|
||||
assert context.budget_reservation is not None
|
||||
assert context.budget_reservation.reserved_cost == reservation["reserved_cost"]
|
||||
assert store.rows["interactions/bg-abc"]["timeout_at"] > datetime.now(timezone.utc)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poller_records_billed_outcome_so_sweep_cannot_rebill():
|
||||
store = _InMemorySettlementStore()
|
||||
logging_obj = _logging_obj()
|
||||
fetch, _ = _fetch_sequence(_response("completed", with_usage=True))
|
||||
|
||||
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch, store=store)
|
||||
|
||||
assert logging_obj.model_call_details["response_cost"] > 0
|
||||
assert store.rows["interactions/bg-abc"]["status"] == "settled"
|
||||
assert store.outcomes["interactions/bg-abc"] == "billed"
|
||||
|
||||
|
||||
def _store_with_row_settled_elsewhere(interaction_id: str = "interactions/bg-abc") -> _InMemorySettlementStore:
|
||||
store = _InMemorySettlementStore()
|
||||
store.rows[interaction_id] = {"status": "settled", "timeout_at": datetime.now(timezone.utc)}
|
||||
return store
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poller_skips_billing_and_finalizes_locally_when_store_claim_lost():
|
||||
store = _store_with_row_settled_elsewhere()
|
||||
reservation = _reservation()
|
||||
logging_obj = _logging_obj_with_reservation(reservation)
|
||||
fetch, _ = _fetch_sequence(_response("completed", with_usage=True))
|
||||
|
||||
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch, store=store)
|
||||
|
||||
assert logging_obj.model_call_details.get("response_cost") is None
|
||||
assert reservation["finalized"] is True
|
||||
assert store.outcomes == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_give_up_does_not_release_when_store_claim_lost():
|
||||
store = _store_with_row_settled_elsewhere()
|
||||
reservation = _reservation()
|
||||
logging_obj = _logging_obj_with_reservation(reservation)
|
||||
fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
|
||||
|
||||
await poll_and_log_background_interaction_cost(
|
||||
_context(logging_obj, timeout_seconds=0.01),
|
||||
fetch_interaction=fetch,
|
||||
store=store,
|
||||
)
|
||||
|
||||
assert reservation["finalized"] is True
|
||||
assert store.outcomes == {}
|
||||
assert logging_obj.model_call_details.get("response_cost") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_give_up_records_abandoned_outcome_when_claim_won():
|
||||
store = _InMemorySettlementStore()
|
||||
reservation = _reservation()
|
||||
logging_obj = _logging_obj_with_reservation(reservation)
|
||||
|
||||
await poll_and_log_background_interaction_cost(
|
||||
_context(logging_obj, timeout_seconds=0.01),
|
||||
fetch_interaction=_fetch_sequence(_response("in_progress", with_usage=False))[0],
|
||||
store=store,
|
||||
)
|
||||
|
||||
assert reservation["finalized"] is True
|
||||
assert store.rows["interactions/bg-abc"]["status"] == "settled"
|
||||
assert store.outcomes["interactions/bg-abc"] == "abandoned"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poller_continues_in_memory_when_persist_fails():
|
||||
store = _InMemorySettlementStore(persist_error=RuntimeError("db unavailable"))
|
||||
logging_obj = _logging_obj()
|
||||
fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
|
||||
|
||||
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch, store=store)
|
||||
|
||||
assert len(calls) == 1
|
||||
assert logging_obj.model_call_details["response_cost"] > 0
|
||||
assert store.rows == {}
|
||||
assert store.outcomes == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poller_exits_early_on_fetch_error_once_settled_elsewhere():
|
||||
store = _store_with_row_settled_elsewhere()
|
||||
reservation = _reservation()
|
||||
logging_obj = _logging_obj_with_reservation(reservation)
|
||||
fetch, calls = _fetch_sequence(RuntimeError("interaction deleted on another pod"))
|
||||
|
||||
await poll_and_log_background_interaction_cost(
|
||||
_context(logging_obj, timeout_seconds=1.0),
|
||||
fetch_interaction=fetch,
|
||||
store=store,
|
||||
)
|
||||
|
||||
assert len(calls) == 1
|
||||
assert reservation["finalized"] is True
|
||||
assert logging_obj.model_call_details.get("response_cost") is None
|
||||
assert store.outcomes == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_registry_miss_delegates_to_store():
|
||||
store = _InMemorySettlementStore()
|
||||
|
||||
await maybe_settle_background_interaction_before_delete(
|
||||
interaction_id="interactions/created-on-another-pod",
|
||||
store=store,
|
||||
)
|
||||
|
||||
assert store.settle_before_delete_calls == ["interactions/created-on-another-pod"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_registry_hit_records_billed_outcome_in_store():
|
||||
store = _InMemorySettlementStore()
|
||||
logging_obj = _logging_obj()
|
||||
task = _register_poll(logging_obj)
|
||||
fetch, _ = _fetch_sequence(_response("completed", with_usage=True))
|
||||
|
||||
await maybe_settle_background_interaction_before_delete(
|
||||
interaction_id="interactions/bg-abc",
|
||||
fetch_interaction=fetch,
|
||||
store=store,
|
||||
)
|
||||
|
||||
assert logging_obj.model_call_details["response_cost"] > 0
|
||||
assert store.outcomes["interactions/bg-abc"] == "billed"
|
||||
await asyncio.wait_for(task, timeout=5)
|
||||
|
||||
|
||||
def test_settlement_reservation_shape_matches_budget_reservation_module():
|
||||
from litellm.interactions.background_cost_polling import SettlementReservation
|
||||
from litellm.proxy.spend_tracking.budget_reservation import _BudgetCounter, _counter_to_reservation_entry
|
||||
|
||||
entry = _counter_to_reservation_entry(
|
||||
counter=_BudgetCounter(
|
||||
counter_key="spend:key:hashed-virtual-key",
|
||||
max_budget=10.0,
|
||||
fallback_spend=0.0,
|
||||
entity_type="key",
|
||||
entity_id="hashed-virtual-key",
|
||||
),
|
||||
reserved_cost=0.05,
|
||||
)
|
||||
reservation = {"reserved_cost": 0.05, "entries": [entry], "finalized": False, "input_cost": 0.001}
|
||||
|
||||
assert SettlementReservation.model_validate(reservation).model_dump() == reservation
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schedule_respects_kill_switch(monkeypatch):
|
||||
import litellm.interactions.background_cost_polling as module
|
||||
|
|
|
|||
|
|
@ -0,0 +1,320 @@
|
|||
import asyncio
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.interactions.background_cost_polling import (
|
||||
BackgroundSettlementContext,
|
||||
SettlementKeyAttribution,
|
||||
SettlementOutcome,
|
||||
SettlementReservation,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.background_settlement import (
|
||||
PendingSettlementRow,
|
||||
_parse_row,
|
||||
_resolve_settlement_credentials,
|
||||
rebuild_logging_for_settlement,
|
||||
settle_row_before_delete,
|
||||
sweep_pending_settlements,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.types.interactions import InteractionsAPIResponse
|
||||
|
||||
USAGE_BLOCK = {
|
||||
"total_tokens": 175,
|
||||
"total_input_tokens": 100,
|
||||
"input_tokens_by_modality": [{"modality": "text", "tokens": 100}],
|
||||
"total_cached_tokens": 0,
|
||||
"total_output_tokens": 50,
|
||||
"output_tokens_by_modality": [{"modality": "text", "tokens": 50}],
|
||||
"total_tool_use_tokens": 0,
|
||||
"total_thought_tokens": 25,
|
||||
}
|
||||
|
||||
INTERACTION_ID = "interactions/bg-foreign"
|
||||
|
||||
|
||||
def _settlement_context(
|
||||
reservation: Optional[SettlementReservation] = None,
|
||||
model_group: Optional[str] = None,
|
||||
) -> BackgroundSettlementContext:
|
||||
return BackgroundSettlementContext(
|
||||
interaction_id=INTERACTION_ID,
|
||||
custom_llm_provider="gemini",
|
||||
model="gemini-2.5-flash",
|
||||
model_group=model_group,
|
||||
litellm_call_id="original-create-call-id",
|
||||
call_type="acreate_interaction",
|
||||
attribution=SettlementKeyAttribution(
|
||||
user_api_key="hashed-virtual-key",
|
||||
user_api_key_user_id="user-123",
|
||||
user_api_key_team_id="team-456",
|
||||
),
|
||||
budget_reservation=reservation,
|
||||
)
|
||||
|
||||
|
||||
def _reservation() -> SettlementReservation:
|
||||
return SettlementReservation(reserved_cost=0.05, entries=[], finalized=False, input_cost=0.001)
|
||||
|
||||
|
||||
def _row(
|
||||
context: Optional[BackgroundSettlementContext] = None,
|
||||
timeout_at: Optional[datetime] = None,
|
||||
) -> PendingSettlementRow:
|
||||
return PendingSettlementRow(
|
||||
interaction_id=INTERACTION_ID,
|
||||
context=context if context is not None else _settlement_context(),
|
||||
timeout_at=timeout_at if timeout_at is not None else datetime.now(timezone.utc) + timedelta(hours=1),
|
||||
)
|
||||
|
||||
|
||||
def _response(status: str, with_usage: bool) -> InteractionsAPIResponse:
|
||||
return InteractionsAPIResponse(
|
||||
id=INTERACTION_ID,
|
||||
model="gemini-2.5-flash",
|
||||
status=status,
|
||||
steps=[],
|
||||
usage=dict(USAGE_BLOCK) if with_usage else None,
|
||||
)
|
||||
|
||||
|
||||
def _fetch_returning(item):
|
||||
calls = []
|
||||
|
||||
async def fetch(context: BackgroundSettlementContext) -> InteractionsAPIResponse:
|
||||
calls.append(context.interaction_id)
|
||||
if isinstance(item, Exception):
|
||||
raise item
|
||||
return item
|
||||
|
||||
return fetch, calls
|
||||
|
||||
|
||||
class _FakeRowStore:
|
||||
def __init__(self, rows: tuple[PendingSettlementRow, ...]) -> None:
|
||||
self.rows = {row.interaction_id: row for row in rows}
|
||||
self.status = {row.interaction_id: "pending" for row in rows}
|
||||
self.outcomes: dict[str, SettlementOutcome] = {}
|
||||
self.claims_won = 0
|
||||
|
||||
async def claim(self, interaction_id: str) -> bool:
|
||||
if self.status.get(interaction_id) != "pending":
|
||||
return False
|
||||
self.status[interaction_id] = "settled"
|
||||
self.claims_won += 1
|
||||
return True
|
||||
|
||||
async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None:
|
||||
self.outcomes[interaction_id] = outcome
|
||||
|
||||
async def list_due(self, older_than: datetime, limit: int) -> tuple[PendingSettlementRow, ...]:
|
||||
due = tuple(row for row_id, row in self.rows.items() if self.status[row_id] == "pending")
|
||||
return due[:limit]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rebuilt_logging_bills_with_original_attribution_and_request_id():
|
||||
logging_obj = rebuild_logging_for_settlement(_settlement_context())
|
||||
response = _response("completed", with_usage=True)
|
||||
|
||||
await logging_obj.async_log_background_interaction_completion(result=response)
|
||||
|
||||
assert logging_obj.model_call_details["response_cost"] > 0
|
||||
payload = get_logging_payload(
|
||||
kwargs=logging_obj.model_call_details,
|
||||
response_obj=response,
|
||||
start_time=datetime.now(timezone.utc),
|
||||
end_time=datetime.now(timezone.utc),
|
||||
)
|
||||
assert payload["request_id"] == INTERACTION_ID
|
||||
assert payload["user"] == "user-123"
|
||||
assert payload["team_id"] == "team-456"
|
||||
assert payload["api_key"]
|
||||
assert payload["spend"] > 0
|
||||
|
||||
|
||||
def test_rebuilt_logging_carries_reservation_for_reconcile():
|
||||
context = _settlement_context(reservation=_reservation())
|
||||
logging_obj = rebuild_logging_for_settlement(context)
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
|
||||
|
||||
metadata = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)
|
||||
assert metadata["user_api_key_budget_reservation"]["reserved_cost"] == 0.05
|
||||
assert metadata["user_api_key_budget_reservation"]["finalized"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sweep_bills_terminal_row_exactly_once_across_concurrent_sweeps():
|
||||
store = _FakeRowStore(rows=(_row(),))
|
||||
fetch, calls = _fetch_returning(_response("completed", with_usage=True))
|
||||
|
||||
await asyncio.gather(
|
||||
sweep_pending_settlements(store=store, fetch=fetch),
|
||||
sweep_pending_settlements(store=store, fetch=fetch),
|
||||
)
|
||||
|
||||
assert store.claims_won == 1
|
||||
assert store.outcomes == {INTERACTION_ID: "billed"}
|
||||
assert store.status[INTERACTION_ID] == "settled"
|
||||
assert len(calls) >= 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sweep_abandons_row_past_timeout_without_fetching():
|
||||
store = _FakeRowStore(rows=(_row(timeout_at=datetime.now(timezone.utc) - timedelta(seconds=1)),))
|
||||
fetch, calls = _fetch_returning(_response("completed", with_usage=True))
|
||||
|
||||
await sweep_pending_settlements(store=store, fetch=fetch)
|
||||
|
||||
assert calls == []
|
||||
assert store.outcomes == {INTERACTION_ID: "abandoned"}
|
||||
assert store.status[INTERACTION_ID] == "settled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sweep_leaves_row_pending_on_fetch_error():
|
||||
store = _FakeRowStore(rows=(_row(),))
|
||||
fetch, calls = _fetch_returning(RuntimeError("provider unavailable"))
|
||||
|
||||
await sweep_pending_settlements(store=store, fetch=fetch)
|
||||
|
||||
assert len(calls) == 1
|
||||
assert store.status[INTERACTION_ID] == "pending"
|
||||
assert store.outcomes == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sweep_leaves_row_pending_while_interaction_still_running():
|
||||
store = _FakeRowStore(rows=(_row(),))
|
||||
fetch, calls = _fetch_returning(_response("in_progress", with_usage=False))
|
||||
|
||||
await sweep_pending_settlements(store=store, fetch=fetch)
|
||||
|
||||
assert len(calls) == 1
|
||||
assert store.status[INTERACTION_ID] == "pending"
|
||||
assert store.outcomes == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_settlement_releases_when_interaction_not_terminal():
|
||||
row = _row(context=_settlement_context(reservation=_reservation()))
|
||||
store = _FakeRowStore(rows=(row,))
|
||||
fetch, _ = _fetch_returning(_response("in_progress", with_usage=False))
|
||||
|
||||
await settle_row_before_delete(row=row, store=store, fetch=fetch)
|
||||
|
||||
assert store.outcomes == {INTERACTION_ID: "released"}
|
||||
assert store.status[INTERACTION_ID] == "settled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_settlement_bills_terminal_interaction():
|
||||
row = _row()
|
||||
store = _FakeRowStore(rows=(row,))
|
||||
fetch, _ = _fetch_returning(_response("completed", with_usage=True))
|
||||
|
||||
await settle_row_before_delete(row=row, store=store, fetch=fetch)
|
||||
|
||||
assert store.outcomes == {INTERACTION_ID: "billed"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_settlement_releases_on_fetch_error():
|
||||
row = _row(context=_settlement_context(reservation=_reservation()))
|
||||
store = _FakeRowStore(rows=(row,))
|
||||
fetch, _ = _fetch_returning(RuntimeError("interaction already deleted"))
|
||||
|
||||
await settle_row_before_delete(row=row, store=store, fetch=fetch)
|
||||
|
||||
assert store.outcomes == {INTERACTION_ID: "released"}
|
||||
assert store.status[INTERACTION_ID] == "settled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_settlement_noop_when_claim_lost():
|
||||
row = _row()
|
||||
store = _FakeRowStore(rows=(row,))
|
||||
store.status[INTERACTION_ID] = "settled"
|
||||
fetch, _ = _fetch_returning(_response("completed", with_usage=True))
|
||||
|
||||
await settle_row_before_delete(row=row, store=store, fetch=fetch)
|
||||
|
||||
assert store.outcomes == {}
|
||||
|
||||
|
||||
class _FakeRouter:
|
||||
def __init__(self, deployment: Optional[dict] = None, error: Optional[Exception] = None) -> None:
|
||||
self.deployment = deployment
|
||||
self.error = error
|
||||
self.requested_models: list[str] = []
|
||||
|
||||
async def async_get_available_deployment(self, model: str, request_kwargs: dict) -> dict:
|
||||
self.requested_models.append(model)
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
assert self.deployment is not None
|
||||
return self.deployment
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credentials_resolved_from_router_deployment():
|
||||
router = _FakeRouter(deployment={"litellm_params": {"api_key": "resolved-key", "api_base": "https://resolved"}})
|
||||
|
||||
credentials = await _resolve_settlement_credentials(
|
||||
_settlement_context(model_group="gemini-3-flash-preview"),
|
||||
router=router,
|
||||
)
|
||||
|
||||
assert credentials.api_key == "resolved-key"
|
||||
assert credentials.api_base == "https://resolved"
|
||||
assert router.requested_models == ["gemini-3-flash-preview"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credentials_fall_back_to_env_when_router_lookup_fails():
|
||||
router = _FakeRouter(error=RuntimeError("no deployments available"))
|
||||
|
||||
credentials = await _resolve_settlement_credentials(
|
||||
_settlement_context(model_group="gemini-3-flash-preview"),
|
||||
router=router,
|
||||
)
|
||||
|
||||
assert credentials.api_key is None
|
||||
assert credentials.api_base is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credentials_fall_back_to_env_without_model_group():
|
||||
credentials = await _resolve_settlement_credentials(_settlement_context(model_group=None))
|
||||
|
||||
assert credentials.api_key is None
|
||||
assert credentials.api_base is None
|
||||
|
||||
|
||||
def test_parse_row_roundtrips_persisted_context_json():
|
||||
context = _settlement_context(reservation=_reservation(), model_group="gemini-3-flash-preview")
|
||||
timeout_at = datetime.now(timezone.utc)
|
||||
|
||||
row = _parse_row(
|
||||
interaction_id=INTERACTION_ID,
|
||||
context_value=context.model_dump_json(),
|
||||
timeout_at=timeout_at,
|
||||
)
|
||||
|
||||
assert row is not None
|
||||
assert row.context == context
|
||||
assert row.timeout_at == timeout_at
|
||||
|
||||
|
||||
def test_parse_row_returns_none_for_corrupt_context():
|
||||
assert (
|
||||
_parse_row(
|
||||
interaction_id=INTERACTION_ID,
|
||||
context_value="{not json",
|
||||
timeout_at=datetime.now(timezone.utc),
|
||||
)
|
||||
is None
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue