diff --git a/litellm/constants.py b/litellm/constants.py index c106be688e4..3900a62b30b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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( diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 44ed0017e42..9d3fdee90c1 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -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), () ) """ 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 71896a18f48..068c0271f10 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,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(