scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix(proxy): cap api_key rollups in the aggregated usage query to the top 100 keys by spend

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-14 22:58:15 +00:00
parent 7fd541efb9
commit d6eac1f974
3 changed files with 152 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/constants.py
View file

@ -65,6 +65,10 @@ DEFAULT_MAX_RETRIES: Final = int(os.getenv("DEFAULT_MAX_RETRIES", 2))
# Max records accepted in one POST /v1/callbacks/logs batch. Bounds the blast
# radius: each record fans out to spend logs + every callback integration.
MAX_CALLBACK_LOG_RECORDS: Final = 1000
# Top keys by spend that get per-key rows in the aggregated usage response.
# Six of its GROUPING SETS are keyed on api_key, so an uncapped result grows
# with every distinct key and the prisma query engine OOMs buffering it.
MAX_API_KEYS_IN_USAGE_BREAKDOWN: Final = 100
DEFAULT_MAX_RECURSE_DEPTH: Final = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH", 100))
DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER", 10))
DEFAULT_FAILURE_THRESHOLD_PERCENT: Final = float(

32
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -9,7 +9,7 @@ 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
from litellm.constants import MAX_API_KEYS_IN_USAGE_BREAKDOWN, PTU_SENTINEL_API_KEY
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.spend_tracking.key_metadata_recovery import (
attach_user_emails,
@ -708,6 +708,9 @@ def _build_aggregated_sql_query(
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.
The api_key rollups only cover the MAX_API_KEYS_IN_USAGE_BREAKDOWN top
keys by spend; every other key lands in a NULL api_key bucket so the
date, model, provider and total rollups still sum every row.
Returns:
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
@ -743,15 +746,23 @@ def _build_aggregated_sql_query(
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
# api_requests rollups are still served from here.
sql_query: Final = f"""
WITH top_api_keys AS (
SELECT api_key AS top_api_key
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY api_key
ORDER BY SUM(spend) DESC, api_key
LIMIT {MAX_API_KEYS_IN_USAGE_BREAKDOWN}
)
SELECT
date,
api_key,
tk.top_api_key AS 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),
GROUPING(date, tk.top_api_key, model, COALESCE(NULLIF(model_group, ''), model),
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
SUM(spend)::float AS spend,
@ -768,21 +779,22 @@ def _build_aggregated_sql_query(
SUM(api_requests)::bigint AS api_requests,
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests
FROM "{pg_table}"
FROM "{pg_table}" t
LEFT JOIN top_api_keys tk ON tk.top_api_key = t.api_key
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date),
(date, api_key),
(date, tk.top_api_key),
(date, model),
(date, model, api_key),
(date, model, tk.top_api_key),
(date, COALESCE(NULLIF(model_group, ''), model)),
(date, COALESCE(NULLIF(model_group, ''), model), api_key),
(date, COALESCE(NULLIF(model_group, ''), model), tk.top_api_key),
(date, custom_llm_provider),
(date, custom_llm_provider, api_key),
(date, custom_llm_provider, tk.top_api_key),
(date, mcp_namespaced_tool_name),
(date, mcp_namespaced_tool_name, api_key),
(date, mcp_namespaced_tool_name, tk.top_api_key),
(date, endpoint),
(date, endpoint, api_key),
(date, endpoint, tk.top_api_key),
()
)
"""

131
tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py
View file

@ -1,13 +1,15 @@
import re
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.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
from litellm.constants import MAX_API_KEYS_IN_USAGE_BREAKDOWN
from litellm.proxy.management_endpoints.common_daily_activity import (
_adjust_dates_for_timezone,
_build_aggregated_sql_query,
@ -19,6 +21,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
get_daily_activity_aggregated,
update_metrics,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
from litellm.types.proxy.management_endpoints.common_daily_activity import (
DailySpendMetadata,
SpendMetrics,
@ -1241,10 +1244,10 @@ class TestBuildAggregatedSqlQuery:
fallback = "COALESCE(NULLIF(model_group, ''), model)"
assert f"{fallback} AS model_group" in normalized
assert (
f"GROUPING(date, api_key, model, {fallback}, "
f"GROUPING(date, tk.top_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 f"(date, {fallback}), (date, {fallback}, tk.top_api_key)," in normalized
assert "(date, model_group)" not in normalized
assert "COALESCE(model_group, model)" not in normalized
@ -1365,6 +1368,124 @@ 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")
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_bounds_api_key_rollups(
_aggregated_postgresql: psycopg.Connection,
):
"""key-004 and key-005 tie on spend at the cutoff: the api_key tiebreaker keeps key-004 only."""
conn: Final = _aggregated_postgresql
with conn.cursor() as cur:
cur.execute(
"""
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
)
"""
)
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)
""",
[
(
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(MAX_API_KEYS_IN_USAGE_BREAKDOWN + 5)
],
)
conn.commit()
row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim
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
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = query_raw
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,
)
keyless_sets: Final = 7
api_key_sets: Final = 6
assert row_counts == [keyless_sets + api_key_sets * (MAX_API_KEYS_IN_USAGE_BREAKDOWN + 1)]
assert result.metadata.total_spend == pytest.approx(5566.0)
assert result.metadata.total_api_requests == MAX_API_KEYS_IN_USAGE_BREAKDOWN + 5
expected_top: Final = {f"key-{i:03d}" for i in range(6, 105)} | {"key-004"}
day: Final = result.results[0]
assert len(day.breakdown.api_keys) == MAX_API_KEYS_IN_USAGE_BREAKDOWN
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 day.breakdown.models["gpt-5"].metrics.spend == pytest.approx(5566.0)
assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == expected_top
assert day.breakdown.providers["openai"].metrics.spend == pytest.approx(5566.0)
assert set(day.breakdown.providers["openai"].api_key_breakdown) == expected_top
assert set(day.breakdown.endpoints["/v1/chat/completions"].api_key_breakdown) == expected_top
def _no_spend_record():
"""A rollup row for a key with no spend, where SUM() returns NULL (None)."""
return SimpleNamespace(