mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
feat(proxy): maintain per-window budget spend rows in the spend writer
Multi-window budgets (budget_limits on keys and teams) enforce off Redis counters with a 60s TTL. Every cold counter, and every authoritative floor check, aggregates LiteLLM_SpendLogs with a range scan over startTime; that table has no index on api_key or team_id, so the scan lands on the highest volume table in the schema and saturates the connection pool (#35766). Every other budget feature reads a maintained running spend value instead. This gives windows the same shape by keeping LiteLLM_BudgetWindowSpend up to date: one row per configured window, whose window_start rolls forward in place. The cost callback already iterates a key's and a team's windows with the window start computed and the actual cost in hand, so it enqueues there, onto a new WindowSpendUpdateQueue. Increments are enqueued even when the cache increment is skipped for a reserved counter, since the reservation only pre-charged an estimate and the row still owes the actual cost. Windows with no reset_at slide with wall clock and cannot be represented by a single row, so they are left to the read path's existing aggregate. The queue flushes alongside the daily spend queues, through the Redis buffer when one is configured (only the pod-lock winner commits) and directly otherwise. A flush selects the primary keys that already exist, seeds the ones that do not from LiteLLM_SpendLogs so a new row cannot undercount spend that predates it, and applies one batch of upserts ordered by primary key. An increment at or behind the stored window_start adds into the row, matching how in-flight requests carry into a window after a reset; a newer one rolls the window and starts from that increment. The conflict arm adds only the increment, never the seeded base, so two pods seeding the same new window cannot double count it. The seed excludes the requests its own batch is about to apply. Spend logs are drained by a separate monitor that fires on a ~2s poll whenever anything is queued, while window increments flush on the much slower batch tick, so by seed time the batch's log rows are normally already in the table; counting them in the aggregate and again in the increments made a fresh row land at exactly twice the true spend. Each increment therefore carries the LiteLLM_SpendLogs request_id it was recorded under, which update_database now returns rather than having the callback re-derive it (a cache hit appends time.time() to that id, so a second derivation would not match). The reset job rolls each expired window's row alongside the counter it zeroes, conditional on the stored window_start still being behind the new one so a pod that already rolled it is not clobbered.
This commit is contained in:
ryan-crabbe-berri
parent
3b1ab1908c
commit
187e0fab60
16 changed files with 2317 additions and 19 deletions
|
|
@ -270,6 +270,7 @@ REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update
|
|||
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_spend_update_buffer"
|
||||
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer"
|
||||
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer"
|
||||
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer"
|
||||
MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
|
||||
# Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth
|
||||
LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000))
|
||||
|
|
|
|||
|
|
@ -2,16 +2,18 @@ import asyncio
|
|||
import json
|
||||
import time
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final, Literal, Protocol, TypeVar
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTableFull,
|
||||
LiteLLM_EndUserTable,
|
||||
Litellm_EntityType,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
LiteLLM_VerificationToken,
|
||||
|
|
@ -21,6 +23,7 @@ from litellm.proxy.common_utils.timezone_utils import (
|
|||
compute_budget_reset_at,
|
||||
get_budget_reset_settings,
|
||||
)
|
||||
from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.table_repositories import (
|
||||
|
|
@ -782,6 +785,9 @@ class ResetBudgetJob:
|
|||
spend_counter_cache: DualCache,
|
||||
now: datetime,
|
||||
reset_settings: BudgetResetSettings,
|
||||
prisma_client: PrismaClient,
|
||||
entity_type: Litellm_EntityType,
|
||||
entity_id: str,
|
||||
) -> bool:
|
||||
"""Reset a single budget window if expired. Returns True if the window was reset."""
|
||||
reset_at_str: Final = window.get("reset_at")
|
||||
|
|
@ -796,11 +802,56 @@ class ResetBudgetJob:
|
|||
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=0.0)
|
||||
except Exception as redis_err:
|
||||
verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err)
|
||||
window["reset_at"] = compute_budget_reset_at(
|
||||
budget_duration=window["budget_duration"], settings=reset_settings
|
||||
).isoformat()
|
||||
budget_duration: Final = window["budget_duration"]
|
||||
next_reset_at: Final = compute_budget_reset_at(budget_duration=budget_duration, settings=reset_settings)
|
||||
window["reset_at"] = next_reset_at.isoformat()
|
||||
await ResetBudgetJob._roll_window_spend_row(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
budget_duration=budget_duration,
|
||||
next_reset_at=next_reset_at,
|
||||
)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
async def _roll_window_spend_row(
|
||||
prisma_client: PrismaClient,
|
||||
entity_type: Litellm_EntityType,
|
||||
entity_id: str,
|
||||
budget_duration: str,
|
||||
next_reset_at: datetime,
|
||||
) -> None:
|
||||
"""Move this window's LiteLLM_BudgetWindowSpend row onto the window
|
||||
that just started, so the maintained total the read path uses starts
|
||||
from zero alongside the counter.
|
||||
|
||||
Best effort: the row is an optimization over aggregating
|
||||
LiteLLM_SpendLogs, so a failure here must not stop the remaining
|
||||
windows from having their counters reset.
|
||||
"""
|
||||
try:
|
||||
window_start: Final = next_reset_at - timedelta(seconds=duration_in_seconds(budget_duration))
|
||||
except Exception as e: # noqa: BLE001 # duration_in_seconds raises bare exceptions on bad input
|
||||
verbose_proxy_logger.warning("Unparseable budget_duration %s: %s", budget_duration, e)
|
||||
return
|
||||
try:
|
||||
await roll_window_spend_row(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=entity_type.value,
|
||||
entity_id=entity_id,
|
||||
window_duration=budget_duration,
|
||||
new_window_start=window_start,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # the row is best effort; counter resets must still land
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to roll budget window spend row for %s=%s window=%s: %s",
|
||||
entity_type.value,
|
||||
entity_id,
|
||||
budget_duration,
|
||||
e,
|
||||
)
|
||||
|
||||
async def reset_budget_windows(self) -> None:
|
||||
"""
|
||||
For keys and teams with budget_limits, reset any individual windows where
|
||||
|
|
@ -836,6 +887,9 @@ class ResetBudgetJob:
|
|||
spend_counter_cache,
|
||||
now,
|
||||
self.reset_settings,
|
||||
prisma_client=self.prisma_client,
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id=row["token"],
|
||||
):
|
||||
changed = True
|
||||
if changed:
|
||||
|
|
@ -865,6 +919,9 @@ class ResetBudgetJob:
|
|||
spend_counter_cache,
|
||||
now,
|
||||
self.reset_settings,
|
||||
prisma_client=self.prisma_client,
|
||||
entity_type=Litellm_EntityType.TEAM,
|
||||
entity_id=row["team_id"],
|
||||
):
|
||||
changed = True
|
||||
if changed:
|
||||
|
|
|
|||
266
litellm/proxy/db/budget_window_spend_writer.py
Normal file
266
litellm/proxy/db/budget_window_spend_writer.py
Normal file
|
|
@ -0,0 +1,266 @@
|
|||
"""
|
||||
Writer for LiteLLM_BudgetWindowSpend.
|
||||
|
||||
The table holds one row per configured budget window whose window_start rolls
|
||||
forward in place, so budget enforcement can read a maintained running total
|
||||
instead of aggregating LiteLLM_SpendLogs every time a window counter goes cold
|
||||
(issue #35766). Raw SQL rather than the Prisma upsert helper because the
|
||||
conditional roll cannot be expressed through the query builder.
|
||||
|
||||
Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, excluding
|
||||
the requests whose increments are in the same batch so neither source counts
|
||||
them twice. One gap survives that exclusion: without the Redis transaction
|
||||
buffer every pod flushes its own increments, so a row seeded by one pod can
|
||||
miss increments still queued on another. That is bounded by a single flush
|
||||
interval, happens at most once per window row, and errs toward under-counting
|
||||
for that interval only.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
WindowSpendTransaction,
|
||||
to_naive_utc,
|
||||
window_spend_group_key,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
_SELECT_EXISTING_ROWS_SQL: Final = (
|
||||
'SELECT entity_type, entity_id, window_duration FROM "LiteLLM_BudgetWindowSpend" '
|
||||
"WHERE (entity_type, entity_id, window_duration) "
|
||||
"IN (SELECT * FROM unnest($1::text[], $2::text[], $3::text[]))"
|
||||
)
|
||||
|
||||
_UPSERT_WINDOW_SPEND_SQL: Final = (
|
||||
'INSERT INTO "LiteLLM_BudgetWindowSpend" '
|
||||
"(entity_type, entity_id, window_duration, window_start, spend, created_at, updated_at) "
|
||||
"VALUES ($1, $2, $3, ($4::timestamptz AT TIME ZONE 'UTC'), $5, "
|
||||
"($7::timestamptz AT TIME ZONE 'UTC'), ($7::timestamptz AT TIME ZONE 'UTC')) "
|
||||
"ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET "
|
||||
"spend = CASE "
|
||||
'WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start '
|
||||
'THEN "LiteLLM_BudgetWindowSpend".spend + $6 '
|
||||
"ELSE EXCLUDED.spend "
|
||||
"END, "
|
||||
'window_start = GREATEST("LiteLLM_BudgetWindowSpend".window_start, EXCLUDED.window_start), '
|
||||
"updated_at = ($7::timestamptz AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
||||
_ROLL_WINDOW_SPEND_SQL: Final = (
|
||||
'UPDATE "LiteLLM_BudgetWindowSpend" SET '
|
||||
"window_start = ($4::timestamptz AT TIME ZONE 'UTC'), "
|
||||
"spend = 0, "
|
||||
"updated_at = ($5::timestamptz AT TIME ZONE 'UTC') "
|
||||
"WHERE entity_type = $1 AND entity_id = $2 AND window_duration = $3 "
|
||||
"AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
||||
_SEED_FROM_SPEND_LOGS_KEY_SQL: Final = (
|
||||
'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" '
|
||||
"WHERE api_key = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') "
|
||||
"AND NOT (request_id = ANY($3::text[]))"
|
||||
)
|
||||
|
||||
_SEED_FROM_SPEND_LOGS_TEAM_SQL: Final = (
|
||||
'SELECT COALESCE(SUM(spend), 0.0) AS total FROM "LiteLLM_SpendLogs" '
|
||||
"WHERE team_id = $1 AND \"startTime\" >= ($2::timestamptz AT TIME ZONE 'UTC') "
|
||||
"AND NOT (request_id = ANY($3::text[]))"
|
||||
)
|
||||
|
||||
_UPSERT_TRANSACTION_TIMEOUT: Final = timedelta(seconds=60)
|
||||
|
||||
|
||||
class WindowSpendLogsAggregate(Protocol):
|
||||
"""Sums LiteLLM_SpendLogs for one entity since window_start, ignoring the
|
||||
requests whose ids are handed in.
|
||||
|
||||
Injected so the flush can be exercised without a database and so the
|
||||
expensive aggregate stays swappable.
|
||||
"""
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
prisma_client: "PrismaClient",
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
exclude_request_ids: Sequence[str],
|
||||
) -> float | None: ...
|
||||
|
||||
|
||||
async def spend_logs_total_excluding(
|
||||
prisma_client: "PrismaClient",
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
exclude_request_ids: Sequence[str],
|
||||
) -> float | None:
|
||||
"""LiteLLM_SpendLogs spend for one entity since window_start, minus the
|
||||
requests already accounted for by the increments being flushed.
|
||||
|
||||
The spend log writer drains its own queue on a ~2s poll whenever anything
|
||||
is queued, while window increments flush on the much slower batch tick, so
|
||||
by the time a window row is seeded its batch's log rows are normally
|
||||
already in the table. Counting them in the seed and again in the increment
|
||||
is what made a fresh row land at twice the true spend.
|
||||
"""
|
||||
if entity_type == Litellm_EntityType.KEY.value:
|
||||
rows = await prisma_client.db.query_raw(
|
||||
_SEED_FROM_SPEND_LOGS_KEY_SQL, entity_id, window_start, tuple(exclude_request_ids)
|
||||
)
|
||||
elif entity_type == Litellm_EntityType.TEAM.value:
|
||||
rows = await prisma_client.db.query_raw(
|
||||
_SEED_FROM_SPEND_LOGS_TEAM_SQL, entity_id, window_start, tuple(exclude_request_ids)
|
||||
)
|
||||
else:
|
||||
return None
|
||||
if not rows:
|
||||
return 0.0
|
||||
return float(rows[0].get("total") or 0.0)
|
||||
|
||||
|
||||
def _primary_key(transaction: WindowSpendTransaction) -> tuple[str, str, str]:
|
||||
return (
|
||||
transaction["entity_type"],
|
||||
transaction["entity_id"],
|
||||
transaction["window_duration"],
|
||||
)
|
||||
|
||||
|
||||
async def _existing_primary_keys(
|
||||
prisma_client: "PrismaClient",
|
||||
transactions: tuple[WindowSpendTransaction, ...],
|
||||
) -> frozenset[tuple[str, str, str]]:
|
||||
rows: Final = await prisma_client.db.query_raw(
|
||||
_SELECT_EXISTING_ROWS_SQL,
|
||||
tuple(transaction["entity_type"] for transaction in transactions),
|
||||
tuple(transaction["entity_id"] for transaction in transactions),
|
||||
tuple(transaction["window_duration"] for transaction in transactions),
|
||||
)
|
||||
return frozenset((row["entity_type"], row["entity_id"], row["window_duration"]) for row in rows or ())
|
||||
|
||||
|
||||
async def _seed_base_for_missing_row(
|
||||
prisma_client: "PrismaClient",
|
||||
transaction: WindowSpendTransaction,
|
||||
existing_primary_keys: frozenset[tuple[str, str, str]],
|
||||
spend_logs_aggregate: WindowSpendLogsAggregate,
|
||||
) -> float:
|
||||
"""Spend already recorded for a window that has no row yet.
|
||||
|
||||
This is the LiteLLM_SpendLogs aggregate the window counter reseed runs on
|
||||
every cold counter today, but here it runs once per window lifetime and off
|
||||
the request path, and it excludes this batch's own requests so they are
|
||||
counted by their increments alone.
|
||||
"""
|
||||
if _primary_key(transaction) in existing_primary_keys:
|
||||
return 0.0
|
||||
base: Final = await spend_logs_aggregate(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=transaction["entity_type"],
|
||||
entity_id=transaction["entity_id"],
|
||||
window_start=datetime.fromisoformat(transaction["window_start"]).replace(tzinfo=timezone.utc),
|
||||
exclude_request_ids=transaction["request_ids"],
|
||||
)
|
||||
return float(base or 0.0)
|
||||
|
||||
|
||||
def _upsert_params(
|
||||
transaction: WindowSpendTransaction,
|
||||
seed_base: float,
|
||||
now: datetime,
|
||||
) -> tuple[str, str, str, datetime, float, float, datetime]:
|
||||
"""$5 is what a brand new row starts at (pre-existing spend plus this
|
||||
increment); $6 is the increment alone, which is all an already-current row
|
||||
may add. They are equal for every row that already existed, so a row is
|
||||
never seeded twice when two pods flush the same new window."""
|
||||
increment: Final = float(transaction["spend"])
|
||||
return (
|
||||
transaction["entity_type"],
|
||||
transaction["entity_id"],
|
||||
transaction["window_duration"],
|
||||
datetime.fromisoformat(transaction["window_start"]),
|
||||
seed_base + increment,
|
||||
increment,
|
||||
now,
|
||||
)
|
||||
|
||||
|
||||
async def commit_window_spend_updates(
|
||||
prisma_client: "PrismaClient",
|
||||
transactions: Sequence[WindowSpendTransaction],
|
||||
spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_excluding,
|
||||
) -> None:
|
||||
"""Apply aggregated window increments to LiteLLM_BudgetWindowSpend.
|
||||
|
||||
An increment at or behind the row's window_start adds into the row (this is
|
||||
how in-flight requests that raced a reset carry into the new window); an
|
||||
increment ahead of it rolls the window and starts from that increment.
|
||||
|
||||
Statements are ordered by primary key so concurrent pods take row locks in
|
||||
the same order, with window_start breaking ties so an older window is
|
||||
applied before the roll that supersedes it.
|
||||
"""
|
||||
if not transactions:
|
||||
return
|
||||
|
||||
ordered: Final = tuple(sorted(transactions, key=window_spend_group_key))
|
||||
existing_primary_keys: Final = await _existing_primary_keys(
|
||||
prisma_client=prisma_client,
|
||||
transactions=ordered,
|
||||
)
|
||||
seed_bases: Final = tuple(
|
||||
[
|
||||
await _seed_base_for_missing_row(
|
||||
prisma_client=prisma_client,
|
||||
transaction=transaction,
|
||||
existing_primary_keys=existing_primary_keys,
|
||||
spend_logs_aggregate=spend_logs_aggregate,
|
||||
)
|
||||
for transaction in ordered
|
||||
]
|
||||
)
|
||||
|
||||
now: Final = to_naive_utc(datetime.now(timezone.utc))
|
||||
verbose_proxy_logger.debug(
|
||||
"Spend tracking - committing %d budget window spend upserts over %d existing rows",
|
||||
len(ordered),
|
||||
len(existing_primary_keys),
|
||||
)
|
||||
async with prisma_client.db.tx(timeout=_UPSERT_TRANSACTION_TIMEOUT) as db_transaction:
|
||||
async with db_transaction.batch_() as batcher:
|
||||
for transaction, seed_base in zip(ordered, seed_bases):
|
||||
batcher.execute_raw(
|
||||
_UPSERT_WINDOW_SPEND_SQL,
|
||||
*_upsert_params(transaction=transaction, seed_base=seed_base, now=now),
|
||||
)
|
||||
|
||||
|
||||
async def roll_window_spend_row(
|
||||
prisma_client: "PrismaClient",
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_duration: str,
|
||||
new_window_start: datetime,
|
||||
) -> None:
|
||||
"""Move a row onto the window that just started and zero its spend.
|
||||
|
||||
Conditional on the stored window_start still being behind the new one so a
|
||||
pod that already rolled the row (or increments that arrived under the new
|
||||
window) are not clobbered.
|
||||
"""
|
||||
await prisma_client.db.execute_raw(
|
||||
_ROLL_WINDOW_SPEND_SQL,
|
||||
entity_type,
|
||||
entity_id,
|
||||
window_duration,
|
||||
to_naive_utc(new_window_start),
|
||||
to_naive_utc(datetime.now(timezone.utc)),
|
||||
)
|
||||
|
|
@ -12,7 +12,7 @@ import os
|
|||
import random
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
|
||||
|
|
@ -50,6 +50,10 @@ from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdate
|
|||
from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
|
||||
ToolDiscoveryQueue,
|
||||
)
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
WindowSpendTransaction,
|
||||
WindowSpendUpdateQueue,
|
||||
)
|
||||
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
|
||||
from litellm.proxy.spend_tracking.compression_savings import (
|
||||
extract_compression_saved_tokens,
|
||||
|
|
@ -132,6 +136,7 @@ class DBSpendUpdateWriter:
|
|||
self.daily_agent_spend_update_queue = DailySpendUpdateQueue()
|
||||
self.daily_org_spend_update_queue = DailySpendUpdateQueue()
|
||||
self.daily_tag_spend_update_queue = DailySpendUpdateQueue()
|
||||
self.window_spend_update_queue = WindowSpendUpdateQueue()
|
||||
|
||||
async def update_database(
|
||||
# LiteLLM management object fields
|
||||
|
|
@ -147,7 +152,11 @@ class DBSpendUpdateWriter:
|
|||
start_time: datetime | None,
|
||||
end_time: datetime | None,
|
||||
response_cost: float | None,
|
||||
):
|
||||
) -> str | None:
|
||||
"""Returns the LiteLLM_SpendLogs request_id this call was recorded
|
||||
under, so the caller can tell the budget-window writer which log rows
|
||||
its increments already cover. None when the payload could not be built.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
disable_spend_logs,
|
||||
litellm_proxy_budget_name,
|
||||
|
|
@ -164,7 +173,7 @@ class DBSpendUpdateWriter:
|
|||
team_id,
|
||||
)
|
||||
if ProxyUpdateSpend.disable_spend_updates() is True:
|
||||
return
|
||||
return None
|
||||
if token is not None and isinstance(token, str) and token.startswith("sk-"):
|
||||
hashed_token = hash_token(token=token)
|
||||
else:
|
||||
|
|
@ -232,6 +241,7 @@ class DBSpendUpdateWriter:
|
|||
)
|
||||
|
||||
verbose_proxy_logger.debug("Runs spend update on all tables")
|
||||
return payload.get("request_id")
|
||||
except Exception:
|
||||
spend_log_error(
|
||||
"Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue "
|
||||
|
|
@ -244,6 +254,7 @@ class DBSpendUpdateWriter:
|
|||
org_id,
|
||||
end_user_id,
|
||||
)
|
||||
return None
|
||||
|
||||
async def _enqueue_tool_usage_transaction(
|
||||
self,
|
||||
|
|
@ -831,6 +842,7 @@ class DBSpendUpdateWriter:
|
|||
daily_org_spend_update_queue=self.daily_org_spend_update_queue,
|
||||
daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue,
|
||||
daily_agent_spend_update_queue=self.daily_agent_spend_update_queue,
|
||||
window_spend_update_queue=self.window_spend_update_queue,
|
||||
)
|
||||
|
||||
# Only commit from redis to db if this pod is the leader
|
||||
|
|
@ -847,6 +859,7 @@ class DBSpendUpdateWriter:
|
|||
daily_org_spend_update_transactions,
|
||||
daily_end_user_spend_update_transactions,
|
||||
daily_agent_spend_update_transactions,
|
||||
window_spend_update_transactions,
|
||||
) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
if db_spend_update_transactions is not None:
|
||||
|
|
@ -906,6 +919,11 @@ class DBSpendUpdateWriter:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_agent_spend_update_transactions,
|
||||
)
|
||||
if window_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter._commit_window_spend_updates(
|
||||
prisma_client=prisma_client,
|
||||
window_spend_transactions=window_spend_update_transactions,
|
||||
)
|
||||
except Exception as e:
|
||||
spend_log_error(
|
||||
"Spend tracking - failed to commit spend updates from Redis to DB. "
|
||||
|
|
@ -1016,6 +1034,17 @@ class DBSpendUpdateWriter:
|
|||
daily_spend_transactions=daily_agent_spend_update_transactions,
|
||||
)
|
||||
|
||||
################## Budget Window Spend Update Transactions ##################
|
||||
# Aggregate all in memory budget window spend transactions and commit to db
|
||||
window_spend_update_transactions: Final = (
|
||||
await self.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
)
|
||||
|
||||
await DBSpendUpdateWriter._commit_window_spend_updates(
|
||||
prisma_client=prisma_client,
|
||||
window_spend_transactions=window_spend_update_transactions,
|
||||
)
|
||||
|
||||
################## Tool Registry Upserts ##################
|
||||
await self._flush_tool_discovery_queue(prisma_client=prisma_client)
|
||||
|
||||
|
|
@ -1086,6 +1115,36 @@ class DBSpendUpdateWriter:
|
|||
cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _commit_window_spend_updates(
|
||||
prisma_client: PrismaClient,
|
||||
window_spend_transactions: Sequence[WindowSpendTransaction],
|
||||
) -> None:
|
||||
"""
|
||||
Commit per-budget-window spend increments to LiteLLM_BudgetWindowSpend.
|
||||
|
||||
Failures are logged rather than raised: these rows exist so budget
|
||||
enforcement can stop aggregating LiteLLM_SpendLogs, and the read path
|
||||
falls back to that aggregate, so a failed commit must not abort the
|
||||
entity and daily spend commits that share this scheduler tick.
|
||||
"""
|
||||
from litellm.proxy.db.budget_window_spend_writer import (
|
||||
commit_window_spend_updates,
|
||||
)
|
||||
|
||||
try:
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=prisma_client,
|
||||
transactions=window_spend_transactions,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # any DB failure here must stay contained to the window rows
|
||||
spend_log_error(
|
||||
"Spend tracking - failed to commit budget window spend updates. %d window increments lost. Error: %s",
|
||||
len(window_spend_transactions),
|
||||
str(e),
|
||||
exc=e,
|
||||
)
|
||||
|
||||
async def _flush_tool_discovery_queue(
|
||||
self,
|
||||
prisma_client: PrismaClient,
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.constants import (
|
|||
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_UPDATE_BUFFER_KEY,
|
||||
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -38,6 +39,10 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
|
|||
DailySpendUpdateQueue,
|
||||
)
|
||||
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
WindowSpendTransaction,
|
||||
WindowSpendUpdateQueue,
|
||||
)
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.caching import (
|
||||
RedisPipelineLpopOperation,
|
||||
|
|
@ -130,6 +135,7 @@ class RedisUpdateBuffer:
|
|||
daily_org_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_agent_spend_update_queue: DailySpendUpdateQueue,
|
||||
window_spend_update_queue: WindowSpendUpdateQueue | None = None,
|
||||
):
|
||||
"""
|
||||
Stores the in-memory spend updates to Redis
|
||||
|
|
@ -196,6 +202,11 @@ class RedisUpdateBuffer:
|
|||
daily_agent_spend_update_transactions: Final = (
|
||||
await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
)
|
||||
window_spend_update_transactions: Final = (
|
||||
await window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
if window_spend_update_queue is not None
|
||||
else ()
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions)
|
||||
verbose_proxy_logger.debug("ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions)
|
||||
|
|
@ -232,6 +243,11 @@ class RedisUpdateBuffer:
|
|||
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
|
||||
ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE,
|
||||
),
|
||||
(
|
||||
window_spend_update_transactions,
|
||||
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
|
||||
ServiceTypes.REDIS_WINDOW_SPEND_UPDATE_QUEUE,
|
||||
),
|
||||
]
|
||||
|
||||
rpush_list: Final[list[RedisPipelineRpushOperation]] = []
|
||||
|
|
@ -272,12 +288,14 @@ class RedisUpdateBuffer:
|
|||
daily_org_spend_update_transactions=daily_org_spend_update_transactions,
|
||||
daily_end_user_spend_update_transactions=daily_end_user_spend_update_transactions,
|
||||
daily_agent_spend_update_transactions=daily_agent_spend_update_transactions,
|
||||
window_spend_update_transactions=window_spend_update_transactions,
|
||||
spend_update_queue=spend_update_queue,
|
||||
daily_spend_update_queue=daily_spend_update_queue,
|
||||
daily_team_spend_update_queue=daily_team_spend_update_queue,
|
||||
daily_org_spend_update_queue=daily_org_spend_update_queue,
|
||||
daily_end_user_spend_update_queue=daily_end_user_spend_update_queue,
|
||||
daily_agent_spend_update_queue=daily_agent_spend_update_queue,
|
||||
window_spend_update_queue=window_spend_update_queue,
|
||||
)
|
||||
return
|
||||
|
||||
|
|
@ -297,12 +315,14 @@ class RedisUpdateBuffer:
|
|||
daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None,
|
||||
daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None,
|
||||
daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None,
|
||||
window_spend_update_transactions: tuple[WindowSpendTransaction, ...] | None,
|
||||
spend_update_queue: SpendUpdateQueue,
|
||||
daily_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_team_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_org_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_agent_spend_update_queue: DailySpendUpdateQueue,
|
||||
window_spend_update_queue: WindowSpendUpdateQueue | None,
|
||||
) -> None:
|
||||
"""
|
||||
Put drained-but-unpushed transactions back into in-memory queues.
|
||||
|
|
@ -372,6 +392,9 @@ class RedisUpdateBuffer:
|
|||
if daily_txns:
|
||||
await daily_queue.update_queue.put(daily_txns)
|
||||
|
||||
if window_spend_update_transactions and window_spend_update_queue is not None:
|
||||
await window_spend_update_queue.update_queue.put(window_spend_update_transactions)
|
||||
|
||||
@staticmethod
|
||||
def _number_of_transactions_to_store_in_redis(
|
||||
db_spend_update_transactions: DBSpendUpdateTransactions,
|
||||
|
|
@ -468,20 +491,22 @@ class RedisUpdateBuffer:
|
|||
dict[str, DailyOrganizationSpendTransaction] | None,
|
||||
dict[str, DailyEndUserSpendTransaction] | None,
|
||||
dict[str, DailyAgentSpendTransaction] | None,
|
||||
tuple[WindowSpendTransaction, ...] | None,
|
||||
]:
|
||||
"""
|
||||
Drains the main 6 Redis buffer queues in a single pipeline round-trip.
|
||||
Drains the main 7 Redis buffer queues in a single pipeline round-trip.
|
||||
|
||||
Returns a 6-tuple of parsed results in this order:
|
||||
Returns a 7-tuple of parsed results in this order:
|
||||
0: DBSpendUpdateTransactions
|
||||
1: daily user spend
|
||||
2: daily team spend
|
||||
3: daily org spend
|
||||
4: daily end-user spend
|
||||
5: daily agent spend
|
||||
6: budget window spend
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
return None, None, None, None, None, None
|
||||
return None, None, None, None, None, None, None
|
||||
|
||||
lpop_list: Final[list[RedisPipelineLpopOperation]] = [
|
||||
RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
|
||||
|
|
@ -505,12 +530,16 @@ class RedisUpdateBuffer:
|
|||
key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
|
||||
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
|
||||
),
|
||||
RedisPipelineLpopOperation(
|
||||
key=REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
|
||||
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
|
||||
),
|
||||
]
|
||||
|
||||
raw_results: Final = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
# Pad with None if pipeline returned fewer results than expected
|
||||
while len(raw_results) < 6:
|
||||
while len(raw_results) < 7:
|
||||
raw_results.append(None)
|
||||
|
||||
# Slot 0: DBSpendUpdateTransactions
|
||||
|
|
@ -531,6 +560,14 @@ class RedisUpdateBuffer:
|
|||
aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily)
|
||||
daily_results.append(aggregated)
|
||||
|
||||
window_spend: Final = (
|
||||
WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(
|
||||
tuple(json.loads(transaction) for transaction in raw_results[6])
|
||||
)
|
||||
if raw_results[6] is not None
|
||||
else None
|
||||
)
|
||||
|
||||
return (
|
||||
db_spend,
|
||||
cast(dict[str, DailyUserSpendTransaction] | None, daily_results[0]),
|
||||
|
|
@ -538,6 +575,7 @@ class RedisUpdateBuffer:
|
|||
cast(dict[str, DailyOrganizationSpendTransaction] | None, daily_results[2]),
|
||||
cast(dict[str, DailyEndUserSpendTransaction] | None, daily_results[3]),
|
||||
cast(dict[str, DailyAgentSpendTransaction] | None, daily_results[4]),
|
||||
window_spend,
|
||||
)
|
||||
|
||||
async def store_in_memory_daily_tag_spend_updates_in_redis(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,160 @@
|
|||
"""
|
||||
In memory buffer for per-budget-window spend increments.
|
||||
|
||||
Kept separate from SpendUpdateQueue: an increment is only meaningful together
|
||||
with the window it landed in, so two increments for the same entity must not be
|
||||
merged when their window_start differs.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime, timezone
|
||||
from itertools import chain, groupby
|
||||
from typing import Final, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE
|
||||
from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue
|
||||
|
||||
|
||||
class WindowSpendTransaction(TypedDict):
|
||||
"""One increment for a single (entity, budget window) pair.
|
||||
|
||||
window_start is an ISO-8601 string rather than a datetime so the
|
||||
transaction survives the JSON round trip through the Redis buffer.
|
||||
|
||||
request_ids carries the LiteLLM_SpendLogs ids this spend came from. The
|
||||
one-time seed for a window that has no row yet subtracts them from its
|
||||
LiteLLM_SpendLogs aggregate, because the spend log writer flushes on its
|
||||
own ~2s poll and will usually have persisted these rows before the window
|
||||
queue flushes; without the exclusion the seed and the increment would each
|
||||
count them.
|
||||
"""
|
||||
|
||||
entity_type: str
|
||||
entity_id: str
|
||||
window_duration: str
|
||||
window_start: str
|
||||
spend: float
|
||||
request_ids: Sequence[str]
|
||||
|
||||
|
||||
def to_naive_utc(value: datetime) -> datetime:
|
||||
"""LiteLLM_BudgetWindowSpend.window_start is TIMESTAMP(3), which holds naive UTC."""
|
||||
if value.tzinfo is None:
|
||||
return value
|
||||
return value.astimezone(timezone.utc).replace(tzinfo=None)
|
||||
|
||||
|
||||
def window_spend_group_key(transaction: WindowSpendTransaction) -> tuple[str, str, str, str]:
|
||||
"""Identity of a window increment: the row's primary key plus the window it
|
||||
belongs to. Two increments only aggregate when all four match."""
|
||||
return (
|
||||
transaction["entity_type"],
|
||||
transaction["entity_id"],
|
||||
transaction["window_duration"],
|
||||
transaction["window_start"],
|
||||
)
|
||||
|
||||
|
||||
def build_window_spend_transaction(
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_duration: str,
|
||||
window_start: datetime,
|
||||
spend: float,
|
||||
request_id: str | None = None,
|
||||
) -> WindowSpendTransaction:
|
||||
return WindowSpendTransaction(
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=to_naive_utc(window_start).isoformat(timespec="microseconds"),
|
||||
spend=spend,
|
||||
request_ids=() if request_id is None else (request_id,),
|
||||
)
|
||||
|
||||
|
||||
def _merge_window_spend_transactions(
|
||||
payloads: tuple[WindowSpendTransaction, ...],
|
||||
) -> WindowSpendTransaction:
|
||||
first: Final = payloads[0]
|
||||
return WindowSpendTransaction(
|
||||
entity_type=first["entity_type"],
|
||||
entity_id=first["entity_id"],
|
||||
window_duration=first["window_duration"],
|
||||
window_start=first["window_start"],
|
||||
spend=math.fsum(payload["spend"] for payload in payloads),
|
||||
request_ids=tuple(sorted(frozenset(chain.from_iterable(payload["request_ids"] for payload in payloads)))),
|
||||
)
|
||||
|
||||
|
||||
class WindowSpendUpdateQueue(BaseUpdateQueue):
|
||||
"""
|
||||
In memory buffer for budget-window spend increments committed to
|
||||
LiteLLM_BudgetWindowSpend.
|
||||
|
||||
Add an update with the payload built by build_window_spend_transaction:
|
||||
window_spend_update_queue.add_update(
|
||||
build_window_spend_transaction(
|
||||
entity_type="key",
|
||||
entity_id="<hashed token>",
|
||||
window_duration="30d",
|
||||
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
spend=0.02,
|
||||
)
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.update_queue: asyncio.Queue[tuple[WindowSpendTransaction, ...]] = asyncio.Queue(
|
||||
maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE
|
||||
)
|
||||
|
||||
async def add_update(self, update: WindowSpendTransaction) -> None:
|
||||
"""Enqueue an update."""
|
||||
verbose_proxy_logger.debug("Adding budget window spend update to queue: %s", update)
|
||||
await self.update_queue.put((update,))
|
||||
if self.update_queue.qsize() >= self.MAX_SIZE_IN_MEMORY_QUEUE:
|
||||
verbose_proxy_logger.warning(
|
||||
"Budget window spend update queue is full. Aggregating all entries in queue to concatenate entries."
|
||||
)
|
||||
await self.aggregate_queue_updates()
|
||||
|
||||
async def aggregate_queue_updates(self) -> None:
|
||||
"""Collapse everything currently queued into a single aggregated update."""
|
||||
updates: Final = await self.flush_all_updates_from_in_memory_queue()
|
||||
await self.update_queue.put(WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates))
|
||||
|
||||
async def flush_and_get_aggregated_window_spend_transactions(
|
||||
self,
|
||||
) -> tuple[WindowSpendTransaction, ...]:
|
||||
"""Drain the queue and return the increments aggregated per window."""
|
||||
updates: Final = await self.flush_all_updates_from_in_memory_queue()
|
||||
if len(updates) > 0:
|
||||
verbose_proxy_logger.info(
|
||||
"Spend tracking - flushed %d budget window spend update batches from in-memory queue",
|
||||
len(updates),
|
||||
)
|
||||
return WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates)
|
||||
|
||||
@staticmethod
|
||||
def get_aggregated_window_spend_transactions(
|
||||
updates: Sequence[Sequence[WindowSpendTransaction]],
|
||||
) -> tuple[WindowSpendTransaction, ...]:
|
||||
"""Sum spend per (entity_type, entity_id, window_duration, window_start).
|
||||
|
||||
Increments belonging to different windows stay separate even when they
|
||||
share a primary key, so a window boundary crossed mid-tick does not fold
|
||||
the new window's spend into the previous window's total.
|
||||
|
||||
The result is ordered by that same key, which is the order the flush
|
||||
needs: primary key first for cross-pod lock ordering, then window_start
|
||||
so an older window is applied before the roll that supersedes it.
|
||||
"""
|
||||
ordered: Final = tuple(sorted(chain.from_iterable(updates), key=window_spend_group_key))
|
||||
return tuple(
|
||||
_merge_window_spend_transactions(tuple(group)) for _, group in groupby(ordered, key=window_spend_group_key)
|
||||
)
|
||||
|
|
@ -498,7 +498,7 @@ async def _update_database_and_spend_counters(
|
|||
request_tags: list[str] | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
spend_log_request_id = await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
token=user_api_key,
|
||||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
|
|
@ -534,6 +534,7 @@ async def _update_database_and_spend_counters(
|
|||
budget_reservation=budget_reservation,
|
||||
end_user_id=end_user_id,
|
||||
tags=request_tags,
|
||||
request_id=spend_log_request_id,
|
||||
)
|
||||
except Exception:
|
||||
if budget_reservation is not None:
|
||||
|
|
|
|||
|
|
@ -349,6 +349,9 @@ from litellm.proxy.config_resolvers.alerting import (
|
|||
from litellm.proxy.container_endpoints.endpoints import router as container_router
|
||||
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
|
||||
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
build_window_spend_transaction,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import (
|
||||
PrismaDBExceptionHandler,
|
||||
call_with_db_reconnect_retry,
|
||||
|
|
@ -2392,6 +2395,7 @@ async def increment_spend_counters(
|
|||
budget_reservation: dict | None = None,
|
||||
end_user_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
request_id: str | None = None,
|
||||
):
|
||||
"""
|
||||
Atomically increment spend counters for budget enforcement.
|
||||
|
|
@ -2446,14 +2450,24 @@ async def increment_spend_counters(
|
|||
for window in key_budget_limits:
|
||||
duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration
|
||||
key_window_counter = f"spend:key:{hashed_token}:window:{duration}"
|
||||
key_window_start = get_budget_window_start(window)
|
||||
if key_window_counter not in reserved_counter_keys:
|
||||
await _init_and_increment_window_spend_counter(
|
||||
counter_key=key_window_counter,
|
||||
entity_type="Key",
|
||||
entity_id=hashed_token,
|
||||
window_start=get_budget_window_start(window),
|
||||
window_start=key_window_start,
|
||||
increment=cost,
|
||||
)
|
||||
await _enqueue_window_spend_row_update(
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id=hashed_token,
|
||||
window=window,
|
||||
window_duration=duration,
|
||||
window_start=key_window_start,
|
||||
increment=cost,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
async def _team_scope(scope_team_id: str) -> None:
|
||||
team_counter_key: Final = f"spend:team:{scope_team_id}"
|
||||
|
|
@ -2477,14 +2491,24 @@ async def increment_spend_counters(
|
|||
for window in team_budget_limits:
|
||||
duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration
|
||||
team_window_counter = f"spend:team:{scope_team_id}:window:{duration}"
|
||||
team_window_start = get_budget_window_start(window)
|
||||
if team_window_counter not in reserved_counter_keys:
|
||||
await _init_and_increment_window_spend_counter(
|
||||
counter_key=team_window_counter,
|
||||
entity_type="Team",
|
||||
entity_id=scope_team_id,
|
||||
window_start=get_budget_window_start(window),
|
||||
window_start=team_window_start,
|
||||
increment=cost,
|
||||
)
|
||||
await _enqueue_window_spend_row_update(
|
||||
entity_type=Litellm_EntityType.TEAM,
|
||||
entity_id=scope_team_id,
|
||||
window=window,
|
||||
window_duration=duration,
|
||||
window_start=team_window_start,
|
||||
increment=cost,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None:
|
||||
team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}"
|
||||
|
|
@ -2669,6 +2693,58 @@ async def _init_and_increment_spend_counter(
|
|||
await _increment_spend_counter_cache(counter_key=counter_key, increment=increment)
|
||||
|
||||
|
||||
async def _enqueue_window_spend_row_update(
|
||||
entity_type: Litellm_EntityType,
|
||||
entity_id: str,
|
||||
window: Any,
|
||||
window_duration: str,
|
||||
window_start: datetime | None,
|
||||
increment: float,
|
||||
request_id: str | None,
|
||||
) -> None:
|
||||
"""Queue this request's cost against the LiteLLM_BudgetWindowSpend row for
|
||||
the window, so enforcement can read a maintained total instead of
|
||||
aggregating LiteLLM_SpendLogs.
|
||||
|
||||
request_id is the LiteLLM_SpendLogs id this cost was recorded under; the
|
||||
flush uses it to keep the one-time seed from counting a request that its
|
||||
increment already covers.
|
||||
|
||||
Enqueued even when the cache increment was skipped for a reserved counter:
|
||||
the reservation only pre-charged the counter, and the row still owes the
|
||||
actual cost.
|
||||
|
||||
Windows with no reset_at slide with wall clock, so their window_start moves
|
||||
on every request and no single row can represent them. Those are left to
|
||||
the read path's LiteLLM_SpendLogs fallback rather than rewritten per
|
||||
request.
|
||||
"""
|
||||
if window_start is None:
|
||||
return
|
||||
reset_at: Final = window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None)
|
||||
if not reset_at:
|
||||
return
|
||||
try:
|
||||
await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update(
|
||||
build_window_spend_transaction(
|
||||
entity_type=entity_type.value,
|
||||
entity_id=entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
spend=increment,
|
||||
request_id=request_id,
|
||||
)
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # spend tracking must never fail the cost callback
|
||||
verbose_proxy_logger.debug(
|
||||
"Unable to enqueue budget window spend update for %s=%s window=%s: %s",
|
||||
entity_type.value,
|
||||
entity_id,
|
||||
window_duration,
|
||||
e,
|
||||
)
|
||||
|
||||
|
||||
async def _init_and_increment_window_spend_counter(
|
||||
counter_key: str,
|
||||
entity_type: str,
|
||||
|
|
|
|||
|
|
@ -42,6 +42,8 @@ class ServiceTypes(str, enum.Enum):
|
|||
# spend update queue - current spend of key, user, team
|
||||
IN_MEMORY_SPEND_UPDATE_QUEUE = "in_memory_spend_update_queue"
|
||||
REDIS_SPEND_UPDATE_QUEUE = "redis_spend_update_queue"
|
||||
# budget window spend queue - per-window spend of key, team
|
||||
REDIS_WINDOW_SPEND_UPDATE_QUEUE = "redis_window_spend_update_queue"
|
||||
|
||||
|
||||
class ServiceConfig(TypedDict):
|
||||
|
|
|
|||
|
|
@ -1082,6 +1082,7 @@ def _make_reset_budget_windows_job(
|
|||
raise AssertionError(f"Unexpected query_raw call: {query}")
|
||||
|
||||
prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw)
|
||||
prisma_client.db.execute_raw = AsyncMock(return_value=1)
|
||||
prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None)
|
||||
prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None)
|
||||
|
||||
|
|
@ -1153,6 +1154,145 @@ def test_reset_budget_windows_resets_expired_key_window(monkeypatch):
|
|||
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0)
|
||||
|
||||
|
||||
def _window_spend_rolls(prisma_client):
|
||||
return [
|
||||
call.args
|
||||
for call in prisma_client.db.execute_raw.await_args_list
|
||||
if "LiteLLM_BudgetWindowSpend" in call.args[0]
|
||||
]
|
||||
|
||||
|
||||
def test_reset_budget_windows_rolls_the_key_window_spend_row(monkeypatch):
|
||||
"""The maintained per-window total has to start the new window at zero
|
||||
alongside the counter, or enforcement keeps reading the old window's spend."""
|
||||
now = datetime.utcnow()
|
||||
expired = (now - timedelta(minutes=5)).isoformat() + "Z"
|
||||
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-expired",
|
||||
"budget_limits": [{"budget_duration": "1d", "reset_at": expired}],
|
||||
}
|
||||
]
|
||||
job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[])
|
||||
|
||||
asyncio.run(job.reset_budget_windows())
|
||||
|
||||
rolls = _window_spend_rolls(prisma_client)
|
||||
assert len(rolls) == 1
|
||||
query, entity_type, entity_id, window_duration, new_window_start, _updated_at = rolls[0]
|
||||
assert (entity_type, entity_id, window_duration) == ("key", "sk-expired", "1d")
|
||||
assert "spend = 0" in " ".join(query.split())
|
||||
|
||||
# window_start is the start of the window that just began: new reset_at minus the duration.
|
||||
written_windows = json.loads(
|
||||
prisma_client.db.litellm_verificationtoken.update.await_args.kwargs["data"]["budget_limits"]
|
||||
)
|
||||
new_reset_at = datetime.fromisoformat(written_windows[0]["reset_at"].replace("Z", "+00:00")).replace(tzinfo=None)
|
||||
assert new_window_start == pytest.approx(
|
||||
new_reset_at - timedelta(days=1),
|
||||
abs=timedelta(seconds=1),
|
||||
)
|
||||
|
||||
|
||||
def test_reset_budget_windows_roll_is_conditional_on_an_older_stored_window(monkeypatch):
|
||||
"""Another pod may already have rolled the row; clobbering it would drop
|
||||
spend that landed under the new window."""
|
||||
now = datetime.utcnow()
|
||||
expired = (now - timedelta(minutes=5)).isoformat() + "Z"
|
||||
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-expired",
|
||||
"budget_limits": [{"budget_duration": "1d", "reset_at": expired}],
|
||||
}
|
||||
]
|
||||
job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[])
|
||||
|
||||
asyncio.run(job.reset_budget_windows())
|
||||
|
||||
query = " ".join(_window_spend_rolls(prisma_client)[0][0].split())
|
||||
assert "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" in query
|
||||
|
||||
|
||||
def test_reset_budget_windows_rolls_the_team_window_spend_row(monkeypatch):
|
||||
now = datetime.utcnow()
|
||||
expired = (now - timedelta(minutes=1)).isoformat() + "Z"
|
||||
|
||||
team_rows = [
|
||||
{
|
||||
"team_id": "team-expired",
|
||||
"budget_limits": [{"budget_duration": "30d", "reset_at": expired}],
|
||||
}
|
||||
]
|
||||
job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=[], team_rows=team_rows)
|
||||
|
||||
asyncio.run(job.reset_budget_windows())
|
||||
|
||||
rolls = _window_spend_rolls(prisma_client)
|
||||
assert len(rolls) == 1
|
||||
assert rolls[0][1:4] == ("team", "team-expired", "30d")
|
||||
|
||||
|
||||
def test_reset_budget_windows_does_not_roll_an_unexpired_window(monkeypatch):
|
||||
now = datetime.utcnow()
|
||||
future = (now + timedelta(hours=1)).isoformat() + "Z"
|
||||
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-future",
|
||||
"budget_limits": [{"budget_duration": "1d", "reset_at": future}],
|
||||
}
|
||||
]
|
||||
job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[])
|
||||
|
||||
asyncio.run(job.reset_budget_windows())
|
||||
|
||||
assert _window_spend_rolls(prisma_client) == []
|
||||
|
||||
|
||||
def test_reset_budget_windows_rolls_only_the_expired_window_of_a_key(monkeypatch):
|
||||
now = datetime.utcnow()
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-mixed",
|
||||
"budget_limits": [
|
||||
{"budget_duration": "1d", "reset_at": (now - timedelta(minutes=5)).isoformat() + "Z"},
|
||||
{"budget_duration": "30d", "reset_at": (now + timedelta(days=2)).isoformat() + "Z"},
|
||||
],
|
||||
}
|
||||
]
|
||||
job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[])
|
||||
|
||||
asyncio.run(job.reset_budget_windows())
|
||||
|
||||
rolls = _window_spend_rolls(prisma_client)
|
||||
assert [roll[3] for roll in rolls] == ["1d"]
|
||||
|
||||
|
||||
def test_reset_budget_windows_survives_a_failed_window_spend_roll(monkeypatch):
|
||||
"""The row is an optimization over aggregating LiteLLM_SpendLogs; a DB
|
||||
failure there must not stop the counter reset from being persisted."""
|
||||
now = datetime.utcnow()
|
||||
expired = (now - timedelta(minutes=5)).isoformat() + "Z"
|
||||
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-expired",
|
||||
"budget_limits": [{"budget_duration": "1d", "reset_at": expired}],
|
||||
}
|
||||
]
|
||||
job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job(
|
||||
monkeypatch, key_rows=key_rows, team_rows=[]
|
||||
)
|
||||
prisma_client.db.execute_raw = AsyncMock(side_effect=Exception("connection reset"))
|
||||
|
||||
asyncio.run(job.reset_budget_windows())
|
||||
|
||||
prisma_client.db.litellm_verificationtoken.update.assert_awaited_once()
|
||||
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0)
|
||||
|
||||
|
||||
def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch):
|
||||
"""If `reset_at` is in the future, no write should happen for that key."""
|
||||
now = datetime.utcnow()
|
||||
|
|
|
|||
|
|
@ -208,7 +208,8 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
|
|||
Verify get_all_transactions_from_redis_buffer_pipeline correctly parses
|
||||
and aggregates results from async_lpop_pipeline.
|
||||
"""
|
||||
# Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories
|
||||
# Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories,
|
||||
# slot 6 = budget window spend
|
||||
db_spend_json = json.dumps(
|
||||
{
|
||||
"key_list_transactions": {"key1": 1.0, "key2": 2.0},
|
||||
|
|
@ -222,6 +223,18 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
|
|||
)
|
||||
daily_user_json = json.dumps({"user_key1": {"spend": 1.0, "api_requests": 1}})
|
||||
daily_team_json = json.dumps({"team_key1": {"spend": 2.0, "api_requests": 2}})
|
||||
window_spend_json = json.dumps(
|
||||
[
|
||||
{
|
||||
"entity_type": "key",
|
||||
"entity_id": "hashed-token",
|
||||
"window_duration": "30d",
|
||||
"window_start": "2026-08-01T00:00:00.000000",
|
||||
"spend": 3.0,
|
||||
"request_ids": ["req-1"],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
mock_redis_cache.async_lpop_pipeline = AsyncMock(
|
||||
return_value=[
|
||||
|
|
@ -231,13 +244,30 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
|
|||
None, # slot 3: daily org (empty)
|
||||
None, # slot 4: daily end-user (empty)
|
||||
None, # slot 5: daily agent (empty)
|
||||
[window_spend_json, window_spend_json], # slot 6: budget window spend
|
||||
]
|
||||
)
|
||||
|
||||
result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
assert len(result) == 6
|
||||
db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent = result
|
||||
assert len(result) == 7
|
||||
(
|
||||
db_spend,
|
||||
daily_user,
|
||||
daily_team,
|
||||
daily_org,
|
||||
daily_end_user,
|
||||
daily_agent,
|
||||
window_spend,
|
||||
) = result
|
||||
|
||||
# Budget window spend from two pods is summed per window, not overwritten,
|
||||
# and both pods' request ids reach the seed exclusion.
|
||||
assert window_spend is not None
|
||||
assert len(window_spend) == 1
|
||||
assert window_spend[0]["spend"] == 6.0
|
||||
assert window_spend[0]["entity_id"] == "hashed-token"
|
||||
assert window_spend[0]["request_ids"] == ("req-1",)
|
||||
|
||||
# Verify db spend was parsed correctly
|
||||
assert db_spend is not None
|
||||
|
|
@ -260,6 +290,12 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
|
|||
|
||||
# Verify pipeline was called once with correct keys
|
||||
mock_redis_cache.async_lpop_pipeline.assert_called_once()
|
||||
from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
|
||||
|
||||
popped_keys = [
|
||||
op["key"] for op in mock_redis_cache.async_lpop_pipeline.call_args.kwargs["lpop_list"]
|
||||
]
|
||||
assert popped_keys[6] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -267,7 +303,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
|
|||
"""When redis_cache is None, should return all Nones"""
|
||||
buffer = RedisUpdateBuffer(redis_cache=None)
|
||||
result = await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
assert result == (None, None, None, None, None, None)
|
||||
assert result == (None, None, None, None, None, None, None)
|
||||
|
||||
|
||||
def test_validate_redis_transaction_buffer_raises_without_redis():
|
||||
|
|
@ -374,3 +410,106 @@ def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch):
|
|||
|
||||
mock_redis_cache.assert_called_once()
|
||||
assert result is mock_redis_cache.return_value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_pushes_budget_window_spend(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
"""The budget window queue has to ride the same rpush as the daily queues,
|
||||
otherwise multi-pod deployments never persist per-window spend."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
WindowSpendUpdateQueue,
|
||||
build_window_spend_transaction,
|
||||
)
|
||||
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1])
|
||||
|
||||
empty_queue = AsyncMock()
|
||||
empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={})
|
||||
empty_daily_queue = AsyncMock()
|
||||
empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={})
|
||||
|
||||
window_queue = WindowSpendUpdateQueue()
|
||||
await window_queue.add_update(
|
||||
build_window_spend_transaction(
|
||||
entity_type="key",
|
||||
entity_id="hashed-token",
|
||||
window_duration="30d",
|
||||
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
spend=1.25,
|
||||
request_id="req-1",
|
||||
)
|
||||
)
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=empty_queue,
|
||||
daily_spend_update_queue=empty_daily_queue,
|
||||
daily_team_spend_update_queue=empty_daily_queue,
|
||||
daily_org_spend_update_queue=empty_daily_queue,
|
||||
daily_end_user_spend_update_queue=empty_daily_queue,
|
||||
daily_agent_spend_update_queue=empty_daily_queue,
|
||||
window_spend_update_queue=window_queue,
|
||||
)
|
||||
|
||||
rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"]
|
||||
assert len(rpush_list) == 1
|
||||
assert rpush_list[0]["key"] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
|
||||
pushed = json.loads(rpush_list[0]["values"][0])
|
||||
assert pushed == [{
|
||||
"entity_type": "key",
|
||||
"entity_id": "hashed-token",
|
||||
"window_duration": "30d",
|
||||
"window_start": "2026-08-01T00:00:00.000000",
|
||||
"spend": 1.25,
|
||||
"request_ids": ["req-1"],
|
||||
}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpush_failure(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
"""The window queue is drained before the rpush, so a Redis hiccup would
|
||||
silently drop per-window spend without the restore."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
WindowSpendUpdateQueue,
|
||||
build_window_spend_transaction,
|
||||
)
|
||||
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away"))
|
||||
|
||||
empty_queue = AsyncMock()
|
||||
empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={})
|
||||
empty_daily_queue = AsyncMock()
|
||||
empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={})
|
||||
|
||||
window_queue = WindowSpendUpdateQueue()
|
||||
await window_queue.add_update(
|
||||
build_window_spend_transaction(
|
||||
entity_type="team",
|
||||
entity_id="team-1",
|
||||
window_duration="7d",
|
||||
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
spend=4.0,
|
||||
)
|
||||
)
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=empty_queue,
|
||||
daily_spend_update_queue=empty_daily_queue,
|
||||
daily_team_spend_update_queue=empty_daily_queue,
|
||||
daily_org_spend_update_queue=empty_daily_queue,
|
||||
daily_end_user_spend_update_queue=empty_daily_queue,
|
||||
daily_agent_spend_update_queue=empty_daily_queue,
|
||||
window_spend_update_queue=window_queue,
|
||||
)
|
||||
|
||||
restored = await window_queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
assert [payload["spend"] for payload in restored] == [4.0]
|
||||
assert [payload["entity_id"] for payload in restored] == ["team-1"]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,249 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
WindowSpendUpdateQueue,
|
||||
build_window_spend_transaction,
|
||||
to_naive_utc,
|
||||
)
|
||||
|
||||
WINDOW_A = datetime(2026, 8, 1, tzinfo=timezone.utc)
|
||||
WINDOW_B = datetime(2026, 8, 31, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def _txn(
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
spend: float,
|
||||
duration: str = "30d",
|
||||
entity_type: str = "key",
|
||||
request_id: str | None = None,
|
||||
):
|
||||
return build_window_spend_transaction(
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
window_duration=duration,
|
||||
window_start=window_start,
|
||||
spend=spend,
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
|
||||
def test_build_window_spend_transaction_stores_naive_utc_iso():
|
||||
"""window_start rides the Redis buffer as a string and lands in a naive-UTC
|
||||
TIMESTAMP(3) column, so a non-UTC input must be converted, not truncated."""
|
||||
non_utc = datetime(2026, 8, 1, 20, 0, tzinfo=timezone(timedelta(hours=-4)))
|
||||
|
||||
assert _txn("k1", non_utc, 1.0, request_id="req-1") == {
|
||||
"entity_type": "key",
|
||||
"entity_id": "k1",
|
||||
"window_duration": "30d",
|
||||
"window_start": "2026-08-02T00:00:00.000000",
|
||||
"spend": 1.0,
|
||||
"request_ids": ("req-1",),
|
||||
}
|
||||
|
||||
|
||||
def test_to_naive_utc_leaves_naive_values_alone():
|
||||
naive = datetime(2026, 8, 1, 12, 0)
|
||||
assert to_naive_utc(naive) == naive
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregation_sums_increments_within_one_window():
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.5))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 2.25))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert len(aggregated) == 1
|
||||
assert aggregated[0]["spend"] == pytest.approx(3.75)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregation_keeps_different_windows_of_same_entity_separate():
|
||||
"""Merging across windows would fold spend from a window that already
|
||||
rolled into the new window's total, over-counting the new window."""
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0))
|
||||
await queue.add_update(_txn("k1", WINDOW_B, 2.0))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert len(aggregated) == 2
|
||||
assert {payload["window_start"]: payload["spend"] for payload in aggregated} == {
|
||||
"2026-08-01T00:00:00.000000": 1.0,
|
||||
"2026-08-31T00:00:00.000000": 2.0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregation_keeps_durations_entities_and_types_separate():
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, duration="30d"))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 2.0, duration="7d"))
|
||||
await queue.add_update(_txn("k2", WINDOW_A, 4.0, duration="30d"))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 8.0, duration="30d", entity_type="team"))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert len(aggregated) == 4
|
||||
assert sorted(payload["spend"] for payload in aggregated) == [1.0, 2.0, 4.0, 8.0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregation_orders_by_primary_key_then_window_start():
|
||||
"""The flush relies on this order: primary key first for cross-pod lock
|
||||
ordering, then window_start so an older window is applied before the roll
|
||||
that supersedes it."""
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("t1", WINDOW_A, 1.0, entity_type="team"))
|
||||
await queue.add_update(_txn("k2", WINDOW_B, 1.0))
|
||||
await queue.add_update(_txn("k2", WINDOW_A, 1.0))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, duration="7d"))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert [
|
||||
(payload["entity_type"], payload["entity_id"], payload["window_duration"], payload["window_start"])
|
||||
for payload in aggregated
|
||||
] == [
|
||||
("key", "k1", "7d", "2026-08-01T00:00:00.000000"),
|
||||
("key", "k2", "30d", "2026-08-01T00:00:00.000000"),
|
||||
("key", "k2", "30d", "2026-08-31T00:00:00.000000"),
|
||||
("team", "t1", "30d", "2026-08-01T00:00:00.000000"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregation_does_not_collide_on_entity_ids_containing_a_separator():
|
||||
"""entity_id is free-form (team ids are user supplied), so grouping must not
|
||||
depend on a flattened string key."""
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("a:30d:2026-08-01T00:00:00.000000:b", WINDOW_A, 1.0))
|
||||
await queue.add_update(_txn("b", WINDOW_A, 2.0))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert len(aggregated) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_empties_the_queue():
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0))
|
||||
|
||||
assert await queue.flush_and_get_aggregated_window_spend_transactions() != ()
|
||||
assert await queue.flush_and_get_aggregated_window_spend_transactions() == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregate_queue_updates_collapses_in_place():
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 2.0))
|
||||
await queue.add_update(_txn("k1", WINDOW_B, 4.0))
|
||||
|
||||
await queue.aggregate_queue_updates()
|
||||
|
||||
assert queue.update_queue.qsize() == 1
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
assert sorted(payload["spend"] for payload in aggregated) == [3.0, 4.0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregation_does_not_mutate_the_queued_payloads():
|
||||
"""The same payload can be re-aggregated after a failed Redis push, so
|
||||
aggregation must not accumulate into the caller's object."""
|
||||
queue = WindowSpendUpdateQueue()
|
||||
update = _txn("k1", WINDOW_A, 1.0)
|
||||
await queue.add_update(update)
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 2.0))
|
||||
|
||||
await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert update["spend"] == 1.0
|
||||
|
||||
|
||||
def test_aggregation_survives_the_redis_json_round_trip():
|
||||
"""The Redis buffer stores transactions as JSON, so the aggregated shape
|
||||
must reload into an equivalent aggregation."""
|
||||
aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(
|
||||
[(_txn("k1", WINDOW_A, 1.0),), (_txn("k1", WINDOW_B, 2.0),)]
|
||||
)
|
||||
|
||||
reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(
|
||||
[json.loads(json.dumps(aggregated))]
|
||||
)
|
||||
|
||||
assert reloaded == aggregated
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregation_unions_the_request_ids_of_merged_increments():
|
||||
"""The seed excludes exactly the requests its batch already covers, so every
|
||||
merged increment's id has to survive aggregation."""
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1"))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 2.0, request_id="req-2"))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert len(aggregated) == 1
|
||||
assert aggregated[0]["request_ids"] == ("req-1", "req-2")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_ids_stay_with_their_own_window():
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a"))
|
||||
await queue.add_update(_txn("k1", WINDOW_B, 2.0, request_id="req-b"))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert {payload["window_start"]: payload["request_ids"] for payload in aggregated} == {
|
||||
"2026-08-01T00:00:00.000000": ("req-a",),
|
||||
"2026-08-31T00:00:00.000000": ("req-b",),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_ids_are_deduplicated_and_ordered():
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-b"))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a"))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-a"))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert aggregated[0]["request_ids"] == ("req-a", "req-b")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_without_a_request_id_carries_no_exclusion():
|
||||
queue = WindowSpendUpdateQueue()
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert aggregated[0]["request_ids"] == ()
|
||||
|
||||
|
||||
def test_request_ids_survive_the_redis_json_round_trip():
|
||||
aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(
|
||||
[(_txn("k1", WINDOW_A, 1.0, request_id="req-1"),)]
|
||||
)
|
||||
|
||||
reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(
|
||||
[json.loads(json.dumps(aggregated))]
|
||||
)
|
||||
|
||||
assert reloaded[0]["request_ids"] == ("req-1",)
|
||||
assert reloaded[0]["spend"] == 1.0
|
||||
535
tests/test_litellm/proxy/db/test_budget_window_spend_writer.py
Normal file
535
tests/test_litellm/proxy/db/test_budget_window_spend_writer.py
Normal file
|
|
@ -0,0 +1,535 @@
|
|||
import math
|
||||
import os
|
||||
import sys
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.db.budget_window_spend_writer import (
|
||||
commit_window_spend_updates,
|
||||
roll_window_spend_row,
|
||||
spend_logs_total_excluding,
|
||||
)
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
build_window_spend_transaction,
|
||||
)
|
||||
|
||||
WINDOW_A = datetime(2026, 8, 1, tzinfo=timezone.utc)
|
||||
WINDOW_B = datetime(2026, 8, 31, tzinfo=timezone.utc)
|
||||
|
||||
ENTITY_TYPE, ENTITY_ID, WINDOW_DURATION, WINDOW_START, INSERT_SPEND, INCREMENT, NOW = range(7)
|
||||
|
||||
|
||||
class _FakeBatcher:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, tuple[Any, ...]]] = []
|
||||
|
||||
def execute_raw(self, query: str, *args: Any) -> None:
|
||||
self.calls.append((query, args))
|
||||
|
||||
|
||||
class _FakeDB:
|
||||
"""Stands in for prisma_client.db; records every statement it is handed."""
|
||||
|
||||
def __init__(self, existing_rows: list[dict[str, str]] | None = None) -> None:
|
||||
self.existing_rows = existing_rows or []
|
||||
self.query_raw_calls: list[tuple[str, tuple[Any, ...]]] = []
|
||||
self.execute_raw_calls: list[tuple[str, tuple[Any, ...]]] = []
|
||||
self.batcher = _FakeBatcher()
|
||||
self.committed = False
|
||||
|
||||
async def query_raw(self, query: str, *args: Any) -> list[dict[str, str]]:
|
||||
self.query_raw_calls.append((query, args))
|
||||
return self.existing_rows
|
||||
|
||||
async def execute_raw(self, query: str, *args: Any) -> int:
|
||||
self.execute_raw_calls.append((query, args))
|
||||
return 1
|
||||
|
||||
@asynccontextmanager
|
||||
async def _tx(self):
|
||||
yield self
|
||||
|
||||
def tx(self, timeout: Any = None):
|
||||
return self._tx()
|
||||
|
||||
@asynccontextmanager
|
||||
async def _batch(self):
|
||||
yield self.batcher
|
||||
self.committed = True
|
||||
|
||||
def batch_(self):
|
||||
return self._batch()
|
||||
|
||||
|
||||
class _FakePrismaClient:
|
||||
def __init__(self, db: _FakeDB) -> None:
|
||||
self.db = db
|
||||
|
||||
|
||||
class _RecordingAggregate:
|
||||
"""Stands in for the LiteLLM_SpendLogs seed aggregate."""
|
||||
|
||||
def __init__(self, value: float = 5.0) -> None:
|
||||
self.value = value
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
prisma_client: Any,
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
exclude_request_ids: Any,
|
||||
) -> float | None:
|
||||
self.calls.append(
|
||||
{
|
||||
"entity_type": entity_type,
|
||||
"entity_id": entity_id,
|
||||
"window_start": window_start,
|
||||
"exclude_request_ids": tuple(exclude_request_ids),
|
||||
}
|
||||
)
|
||||
return self.value
|
||||
|
||||
|
||||
class _SpendLogsFake:
|
||||
"""Sums the LiteLLM_SpendLogs rows it holds, honouring the request-id
|
||||
exclusion exactly as the real aggregate's NOT (request_id = ANY(...)) does."""
|
||||
|
||||
def __init__(self, rows: tuple[tuple[str, float], ...]) -> None:
|
||||
self.rows = rows
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
prisma_client: Any,
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
exclude_request_ids: Any,
|
||||
) -> float | None:
|
||||
excluded = frozenset(exclude_request_ids)
|
||||
return math.fsum(spend for request_id, spend in self.rows if request_id not in excluded)
|
||||
|
||||
|
||||
def _existing(entity_type: str, entity_id: str, window_duration: str) -> dict[str, str]:
|
||||
return {"entity_type": entity_type, "entity_id": entity_id, "window_duration": window_duration}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_transactions_touches_no_database():
|
||||
db = _FakeDB()
|
||||
|
||||
await commit_window_spend_updates(prisma_client=_FakePrismaClient(db), transactions=())
|
||||
|
||||
assert db.query_raw_calls == []
|
||||
assert db.batcher.calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_row_is_seeded_from_spend_logs_once():
|
||||
"""A row created mid-window would undercount everything spent before it
|
||||
existed, so a brand new primary key inserts the SpendLogs total plus this
|
||||
increment."""
|
||||
db = _FakeDB(existing_rows=[])
|
||||
aggregate = _RecordingAggregate(value=5.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
assert len(aggregate.calls) == 1
|
||||
assert aggregate.calls[0]["entity_type"] == "key"
|
||||
assert aggregate.calls[0]["entity_id"] == "k1"
|
||||
assert aggregate.calls[0]["window_start"] == WINDOW_A
|
||||
|
||||
(_, params), = db.batcher.calls
|
||||
assert params[ENTITY_TYPE] == "key"
|
||||
assert params[ENTITY_ID] == "k1"
|
||||
assert params[WINDOW_DURATION] == "30d"
|
||||
assert params[WINDOW_START] == datetime(2026, 8, 1)
|
||||
assert params[INSERT_SPEND] == pytest.approx(6.0)
|
||||
assert params[INCREMENT] == pytest.approx(1.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_existing_row_is_never_reseeded():
|
||||
"""The seed is a full LiteLLM_SpendLogs scan; running it for a row that is
|
||||
already maintained would both cost a scan and double count."""
|
||||
db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")])
|
||||
aggregate = _RecordingAggregate(value=5.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
assert aggregate.calls == []
|
||||
(_, params), = db.batcher.calls
|
||||
assert params[INSERT_SPEND] == pytest.approx(1.0)
|
||||
assert params[INCREMENT] == pytest.approx(1.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_runs_only_for_the_primary_keys_that_are_missing():
|
||||
db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")])
|
||||
aggregate = _RecordingAggregate(value=5.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(
|
||||
build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),
|
||||
build_window_spend_transaction("team", "t1", "30d", WINDOW_A, 2.0),
|
||||
),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
assert [call["entity_id"] for call in aggregate.calls] == ["t1"]
|
||||
assert [call["entity_type"] for call in aggregate.calls] == ["team"]
|
||||
by_entity = {params[ENTITY_ID]: params for _, params in db.batcher.calls}
|
||||
assert by_entity["k1"][INSERT_SPEND] == pytest.approx(1.0)
|
||||
assert by_entity["t1"][INSERT_SPEND] == pytest.approx(7.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_spend_and_increment_differ_only_when_a_row_is_seeded():
|
||||
"""The conflict arm adds the increment alone so two pods that both seed the
|
||||
same new window cannot add the SpendLogs base twice."""
|
||||
db = _FakeDB(existing_rows=[])
|
||||
aggregate = _RecordingAggregate(value=9.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 0.25),),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
(_, params), = db.batcher.calls
|
||||
assert params[INSERT_SPEND] == pytest.approx(9.25)
|
||||
assert params[INCREMENT] == pytest.approx(0.25)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_sql_adds_for_a_current_window_and_replaces_for_a_newer_one():
|
||||
"""The CASE is the whole contract: an increment at or behind the stored
|
||||
window_start accumulates, a newer one restarts the window."""
|
||||
db = _FakeDB(existing_rows=[_existing("key", "k1", "30d")])
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),),
|
||||
)
|
||||
|
||||
(query, _), = db.batcher.calls
|
||||
normalized = " ".join(query.split())
|
||||
assert (
|
||||
'spend = CASE WHEN "LiteLLM_BudgetWindowSpend".window_start >= EXCLUDED.window_start '
|
||||
'THEN "LiteLLM_BudgetWindowSpend".spend + $6 ELSE EXCLUDED.spend END' in normalized
|
||||
)
|
||||
assert 'window_start = GREATEST("LiteLLM_BudgetWindowSpend".window_start, EXCLUDED.window_start)' in normalized
|
||||
assert "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET" in normalized
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_never_interpolates_values_into_the_sql():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
aggregate = _RecordingAggregate(value=0.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(build_window_spend_transaction("key", "'; DROP TABLE x; --", "30d", WINDOW_A, 1.0),),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
(query, params), = db.batcher.calls
|
||||
assert "DROP TABLE" not in query
|
||||
assert params[ENTITY_ID] == "'; DROP TABLE x; --"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upserts_are_ordered_by_primary_key_then_window_start():
|
||||
"""Cross-pod lock ordering, plus an older window must be applied before the
|
||||
roll that supersedes it or the roll would be undone."""
|
||||
db = _FakeDB(existing_rows=[])
|
||||
aggregate = _RecordingAggregate(value=0.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(
|
||||
build_window_spend_transaction("team", "t1", "30d", WINDOW_A, 1.0),
|
||||
build_window_spend_transaction("key", "k2", "30d", WINDOW_B, 1.0),
|
||||
build_window_spend_transaction("key", "k2", "30d", WINDOW_A, 1.0),
|
||||
build_window_spend_transaction("key", "k1", "7d", WINDOW_A, 1.0),
|
||||
),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
ordered = [(params[ENTITY_TYPE], params[ENTITY_ID], params[WINDOW_DURATION], params[WINDOW_START]) for _, params in db.batcher.calls]
|
||||
assert ordered == [
|
||||
("key", "k1", "7d", datetime(2026, 8, 1)),
|
||||
("key", "k2", "30d", datetime(2026, 8, 1)),
|
||||
("key", "k2", "30d", datetime(2026, 8, 31)),
|
||||
("team", "t1", "30d", datetime(2026, 8, 1)),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_existing_row_lookup_sends_every_primary_key_as_array_params():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
aggregate = _RecordingAggregate(value=0.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(
|
||||
build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),
|
||||
build_window_spend_transaction("team", "t1", "7d", WINDOW_A, 1.0),
|
||||
),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
(query, params), = db.query_raw_calls
|
||||
assert "unnest($1::text[], $2::text[], $3::text[])" in query
|
||||
assert params == (("key", "team"), ("k1", "t1"), ("30d", "7d"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_upserts_are_committed_in_one_transaction():
|
||||
db = _FakeDB(existing_rows=[_existing("key", "k1", "30d"), _existing("key", "k2", "30d")])
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(
|
||||
build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),
|
||||
build_window_spend_transaction("key", "k2", "30d", WINDOW_A, 2.0),
|
||||
),
|
||||
)
|
||||
|
||||
assert len(db.batcher.calls) == 2
|
||||
assert db.committed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_entity_type_contributes_no_seed():
|
||||
"""Only key and team windows have a LiteLLM_SpendLogs column to aggregate;
|
||||
anything else starts from its increment alone."""
|
||||
db = _FakeDB(existing_rows=[])
|
||||
|
||||
async def no_such_column(prisma_client, entity_type, entity_id, window_start, exclude_request_ids):
|
||||
return None
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(build_window_spend_transaction("user", "u1", "30d", WINDOW_A, 1.0),),
|
||||
spend_logs_aggregate=no_such_column,
|
||||
)
|
||||
|
||||
(_, params), = db.batcher.calls
|
||||
assert params[INSERT_SPEND] == pytest.approx(1.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unavailable_spend_logs_aggregate_seeds_zero_rather_than_failing():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
|
||||
async def unavailable(prisma_client, entity_type, entity_id, window_start, exclude_request_ids):
|
||||
return None
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(build_window_spend_transaction("key", "k1", "30d", WINDOW_A, 1.0),),
|
||||
spend_logs_aggregate=unavailable,
|
||||
)
|
||||
|
||||
(_, params), = db.batcher.calls
|
||||
assert params[INSERT_SPEND] == pytest.approx(1.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_older():
|
||||
"""Unconditional zeroing would wipe increments a pod already applied under
|
||||
the new window."""
|
||||
db = _FakeDB()
|
||||
|
||||
await roll_window_spend_row(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
entity_type="team",
|
||||
entity_id="t1",
|
||||
window_duration="30d",
|
||||
new_window_start=WINDOW_B,
|
||||
)
|
||||
|
||||
(query, params), = db.execute_raw_calls
|
||||
normalized = " ".join(query.split())
|
||||
assert "SET window_start = ($4::timestamptz AT TIME ZONE 'UTC'), spend = 0" in normalized
|
||||
assert "WHERE entity_type = $1 AND entity_id = $2 AND window_duration = $3" in normalized
|
||||
assert "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" in normalized
|
||||
assert params[:4] == ("team", "t1", "30d", datetime(2026, 8, 31))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_receives_the_batch_request_ids_to_exclude():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
aggregate = _RecordingAggregate(value=0.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(
|
||||
{
|
||||
"entity_type": "key",
|
||||
"entity_id": "k1",
|
||||
"window_duration": "30d",
|
||||
"window_start": "2026-08-01T00:00:00.000000",
|
||||
"spend": 3.0,
|
||||
"request_ids": ("req-1", "req-2", "req-3"),
|
||||
},
|
||||
),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
assert aggregate.calls[0]["exclude_request_ids"] == ("req-1", "req-2", "req-3")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed():
|
||||
"""The spend log writer drains on a ~2s poll while window increments flush
|
||||
on the ~10s batch tick, so a new row is normally seeded from a table that
|
||||
already holds this batch's rows. Counting them in both places is what made
|
||||
a fresh row land at exactly twice the true spend."""
|
||||
db = _FakeDB(existing_rows=[])
|
||||
already_flushed = _SpendLogsFake(
|
||||
rows=(("req-1", 0.000047), ("req-2", 0.000047), ("req-3", 0.000047)),
|
||||
)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(
|
||||
{
|
||||
"entity_type": "key",
|
||||
"entity_id": "k1",
|
||||
"window_duration": "30d",
|
||||
"window_start": "2026-08-01T00:00:00.000000",
|
||||
"spend": 0.000141,
|
||||
"request_ids": ("req-1", "req-2", "req-3"),
|
||||
},
|
||||
),
|
||||
spend_logs_aggregate=already_flushed,
|
||||
)
|
||||
|
||||
(_, params), = db.batcher.calls
|
||||
assert params[INSERT_SPEND] == pytest.approx(0.000141)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_row_still_covers_spend_that_predates_the_batch():
|
||||
"""The exclusion must not throw away the pre-existing spend the seed is for."""
|
||||
db = _FakeDB(existing_rows=[])
|
||||
spend_logs = _SpendLogsFake(rows=(("older", 0.5), ("req-1", 0.000047)))
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(
|
||||
{
|
||||
"entity_type": "key",
|
||||
"entity_id": "k1",
|
||||
"window_duration": "30d",
|
||||
"window_start": "2026-08-01T00:00:00.000000",
|
||||
"spend": 0.000047,
|
||||
"request_ids": ("req-1",),
|
||||
},
|
||||
),
|
||||
spend_logs_aggregate=spend_logs,
|
||||
)
|
||||
|
||||
(_, params), = db.batcher.calls
|
||||
assert params[INSERT_SPEND] == pytest.approx(0.500047)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet():
|
||||
"""The other side of the race: rows absent from the aggregate are still
|
||||
counted exactly once, by their increment."""
|
||||
db = _FakeDB(existing_rows=[])
|
||||
nothing_flushed = _SpendLogsFake(rows=())
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(
|
||||
{
|
||||
"entity_type": "key",
|
||||
"entity_id": "k1",
|
||||
"window_duration": "30d",
|
||||
"window_start": "2026-08-01T00:00:00.000000",
|
||||
"spend": 0.000141,
|
||||
"request_ids": ("req-1", "req-2", "req-3"),
|
||||
},
|
||||
),
|
||||
spend_logs_aggregate=nothing_flushed,
|
||||
)
|
||||
|
||||
(_, params), = db.batcher.calls
|
||||
assert params[INSERT_SPEND] == pytest.approx(0.000141)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"entity_type, expected_column",
|
||||
[("key", "api_key = $1"), ("team", "team_id = $1")],
|
||||
)
|
||||
async def test_seed_aggregate_sql_excludes_the_request_ids_by_parameter(entity_type, expected_column):
|
||||
db = _FakeDB(existing_rows=[{"total": 1.25}])
|
||||
|
||||
total = await spend_logs_total_excluding(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
entity_type=entity_type,
|
||||
entity_id="e1",
|
||||
window_start=WINDOW_A,
|
||||
exclude_request_ids=("req-1", "req-2"),
|
||||
)
|
||||
|
||||
assert total == pytest.approx(1.25)
|
||||
(query, params), = db.query_raw_calls
|
||||
normalized = " ".join(query.split())
|
||||
assert expected_column in normalized
|
||||
assert "NOT (request_id = ANY($3::text[]))" in normalized
|
||||
assert 'FROM "LiteLLM_SpendLogs"' in normalized
|
||||
assert params == ("e1", WINDOW_A, ("req-1", "req-2"))
|
||||
# The ids are bound, never spliced into the statement.
|
||||
assert "req-1" not in query
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs_column():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
|
||||
total = await spend_logs_total_excluding(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
entity_type="user",
|
||||
entity_id="u1",
|
||||
window_start=WINDOW_A,
|
||||
exclude_request_ids=(),
|
||||
)
|
||||
|
||||
assert total is None
|
||||
assert db.query_raw_calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
|
||||
total = await spend_logs_total_excluding(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
entity_type="key",
|
||||
entity_id="k-unknown",
|
||||
window_start=WINDOW_A,
|
||||
exclude_request_ids=(),
|
||||
)
|
||||
|
||||
assert total == 0.0
|
||||
|
|
@ -9,6 +9,7 @@ sys.path.insert(
|
|||
) # Adds the parent directory to the system path
|
||||
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
|
||||
|
|
@ -18,6 +19,9 @@ from redis.exceptions import DataError
|
|||
import litellm
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
build_window_spend_transaction,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -2236,3 +2240,258 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent():
|
|||
assert transaction["compression_saved_tokens"] == 0
|
||||
assert transaction["compression_savings_spend"] == 0
|
||||
assert transaction["prompt_caching_savings_spend"] == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Budget window spend flush (LiteLLM_BudgetWindowSpend)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _WindowSpendFakeBatcher:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def execute_raw(self, query, *args):
|
||||
self.calls.append((query, args))
|
||||
|
||||
|
||||
class _WindowSpendFakeDB:
|
||||
"""Minimal prisma_client.db that records the raw statements it is handed."""
|
||||
|
||||
def __init__(self, existing_rows=None):
|
||||
self.existing_rows = existing_rows or []
|
||||
self.query_raw_calls = []
|
||||
self.batcher = _WindowSpendFakeBatcher()
|
||||
|
||||
async def query_raw(self, query, *args):
|
||||
self.query_raw_calls.append((query, args))
|
||||
if "LiteLLM_BudgetWindowSpend" in query:
|
||||
return self.existing_rows
|
||||
return []
|
||||
|
||||
@asynccontextmanager
|
||||
async def _tx(self):
|
||||
yield self
|
||||
|
||||
def tx(self, timeout=None):
|
||||
return self._tx()
|
||||
|
||||
@asynccontextmanager
|
||||
async def _batch(self):
|
||||
yield self.batcher
|
||||
|
||||
def batch_(self):
|
||||
return self._batch()
|
||||
|
||||
|
||||
class _WindowSpendFakePrisma:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
|
||||
|
||||
def _window_spend_upserts(db):
|
||||
return [
|
||||
params for query, params in db.batcher.calls if "LiteLLM_BudgetWindowSpend" in query
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_spend_queue_is_flushed_without_redis_buffer():
|
||||
"""The in-memory window queue must reach the DB on the same scheduler tick
|
||||
as the other spend queues when the Redis buffer is off."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
await db_writer.window_spend_update_queue.add_update(
|
||||
build_window_spend_transaction(
|
||||
entity_type="key",
|
||||
entity_id="hashed-token",
|
||||
window_duration="30d",
|
||||
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
spend=0.5,
|
||||
)
|
||||
)
|
||||
db = _WindowSpendFakeDB(
|
||||
existing_rows=[
|
||||
{"entity_type": "key", "entity_id": "hashed-token", "window_duration": "30d"}
|
||||
]
|
||||
)
|
||||
|
||||
await db_writer._commit_spend_updates_to_db_without_redis_buffer(
|
||||
prisma_client=_WindowSpendFakePrisma(db),
|
||||
n_retry_times=0,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
upserts = _window_spend_upserts(db)
|
||||
assert len(upserts) == 1
|
||||
assert upserts[0][0] == "key"
|
||||
assert upserts[0][1] == "hashed-token"
|
||||
assert upserts[0][2] == "30d"
|
||||
assert upserts[0][5] == pytest.approx(0.5)
|
||||
assert db_writer.window_spend_update_queue.update_queue.qsize() == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_spend_queue_is_handed_to_the_redis_buffer():
|
||||
"""Multi-pod deployments buffer through Redis, so the window queue has to
|
||||
ride the same rpush path as the daily queues."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
mock_redis_update_buffer = AsyncMock()
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
|
||||
return_value=(None, None, None, None, None, None, None)
|
||||
)
|
||||
db_writer.redis_update_buffer = mock_redis_update_buffer
|
||||
db_writer.pod_lock_manager = AsyncMock()
|
||||
db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
|
||||
await db_writer._commit_spend_updates_to_db_with_redis(
|
||||
prisma_client=MagicMock(),
|
||||
n_retry_times=0,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
stored = mock_redis_update_buffer.store_in_memory_spend_updates_in_redis.call_args[1]
|
||||
assert stored["window_spend_update_queue"] is db_writer.window_spend_update_queue
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_spend_transactions_from_redis_are_committed_by_the_lock_winner():
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
window_transactions = (
|
||||
build_window_spend_transaction(
|
||||
entity_type="team",
|
||||
entity_id="team-1",
|
||||
window_duration="7d",
|
||||
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
spend=2.0,
|
||||
),
|
||||
)
|
||||
mock_redis_update_buffer = AsyncMock()
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
|
||||
return_value=(None, None, None, None, None, None, window_transactions)
|
||||
)
|
||||
db_writer.redis_update_buffer = mock_redis_update_buffer
|
||||
db_writer.pod_lock_manager = AsyncMock()
|
||||
db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
db = _WindowSpendFakeDB(
|
||||
existing_rows=[{"entity_type": "team", "entity_id": "team-1", "window_duration": "7d"}]
|
||||
)
|
||||
|
||||
await db_writer._commit_spend_updates_to_db_with_redis(
|
||||
prisma_client=_WindowSpendFakePrisma(db),
|
||||
n_retry_times=0,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
upserts = _window_spend_upserts(db)
|
||||
assert len(upserts) == 1
|
||||
assert upserts[0][:3] == ("team", "team-1", "7d")
|
||||
assert upserts[0][5] == pytest.approx(2.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_spend_transactions_are_not_committed_without_the_pod_lock():
|
||||
"""Every pod buffers to Redis but only the lock winner may drain it."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
mock_redis_update_buffer = AsyncMock()
|
||||
db_writer.redis_update_buffer = mock_redis_update_buffer
|
||||
db_writer.pod_lock_manager = AsyncMock()
|
||||
db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
|
||||
db = _WindowSpendFakeDB()
|
||||
|
||||
await db_writer._commit_spend_updates_to_db_with_redis(
|
||||
prisma_client=_WindowSpendFakePrisma(db),
|
||||
n_retry_times=0,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline.assert_not_called()
|
||||
assert _window_spend_upserts(db) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_window_spend_commit_does_not_abort_the_rest_of_the_flush():
|
||||
"""Window rows are an optimization over aggregating LiteLLM_SpendLogs, so a
|
||||
failure must not take the tool registry flush down with it."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
await db_writer.window_spend_update_queue.add_update(
|
||||
build_window_spend_transaction(
|
||||
entity_type="key",
|
||||
entity_id="hashed-token",
|
||||
window_duration="30d",
|
||||
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
spend=0.5,
|
||||
)
|
||||
)
|
||||
db = _WindowSpendFakeDB()
|
||||
db.query_raw = AsyncMock(side_effect=Exception("connection reset"))
|
||||
db_writer._flush_tool_discovery_queue = AsyncMock()
|
||||
|
||||
await db_writer._commit_spend_updates_to_db_without_redis_buffer(
|
||||
prisma_client=_WindowSpendFakePrisma(db),
|
||||
n_retry_times=0,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
db_writer._flush_tool_discovery_queue.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_returns_the_spend_log_request_id():
|
||||
"""The budget-window seed excludes the log rows its increments already
|
||||
cover, so the caller needs the id this call was recorded under. It cannot
|
||||
be re-derived: cache hits append time.time() to the id."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db_writer._insert_spend_log_to_db = AsyncMock()
|
||||
db_writer._enqueue_tool_usage_transaction = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.disable_spend_logs", False),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"),
|
||||
):
|
||||
request_id = await db_writer.update_database(
|
||||
token="test-token",
|
||||
user_id="test-user",
|
||||
end_user_id=None,
|
||||
team_id="test-team",
|
||||
org_id=None,
|
||||
kwargs={"model": "gpt-4", "custom_llm_provider": "openai", "litellm_call_id": "call-xyz"},
|
||||
completion_response=MagicMock(),
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
response_cost=0.1,
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert request_id is not None
|
||||
# Same id the spend log row was queued under.
|
||||
assert request_id == db_writer._insert_spend_log_to_db.call_args[1]["payload"]["request_id"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_returns_none_when_the_payload_cannot_be_built():
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.disable_spend_logs", False),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"),
|
||||
patch(
|
||||
"litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload",
|
||||
side_effect=Exception("payload boom"),
|
||||
),
|
||||
):
|
||||
request_id = await db_writer.update_database(
|
||||
token="test-token",
|
||||
user_id="test-user",
|
||||
end_user_id=None,
|
||||
team_id="test-team",
|
||||
org_id=None,
|
||||
kwargs={"model": "gpt-4"},
|
||||
completion_response=MagicMock(),
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
response_cost=0.1,
|
||||
)
|
||||
|
||||
assert request_id is None
|
||||
|
|
|
|||
|
|
@ -440,7 +440,9 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re
|
|||
@pytest.mark.asyncio
|
||||
async def test_update_database_and_spend_counters_updates_counters_after_db_update():
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock()
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(
|
||||
return_value="chatcmpl-abc123"
|
||||
)
|
||||
increment_spend_counters = AsyncMock()
|
||||
budget_reservation = {"reserved_cost": 0.5, "entries": []}
|
||||
|
||||
|
|
@ -471,6 +473,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
|
|||
budget_reservation=budget_reservation,
|
||||
end_user_id="test_end_user_id",
|
||||
tags=["tag-a"],
|
||||
request_id="chatcmpl-abc123",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1263,3 +1266,58 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(
|
|||
assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (
|
||||
1 if expect_spend_log else 0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_and_spend_counters_forwards_the_spend_log_request_id():
|
||||
"""The budget-window flush excludes the log rows its increments already
|
||||
cover. That only works if the id update_database recorded the row under is
|
||||
handed to the counter update, so this seam is load-bearing."""
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(
|
||||
return_value="chatcmpl-abc123"
|
||||
)
|
||||
increment_spend_counters = AsyncMock()
|
||||
|
||||
await _update_database_and_spend_counters(
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
increment_spend_counters=increment_spend_counters,
|
||||
user_api_key="test_api_key",
|
||||
user_id="test_user_id",
|
||||
end_user_id=None,
|
||||
team_id="test_team_id",
|
||||
org_id="test_org_id",
|
||||
kwargs={},
|
||||
completion_response=None,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
response_cost=0.2,
|
||||
budget_reservation=None,
|
||||
)
|
||||
|
||||
assert increment_spend_counters.await_args.kwargs["request_id"] == "chatcmpl-abc123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_and_spend_counters_forwards_a_missing_request_id_as_none():
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(return_value=None)
|
||||
increment_spend_counters = AsyncMock()
|
||||
|
||||
await _update_database_and_spend_counters(
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
increment_spend_counters=increment_spend_counters,
|
||||
user_api_key="test_api_key",
|
||||
user_id="test_user_id",
|
||||
end_user_id=None,
|
||||
team_id="test_team_id",
|
||||
org_id="test_org_id",
|
||||
kwargs={},
|
||||
completion_response=None,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
response_cost=0.2,
|
||||
budget_reservation=None,
|
||||
)
|
||||
|
||||
assert increment_spend_counters.await_args.kwargs["request_id"] is None
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
|
|
@ -11016,3 +11017,260 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog):
|
|||
ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings={})
|
||||
|
||||
assert MOCK_TESTING_CONFIG_KEY not in caplog.text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Budget window spend row enqueue (LiteLLM_BudgetWindowSpend writer)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _window_spend_enqueue_env(cached_objects: dict):
|
||||
"""Point increment_spend_counters at throwaway caches and a real
|
||||
WindowSpendUpdateQueue, and hand back the queue to inspect."""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
WindowSpendUpdateQueue,
|
||||
)
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_get_cache = AsyncMock(side_effect=lambda key, **_: cached_objects.get(key))
|
||||
|
||||
queue = WindowSpendUpdateQueue()
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.window_spend_update_queue = queue
|
||||
|
||||
originals = (
|
||||
ps.user_api_key_cache,
|
||||
ps.spend_counter_cache,
|
||||
ps.prisma_client,
|
||||
ps.proxy_logging_obj,
|
||||
)
|
||||
ps.user_api_key_cache = user_api_key_cache
|
||||
ps.spend_counter_cache = DualCache()
|
||||
ps.prisma_client = None
|
||||
ps.proxy_logging_obj = proxy_logging_obj
|
||||
try:
|
||||
yield queue
|
||||
finally:
|
||||
(
|
||||
ps.user_api_key_cache,
|
||||
ps.spend_counter_cache,
|
||||
ps.prisma_client,
|
||||
ps.proxy_logging_obj,
|
||||
) = originals
|
||||
|
||||
|
||||
async def _drain(queue):
|
||||
return list(await queue.flush_and_get_aggregated_window_spend_transactions())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_window_spend_row_is_enqueued_with_the_actual_cost():
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
reset_at = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
key_obj = MagicMock()
|
||||
key_obj.budget_limits = [
|
||||
{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}
|
||||
]
|
||||
|
||||
with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue:
|
||||
await increment_spend_counters(
|
||||
token="hashed-token", team_id=None, user_id=None, response_cost=0.25
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert len(enqueued) == 1
|
||||
assert enqueued[0]["entity_type"] == "key"
|
||||
assert enqueued[0]["entity_id"] == "hashed-token"
|
||||
assert enqueued[0]["window_duration"] == "30d"
|
||||
assert enqueued[0]["spend"] == pytest.approx(0.25)
|
||||
assert enqueued[0]["window_start"] == (reset_at - timedelta(days=30)).astimezone(timezone.utc).replace(
|
||||
tzinfo=None
|
||||
).isoformat(timespec="microseconds")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_window_spend_row_is_enqueued():
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
reset_at = datetime.now(timezone.utc) + timedelta(days=3)
|
||||
team_obj = MagicMock()
|
||||
team_obj.budget_limits = [
|
||||
{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}
|
||||
]
|
||||
|
||||
with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue:
|
||||
await increment_spend_counters(
|
||||
token=None, team_id="team-1", user_id=None, response_cost=1.5
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert len(enqueued) == 1
|
||||
assert enqueued[0]["entity_type"] == "team"
|
||||
assert enqueued[0]["entity_id"] == "team-1"
|
||||
assert enqueued[0]["window_duration"] == "7d"
|
||||
assert enqueued[0]["spend"] == pytest.approx(1.5)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_spend_row_is_enqueued_even_when_the_counter_was_reserved():
|
||||
"""A reservation only pre-charged the cache counter with an estimate; the
|
||||
row still owes the actual cost, so the enqueue must not be skipped."""
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
import litellm.proxy.spend_tracking.budget_reservation as br
|
||||
|
||||
reset_at = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
key_obj = MagicMock()
|
||||
key_obj.budget_limits = [
|
||||
{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}
|
||||
]
|
||||
reservation = {
|
||||
"entries": [
|
||||
{"counter_key": "spend:key:hashed-token", "reserved": 1.0},
|
||||
{"counter_key": "spend:key:hashed-token:window:30d", "reserved": 1.0},
|
||||
]
|
||||
}
|
||||
|
||||
original_reconcile = br.reconcile_budget_reservation
|
||||
br.reconcile_budget_reservation = AsyncMock(return_value=None)
|
||||
try:
|
||||
with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue:
|
||||
await increment_spend_counters(
|
||||
token="hashed-token",
|
||||
team_id=None,
|
||||
user_id=None,
|
||||
response_cost=0.25,
|
||||
budget_reservation=reservation,
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
finally:
|
||||
br.reconcile_budget_reservation = original_reconcile
|
||||
|
||||
assert len(enqueued) == 1
|
||||
assert enqueued[0]["spend"] == pytest.approx(0.25)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sliding_window_without_reset_at_is_not_enqueued():
|
||||
"""Windows with no reset_at slide with wall clock, so window_start moves on
|
||||
every request and no single row can represent them; the read path keeps
|
||||
using its LiteLLM_SpendLogs fallback instead."""
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
key_obj = MagicMock()
|
||||
key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0}]
|
||||
|
||||
with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue:
|
||||
await increment_spend_counters(
|
||||
token="hashed-token", team_id=None, user_id=None, response_cost=0.25
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert enqueued == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_configured_window_gets_its_own_row_enqueue():
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
key_obj = MagicMock()
|
||||
key_obj.budget_limits = [
|
||||
{"budget_duration": "1d", "max_budget": 5.0, "reset_at": (now + timedelta(hours=5)).isoformat()},
|
||||
{"budget_duration": "30d", "max_budget": 100.0, "reset_at": (now + timedelta(days=10)).isoformat()},
|
||||
]
|
||||
|
||||
with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue:
|
||||
await increment_spend_counters(
|
||||
token="hashed-token", team_id=None, user_id=None, response_cost=0.25
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert sorted(item["window_duration"] for item in enqueued) == ["1d", "30d"]
|
||||
assert all(item["spend"] == pytest.approx(0.25) for item in enqueued)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_window_spend_row_enqueued_without_budget_limits():
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
key_obj = MagicMock()
|
||||
key_obj.budget_limits = None
|
||||
|
||||
with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue:
|
||||
await increment_spend_counters(
|
||||
token="hashed-token", team_id=None, user_id=None, response_cost=0.25
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert enqueued == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_spend_row_carries_the_spend_log_request_id():
|
||||
"""The flush excludes these ids from its one-time seed, so the id threaded
|
||||
here has to be the same one the LiteLLM_SpendLogs row was written under."""
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
reset_at = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
key_obj = MagicMock()
|
||||
key_obj.budget_limits = [
|
||||
{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}
|
||||
]
|
||||
|
||||
with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue:
|
||||
await increment_spend_counters(
|
||||
token="hashed-token",
|
||||
team_id=None,
|
||||
user_id=None,
|
||||
response_cost=0.25,
|
||||
request_id="chatcmpl-abc123",
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert enqueued[0]["request_ids"] == ("chatcmpl-abc123",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_spend_row_without_a_request_id_excludes_nothing():
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
reset_at = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
key_obj = MagicMock()
|
||||
key_obj.budget_limits = [
|
||||
{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}
|
||||
]
|
||||
|
||||
with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue:
|
||||
await increment_spend_counters(
|
||||
token="hashed-token", team_id=None, user_id=None, response_cost=0.25
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert enqueued[0]["request_ids"] == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_window_spend_row_carries_the_request_id():
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
reset_at = datetime.now(timezone.utc) + timedelta(days=3)
|
||||
team_obj = MagicMock()
|
||||
team_obj.budget_limits = [
|
||||
{"budget_duration": "7d", "max_budget": 50.0, "reset_at": reset_at.isoformat()}
|
||||
]
|
||||
|
||||
with _window_spend_enqueue_env({"team_id:team-1": team_obj}) as queue:
|
||||
await increment_spend_counters(
|
||||
token=None,
|
||||
team_id="team-1",
|
||||
user_id=None,
|
||||
response_cost=1.5,
|
||||
request_id="chatcmpl-team",
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert enqueued[0]["request_ids"] == ("chatcmpl-team",)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue