mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
revert(usage): remove the top-N aggregated key cap and daily global spend rollup code from rc/1.103.0 (#43326)
Reverts the code from #41324 (merge87694c26ef) and #41293 (merge2e46b10320). The LiteLLM_DailyGlobalSpend model and migration stay so databases that applied rc.1 keep a consistent table and rollup marker when the feature returns.
This commit is contained in:
parent
ed9ae4fc8c
commit
a78a439c68
21 changed files with 218 additions and 1877 deletions
|
|
@ -2084,16 +2084,12 @@ MCP_SPEND_LOG_MODEL_PREFIX: Final[str] = "MCP: "
|
|||
PTU_SENTINEL_API_KEY: Final[str] = "__ptu_flat_cost__"
|
||||
PTU_ROLLUP_JOB_ID: Final[str] = "ptu_flat_cost_rollup_job"
|
||||
PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900
|
||||
USAGE_TOP_API_KEYS_LIMIT: Final[int] = int(os.getenv("USAGE_TOP_API_KEYS_LIMIT", "100"))
|
||||
# Furthest back the catch-up pass looks for unpriced PTU days when a deployment
|
||||
# declares no ptu_effective_from, bounding the scan for an open-ended window.
|
||||
PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90
|
||||
# Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide
|
||||
# expiry cannot produce an alert too large for the channel delivering it.
|
||||
PTU_LAPSED_ALERT_LIMIT: Final[int] = 10
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID: Final[str] = "daily_global_spend_reconcile_job"
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS: Final[int] = 3600
|
||||
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM: Final[str] = "daily_global_spend_reconciled_through"
|
||||
# Slack allowed when deciding a sentinel row is stale. The row's updated_at and the
|
||||
# run's cutoff are stamped by different hosts, so clock skew between them must not let
|
||||
# one run delete a charge another just wrote. A stale row is hours old and a concurrent
|
||||
|
|
|
|||
|
|
@ -3050,18 +3050,6 @@
|
|||
},
|
||||
"DailySpendMetadata": {
|
||||
"properties": {
|
||||
"api_key_limit": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "When set, api_keys and every api_key_breakdown list at most this many keys, ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.",
|
||||
"title": "Api Key Limit"
|
||||
},
|
||||
"has_more": {
|
||||
"default": false,
|
||||
"title": "Has More",
|
||||
|
|
@ -3072,18 +3060,6 @@
|
|||
"title": "Page",
|
||||
"type": "integer"
|
||||
},
|
||||
"total_api_keys": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "integer"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"description": "Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key lists are truncated to the highest-spend keys.",
|
||||
"title": "Total Api Keys"
|
||||
},
|
||||
"total_api_requests": {
|
||||
"default": 0,
|
||||
"title": "Total Api Requests",
|
||||
|
|
|
|||
|
|
@ -9,9 +9,8 @@ from fastapi import HTTPException, status
|
|||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.spend_tracking.daily_global_spend_rollup import GLOBAL_SPEND_TABLE_NAME, reconciled_through
|
||||
from litellm.proxy.spend_tracking.key_metadata_recovery import (
|
||||
attach_user_emails,
|
||||
recover_double_hashed_key_metadata,
|
||||
|
|
@ -149,9 +148,15 @@ class _AggregatedSpendData(TypedDict):
|
|||
totals: SpendMetrics
|
||||
|
||||
|
||||
class _RollupMetricsRow(SimpleNamespace):
|
||||
class _GroupingSetsRow(SimpleNamespace):
|
||||
date: str
|
||||
api_key: str | None
|
||||
model: str | None
|
||||
model_group: str | None
|
||||
custom_llm_provider: str | None
|
||||
mcp_namespaced_tool_name: str | None
|
||||
endpoint: str | None
|
||||
group_level: int
|
||||
spend: float | None
|
||||
prompt_tokens: int | None
|
||||
completion_tokens: int | None
|
||||
|
|
@ -169,46 +174,12 @@ class _RollupMetricsRow(SimpleNamespace):
|
|||
timed_requests: int | None
|
||||
|
||||
|
||||
class _GroupingSetsRow(_RollupMetricsRow):
|
||||
model: str | None
|
||||
model_group: str | None
|
||||
custom_llm_provider: str | None
|
||||
mcp_namespaced_tool_name: str | None
|
||||
endpoint: str | None
|
||||
group_level: int
|
||||
distinct_api_keys: int | None
|
||||
|
||||
|
||||
class _EntityRollupRow(_RollupMetricsRow):
|
||||
class _EntityRollupRow(_GroupingSetsRow):
|
||||
entity_id: str | None
|
||||
api_key_rolled: int
|
||||
|
||||
|
||||
class _AggregatedQueryKwargs(TypedDict):
|
||||
table_name: ReadOnly[str]
|
||||
entity_id_field: ReadOnly[str]
|
||||
entity_id: ReadOnly[str | list[str] | None]
|
||||
start_date: ReadOnly[str]
|
||||
end_date: ReadOnly[str]
|
||||
model: ReadOnly[str | None]
|
||||
api_key: ReadOnly[str | list[str] | None]
|
||||
exclude_entity_ids: ReadOnly[list[str] | None]
|
||||
timezone_offset_minutes: ReadOnly[int | None]
|
||||
include_current_utc_day: ReadOnly[bool]
|
||||
|
||||
|
||||
_SqlQuery = tuple[str, list[str]]
|
||||
|
||||
|
||||
async def _query_raw_optional(
|
||||
prisma_client: PrismaClient, query: _SqlQuery | None
|
||||
) -> list[dict[str, object]] | None: # mutable-ok: prisma query_raw return shape
|
||||
if query is None:
|
||||
return None
|
||||
return await prisma_client.db.query_raw(query[0], *query[1])
|
||||
|
||||
|
||||
def _reported_flat_cost(record: DailySpendRecord | _RollupMetricsRow) -> float:
|
||||
def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float:
|
||||
"""Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled.
|
||||
|
||||
Both read paths funnel through here: the paginated path reads the ``ptu_flat_cost``
|
||||
|
|
@ -731,8 +702,71 @@ def _ptu_flat_cost_select(table_name: str) -> str:
|
|||
return "0::float AS ptu_flat_cost"
|
||||
|
||||
|
||||
def _rollup_metric_select(table_name: str) -> str:
|
||||
return f"""
|
||||
def _build_aggregated_sql_query(
|
||||
*,
|
||||
table_name: str,
|
||||
entity_id_field: str,
|
||||
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
model: str | None,
|
||||
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
|
||||
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
|
||||
timezone_offset_minutes: int | None = None,
|
||||
include_current_utc_day: bool = False,
|
||||
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
|
||||
"""Build a parameterized SQL GROUP BY query for aggregated daily activity.
|
||||
|
||||
Groups by (date, api_key, model, model_group, custom_llm_provider,
|
||||
mcp_namespaced_tool_name, endpoint) with SUMs on all metric columns.
|
||||
The entity_id column is intentionally omitted from GROUP BY to collapse
|
||||
rows across entities — this is where the biggest row reduction comes from.
|
||||
|
||||
Returns:
|
||||
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
|
||||
"""
|
||||
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
|
||||
if pg_table is None:
|
||||
raise ValueError(f"Unknown table name: {table_name}")
|
||||
|
||||
adjusted_start, adjusted_end = _adjust_dates_for_timezone(
|
||||
start_date, end_date, timezone_offset_minutes, include_current_utc_day
|
||||
)
|
||||
|
||||
where_clause, sql_params = _build_aggregated_where_clause(
|
||||
entity_id_field=entity_id_field,
|
||||
entity_id=entity_id,
|
||||
adjusted_start=adjusted_start,
|
||||
adjusted_end=adjusted_end,
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
exclude_entity_ids=exclude_entity_ids,
|
||||
)
|
||||
|
||||
# Postgres computes every rollup level the response needs — per-date
|
||||
# totals, per-(date, model), per-(date, model, api_key), per-provider,
|
||||
# etc. — in a single pass via GROUPING SETS. The GROUPING() bitmask
|
||||
# encodes which level a row belongs to so Python can dispatch rows
|
||||
# straight into their buckets without re-summing. The leaf grouping
|
||||
# is omitted on purpose: nothing in the response shape needs it once
|
||||
# all the rollups are present.
|
||||
#
|
||||
# TODO: drop the successful_requests/failed_requests aggregates (and the
|
||||
# total_successful_requests metadata they feed) once the admin UI reads SGR
|
||||
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
|
||||
# api_requests rollups are still served from here.
|
||||
sql_query: Final = f"""
|
||||
SELECT
|
||||
date,
|
||||
api_key,
|
||||
model,
|
||||
COALESCE(NULLIF(model_group, ''), model) AS model_group,
|
||||
custom_llm_provider,
|
||||
mcp_namespaced_tool_name,
|
||||
endpoint,
|
||||
GROUPING(date, api_key, model, COALESCE(NULLIF(model_group, ''), model),
|
||||
custom_llm_provider, mcp_namespaced_tool_name,
|
||||
endpoint) AS group_level,
|
||||
SUM(spend)::float AS spend,
|
||||
{_ptu_flat_cost_select(table_name)},
|
||||
SUM(prompt_tokens)::bigint AS prompt_tokens,
|
||||
|
|
@ -748,175 +782,27 @@ def _rollup_metric_select(table_name: str) -> str:
|
|||
SUM(successful_requests)::bigint AS successful_requests,
|
||||
SUM(failed_requests)::bigint AS failed_requests,
|
||||
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
|
||||
SUM(timed_requests)::bigint AS timed_requests"""
|
||||
|
||||
|
||||
_MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)"
|
||||
|
||||
|
||||
_KEY_FREE_SOURCE_COLUMNS: Final = (
|
||||
"date",
|
||||
"model",
|
||||
"model_group",
|
||||
"custom_llm_provider",
|
||||
"mcp_namespaced_tool_name",
|
||||
"endpoint",
|
||||
"spend",
|
||||
"prompt_tokens",
|
||||
"completion_tokens",
|
||||
"cache_read_input_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
"compression_saved_tokens",
|
||||
"compression_savings_spend",
|
||||
"prompt_caching_savings_spend",
|
||||
"gateway_injected_caching_savings_spend",
|
||||
"autorouter_savings_spend",
|
||||
"api_requests",
|
||||
"successful_requests",
|
||||
"failed_requests",
|
||||
"total_response_time_ms",
|
||||
"timed_requests",
|
||||
)
|
||||
|
||||
|
||||
async def global_rollup_reconciled_through(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None:
|
||||
"""The last day ``LiteLLM_DailyGlobalSpend`` can answer the key-free arm for, or None to
|
||||
read it all from the per-key table.
|
||||
|
||||
Only an unfiltered read of the user table sums to the same rows as the global table. The
|
||||
marker read is served from the config cache, so this is not a database round trip per request.
|
||||
"""
|
||||
if query["table_name"] != "litellm_dailyuserspend":
|
||||
return None
|
||||
if query["entity_id"] is not None or query["api_key"] is not None or query["exclude_entity_ids"]:
|
||||
return None
|
||||
try:
|
||||
return await reconciled_through(prisma_client)
|
||||
except Exception as exc: # noqa: BLE001 # the per-key table is always a correct answer, so never fail the read
|
||||
verbose_proxy_logger.warning("Could not read the daily global spend marker, using the per-key table: %s", exc)
|
||||
return None
|
||||
|
||||
|
||||
def _key_free_source(pg_table: str, where_clause: str, marker_param: str | None) -> str:
|
||||
"""The relation the key-free arm aggregates: the per-key table alone, or the global rollup
|
||||
for days through the marker plus the per-key table for the days still open after it."""
|
||||
if marker_param is None:
|
||||
return f'"{pg_table}"\n WHERE {where_clause}'
|
||||
columns: Final = ", ".join(_KEY_FREE_SOURCE_COLUMNS)
|
||||
return f"""(
|
||||
SELECT {columns}
|
||||
FROM "{GLOBAL_SPEND_TABLE_NAME}"
|
||||
WHERE {where_clause} AND date <= {marker_param}
|
||||
UNION ALL
|
||||
SELECT {columns}
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause} AND date > {marker_param}
|
||||
) AS key_free_source"""
|
||||
|
||||
|
||||
def _build_aggregated_sql_query(
|
||||
*,
|
||||
table_name: str,
|
||||
entity_id_field: str,
|
||||
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
model: str | None,
|
||||
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
|
||||
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
|
||||
timezone_offset_minutes: int | None = None,
|
||||
include_current_utc_day: bool = False,
|
||||
global_rollup_through: str | None = None,
|
||||
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
|
||||
"""Build the GROUPING SETS query for aggregated daily activity.
|
||||
|
||||
Returns:
|
||||
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
|
||||
"""
|
||||
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
|
||||
if pg_table is None:
|
||||
raise ValueError(f"Unknown table name: {table_name}")
|
||||
|
||||
adjusted_start, adjusted_end = _adjust_dates_for_timezone(
|
||||
start_date, end_date, timezone_offset_minutes, include_current_utc_day
|
||||
)
|
||||
|
||||
where_clause, where_params = _build_aggregated_where_clause(
|
||||
entity_id_field=entity_id_field,
|
||||
entity_id=entity_id,
|
||||
adjusted_start=adjusted_start,
|
||||
adjusted_end=adjusted_end,
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
exclude_entity_ids=exclude_entity_ids,
|
||||
)
|
||||
sentinel_param: Final = f"${len(where_params) + 1}"
|
||||
marker_param: Final = None if global_rollup_through is None else f"${len(where_params) + 2}"
|
||||
metric_select: Final = _rollup_metric_select(table_name)
|
||||
|
||||
# TODO: drop the successful_requests/failed_requests aggregates (and the
|
||||
# total_successful_requests metadata they feed) once the admin UI reads SGR
|
||||
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
|
||||
# api_requests rollups are still served from here.
|
||||
sql_query: Final = f"""
|
||||
(SELECT
|
||||
date,
|
||||
NULL::text AS api_key,
|
||||
model,
|
||||
{_MODEL_GROUP_EXPR} AS model_group,
|
||||
custom_llm_provider,
|
||||
mcp_namespaced_tool_name,
|
||||
endpoint,
|
||||
(GROUPING(date) << 6) | {_API_KEY_ROLLED_UP_BIT}
|
||||
| GROUPING(model, {_MODEL_GROUP_EXPR},
|
||||
custom_llm_provider, mcp_namespaced_tool_name,
|
||||
endpoint) AS group_level,
|
||||
NULL::bigint AS distinct_api_keys,{metric_select}
|
||||
FROM {_key_free_source(pg_table, where_clause, marker_param)}
|
||||
GROUP BY GROUPING SETS (
|
||||
(date),
|
||||
(date, model),
|
||||
(date, {_MODEL_GROUP_EXPR}),
|
||||
(date, custom_llm_provider),
|
||||
(date, mcp_namespaced_tool_name),
|
||||
(date, endpoint),
|
||||
()
|
||||
))
|
||||
UNION ALL
|
||||
(WITH top_api_keys AS (
|
||||
SELECT api_key, COUNT(*) OVER () AS distinct_api_keys
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause} AND api_key <> {sentinel_param}
|
||||
GROUP BY api_key
|
||||
ORDER BY SUM(spend) DESC, api_key
|
||||
LIMIT {USAGE_TOP_API_KEYS_LIMIT}
|
||||
)
|
||||
SELECT
|
||||
date,
|
||||
api_key,
|
||||
model,
|
||||
{_MODEL_GROUP_EXPR} AS model_group,
|
||||
custom_llm_provider,
|
||||
mcp_namespaced_tool_name,
|
||||
endpoint,
|
||||
GROUPING(date, api_key, model, {_MODEL_GROUP_EXPR},
|
||||
custom_llm_provider, mcp_namespaced_tool_name,
|
||||
endpoint) AS group_level,
|
||||
MAX(top_api_keys.distinct_api_keys) AS distinct_api_keys,{metric_select}
|
||||
FROM "{pg_table}" JOIN top_api_keys USING (api_key)
|
||||
SUM(timed_requests)::bigint AS timed_requests
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause}
|
||||
GROUP BY GROUPING SETS (
|
||||
(date),
|
||||
(date, api_key),
|
||||
(date, model),
|
||||
(date, model, api_key),
|
||||
(date, {_MODEL_GROUP_EXPR}, api_key),
|
||||
(date, COALESCE(NULLIF(model_group, ''), model)),
|
||||
(date, COALESCE(NULLIF(model_group, ''), model), api_key),
|
||||
(date, custom_llm_provider),
|
||||
(date, custom_llm_provider, api_key),
|
||||
(date, mcp_namespaced_tool_name),
|
||||
(date, mcp_namespaced_tool_name, api_key),
|
||||
(date, endpoint, api_key)
|
||||
))
|
||||
(date, endpoint),
|
||||
(date, endpoint, api_key),
|
||||
()
|
||||
)
|
||||
"""
|
||||
|
||||
marker_params: Final = () if global_rollup_through is None else (global_rollup_through,)
|
||||
return sql_query, [*where_params, PTU_SENTINEL_API_KEY, *marker_params]
|
||||
return sql_query, sql_params
|
||||
|
||||
|
||||
def _build_entity_rollup_sql_query(
|
||||
|
|
@ -961,7 +847,23 @@ def _build_entity_rollup_sql_query(
|
|||
"{entity_id_field}" AS entity_id,
|
||||
date,
|
||||
api_key,
|
||||
GROUPING(api_key) AS api_key_rolled,{_rollup_metric_select(table_name)}
|
||||
GROUPING(api_key) AS api_key_rolled,
|
||||
SUM(spend)::float AS spend,
|
||||
{_ptu_flat_cost_select(table_name)},
|
||||
SUM(prompt_tokens)::bigint AS prompt_tokens,
|
||||
SUM(completion_tokens)::bigint AS completion_tokens,
|
||||
SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens,
|
||||
SUM(cache_creation_input_tokens)::bigint AS cache_creation_input_tokens,
|
||||
SUM(compression_saved_tokens)::bigint AS compression_saved_tokens,
|
||||
SUM(compression_savings_spend)::float AS compression_savings_spend,
|
||||
SUM(prompt_caching_savings_spend)::float AS prompt_caching_savings_spend,
|
||||
SUM(gateway_injected_caching_savings_spend)::float AS gateway_injected_caching_savings_spend,
|
||||
SUM(autorouter_savings_spend)::float AS autorouter_savings_spend,
|
||||
SUM(api_requests)::bigint AS api_requests,
|
||||
SUM(successful_requests)::bigint AS successful_requests,
|
||||
SUM(failed_requests)::bigint AS failed_requests,
|
||||
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
|
||||
SUM(timed_requests)::bigint AS timed_requests
|
||||
FROM "{pg_table}"
|
||||
WHERE {where_clause}
|
||||
GROUP BY GROUPING SETS (
|
||||
|
|
@ -1063,7 +965,6 @@ async def _aggregate_spend_records(
|
|||
# current grouping set's key), 0 when the column is part of the key.
|
||||
_GROUP_GRAND_TOTAL: Final = 127 # 0b1111111 — all rolled up
|
||||
_GROUP_DATE: Final = 63 # 0b0111111 — only date kept
|
||||
_API_KEY_ROLLED_UP_BIT: Final = 32 # 0b0100000
|
||||
_GROUP_DATE_API_KEY: Final = 31 # 0b0011111
|
||||
_GROUP_DATE_MODEL: Final = 47 # 0b0101111
|
||||
_GROUP_DATE_MODEL_API_KEY: Final = 15 # 0b0001111
|
||||
|
|
@ -1077,7 +978,7 @@ _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110
|
|||
_GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110
|
||||
|
||||
|
||||
def _record_to_spend_metrics(record: _RollupMetricsRow) -> SpendMetrics:
|
||||
def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
|
||||
"""Build a SpendMetrics directly from one already-aggregated rollup row.
|
||||
|
||||
SUM() over zero rows is SQL NULL, so rollup rows (notably the grand-total
|
||||
|
|
@ -1431,6 +1332,10 @@ async def get_daily_activity_aggregated(
|
|||
) -> SpendAnalyticsPaginatedResponse:
|
||||
"""Aggregated variant that returns the full result set (no pagination).
|
||||
|
||||
Uses SQL GROUP BY to aggregate rows in the database rather than fetching
|
||||
all individual rows into Python. This collapses rows across entities
|
||||
(users/teams/orgs), reducing ~150k rows to ~2-3k grouped rows.
|
||||
|
||||
include_entity_breakdown runs a small companion rollup query and folds
|
||||
`breakdown.entities` onto the response, as entity-scoped views like Team Usage need.
|
||||
|
||||
|
|
@ -1449,7 +1354,7 @@ async def get_daily_activity_aggregated(
|
|||
)
|
||||
|
||||
try:
|
||||
query_kwargs: Final = _AggregatedQueryKwargs(
|
||||
sql_query, sql_params = _build_aggregated_sql_query(
|
||||
table_name=table_name,
|
||||
entity_id_field=entity_id_field,
|
||||
entity_id=entity_id,
|
||||
|
|
@ -1461,19 +1366,36 @@ async def get_daily_activity_aggregated(
|
|||
timezone_offset_minutes=timezone_offset_minutes,
|
||||
include_current_utc_day=include_current_utc_day,
|
||||
)
|
||||
sql_query, sql_params = _build_aggregated_sql_query(
|
||||
**query_kwargs,
|
||||
global_rollup_through=await global_rollup_reconciled_through(prisma_client, query_kwargs),
|
||||
)
|
||||
entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None
|
||||
|
||||
raw_rows, raw_entity_rows = await asyncio.gather(
|
||||
prisma_client.db.query_raw(sql_query, *sql_params),
|
||||
_query_raw_optional(prisma_client, entity_query),
|
||||
entity_query: Final = (
|
||||
_build_entity_rollup_sql_query(
|
||||
table_name=table_name,
|
||||
entity_id_field=entity_id_field,
|
||||
entity_id=entity_id,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
model=model,
|
||||
api_key=api_key,
|
||||
exclude_entity_ids=exclude_entity_ids,
|
||||
timezone_offset_minutes=timezone_offset_minutes,
|
||||
include_current_utc_day=include_current_utc_day,
|
||||
)
|
||||
if include_entity_breakdown
|
||||
else None
|
||||
)
|
||||
|
||||
records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or ())]
|
||||
total_api_keys: Final = next((r.distinct_api_keys for r in records if r.distinct_api_keys is not None), 0)
|
||||
# Execute the GROUPING SETS query (one row per rollup level), alongside
|
||||
# the per-entity companion rollup when the caller wants entities.
|
||||
raw_rows, raw_entity_rows = (
|
||||
await asyncio.gather(
|
||||
prisma_client.db.query_raw(sql_query, *sql_params),
|
||||
prisma_client.db.query_raw(entity_query[0], *entity_query[1]),
|
||||
)
|
||||
if entity_query is not None
|
||||
else (await prisma_client.db.query_raw(sql_query, *sql_params), None)
|
||||
)
|
||||
|
||||
records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or [])]
|
||||
|
||||
# The grouping-sets dispatcher places each row directly in its bucket
|
||||
# using the row's GROUPING() bitmask. No Python-side summing needed.
|
||||
|
|
@ -1527,8 +1449,6 @@ async def get_daily_activity_aggregated(
|
|||
page=1,
|
||||
total_pages=1,
|
||||
has_more=False,
|
||||
api_key_limit=USAGE_TOP_API_KEYS_LIMIT,
|
||||
total_api_keys=total_api_keys,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -265,7 +265,6 @@ from litellm.constants import (
|
|||
APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
APSCHEDULER_REPLACE_EXISTING,
|
||||
CLI_SSO_SESSION_TTL_SECONDS,
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID,
|
||||
DAYS_IN_A_MONTH,
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
|
|
@ -696,9 +695,6 @@ from litellm.proxy.route_priority import hot_routes_first
|
|||
from litellm.proxy.search_endpoints.endpoints import router as search_router
|
||||
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
|
||||
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
|
||||
from litellm.proxy.spend_tracking.daily_global_spend_rollup import (
|
||||
run_scheduled_daily_global_spend_reconcile,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import (
|
||||
PendingSpendIncrement,
|
||||
active_spend_counter_batch,
|
||||
|
|
@ -10229,12 +10225,6 @@ class ProxyStartupEvent:
|
|||
|
||||
await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler)
|
||||
|
||||
cls._initialize_daily_global_spend_reconcile_job(
|
||||
scheduler=scheduler,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
### PTU DAILY ROLLUP ###
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
||||
is_ptu_cost_attribution_enabled,
|
||||
|
|
@ -10576,39 +10566,6 @@ class ProxyStartupEvent:
|
|||
"LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _initialize_daily_global_spend_reconcile_job(
|
||||
cls,
|
||||
scheduler: AsyncIOScheduler,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
async def alert(message: str) -> None:
|
||||
await proxy_logging_obj.alerting_handler(
|
||||
message=message,
|
||||
level="High",
|
||||
alert_type=AlertType.failed_tracking_spend,
|
||||
)
|
||||
|
||||
async def reconcile() -> None:
|
||||
await run_scheduled_daily_global_spend_reconcile(
|
||||
prisma_client,
|
||||
pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager,
|
||||
alert=alert,
|
||||
)
|
||||
|
||||
scheduler.add_job(
|
||||
reconcile,
|
||||
"cron",
|
||||
hour=0,
|
||||
minute=30,
|
||||
timezone="UTC",
|
||||
id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID,
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
next_run_time=datetime.now(timezone.utc) + timedelta(minutes=2),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def _initialize_slack_alerting_jobs(
|
||||
cls,
|
||||
|
|
|
|||
|
|
@ -1,293 +0,0 @@
|
|||
"""Roll closed UTC days of ``LiteLLM_DailyUserSpend`` up into ``LiteLLM_DailyGlobalSpend``.
|
||||
|
||||
Only days that are over get rolled up, so a pod still flushing per-key spend for the current
|
||||
day can never leave the global table short; usage reads serve days through the recorded
|
||||
marker from the global table and later days live from the per-key table. Per-key rows are
|
||||
dated by request start, so spend can land on a day that was already rolled up (a flush
|
||||
straddling midnight, a retry after an outage). Each run therefore also rewrites every closed
|
||||
day that has rows touched since the previous run's scan, whatever the date. The marker lives
|
||||
in ``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on
|
||||
a large deployment the first backfill is minutes of work.
|
||||
"""
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, timedelta
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID,
|
||||
DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS,
|
||||
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend"
|
||||
# The unique constraint, in constraint order. NULL never matches itself in a unique index, so
|
||||
# every column is normalized to '' or the same group would be inserted again on every run.
|
||||
_KEY_COLUMNS: Final = ("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint")
|
||||
_METRIC_COLUMNS: Final = (
|
||||
"prompt_tokens",
|
||||
"completion_tokens",
|
||||
"cache_read_input_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
"compression_saved_tokens",
|
||||
"api_requests",
|
||||
"successful_requests",
|
||||
"failed_requests",
|
||||
"total_response_time_ms",
|
||||
"timed_requests",
|
||||
"compression_savings_spend",
|
||||
"prompt_caching_savings_spend",
|
||||
"gateway_injected_caching_savings_spend",
|
||||
"autorouter_savings_spend",
|
||||
"spend",
|
||||
)
|
||||
|
||||
|
||||
def _quoted(columns: tuple[str, ...]) -> str:
|
||||
return ", ".join(f'"{column}"' for column in columns)
|
||||
|
||||
|
||||
def _reconcile_day_sql() -> str:
|
||||
normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in _KEY_COLUMNS)
|
||||
sums: Final = ", ".join(f'SUM("{column}")' for column in _METRIC_COLUMNS)
|
||||
overwrite: Final = ", ".join(f'"{column}" = EXCLUDED."{column}"' for column in _METRIC_COLUMNS)
|
||||
return (
|
||||
f'INSERT INTO "{GLOBAL_SPEND_TABLE_NAME}" ("id", {_quoted(_KEY_COLUMNS)}, {_quoted(_METRIC_COLUMNS)}, '
|
||||
'"updated_at")\n'
|
||||
f"SELECT gen_random_uuid()::text, {normalized_keys}, {sums}, (NOW() AT TIME ZONE 'UTC')\n"
|
||||
'FROM "LiteLLM_DailyUserSpend" WHERE "date" = $1\n'
|
||||
f"GROUP BY {normalized_keys}\n"
|
||||
f"ON CONFLICT ({_quoted(_KEY_COLUMNS)}) DO UPDATE SET {overwrite}, "
|
||||
"\"updated_at\" = (NOW() AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
||||
|
||||
RECONCILE_DAY_SQL: Final = _reconcile_day_sql()
|
||||
_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now, (NOW() AT TIME ZONE 'UTC')::date::text AS today"
|
||||
_ALL_CLOSED_DAYS_SQL: Final = 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ORDER BY "date"'
|
||||
# Pod clocks drift from the database clock and from each other, so rows are picked up from a
|
||||
# little before the previous scan; rewriting a day twice is idempotent.
|
||||
_PENDING_DAYS_SQL: Final = (
|
||||
'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 '
|
||||
'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') '
|
||||
'ORDER BY "date"'
|
||||
)
|
||||
# Runs can overlap (Redis unreachable, lock expired on a long backfill), so the database keeps the
|
||||
# later of the stored and the incoming day and scan time in one statement; GREATEST skips NULL.
|
||||
_ADVANCE_MARKER_SQL: Final = (
|
||||
'INSERT INTO "LiteLLM_Config" ("param_name", "param_value") '
|
||||
"VALUES ($1, jsonb_build_object('reconciled_through', $2::text, 'scanned_at', $3::text)) "
|
||||
'ON CONFLICT ("param_name") DO UPDATE SET "param_value" = jsonb_build_object('
|
||||
"'reconciled_through', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'reconciled_through', "
|
||||
"EXCLUDED.\"param_value\" ->> 'reconciled_through'), "
|
||||
"'scanned_at', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'scanned_at', "
|
||||
"EXCLUDED.\"param_value\" ->> 'scanned_at'))"
|
||||
)
|
||||
|
||||
|
||||
class ReconciledThrough(BaseModel):
|
||||
"""``reconciled_through`` is the last closed UTC day the global table covers. ``scanned_at`` is
|
||||
the database clock when the scan behind the last fully successful run started: every per-key
|
||||
row written before it, on any day through the marker, is in the global table."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
reconciled_through: str
|
||||
scanned_at: str | None = None
|
||||
|
||||
|
||||
class _MarkerRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True)
|
||||
|
||||
param_value: object = None
|
||||
|
||||
|
||||
class _DateRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
date: str
|
||||
|
||||
|
||||
class _NowRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
now: str
|
||||
today: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ReconcileResult:
|
||||
days_reconciled: tuple[str, ...]
|
||||
reconciled_through: str | None
|
||||
failed_day: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PendingScan:
|
||||
marker: ReconciledThrough | None
|
||||
scanned_at: str
|
||||
days: tuple[str, ...]
|
||||
|
||||
|
||||
def _marker_from_param_value(value: object) -> ReconciledThrough | None:
|
||||
try:
|
||||
return (
|
||||
ReconciledThrough.model_validate_json(value)
|
||||
if isinstance(value, str)
|
||||
else ReconciledThrough.model_validate(value)
|
||||
)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
async def read_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None:
|
||||
from litellm.proxy.utils import get_config_param
|
||||
|
||||
row: Final = await get_config_param(prisma_client, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
return None if row is None else _marker_from_param_value(_MarkerRow.model_validate(row).param_value)
|
||||
|
||||
|
||||
async def reconciled_through(prisma_client: "PrismaClient") -> str | None:
|
||||
"""The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run."""
|
||||
marker: Final = await read_marker(prisma_client)
|
||||
return None if marker is None else marker.reconciled_through
|
||||
|
||||
|
||||
async def _advance_marker(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None:
|
||||
"""Move the stored marker to the last of ``days`` and to ``scanned_at`` where those are later
|
||||
than what is stored, so a slower overlapping run can only add to a faster run's marker."""
|
||||
from litellm.proxy.utils import invalidate_config_param
|
||||
|
||||
await prisma_client.db.execute_raw(
|
||||
_ADVANCE_MARKER_SQL,
|
||||
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM,
|
||||
max(days) if days else None,
|
||||
scanned_at,
|
||||
)
|
||||
await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
|
||||
|
||||
async def _db_now(prisma_client: "PrismaClient") -> _NowRow:
|
||||
rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL)
|
||||
return _NowRow.model_validate(rows[0])
|
||||
|
||||
|
||||
async def _scan_pending(prisma_client: "PrismaClient") -> _PendingScan:
|
||||
"""Every closed UTC day (strictly before the database's today) still to roll up, oldest first:
|
||||
days past the marker, plus any day with per-key rows written since the scan behind the marker.
|
||||
Before a run has fully succeeded there is no such scan, so every closed day is rolled up."""
|
||||
marker: Final = await read_marker(prisma_client)
|
||||
db_now: Final = await _db_now(prisma_client)
|
||||
last_closed_day: Final = (date.fromisoformat(db_now.today) - timedelta(days=1)).isoformat()
|
||||
rows: Final = (
|
||||
await prisma_client.db.query_raw(_ALL_CLOSED_DAYS_SQL, last_closed_day)
|
||||
if marker is None or marker.scanned_at is None
|
||||
else await prisma_client.db.query_raw(
|
||||
_PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at
|
||||
)
|
||||
)
|
||||
return _PendingScan(marker, db_now.now, tuple(_DateRow.model_validate(row).date for row in rows))
|
||||
|
||||
|
||||
async def pending_days(prisma_client: "PrismaClient") -> tuple[str, ...]:
|
||||
return (await _scan_pending(prisma_client)).days
|
||||
|
||||
|
||||
async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None:
|
||||
"""Rewrite one day of the global table from the per-key sums. Idempotent: a rerun
|
||||
overwrites every group with the same totals."""
|
||||
await prisma_client.db.execute_raw(RECONCILE_DAY_SQL, day)
|
||||
|
||||
|
||||
async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> ReconcileResult:
|
||||
"""Roll up every pending day, advancing the marker after each; a failing day stops the run
|
||||
with the marker on the last good day so the next run resumes there. The scan time is only
|
||||
recorded once every pending day is done, so late rows a failed run saw are found again."""
|
||||
scan: Final = await _scan_pending(prisma_client)
|
||||
done: Final = await _reconcile_until_failure(prisma_client, scan)
|
||||
if len(done) < len(scan.days):
|
||||
marker: Final = await reconciled_through(prisma_client)
|
||||
return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)])
|
||||
if scan.marker is not None or done:
|
||||
await _advance_marker(prisma_client, done, scanned_at=scan.scanned_at)
|
||||
return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client))
|
||||
|
||||
|
||||
async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]:
|
||||
for index, day in enumerate(scan.days):
|
||||
if not await _reconcile_and_record(prisma_client, scan.days[: index + 1]):
|
||||
return scan.days[:index]
|
||||
return scan.days
|
||||
|
||||
|
||||
async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: tuple[str, ...]) -> bool:
|
||||
day: Final = done_with_this[-1]
|
||||
try:
|
||||
await reconcile_day(prisma_client, day)
|
||||
await _advance_marker(prisma_client, done_with_this, scanned_at=None)
|
||||
except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done
|
||||
verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
async def run_scheduled_daily_global_spend_reconcile(
|
||||
prisma_client: "PrismaClient",
|
||||
pod_lock_manager: "PodLockManager | None" = None,
|
||||
alert: Callable[[str], Awaitable[None]] | None = None,
|
||||
) -> ReconcileResult | None:
|
||||
"""Run the reconcile under a cross-pod lock so one proxy does the work; the lock only saves
|
||||
effort (each day is an idempotent rewrite), so an unreachable Redis runs unguarded rather than skipping."""
|
||||
redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache
|
||||
if pod_lock_manager is None or redis_cache is None:
|
||||
return await _run_and_alert(prisma_client, alert=alert)
|
||||
|
||||
acquired: Final = await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, ttl=DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS
|
||||
)
|
||||
if not acquired and await _lock_is_held(pod_lock_manager, redis_cache):
|
||||
verbose_proxy_logger.info("Daily global spend reconcile: another pod holds the lock, skipping this run")
|
||||
return None
|
||||
try:
|
||||
return await _run_and_alert(prisma_client, alert=alert)
|
||||
finally:
|
||||
if acquired:
|
||||
await pod_lock_manager.release_lock(cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID)
|
||||
|
||||
|
||||
async def _lock_is_held(pod_lock_manager: "PodLockManager", redis_cache: "RedisCache") -> bool:
|
||||
try:
|
||||
lock_key: Final = pod_lock_manager.get_redis_lock_key(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID)
|
||||
return bool(await redis_cache.async_get_cache(lock_key))
|
||||
except Exception as exc: # noqa: BLE001 # an unreadable lock must not skip the run
|
||||
verbose_proxy_logger.warning("Daily global spend reconcile: could not read the lock: %s", exc)
|
||||
return False
|
||||
|
||||
|
||||
async def _run_and_alert(
|
||||
prisma_client: "PrismaClient",
|
||||
*,
|
||||
alert: Callable[[str], Awaitable[None]] | None,
|
||||
) -> ReconcileResult:
|
||||
result: Final = await run_daily_global_spend_reconcile(prisma_client)
|
||||
if result.days_reconciled:
|
||||
verbose_proxy_logger.info(
|
||||
"Daily global spend reconcile: rolled up %d day(s), reconciled through %s",
|
||||
len(result.days_reconciled),
|
||||
result.reconciled_through,
|
||||
)
|
||||
if result.failed_day is not None and alert is not None:
|
||||
await alert(
|
||||
f"Daily global spend reconcile stopped at {result.failed_day}; usage totals keep reading the per-key "
|
||||
f"table for ranges past {result.reconciled_through or 'the beginning'} until the next run succeeds."
|
||||
)
|
||||
return result
|
||||
|
|
@ -101,16 +101,6 @@ class DailySpendMetadata(BaseModel):
|
|||
page: int = Field(default=1)
|
||||
total_pages: int = Field(default=1)
|
||||
has_more: bool = Field(default=False)
|
||||
api_key_limit: int | None = Field(
|
||||
default=None,
|
||||
description="When set, api_keys and every api_key_breakdown list at most this many keys, "
|
||||
"ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.",
|
||||
)
|
||||
total_api_keys: int | None = Field(
|
||||
default=None,
|
||||
description="Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key "
|
||||
"lists are truncated to the highest-spend keys.",
|
||||
)
|
||||
|
||||
|
||||
class SpendAnalyticsPaginatedResponse(BaseModel):
|
||||
|
|
|
|||
|
|
@ -1,21 +1,13 @@
|
|||
import pathlib
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from psycopg.rows import dict_row
|
||||
from pytest_postgresql import factories
|
||||
|
||||
from litellm.constants import (
|
||||
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM,
|
||||
PTU_SENTINEL_API_KEY,
|
||||
USAGE_TOP_API_KEYS_LIMIT,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
|
||||
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
_adjust_dates_for_timezone,
|
||||
_build_aggregated_sql_query,
|
||||
|
|
@ -25,12 +17,8 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
|
|||
get_api_key_metadata,
|
||||
get_daily_activity,
|
||||
get_daily_activity_aggregated,
|
||||
global_rollup_reconciled_through,
|
||||
update_metrics,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.daily_global_spend_rollup import RECONCILE_DAY_SQL
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
from litellm.proxy.utils import evict_config_param
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
DailySpendMetadata,
|
||||
SpendMetrics,
|
||||
|
|
@ -181,7 +169,6 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
|
|||
"endpoint": "/v1/chat/completions",
|
||||
"api_key": None,
|
||||
"group_level": 62,
|
||||
"distinct_api_keys": None,
|
||||
"spend": 15.0,
|
||||
"prompt_tokens": 150,
|
||||
"completion_tokens": 75,
|
||||
|
|
@ -194,7 +181,31 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
|
|||
"endpoint": "/v1/embeddings",
|
||||
"api_key": None,
|
||||
"group_level": 62,
|
||||
"distinct_api_keys": None,
|
||||
"spend": 3.0,
|
||||
"prompt_tokens": 30,
|
||||
"completion_tokens": 0,
|
||||
"api_requests": 1,
|
||||
"successful_requests": 1,
|
||||
},
|
||||
# (date, endpoint, api_key) — populates the per-key sub-bucket
|
||||
{
|
||||
**base,
|
||||
"date": "2024-01-01",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"api_key": "key-1",
|
||||
"group_level": 30,
|
||||
"spend": 15.0,
|
||||
"prompt_tokens": 150,
|
||||
"completion_tokens": 75,
|
||||
"api_requests": 2,
|
||||
"successful_requests": 2,
|
||||
},
|
||||
{
|
||||
**base,
|
||||
"date": "2024-01-01",
|
||||
"endpoint": "/v1/embeddings",
|
||||
"api_key": "key-2",
|
||||
"group_level": 30,
|
||||
"spend": 3.0,
|
||||
"prompt_tokens": 30,
|
||||
"completion_tokens": 0,
|
||||
|
|
@ -208,7 +219,6 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
|
|||
"endpoint": None,
|
||||
"api_key": None,
|
||||
"group_level": 63,
|
||||
"distinct_api_keys": None,
|
||||
"spend": 18.0,
|
||||
"prompt_tokens": 180,
|
||||
"completion_tokens": 75,
|
||||
|
|
@ -222,40 +232,12 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
|
|||
"endpoint": None,
|
||||
"api_key": None,
|
||||
"group_level": 127,
|
||||
"distinct_api_keys": None,
|
||||
"spend": 18.0,
|
||||
"prompt_tokens": 180,
|
||||
"completion_tokens": 75,
|
||||
"api_requests": 3,
|
||||
"successful_requests": 3,
|
||||
},
|
||||
# (date, endpoint, api_key) — populates the per-key sub-bucket
|
||||
{
|
||||
**base,
|
||||
"date": "2024-01-01",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"api_key": "key-1",
|
||||
"group_level": 30,
|
||||
"distinct_api_keys": 2,
|
||||
"spend": 15.0,
|
||||
"prompt_tokens": 150,
|
||||
"completion_tokens": 75,
|
||||
"api_requests": 2,
|
||||
"successful_requests": 2,
|
||||
},
|
||||
{
|
||||
**base,
|
||||
"date": "2024-01-01",
|
||||
"endpoint": "/v1/embeddings",
|
||||
"api_key": "key-2",
|
||||
"group_level": 30,
|
||||
"distinct_api_keys": 2,
|
||||
"spend": 3.0,
|
||||
"prompt_tokens": 30,
|
||||
"completion_tokens": 0,
|
||||
"api_requests": 1,
|
||||
"successful_requests": 1,
|
||||
},
|
||||
]
|
||||
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=mock_rows)
|
||||
|
|
@ -492,7 +474,9 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash(
|
|||
return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")]
|
||||
)
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}]
|
||||
return_value=[
|
||||
{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}
|
||||
]
|
||||
)
|
||||
|
||||
result = await get_api_key_metadata(
|
||||
|
|
@ -869,7 +853,6 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys():
|
|||
"endpoint": "/v1/chat/completions",
|
||||
"api_key": None,
|
||||
"group_level": 62,
|
||||
"distinct_api_keys": None,
|
||||
"spend": 10.0,
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
|
|
@ -882,7 +865,6 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys():
|
|||
"endpoint": "/v1/chat/completions",
|
||||
"api_key": "deleted-key-hash",
|
||||
"group_level": 30,
|
||||
"distinct_api_keys": 1,
|
||||
"spend": 10.0,
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
|
|
@ -1327,11 +1309,42 @@ class TestBuildAggregatedSqlQuery:
|
|||
"user-1",
|
||||
"bedrock/global.anthropic.claude-opus-4-8",
|
||||
"sk-test",
|
||||
PTU_SENTINEL_API_KEY,
|
||||
]
|
||||
assert "model = $4" in sql
|
||||
assert "api_key = $5" in sql
|
||||
|
||||
def test_model_group_rollups_fall_back_to_model_name(self):
|
||||
"""Aggregated model_groups rollups must fall back to model for group-less rows.
|
||||
|
||||
The (date, model_group) grouping level cannot recover the model column
|
||||
after the fact (it is rolled up), so the fallback has to happen in SQL;
|
||||
without it, group-less rows silently vanish from the model_groups
|
||||
breakdown that the usage UI now renders by default. Group-less rows are
|
||||
stored as empty strings, not NULL (spend_tracking_utils defaults
|
||||
model_group to ""), so a plain COALESCE is not enough: the fallback must
|
||||
be NULLIF-wrapped to catch both
|
||||
"""
|
||||
sql, _ = _build_aggregated_sql_query(
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=None,
|
||||
start_date="2026-07-01",
|
||||
end_date="2026-07-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
normalized = " ".join(sql.split())
|
||||
fallback = "COALESCE(NULLIF(model_group, ''), model)"
|
||||
assert f"{fallback} AS model_group" in normalized
|
||||
assert (
|
||||
f"GROUPING(date, api_key, model, {fallback}, "
|
||||
"custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level" in normalized
|
||||
)
|
||||
assert f"(date, {fallback}), (date, {fallback}, api_key)," in normalized
|
||||
assert "(date, model_group)" not in normalized
|
||||
assert "COALESCE(model_group, model)" not in normalized
|
||||
|
||||
|
||||
class TestAggregatedEmptyEntityFilter:
|
||||
_BUILDERS: Final = (_build_aggregated_sql_query, _build_entity_rollup_sql_query)
|
||||
|
|
@ -1351,8 +1364,7 @@ class TestAggregatedEmptyEntityFilter:
|
|||
normalized = " ".join(sql.split())
|
||||
assert "IN ()" not in normalized
|
||||
assert '"team_id" IN' not in normalized
|
||||
sentinel_params = [PTU_SENTINEL_API_KEY] if build is _build_aggregated_sql_query else []
|
||||
assert params == ["2026-08-01", "2026-08-19", *sentinel_params]
|
||||
assert params == ["2026-08-01", "2026-08-19"]
|
||||
|
||||
@pytest.mark.parametrize("build", _BUILDERS)
|
||||
def test_empty_entity_list_matches_nothing_rather_than_everything(self, build):
|
||||
|
|
@ -1383,8 +1395,7 @@ class TestAggregatedEmptyEntityFilter:
|
|||
normalized = " ".join(sql.split())
|
||||
assert '"team_id" IN ($3, $4)' in normalized
|
||||
assert "FALSE" not in normalized
|
||||
sentinel_params = [PTU_SENTINEL_API_KEY] if build is _build_aggregated_sql_query else []
|
||||
assert params == ["2026-08-01", "2026-08-19", "team-alpha", "team-beta", *sentinel_params]
|
||||
assert params == ["2026-08-01", "2026-08-19", "team-alpha", "team-beta"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1409,7 +1420,6 @@ async def test_get_daily_activity_aggregated_empty_result_set():
|
|||
"mcp_namespaced_tool_name": None,
|
||||
"endpoint": None,
|
||||
"group_level": 127,
|
||||
"distinct_api_keys": None,
|
||||
"spend": None,
|
||||
"prompt_tokens": None,
|
||||
"completion_tokens": None,
|
||||
|
|
@ -1454,484 +1464,6 @@ async def test_get_daily_activity_aggregated_empty_result_set():
|
|||
assert result.metadata.total_compression_saved_tokens == 0
|
||||
|
||||
|
||||
_aggregated_postgresql_proc: Final = factories.postgresql_proc()
|
||||
_aggregated_postgresql: Final = factories.postgresql("_aggregated_postgresql_proc")
|
||||
|
||||
_DAILY_USER_SPEND_DDL: Final = """
|
||||
CREATE TABLE "LiteLLM_DailyUserSpend" (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT,
|
||||
date TEXT NOT NULL,
|
||||
api_key TEXT NOT NULL,
|
||||
model TEXT,
|
||||
model_group TEXT,
|
||||
custom_llm_provider TEXT,
|
||||
mcp_namespaced_tool_name TEXT,
|
||||
endpoint TEXT,
|
||||
prompt_tokens BIGINT DEFAULT 0,
|
||||
completion_tokens BIGINT DEFAULT 0,
|
||||
cache_read_input_tokens BIGINT DEFAULT 0,
|
||||
cache_creation_input_tokens BIGINT DEFAULT 0,
|
||||
compression_saved_tokens BIGINT DEFAULT 0,
|
||||
compression_savings_spend DOUBLE PRECISION DEFAULT 0,
|
||||
prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0,
|
||||
gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0,
|
||||
autorouter_savings_spend DOUBLE PRECISION DEFAULT 0,
|
||||
spend DOUBLE PRECISION DEFAULT 0,
|
||||
api_requests BIGINT DEFAULT 0,
|
||||
successful_requests BIGINT DEFAULT 0,
|
||||
failed_requests BIGINT DEFAULT 0,
|
||||
total_response_time_ms BIGINT DEFAULT 0,
|
||||
timed_requests BIGINT DEFAULT 0
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
def _seed_daily_user_spend(conn: psycopg.Connection, rows: Sequence[tuple[object, ...]]) -> None:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(_DAILY_USER_SPEND_DDL)
|
||||
cur.executemany(
|
||||
"""
|
||||
INSERT INTO "LiteLLM_DailyUserSpend"
|
||||
(id, user_id, date, api_key, model, model_group, custom_llm_provider,
|
||||
endpoint, prompt_tokens, spend, api_requests, successful_requests)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
rows,
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _psycopg_query_raw(conn: psycopg.Connection, row_counts: list[int]):
|
||||
"""Run the proxy's $N-parameterized SQL through psycopg, recording each result size."""
|
||||
|
||||
async def query_raw(sql: str, *params: str) -> list[dict[str, object]]:
|
||||
converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql)
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
cur.execute(
|
||||
converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query
|
||||
{f"p{i}": v for i, v in enumerate(params, start=1)},
|
||||
)
|
||||
rows: Final = cur.fetchall()
|
||||
row_counts.append(len(rows))
|
||||
return rows
|
||||
|
||||
return query_raw
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_daily_activity_aggregated_bounds_api_key_rollups(
|
||||
_aggregated_postgresql: psycopg.Connection,
|
||||
):
|
||||
"""Run the GROUPING SETS statement against real Postgres with more keys than the cap.
|
||||
|
||||
key-004 and key-005 tie on spend exactly at the USAGE_TOP_API_KEYS_LIMIT
|
||||
cutoff; the api_key tiebreaker must keep key-004 and drop key-005. The PTU
|
||||
sentinel outspends every key but must not take a slot. Excluded keys and the
|
||||
sentinel still count toward the totals and the model rollup, which come from
|
||||
the key-free arm.
|
||||
"""
|
||||
n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 5
|
||||
key_rows: Final = [
|
||||
(
|
||||
f"row-{i:03d}",
|
||||
f"user-{i:03d}",
|
||||
"2026-06-01",
|
||||
f"key-{i:03d}",
|
||||
"gpt-5",
|
||||
"",
|
||||
"openai",
|
||||
"/v1/chat/completions",
|
||||
10,
|
||||
6.0 if i == 4 else float(i + 1),
|
||||
1,
|
||||
1,
|
||||
)
|
||||
for i in range(n_keys)
|
||||
]
|
||||
sentinel_row: Final = (
|
||||
"row-ptu",
|
||||
None,
|
||||
"2026-06-01",
|
||||
PTU_SENTINEL_API_KEY,
|
||||
"gpt-5",
|
||||
"",
|
||||
"azure",
|
||||
None,
|
||||
0,
|
||||
1000.0,
|
||||
0,
|
||||
0,
|
||||
)
|
||||
_seed_daily_user_spend(_aggregated_postgresql, [*key_rows, sentinel_row])
|
||||
key_spend: Final = sum(6.0 if i == 4 else float(i + 1) for i in range(n_keys))
|
||||
|
||||
row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts)
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await get_daily_activity_aggregated(
|
||||
prisma_client=mock_prisma,
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=None,
|
||||
entity_metadata_field=None,
|
||||
start_date="2026-06-01",
|
||||
end_date="2026-06-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
# Key-free arm: (), (date), (date, model), (date, model_group), two providers,
|
||||
# one mcp NULL bucket, endpoint plus its NULL bucket = 9 rows regardless of key count.
|
||||
# Per-key arm: six per-key grouping sets, each capped at the limit.
|
||||
assert row_counts == [9 + 6 * USAGE_TOP_API_KEYS_LIMIT]
|
||||
|
||||
assert result.metadata.total_spend == pytest.approx(key_spend + 1000.0)
|
||||
assert result.metadata.total_api_requests == n_keys
|
||||
assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT
|
||||
assert result.metadata.total_api_keys == n_keys
|
||||
|
||||
expected_top: Final = {f"key-{i:03d}" for i in range(6, n_keys)} | {"key-004"}
|
||||
day: Final = result.results[0]
|
||||
assert day.metrics.spend == pytest.approx(key_spend + 1000.0)
|
||||
assert set(day.breakdown.api_keys) == expected_top
|
||||
assert day.breakdown.api_keys["key-004"].metrics.spend == 6.0
|
||||
assert "key-005" not in day.breakdown.api_keys
|
||||
assert PTU_SENTINEL_API_KEY not in day.breakdown.api_keys
|
||||
|
||||
assert day.breakdown.models["gpt-5"].metrics.spend == pytest.approx(key_spend + 1000.0)
|
||||
assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == expected_top
|
||||
assert day.breakdown.providers["openai"].metrics.spend == pytest.approx(key_spend)
|
||||
assert set(day.breakdown.providers["openai"].api_key_breakdown) == expected_top
|
||||
assert day.breakdown.endpoints["/v1/chat/completions"].metrics.api_requests == n_keys
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both_arms(
|
||||
_aggregated_postgresql: psycopg.Connection,
|
||||
):
|
||||
"""An explicit api_key filter must scope the key-free totals and the per-key
|
||||
rollups to that key alone, so the two arms never disagree."""
|
||||
rows: Final = [
|
||||
(
|
||||
f"row-{i}",
|
||||
f"user-{i}",
|
||||
"2026-06-01",
|
||||
f"key-{i}",
|
||||
"gpt-5",
|
||||
"",
|
||||
"openai",
|
||||
"/v1/chat/completions",
|
||||
10,
|
||||
float(i + 1),
|
||||
1,
|
||||
1,
|
||||
)
|
||||
for i in range(3)
|
||||
]
|
||||
_seed_daily_user_spend(_aggregated_postgresql, rows)
|
||||
|
||||
row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts)
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await get_daily_activity_aggregated(
|
||||
prisma_client=mock_prisma,
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=None,
|
||||
entity_metadata_field=None,
|
||||
start_date="2026-06-01",
|
||||
end_date="2026-06-01",
|
||||
model=None,
|
||||
api_key="key-1",
|
||||
)
|
||||
|
||||
assert result.metadata.total_spend == 2.0
|
||||
assert result.metadata.total_api_keys == 1
|
||||
day: Final = result.results[0]
|
||||
assert set(day.breakdown.api_keys) == {"key-1"}
|
||||
assert day.breakdown.api_keys["key-1"].metrics.spend == 2.0
|
||||
assert day.breakdown.models["gpt-5"].metrics.spend == 2.0
|
||||
assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == {"key-1"}
|
||||
|
||||
|
||||
def _prisma_with_marker(marker: str | None) -> MagicMock:
|
||||
prisma = MagicMock()
|
||||
prisma.db = MagicMock()
|
||||
prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
row = (
|
||||
None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}')
|
||||
)
|
||||
prisma.get_generic_data = AsyncMock(return_value=row)
|
||||
return prisma
|
||||
|
||||
|
||||
def _unfiltered_user_query(**overrides):
|
||||
return {
|
||||
"table_name": "litellm_dailyuserspend",
|
||||
"entity_id_field": "user_id",
|
||||
"entity_id": None,
|
||||
"start_date": "2026-06-01",
|
||||
"end_date": "2026-06-02",
|
||||
"model": None,
|
||||
"api_key": None,
|
||||
"exclude_entity_ids": None,
|
||||
"timezone_offset_minutes": None,
|
||||
"include_current_utc_day": False,
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("marker", "overrides", "expected"),
|
||||
[
|
||||
("2026-06-02", {}, "2026-06-02"),
|
||||
("2026-06-02", {"model": "gpt-5"}, "2026-06-02"),
|
||||
("2026-05-01", {}, "2026-05-01"),
|
||||
(None, {}, None),
|
||||
("2026-06-02", {"api_key": "sk-1"}, None),
|
||||
("2026-06-02", {"api_key": []}, None),
|
||||
("2026-06-02", {"entity_id": "u-1"}, None),
|
||||
("2026-06-02", {"exclude_entity_ids": ["u-1"]}, None),
|
||||
("2026-06-02", {"table_name": "litellm_dailyteamspend", "entity_id_field": "team_id"}, None),
|
||||
],
|
||||
)
|
||||
async def test_global_rollup_marker_is_used_only_for_unfiltered_user_reads(marker, overrides, expected):
|
||||
"""Anything that filters by key or entity has no counterpart in the global table; the
|
||||
SQL splits the range at the marker itself, so the marker passes through unchanged."""
|
||||
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
prisma = _prisma_with_marker(marker)
|
||||
|
||||
assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query(**overrides)) == expected
|
||||
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_rollup_marker_read_failure_falls_back_to_the_per_key_table():
|
||||
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
prisma = _prisma_with_marker(None)
|
||||
prisma.get_generic_data = AsyncMock(side_effect=RuntimeError("db down"))
|
||||
|
||||
assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query()) is None
|
||||
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
|
||||
|
||||
_GLOBAL_SPEND_MIGRATION: Final = (
|
||||
pathlib.Path(__file__).resolve().parents[4]
|
||||
/ "litellm-proxy-extras"
|
||||
/ "litellm_proxy_extras"
|
||||
/ "migrations"
|
||||
/ "20260915000000_add_daily_global_spend"
|
||||
/ "migration.sql"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_table_and_open_days_live(
|
||||
_aggregated_postgresql: psycopg.Connection,
|
||||
):
|
||||
"""Day 1 is rolled up and day 2 is still open (never rolled up), so a marker of day 1 must
|
||||
give the same response as reading everything per-key: day 1 from the global table, day 2
|
||||
live, one grand total across both. Per-key rows that land after the rollup then tell the
|
||||
two sources apart: a late day 1 row is invisible to totals until the next reconcile while a
|
||||
late day 2 row shows up at once, and both keys rank in the key breakdown, which stays
|
||||
per-key throughout."""
|
||||
n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3
|
||||
rows: Final = [
|
||||
(
|
||||
f"row-{day}-{i:03d}",
|
||||
f"user-{i % 7}",
|
||||
day,
|
||||
f"key-{i:03d}",
|
||||
"gpt-5" if i % 2 else "claude",
|
||||
"" if i % 3 else "gpt-5",
|
||||
"openai" if i % 2 else None,
|
||||
"/v1/chat/completions" if i % 5 else None,
|
||||
10,
|
||||
float(i + 1),
|
||||
1,
|
||||
1,
|
||||
)
|
||||
for day in ("2026-06-01", "2026-06-02")
|
||||
for i in range(n_keys)
|
||||
]
|
||||
_seed_daily_user_spend(_aggregated_postgresql, rows)
|
||||
with _aggregated_postgresql.cursor() as cur:
|
||||
cur.execute(
|
||||
'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms = prompt_tokens * 25, '
|
||||
"timed_requests = api_requests"
|
||||
)
|
||||
cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal
|
||||
cur.execute(
|
||||
re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg
|
||||
{"p1": "2026-06-01"},
|
||||
)
|
||||
_aggregated_postgresql.commit()
|
||||
|
||||
async def read(marker: str | None):
|
||||
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
prisma = _prisma_with_marker(marker)
|
||||
prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, [])
|
||||
return await get_daily_activity_aggregated(
|
||||
prisma_client=prisma,
|
||||
entity_metadata_field=None,
|
||||
**_unfiltered_user_query(),
|
||||
)
|
||||
|
||||
from_per_key = await read(None)
|
||||
from_global = await read("2026-06-01")
|
||||
|
||||
assert from_global.model_dump() == from_per_key.model_dump()
|
||||
seeded_spend: Final = 2 * sum(float(i + 1) for i in range(n_keys))
|
||||
assert from_global.metadata.total_spend == pytest.approx(seeded_spend)
|
||||
assert from_global.metadata.total_response_time_ms == 2 * n_keys * 10 * 25
|
||||
assert from_global.metadata.total_timed_requests == 2 * n_keys
|
||||
assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"}
|
||||
assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT
|
||||
assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"}
|
||||
|
||||
with _aggregated_postgresql.cursor() as cur:
|
||||
cur.executemany(
|
||||
"""
|
||||
INSERT INTO "LiteLLM_DailyUserSpend"
|
||||
(id, user_id, date, api_key, model, model_group, custom_llm_provider,
|
||||
endpoint, prompt_tokens, spend, api_requests, successful_requests)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
[
|
||||
("late-1", "user-late", "2026-06-01", "key-late-1", "gpt-5", "", "openai", None, 10, 1000.0, 1, 1),
|
||||
("late-2", "user-late", "2026-06-02", "key-late-2", "gpt-5", "", "openai", None, 10, 500.0, 1, 1),
|
||||
],
|
||||
)
|
||||
_aggregated_postgresql.commit()
|
||||
|
||||
late_per_key = await read(None)
|
||||
late_global = await read("2026-06-01")
|
||||
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
|
||||
assert late_per_key.metadata.total_spend == pytest.approx(seeded_spend + 1000.0 + 500.0)
|
||||
assert late_global.metadata.total_spend == pytest.approx(seeded_spend + 500.0)
|
||||
by_day: Final = {day.date.isoformat(): day for day in late_global.results}
|
||||
assert by_day["2026-06-01"].metrics.spend == pytest.approx(seeded_spend / 2)
|
||||
assert by_day["2026-06-02"].metrics.spend == pytest.approx(seeded_spend / 2 + 500.0)
|
||||
assert by_day["2026-06-01"].breakdown.api_keys["key-late-1"].metrics.spend == pytest.approx(1000.0)
|
||||
assert by_day["2026-06-02"].breakdown.api_keys["key-late-2"].metrics.spend == pytest.approx(500.0)
|
||||
assert late_global.metadata.total_api_keys == n_keys + 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_daily_activity_aggregated_reports_exact_limit_key_count_as_complete(
|
||||
_aggregated_postgresql: psycopg.Connection,
|
||||
):
|
||||
"""With exactly USAGE_TOP_API_KEYS_LIMIT keys nothing is dropped, and the
|
||||
response must say so: total_api_keys equals the limit rather than exceeding it."""
|
||||
rows: Final = [
|
||||
(
|
||||
f"row-{i:03d}",
|
||||
f"user-{i:03d}",
|
||||
"2026-06-01",
|
||||
f"key-{i:03d}",
|
||||
"gpt-5",
|
||||
"",
|
||||
"openai",
|
||||
"/v1/chat/completions",
|
||||
10,
|
||||
float(i + 1),
|
||||
1,
|
||||
1,
|
||||
)
|
||||
for i in range(USAGE_TOP_API_KEYS_LIMIT)
|
||||
]
|
||||
_seed_daily_user_spend(_aggregated_postgresql, rows)
|
||||
|
||||
row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts)
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await get_daily_activity_aggregated(
|
||||
prisma_client=mock_prisma,
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=None,
|
||||
entity_metadata_field=None,
|
||||
start_date="2026-06-01",
|
||||
end_date="2026-06-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
assert result.metadata.total_api_keys == USAGE_TOP_API_KEYS_LIMIT
|
||||
assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT
|
||||
assert len(result.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_model_name(
|
||||
_aggregated_postgresql: psycopg.Connection,
|
||||
):
|
||||
"""Rows stored with an empty or NULL model_group must land in the model_groups
|
||||
breakdown under their model name instead of vanishing from the usage UI."""
|
||||
rows: Final = [
|
||||
(
|
||||
"row-0",
|
||||
"user-0",
|
||||
"2026-06-01",
|
||||
"key-0",
|
||||
"gpt-5",
|
||||
"gpt-5-eu",
|
||||
"openai",
|
||||
"/v1/chat/completions",
|
||||
10,
|
||||
7.0,
|
||||
1,
|
||||
1,
|
||||
),
|
||||
("row-1", "user-1", "2026-06-01", "key-1", "gpt-5", "", "openai", "/v1/chat/completions", 10, 3.0, 1, 1),
|
||||
("row-2", "user-2", "2026-06-01", "key-2", "claude-x", None, "anthropic", "/v1/messages", 10, 2.0, 1, 1),
|
||||
]
|
||||
_seed_daily_user_spend(_aggregated_postgresql, rows)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, [])
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
|
||||
result = await get_daily_activity_aggregated(
|
||||
prisma_client=mock_prisma,
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=None,
|
||||
entity_metadata_field=None,
|
||||
start_date="2026-06-01",
|
||||
end_date="2026-06-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
breakdown: Final = result.results[0].breakdown
|
||||
assert set(breakdown.model_groups) == {"gpt-5-eu", "gpt-5", "claude-x"}
|
||||
assert breakdown.model_groups["gpt-5-eu"].metrics.spend == 7.0
|
||||
assert breakdown.model_groups["gpt-5"].metrics.spend == 3.0
|
||||
assert breakdown.model_groups["claude-x"].metrics.spend == 2.0
|
||||
assert set(breakdown.model_groups["gpt-5"].api_key_breakdown) == {"key-1"}
|
||||
assert set(breakdown.models) == {"gpt-5", "claude-x"}
|
||||
assert breakdown.models["gpt-5"].metrics.spend == 10.0
|
||||
|
||||
|
||||
def _no_spend_record():
|
||||
"""A rollup row for a key with no spend, where SUM() returns NULL (None)."""
|
||||
return SimpleNamespace(
|
||||
|
|
@ -2717,7 +2249,7 @@ def test_entity_rollup_sql_query_and_api_key_list_filter():
|
|||
api_key=[],
|
||||
)
|
||||
assert "FALSE" in empty_sql
|
||||
assert empty_params == ["2024-01-01", "2024-01-31", PTU_SENTINEL_API_KEY]
|
||||
assert empty_params == ["2024-01-01", "2024-01-31"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -2751,10 +2283,10 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown():
|
|||
"successful_requests": 0,
|
||||
}
|
||||
main_rows = [
|
||||
{**base, "date": None, "group_level": 127, "distinct_api_keys": None, "spend": 18.0},
|
||||
{**base, "date": "2024-01-01", "group_level": 63, "distinct_api_keys": None, "spend": 18.0},
|
||||
{**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 47, "distinct_api_keys": None, "spend": 18.0},
|
||||
{**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 31, "distinct_api_keys": 1, "spend": 12.0},
|
||||
{**base, "date": None, "group_level": 127, "spend": 18.0},
|
||||
{**base, "date": "2024-01-01", "group_level": 63, "spend": 18.0},
|
||||
{**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 47, "spend": 18.0},
|
||||
{**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 31, "spend": 12.0},
|
||||
]
|
||||
entity_base = {
|
||||
key: value
|
||||
|
|
|
|||
|
|
@ -28,7 +28,6 @@ from typing import List, Optional, Union
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from fastapi import FastAPI
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypedDict
|
||||
|
|
@ -1043,57 +1042,6 @@ async def test_spend_report_locks_are_never_released():
|
|||
proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited()
|
||||
|
||||
|
||||
def _init_daily_global_spend_reconcile_job() -> tuple[AsyncIOScheduler, MagicMock, MagicMock]:
|
||||
scheduler = AsyncIOScheduler()
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.alerting_handler = AsyncMock()
|
||||
prisma_client = MagicMock()
|
||||
ProxyStartupEvent._initialize_daily_global_spend_reconcile_job(
|
||||
scheduler=scheduler,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
return scheduler, proxy_logging_obj, prisma_client
|
||||
|
||||
|
||||
def test_daily_global_spend_reconcile_job_is_scheduled_nightly_with_an_immediate_catch_up_run():
|
||||
"""Startup schedules the LiteLLM_DailyGlobalSpend backfill a couple of minutes out, so a
|
||||
fresh deploy switches usage reads to the global table without waiting for the nightly
|
||||
run, and after that it fires once a day at 00:30 UTC, when the previous UTC day is closed."""
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID
|
||||
|
||||
scheduler, _, _ = _init_daily_global_spend_reconcile_job()
|
||||
job = scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID)
|
||||
assert job is not None
|
||||
|
||||
assert timedelta(0) < job.next_run_time - datetime.now(timezone.utc) <= timedelta(minutes=2)
|
||||
after_catch_up = datetime(2026, 9, 16, 12, 0, tzinfo=timezone.utc)
|
||||
assert job.trigger.get_next_fire_time(None, after_catch_up) == datetime(2026, 9, 17, 0, 30, tzinfo=timezone.utc)
|
||||
just_after_a_run = datetime(2026, 9, 17, 0, 30, 1, tzinfo=timezone.utc)
|
||||
assert job.trigger.get_next_fire_time(None, just_after_a_run) == datetime(2026, 9, 18, 0, 30, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_global_spend_reconcile_job_runs_under_the_pod_lock_and_alerts_through_the_proxy(monkeypatch):
|
||||
from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID
|
||||
|
||||
scheduler, proxy_logging_obj, prisma_client = _init_daily_global_spend_reconcile_job()
|
||||
run = AsyncMock()
|
||||
monkeypatch.setattr(ps, "run_scheduled_daily_global_spend_reconcile", run)
|
||||
|
||||
await scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID).func()
|
||||
|
||||
run.assert_awaited_once()
|
||||
assert run.await_args.args == (prisma_client,)
|
||||
assert run.await_args.kwargs["pod_lock_manager"] is proxy_logging_obj.db_spend_update_writer.pod_lock_manager
|
||||
await run.await_args.kwargs["alert"]("day 2026-09-01 failed")
|
||||
proxy_logging_obj.alerting_handler.assert_awaited_once()
|
||||
assert proxy_logging_obj.alerting_handler.await_args.kwargs["message"] == "day 2026-09-01 failed"
|
||||
assert proxy_logging_obj.alerting_handler.await_args.kwargs["level"] == "High"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch):
|
||||
"""The boot-time send goes through the same gate, so a losing pod sends nothing at all:
|
||||
|
|
|
|||
|
|
@ -1,532 +0,0 @@
|
|||
"""Tests for the LiteLLM_DailyGlobalSpend reconcile job (LIT-7818)."""
|
||||
|
||||
import json
|
||||
import pathlib
|
||||
import re
|
||||
from datetime import date
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from psycopg.rows import dict_row
|
||||
from pytest_postgresql import factories
|
||||
|
||||
from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM
|
||||
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key
|
||||
from litellm.proxy.spend_tracking.daily_global_spend_rollup import (
|
||||
_ADVANCE_MARKER_SQL,
|
||||
RECONCILE_DAY_SQL,
|
||||
read_marker,
|
||||
reconciled_through,
|
||||
run_daily_global_spend_reconcile,
|
||||
run_scheduled_daily_global_spend_reconcile,
|
||||
)
|
||||
from litellm.proxy.utils import evict_config_param
|
||||
|
||||
USER_TABLE: Final = DAILY_SPEND_TABLES["user"]
|
||||
TODAY: Final = date(2026, 9, 15)
|
||||
|
||||
|
||||
class _FakeConfigRow:
|
||||
def __init__(self, param_name: str, param_value: object) -> None:
|
||||
self.param_name = param_name
|
||||
self.param_value = param_value
|
||||
|
||||
|
||||
class _FakeConfigTable:
|
||||
def __init__(self) -> None:
|
||||
self.rows: dict[str, object] = {}
|
||||
|
||||
def advance(self, param_name: str, through: str | None, scanned_at: str | None) -> None:
|
||||
"""What ``_ADVANCE_MARKER_SQL`` does in Postgres: keep the later of stored and incoming per field."""
|
||||
stored = self.rows.get(param_name)
|
||||
current: dict[str, str | None] = json.loads(stored) if isinstance(stored, str) else {}
|
||||
self.rows[param_name] = json.dumps(
|
||||
{
|
||||
"reconciled_through": _greatest(current.get("reconciled_through"), through),
|
||||
"scanned_at": _greatest(current.get("scanned_at"), scanned_at),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _greatest(stored: str | None, incoming: str | None) -> str | None:
|
||||
present = [value for value in (stored, incoming) if value is not None]
|
||||
return max(present) if present else None
|
||||
|
||||
|
||||
class _FakeDb:
|
||||
"""Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query,
|
||||
so "rows written since the last scan" behaves like Postgres would. The database's own
|
||||
date decides which day is still open, never the pod's clock."""
|
||||
|
||||
def __init__(self, prisma: "_FakePrisma") -> None:
|
||||
self._prisma = prisma
|
||||
self.litellm_config = _FakeConfigTable()
|
||||
|
||||
async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]:
|
||||
if sql.startswith("SELECT (NOW()"):
|
||||
self._prisma.clock += 1
|
||||
return [{"now": f"clock-{self._prisma.clock:04d}", "today": self._prisma.today.isoformat()}]
|
||||
rows = self._prisma.user_rows
|
||||
if len(params) == 1:
|
||||
(last,) = params
|
||||
return [{"date": d} for d in sorted(rows) if d <= last]
|
||||
last, marker, scanned_at = params
|
||||
return [
|
||||
{"date": d} for d, written in sorted(rows.items()) if d <= last and (d > marker or written >= scanned_at)
|
||||
]
|
||||
|
||||
async def execute_raw(self, sql: str, *params: str | None) -> int:
|
||||
if sql == _ADVANCE_MARKER_SQL:
|
||||
param_name, through, scanned_at = params
|
||||
assert param_name is not None
|
||||
self.litellm_config.advance(param_name, through, scanned_at)
|
||||
return 1
|
||||
(day,) = params
|
||||
if day is None or day in self._prisma.failing_days:
|
||||
raise RuntimeError(f"day {day} exploded")
|
||||
self._prisma.reconciled.append(day)
|
||||
landing = self._prisma.marker_landing_on_day.get(day)
|
||||
if landing is not None:
|
||||
self.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = landing
|
||||
return 1
|
||||
|
||||
|
||||
class _FakePrisma:
|
||||
"""Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.
|
||||
``marker_landing_on_day`` stores another pod's marker the moment this run rewrites that day."""
|
||||
|
||||
def __init__(
|
||||
self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset(), today: date = TODAY
|
||||
) -> None:
|
||||
self.clock = 0
|
||||
self.today = today
|
||||
self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days}
|
||||
self.failing_days = failing_days
|
||||
self.marker_landing_on_day: dict[str, str] = {}
|
||||
self.reconciled: list[str] = []
|
||||
self.db = _FakeDb(self)
|
||||
|
||||
def write_late_row(self, day: str) -> None:
|
||||
"""A per-key row for ``day`` lands now, after whatever scans already happened."""
|
||||
self.clock += 1
|
||||
self.user_rows[day] = f"clock-{self.clock:04d}"
|
||||
|
||||
async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None:
|
||||
stored = self.db.litellm_config.rows.get(value)
|
||||
return None if stored is None else _FakeConfigRow(value, stored)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def _fresh_marker_cache():
|
||||
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
yield
|
||||
await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_run_rolls_up_every_closed_day_and_never_the_database_s_today():
|
||||
"""Before any marker exists every closed day with per-key rows is rolled up. Today is left
|
||||
out: pods are still flushing it, so it is served live from the per-key table until it closes.
|
||||
The database clock says which day that is; a pod booting with its clock a day ahead must not
|
||||
roll the open day up and mark it reconciled."""
|
||||
prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15"))
|
||||
|
||||
result = await run_daily_global_spend_reconcile(prisma)
|
||||
|
||||
assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14")
|
||||
assert result.failed_day is None
|
||||
assert result.reconciled_through == "2026-09-14"
|
||||
assert await reconciled_through(prisma) == "2026-09-14"
|
||||
assert "2026-09-15" not in prisma.reconciled
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed():
|
||||
prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14"), today=date(2026, 9, 14))
|
||||
await run_daily_global_spend_reconcile(prisma)
|
||||
prisma.reconciled.clear()
|
||||
prisma.today = TODAY
|
||||
|
||||
result = await run_daily_global_spend_reconcile(prisma)
|
||||
|
||||
assert result.days_reconciled == ("2026-09-14",)
|
||||
assert await reconciled_through(prisma) == "2026-09-14"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_run():
|
||||
"""Per-key rows carry the request start date, so a delayed flush or retry can add spend to a
|
||||
day far behind the marker. That day is rewritten, and the marker never moves back for it."""
|
||||
prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13"), today=date(2026, 9, 14))
|
||||
await run_daily_global_spend_reconcile(prisma)
|
||||
prisma.reconciled.clear()
|
||||
prisma.today = TODAY
|
||||
prisma.write_late_row("2026-09-01")
|
||||
prisma.write_late_row("2026-09-03")
|
||||
|
||||
result = await run_daily_global_spend_reconcile(prisma)
|
||||
|
||||
assert result.days_reconciled == ("2026-09-01", "2026-09-03")
|
||||
assert "2026-09-05" not in prisma.reconciled
|
||||
assert await reconciled_through(prisma) == "2026-09-13"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_late_row_seen_by_a_failed_run_is_seen_again_by_the_next_one():
|
||||
"""The scan time only advances when every pending day was rewritten, otherwise a late row
|
||||
found by the failed run would be counted as handled."""
|
||||
prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13"), today=date(2026, 9, 14))
|
||||
await run_daily_global_spend_reconcile(prisma)
|
||||
prisma.today = TODAY
|
||||
prisma.write_late_row("2026-09-01")
|
||||
prisma.failing_days = frozenset({"2026-09-01"})
|
||||
failed = await run_daily_global_spend_reconcile(prisma)
|
||||
prisma.failing_days = frozenset()
|
||||
prisma.reconciled.clear()
|
||||
|
||||
result = await run_daily_global_spend_reconcile(prisma)
|
||||
|
||||
assert failed.failed_day == "2026-09-01"
|
||||
assert failed.reconciled_through == "2026-09-13"
|
||||
assert result.days_reconciled == ("2026-09-01",)
|
||||
assert result.failed_day is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again():
|
||||
prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13"))
|
||||
prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-13"}'
|
||||
|
||||
result = await run_daily_global_spend_reconcile(prisma)
|
||||
|
||||
assert result.days_reconciled == ("2026-09-01", "2026-09-13")
|
||||
marker = await read_marker(prisma)
|
||||
assert marker is not None and marker.reconciled_through == "2026-09-13" and marker.scanned_at is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_run_with_no_new_closed_days_keeps_the_marker():
|
||||
prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14))
|
||||
await run_daily_global_spend_reconcile(prisma)
|
||||
prisma.reconciled.clear()
|
||||
|
||||
result = await run_daily_global_spend_reconcile(prisma)
|
||||
|
||||
assert result.days_reconciled == ()
|
||||
assert result.reconciled_through == "2026-09-13"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failing_day_stops_the_run_and_leaves_the_marker_on_the_last_good_day():
|
||||
"""The marker may never claim a day that was not rewritten: reads past it would then trust
|
||||
a global table missing that day's spend."""
|
||||
prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"}))
|
||||
|
||||
result = await run_daily_global_spend_reconcile(prisma)
|
||||
|
||||
assert result.days_reconciled == ("2026-09-01",)
|
||||
assert result.failed_day == "2026-09-02"
|
||||
assert result.reconciled_through == "2026-09-01"
|
||||
assert prisma.reconciled == ["2026-09-01"]
|
||||
assert await reconciled_through(prisma) == "2026-09-01"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_next_run_resumes_from_the_failed_day():
|
||||
prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"}))
|
||||
await run_daily_global_spend_reconcile(prisma)
|
||||
prisma.failing_days = frozenset()
|
||||
|
||||
result = await run_daily_global_spend_reconcile(prisma)
|
||||
|
||||
assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03")
|
||||
assert await reconciled_through(prisma) == "2026-09-03"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_slower_overlapping_run_never_rewinds_the_marker_a_faster_run_stored():
|
||||
"""Two pods can reconcile at once (Redis unreachable, or the lock expired on a long backfill).
|
||||
When the faster one has already stored a later marker, the slower one may only add to it. Putting
|
||||
its own older prefix back, or dropping the scan time, would send usage reads for every day in
|
||||
between back to the per-key table until the next run."""
|
||||
prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-03"}))
|
||||
prisma.marker_landing_on_day = {
|
||||
"2026-09-02": '{"reconciled_through": "2026-09-14", "scanned_at": "clock-0009"}',
|
||||
}
|
||||
|
||||
result = await run_daily_global_spend_reconcile(prisma)
|
||||
|
||||
assert result.days_reconciled == ("2026-09-01", "2026-09-02")
|
||||
assert result.reconciled_through == "2026-09-14"
|
||||
marker = await read_marker(prisma)
|
||||
assert marker is not None and (marker.reconciled_through, marker.scanned_at) == ("2026-09-14", "clock-0009")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts():
|
||||
"""When the rewrite of a late day fails the marker must stay put and the operator must hear about it."""
|
||||
prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14))
|
||||
await run_daily_global_spend_reconcile(prisma)
|
||||
prisma.today = TODAY
|
||||
prisma.write_late_row("2026-09-12")
|
||||
prisma.failing_days = frozenset({"2026-09-12"})
|
||||
alert = AsyncMock()
|
||||
|
||||
result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert)
|
||||
|
||||
assert result is not None
|
||||
assert result.days_reconciled == ()
|
||||
assert result.failed_day == "2026-09-12"
|
||||
assert result.reconciled_through == "2026-09-13"
|
||||
alert.assert_awaited_once()
|
||||
assert "2026-09-12" in alert.await_args.args[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_clean_run_does_not_alert():
|
||||
prisma = _FakePrisma(user_days=("2026-09-13",))
|
||||
alert = AsyncMock()
|
||||
|
||||
await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert)
|
||||
|
||||
alert.assert_not_awaited()
|
||||
|
||||
|
||||
def _pod_lock(acquired: bool) -> MagicMock:
|
||||
lock = MagicMock()
|
||||
lock.redis_cache = MagicMock()
|
||||
lock.redis_cache.async_get_cache = AsyncMock(return_value="other-pod")
|
||||
lock.get_redis_lock_key = MagicMock(return_value="lock-key")
|
||||
lock.acquire_lock = AsyncMock(return_value=acquired)
|
||||
lock.release_lock = AsyncMock()
|
||||
return lock
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_run_skips_when_another_pod_holds_the_lock():
|
||||
prisma = _FakePrisma(user_days=("2026-09-13",))
|
||||
lock = _pod_lock(acquired=False)
|
||||
|
||||
result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock)
|
||||
|
||||
assert result is None
|
||||
assert prisma.reconciled == []
|
||||
lock.release_lock.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins():
|
||||
prisma = _FakePrisma(user_days=("2026-09-13",))
|
||||
lock = _pod_lock(acquired=True)
|
||||
|
||||
result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock)
|
||||
|
||||
assert result is not None and result.days_reconciled == ("2026-09-13",)
|
||||
lock.release_lock.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read():
|
||||
"""A Redis outage must not stall the backfill: the day rewrite is idempotent, so running
|
||||
twice is only wasted effort while skipping forever leaves usage on the slow path."""
|
||||
prisma = _FakePrisma(user_days=("2026-09-13",))
|
||||
lock = _pod_lock(acquired=False)
|
||||
lock.redis_cache.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down"))
|
||||
|
||||
result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock)
|
||||
|
||||
assert result is not None and result.days_reconciled == ("2026-09-13",)
|
||||
lock.release_lock.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_marker_is_read_back_from_the_json_string_the_config_table_stores():
|
||||
prisma = _FakePrisma(user_days=())
|
||||
prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-10"}'
|
||||
|
||||
assert await reconciled_through(prisma) == "2026-09-10"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_unparseable_marker_reads_as_never_reconciled():
|
||||
prisma = _FakePrisma(user_days=())
|
||||
prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"something_else": 1}'
|
||||
|
||||
assert await reconciled_through(prisma) is None
|
||||
|
||||
|
||||
_rollup_postgresql_proc: Final = factories.postgresql_proc()
|
||||
_rollup_postgresql: Final = factories.postgresql("_rollup_postgresql_proc")
|
||||
|
||||
_MIGRATIONS_DIR: Final = (
|
||||
pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
|
||||
)
|
||||
_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql"
|
||||
|
||||
_DAILY_USER_SPEND_DDL: Final = """
|
||||
CREATE TABLE "LiteLLM_DailyUserSpend" (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT,
|
||||
date TEXT NOT NULL,
|
||||
api_key TEXT NOT NULL,
|
||||
model TEXT,
|
||||
model_group TEXT,
|
||||
custom_llm_provider TEXT,
|
||||
mcp_namespaced_tool_name TEXT,
|
||||
endpoint TEXT,
|
||||
prompt_tokens BIGINT DEFAULT 0,
|
||||
completion_tokens BIGINT DEFAULT 0,
|
||||
cache_read_input_tokens BIGINT DEFAULT 0,
|
||||
cache_creation_input_tokens BIGINT DEFAULT 0,
|
||||
compression_saved_tokens BIGINT DEFAULT 0,
|
||||
compression_savings_spend DOUBLE PRECISION DEFAULT 0,
|
||||
prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0,
|
||||
gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0,
|
||||
autorouter_savings_spend DOUBLE PRECISION DEFAULT 0,
|
||||
spend DOUBLE PRECISION DEFAULT 0,
|
||||
api_requests BIGINT DEFAULT 0,
|
||||
successful_requests BIGINT DEFAULT 0,
|
||||
failed_requests BIGINT DEFAULT 0,
|
||||
total_response_time_ms BIGINT DEFAULT 0,
|
||||
timed_requests BIGINT DEFAULT 0,
|
||||
created_at TIMESTAMP DEFAULT now(),
|
||||
updated_at TIMESTAMP,
|
||||
UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint)
|
||||
)
|
||||
"""
|
||||
|
||||
_PER_KEY_SUMS_SQL: Final = """
|
||||
SELECT COALESCE(model, '') AS model, COALESCE(model_group, '') AS model_group,
|
||||
COALESCE(custom_llm_provider, '') AS custom_llm_provider,
|
||||
SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, SUM(api_requests) AS api_requests,
|
||||
SUM(total_response_time_ms) AS total_response_time_ms, SUM(timed_requests) AS timed_requests
|
||||
FROM "LiteLLM_DailyUserSpend" WHERE date = %s
|
||||
GROUP BY 1, 2, 3 ORDER BY 1, 2, 3
|
||||
"""
|
||||
_GLOBAL_ROWS_SQL: Final = """
|
||||
SELECT model, model_group, custom_llm_provider, spend, prompt_tokens, api_requests,
|
||||
total_response_time_ms, timed_requests
|
||||
FROM "LiteLLM_DailyGlobalSpend" WHERE date = %s ORDER BY 1, 2, 3
|
||||
"""
|
||||
|
||||
|
||||
def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None:
|
||||
converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql)
|
||||
conn.execute(
|
||||
converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query
|
||||
{f"p{i}": v for i, v in enumerate(params, start=1)},
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
|
||||
def _user_txn(**overrides):
|
||||
return {
|
||||
"user_id": "u-1",
|
||||
"date": "2026-09-14",
|
||||
"api_key": "sk-1",
|
||||
"model": "gpt-5",
|
||||
"model_group": "gpt-5",
|
||||
"custom_llm_provider": "openai",
|
||||
"mcp_namespaced_tool_name": "",
|
||||
"endpoint": "/chat/completions",
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"spend": 1.0,
|
||||
"api_requests": 1,
|
||||
"successful_requests": 1,
|
||||
"failed_requests": 0,
|
||||
"total_response_time_ms": 800,
|
||||
"timed_requests": 1,
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]:
|
||||
return [
|
||||
(
|
||||
r["model"],
|
||||
r["model_group"],
|
||||
r["custom_llm_provider"],
|
||||
float(r["spend"]),
|
||||
int(r["prompt_tokens"]),
|
||||
int(r["api_requests"]),
|
||||
int(r["total_response_time_ms"]),
|
||||
int(r["timed_requests"]),
|
||||
) # pyright: ignore[reportArgumentType] # dict_row values are untyped
|
||||
for r in rows
|
||||
]
|
||||
|
||||
|
||||
def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_postgresql: psycopg.Connection):
|
||||
"""Against real Postgres and the shipped migration: writer-shaped rows and legacy rows
|
||||
(NULL and '' dimension spellings) fold into one global day, running the day twice changes
|
||||
nothing, and other days are left alone."""
|
||||
conn: Final = _rollup_postgresql
|
||||
conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal
|
||||
conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal
|
||||
conn.commit()
|
||||
|
||||
written_batch = merge_by_conflict_key(
|
||||
USER_TABLE,
|
||||
(_user_txn(api_key="sk-1", spend=1.0), _user_txn(api_key="sk-2", user_id="u-2", spend=2.0, prompt_tokens=20)),
|
||||
)
|
||||
_execute_dollar_sql(conn, *build_bulk_upsert(USER_TABLE, written_batch))
|
||||
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO "LiteLLM_DailyUserSpend"
|
||||
(id, user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name,
|
||||
endpoint, prompt_tokens, spend, api_requests)
|
||||
VALUES
|
||||
('legacy-1', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', NULL, 'openai', NULL, NULL, 5, 4.0, 1),
|
||||
('legacy-2', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', '', 'openai', '', '', 5, 8.0, 1),
|
||||
('legacy-3', 'u-9', '2026-09-13', 'sk-9', 'claude', '', 'anthropic', '', '', 7, 16.0, 1)
|
||||
"""
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
_execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",))
|
||||
_execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",))
|
||||
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
global_rows = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-14",)).fetchall()
|
||||
per_key = cur.execute(_PER_KEY_SUMS_SQL, ("2026-09-14",)).fetchall()
|
||||
untouched = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-13",)).fetchall()
|
||||
|
||||
assert _normalized(global_rows) == _normalized(per_key)
|
||||
assert sum(float(r["spend"]) for r in global_rows) == pytest.approx(15.0) # pyright: ignore[reportArgumentType] # dict_row values are untyped
|
||||
assert sum(int(r["total_response_time_ms"]) for r in global_rows) == 1600 # pyright: ignore[reportArgumentType] # dict_row values are untyped
|
||||
assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")]
|
||||
assert untouched == []
|
||||
|
||||
|
||||
_CONFIG_DDL: Final = 'CREATE TABLE "LiteLLM_Config" (param_name TEXT PRIMARY KEY, param_value JSONB)'
|
||||
_MARKER_SQL: Final = 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s'
|
||||
|
||||
|
||||
def test_advance_marker_sql_only_ever_moves_the_stored_marker_forward(_rollup_postgresql: psycopg.Connection):
|
||||
"""Against real Postgres: the statement a slower overlapping run issues after the faster run
|
||||
already stored a later marker leaves that marker alone, whether it carries an older scan time or
|
||||
none at all, while a run that is further along moves both fields on."""
|
||||
conn: Final = _rollup_postgresql
|
||||
conn.execute(_CONFIG_DDL) # pyright: ignore[reportArgumentType] # DDL literal
|
||||
conn.commit()
|
||||
param: Final = DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM
|
||||
|
||||
def stored() -> object:
|
||||
with conn.cursor(row_factory=dict_row) as cur:
|
||||
row = cur.execute(_MARKER_SQL, (param,)).fetchone()
|
||||
return None if row is None else row["param_value"]
|
||||
|
||||
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-01", None))
|
||||
assert stored() == {"reconciled_through": "2026-09-01", "scanned_at": None}
|
||||
|
||||
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-14", "2026-09-15 00:30:02.5"))
|
||||
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-02", None))
|
||||
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-03", "2026-09-15 00:30:01.25"))
|
||||
assert stored() == {"reconciled_through": "2026-09-14", "scanned_at": "2026-09-15 00:30:02.5"}
|
||||
|
||||
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-15", "2026-09-16 00:30:00.75"))
|
||||
assert stored() == {"reconciled_through": "2026-09-15", "scanned_at": "2026-09-16 00:30:00.75"}
|
||||
|
|
@ -178,28 +178,4 @@ describe("CacheLeakageCard", () => {
|
|||
screen.queryByText("Data is still loading; rows and totals will update as the rest of the range arrives."),
|
||||
).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("says which keys are missing from the key ranking when the proxy capped the per-key lists", () => {
|
||||
const day = dayWithKeys("2026-07-12", {
|
||||
"hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }),
|
||||
});
|
||||
renderWith([day], { apiKeyTruncation: { limit: 100, total: 3000 } });
|
||||
|
||||
expect(screen.getByRole("note")).toHaveTextContent(
|
||||
"Only the 100 highest-spend keys of 3,000 are loaded, so a lower-spend key that leaks more is not listed here.",
|
||||
);
|
||||
|
||||
fireEvent.click(screen.getByRole("tab", { name: "By model" }));
|
||||
|
||||
expect(screen.queryByRole("note")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the key ranking note off when every key was loaded", () => {
|
||||
const day = dayWithKeys("2026-07-12", {
|
||||
"hash-leaky": key("leaky-key", { prompt_tokens: 10000, cache_read_input_tokens: 0 }),
|
||||
});
|
||||
renderWith([day]);
|
||||
|
||||
expect(screen.queryByRole("note")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -81,7 +81,7 @@ const SortableHead = ({
|
|||
};
|
||||
|
||||
const CacheLeakageCard: React.FC<CacheLeakageCardProps> = ({ activity }) => {
|
||||
const { dateValue, onDateChange, results, loading, isFetchingMore, apiKeyTruncation } = activity;
|
||||
const { dateValue, onDateChange, results, loading, isFetchingMore } = activity;
|
||||
const [dimension, setDimension] = useState<CacheLeakageDimension>("key");
|
||||
const [sort, setSort] = useState<SortState>({ column: "potentialSavings", dir: "desc" });
|
||||
const leakage = useMemo(() => computeCacheLeakage(results, dimension), [results, dimension]);
|
||||
|
|
@ -123,13 +123,6 @@ const CacheLeakageCard: React.FC<CacheLeakageCardProps> = ({ activity }) => {
|
|||
</Tabs>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
{dimension === "key" && apiKeyTruncation !== undefined && (
|
||||
<p className="mb-2 text-sm text-muted-foreground" role="note">
|
||||
Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "}
|
||||
{apiKeyTruncation.total.toLocaleString()} are loaded, so a lower-spend key that leaks more is not listed
|
||||
here. Raise USAGE_TOP_API_KEYS_LIMIT on the proxy to load more keys.
|
||||
</p>
|
||||
)}
|
||||
{rows.length > 0 && isFetchingMore && (
|
||||
<p className="mb-2 text-sm text-muted-foreground">
|
||||
Data is still loading; rows and totals will update as the rest of the range arrives.
|
||||
|
|
|
|||
|
|
@ -4,13 +4,12 @@ import { describe, expect, it, vi } from "vitest";
|
|||
const mockUsePaginatedDailyActivity = vi.fn();
|
||||
|
||||
const mockCancel = vi.fn();
|
||||
let mockMetadata: Record<string, number> = {};
|
||||
|
||||
vi.mock("@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity", () => ({
|
||||
usePaginatedDailyActivity: (args: unknown) => {
|
||||
mockUsePaginatedDailyActivity(args);
|
||||
return {
|
||||
data: { results: [], metadata: mockMetadata },
|
||||
data: { results: [] },
|
||||
loading: false,
|
||||
isFetchingMore: false,
|
||||
progress: { currentPage: 4, totalPages: 9 },
|
||||
|
|
@ -81,18 +80,4 @@ describe("useDailyActivityRange", () => {
|
|||
|
||||
expect(mockUsePaginatedDailyActivity).toHaveBeenLastCalledWith(expect.objectContaining({ enabled: false }));
|
||||
});
|
||||
|
||||
it("reports how many keys the proxy left out of the per-key lists", () => {
|
||||
mockMetadata = { api_key_limit: 100, total_api_keys: 3000 };
|
||||
const { result } = renderHook(() => useDailyActivityRange("test-token", "u1", "proxy_admin"));
|
||||
|
||||
expect(result.current.apiKeyTruncation).toEqual({ limit: 100, total: 3000 });
|
||||
});
|
||||
|
||||
it("reports no key truncation when every key fit under the proxy limit", () => {
|
||||
mockMetadata = { api_key_limit: 100, total_api_keys: 100 };
|
||||
const { result } = renderHook(() => useDailyActivityRange("test-token", "u1", "proxy_admin"));
|
||||
|
||||
expect(result.current.apiKeyTruncation).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
import { useMemo, useState } from "react";
|
||||
|
||||
import { userDailyActivityAggregatedCall, userDailyActivityCall } from "@/components/networking";
|
||||
import { ApiKeyTruncation, getApiKeyTruncation } from "@/components/EntityUsageExport/exportBlockedReason";
|
||||
import { DailyData } from "@/components/UsagePage/types";
|
||||
import { spendScopeUserId } from "@/utils/roles";
|
||||
import { usePaginatedDailyActivity } from "@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity";
|
||||
|
|
@ -23,7 +22,6 @@ export interface DailyActivityRange {
|
|||
cancelled: boolean;
|
||||
failed: boolean;
|
||||
cancel: () => void;
|
||||
apiKeyTruncation?: ApiKeyTruncation;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -80,7 +78,6 @@ export const useScopedDailyActivityRange = (
|
|||
cancelled,
|
||||
failed,
|
||||
cancel,
|
||||
apiKeyTruncation: getApiKeyTruncation(data.metadata?.api_key_limit, data.metadata?.total_api_keys),
|
||||
};
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -650,23 +650,6 @@ describe("EntityUsage", () => {
|
|||
expect(screen.getAllByText("Activity Metrics")[1]).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("tells the team view how many keys the proxy left out of the per-key lists", async () => {
|
||||
mockTeamDailyActivityAggregatedCall.mockResolvedValue({
|
||||
...mockSpendData,
|
||||
metadata: { ...mockSpendData.metadata, api_key_limit: 100, total_api_keys: 3000 },
|
||||
});
|
||||
render(<EntityUsage {...defaultProps} entityType="team" />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockTeamDailyActivityAggregatedCall).toHaveBeenCalled();
|
||||
});
|
||||
act(() => {
|
||||
fireEvent.click(screen.getByText("Key Activity"));
|
||||
});
|
||||
|
||||
expect(await screen.findByRole("note")).toHaveTextContent("Only the 100 highest-spend keys of 3,000 are loaded");
|
||||
});
|
||||
|
||||
// An inactive tab panel is marked aria-selected="false" by one tab library and hidden by the
|
||||
// other, so treat either as "not on screen" and the assertion holds whichever one is rendering.
|
||||
const isShowing = (element: HTMLElement): boolean => {
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ import TeamMultiSelect from "@/components/common_components/team_multi_select";
|
|||
import UserDropdown from "@/components/common_components/UserDropdown";
|
||||
import { ActivityMetrics, processActivityData } from "@/components/activity_metrics";
|
||||
import { UsageExportHeader } from "@/components/EntityUsageExport";
|
||||
import { getApiKeyTruncation, getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason";
|
||||
import { getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason";
|
||||
import type { EntityType } from "@/components/EntityUsageExport/types";
|
||||
import {
|
||||
agentDailyActivityCall,
|
||||
|
|
@ -71,8 +71,6 @@ interface EntitySpendData {
|
|||
total_successful_requests: number;
|
||||
total_failed_requests: number;
|
||||
total_tokens: number;
|
||||
api_key_limit?: number | null;
|
||||
total_api_keys?: number | null;
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -162,7 +160,6 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
|
|||
});
|
||||
|
||||
const spendData = spendDataRaw as unknown as EntitySpendData;
|
||||
const apiKeyTruncation = getApiKeyTruncation(spendData.metadata?.api_key_limit, spendData.metadata?.total_api_keys);
|
||||
|
||||
const {
|
||||
data: agentSpendDataRaw,
|
||||
|
|
@ -662,18 +659,12 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
|
|||
{
|
||||
key: "keys",
|
||||
label: "Key Activity",
|
||||
content: (
|
||||
<KeyActivityPanel
|
||||
keyMetrics={keyMetrics}
|
||||
hidePromptCachingMetrics={entityType === "agent"}
|
||||
apiKeyTruncation={apiKeyTruncation}
|
||||
/>
|
||||
),
|
||||
content: <KeyActivityPanel keyMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} />,
|
||||
},
|
||||
{ key: "endpoints", label: "Endpoint Activity", content: <EndpointUsage userSpendData={spendData} /> },
|
||||
];
|
||||
|
||||
const spendFetchState = { coversRange, cancelled, failed, apiKeyTruncation };
|
||||
const spendFetchState = { coversRange, cancelled, failed };
|
||||
|
||||
return (
|
||||
<div style={{ width: "100%" }} className="relative">
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ import { ActivityMetrics, processActivityData } from "@/components/activity_metr
|
|||
import CloudZeroExportModal from "@/components/cloudzero_export_modal";
|
||||
import UserDropdown from "@/components/common_components/UserDropdown";
|
||||
import EntityUsageExportModal from "@/components/EntityUsageExport";
|
||||
import { getApiKeyTruncation, getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason";
|
||||
import { getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason";
|
||||
import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel";
|
||||
import { Team } from "@/components/key_team_helpers/key_list";
|
||||
import {
|
||||
|
|
@ -256,10 +256,6 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
|
|||
coversRange: activeAggregated !== null || paginatedResult.coversRange,
|
||||
cancelled: paginatedResult.cancelled,
|
||||
failed: paginatedResult.failed,
|
||||
apiKeyTruncation: getApiKeyTruncation(
|
||||
userSpendData.metadata?.api_key_limit,
|
||||
userSpendData.metadata?.total_api_keys,
|
||||
),
|
||||
};
|
||||
const exportBlockedReason = getExportBlockedReason(spendFetchState);
|
||||
|
||||
|
|
@ -865,7 +861,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
|
|||
<ActivityMetrics modelMetrics={modelMetrics} />
|
||||
</TabsContent>
|
||||
<TabsContent value="keys" keepMounted>
|
||||
<KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={spendFetchState.apiKeyTruncation} />
|
||||
<KeyActivityPanel keyMetrics={keyMetrics} />
|
||||
</TabsContent>
|
||||
<TabsContent value="mcp" keepMounted>
|
||||
<ActivityMetrics modelMetrics={mcpServerMetrics} />
|
||||
|
|
|
|||
|
|
@ -1,12 +1,11 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
|
||||
import { getApiKeyTruncation, getExportBlockedReason, type UsageFetchState } from "./exportBlockedReason";
|
||||
import { getExportBlockedReason, type UsageFetchState } from "./exportBlockedReason";
|
||||
|
||||
const state = (overrides: Partial<UsageFetchState> = {}): UsageFetchState => ({
|
||||
coversRange: true,
|
||||
cancelled: false,
|
||||
failed: false,
|
||||
apiKeyTruncation: undefined,
|
||||
...overrides,
|
||||
});
|
||||
|
||||
|
|
@ -32,27 +31,4 @@ describe("getExportBlockedReason", () => {
|
|||
expect(reason).toMatch(/failed to load/i);
|
||||
expect(reason).not.toMatch(/stopped/i);
|
||||
});
|
||||
|
||||
it("blocks when the aggregated endpoint dropped keys, since a per-team CSV would miss them", () => {
|
||||
const reason = getExportBlockedReason(state({ apiKeyTruncation: { limit: 100, total: 3000 } }));
|
||||
|
||||
expect(reason).toMatch(/100 highest-spend keys of 3000/);
|
||||
expect(reason).toMatch(/USAGE_TOP_API_KEYS_LIMIT/);
|
||||
});
|
||||
});
|
||||
|
||||
describe("getApiKeyTruncation", () => {
|
||||
it("reports truncation once the proxy saw more keys than it returned", () => {
|
||||
expect(getApiKeyTruncation(100, 101)).toEqual({ limit: 100, total: 101 });
|
||||
});
|
||||
|
||||
it("stays quiet when exactly the cap exists, since every key is on screen", () => {
|
||||
expect(getApiKeyTruncation(100, 100)).toBeUndefined();
|
||||
expect(getApiKeyTruncation(100, 7)).toBeUndefined();
|
||||
});
|
||||
|
||||
it("stays quiet when the response carries no cap, as the paginated fallback does", () => {
|
||||
expect(getApiKeyTruncation(undefined, undefined)).toBeUndefined();
|
||||
expect(getApiKeyTruncation(100, null)).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,31 +1,13 @@
|
|||
export interface ApiKeyTruncation {
|
||||
limit: number;
|
||||
total: number;
|
||||
}
|
||||
|
||||
export interface UsageFetchState {
|
||||
coversRange: boolean;
|
||||
cancelled: boolean;
|
||||
failed: boolean;
|
||||
apiKeyTruncation: ApiKeyTruncation | undefined;
|
||||
}
|
||||
|
||||
export const getApiKeyTruncation = (apiKeyLimit: unknown, totalApiKeys: unknown): ApiKeyTruncation | undefined => {
|
||||
if (typeof apiKeyLimit !== "number" || typeof totalApiKeys !== "number") return undefined;
|
||||
return totalApiKeys > apiKeyLimit ? { limit: apiKeyLimit, total: totalApiKeys } : undefined;
|
||||
};
|
||||
|
||||
export const getExportBlockedReason = ({
|
||||
coversRange,
|
||||
cancelled,
|
||||
failed,
|
||||
apiKeyTruncation,
|
||||
}: UsageFetchState): string | undefined => {
|
||||
export const getExportBlockedReason = ({ coversRange, cancelled, failed }: UsageFetchState): string | undefined => {
|
||||
if (failed) return "Some spend data failed to load, so an export would under-report. Reload the page to try again.";
|
||||
if (cancelled)
|
||||
return "Loading was stopped before the whole range arrived, so an export would under-report. Reload the page to load it all.";
|
||||
if (!coversRange) return "Spend data is still loading, so an export would under-report. Wait for it to finish.";
|
||||
if (apiKeyTruncation !== undefined)
|
||||
return `Only the ${apiKeyTruncation.limit} highest-spend keys of ${apiKeyTruncation.total} were loaded, so a per-team export would under-report. Raise USAGE_TOP_API_KEYS_LIMIT on the proxy to load more keys.`;
|
||||
return undefined;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -68,14 +68,4 @@ describe("KeyActivityPanel", () => {
|
|||
expect(screen.getByLabelText("Search keys")).toHaveValue("");
|
||||
expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alicehash-bob");
|
||||
});
|
||||
|
||||
it("says how many keys the proxy left out when only the top spenders were loaded", () => {
|
||||
render(<KeyActivityPanel keyMetrics={keyMetrics} apiKeyTruncation={{ limit: 2, total: 3000 }} />);
|
||||
expect(screen.getByRole("note")).toHaveTextContent("Only the 2 highest-spend keys of 3,000 are loaded");
|
||||
});
|
||||
|
||||
it("shows no truncation note when every key is loaded", () => {
|
||||
render(<KeyActivityPanel keyMetrics={keyMetrics} />);
|
||||
expect(screen.queryByRole("note")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@ import { Search, X } from "lucide-react";
|
|||
import React, { useMemo, useState } from "react";
|
||||
|
||||
import { ActivityMetrics } from "@/components/activity_metrics";
|
||||
import type { ApiKeyTruncation } from "@/components/EntityUsageExport/exportBlockedReason";
|
||||
import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group";
|
||||
|
||||
import { filterKeyActivity } from "../keyActivityFilter";
|
||||
|
|
@ -11,14 +10,9 @@ import type { ModelActivityData } from "../types";
|
|||
interface KeyActivityPanelProps {
|
||||
keyMetrics: Record<string, ModelActivityData>;
|
||||
hidePromptCachingMetrics?: boolean;
|
||||
apiKeyTruncation?: ApiKeyTruncation;
|
||||
}
|
||||
|
||||
const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({
|
||||
keyMetrics,
|
||||
hidePromptCachingMetrics = false,
|
||||
apiKeyTruncation,
|
||||
}) => {
|
||||
const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({ keyMetrics, hidePromptCachingMetrics = false }) => {
|
||||
const [query, setQuery] = useState("");
|
||||
const filtered = useMemo(() => filterKeyActivity(keyMetrics, query), [keyMetrics, query]);
|
||||
const totalKeys = Object.keys(keyMetrics).length;
|
||||
|
|
@ -49,12 +43,6 @@ const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({
|
|||
<span className="text-sm text-muted-foreground">
|
||||
Showing {shownKeys.toLocaleString()} of {totalKeys.toLocaleString()} keys
|
||||
</span>
|
||||
{apiKeyTruncation !== undefined && (
|
||||
<span className="text-sm text-muted-foreground" role="note">
|
||||
Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "}
|
||||
{apiKeyTruncation.total.toLocaleString()} are loaded
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
{isFiltering && totalKeys > 0 && shownKeys === 0 ? (
|
||||
<p className="rounded-lg border p-6 text-center text-sm text-muted-foreground">
|
||||
|
|
|
|||
10
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
10
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -27961,11 +27961,6 @@ export interface components {
|
|||
};
|
||||
/** DailySpendMetadata */
|
||||
DailySpendMetadata: {
|
||||
/**
|
||||
* Api Key Limit
|
||||
* @description When set, api_keys and every api_key_breakdown list at most this many keys, ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.
|
||||
*/
|
||||
api_key_limit?: number | null;
|
||||
/**
|
||||
* Has More
|
||||
* @default false
|
||||
|
|
@ -27976,11 +27971,6 @@ export interface components {
|
|||
* @default 1
|
||||
*/
|
||||
page: number;
|
||||
/**
|
||||
* Total Api Keys
|
||||
* @description Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key lists are truncated to the highest-spend keys.
|
||||
*/
|
||||
total_api_keys?: number | null;
|
||||
/**
|
||||
* Total Api Requests
|
||||
* @default 0
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue