scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(proxy): bound api_key rollups in aggregated usage query to top keys

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus 2026-09-08 15:52:30 +00:00
parent 1af7a403c6
commit 49fde3761b
2 changed files with 80 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -710,20 +710,33 @@ def _build_aggregated_sql_query(
# is omitted on purpose: nothing in the response shape needs it once
# all the rollups are present.
#
# api_key appears as tk.top_api_key from the top_api_keys CTE, bounding
# the api_key-keyed sets to the top _MAX_API_KEYS_IN_BREAKDOWN keys by
# spend instead of every distinct key in the window. The same
# where_clause/$N params run in both the CTE and the outer query.
#
# 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"""
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
LIMIT {_MAX_API_KEYS_IN_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,
@ -740,21 +753,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),
()
)
"""
@ -930,6 +944,16 @@ _GROUP_DATE_MCP_API_KEY: Final = 29 # 0b0011101
_GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110
_GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110
# Cap on distinct api_keys carried into the api_key-keyed grouping sets of
# _build_aggregated_sql_query. Six of the thirteen sets include api_key, so
# result rows scale with distinct-key count; on large deployments the
# prisma-query-engine buffers the whole result and gets OOM-killed. The UI
# only renders up to 50 top keys (TOP_KEYS_LIMITS), so 100 is generous.
# Keys outside the top N group into a NULL api_key bucket that the
# dispatcher skips, and non-keyed totals are unaffected (LEFT JOIN keeps
# every row).
_MAX_API_KEYS_IN_BREAKDOWN: Final = 100
def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
"""Build a SpendMetrics directly from one already-aggregated rollup row.

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

@ -9,6 +9,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_E
from litellm.proxy.management_endpoints.common_daily_activity import (
_MAX_API_KEYS_IN_BREAKDOWN,
_adjust_dates_for_timezone,
_build_aggregated_sql_query,
_build_entity_rollup_sql_query,
@ -1149,13 +1150,57 @@ 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
def test_api_key_rollups_are_bounded_to_top_keys(self):
"""api_key-keyed grouping sets must group on the bounded top-N CTE, not the raw column.
Six of the thirteen grouping sets include api_key, so result rows used
to scale with the total distinct-key count in the window; the
prisma-query-engine buffered the whole result and got OOM-killed. The
top_api_keys CTE bounds the dimension to the top
_MAX_API_KEYS_IN_BREAKDOWN keys by spend; every api_key grouping must
reference tk.top_api_key, and grouping on the raw api_key column is the
regression this guards.
"""
sql, params = _build_aggregated_sql_query(
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id="user-1",
start_date="2026-05-29",
end_date="2026-06-02",
model="bedrock/global.anthropic.claude-opus-4-8",
api_key="sk-test",
timezone_offset_minutes=-330,
)
normalized = " ".join(sql.split())
assert "WITH top_api_keys AS ( SELECT api_key AS top_api_key" in normalized
assert f"ORDER BY SUM(spend) DESC LIMIT {_MAX_API_KEYS_IN_BREAKDOWN}" in normalized
grouping_block = normalized.split("GROUP BY GROUPING SETS (", 1)[1]
grouping_sets = {part.strip() for part in grouping_block.split("),")}
for grouping_set in grouping_sets:
assert "api_key" not in grouping_set.replace("tk.top_api_key", ""), (
f"grouping set uses the unbounded raw api_key column: {grouping_set}"
)
assert grouping_block.count("tk.top_api_key") == 6
# The CTE and the outer WHERE reuse the same $N placeholders, so the
# params list is unchanged: date bounds, entity, model, api_key filter.
assert params == [
"2026-05-29",
"2026-06-02",
"user-1",
"bedrock/global.anthropic.claude-opus-4-8",
"sk-test",
]
class TestAggregatedEmptyEntityFilter:
_BUILDERS: Final = (_build_aggregated_sql_query, _build_entity_rollup_sql_query)