diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index a8aef30107c..bd79501847a 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -742,15 +742,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_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, @@ -767,21 +775,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), () ) """ @@ -959,6 +968,11 @@ _GROUP_DATE_MCP_API_KEY: Final = 29 # 0b0011101 _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110 _GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110 +# Six of the thirteen grouping sets in _build_aggregated_sql_query are keyed on api_key, so +# without a cap the result grows with every distinct key and the prisma query engine OOMs +# buffering it. Keys outside the top N fall into a NULL api_key bucket the dispatcher skips. +_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 6cd900cb041..9131d9195f1 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 @@ -1,14 +1,16 @@ +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 litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR - +from psycopg.rows import dict_row +from pytest_postgresql import factories 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, @@ -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, @@ -472,9 +475,7 @@ 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( @@ -1236,10 +1237,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 @@ -1360,6 +1361,128 @@ 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, +): + """Exercise the GROUPING SETS query against real Postgres with more keys than the cap. + + key-004 and key-005 tie on spend exactly at the _MAX_API_KEYS_IN_BREAKDOWN + cutoff; the deterministic api_key tiebreaker must keep key-004 and drop + key-005. Excluded keys still count toward the totals via the grand-total + and date-level rollup rows. + """ + 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_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, + ) + + # 13 grouping sets: 7 key-less sets emit one row each; the 6 api_key sets + # emit the top N keys plus one NULL bucket for the excluded keys. + assert row_counts == [7 + 6 * (_MAX_API_KEYS_IN_BREAKDOWN + 1)] + + # Totals still cover every key, capped or not. + assert result.metadata.total_spend == pytest.approx(5566.0) + assert result.metadata.total_api_requests == _MAX_API_KEYS_IN_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_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 + + def _no_spend_record(): """A rollup row for a key with no spend, where SUM() returns NULL (None).""" return SimpleNamespace(