diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index ce6a97708ab..e6e06accdb6 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -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. diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 37a54c4901a..173c4b4cccd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -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)