mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
refactor(proxy): bound the budget window seed by time instead of request ids
The one-time seed for a budget window row subtracted the batch's own LiteLLM_SpendLogs rows by request_id, and request_id is the client's x-litellm-call-id whenever the response carries no id of its own. Carrying that set through the queue meant an unbounded, client-controlled aggregate that the commit-failure requeue kept alive across retries. Every log row at or after a batch's earliest start is owed by an increment that still reaches the row, so summing only rows before it needs nothing from the request. That drops request_ids end to end and closes the cross-pod double count the id list could not see.
This commit is contained in:
ryan-crabbe-berri
parent
2a79a81b46
commit
abfb6adc2b
11 changed files with 104 additions and 379 deletions
|
|
@ -7,17 +7,14 @@ 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.
|
||||
Seeding a row that does not exist yet reads LiteLLM_SpendLogs once, summing
|
||||
only rows that started before the batch being flushed so neither source counts
|
||||
the same request twice. Anything at or after that cutoff is owed by an
|
||||
increment that still reaches the row, on this pod's next flush or another
|
||||
pod's, so a row 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. A request whose increment is lost before it flushes,
|
||||
which today means the pod dying, is missed by both sources and stays missing.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
|
@ -69,13 +66,13 @@ _ROLL_WINDOW_SPEND_SQL: Final = (
|
|||
_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'))"
|
||||
"AND \"startTime\" < ($3::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'))"
|
||||
"AND \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
||||
_SEED_FROM_SPEND_LOGS_KEY_UNBOUNDED_SQL: Final = (
|
||||
|
|
@ -92,8 +89,8 @@ _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.
|
||||
"""Sums LiteLLM_SpendLogs for one entity between window_start and the
|
||||
batch's earliest request.
|
||||
|
||||
Injected so the flush can be exercised without a database and so the
|
||||
expensive aggregate stays swappable.
|
||||
|
|
@ -105,21 +102,19 @@ class WindowSpendLogsAggregate(Protocol):
|
|||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
exclude_request_ids: Sequence[str],
|
||||
exclude_started_at: datetime | None,
|
||||
batch_started_at: datetime | None,
|
||||
) -> float | None: ...
|
||||
|
||||
|
||||
async def spend_logs_total_excluding(
|
||||
async def spend_logs_total_before_batch(
|
||||
prisma_client: "PrismaClient",
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
exclude_request_ids: Sequence[str],
|
||||
exclude_started_at: datetime | None,
|
||||
batch_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.
|
||||
"""LiteLLM_SpendLogs spend for one entity since window_start, stopping
|
||||
before the requests the increments being flushed already cover.
|
||||
|
||||
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
|
||||
|
|
@ -127,13 +122,11 @@ async def spend_logs_total_excluding(
|
|||
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.
|
||||
Every log row at or after the cutoff belongs to a request whose own
|
||||
increment still reaches this row, on this pod's next flush or another pod's,
|
||||
so bounding the sum by time needs nothing from the request itself. Without a
|
||||
known start the whole window is summed: 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
|
||||
|
|
@ -143,13 +136,12 @@ async def spend_logs_total_excluding(
|
|||
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
|
||||
if batch_started_at is None
|
||||
else await prisma_client.db.query_raw(
|
||||
bounded_sql,
|
||||
entity_id,
|
||||
window_start,
|
||||
tuple(exclude_request_ids),
|
||||
_exclusion_lower_bound(exclude_started_at),
|
||||
_exclusion_upper_bound(batch_started_at),
|
||||
)
|
||||
)
|
||||
if not rows:
|
||||
|
|
@ -157,7 +149,7 @@ async def spend_logs_total_excluding(
|
|||
return float(rows[0].get("total") or 0.0)
|
||||
|
||||
|
||||
def _exclusion_lower_bound(started_at: datetime) -> datetime:
|
||||
def _exclusion_upper_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)
|
||||
|
|
@ -194,8 +186,8 @@ async def _seed_base_for_missing_row(
|
|||
|
||||
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.
|
||||
the request path, and it stops before the queued increments so they are
|
||||
counted once.
|
||||
"""
|
||||
if _primary_key(transaction) in existing_primary_keys:
|
||||
return 0.0
|
||||
|
|
@ -204,8 +196,7 @@ async def _seed_base_for_missing_row(
|
|||
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),
|
||||
batch_started_at=_transaction_started_at(transaction),
|
||||
)
|
||||
return float(base or 0.0)
|
||||
|
||||
|
|
@ -241,7 +232,7 @@ def _upsert_params(
|
|||
async def commit_window_spend_updates(
|
||||
prisma_client: "PrismaClient",
|
||||
transactions: Sequence[WindowSpendTransaction],
|
||||
spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_excluding,
|
||||
spend_logs_aggregate: WindowSpendLogsAggregate = spend_logs_total_before_batch,
|
||||
) -> None:
|
||||
"""Apply aggregated window increments to LiteLLM_BudgetWindowSpend.
|
||||
|
||||
|
|
|
|||
|
|
@ -215,11 +215,7 @@ 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.
|
||||
"""
|
||||
) -> None:
|
||||
from litellm.proxy.proxy_server import (
|
||||
disable_spend_logs,
|
||||
litellm_proxy_budget_name,
|
||||
|
|
@ -310,7 +306,6 @@ 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 "
|
||||
|
|
|
|||
|
|
@ -26,17 +26,11 @@ class WindowSpendTransaction(TypedDict):
|
|||
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.
|
||||
started_at is the earliest request start in the batch. The one-time seed for
|
||||
a window that has no row yet sums only LiteLLM_SpendLogs rows that started
|
||||
before it, because the spend log writer flushes on its own ~2s poll and will
|
||||
usually have persisted this batch's rows before the window queue flushes;
|
||||
without the bound the seed and the increment would each count them.
|
||||
"""
|
||||
|
||||
entity_type: ReadOnly[str]
|
||||
|
|
@ -44,7 +38,6 @@ class WindowSpendTransaction(TypedDict):
|
|||
window_duration: ReadOnly[str]
|
||||
window_start: ReadOnly[str]
|
||||
spend: ReadOnly[float]
|
||||
request_ids: ReadOnly[Sequence[str]]
|
||||
started_at: ReadOnly[str | None]
|
||||
|
||||
|
||||
|
|
@ -72,7 +65,6 @@ def build_window_spend_transaction(
|
|||
window_duration: str,
|
||||
window_start: datetime,
|
||||
spend: float,
|
||||
request_id: str | None = None,
|
||||
started_at: datetime | None = None,
|
||||
) -> WindowSpendTransaction:
|
||||
return WindowSpendTransaction(
|
||||
|
|
@ -81,7 +73,6 @@ def build_window_spend_transaction(
|
|||
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"),
|
||||
|
|
@ -101,7 +92,6 @@ def _merge_window_spend_transactions(
|
|||
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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -587,7 +587,7 @@ async def _update_database_and_spend_counters(
|
|||
model_access_groups: Sequence[str] | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
spend_log_request_id = await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
token=user_api_key,
|
||||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
|
|
@ -623,7 +623,6 @@ 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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2658,7 +2658,6 @@ 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,
|
||||
):
|
||||
|
|
@ -2733,7 +2732,6 @@ async def increment_spend_counters(
|
|||
window_duration=duration,
|
||||
window_start=key_window_start,
|
||||
increment=cost,
|
||||
request_id=request_id,
|
||||
request_started_at=request_started_at,
|
||||
)
|
||||
|
||||
|
|
@ -2777,7 +2775,6 @@ async def increment_spend_counters(
|
|||
window_duration=duration,
|
||||
window_start=team_window_start,
|
||||
increment=cost,
|
||||
request_id=request_id,
|
||||
request_started_at=request_started_at,
|
||||
)
|
||||
|
||||
|
|
@ -3005,16 +3002,15 @@ async def _enqueue_window_spend_row_update(
|
|||
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.
|
||||
request_started_at is this request's LiteLLM_SpendLogs startTime; the flush
|
||||
stops the one-time seed there so a request its increment already covers is
|
||||
not counted twice.
|
||||
|
||||
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
|
||||
|
|
@ -3035,7 +3031,6 @@ async def _enqueue_window_spend_row_update(
|
|||
window_duration=window_duration,
|
||||
window_start=window_start,
|
||||
spend=increment,
|
||||
request_id=request_id,
|
||||
started_at=request_started_at,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -203,7 +203,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff
|
|||
"window_duration": "30d",
|
||||
"window_start": "2026-08-01T00:00:00.000000",
|
||||
"spend": 3.0,
|
||||
"request_ids": ["req-1"],
|
||||
"started_at": None,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
|
@ -233,13 +233,11 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff
|
|||
window_spend,
|
||||
) = result
|
||||
|
||||
# Budget window spend from two pods is summed per window, not overwritten,
|
||||
# and both pods' request ids reach the seed exclusion.
|
||||
# Budget window spend from two pods is summed per window, not overwritten.
|
||||
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
|
||||
|
|
@ -326,7 +324,6 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(redis_upd
|
|||
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),
|
||||
),
|
||||
)
|
||||
|
|
@ -500,7 +497,6 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up
|
|||
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),
|
||||
)
|
||||
)
|
||||
|
|
@ -526,7 +522,6 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up
|
|||
"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",
|
||||
}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -19,7 +19,6 @@ def _txn(
|
|||
spend: float,
|
||||
duration: str = "30d",
|
||||
entity_type: str = "key",
|
||||
request_id: str | None = None,
|
||||
started_at: datetime | None = None,
|
||||
):
|
||||
return build_window_spend_transaction(
|
||||
|
|
@ -28,7 +27,6 @@ def _txn(
|
|||
window_duration=duration,
|
||||
window_start=window_start,
|
||||
spend=spend,
|
||||
request_id=request_id,
|
||||
started_at=started_at,
|
||||
)
|
||||
|
||||
|
|
@ -38,13 +36,12 @@ def test_build_window_spend_transaction_stores_naive_utc_iso():
|
|||
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") == {
|
||||
assert _txn("k1", non_utc, 1.0) == {
|
||||
"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,
|
||||
}
|
||||
|
||||
|
|
@ -59,19 +56,19 @@ def test_build_window_spend_transaction_stores_started_at_as_naive_utc_iso():
|
|||
|
||||
@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."""
|
||||
"""The seed stops at the batch's earliest start, so a later start must never
|
||||
win the merge: it would push the cutoff forward and count a request the
|
||||
increments already cover."""
|
||||
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"))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, started_at=earliest + timedelta(seconds=5)))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0, started_at=earliest))
|
||||
await queue.add_update(_txn("k1", WINDOW_A, 1.0))
|
||||
|
||||
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():
|
||||
|
|
@ -208,62 +205,12 @@ def test_aggregation_survives_the_redis_json_round_trip():
|
|||
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():
|
||||
def test_started_at_survives_the_redis_json_round_trip():
|
||||
aggregated = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(
|
||||
[(_txn("k1", WINDOW_A, 1.0, request_id="req-1"),)]
|
||||
[(_txn("k1", WINDOW_A, 1.0, started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc)),)]
|
||||
)
|
||||
|
||||
reloaded = WindowSpendUpdateQueue.get_aggregated_window_spend_transactions([json.loads(json.dumps(aggregated))])
|
||||
|
||||
assert reloaded[0]["request_ids"] == ("req-1",)
|
||||
assert reloaded[0]["started_at"] == "2026-08-10T12:00:00.000000"
|
||||
assert reloaded[0]["spend"] == 1.0
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import pytest
|
|||
from litellm.proxy.db.budget_window_spend_writer import (
|
||||
commit_window_spend_updates,
|
||||
roll_window_spend_row,
|
||||
spend_logs_total_excluding,
|
||||
spend_logs_total_before_batch,
|
||||
)
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
build_window_spend_transaction,
|
||||
|
|
@ -82,16 +82,14 @@ class _RecordingAggregate:
|
|||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
exclude_request_ids: Any,
|
||||
exclude_started_at: datetime | None,
|
||||
batch_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,
|
||||
"batch_started_at": batch_started_at,
|
||||
}
|
||||
)
|
||||
return self.value
|
||||
|
|
@ -99,8 +97,8 @@ class _RecordingAggregate:
|
|||
|
||||
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."""
|
||||
honouring the cutoff exactly as the real aggregate's
|
||||
startTime < bound does."""
|
||||
|
||||
def __init__(self, rows: tuple[tuple[str, float, datetime], ...]) -> None:
|
||||
self.rows = rows
|
||||
|
|
@ -111,25 +109,22 @@ class _SpendLogsFake:
|
|||
entity_type: str,
|
||||
entity_id: str,
|
||||
window_start: datetime,
|
||||
exclude_request_ids: Any,
|
||||
exclude_started_at: datetime | None,
|
||||
batch_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)
|
||||
for _request_id, spend, started_at in self.rows
|
||||
if batch_started_at is None or started_at < batch_started_at
|
||||
)
|
||||
|
||||
|
||||
def _batch(request_ids: tuple[str, ...], spend: float, started_at: datetime | None = BATCH_STARTED_AT) -> dict:
|
||||
def _batch(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"),
|
||||
|
|
@ -345,7 +340,7 @@ async def test_unknown_entity_type_contributes_no_seed():
|
|||
db = _FakeDB(existing_rows=[])
|
||||
|
||||
async def no_such_column(
|
||||
prisma_client, entity_type, entity_id, window_start, exclude_request_ids, exclude_started_at
|
||||
prisma_client, entity_type, entity_id, window_start, batch_started_at
|
||||
):
|
||||
return None
|
||||
|
||||
|
|
@ -363,7 +358,7 @@ async def test_unknown_entity_type_contributes_no_seed():
|
|||
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):
|
||||
async def unavailable(prisma_client, entity_type, entity_id, window_start, batch_started_at):
|
||||
return None
|
||||
|
||||
await commit_window_spend_updates(
|
||||
|
|
@ -399,18 +394,17 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_receives_the_batch_request_ids_and_earliest_start_to_exclude():
|
||||
async def test_seed_receives_the_batch_earliest_start_as_its_cutoff():
|
||||
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),),
|
||||
transactions=(_batch(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
|
||||
assert aggregate.calls[0]["batch_started_at"] == BATCH_STARTED_AT
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -420,11 +414,11 @@ async def test_seed_passes_no_start_bound_when_the_batch_has_none():
|
|||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("req-1",), 1.0, started_at=None),),
|
||||
transactions=(_batch(1.0, started_at=None),),
|
||||
spend_logs_aggregate=aggregate,
|
||||
)
|
||||
|
||||
assert aggregate.calls[0]["exclude_started_at"] is None
|
||||
assert aggregate.calls[0]["batch_started_at"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -444,7 +438,7 @@ async def test_new_row_is_not_double_counted_when_the_batch_logs_already_flushed
|
|||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),),
|
||||
transactions=(_batch(0.000141),),
|
||||
spend_logs_aggregate=already_flushed,
|
||||
)
|
||||
|
||||
|
|
@ -460,7 +454,7 @@ async def test_new_row_still_covers_spend_that_predates_the_batch():
|
|||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("req-1",), 0.000047),),
|
||||
transactions=(_batch(0.000047),),
|
||||
spend_logs_aggregate=spend_logs,
|
||||
)
|
||||
|
||||
|
|
@ -469,17 +463,20 @@ async def test_new_row_still_covers_spend_that_predates_the_batch():
|
|||
|
||||
|
||||
@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."""
|
||||
async def test_seed_skips_logs_from_requests_this_batch_never_saw():
|
||||
"""A concurrent request on another pod can land its spend log before this
|
||||
pod seeds the row. Its increment is still queued over there, so the cutoff
|
||||
has to drop it from the seed even though this batch has no way to know its
|
||||
id; counting it here and again on that pod's flush is the double count the
|
||||
old id list could not catch."""
|
||||
db = _FakeDB(existing_rows=[])
|
||||
spend_logs = _SpendLogsFake(rows=(("replayed", 0.5, BEFORE_BATCH),))
|
||||
spend_logs = _SpendLogsFake(
|
||||
rows=(("older", 0.5, BEFORE_BATCH), ("other-pod", 0.25, BATCH_STARTED_AT + timedelta(seconds=1))),
|
||||
)
|
||||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("replayed",), 0.000047),),
|
||||
transactions=(_batch(0.000047),),
|
||||
spend_logs_aggregate=spend_logs,
|
||||
)
|
||||
|
||||
|
|
@ -496,7 +493,7 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet():
|
|||
|
||||
await commit_window_spend_updates(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
transactions=(_batch(("req-1", "req-2", "req-3"), 0.000141),),
|
||||
transactions=(_batch(0.000141),),
|
||||
spend_logs_aggregate=nothing_flushed,
|
||||
)
|
||||
|
||||
|
|
@ -509,57 +506,48 @@ async def test_new_row_is_correct_when_the_batch_logs_have_not_flushed_yet():
|
|||
"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
|
||||
):
|
||||
async def test_seed_aggregate_sql_stops_at_the_batch_start(entity_type, expected_column):
|
||||
db = _FakeDB(existing_rows=[{"total": 1.25}])
|
||||
|
||||
total = await spend_logs_total_excluding(
|
||||
total = await spend_logs_total_before_batch(
|
||||
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,
|
||||
batch_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 "AND \"startTime\" < ($3::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
|
||||
assert params == ("e1", WINDOW_A, datetime(2026, 8, 10, 12, 0, 0))
|
||||
# Nothing the caller supplied reaches the statement text.
|
||||
assert "e1" 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."""
|
||||
async def test_seed_aggregate_sums_the_whole_window_without_a_start_bound():
|
||||
"""A batch with no known start cannot place the cutoff, so the seed counts
|
||||
everything; at worst that over-counts one batch, which enforcement
|
||||
tolerates, where under-counting is a budget bypass."""
|
||||
db = _FakeDB(existing_rows=[{"total": 1.25}])
|
||||
|
||||
total = await spend_logs_total_excluding(
|
||||
total = await spend_logs_total_before_batch(
|
||||
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,
|
||||
batch_started_at=None,
|
||||
)
|
||||
|
||||
assert total == pytest.approx(1.25)
|
||||
((query, params),) = db.query_raw_calls
|
||||
assert "request_id" not in query
|
||||
assert '"startTime" <' not in query
|
||||
assert params == ("e1", WINDOW_A)
|
||||
|
||||
|
||||
|
|
@ -567,13 +555,12 @@ async def test_seed_aggregate_excludes_nothing_without_both_ids_and_a_start_boun
|
|||
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(
|
||||
total = await spend_logs_total_before_batch(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
entity_type="user",
|
||||
entity_id="u1",
|
||||
window_start=WINDOW_A,
|
||||
exclude_request_ids=(),
|
||||
exclude_started_at=None,
|
||||
batch_started_at=None,
|
||||
)
|
||||
|
||||
assert total is None
|
||||
|
|
@ -584,13 +571,12 @@ async def test_seed_aggregate_returns_none_for_an_entity_type_with_no_spend_logs
|
|||
async def test_seed_aggregate_treats_an_entity_with_no_rows_as_zero():
|
||||
db = _FakeDB(existing_rows=[])
|
||||
|
||||
total = await spend_logs_total_excluding(
|
||||
total = await spend_logs_total_before_batch(
|
||||
prisma_client=_FakePrismaClient(db),
|
||||
entity_type="key",
|
||||
entity_id="k-unknown",
|
||||
window_start=WINDOW_A,
|
||||
exclude_request_ids=(),
|
||||
exclude_started_at=None,
|
||||
batch_started_at=None,
|
||||
)
|
||||
|
||||
assert total == 0.0
|
||||
|
|
|
|||
|
|
@ -2581,7 +2581,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_
|
|||
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()
|
||||
|
|
@ -2611,7 +2610,6 @@ async def test_failed_window_spend_commit_from_redis_is_restored_to_redis():
|
|||
window_duration="7d",
|
||||
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
spend=2.0,
|
||||
request_id="req-1",
|
||||
),
|
||||
)
|
||||
mock_redis_update_buffer = AsyncMock()
|
||||
|
|
@ -2638,74 +2636,6 @@ async def test_failed_window_spend_commit_from_redis_is_restored_to_redis():
|
|||
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
|
||||
|
|
|
|||
|
|
@ -567,9 +567,7 @@ 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(
|
||||
return_value="chatcmpl-abc123"
|
||||
)
|
||||
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock()
|
||||
increment_spend_counters = AsyncMock()
|
||||
budget_reservation = {"reserved_cost": 0.5, "entries": []}
|
||||
start_time = datetime.now()
|
||||
|
|
@ -602,7 +600,6 @@ 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",),
|
||||
)
|
||||
|
|
@ -1884,61 +1881,6 @@ 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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -11399,35 +11399,10 @@ async def test_no_window_spend_row_enqueued_without_budget_limits():
|
|||
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."""
|
||||
"""The seed sums LiteLLM_SpendLogs only up to this point, so it must be the
|
||||
same start the spend log row was written with."""
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
reset_at = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
|
|
@ -11442,7 +11417,6 @@ async def test_window_spend_row_carries_the_request_start_time():
|
|||
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)
|
||||
|
|
@ -11451,26 +11425,7 @@ async def test_window_spend_row_carries_the_request_start_time():
|
|||
|
||||
|
||||
@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():
|
||||
async def test_team_window_spend_row_carries_the_request_start_time():
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
reset_at = datetime.now(timezone.utc) + timedelta(days=3)
|
||||
|
|
@ -11485,11 +11440,11 @@ async def test_team_window_spend_row_carries_the_request_id():
|
|||
team_id="team-1",
|
||||
user_id=None,
|
||||
response_cost=1.5,
|
||||
request_id="chatcmpl-team",
|
||||
request_started_at=datetime(2026, 8, 10, 12, 0, 0, 500_000, tzinfo=timezone.utc),
|
||||
)
|
||||
enqueued = await _drain(queue)
|
||||
|
||||
assert enqueued[0]["request_ids"] == ("chatcmpl-team",)
|
||||
assert enqueued[0]["started_at"] == "2026-08-10T12:00:00.500000"
|
||||
|
||||
|
||||
def _mock_startup_prisma_client(health_check_error=None, connect_error=None):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue