revert: "perf(proxy): split aggregated usage query into key-free rollups and bounded top-N keys (#41293)"

litellm/constants.py

@@ -2140,7 +2140,6 @@
 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

litellm/proxy/_lazy_openapi_snapshot.json

@@ -3313,18 +3313,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",
@@ -3335,18 +3323,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",

litellm/proxy/management_endpoints/common_daily_activity.py

@@ -9,7 +9,7 @@
 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.key_metadata_recovery import (
     attach_user_details,
@@ -150,9 +150,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
@@ -170,46 +176,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``
@@ -733,29 +705,6 @@ 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"""
-            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"""
-
-
-_MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)"
-
-
 def _build_aggregated_sql_query(
     *,
     table_name: str,
@@ -769,7 +718,12 @@ def _build_aggregated_sql_query(
     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 the GROUPING SETS query for aggregated daily activity.
+    """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().
@@ -782,7 +736,7 @@ def _build_aggregated_sql_query(
         start_date, end_date, timezone_offset_minutes, include_current_utc_day
     )
 
-    where_clause, where_params = _build_aggregated_where_clause(
+    where_clause, sql_params = _build_aggregated_where_clause(
         entity_id_field=entity_id_field,
         entity_id=entity_id,
         adjusted_start=adjusted_start,
@@ -791,72 +745,67 @@ def _build_aggregated_sql_query(
         api_key=api_key,
         exclude_entity_ids=exclude_entity_ids,
     )
-    sentinel_param: Final = f"${len(where_params) + 1}"
-    metric_select: Final = _rollup_metric_select(table_name)
 
+    # 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,
-            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 "{pg_table}"
-        WHERE {where_clause}
-        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,
+            COALESCE(NULLIF(model_group, ''), model) AS model_group,
             custom_llm_provider,
             mcp_namespaced_tool_name,
             endpoint,
-            GROUPING(date, api_key, model, {_MODEL_GROUP_EXPR},
+            GROUPING(date, api_key, model, COALESCE(NULLIF(model_group, ''), model),
                      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(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 (
+            (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),
+            ()
+        )
     """
 
-    return sql_query, [*where_params, PTU_SENTINEL_API_KEY]
+    return sql_query, sql_params
 
 
 def _build_entity_rollup_sql_query(
@@ -901,7 +850,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 (
@@ -1003,7 +968,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
@@ -1017,7 +981,7 @@ async def _aggregate_spend_records(
 _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
@@ -1373,6 +1337,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.
 
@@ -1391,7 +1359,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,
@@ -1403,16 +1371,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)
-        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
+        )
+
+        # 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 ())]
-        total_api_keys: Final = next((r.distinct_api_keys for r in records if r.distinct_api_keys is not None), 0)
+        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.
@@ -1466,8 +1454,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,
             ),
         )