From 7050e0c925388339f4e37da7adf437cf776882e6 Mon Sep 17 00:00:00 2001 From: Abhimanyu Kapur <38531241+akapur99@users.noreply.github.com> Date: Sat, 8 Aug 2026 15:34:05 -0700 Subject: [PATCH] perf: read shadow eval volume estimate from the daily rollup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The start endpoint counted the key's trailing 7-day requests directly against LiteLLM_SpendLogs, which has no api_key index — on a busy proxy that is a scan over every request in the window (potentially tens of millions of rows) to answer one count, holding a multi-second query per 'Start shadow eval' click. Read SUM(api_requests) from LiteLLM_DailyUserSpend instead: one indexed row per key/day, the same table the usage dashboards already use for this question. A regression test asserts the estimate queries the rollup and never touches LiteLLM_SpendLogs. Co-Authored-By: Claude --- .../auto_router_endpoints.py | 36 +++++++++++++++---- .../test_auto_router_endpoints.py | 18 +++++++++- 2 files changed, 46 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 9ff8bf001ec..437bee9fc69 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -51,6 +51,7 @@ from litellm.types.management_endpoints.auto_router_endpoints import ( if TYPE_CHECKING: from fastapi import APIRouter, Depends, HTTPException, Query, status + from litellm.proxy.utils import PrismaClient from litellm.router import Router else: try: @@ -453,6 +454,33 @@ _JUDGE_PROMPT_TOKENS_ESTIMATE: Final = 4000 # The estimate projects from the key's request volume over this many trailing days. _ESTIMATE_LOOKBACK_DAYS: Final = 7 +_ESTIMATE_VOLUME_SQL: Final = """ +SELECT COALESCE(SUM(api_requests), 0)::bigint AS request_count +FROM "LiteLLM_DailyUserSpend" +WHERE api_key = $1 AND date >= $2 +""" + + +class _EstimateVolumeRow(BaseModel): + request_count: int + + +_ESTIMATE_VOLUME_ROWS: Final = TypeAdapter(list[_EstimateVolumeRow]) + + +async def _recent_request_volume(prisma_client: "PrismaClient", api_key_id: str) -> int: + """The key's request count over the estimate lookback window. + + Read from the LiteLLM_DailyUserSpend rollup, never LiteLLM_SpendLogs: the raw log + table has no api_key index, so a per-key count there scans every request the proxy + served in the window. The rollup is a handful of indexed rows per key/day and is + how the usage dashboards answer the same question. + """ + lookback_date: Final = (datetime.now(timezone.utc) - timedelta(days=_ESTIMATE_LOOKBACK_DAYS)).strftime("%Y-%m-%d") + raw_rows: Final = await prisma_client.db.query_raw(_ESTIMATE_VOLUME_SQL, api_key_id, lookback_date) + rows: Final = _ESTIMATE_VOLUME_ROWS.validate_python(raw_rows or []) + return rows[0].request_count if rows else 0 + def _require_admin_viewer(user_api_key_dict: UserAPIKeyAuth, action: str) -> None: if user_api_key_dict.user_role not in ( @@ -634,13 +662,7 @@ async def start_shadow_eval( detail=f"Key already has an active shadow eval job ({existing.id}). Stop it first.", ) - lookback_start: Final = datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(days=_ESTIMATE_LOOKBACK_DAYS) - recent_requests: Final = await prisma_client.db.litellm_spendlogs.count( - where={ # mutable-ok: Prisma filter - "api_key": data.api_key_id, - "startTime": {"gte": lookback_start}, # mutable-ok: Prisma filter - }, - ) + recent_requests: Final = await _recent_request_volume(prisma_client, data.api_key_id) sampled: Final = int( recent_requests * (data.duration_days / _ESTIMATE_LOOKBACK_DAYS) * data.shadow_percentage / 100.0 ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 0da0bac5691..ae19fa56c7c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -574,7 +574,7 @@ class TestShadowEvalJobsAreTimeBound: router.adaptive_routers = {} router.quality_routers = {} prisma = MagicMock() - prisma.db.litellm_spendlogs.count = AsyncMock(return_value=recent_requests) + prisma.db.query_raw = AsyncMock(return_value=[{"request_count": recent_requests}]) prisma.db.litellm_shadowevaljob.find_first = AsyncMock(return_value=None) created = MagicMock() created.id = "job-1" @@ -583,6 +583,22 @@ class TestShadowEvalJobsAreTimeBound: monkeypatch.setattr(proxy_server, "prisma_client", prisma) return prisma + @pytest.mark.asyncio + async def test_estimate_reads_the_daily_rollup_not_the_raw_spend_log_table(self, monkeypatch: pytest.MonkeyPatch): + """LiteLLM_SpendLogs has no api_key index; a per-key count there scans every + request in the window. The estimate must come from LiteLLM_DailyUserSpend.""" + from litellm.proxy.management_endpoints.auto_router_endpoints import start_shadow_eval + + prisma = self._proxy_mocks(monkeypatch, recent_requests=700) + + response = await start_shadow_eval(self._start_request(), ADMIN) + + assert response.estimated_request_count == 70 + prisma.db.litellm_spendlogs.count.assert_not_called() + sql = prisma.db.query_raw.call_args.args[0] + assert 'FROM "LiteLLM_DailyUserSpend"' in sql + assert "LiteLLM_SpendLogs" not in sql + def test_duration_defaults_to_a_week_and_rejects_zero_and_over_a_month(self): assert self._start_request().duration_days == 7 with pytest.raises(ValidationError):