mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge pull request #35887 from BerriAI/litellm_window_spend_reader
perf(proxy): read budget-window spend from the maintained window table
This commit is contained in:
ryan-crabbe-berri
GitHub
commit
2ffd3d0e7d
23 changed files with 3102 additions and 132 deletions
|
|
@ -289,6 +289,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))
|
||||
|
|
|
|||
|
|
@ -4826,6 +4826,7 @@ async def _virtual_key_multi_budget_check(
|
|||
max_budget=w["max_budget"],
|
||||
window_entity_type="Key",
|
||||
window_entity_id=valid_token.token,
|
||||
window_duration=str(w["budget_duration"]),
|
||||
window_start=get_budget_window_start(w),
|
||||
)
|
||||
if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]:
|
||||
|
|
@ -5199,6 +5200,7 @@ async def _team_multi_budget_check(
|
|||
max_budget=w["max_budget"],
|
||||
window_entity_type="Team",
|
||||
window_entity_id=team_object.team_id,
|
||||
window_duration=str(w["budget_duration"]),
|
||||
window_start=get_budget_window_start(w),
|
||||
)
|
||||
if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]:
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import math
|
|||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, Protocol, TypeVar, assert_never
|
||||
|
|
@ -20,10 +20,12 @@ from litellm.constants import (
|
|||
RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN,
|
||||
RESET_BUDGET_JOB_NAME,
|
||||
)
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.proxy._types import (
|
||||
DB_RETRY_SAFE_ERROR_TYPES,
|
||||
LiteLLM_BudgetTableFull,
|
||||
LiteLLM_EndUserTable,
|
||||
Litellm_EntityType,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
LiteLLM_VerificationToken,
|
||||
|
|
@ -38,6 +40,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
model_access_group_spend_counter_key,
|
||||
tag_cache_key,
|
||||
)
|
||||
from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
|
@ -347,6 +350,7 @@ class _WindowSource:
|
|||
|
||||
table: str
|
||||
id_column: str
|
||||
entity_type: Litellm_EntityType
|
||||
counter_prefix: str
|
||||
log_subject: str
|
||||
retry_subject: str
|
||||
|
|
@ -371,6 +375,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = (
|
|||
_WindowSource(
|
||||
table="LiteLLM_VerificationToken",
|
||||
id_column="token",
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
counter_prefix="spend:key",
|
||||
log_subject="keys",
|
||||
retry_subject="key",
|
||||
|
|
@ -379,6 +384,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = (
|
|||
_WindowSource(
|
||||
table="LiteLLM_TeamTable",
|
||||
id_column="team_id",
|
||||
entity_type=Litellm_EntityType.TEAM,
|
||||
counter_prefix="spend:team",
|
||||
log_subject="teams",
|
||||
retry_subject="team",
|
||||
|
|
@ -1241,6 +1247,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")
|
||||
|
|
@ -1256,11 +1265,56 @@ class ResetBudgetJob:
|
|||
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_value)
|
||||
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,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _window_carried_spend(
|
||||
window: Mapping[str, object], counter_key: str, spend_counter_cache: DualCache
|
||||
|
|
@ -1356,6 +1410,9 @@ class ResetBudgetJob:
|
|||
spend_counter_cache,
|
||||
now,
|
||||
self.reset_settings,
|
||||
prisma_client=self.prisma_client,
|
||||
entity_type=source.entity_type,
|
||||
entity_id=row_id,
|
||||
):
|
||||
changed = True
|
||||
if changed:
|
||||
|
|
|
|||
313
litellm/proxy/db/budget_window_spend_writer.py
Normal file
313
litellm/proxy/db/budget_window_spend_writer.py
Normal file
|
|
@ -0,0 +1,313 @@
|
|||
"""
|
||||
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
|
||||
include spend logs whose increments are still queued on another pod, and those
|
||||
increments are added again when that pod flushes. That is bounded by a single
|
||||
flush interval, happens at most once per window row, and only ever over-counts:
|
||||
the seed never omits spend, because every increment not yet in the row still
|
||||
reaches it on its own pod's next flush. A row therefore lags real spend by at
|
||||
most one flush interval of queued increments, the same lag the SpendLogs
|
||||
aggregate it replaces (and every other spend column) already has.
|
||||
"""
|
||||
|
||||
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[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))"
|
||||
)
|
||||
|
||||
_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[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))"
|
||||
)
|
||||
|
||||
_SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_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')"
|
||||
)
|
||||
|
||||
_SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_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')"
|
||||
)
|
||||
|
||||
_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],
|
||||
exclude_started_at: datetime | None,
|
||||
) -> 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],
|
||||
exclude_started_at: datetime | None,
|
||||
) -> 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.
|
||||
|
||||
The exclusion is bounded to rows that started at or after the batch's
|
||||
earliest request. request_id can be chosen by the client
|
||||
(x-litellm-call-id), so an unbounded exclusion would let a replayed old id
|
||||
erase a historical row from the seed while its increment still lands.
|
||||
Without a known start the batch's ids are not excluded at all: that can
|
||||
only over-count once, which enforcement tolerates, whereas under-counting
|
||||
is a budget bypass.
|
||||
"""
|
||||
if entity_type == Litellm_EntityType.KEY.value:
|
||||
bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_KEY_SQL, _SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL
|
||||
elif entity_type == Litellm_EntityType.TEAM.value:
|
||||
bounded_sql, unbounded_sql = _SEED_FROM_SPEND_LOGS_TEAM_SQL, _SEED_FROM_SPEND_LOGS_TEAM_UNBOUNDED_SQL
|
||||
else:
|
||||
return None
|
||||
rows: Final = (
|
||||
await prisma_client.db.query_raw(unbounded_sql, entity_id, window_start)
|
||||
if exclude_started_at is None or not exclude_request_ids
|
||||
else await prisma_client.db.query_raw(
|
||||
bounded_sql,
|
||||
entity_id,
|
||||
window_start,
|
||||
tuple(exclude_request_ids),
|
||||
_exclusion_lower_bound(exclude_started_at),
|
||||
)
|
||||
)
|
||||
if not rows:
|
||||
return 0.0
|
||||
return float(rows[0].get("total") or 0.0)
|
||||
|
||||
|
||||
def _exclusion_lower_bound(started_at: datetime) -> datetime:
|
||||
"""LiteLLM_SpendLogs.startTime is TIMESTAMP(3); floor to the second so a
|
||||
millisecond rounding of the batch's own earliest row cannot slip under it."""
|
||||
return to_naive_utc(started_at).replace(microsecond=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"],
|
||||
exclude_started_at=_transaction_started_at(transaction),
|
||||
)
|
||||
return float(base or 0.0)
|
||||
|
||||
|
||||
def _transaction_started_at(transaction: WindowSpendTransaction) -> datetime | None:
|
||||
started_at: Final = transaction.get("started_at")
|
||||
if started_at is None:
|
||||
return None
|
||||
return datetime.fromisoformat(started_at).replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
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,
|
||||
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)),
|
||||
)
|
||||
|
|
@ -56,6 +56,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,
|
||||
|
|
@ -195,6 +199,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
|
||||
|
|
@ -210,7 +215,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,
|
||||
|
|
@ -227,7 +236,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:
|
||||
|
|
@ -301,6 +310,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 "
|
||||
|
|
@ -313,6 +323,7 @@ class DBSpendUpdateWriter:
|
|||
org_id,
|
||||
end_user_id,
|
||||
)
|
||||
return None
|
||||
|
||||
async def _enqueue_tool_usage_transaction(
|
||||
self,
|
||||
|
|
@ -999,6 +1010,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
|
||||
|
|
@ -1017,6 +1029,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()
|
||||
|
||||
uncommitted = { # mutable-ok: drives which popped categories still need re-queuing
|
||||
|
|
@ -1026,6 +1039,7 @@ class DBSpendUpdateWriter:
|
|||
"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,
|
||||
}
|
||||
|
||||
if db_spend_update_transactions is not None:
|
||||
|
|
@ -1095,6 +1109,12 @@ class DBSpendUpdateWriter:
|
|||
daily_spend_transactions=daily_agent_spend_update_transactions,
|
||||
)
|
||||
uncommitted.pop("daily_agent_spend_update_transactions", None)
|
||||
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,
|
||||
)
|
||||
uncommitted.pop("window_spend_update_transactions", None)
|
||||
except Exception as e:
|
||||
spend_log_error(
|
||||
"Spend tracking - failed to commit spend updates from Redis to DB. "
|
||||
|
|
@ -1210,6 +1230,27 @@ 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()
|
||||
)
|
||||
|
||||
try:
|
||||
await DBSpendUpdateWriter._commit_window_spend_updates(
|
||||
prisma_client=prisma_client,
|
||||
window_spend_transactions=window_spend_update_transactions,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run
|
||||
spend_log_error(
|
||||
"Spend tracking - failed to commit budget window spend updates. "
|
||||
"Re-queued %d window increments for retry on next tick. Error: %s",
|
||||
len(window_spend_update_transactions),
|
||||
str(e),
|
||||
exc=e,
|
||||
)
|
||||
await self.window_spend_update_queue.update_queue.put(window_spend_update_transactions)
|
||||
|
||||
################## Tool Registry Upserts ##################
|
||||
await self._flush_tool_discovery_queue(prisma_client=prisma_client)
|
||||
|
||||
|
|
@ -1274,6 +1315,28 @@ 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.
|
||||
|
||||
Raises on failure so the caller re-queues the increments: budget
|
||||
enforcement trusts a current row without reconciling it against
|
||||
LiteLLM_SpendLogs, so a dropped increment would let the entity spend
|
||||
past its window limit after the next counter reseed.
|
||||
"""
|
||||
from litellm.proxy.db.budget_window_spend_writer import (
|
||||
commit_window_spend_updates,
|
||||
)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=prisma_client,
|
||||
transactions=window_spend_transactions,
|
||||
)
|
||||
|
||||
async def _drain_and_commit_daily_tag_spend_from_redis(
|
||||
self,
|
||||
prisma_client: PrismaClient,
|
||||
|
|
|
|||
|
|
@ -23,6 +23,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 (
|
||||
|
|
@ -42,6 +43,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,
|
||||
|
|
@ -182,6 +187,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
|
||||
|
|
@ -250,6 +256,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)
|
||||
|
|
@ -286,6 +297,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]] = []
|
||||
|
|
@ -326,12 +342,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
|
||||
|
||||
|
|
@ -351,12 +369,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.
|
||||
|
|
@ -430,6 +450,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)
|
||||
|
||||
async def restore_transactions_to_redis(
|
||||
self,
|
||||
db_spend_update_transactions: DBSpendUpdateTransactions | None = None,
|
||||
|
|
@ -439,6 +462,7 @@ class RedisUpdateBuffer:
|
|||
daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
|
||||
daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
|
||||
daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
|
||||
window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Re-push transactions that were popped from Redis but not committed to the DB.
|
||||
|
|
@ -460,6 +484,7 @@ class RedisUpdateBuffer:
|
|||
(daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY),
|
||||
(daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY),
|
||||
(daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY),
|
||||
(window_spend_update_transactions, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY),
|
||||
)
|
||||
|
||||
rpush_list: Final = tuple(
|
||||
|
|
@ -577,20 +602,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),
|
||||
|
|
@ -614,12 +641,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
|
||||
|
|
@ -640,6 +671,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]),
|
||||
|
|
@ -647,6 +686,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,176 @@
|
|||
"""
|
||||
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 typing_extensions import ReadOnly
|
||||
|
||||
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.
|
||||
|
||||
started_at is the earliest request start in the batch. The seed only
|
||||
subtracts a request_id whose LiteLLM_SpendLogs.startTime is at or after it,
|
||||
so a client that replays an old id through x-litellm-call-id cannot make the
|
||||
seed drop the historical row that id already paid for.
|
||||
"""
|
||||
|
||||
entity_type: ReadOnly[str]
|
||||
entity_id: ReadOnly[str]
|
||||
window_duration: ReadOnly[str]
|
||||
window_start: ReadOnly[str]
|
||||
spend: ReadOnly[float]
|
||||
request_ids: ReadOnly[Sequence[str]]
|
||||
started_at: ReadOnly[str | None]
|
||||
|
||||
|
||||
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,
|
||||
started_at: datetime | 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,),
|
||||
started_at=None
|
||||
if started_at is None
|
||||
else to_naive_utc(started_at.astimezone(timezone.utc)).isoformat(timespec="microseconds"),
|
||||
)
|
||||
|
||||
|
||||
def _merge_window_spend_transactions(
|
||||
payloads: tuple[WindowSpendTransaction, ...],
|
||||
) -> WindowSpendTransaction:
|
||||
first: Final = payloads[0]
|
||||
started_ats: Final = tuple(
|
||||
started_at for payload in payloads if (started_at := payload.get("started_at")) is not None
|
||||
)
|
||||
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)))),
|
||||
started_at=min(started_ats) if started_ats else None,
|
||||
)
|
||||
|
||||
|
||||
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) -> None:
|
||||
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)
|
||||
)
|
||||
|
|
@ -14,14 +14,18 @@ memory in long-lived deployments.
|
|||
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from datetime import datetime
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.table_repositories import (
|
||||
BudgetWindowSpendRepository,
|
||||
SpendLogsRepository,
|
||||
TeamMembershipRepository,
|
||||
)
|
||||
|
|
@ -36,6 +40,25 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
_WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"Key": Litellm_EntityType.KEY.value,
|
||||
"Team": Litellm_EntityType.TEAM.value,
|
||||
}
|
||||
)
|
||||
|
||||
_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"Key": "api_key",
|
||||
"Team": "team_id",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _as_utc(value: datetime) -> datetime:
|
||||
return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
class SpendCounterReseed:
|
||||
"""
|
||||
Reseeds spend counters from the authoritative DB and warms the cache,
|
||||
|
|
@ -205,6 +228,92 @@ class SpendCounterReseed:
|
|||
raise
|
||||
return current_value
|
||||
|
||||
@staticmethod
|
||||
async def window_from_table(
|
||||
prisma_client: Optional["PrismaClient"],
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_duration: str,
|
||||
expected_window_start: datetime,
|
||||
) -> float | None:
|
||||
"""
|
||||
Read the maintained per-window spend row by primary key.
|
||||
|
||||
Returns the row's spend only when the row belongs to the window the
|
||||
caller is enforcing, i.e. ``row.window_start >= expected_window_start``.
|
||||
A row at or past the expected start was rolled by a pod whose reset_at
|
||||
was at least as fresh as this caller's, so it is trusted; an older row
|
||||
means the window boundary was crossed and nothing has rolled the row
|
||||
yet, so its spend belongs to a previous window.
|
||||
|
||||
Returns None for a missing, stale or unreadable row so the caller falls
|
||||
back to the spend-logs aggregate. ``entity_type`` is the counter-facing
|
||||
label ("Key"/"Team"); anything else has no row and returns None.
|
||||
"""
|
||||
if prisma_client is None:
|
||||
return None
|
||||
row_entity_type: Final = _WINDOW_SPEND_ENTITY_TYPES.get(entity_type)
|
||||
if row_entity_type is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique(
|
||||
where={
|
||||
"entity_type_entity_id_window_duration": {
|
||||
"entity_type": row_entity_type,
|
||||
"entity_id": entity_id,
|
||||
"window_duration": window_duration,
|
||||
}
|
||||
}
|
||||
)
|
||||
except Exception: # noqa: BLE001 # any read failure (DB, stale prisma client) must degrade to the aggregate path
|
||||
verbose_proxy_logger.exception(
|
||||
"SpendCounterReseed.window_from_table: failed for %s=%s window=%s",
|
||||
entity_type,
|
||||
entity_id,
|
||||
window_duration,
|
||||
)
|
||||
return None
|
||||
|
||||
if row is None:
|
||||
return None
|
||||
if _as_utc(row.window_start) < _as_utc(expected_window_start):
|
||||
return None
|
||||
return float(row.spend or 0.0)
|
||||
|
||||
@staticmethod
|
||||
async def window_from_db(
|
||||
prisma_client: Optional["PrismaClient"],
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_duration: str | None,
|
||||
window_start: datetime,
|
||||
) -> float | None:
|
||||
"""
|
||||
Authoritative window spend: the maintained row first, falling back to
|
||||
the spend-logs aggregate only when no current row exists.
|
||||
|
||||
The aggregate range-scans an unindexed table, so it must stay a
|
||||
transitional path (window configured before the row existed) rather
|
||||
than a steady-state read.
|
||||
"""
|
||||
if window_duration is not None:
|
||||
from_table: Final = await SpendCounterReseed.window_from_table(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
window_duration=window_duration,
|
||||
expected_window_start=window_start,
|
||||
)
|
||||
if from_table is not None:
|
||||
return from_table
|
||||
return await SpendCounterReseed.window_from_spend_logs(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
window_start=window_start,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def window_from_spend_logs(
|
||||
prisma_client: Optional["PrismaClient"],
|
||||
|
|
@ -215,20 +324,13 @@ class SpendCounterReseed:
|
|||
if prisma_client is None:
|
||||
return None
|
||||
|
||||
if entity_type == "Key":
|
||||
group_field = "api_key"
|
||||
where = {
|
||||
"api_key": entity_id,
|
||||
"startTime": {"gte": window_start},
|
||||
}
|
||||
elif entity_type == "Team":
|
||||
group_field = "team_id"
|
||||
where = {
|
||||
"team_id": entity_id,
|
||||
"startTime": {"gte": window_start},
|
||||
}
|
||||
else:
|
||||
group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type)
|
||||
if group_field is None:
|
||||
return None
|
||||
where: Final = {
|
||||
group_field: entity_id,
|
||||
"startTime": {"gte": window_start},
|
||||
}
|
||||
|
||||
try:
|
||||
response: Final = await SpendLogsRepository(prisma_client).table.group_by(
|
||||
|
|
@ -258,6 +360,7 @@ class SpendCounterReseed:
|
|||
counter_key: str,
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_duration: str | None,
|
||||
window_start: datetime,
|
||||
) -> float | None:
|
||||
lock: Final = await SpendCounterReseed._get_lock(counter_key)
|
||||
|
|
@ -276,10 +379,11 @@ class SpendCounterReseed:
|
|||
if val is not None:
|
||||
return float(val)
|
||||
|
||||
window_spend: Final = await SpendCounterReseed.window_from_spend_logs(
|
||||
window_spend: Final = await SpendCounterReseed.window_from_db(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
)
|
||||
if window_spend is None:
|
||||
|
|
|
|||
|
|
@ -587,7 +587,7 @@ async def _update_database_and_spend_counters(
|
|||
model_access_groups: Sequence[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,
|
||||
|
|
@ -623,6 +623,8 @@ 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,
|
||||
request_started_at=start_time,
|
||||
model_access_groups=model_access_groups,
|
||||
)
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -398,6 +398,9 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
|
|||
SPEND_LOG_CLEANUP_BOUND_SETTINGS,
|
||||
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,
|
||||
|
|
@ -2433,6 +2436,7 @@ async def get_current_spend(
|
|||
max_budget: float | None = None,
|
||||
window_entity_type: str | None = None,
|
||||
window_entity_id: str | None = None,
|
||||
window_duration: str | None = None,
|
||||
window_start: datetime | None = None,
|
||||
fallback_authoritative: bool = False,
|
||||
) -> float:
|
||||
|
|
@ -2457,7 +2461,8 @@ async def get_current_spend(
|
|||
runs and a key can leak spend past ``max_budget`` indefinitely. The
|
||||
authoritative source depends on the counter: primary key/team/user/org
|
||||
counters read the DB row; per-window counters (``window_start`` supplied)
|
||||
aggregate spend logs; end-user/tag counters have no DB row, so the caller's
|
||||
read the maintained window-spend row and only aggregate spend logs when
|
||||
that row is missing or stale; end-user/tag counters have no DB row, so the caller's
|
||||
``fallback_spend`` (loaded fresh in auth) is authoritative. The DB read is
|
||||
skipped for healthy primary counters (counter at or above recorded spend)
|
||||
and cached in-process for a few seconds, so a persistently stale counter
|
||||
|
|
@ -2482,6 +2487,7 @@ async def get_current_spend(
|
|||
counter_key=counter_key,
|
||||
window_entity_type=window_entity_type,
|
||||
window_entity_id=window_entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
)
|
||||
if authoritative is not None:
|
||||
|
|
@ -2563,6 +2569,7 @@ async def _authoritative_floor_spend(
|
|||
counter_key: str,
|
||||
window_entity_type: str | None = None,
|
||||
window_entity_id: str | None = None,
|
||||
window_duration: str | None = None,
|
||||
window_start: datetime | None = None,
|
||||
) -> float | None:
|
||||
marker_key: Final = f"spend_db_floor:{counter_key}"
|
||||
|
|
@ -2577,10 +2584,11 @@ async def _authoritative_floor_spend(
|
|||
and window_entity_id is not None
|
||||
and window_start is not None
|
||||
):
|
||||
db_spend = await SpendCounterReseed.window_from_spend_logs(
|
||||
db_spend = await SpendCounterReseed.window_from_db(
|
||||
prisma_client=prisma_client,
|
||||
entity_type=window_entity_type,
|
||||
entity_id=window_entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
)
|
||||
if db_spend is None:
|
||||
|
|
@ -2650,6 +2658,8 @@ 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,
|
||||
request_started_at: datetime | None = None,
|
||||
model_access_groups: Sequence[str] | None = None,
|
||||
):
|
||||
"""
|
||||
|
|
@ -2704,15 +2714,28 @@ async def increment_spend_counters(
|
|||
return
|
||||
for window in key_budget_limits:
|
||||
duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration
|
||||
key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at
|
||||
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_duration=duration,
|
||||
window_start=key_window_start,
|
||||
increment=cost,
|
||||
)
|
||||
await _enqueue_window_spend_row_update(
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id=hashed_token,
|
||||
reset_at=key_window_reset_at,
|
||||
window_duration=duration,
|
||||
window_start=key_window_start,
|
||||
increment=cost,
|
||||
request_id=request_id,
|
||||
request_started_at=request_started_at,
|
||||
)
|
||||
|
||||
async def _team_scope(scope_team_id: str) -> None:
|
||||
team_counter_key: Final = f"spend:team:{scope_team_id}"
|
||||
|
|
@ -2735,15 +2758,28 @@ async def increment_spend_counters(
|
|||
return
|
||||
for window in team_budget_limits:
|
||||
duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration
|
||||
team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at
|
||||
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_duration=duration,
|
||||
window_start=team_window_start,
|
||||
increment=cost,
|
||||
)
|
||||
await _enqueue_window_spend_row_update(
|
||||
entity_type=Litellm_EntityType.TEAM,
|
||||
entity_id=scope_team_id,
|
||||
reset_at=team_window_reset_at,
|
||||
window_duration=duration,
|
||||
window_start=team_window_start,
|
||||
increment=cost,
|
||||
request_id=request_id,
|
||||
request_started_at=request_started_at,
|
||||
)
|
||||
|
||||
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}"
|
||||
|
|
@ -2962,10 +2998,62 @@ 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,
|
||||
reset_at: datetime | str | None,
|
||||
window_duration: str,
|
||||
window_start: datetime | None,
|
||||
increment: float,
|
||||
request_id: str | None,
|
||||
request_started_at: datetime | 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 and
|
||||
request_started_at its startTime; the flush uses them 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 or 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,
|
||||
started_at=request_started_at,
|
||||
)
|
||||
)
|
||||
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,
|
||||
entity_id: str,
|
||||
window_duration: str | None,
|
||||
window_start: datetime | None,
|
||||
increment: float,
|
||||
):
|
||||
|
|
@ -2980,6 +3068,7 @@ async def _init_and_increment_window_spend_counter(
|
|||
counter_key=counter_key,
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
)
|
||||
if initialized is False:
|
||||
|
|
@ -3025,6 +3114,7 @@ async def _ensure_window_spend_counter_initialized(
|
|||
counter_key: str,
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_duration: str | None,
|
||||
window_start: datetime,
|
||||
) -> bool:
|
||||
is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key)
|
||||
|
|
@ -3037,6 +3127,7 @@ async def _ensure_window_spend_counter_initialized(
|
|||
counter_key=counter_key,
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
)
|
||||
if window_spend is None:
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ class _BudgetCounter:
|
|||
entity_id: str
|
||||
source_cache_key: str | None = None
|
||||
spend_log_entity_id: str | None = None
|
||||
window_duration: str | None = None
|
||||
window_start: datetime | None = None
|
||||
|
||||
|
||||
|
|
@ -709,6 +710,7 @@ def _get_budget_limit_counters(
|
|||
entity_type=entity_type,
|
||||
entity_id=f"{entity_id}:{budget_duration}",
|
||||
spend_log_entity_id=entity_id,
|
||||
window_duration=str(budget_duration),
|
||||
window_start=window_start,
|
||||
)
|
||||
)
|
||||
|
|
@ -752,6 +754,7 @@ async def _reserve_counter(
|
|||
counter_key=counter.counter_key,
|
||||
entity_type=counter.entity_type,
|
||||
entity_id=counter.spend_log_entity_id,
|
||||
window_duration=counter.window_duration,
|
||||
window_start=counter.window_start,
|
||||
)
|
||||
if initialized is False:
|
||||
|
|
|
|||
|
|
@ -68,6 +68,10 @@ class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs
|
|||
table_name = "litellm_spendlogs"
|
||||
|
||||
|
||||
class BudgetWindowSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_BudgetWindowSpend"]):
|
||||
table_name = "litellm_budgetwindowspend"
|
||||
|
||||
|
||||
class ClaudeCodePluginRepository(PrismaTableRepository["prisma_models.LiteLLM_ClaudeCodePluginTable"]):
|
||||
table_name = "litellm_claudecodeplugintable"
|
||||
|
||||
|
|
|
|||
|
|
@ -40,6 +40,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):
|
||||
|
|
|
|||
|
|
@ -833,6 +833,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)
|
||||
|
||||
|
|
@ -904,6 +905,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()
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import json
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -22,9 +23,7 @@ def redis_update_buffer(mock_redis_cache):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_uses_pipeline(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache):
|
||||
"""
|
||||
Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once
|
||||
with the correct operations and skips empty queues.
|
||||
|
|
@ -33,35 +32,29 @@ async def test_store_in_memory_spend_updates_uses_pipeline(
|
|||
|
||||
# Create mock queues - only 3 of 6 have data
|
||||
spend_update_queue = AsyncMock()
|
||||
spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = (
|
||||
AsyncMock(return_value={"key_list_transactions": {"key1": 1.0}})
|
||||
spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
|
||||
return_value={"key_list_transactions": {"key1": 1.0}}
|
||||
)
|
||||
|
||||
daily_spend_queue = AsyncMock()
|
||||
daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = (
|
||||
AsyncMock(return_value={"user_key1": {"spend": 1.0}})
|
||||
daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={"user_key1": {"spend": 1.0}}
|
||||
)
|
||||
|
||||
daily_team_queue = AsyncMock()
|
||||
daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = (
|
||||
AsyncMock(return_value={"team_key1": {"spend": 2.0}})
|
||||
daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={"team_key1": {"spend": 2.0}}
|
||||
)
|
||||
|
||||
# Empty queues
|
||||
daily_org_queue = AsyncMock()
|
||||
daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = (
|
||||
AsyncMock(return_value={})
|
||||
)
|
||||
daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={})
|
||||
|
||||
daily_end_user_queue = AsyncMock()
|
||||
daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = (
|
||||
AsyncMock(return_value=None)
|
||||
)
|
||||
daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value=None)
|
||||
|
||||
daily_agent_queue = AsyncMock()
|
||||
daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = (
|
||||
AsyncMock(return_value={})
|
||||
)
|
||||
daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={})
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=spend_update_queue,
|
||||
|
|
@ -82,9 +75,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_update_buffer, mock_redis_cache):
|
||||
"""
|
||||
If async_rpush_pipeline raises, the already-drained transactions must be
|
||||
put back into the in-memory queues so the next scheduler tick retries.
|
||||
|
|
@ -98,9 +89,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
|
|||
SpendUpdateQueue,
|
||||
)
|
||||
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(
|
||||
side_effect=ConnectionError("redis went away")
|
||||
)
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away"))
|
||||
|
||||
spend_queue = SpendUpdateQueue()
|
||||
daily_user_queue = DailySpendUpdateQueue()
|
||||
|
|
@ -145,16 +134,12 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
|
|||
|
||||
# After restore, the main spend queue should hold one item per
|
||||
# (entity_type, entity_id) pair with the aggregated cost
|
||||
restored_spend = (
|
||||
await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
|
||||
)
|
||||
restored_spend = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
|
||||
assert restored_spend["key_list_transactions"] == {"key-abc": 1.5}
|
||||
assert restored_spend["team_list_transactions"] == {"team-xyz": 2.5}
|
||||
|
||||
# Daily user queue should hold the same aggregated dict
|
||||
restored_daily = (
|
||||
await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
)
|
||||
restored_daily = await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
assert restored_daily == {
|
||||
"user1_day_model": {
|
||||
"spend": 1.0,
|
||||
|
|
@ -165,9 +150,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_all_empty_returns_early(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_update_buffer, mock_redis_cache):
|
||||
"""
|
||||
When all queues are empty, pipeline should never be called.
|
||||
"""
|
||||
|
|
@ -175,13 +158,9 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(
|
|||
|
||||
# All queues return empty
|
||||
empty_queue = AsyncMock()
|
||||
empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
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={})
|
||||
)
|
||||
empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={})
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=empty_queue,
|
||||
|
|
@ -196,14 +175,13 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_transactions_from_redis_buffer_pipeline(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache):
|
||||
"""
|
||||
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},
|
||||
|
|
@ -217,6 +195,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=[
|
||||
|
|
@ -226,13 +216,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
|
||||
|
|
@ -255,6 +262,10 @@ 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
|
||||
|
|
@ -262,13 +273,11 @@ 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)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restore_transactions_to_redis_pushes_only_provided(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache):
|
||||
"""
|
||||
restore_transactions_to_redis re-pushes only the transaction sets it was
|
||||
given, to their matching buffer keys, so uncommitted spend can be retried.
|
||||
|
|
@ -302,9 +311,42 @@ async def test_restore_transactions_to_redis_pushes_only_provided(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restore_transactions_to_redis_noop_when_empty(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache):
|
||||
"""A window commit that fails after the destructive lpop must be re-pushed
|
||||
in the store path's encoding, so the next drain returns the same increments."""
|
||||
from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
build_window_spend_transaction,
|
||||
)
|
||||
|
||||
window_transactions = (
|
||||
build_window_spend_transaction(
|
||||
entity_type="key",
|
||||
entity_id="hashed-token",
|
||||
window_duration="30d",
|
||||
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
spend=3.0,
|
||||
request_id="req-1",
|
||||
started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc),
|
||||
),
|
||||
)
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1])
|
||||
|
||||
await redis_update_buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions)
|
||||
|
||||
rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"]
|
||||
assert [op["key"] for op in rpush_list] == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY]
|
||||
|
||||
mock_redis_cache.async_lpop_pipeline = AsyncMock(
|
||||
return_value=[None, None, None, None, None, None, list(rpush_list[0]["values"])]
|
||||
)
|
||||
drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
assert drained[6] == window_transactions
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache):
|
||||
"""Nothing to restore -> no Redis call."""
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock()
|
||||
await redis_update_buffer.restore_transactions_to_redis()
|
||||
|
|
@ -312,15 +354,11 @@ async def test_restore_transactions_to_redis_noop_when_empty(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restore_transactions_to_redis_swallows_redis_error(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache):
|
||||
"""A Redis failure during restore must not propagate to the caller's finally block."""
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(
|
||||
side_effect=RedisError("redis down")
|
||||
)
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down"))
|
||||
|
||||
await redis_update_buffer.restore_transactions_to_redis(
|
||||
db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}},
|
||||
|
|
@ -433,3 +471,108 @@ 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",
|
||||
started_at=datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
)
|
||||
|
||||
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"],
|
||||
"started_at": "2026-08-10T12:00:00.000000",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@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,269 @@
|
|||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
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,
|
||||
started_at: datetime | 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,
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
|
||||
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",),
|
||||
"started_at": None,
|
||||
}
|
||||
|
||||
|
||||
def test_build_window_spend_transaction_stores_started_at_as_naive_utc_iso():
|
||||
"""started_at is compared against LiteLLM_SpendLogs.startTime, which the
|
||||
spend log writer stores after converting the request start to UTC."""
|
||||
non_utc = datetime(2026, 8, 10, 8, 30, 15, 123456, tzinfo=timezone(timedelta(hours=-4)))
|
||||
|
||||
assert _txn("k1", WINDOW_A, 1.0, started_at=non_utc)["started_at"] == "2026-08-10T12:30:15.123456"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aggregation_keeps_the_earliest_started_at_of_the_batch():
|
||||
"""The seed bounds its request-id exclusion at the batch's earliest start,
|
||||
so a later start must never win the merge."""
|
||||
queue = WindowSpendUpdateQueue()
|
||||
earliest = datetime(2026, 8, 10, 12, 0, 0, tzinfo=timezone.utc)
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-2", started_at=earliest + timedelta(seconds=5)))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-1", started_at=earliest))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, request_id="req-3"))
|
||||
|
||||
aggregated = await queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
|
||||
assert len(aggregated) == 1
|
||||
assert aggregated[0]["started_at"] == "2026-08-10T12:00:00.000000"
|
||||
assert aggregated[0]["request_ids"] == ("req-1", "req-2", "req-3")
|
||||
|
||||
|
||||
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
|
||||
596
tests/test_litellm/proxy/db/test_budget_window_spend_writer.py
Normal file
596
tests/test_litellm/proxy/db/test_budget_window_spend_writer.py
Normal file
|
|
@ -0,0 +1,596 @@
|
|||
import math
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
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)
|
||||
BATCH_STARTED_AT = datetime(2026, 8, 10, 12, 0, 0, 250_000, tzinfo=timezone.utc)
|
||||
BEFORE_BATCH = BATCH_STARTED_AT - timedelta(hours=1)
|
||||
|
||||
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,
|
||||
exclude_started_at: datetime | None,
|
||||
) -> float | None:
|
||||
self.calls.append(
|
||||
{
|
||||
"entity_type": entity_type,
|
||||
"entity_id": entity_id,
|
||||
"window_start": window_start,
|
||||
"exclude_request_ids": tuple(exclude_request_ids),
|
||||
"exclude_started_at": exclude_started_at,
|
||||
}
|
||||
)
|
||||
return self.value
|
||||
|
||||
|
||||
class _SpendLogsFake:
|
||||
"""Sums the LiteLLM_SpendLogs rows (request_id, spend, startTime) it holds,
|
||||
honouring the exclusion exactly as the real aggregate's
|
||||
NOT (request_id = ANY(...) AND startTime >= bound) does."""
|
||||
|
||||
def __init__(self, rows: tuple[tuple[str, float, datetime], ...]) -> None:
|
||||
self.rows = rows
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
prisma_client: Any,
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
exclude_request_ids: Any,
|
||||
exclude_started_at: datetime | None,
|
||||
) -> float | None:
|
||||
excluded = frozenset(exclude_request_ids) if exclude_started_at is not None else frozenset()
|
||||
return math.fsum(
|
||||
spend
|
||||
for request_id, spend, started_at in self.rows
|
||||
if not (request_id in excluded and started_at >= exclude_started_at)
|
||||
)
|
||||
|
||||
|
||||
def _batch(request_ids: tuple[str, ...], spend: float, started_at: datetime | None = BATCH_STARTED_AT) -> dict:
|
||||
return {
|
||||
"entity_type": "key",
|
||||
"entity_id": "k1",
|
||||
"window_duration": "30d",
|
||||
"window_start": "2026-08-01T00:00:00.000000",
|
||||
"spend": spend,
|
||||
"request_ids": request_ids,
|
||||
"started_at": None
|
||||
if started_at is None
|
||||
else started_at.replace(tzinfo=None).isoformat(timespec="microseconds"),
|
||||
}
|
||||
|
||||
|
||||
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, exclude_started_at
|
||||
):
|
||||
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, exclude_started_at):
|
||||
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_and_earliest_start_to_exclude():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
aggregate = _RecordingAggregate(value=0.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("req-1", "req-2", "req-3"), 3.0),),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
assert aggregate.calls[0]["exclude_request_ids"] == ("req-1", "req-2", "req-3")
|
||||
assert aggregate.calls[0]["exclude_started_at"] == BATCH_STARTED_AT
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_passes_no_start_bound_when_the_batch_has_none():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
aggregate = _RecordingAggregate(value=0.0)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("req-1",), 1.0, started_at=None),),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
assert aggregate.calls[0]["exclude_started_at"] is None
|
||||
|
||||
|
||||
@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, BATCH_STARTED_AT),
|
||||
("req-2", 0.000047, BATCH_STARTED_AT + timedelta(seconds=1)),
|
||||
("req-3", 0.000047, BATCH_STARTED_AT + timedelta(seconds=2)),
|
||||
),
|
||||
)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),),
|
||||
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, BEFORE_BATCH), ("req-1", 0.000047, BATCH_STARTED_AT)))
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("req-1",), 0.000047),),
|
||||
spend_logs_aggregate=spend_logs,
|
||||
)
|
||||
|
||||
((_, params),) = db.batcher.calls
|
||||
assert params[INSERT_SPEND] == pytest.approx(0.500047)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replayed_request_id_cannot_erase_historical_spend_from_the_seed():
|
||||
"""request_id can be chosen by the client via x-litellm-call-id. A request
|
||||
that replays an id from before this batch writes no new LiteLLM_SpendLogs
|
||||
row (the insert skips duplicates), so the seed must keep counting the
|
||||
historical row that id belongs to; only its increment is new."""
|
||||
db = _FakeDB(existing_rows=[])
|
||||
spend_logs = _SpendLogsFake(rows=(("replayed", 0.5, BEFORE_BATCH),))
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("replayed",), 0.000047),),
|
||||
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=(_batch(("req-1", "req-2", "req-3"), 0.000141),),
|
||||
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_only_within_the_batch_start_bound(
|
||||
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"),
|
||||
exclude_started_at=BATCH_STARTED_AT,
|
||||
)
|
||||
|
||||
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[]) AND \"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC'))" in normalized
|
||||
assert 'FROM "LiteLLM_SpendLogs"' in normalized
|
||||
# startTime is TIMESTAMP(3): the bound is floored to the second so the
|
||||
# batch's own earliest row cannot round under it.
|
||||
assert params == ("e1", WINDOW_A, ("req-1", "req-2"), datetime(2026, 8, 10, 12, 0, 0))
|
||||
# The ids are bound, never spliced into the statement.
|
||||
assert "req-1" not in query
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"exclude_request_ids, exclude_started_at",
|
||||
[(("req-1",), None), ((), BATCH_STARTED_AT)],
|
||||
)
|
||||
async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_bound(
|
||||
exclude_request_ids, exclude_started_at
|
||||
):
|
||||
"""Ids without a start bound would reopen the replayed-id hole, so the
|
||||
seed counts everything instead; at worst that over-counts one batch."""
|
||||
db = _FakeDB(existing_rows=[{"total": 1.25}])
|
||||
|
||||
total = await spend_logs_total_excluding(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
entity_type="key",
|
||||
entity_id="e1",
|
||||
window_start=WINDOW_A,
|
||||
exclude_request_ids=exclude_request_ids,
|
||||
exclude_started_at=exclude_started_at,
|
||||
)
|
||||
|
||||
assert total == pytest.approx(1.25)
|
||||
((query, params),) = db.query_raw_calls
|
||||
assert "request_id" not in query
|
||||
assert params == ("e1", WINDOW_A)
|
||||
|
||||
|
||||
@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=(),
|
||||
exclude_started_at=None,
|
||||
)
|
||||
|
||||
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=(),
|
||||
exclude_started_at=None,
|
||||
)
|
||||
|
||||
assert total == 0.0
|
||||
|
|
@ -4,8 +4,8 @@ import json
|
|||
import re
|
||||
|
||||
|
||||
|
||||
from collections.abc import Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
|
||||
|
|
@ -15,6 +15,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
|
||||
|
|
@ -64,9 +67,7 @@ async def test_daily_spend_tracking_with_disabled_spend_logs():
|
|||
assert db_writer.add_spend_log_transaction_to_daily_user_transaction.called
|
||||
|
||||
# Verify the payload passed to add_spend_log_transaction_to_daily_user_transaction
|
||||
call_args = (
|
||||
db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1]
|
||||
)
|
||||
call_args = db_writer.add_spend_log_transaction_to_daily_user_transaction.call_args[1]
|
||||
assert "payload" in call_args
|
||||
assert call_args["payload"]["spend"] == 0.1
|
||||
assert call_args["payload"]["model"] == "gpt-4"
|
||||
|
|
@ -406,7 +407,7 @@ async def test_update_daily_spend_sorting():
|
|||
# fields, but entity_id is sufficient to test sorting.
|
||||
daily_spend_transactions = {
|
||||
f"test_key_{i}": {
|
||||
"user_id": f"user{60-i}", # user60 ... user11, reverse order
|
||||
"user_id": f"user{60 - i}", # user60 ... user11, reverse order
|
||||
"date": "2024-01-01",
|
||||
"api_key": "test-api-key",
|
||||
"model": "gpt-4",
|
||||
|
|
@ -985,9 +986,9 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
|
|||
transaction_dict = call[1]["update"]
|
||||
# Each transaction should have one key with the format tag_date_api_key_model_provider
|
||||
for key, transaction in transaction_dict.items():
|
||||
assert (
|
||||
transaction["request_id"] == request_id
|
||||
), f"request_id should be {request_id} but got {transaction.get('request_id')}"
|
||||
assert transaction["request_id"] == request_id, (
|
||||
f"request_id should be {request_id} but got {transaction.get('request_id')}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1213,21 +1214,15 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_calls_common
|
|||
}
|
||||
|
||||
writer.daily_agent_spend_update_queue.add_update = AsyncMock()
|
||||
original_common_helper = (
|
||||
writer._common_add_spend_log_transaction_to_daily_transaction
|
||||
)
|
||||
writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock(
|
||||
wraps=original_common_helper
|
||||
)
|
||||
original_common_helper = writer._common_add_spend_log_transaction_to_daily_transaction
|
||||
writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock(wraps=original_common_helper)
|
||||
|
||||
await writer.add_spend_log_transaction_to_daily_agent_transaction(
|
||||
payload=payload,
|
||||
prisma_client=mock_prisma,
|
||||
)
|
||||
|
||||
assert (
|
||||
writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1
|
||||
)
|
||||
assert writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1382,6 +1377,7 @@ async def test_update_daily_spend_re_raises_exception_after_logging():
|
|||
Test that when batch upsert fails, the exception is properly re-raised after logging.
|
||||
This ensures that error handling continues to work correctly upstream.
|
||||
"""
|
||||
|
||||
def raise_connection_lost():
|
||||
raise ValueError("Database connection lost")
|
||||
|
||||
|
|
@ -1562,9 +1558,7 @@ async def test_update_database_creates_single_task():
|
|||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"),
|
||||
patch(
|
||||
"litellm.proxy.db.db_spend_update_writer.asyncio.create_task"
|
||||
) as mock_create_task,
|
||||
patch("litellm.proxy.db.db_spend_update_writer.asyncio.create_task") as mock_create_task,
|
||||
):
|
||||
await db_writer.update_database(
|
||||
token="test-token",
|
||||
|
|
@ -1663,9 +1657,7 @@ async def test_daily_agent_receives_deepcopied_payload():
|
|||
db_writer._update_agent_db = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(
|
||||
side_effect=capture_agent_payload
|
||||
)
|
||||
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(side_effect=capture_agent_payload)
|
||||
db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock()
|
||||
|
|
@ -1727,8 +1719,8 @@ async def test_commit_spend_updates_uses_pipeline():
|
|||
mock_redis_update_buffer = AsyncMock()
|
||||
mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock()
|
||||
# Return all-None tuple (no data to commit); the pipeline yields 6 slots
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = (
|
||||
AsyncMock(return_value=(None, None, None, None, None, None))
|
||||
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
|
||||
|
||||
|
|
@ -1782,7 +1774,7 @@ async def test_commit_with_redis_requeues_all_on_db_failure():
|
|||
|
||||
mock_redis_update_buffer = AsyncMock()
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
|
||||
return_value=(db_spend, daily_user, None, None, None, None)
|
||||
return_value=(db_spend, daily_user, None, None, None, None, None)
|
||||
)
|
||||
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
|
||||
db_writer.redis_update_buffer = mock_redis_update_buffer
|
||||
|
|
@ -1837,7 +1829,7 @@ async def test_commit_with_redis_only_requeues_failed_category():
|
|||
|
||||
mock_redis_update_buffer = AsyncMock()
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
|
||||
return_value=(db_spend, daily_user, None, None, None, None)
|
||||
return_value=(db_spend, daily_user, None, None, None, None, None)
|
||||
)
|
||||
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
|
||||
db_writer.redis_update_buffer = mock_redis_update_buffer
|
||||
|
|
@ -1885,7 +1877,7 @@ async def test_commit_with_redis_no_requeue_on_success():
|
|||
|
||||
mock_redis_update_buffer = AsyncMock()
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
|
||||
return_value=(db_spend, None, None, None, None, None)
|
||||
return_value=(db_spend, None, None, None, None, None, None)
|
||||
)
|
||||
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
|
||||
db_writer.redis_update_buffer = mock_redis_update_buffer
|
||||
|
|
@ -2156,9 +2148,7 @@ async def test_update_database_does_not_deepcopy_on_request_path():
|
|||
db_writer._update_org_db = AsyncMock()
|
||||
db_writer._update_tag_db = AsyncMock()
|
||||
db_writer._update_agent_db = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock(
|
||||
side_effect=capture_batch_payload
|
||||
)
|
||||
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock(side_effect=capture_batch_payload)
|
||||
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock()
|
||||
|
|
@ -2250,9 +2240,7 @@ async def test_spend_update_path_never_queries_user_cache_with_none_user_id():
|
|||
db_writer = DBSpendUpdateWriter()
|
||||
|
||||
strict_redis_backed_cache = MagicMock()
|
||||
strict_redis_backed_cache.async_get_cache = AsyncMock(
|
||||
side_effect=DataError("Invalid input of type: 'NoneType'")
|
||||
)
|
||||
strict_redis_backed_cache.async_get_cache = AsyncMock(side_effect=DataError("Invalid input of type: 'NoneType'"))
|
||||
|
||||
with (
|
||||
patch.object(litellm, "max_budget", 0),
|
||||
|
|
@ -2380,8 +2368,7 @@ async def test_daily_transaction_carries_compression_saved_tokens():
|
|||
cache_write_cost = model_info.get("cache_creation_input_token_cost") or input_cost
|
||||
assert transaction["compression_savings_spend"] == pytest.approx(7600 * input_cost)
|
||||
assert transaction["prompt_caching_savings_spend"] == pytest.approx(
|
||||
40 * max(input_cost - cache_read_cost, 0.0)
|
||||
- 15 * (cache_write_cost - input_cost)
|
||||
40 * max(input_cost - cache_read_cost, 0.0) - 15 * (cache_write_cost - input_cost)
|
||||
)
|
||||
assert transaction["compression_savings_spend"] > 0
|
||||
assert transaction["prompt_caching_savings_spend"] > 0
|
||||
|
|
@ -2421,6 +2408,304 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent():
|
|||
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_requeues_the_increments_and_continues_the_flush():
|
||||
"""Budget enforcement trusts a current window row without reconciling it
|
||||
against LiteLLM_SpendLogs, so a dropped increment would let the key spend
|
||||
past its limit after the next reseed. The increments must go back on the
|
||||
queue, and the tool registry flush must still run."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
transaction = 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,
|
||||
request_id="req-1",
|
||||
)
|
||||
await db_writer.window_spend_update_queue.add_update(transaction)
|
||||
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()
|
||||
requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
assert requeued == (transaction,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_window_spend_commit_from_redis_is_restored_to_redis():
|
||||
"""The Redis drain is destructive, so a failed window commit has to push
|
||||
the popped increments back exactly like the other spend categories."""
|
||||
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,
|
||||
request_id="req-1",
|
||||
),
|
||||
)
|
||||
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)
|
||||
)
|
||||
mock_redis_update_buffer.restore_transactions_to_redis = 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=True)
|
||||
db = _WindowSpendFakeDB()
|
||||
db.query_raw = AsyncMock(side_effect=Exception("connection reset"))
|
||||
|
||||
await db_writer._commit_spend_updates_to_db_with_redis(
|
||||
prisma_client=_WindowSpendFakePrisma(db),
|
||||
n_retry_times=0,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert _window_spend_upserts(db) == []
|
||||
mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once_with(
|
||||
window_spend_update_transactions=window_transactions
|
||||
)
|
||||
db_writer.pod_lock_manager.release_lock.assert_awaited_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.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam
|
||||
"litellm.proxy.proxy_server",
|
||||
disable_spend_logs=False,
|
||||
prisma_client=MagicMock(),
|
||||
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.multiple( # test-quality-ok: update_database lazily imports these proxy_server globals; no injection seam
|
||||
"litellm.proxy.proxy_server",
|
||||
disable_spend_logs=False,
|
||||
prisma_client=MagicMock(),
|
||||
litellm_proxy_budget_name="test-budget",
|
||||
),
|
||||
patch( # test-quality-ok: the payload builder is called by name inside update_database; no injection seam
|
||||
"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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at():
|
||||
"""Spend flushes must leave settings_updated_at alone, or it decays into
|
||||
|
|
@ -2721,9 +3006,7 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey
|
|||
"call_type, expects_flush",
|
||||
[("aresponses", True), ("responses", True), ("acompletion", False)],
|
||||
)
|
||||
async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(
|
||||
call_type: str, expects_flush: bool
|
||||
):
|
||||
async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(call_type: str, expects_flush: bool):
|
||||
"""
|
||||
A `previous_response_id` chained straight off the previous turn reads the DB, so a
|
||||
Responses row cannot sit in this worker's queue until the monitor's next poll.
|
||||
|
|
@ -2753,9 +3036,7 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(
|
|||
pytest.param("", True, id="injected-before-a-deployment-was-chosen"),
|
||||
],
|
||||
)
|
||||
async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected(
|
||||
injected_deployment, attributed
|
||||
):
|
||||
async def test_caching_savings_are_attributed_to_the_deployment_that_was_injected(injected_deployment, attributed):
|
||||
"""Retries, same-group failover and cross-model-group fallbacks all reuse one metadata
|
||||
bucket and one litellm_call_id, so a marker written by the leg that injected is
|
||||
visible to every sibling and nothing request-scoped can tell them apart.
|
||||
|
|
|
|||
|
|
@ -414,12 +414,14 @@ async def test_redis_buffer_requeues_access_group_transactions_as_queue_items():
|
|||
daily_org_spend_update_transactions=None,
|
||||
daily_end_user_spend_update_transactions=None,
|
||||
daily_agent_spend_update_transactions=None,
|
||||
window_spend_update_transactions=None,
|
||||
spend_update_queue=queue,
|
||||
daily_spend_update_queue=daily_queue,
|
||||
daily_team_spend_update_queue=daily_queue,
|
||||
daily_org_spend_update_queue=daily_queue,
|
||||
daily_end_user_spend_update_queue=daily_queue,
|
||||
daily_agent_spend_update_queue=daily_queue,
|
||||
window_spend_update_queue=None,
|
||||
)
|
||||
|
||||
updates = await _drain(queue)
|
||||
|
|
|
|||
250
tests/test_litellm/proxy/db/test_spend_counter_reseed.py
Normal file
250
tests/test_litellm/proxy/db/test_spend_counter_reseed.py
Normal file
|
|
@ -0,0 +1,250 @@
|
|||
"""Window-spend reads in ``SpendCounterReseed``.
|
||||
|
||||
The maintained ``LiteLLM_BudgetWindowSpend`` row replaces a per-request
|
||||
``LiteLLM_SpendLogs`` range scan, so these pin *when* the aggregate is still
|
||||
allowed to run: only when the row is missing or belongs to an older window.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
||||
|
||||
WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
class _FakeWindowSpendTable:
|
||||
def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None:
|
||||
self._row = row
|
||||
self._error = error
|
||||
self.where_clauses: list[dict] = []
|
||||
|
||||
async def find_unique(self, where: dict):
|
||||
self.where_clauses.append(where)
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return self._row
|
||||
|
||||
|
||||
class _FakeSpendLogsTable:
|
||||
def __init__(self, total: float) -> None:
|
||||
self._total = total
|
||||
self.call_count = 0
|
||||
|
||||
async def group_by(self, by: list[str], where: dict, sum: dict):
|
||||
self.call_count += 1
|
||||
return [{by[0]: where.get(by[0]), "_sum": {"spend": self._total}}]
|
||||
|
||||
|
||||
class _FakePrismaClient:
|
||||
def __init__(
|
||||
self,
|
||||
row: SimpleNamespace | None = None,
|
||||
spend_logs_total: float = 0.0,
|
||||
error: Exception | None = None,
|
||||
) -> None:
|
||||
self.db = SimpleNamespace(
|
||||
litellm_budgetwindowspend=_FakeWindowSpendTable(row=row, error=error),
|
||||
litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total),
|
||||
)
|
||||
|
||||
|
||||
def _row(window_start: datetime, spend: float) -> SimpleNamespace:
|
||||
return SimpleNamespace(window_start=window_start, spend=spend)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_from_table_reads_row_by_primary_key():
|
||||
"""The lookup must use the table's own entity_type values ("key"), not the
|
||||
"Key"/"Team" labels the counter keys and spend-log aggregates use."""
|
||||
prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5))
|
||||
|
||||
result = await SpendCounterReseed.window_from_table(
|
||||
prisma_client=prisma,
|
||||
entity_type="Key",
|
||||
entity_id="tok-1",
|
||||
window_duration="30d",
|
||||
expected_window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result == 4.5
|
||||
assert prisma.db.litellm_budgetwindowspend.where_clauses == [
|
||||
{
|
||||
"entity_type_entity_id_window_duration": {
|
||||
"entity_type": "key",
|
||||
"entity_id": "tok-1",
|
||||
"window_duration": "30d",
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_from_table_maps_team_entity_type():
|
||||
prisma = _FakePrismaClient(row=_row(WINDOW_START, 9.0))
|
||||
|
||||
result = await SpendCounterReseed.window_from_table(
|
||||
prisma_client=prisma,
|
||||
entity_type="Team",
|
||||
entity_id="team-1",
|
||||
window_duration="1d",
|
||||
expected_window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result == 9.0
|
||||
inner = prisma.db.litellm_budgetwindowspend.where_clauses[0]["entity_type_entity_id_window_duration"]
|
||||
assert inner["entity_type"] == "team"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_from_table_trusts_row_newer_than_expected_window():
|
||||
"""Regression: a pod holding a stale ``reset_at`` computes an expected start
|
||||
behind a window another pod already rolled. Trusting only an exact match
|
||||
would make it re-add the previous window's spend to the current one."""
|
||||
prisma = _FakePrismaClient(row=_row(WINDOW_START + timedelta(days=1), 2.0))
|
||||
|
||||
result = await SpendCounterReseed.window_from_table(
|
||||
prisma_client=prisma,
|
||||
entity_type="Key",
|
||||
entity_id="tok-1",
|
||||
window_duration="30d",
|
||||
expected_window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result == 2.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_from_table_rejects_row_from_previous_window():
|
||||
prisma = _FakePrismaClient(row=_row(WINDOW_START - timedelta(seconds=1), 99.0))
|
||||
|
||||
result = await SpendCounterReseed.window_from_table(
|
||||
prisma_client=prisma,
|
||||
entity_type="Key",
|
||||
entity_id="tok-1",
|
||||
window_duration="30d",
|
||||
expected_window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_from_table_treats_naive_row_timestamp_as_utc():
|
||||
"""The column is ``timestamp(3)``, so a driver that hands back a naive value
|
||||
must still compare against the tz-aware expected start."""
|
||||
prisma = _FakePrismaClient(row=_row(WINDOW_START.replace(tzinfo=None), 3.0))
|
||||
|
||||
result = await SpendCounterReseed.window_from_table(
|
||||
prisma_client=prisma,
|
||||
entity_type="Key",
|
||||
entity_id="tok-1",
|
||||
window_duration="30d",
|
||||
expected_window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result == 3.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"prisma, entity_type",
|
||||
[
|
||||
(_FakePrismaClient(row=None), "Key"),
|
||||
(_FakePrismaClient(row=_row(WINDOW_START, 1.0)), "User"),
|
||||
(_FakePrismaClient(error=RuntimeError("connection reset")), "Key"),
|
||||
(None, "Key"),
|
||||
],
|
||||
)
|
||||
async def test_window_from_table_returns_none_without_a_usable_row(prisma, entity_type):
|
||||
result = await SpendCounterReseed.window_from_table(
|
||||
prisma_client=prisma,
|
||||
entity_type=entity_type,
|
||||
entity_id="tok-1",
|
||||
window_duration="30d",
|
||||
expected_window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_from_db_prefers_the_row_over_the_spend_logs_aggregate():
|
||||
"""The aggregate range-scans an unindexed table; a current row must keep it
|
||||
from running at all."""
|
||||
prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0)
|
||||
|
||||
result = await SpendCounterReseed.window_from_db(
|
||||
prisma_client=prisma,
|
||||
entity_type="Key",
|
||||
entity_id="tok-1",
|
||||
window_duration="30d",
|
||||
window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result == 4.5
|
||||
assert prisma.db.litellm_spendlogs.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"row",
|
||||
[None, _row(WINDOW_START - timedelta(seconds=1), 99.0)],
|
||||
ids=["missing_row", "previous_window_row"],
|
||||
)
|
||||
async def test_window_from_db_falls_back_to_spend_logs(row):
|
||||
prisma = _FakePrismaClient(row=row, spend_logs_total=7.25)
|
||||
|
||||
result = await SpendCounterReseed.window_from_db(
|
||||
prisma_client=prisma,
|
||||
entity_type="Key",
|
||||
entity_id="tok-1",
|
||||
window_duration="30d",
|
||||
window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result == 7.25
|
||||
assert prisma.db.litellm_spendlogs.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_window_from_db_without_a_duration_skips_the_row_lookup():
|
||||
"""Callers that cannot name the window (no PK) keep the pre-table behavior."""
|
||||
prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=7.25)
|
||||
|
||||
result = await SpendCounterReseed.window_from_db(
|
||||
prisma_client=prisma,
|
||||
entity_type="Key",
|
||||
entity_id="tok-1",
|
||||
window_duration=None,
|
||||
window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result == 7.25
|
||||
assert prisma.db.litellm_budgetwindowspend.where_clauses == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_coalesced_window_seeds_a_cold_counter_from_the_row():
|
||||
prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0)
|
||||
cache = DualCache()
|
||||
counter_key = "spend:key:tok-1:window:30d"
|
||||
|
||||
result = await SpendCounterReseed.coalesced_window(
|
||||
prisma_client=prisma,
|
||||
spend_counter_cache=cache,
|
||||
counter_key=counter_key,
|
||||
entity_type="Key",
|
||||
entity_id="tok-1",
|
||||
window_duration="30d",
|
||||
window_start=WINDOW_START,
|
||||
)
|
||||
|
||||
assert result == 4.5
|
||||
assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5
|
||||
assert prisma.db.litellm_spendlogs.call_count == 0
|
||||
|
|
@ -567,9 +567,12 @@ 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": []}
|
||||
start_time = datetime.now()
|
||||
|
||||
await _update_database_and_spend_counters(
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
|
|
@ -581,7 +584,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
|
|||
org_id="test_org_id",
|
||||
kwargs={},
|
||||
completion_response=None,
|
||||
start_time=datetime.now(),
|
||||
start_time=start_time,
|
||||
end_time=datetime.now(),
|
||||
response_cost=0.2,
|
||||
budget_reservation=budget_reservation,
|
||||
|
|
@ -599,6 +602,8 @@ 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",
|
||||
request_started_at=start_time,
|
||||
model_access_groups=("premium",),
|
||||
)
|
||||
|
||||
|
|
@ -1879,6 +1884,61 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(
|
|||
)
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
class _FakeDeploymentLookup:
|
||||
"""Deployment lookup returning the access groups each deployment declares."""
|
||||
|
||||
|
|
|
|||
|
|
@ -271,6 +271,81 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch):
|
|||
)
|
||||
|
||||
|
||||
def _make_window_spend_prisma(row=None, spend_logs_total=0.0):
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=row)
|
||||
prisma.db.litellm_spendlogs.group_by = AsyncMock(
|
||||
return_value=[{"api_key": "tok", "_sum": {"spend": spend_logs_total}}]
|
||||
)
|
||||
return prisma
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_floors_window_against_maintained_row(monkeypatch):
|
||||
"""The floor re-check runs every few seconds per pod, so the window branch
|
||||
must read the maintained row and leave the unindexed spend-logs scan alone."""
|
||||
from datetime import timezone
|
||||
from types import SimpleNamespace
|
||||
|
||||
window_start = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
fake_prisma = _make_window_spend_prisma(
|
||||
row=SimpleNamespace(window_start=window_start, spend=15.0),
|
||||
spend_logs_total=100.0,
|
||||
)
|
||||
fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
|
||||
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
|
||||
monkeypatch.setattr(ps, "prisma_client", fake_prisma)
|
||||
|
||||
counter_key = "spend:key:tok:window:7d"
|
||||
result = await ps.get_current_spend(
|
||||
counter_key=counter_key,
|
||||
fallback_spend=0.0,
|
||||
max_budget=10.0,
|
||||
window_entity_type="Key",
|
||||
window_entity_id="tok",
|
||||
window_duration="7d",
|
||||
window_start=window_start,
|
||||
)
|
||||
|
||||
assert result == 15.0
|
||||
fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited()
|
||||
fake_cache.redis_cache.async_set_max.assert_awaited_once_with(
|
||||
key=counter_key, value=15.0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_floors_window_against_logs_when_row_stale(monkeypatch):
|
||||
"""A row left behind at a crossed window boundary must not be read as the
|
||||
current window's spend; the aggregate stays the fallback."""
|
||||
from datetime import timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
|
||||
window_start = datetime(2026, 1, 8, tzinfo=timezone.utc)
|
||||
fake_prisma = _make_window_spend_prisma(
|
||||
row=SimpleNamespace(
|
||||
window_start=window_start - timedelta(days=7), spend=999.0
|
||||
),
|
||||
spend_logs_total=15.0,
|
||||
)
|
||||
fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
|
||||
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
|
||||
monkeypatch.setattr(ps, "prisma_client", fake_prisma)
|
||||
|
||||
result = await ps.get_current_spend(
|
||||
counter_key="spend:key:tok:window:7d",
|
||||
fallback_spend=0.0,
|
||||
max_budget=10.0,
|
||||
window_entity_type="Key",
|
||||
window_entity_id="tok",
|
||||
window_duration="7d",
|
||||
window_start=window_start,
|
||||
)
|
||||
|
||||
assert result == 15.0
|
||||
fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypatch):
|
||||
"""With fail_closed_budget_enforcement on, an admit decision backed only by a
|
||||
|
|
@ -895,6 +970,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ
|
|||
counter_key="spend:key:k:window:1d",
|
||||
entity_type="Key",
|
||||
entity_id="k",
|
||||
window_duration="1d",
|
||||
window_start=datetime(2024, 1, 1),
|
||||
increment=5.0,
|
||||
)
|
||||
|
|
@ -922,6 +998,7 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva
|
|||
counter_key="spend:key:k:window:1d",
|
||||
entity_type="Key",
|
||||
entity_id="k",
|
||||
window_duration="1d",
|
||||
window_start=None,
|
||||
increment=5.0,
|
||||
)
|
||||
|
|
@ -1059,6 +1136,7 @@ async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeyp
|
|||
counter_key="spend:key:k:window:1d",
|
||||
entity_type="Key",
|
||||
entity_id="k",
|
||||
window_duration="1d",
|
||||
window_start=datetime(2024, 1, 1),
|
||||
)
|
||||
|
||||
|
|
@ -1091,6 +1169,7 @@ async def test_ensure_window_spend_counter_initialized_db_failure_invalid_return
|
|||
counter_key="spend:key:k:window:1d",
|
||||
entity_type="Key",
|
||||
entity_id="k",
|
||||
window_duration="1d",
|
||||
window_start=datetime(2024, 1, 1),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
|
|
@ -7816,6 +7817,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss():
|
|||
counter_cache = DualCache()
|
||||
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
|
||||
fake_prisma.db.litellm_spendlogs.group_by = AsyncMock(
|
||||
return_value=[{"api_key": "key-window", "_sum": {"spend": 2.25}}]
|
||||
)
|
||||
|
|
@ -7830,6 +7832,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss():
|
|||
counter_key="spend:key:key-window:window:1h",
|
||||
entity_type="Key",
|
||||
entity_id="key-window",
|
||||
window_duration="1h",
|
||||
window_start=window_start,
|
||||
increment=0.5,
|
||||
)
|
||||
|
|
@ -7933,6 +7936,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory():
|
|||
counter_cache.redis_cache = fake_redis
|
||||
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
|
||||
fake_prisma.db.litellm_spendlogs.group_by = AsyncMock(
|
||||
return_value=[{"api_key": "key-window-stale-local", "_sum": {"spend": 2.25}}]
|
||||
)
|
||||
|
|
@ -7947,6 +7951,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory():
|
|||
counter_key=counter_key,
|
||||
entity_type="Key",
|
||||
entity_id="key-window-stale-local",
|
||||
window_duration="1h",
|
||||
window_start=window_start,
|
||||
increment=0.5,
|
||||
)
|
||||
|
|
@ -7995,6 +8000,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed()
|
|||
counter_cache.redis_cache = fake_redis
|
||||
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
|
||||
fake_prisma.db.litellm_spendlogs.group_by = AsyncMock(
|
||||
return_value=[{"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}}]
|
||||
)
|
||||
|
|
@ -8009,6 +8015,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed()
|
|||
counter_key=counter_key,
|
||||
entity_type="Key",
|
||||
entity_id="key-window-concurrent-seed",
|
||||
window_duration="1h",
|
||||
window_start=window_start,
|
||||
increment=0.5,
|
||||
)
|
||||
|
|
@ -8041,6 +8048,7 @@ async def test_window_spend_counter_skips_invalid_window_start():
|
|||
counter_key="spend:key:key-invalid-window:window:not-a-duration",
|
||||
entity_type="Key",
|
||||
entity_id="key-invalid-window",
|
||||
window_duration="not-a-duration",
|
||||
window_start=None,
|
||||
increment=0.5,
|
||||
)
|
||||
|
|
@ -8068,6 +8076,7 @@ async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable():
|
|||
counter_key=counter_key,
|
||||
entity_type="Key",
|
||||
entity_id="key-window-db-unavailable",
|
||||
window_duration="1h",
|
||||
window_start=datetime.now(timezone.utc) - timedelta(hours=1),
|
||||
)
|
||||
|
||||
|
|
@ -11200,6 +11209,289 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog):
|
|||
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_carries_the_request_start_time():
|
||||
"""The seed only excludes a batch id whose LiteLLM_SpendLogs.startTime is at
|
||||
or after this, so it must be the same start the spend log was written with."""
|
||||
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",
|
||||
request_started_at=datetime(2026, 8, 10, 12, 0, 0, 500_000, tzinfo=timezone.utc),
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert enqueued[0]["started_at"] == "2026-08-10T12:00:00.500000"
|
||||
|
||||
|
||||
@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",)
|
||||
|
||||
|
||||
def _mock_startup_prisma_client(health_check_error=None, connect_error=None):
|
||||
client = MagicMock()
|
||||
client.connect = AsyncMock(side_effect=connect_error)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue