diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_shadow_eval_multi_key/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_shadow_eval_multi_key/migration.sql new file mode 100644 index 00000000000..fb4fb1b47de --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260814000000_shadow_eval_multi_key/migration.sql @@ -0,0 +1,59 @@ +CREATE TABLE IF NOT EXISTS "LiteLLM_ShadowEvalJobKey" ( + "id" TEXT NOT NULL, + "job_id" TEXT NOT NULL, + "direction" TEXT NOT NULL, + "api_key_id" TEXT NOT NULL, + "max_turns" INTEGER NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "stopped_at" TIMESTAMP(3), + + CONSTRAINT "LiteLLM_ShadowEvalJobKey_pkey" PRIMARY KEY ("id") +); + +CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJobKey_job_id_idx" ON "LiteLLM_ShadowEvalJobKey"("job_id"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJobKey_api_key_id_idx" ON "LiteLLM_ShadowEvalJobKey"("api_key_id"); + +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_id_direction_key" ON "LiteLLM_ShadowEvalJob"("id", "direction"); + +DO $$ BEGIN + ALTER TABLE "LiteLLM_ShadowEvalJobKey" ADD CONSTRAINT "LiteLLM_ShadowEvalJobKey_job_id_direction_fkey" + FOREIGN KEY ("job_id", "direction") REFERENCES "LiteLLM_ShadowEvalJob"("id", "direction") + ON DELETE CASCADE ON UPDATE CASCADE; +EXCEPTION + WHEN duplicate_object THEN NULL; +END $$; + +-- One active job per key per direction, enforced by the database rather than a read-then-create in +-- the start endpoint, which races against a concurrent start on another pod. Partial indexes are not +-- expressible in schema.prisma, so this lives here only. Active means not yet stopped; the start +-- endpoint stamps stopped_at on exhausted key rows before creating. +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJobKey_one_active_per_key_direction" + ON "LiteLLM_ShadowEvalJobKey"("api_key_id", "direction") WHERE "stopped_at" IS NULL; + +INSERT INTO "LiteLLM_ShadowEvalJobKey" ("id", "job_id", "direction", "api_key_id", "max_turns", "created_at", "stopped_at") +SELECT 'jobkey_' || "id", "id", "direction", "api_key_id", "max_turns", "created_at", "stopped_at" +FROM "LiteLLM_ShadowEvalJob" +ON CONFLICT DO NOTHING; + +ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "api_key_id" TEXT; + +UPDATE "LiteLLM_ShadowEvalAttempt" a SET "api_key_id" = j."api_key_id" +FROM "LiteLLM_ShadowEvalJob" j WHERE a."job_id" = j."id" AND a."api_key_id" IS NULL; + +ALTER TABLE "LiteLLM_ShadowEvalAttempt" ALTER COLUMN "api_key_id" SET NOT NULL; + +CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalAttempt_job_id_api_key_id_idx" ON "LiteLLM_ShadowEvalAttempt"("job_id", "api_key_id"); + +DROP INDEX IF EXISTS "LiteLLM_ShadowEvalAttempt_job_id_idx"; + +DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_one_active_per_key_direction"; + +DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_api_key_id_idx"; + +-- Key scope, per-key budget and stop state moved into LiteLLM_ShadowEvalJobKey above. Dropping them +-- in the same migration is safe because no tagged release ships shadow eval, so no pod is serving +-- the single key code and there is no mixed version window that needs these columns readable. +ALTER TABLE "LiteLLM_ShadowEvalJob" DROP COLUMN IF EXISTS "api_key_id", + DROP COLUMN IF EXISTS "max_turns", + DROP COLUMN IF EXISTS "stopped_at"; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 71345d2ccde..a035e8ec623 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1450,36 +1450,57 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } -// Shadow eval: evaluation of an auto-router against a key's live traffic, in either -// direction. forward duplicates the requests the key did not route through the router -// through it, answering whether the key should adopt it; reverse duplicates the requests +// Shadow eval: evaluation of an auto-router against a set of keys' live traffic, in either +// direction. forward duplicates the requests the keys did not route through the router +// through it, answering whether they should adopt it; reverse duplicates the requests // the router did serve against a fixed baseline model, answering whether a key already on // it still benefits. Either way a sampled slice runs in a detached task and an LLM judge -// compares real vs shadow responses blind. The job row is immutable config plus -// stopped_at; every count, status, and spend figure is derived from the append-only -// attempt rows, so nothing can disagree across pods or stop races. +// compares real vs shadow responses blind. The job row is immutable config; the keys it +// scopes, their budgets and their stop state live on LiteLLM_ShadowEvalJobKey, and every +// count, status, and spend figure is derived from the append-only attempt rows, so nothing +// can disagree across pods or stop races. model LiteLLM_ShadowEvalJob { id String @id @default(cuid()) - api_key_id String // hashed virtual key whose traffic is shadowed router_name String // the auto-router under evaluation, in either direction direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String shadow_percentage Float - max_turns Int // sample budget: judge at most this many turns created_at DateTime @default(now()) created_by String? ends_at DateTime - stopped_at DateTime? + keys LiteLLM_ShadowEvalJobKey[] - @@index([api_key_id]) + @@unique([id, direction]) @@index([created_at]) } +// One row per key a job shadows. Budget and stop state are per key, so one key exhausting +// its turns never ends a sibling's sampling. "One active job per (key, direction)" is a +// partial unique index on (api_key_id, direction) WHERE stopped_at IS NULL, expressed only +// in the migration because schema.prisma cannot state partial indexes; it is what makes a +// concurrent start on another pod race-safe rather than read-then-create. direction is +// carried here rather than read off the job so that index can exist, and the composite +// foreign key makes it provably the job's own direction. +model LiteLLM_ShadowEvalJobKey { + id String @id @default(cuid()) + job_id String + direction String + api_key_id String // hashed virtual key whose traffic is shadowed + max_turns Int // this key's sample budget: judge at most this many turns + created_at DateTime @default(now()) + stopped_at DateTime? + job LiteLLM_ShadowEvalJob @relation(fields: [job_id, direction], references: [id, direction], onDelete: Cascade) + + @@index([job_id]) + @@index([api_key_id]) +} + // One row per sampled pipeline: a blind verdict (real | shadow | tie) or an error. model LiteLLM_ShadowEvalAttempt { id String @id @default(cuid()) job_id String + api_key_id String // which of the job's scoped keys produced this turn request_id String // the judged real request outcome String // real | shadow | tie | error tier String? // router's tier for the prompt, when classified @@ -1490,7 +1511,7 @@ model LiteLLM_ShadowEvalAttempt { error String? created_at DateTime @default(now()) - @@index([job_id]) + @@index([job_id, api_key_id]) } // --------------------------------------------------------------------------- diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 99d5ab47f1a..9fc7936ed43 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -4,7 +4,9 @@ each against the job's other arm in a detached task (the auto-router for a forwa fixed baseline model for a reverse one), blind-judges real vs shadow, and appends one ``LiteLLM_ShadowEvalAttempt`` row (verdict or error) as the feature's only hot-path write. Counts, status, and spend derive from those rows at read time, so nothing can disagree -across pods or stop races; the hook reads active jobs through a short-TTL cache.""" +across pods or stop races; the hook reads active keys through a short-TTL cache. A job can +scope several keys, and each carries its own turn budget and stop state, so the unit this +module samples and budgets against is the key, not the job.""" import asyncio import hashlib @@ -449,14 +451,16 @@ class _JudgeVerdict: class ActiveShadowEvalJob(BaseModel): - """One active job as the sampling path needs it, validated straight off the untyped - job row: immutable config plus the attempt count as of the cache fill (the turn - budget's staleness is bounded by the cache TTL). Every way a row can be unsamplable - is a validation error here, so a bad row is skipped rather than sampled wrongly.""" + """One key's active sampling as the path needs it: immutable job config plus that key's + own budget and attempt count as of the cache fill (the budget's staleness is bounded by + the cache TTL). A job over N keys yields N of these, one per key, with independent + budgets, so exhausting one leaves its siblings sampling. Every way a row can be + unsamplable is a validation error here, so a bad row is skipped rather than sampled wrongly.""" - model_config = ConfigDict(frozen=True, from_attributes=True) + model_config = ConfigDict(frozen=True) id: str + api_key_id: str router_name: str direction: ShadowEvalDirection = "forward" baseline_model: str | None = None @@ -485,16 +489,15 @@ class ActiveShadowEvalJob(BaseModel): return self.baseline_model or self.router_name -def _as_active_job(record: object, attempts: int) -> ActiveShadowEvalJob | None: - """The sampling path's view of one job row, or None for a row it cannot sample: an - unknown direction, or a reverse job with no baseline model to duplicate against. - Failing closed here is what keeps the dispatch path total.""" +def _as_active_job(payload: Mapping[str, object]) -> ActiveShadowEvalJob | None: + """The sampling path's view of one key row and the job it belongs to, or None for a row + it cannot sample: an unknown direction, or a reverse job with no baseline model to + duplicate against. Failing closed here is what keeps the dispatch path total.""" try: - job: Final = ActiveShadowEvalJob.model_validate(record) + return ActiveShadowEvalJob.model_validate(payload) except ValidationError as e: - verbose_logger.debug("shadow_eval: skipping unsamplable job row: %s", e) + verbose_logger.debug("shadow_eval: skipping unsamplable job key row: %s", e) return None - return job.model_copy(update={"attempts": attempts}) _jobs_cache: Final = InMemoryCache(max_size_in_memory=4, default_ttl=_JOBS_CACHE_TTL_SECONDS) @@ -516,14 +519,16 @@ class ShadowEvalLogger(CustomLogger): self._prisma_provider = prisma_provider or _default_prisma_provider self._jobs_cache = jobs_cache or _jobs_cache self._inflight_shadow_tasks: int = 0 - # Starts per job since the last cache fill, never decremented within a - # generation; the refill absorbs written rows and resets. - self._job_starts: dict[str, int] = {} # mutable-ok: per-generation counter + # Starts per shadowed key row since the last cache fill, never decremented within a + # generation; the refill absorbs written rows and resets. Keyed by (job, key) rather + # than by either alone, since the turn budget it guards is one key's within one job. + self._job_starts: dict[tuple[str, str], int] = {} # mutable-ok: per-generation counter async def _active_jobs(self) -> Mapping[str, tuple[ActiveShadowEvalJob, ...]]: - """Active jobs by api_key_id, cache-first. A key holds at most one job per - direction, so the value is a collection. A DB fault returns empty without - caching, so sampling pauses for that request and the next one retries.""" + """Sampling key rows by api_key_id, cache-first. A key holds at most one unstopped + row per direction (partial unique index), so the value is a collection of at most + two. A DB fault returns empty without caching, so sampling pauses for that request + and the next one retries.""" cached: Final = await self._jobs_cache.async_get_cache(_JOBS_CACHE_KEY) if cached is not None: return cached # pyright: ignore[reportReturnType] # cache stores exactly this mapping shape @@ -531,32 +536,48 @@ class ShadowEvalLogger(CustomLogger): if prisma is None: return _EMPTY_JOBS try: - records: Final = await prisma.db.litellm_shadowevaljob.find_many( + records: Final = await prisma.db.litellm_shadowevaljobkey.find_many( where={ # mutable-ok: Prisma filter "stopped_at": None, - "ends_at": {"gt": datetime.now(timezone.utc)}, # mutable-ok: Prisma filter + "job": {"is": {"ends_at": {"gt": datetime.now(timezone.utc)}}}, # mutable-ok: Prisma filter }, + include={"job": True}, # mutable-ok: Prisma payload ) + pairs: Final = [{"job_id": r.job_id, "api_key_id": r.api_key_id} for r in records] # mutable-ok: Prisma OR grouped: Final = ( await prisma.db.litellm_shadowevalattempt.group_by( - by=["job_id"], + by=["job_id", "api_key_id"], count=True, - where={"job_id": {"in": [str(record.id) for record in records]}}, # mutable-ok: Prisma filter + where={"OR": pairs}, # mutable-ok: Prisma filter ) if records else () ) - attempt_counts: Final = {str(row["job_id"]): int(row["_count"]["_all"]) for row in grouped or []} - by_key: Final = tuple( - sorted( - ( - (str(record.api_key_id), job) - for record in records or [] - if (job := _as_active_job(record, attempt_counts.get(str(record.id), 0))) is not None - ), - key=itemgetter(0), - ) + attempt_counts: Final = MappingProxyType( + {(str(row["job_id"]), str(row["api_key_id"])): int(row["_count"]["_all"]) for row in grouped or ()} ) + candidates: Final = tuple( + ( + str(record.api_key_id), + _as_active_job( + { + "id": record.job_id, + "api_key_id": record.api_key_id, + "router_name": job.router_name, + "direction": job.direction, + "baseline_model": job.baseline_model, + "shadow_percentage": job.shadow_percentage, + "judge_model": job.judge_model, + "max_turns": record.max_turns, + "ends_at": job.ends_at, + "attempts": attempt_counts.get((str(record.job_id), str(record.api_key_id)), 0), + } + ), + ) + for record in records or () + if (job := record.job) is not None + ) + by_key: Final = tuple(sorted(((key, job) for key, job in candidates if job is not None), key=itemgetter(0))) jobs: Final = MappingProxyType( {key: tuple(job for _, job in group) for key, group in groupby(by_key, key=itemgetter(0))} ) @@ -608,7 +629,7 @@ class ShadowEvalLogger(CustomLogger): job for job in (await self._active_jobs()).get(str(api_key_hash), ()) if datetime.now(timezone.utc) < job.ends_at - and job.attempts + self._job_starts.get(job.id, 0) < job.max_turns + and job.attempts + self._job_starts.get((job.id, job.api_key_id), 0) < job.max_turns and _sample_hits(request_id, job.id, job.shadow_percentage) and _request_was_routed_by(request_metadata, job.router_name) == (job.direction == "reverse") ) @@ -627,7 +648,7 @@ class ShadowEvalLogger(CustomLogger): for job in eligible: if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS: return - self._job_starts[job.id] = self._job_starts.get(job.id, 0) + 1 + self._job_starts[(job.id, job.api_key_id)] = self._job_starts.get((job.id, job.api_key_id), 0) + 1 self._inflight_shadow_tasks += 1 asyncio.create_task( self._run_shadow_eval( @@ -731,6 +752,7 @@ class ShadowEvalLogger(CustomLogger): await prisma.db.litellm_shadowevalattempt.create( data={ # mutable-ok: Prisma payload "job_id": job.id, + "api_key_id": job.api_key_id, "request_id": request_id, "outcome": outcome, "tier": control_tier if job.direction == "reverse" else (shadow.tier if shadow else None), diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 4b2569fa9fa..65514de05be 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -528,12 +528,18 @@ GROUP BY 1 _ATTEMPT_AGG_BY_TIER_SQL: Final = "SELECT COALESCE(tier, 'UNCLASSIFIED') AS grp," + _ATTEMPT_AGG_SELECT _ATTEMPT_AGG_BY_MODEL_SQL: Final = "SELECT COALESCE(real_model, 'unknown') AS grp," + _ATTEMPT_AGG_SELECT -_SWEEP_FINISHED_JOBS_SQL: Final = """ -UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = NOW() -WHERE j.api_key_id = $1 AND j.stopped_at IS NULL +_ATTEMPT_AGG_BY_KEY_SQL: Final = "SELECT api_key_id AS grp," + _ATTEMPT_AGG_SELECT + +_SWEEP_FINISHED_KEYS_SQL: Final = """ +UPDATE "LiteLLM_ShadowEvalJobKey" k SET stopped_at = NOW() +FROM "LiteLLM_ShadowEvalJob" j +WHERE k.job_id = j.id AND k.api_key_id = ANY($1::text[]) AND k.stopped_at IS NULL AND ( j.ends_at <= NOW() - OR (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = j.id) >= j.max_turns + OR ( + SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a + WHERE a.job_id = k.job_id AND a.api_key_id = k.api_key_id + ) >= k.max_turns ) """ @@ -575,11 +581,12 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]: async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> ShadowEvalResult | None: - """Both stratifications of one job's verdicts. Tier answers "where does the router do - well"; the model stratification groups by whichever model served the real arm, so it - answers "which of the models this key uses today would the router beat" forward, and - "for the turns the router sent to X, did X beat the baseline" in reverse. Reads are - bounded by the job's own attempts (<= max_turns) via the job_id index.""" + """All three stratifications of one job's verdicts. Tier answers "where does the router + do well"; the model stratification groups by whichever model served the real arm, so it + answers "which of the models these keys use today would the router beat" forward, and + "for the turns the router sent to X, did X beat the baseline" in reverse; key answers + "which key's traffic does the router suit". Reads are bounded by the job's own attempts + (<= the sum of its keys' max_turns) via the job_id index.""" by_tier: Final = _ATTEMPT_AGG_ROWS.validate_python( await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_TIER_SQL, job_id) or () ) @@ -588,10 +595,14 @@ async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> Sh by_model: Final = _ATTEMPT_AGG_ROWS.validate_python( await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_MODEL_SQL, job_id) or () ) + by_key: Final = _ATTEMPT_AGG_ROWS.validate_python( + await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_KEY_SQL, job_id) or () + ) total_turns: Final = sum(r.turn_count for r in by_tier) return ShadowEvalResult( by_tier=_slices(by_tier), by_current_model=_slices(by_model), + by_key=_slices(by_key), overall_shadow_win_rate_pct=_pct_of(sum(r.shadow_wins for r in by_tier), total_turns), overall_tie_rate_pct=_pct_of(sum(r.ties for r in by_tier), total_turns), ) @@ -609,20 +620,21 @@ async def start_shadow_eval( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ) -> ShadowEvalJobResponse: """ - Start a shadow eval: duplicate a sampled slice of a key's live traffic against a second - arm, judge the two responses blind, and stratify win rates by tier and by the model that - served the real arm. + Start a shadow eval: duplicate a sampled slice of one or more keys' live traffic against + a second arm, judge the two responses blind, and stratify win rates by tier, by the model + that served the real arm, and by key. - A forward job answers whether the key should adopt router_name: it samples the requests + A forward job answers whether the keys should adopt router_name: it samples the requests the router did not serve and duplicates them through it. A reverse job answers whether a key already on the router still gains from it: it samples the requests the router did serve and duplicates them against baseline_model. A key can hold one active job per direction, so both questions can run at once. - Shadow responses are never served to users. The job samples until it has judged - max_turns turns, reaches the end of its window, or is stopped; sampling changes - propagate to pods within about 10 seconds. Shadow and judge calls bill to the - shadowed key but are excluded from request counts and auto-router adoption metrics. + Shadow responses are never served to users. Each key samples until it has judged + max_turns turns of its own traffic, the job's window ends, or the job is stopped, so a + busy key running out of budget does not end sampling for the others; sampling changes + propagate to pods within about 10 seconds. Shadow and judge calls bill to the shadowed + key but are excluded from request counts and auto-router adoption metrics. """ from litellm.proxy.proxy_server import llm_router, prisma_client @@ -634,48 +646,58 @@ async def start_shadow_eval( _validate_plain_model(llm_router, data.judge_model, "judge_model") if data.baseline_model is not None: _validate_plain_model(llm_router, data.baseline_model, "baseline_model") - key_row: Final = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": data.api_key_id} # mutable-ok: Prisma filter + token_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many( + where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter ) - if key_row is None: + unknown: Final = tuple(sorted(frozenset(data.api_key_ids) - frozenset(row.token for row in token_rows))) + if unknown: raise HTTPException( status_code=400, detail=( - f"api_key_id '{data.api_key_id}' is not a key on this proxy; pass the key's token hash, " + f"api_key_ids not on this proxy: {', '.join(unknown)}; pass each key's token hash, " "the value the key list and key info endpoints report" ), ) - # A job that expired or exhausted its turn budget stopped sampling on its own, but - # still holds its slot in the per-key, per-direction partial unique index until - # stamped; free it so a new eval can start. Sweeping both directions is deliberate. - await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, data.api_key_id) - active: Final = await prisma_client.db.litellm_shadowevaljob.find_first( + # A key whose job expired or whose own turn budget ran out stopped sampling on its own, but + # still holds its slot in the per-key, per-direction partial unique index until stamped; + # free it so a new eval can start. Sweeping both directions is deliberate. + await prisma_client.db.execute_raw(_SWEEP_FINISHED_KEYS_SQL, list(data.api_key_ids)) # mutable-ok: query param + claimed: Final = await prisma_client.db.litellm_shadowevaljobkey.find_many( where={ # mutable-ok: Prisma filter - "api_key_id": data.api_key_id, + "api_key_id": {"in": list(data.api_key_ids)}, "direction": data.direction, "stopped_at": None, }, ) - if active is not None: + if claimed: raise HTTPException( status_code=409, - detail=f"Key already has an active {data.direction} shadow eval job ({active.id}). Stop it first.", + detail=( + f"Already in an active {data.direction} shadow eval job: " + + ", ".join(sorted(f"{row.api_key_id} (job {row.job_id})" for row in claimed)) + + ". Stop it first." + ), ) now: Final = datetime.now(timezone.utc) + key_rows: Final = tuple( + {"api_key_id": api_key_id, "max_turns": data.max_turns} # mutable-ok: Prisma create payload + for api_key_id in data.api_key_ids + ) try: job: Final = await prisma_client.db.litellm_shadowevaljob.create( data={ # mutable-ok: Prisma payload - "api_key_id": data.api_key_id, "router_name": data.router_name, "direction": data.direction, "baseline_model": data.baseline_model, "judge_model": data.judge_model, "shadow_percentage": data.shadow_percentage, - "max_turns": data.max_turns, "created_by": user_api_key_dict.user_id, "ends_at": now + timedelta(days=data.duration_days), - } + # direction is injected into each key row by Prisma from the composite FK + "keys": {"create": list(key_rows)}, # mutable-ok: Prisma nested create + }, + include={"keys": True}, # mutable-ok: Prisma payload ) except Exception as e: if not _is_unique_violation(e): @@ -683,7 +705,7 @@ async def start_shadow_eval( raise HTTPException( status_code=409, detail=( - f"Key already has an active {data.direction} shadow eval job (started concurrently). Stop it first." + f"A requested key was claimed by another {data.direction} shadow eval job concurrently. Stop it first." ), ) from e return ShadowEvalJobResponse.model_validate(job, from_attributes=True) @@ -697,7 +719,9 @@ async def start_shadow_eval( ) async def list_shadow_eval_jobs( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], - api_key_id: Annotated[str | None, Query(description="Filter to jobs shadowing this key")] = None, + api_key_id: Annotated[ + str | None, Query(description="Filter to jobs that shadow this key, alone or alongside others") + ] = None, limit: Annotated[int, Query(ge=1, le=200, description="Newest jobs to return")] = 50, ) -> tuple[ShadowEvalJobResponse, ...]: """List shadow eval jobs, newest first. Counts and results ride the detail endpoint only.""" @@ -707,7 +731,8 @@ async def list_shadow_eval_jobs( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) records: Final = await prisma_client.db.litellm_shadowevaljob.find_many( - where={"api_key_id": api_key_id} if api_key_id else {}, # mutable-ok: Prisma filter + where={"keys": {"some": {"api_key_id": api_key_id}}} if api_key_id else {}, # mutable-ok: Prisma filter + include={"keys": True}, # mutable-ok: Prisma payload order={"created_at": "desc"}, # mutable-ok: Prisma order take=limit, ) @@ -731,7 +756,8 @@ async def get_shadow_eval_job( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique( - where={"id": job_id} # mutable-ok: Prisma filter + where={"id": job_id}, # mutable-ok: Prisma filter + include={"keys": True}, # mutable-ok: Prisma payload ) if record is None: raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}") @@ -763,22 +789,31 @@ async def stop_shadow_eval_job( job_id: str, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ) -> ShadowEvalJobResponse: - """Stop an active shadow eval job. Attempts are kept; sampling halts within ~10s.""" + """Stop an active shadow eval job, every key it scopes at once. Attempts are kept; + sampling halts within ~10s. Keys that already stopped on their own budget keep the + stopped_at they earned, so one write records the outcome for the whole job.""" from litellm.proxy.proxy_server import prisma_client _require_admin_writer(user_api_key_dict, "stop a shadow eval") if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique( - where={"id": job_id} # mutable-ok: Prisma filter + where={"id": job_id}, # mutable-ok: Prisma filter + include={"keys": True}, # mutable-ok: Prisma payload ) if record is None: raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}") current: Final = ShadowEvalJobResponse.model_validate(record, from_attributes=True) if current.status != "running": raise HTTPException(status_code=400, detail=f"Job {job_id} is already {current.status}") - updated: Final = await prisma_client.db.litellm_shadowevaljob.update( - where={"id": job_id}, # mutable-ok: Prisma filter + await prisma_client.db.litellm_shadowevaljobkey.update_many( + where={"job_id": job_id, "stopped_at": None}, # mutable-ok: Prisma filter data={"stopped_at": datetime.now(timezone.utc)}, # mutable-ok: Prisma payload ) + updated: Final = await prisma_client.db.litellm_shadowevaljob.find_unique( + where={"id": job_id}, # mutable-ok: Prisma filter + include={"keys": True}, # mutable-ok: Prisma payload + ) + if updated is None: + raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}") return ShadowEvalJobResponse.model_validate(updated, from_attributes=True) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 71345d2ccde..a035e8ec623 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1450,36 +1450,57 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } -// Shadow eval: evaluation of an auto-router against a key's live traffic, in either -// direction. forward duplicates the requests the key did not route through the router -// through it, answering whether the key should adopt it; reverse duplicates the requests +// Shadow eval: evaluation of an auto-router against a set of keys' live traffic, in either +// direction. forward duplicates the requests the keys did not route through the router +// through it, answering whether they should adopt it; reverse duplicates the requests // the router did serve against a fixed baseline model, answering whether a key already on // it still benefits. Either way a sampled slice runs in a detached task and an LLM judge -// compares real vs shadow responses blind. The job row is immutable config plus -// stopped_at; every count, status, and spend figure is derived from the append-only -// attempt rows, so nothing can disagree across pods or stop races. +// compares real vs shadow responses blind. The job row is immutable config; the keys it +// scopes, their budgets and their stop state live on LiteLLM_ShadowEvalJobKey, and every +// count, status, and spend figure is derived from the append-only attempt rows, so nothing +// can disagree across pods or stop races. model LiteLLM_ShadowEvalJob { id String @id @default(cuid()) - api_key_id String // hashed virtual key whose traffic is shadowed router_name String // the auto-router under evaluation, in either direction direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String shadow_percentage Float - max_turns Int // sample budget: judge at most this many turns created_at DateTime @default(now()) created_by String? ends_at DateTime - stopped_at DateTime? + keys LiteLLM_ShadowEvalJobKey[] - @@index([api_key_id]) + @@unique([id, direction]) @@index([created_at]) } +// One row per key a job shadows. Budget and stop state are per key, so one key exhausting +// its turns never ends a sibling's sampling. "One active job per (key, direction)" is a +// partial unique index on (api_key_id, direction) WHERE stopped_at IS NULL, expressed only +// in the migration because schema.prisma cannot state partial indexes; it is what makes a +// concurrent start on another pod race-safe rather than read-then-create. direction is +// carried here rather than read off the job so that index can exist, and the composite +// foreign key makes it provably the job's own direction. +model LiteLLM_ShadowEvalJobKey { + id String @id @default(cuid()) + job_id String + direction String + api_key_id String // hashed virtual key whose traffic is shadowed + max_turns Int // this key's sample budget: judge at most this many turns + created_at DateTime @default(now()) + stopped_at DateTime? + job LiteLLM_ShadowEvalJob @relation(fields: [job_id, direction], references: [id, direction], onDelete: Cascade) + + @@index([job_id]) + @@index([api_key_id]) +} + // One row per sampled pipeline: a blind verdict (real | shadow | tie) or an error. model LiteLLM_ShadowEvalAttempt { id String @id @default(cuid()) job_id String + api_key_id String // which of the job's scoped keys produced this turn request_id String // the judged real request outcome String // real | shadow | tie | error tier String? // router's tier for the prompt, when classified @@ -1490,7 +1511,7 @@ model LiteLLM_ShadowEvalAttempt { error String? created_at DateTime @default(now()) - @@index([job_id]) + @@index([job_id, api_key_id]) } // --------------------------------------------------------------------------- diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 1b0c7476fc3..2983b670d81 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -152,13 +152,15 @@ DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5" class StartShadowEvalRequest(BaseModel): - """Start duplicating a key's traffic for blind comparison against an auto-router.""" + """Start duplicating one or more keys' traffic for blind comparison against an auto-router.""" - api_key_id: str = Field( + api_key_ids: tuple[str, ...] = Field( + min_length=1, description=( - "The hashed virtual key whose traffic will be shadowed. Shadow evaluation runs ONLY on this " - "key's traffic; requests made with any other key are not sampled." - ) + "The hashed virtual keys whose traffic will be shadowed. Shadow evaluation runs ONLY on these " + "keys' traffic; requests made with any other key are not sampled. Each key gets its own " + "max_turns budget, so one key exhausting its budget does not end sampling for the others." + ), ) router_name: str = Field(description="The auto-router under evaluation, in either direction") direction: ShadowEvalDirection = Field( @@ -201,8 +203,9 @@ class StartShadowEvalRequest(BaseModel): ge=1, le=2000, description=( - "Sample budget: the job judges at most this many turns, then completes. This is also the spend " - "bound; expected judge cost is roughly max_turns times one judge call" + "Per-key sample budget: the job judges at most this many turns of EACH scoped key's traffic, so a " + "job over N keys judges at most N times max_turns turns. This is also the spend bound; expected " + "judge cost is roughly that turn ceiling times one judge call" ), ) @@ -211,6 +214,13 @@ class StartShadowEvalRequest(BaseModel): def _round_percentage(cls, value: float) -> float: return round(value, 2) + @field_validator("api_key_ids") + @classmethod + def _dedupe_keys(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Naming a key twice means the same thing as naming it once, and the second key row + would collide with the first on the one-active-per-(key, direction) index.""" + return tuple(dict.fromkeys(value)) + @model_validator(mode="after") def _baseline_model_matches_direction(self) -> "StartShadowEvalRequest": if self.direction == "reverse" and self.baseline_model is None: @@ -221,8 +231,8 @@ class StartShadowEvalRequest(BaseModel): class ShadowEvalSlice(BaseModel): - """Judge outcomes for one slice of a job's verdicts (a router tier, or one of the - models that served the real arm).""" + """Judge outcomes for one slice of a job's verdicts (a router tier, one of the models + that served the real arm, or one of the keys the job is scoped to).""" group: str turn_count: int @@ -248,33 +258,56 @@ class ShadowEvalResult(BaseModel): by_tier: tuple[ShadowEvalSlice, ...] by_current_model: tuple[ShadowEvalSlice, ...] = Field( description=( - "Sliced by the model that served the real arm: the key's incumbent models in forward mode, " + "Sliced by the model that served the real arm: the keys' incumbent models in forward mode, " "and in reverse the models the router itself picked" ) ) + by_key: tuple[ShadowEvalSlice, ...] = Field( + default=(), + description=( + "One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job " + "scopes but has not yet judged a turn for are absent rather than reported as zero" + ), + ) overall_shadow_win_rate_pct: float overall_tie_rate_pct: float +class ShadowEvalJobKeyResponse(BaseModel): + """One key a job shadows, with its own budget and stop state.""" + + api_key_id: str = Field(description="The hashed virtual key whose traffic this row scopes") + max_turns: int = Field(description="This key's own sample budget, independent of its siblings'") + stopped_at: datetime | None = Field( + default=None, + description=( + "When this key stopped sampling, whether its own budget ran out, the window closed, or an " + "operator stopped the job. The key reads completed whenever the job does, otherwise stopped " + "once this is set and running until then" + ), + ) + + class ShadowEvalJobResponse(BaseModel): """A shadow-eval job. Validates directly from the prisma record (job_id reads the - row's id); status is derived from stopped_at and ends_at, never stored, so no writer - anywhere can produce an inconsistent one. Aggregate fields are populated by the + row's id); status is derived from the keys' stopped_at and ends_at, never stored, so no + writer anywhere can produce an inconsistent one. Aggregate fields are populated by the detail endpoint only and stay None on list responses.""" model_config = ConfigDict(from_attributes=True, populate_by_name=True) job_id: str = Field(validation_alias=AliasChoices("id", "job_id")) - api_key_id: str = Field(description="The hashed virtual key whose traffic this job evaluates, and only that key's") + keys: tuple[ShadowEvalJobKeyResponse, ...] = Field( + min_length=1, + description="The keys whose traffic this job evaluates, and only those keys', each with its own budget", + ) router_name: str direction: ShadowEvalDirection = "forward" baseline_model: str | None = None judge_model: str shadow_percentage: float - max_turns: int created_at: datetime ends_at: datetime - stopped_at: datetime | None = None judged_count: int | None = Field(default=None, description="Verdicts recorded; detail endpoint only") error_count: int | None = Field(default=None, description="Sampled attempts that errored; detail endpoint only") @@ -285,12 +318,13 @@ class ShadowEvalJobResponse(BaseModel): @computed_field @property def status(self) -> ShadowEvalStatus: - """A job whose window has passed reads completed even if a later sweep stamped - stopped_at; stopped means sampling ended before the window did.""" + """A job whose window has passed reads completed even if a sweep stamped its keys + first; stopped means every key ended sampling before the window did. One key + exhausting its own budget leaves the job running while any sibling still samples.""" if datetime.now(timezone.utc) >= ( self.ends_at if self.ends_at.tzinfo else self.ends_at.replace(tzinfo=timezone.utc) ): return "completed" - if self.stopped_at is not None: + if all(key.stopped_at is not None for key in self.keys): return "stopped" return "running" diff --git a/schema.prisma b/schema.prisma index 71345d2ccde..a035e8ec623 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1450,36 +1450,57 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } -// Shadow eval: evaluation of an auto-router against a key's live traffic, in either -// direction. forward duplicates the requests the key did not route through the router -// through it, answering whether the key should adopt it; reverse duplicates the requests +// Shadow eval: evaluation of an auto-router against a set of keys' live traffic, in either +// direction. forward duplicates the requests the keys did not route through the router +// through it, answering whether they should adopt it; reverse duplicates the requests // the router did serve against a fixed baseline model, answering whether a key already on // it still benefits. Either way a sampled slice runs in a detached task and an LLM judge -// compares real vs shadow responses blind. The job row is immutable config plus -// stopped_at; every count, status, and spend figure is derived from the append-only -// attempt rows, so nothing can disagree across pods or stop races. +// compares real vs shadow responses blind. The job row is immutable config; the keys it +// scopes, their budgets and their stop state live on LiteLLM_ShadowEvalJobKey, and every +// count, status, and spend figure is derived from the append-only attempt rows, so nothing +// can disagree across pods or stop races. model LiteLLM_ShadowEvalJob { id String @id @default(cuid()) - api_key_id String // hashed virtual key whose traffic is shadowed router_name String // the auto-router under evaluation, in either direction direction String @default("forward") // forward | reverse baseline_model String? // reverse only: the fixed model the router is judged against judge_model String shadow_percentage Float - max_turns Int // sample budget: judge at most this many turns created_at DateTime @default(now()) created_by String? ends_at DateTime - stopped_at DateTime? + keys LiteLLM_ShadowEvalJobKey[] - @@index([api_key_id]) + @@unique([id, direction]) @@index([created_at]) } +// One row per key a job shadows. Budget and stop state are per key, so one key exhausting +// its turns never ends a sibling's sampling. "One active job per (key, direction)" is a +// partial unique index on (api_key_id, direction) WHERE stopped_at IS NULL, expressed only +// in the migration because schema.prisma cannot state partial indexes; it is what makes a +// concurrent start on another pod race-safe rather than read-then-create. direction is +// carried here rather than read off the job so that index can exist, and the composite +// foreign key makes it provably the job's own direction. +model LiteLLM_ShadowEvalJobKey { + id String @id @default(cuid()) + job_id String + direction String + api_key_id String // hashed virtual key whose traffic is shadowed + max_turns Int // this key's sample budget: judge at most this many turns + created_at DateTime @default(now()) + stopped_at DateTime? + job LiteLLM_ShadowEvalJob @relation(fields: [job_id, direction], references: [id, direction], onDelete: Cascade) + + @@index([job_id]) + @@index([api_key_id]) +} + // One row per sampled pipeline: a blind verdict (real | shadow | tie) or an error. model LiteLLM_ShadowEvalAttempt { id String @id @default(cuid()) job_id String + api_key_id String // which of the job's scoped keys produced this turn request_id String // the judged real request outcome String // real | shadow | tie | error tier String? // router's tier for the prompt, when classified @@ -1490,7 +1511,7 @@ model LiteLLM_ShadowEvalAttempt { error String? created_at DateTime @default(now()) - @@index([job_id]) + @@index([job_id, api_key_id]) } // --------------------------------------------------------------------------- diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 6f4979b9b84..f9de6abab0c 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -230,6 +230,7 @@ class TestMigrationSQLIdempotency: "20250918083359_drop_spec_version_column_from_mcp_table", "20260213170952_access_group_change_to_model_name", "20260224203854_add_agent_object_permissions_table", + "20260814000000_shadow_eval_multi_key", } def test_no_drop_column_statements(self, all_migrations): diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 8d2f9482fa7..86e77d5bb1e 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -27,6 +27,7 @@ from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTE def _job(**overrides) -> ActiveShadowEvalJob: defaults = dict( id="job-1", + api_key_id="key-hash", router_name="my-router", shadow_percentage=100.0, judge_model="judge-model", @@ -37,30 +38,34 @@ def _job(**overrides) -> ActiveShadowEvalJob: return ActiveShadowEvalJob(**{**defaults, **overrides}) -def _prisma(jobs=(), attempt_counts=()) -> MagicMock: +def _prisma(keys=(), attempt_counts=()) -> MagicMock: prisma = MagicMock() - prisma.db.litellm_shadowevaljob.find_many = AsyncMock(return_value=list(jobs)) + prisma.db.litellm_shadowevaljobkey.find_many = AsyncMock(return_value=list(keys)) prisma.db.litellm_shadowevalattempt.group_by = AsyncMock( - return_value=[{"job_id": job_id, "_count": {"_all": count}} for job_id, count in attempt_counts] + return_value=[ + {"job_id": job_id, "api_key_id": api_key_id, "_count": {"_all": count}} + for job_id, api_key_id, count in attempt_counts + ] ) prisma.db.litellm_shadowevalattempt.create = AsyncMock() return prisma -def _job_record(job: ActiveShadowEvalJob, api_key_id="key-hash") -> MagicMock: +def _key_record(job: ActiveShadowEvalJob) -> MagicMock: + """One LiteLLM_ShadowEvalJobKey row with its job relation loaded, as `include={"job": True}` + hands it back: per-key budget on the row, shared config on the relation.""" record = MagicMock() + for field, value in dict(job_id=job.id, api_key_id=job.api_key_id, max_turns=job.max_turns).items(): + setattr(record, field, value) for field, value in dict( - id=job.id, - api_key_id=api_key_id, router_name=job.router_name, direction=job.direction, baseline_model=job.baseline_model, shadow_percentage=job.shadow_percentage, judge_model=job.judge_model, - max_turns=job.max_turns, ends_at=job.ends_at, ).items(): - setattr(record, field, value) + setattr(record.job, field, value) return record @@ -97,7 +102,13 @@ def _logger(router=None, prisma=None, jobs=()) -> ShadowEvalLogger: jobs_cache=cache, ) if jobs: - cache.set_cache("shadow_eval:active_jobs", {"key-hash": tuple(jobs)}) + cache.set_cache( + "shadow_eval:active_jobs", + { + key: tuple(job for job in jobs if job.api_key_id == key) + for key in dict.fromkeys(job.api_key_id for job in jobs) + }, + ) return logger @@ -466,6 +477,7 @@ class TestSuccessHookSkipChain: create.assert_awaited_once() row = create.call_args.kwargs["data"] assert row["job_id"] == "job-1" + assert row["api_key_id"] == "key-hash" assert row["request_id"] == "req-1" assert row["outcome"] in ("real", "shadow") assert row["tier"] == "SIMPLE" @@ -474,7 +486,7 @@ class TestSuccessHookSkipChain: assert row["confidence"] == 0.9 assert row["judge_cost"] == 0.005 assert row["error"] is None - assert prisma.db.litellm_shadowevaljob.find_many.await_count == 0 + assert prisma.db.litellm_shadowevaljobkey.find_many.await_count == 0 @pytest.mark.parametrize( "kwargs_mutation,job_mutation", @@ -503,13 +515,13 @@ class TestSuccessHookSkipChain: starts = job_mutation.pop("_starts", 0) prisma = _prisma() logger = _logger(router=_router(), prisma=prisma, jobs=(_job(**job_mutation),)) - logger._job_starts = {"job-1": starts} + logger._job_starts = {("job-1", "key-hash"): starts} await logger.async_log_success_event(_success_kwargs(**kwargs_mutation), RESPONSE, None, None) await _drain(logger) prisma.db.litellm_shadowevalattempt.create.assert_not_called() - assert logger._job_starts.get("job-1", 0) == starts + assert logger._job_starts.get(("job-1", "key-hash"), 0) == starts async def test_completed_pipelines_hold_turn_budget_within_a_cache_generation(self): """A finished pipeline frees its concurrency slot but not its slice of the turn @@ -524,6 +536,50 @@ class TestSuccessHookSkipChain: assert prisma.db.litellm_shadowevalattempt.create.await_count == 1 + async def test_one_keys_exhausted_budget_leaves_its_siblings_sampling(self): + """The turn budget belongs to the key, not to the job, so a job scoping two keys + keeps shadowing the second after the first has spent its own max_turns.""" + prisma = _prisma() + logger = _logger( + router=_router(), + prisma=prisma, + jobs=( + _job(api_key_id="spent-key", attempts=200, max_turns=200), + _job(api_key_id="fresh-key", attempts=0, max_turns=200), + ), + ) + + await logger.async_log_success_event(_success_kwargs(api_key_hash="spent-key"), RESPONSE, None, None) + await _drain(logger) + await logger.async_log_success_event(_success_kwargs(api_key_hash="fresh-key"), RESPONSE, None, None) + await _drain(logger) + + create = prisma.db.litellm_shadowevalattempt.create + create.assert_awaited_once() + assert create.call_args.kwargs["data"]["api_key_id"] == "fresh-key" + + async def test_started_turns_are_held_against_the_starting_key_only(self): + """The in-generation starts counter is keyed by key too. Spending the sibling's + last turn must not close the budget of a key that has one left.""" + prisma = _prisma() + logger = _logger( + router=_router(), + prisma=prisma, + jobs=( + _job(api_key_id="key-a", attempts=199, max_turns=200), + _job(api_key_id="key-b", attempts=199, max_turns=200), + ), + ) + + for api_key_hash in ("key-a", "key-a", "key-b"): + await logger.async_log_success_event( + _success_kwargs(request_id=f"req-{api_key_hash}", api_key_hash=api_key_hash), RESPONSE, None, None + ) + await _drain(logger) + + attributed = [call.kwargs["data"]["api_key_id"] for call in prisma.db.litellm_shadowevalattempt.create.call_args_list] + assert attributed == ["key-a", "key-b"] + async def test_v1_messages_surface_forwards_identity_from_litellm_metadata(self): """/v1/messages stores identity in litellm_params.litellm_metadata, so the hook resolves the bucket through the shared helper; every surface forwards the same @@ -574,7 +630,7 @@ class TestSuccessHookSkipChain: class TestActiveJobsCache: async def test_cache_miss_reads_db_once_then_serves_from_cache(self): job = _job() - prisma = _prisma(jobs=[_job_record(job)], attempt_counts=[("job-1", 7)]) + prisma = _prisma(keys=[_key_record(job)], attempt_counts=[("job-1", "key-hash", 7)]) logger = ShadowEvalLogger( router_provider=lambda: None, prisma_provider=lambda: prisma, @@ -586,15 +642,58 @@ class TestActiveJobsCache: assert [job.id for job in first["key-hash"]] == ["job-1"] assert second["key-hash"][0].attempts == 7 - assert prisma.db.litellm_shadowevaljob.find_many.await_count == 1 - where = prisma.db.litellm_shadowevaljob.find_many.call_args.kwargs["where"] + assert prisma.db.litellm_shadowevaljobkey.find_many.await_count == 1 + where = prisma.db.litellm_shadowevaljobkey.find_many.call_args.kwargs["where"] assert where["stopped_at"] is None - assert "gt" in where["ends_at"] + assert "gt" in where["job"]["is"]["ends_at"] count_where = prisma.db.litellm_shadowevalattempt.group_by.call_args.kwargs["where"] - assert count_where == {"job_id": {"in": ["job-1"]}} + assert count_where == {"OR": [{"job_id": "job-1", "api_key_id": "key-hash"}]} + + async def test_the_attempt_count_filter_matches_exact_pairs_not_a_cross_product(self): + """Two active keys under two jobs are two pairs. Filtering job_id IN (...) AND + api_key_id IN (...) also selects the pairs that never ran together, so every + other job's rows for the same key get scanned and grouped for nothing.""" + alpha = _key_record(_job(id="job-a", api_key_id="key-a")) + beta = _key_record(_job(id="job-b", api_key_id="key-b")) + prisma = _prisma(keys=[alpha, beta]) + logger = ShadowEvalLogger( + router_provider=lambda: None, + prisma_provider=lambda: prisma, + jobs_cache=InMemoryCache(max_size_in_memory=4, default_ttl=60), + ) + + await logger._active_jobs() + + count_where = prisma.db.litellm_shadowevalattempt.group_by.call_args.kwargs["where"] + assert count_where == { + "OR": [{"job_id": "job-a", "api_key_id": "key-a"}, {"job_id": "job-b", "api_key_id": "key-b"}] + } + + async def test_one_job_fills_one_entry_per_key_with_that_keys_own_count(self): + """Two keys of one job carry independent budgets, so the fill must group attempts + by key as well as by job. Grouping by job alone hands both keys the same count.""" + spent = _job(api_key_id="spent-key", max_turns=200) + fresh = _job(api_key_id="fresh-key", max_turns=50) + prisma = _prisma( + keys=[_key_record(spent), _key_record(fresh)], + attempt_counts=[("job-1", "spent-key", 200), ("job-1", "fresh-key", 3)], + ) + logger = ShadowEvalLogger( + router_provider=lambda: None, + prisma_provider=lambda: prisma, + jobs_cache=InMemoryCache(max_size_in_memory=4, default_ttl=60), + ) + + jobs = await logger._active_jobs() + + assert set(jobs) == {"spent-key", "fresh-key"} + assert jobs["spent-key"][0].attempts == 200 and jobs["spent-key"][0].max_turns == 200 + assert jobs["fresh-key"][0].attempts == 3 and jobs["fresh-key"][0].max_turns == 50 + assert jobs["fresh-key"][0].id == "job-1" + assert prisma.db.litellm_shadowevalattempt.group_by.call_args.kwargs["by"] == ["job_id", "api_key_id"] async def test_no_active_jobs_is_cached_too(self): - prisma = _prisma(jobs=[]) + prisma = _prisma(keys=[]) logger = ShadowEvalLogger( router_provider=lambda: None, prisma_provider=lambda: prisma, @@ -603,12 +702,12 @@ class TestActiveJobsCache: assert await logger._active_jobs() == {} assert await logger._active_jobs() == {} - assert prisma.db.litellm_shadowevaljob.find_many.await_count == 1 + assert prisma.db.litellm_shadowevaljobkey.find_many.await_count == 1 prisma.db.litellm_shadowevalattempt.group_by.assert_not_called() async def test_db_fault_returns_empty_without_caching_the_fault(self): prisma = _prisma() - prisma.db.litellm_shadowevaljob.find_many = AsyncMock(side_effect=RuntimeError("db blip")) + prisma.db.litellm_shadowevaljobkey.find_many = AsyncMock(side_effect=RuntimeError("db blip")) logger = ShadowEvalLogger( router_provider=lambda: None, prisma_provider=lambda: prisma, @@ -617,17 +716,17 @@ class TestActiveJobsCache: assert await logger._active_jobs() == {} assert await logger._active_jobs() == {} - assert prisma.db.litellm_shadowevaljob.find_many.await_count == 2 + assert prisma.db.litellm_shadowevaljobkey.find_many.await_count == 2 async def test_cache_refill_resets_the_starts_counter(self): job = _job() - prisma = _prisma(jobs=[_job_record(job)], attempt_counts=[("job-1", 7)]) + prisma = _prisma(keys=[_key_record(job)], attempt_counts=[("job-1", "key-hash", 7)]) logger = ShadowEvalLogger( router_provider=lambda: None, prisma_provider=lambda: prisma, jobs_cache=InMemoryCache(max_size_in_memory=4, default_ttl=60), ) - logger._job_starts = {"job-1": 5} + logger._job_starts = {("job-1", "key-hash"): 5} await logger._active_jobs() @@ -872,7 +971,7 @@ class TestDirection: rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.call_args_list] assert sorted(row["job_id"] for row in rows) == ["forward-job", "reverse-job"] - assert logger._job_starts == {"forward-job": 1, "reverse-job": 1} + assert logger._job_starts == {("forward-job", "key-hash"): 1, ("reverse-job", "key-hash"): 1} @pytest.mark.asyncio @@ -880,10 +979,12 @@ class TestActiveJobsFailClosed: async def test_a_row_the_sampler_cannot_read_is_dropped_not_guessed(self): """A reverse row with no baseline model has no second arm to call, so it is skipped rather than silently dispatched at the router it is supposed to be judging.""" - broken = _job_record(_job(id="job-broken")) - broken.direction = "reverse" - broken.baseline_model = None - prisma = _prisma(jobs=[broken, _job_record(_job(id="job-ok"))], attempt_counts=[("job-ok", 1)]) + broken = _key_record(_job(id="job-broken")) + broken.job.direction = "reverse" + broken.job.baseline_model = None + prisma = _prisma( + keys=[broken, _key_record(_job(id="job-ok"))], attempt_counts=[("job-ok", "key-hash", 1)] + ) logger = ShadowEvalLogger( router_provider=lambda: None, prisma_provider=lambda: prisma, @@ -894,11 +995,11 @@ class TestActiveJobsFailClosed: async def test_both_of_a_key_s_jobs_survive_the_lookup(self): records = [ - _job_record(_job(id="job-forward")), - _job_record(_reverse_job(id="job-reverse")), - _job_record(_job(id="job-other"), api_key_id="other-key"), + _key_record(_job(id="job-forward")), + _key_record(_reverse_job(id="job-reverse")), + _key_record(_job(id="job-other", api_key_id="other-key")), ] - prisma = _prisma(jobs=records, attempt_counts=[("job-reverse", 3)]) + prisma = _prisma(keys=records, attempt_counts=[("job-reverse", "key-hash", 3)]) logger = ShadowEvalLogger( router_provider=lambda: None, prisma_provider=lambda: prisma, 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 dbde7c461b8..f518ba5fdb5 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 @@ -318,9 +318,7 @@ class TestAutoRouterBenchmarks: from litellm.proxy.management_endpoints.auto_router_endpoints import _benchmark_totals totals = _benchmark_totals(self.ROW) - bucket_hits = ( - totals.cache.same_model.hits + totals.cache.first_visit.hits + totals.cache.return_to_tier.hits - ) + bucket_hits = totals.cache.same_model.hits + totals.cache.first_visit.hits + totals.cache.return_to_tier.hits assert bucket_hits == 27 assert totals.cache.hit_rate_pct == pytest.approx(100.0 * 28 / 38, abs=0.1) @@ -500,43 +498,71 @@ def _shadow_router() -> MagicMock: return router -def _job_record(**overrides: object) -> MagicMock: +def _spec_record(fields: dict) -> MagicMock: """Spec'd like a real prisma row: only the table's columns exist as attributes, so from_attributes validation falls back to model defaults for everything else.""" - defaults = { - "id": "job-1", - "api_key_id": "key-hash", - "router_name": "my-router", - "judge_model": "anthropic/claude-sonnet-5", - "shadow_percentage": 10.0, - "max_turns": 200, - "created_at": datetime(2026, 8, 11, tzinfo=timezone.utc), - "ends_at": datetime.now(timezone.utc) + timedelta(days=7), - "stopped_at": None, - } - fields = {**defaults, **overrides} record = MagicMock(spec=list(fields)) for key, value in fields.items(): setattr(record, key, value) return record -def _shadow_prisma(active_job=None, agg_rows=None) -> MagicMock: +def _key_row(api_key_id: str = "key-hash", max_turns: int = 200, stopped_at=None) -> MagicMock: + return _spec_record({"api_key_id": api_key_id, "max_turns": max_turns, "stopped_at": stopped_at}) + + +def _claim(api_key_id: str, job_id: str = "job-9", direction: str = "forward", stopped_at=None) -> MagicMock: + """A key row a prior job holds, still claiming the one-active-per-key-and-direction + index unless stopped.""" + return _spec_record( + {"api_key_id": api_key_id, "job_id": job_id, "direction": direction, "stopped_at": stopped_at} + ) + + +def _job_record(**overrides: object) -> MagicMock: + defaults = { + "id": "job-1", + "router_name": "my-router", + "direction": "forward", + "baseline_model": None, + "judge_model": "anthropic/claude-sonnet-5", + "shadow_percentage": 10.0, + "created_at": datetime(2026, 8, 11, tzinfo=timezone.utc), + "ends_at": datetime.now(timezone.utc) + timedelta(days=7), + "keys": (_key_row(),), + } + return _spec_record({**defaults, **overrides}) + + +def _shadow_prisma(claimed=(), agg_rows=None, by_key_rows=None, known_keys=("key-hash", "key-hash-2")) -> MagicMock: prisma = MagicMock() - prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=MagicMock()) + prisma.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[MagicMock(token=token) for token in known_keys] + ) prisma.db.execute_raw = AsyncMock(return_value=0) - prisma.db.litellm_shadowevaljob.find_first = AsyncMock(return_value=active_job) + + async def find_claimed_keys(where: dict, **_: object): + """Honours the filter it is handed, so a claim read that forgets stopped_at reads the + stopped rows a real partial index would have released, and one that forgets direction + reads the opposite-direction rows a key is entitled to hold at the same time.""" + scoped = [row for row in claimed if row.api_key_id in where["api_key_id"]["in"]] + by_direction = [row for row in scoped if row.direction == where.get("direction", row.direction)] + if "stopped_at" not in where: + return by_direction + return [row for row in by_direction if row.stopped_at is where["stopped_at"]] + + prisma.db.litellm_shadowevaljobkey.find_many = AsyncMock(side_effect=find_claimed_keys) + prisma.db.litellm_shadowevaljobkey.update_many = AsyncMock(return_value=1) prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=None) prisma.db.litellm_shadowevaljob.find_many = AsyncMock(return_value=[]) prisma.db.litellm_shadowevaljob.create = AsyncMock(return_value=_job_record()) - prisma.db.litellm_shadowevaljob.update = AsyncMock( - return_value=_job_record(stopped_at=datetime.now(timezone.utc)) - ) prisma.db.litellm_shadowevalattempt.find_first = AsyncMock(return_value=None) async def query_raw(sql: str, *params: object): if "FILTER (WHERE outcome != 'error')::int AS judged_count" in sql: return [{"judged_count": 10, "error_count": 2, "judge_spend": 0.031}] + if "api_key_id AS grp" in sql: + return by_key_rows if by_key_rows is not None else [] return agg_rows if agg_rows is not None else [] prisma.db.query_raw = AsyncMock(side_effect=query_raw) @@ -545,7 +571,7 @@ def _shadow_prisma(active_job=None, agg_rows=None) -> MagicMock: def _start_request(**overrides: object) -> StartShadowEvalRequest: payload = { - "api_key_id": "key-hash", + "api_key_ids": ("key-hash",), "router_name": "my-router", "shadow_percentage": 10.0, "judge_model": "anthropic/claude-sonnet-5", @@ -559,27 +585,33 @@ def _start_request(**overrides: object) -> StartShadowEvalRequest: @pytest.mark.asyncio async def test_start_shadow_eval_creates_job_and_frees_expired_or_exhausted_ones(monkeypatch: pytest.MonkeyPatch): """Expiry and turn-budget exhaustion both end sampling on their own; either must - release the key's slot in the active-job index so a new eval can start.""" + release the key's slot in the active-job index so a new eval can start. Every requested + key is swept and claimed on its own row, and each gets its own copy of the turn budget.""" import litellm.proxy.proxy_server as proxy_server prisma = _shadow_prisma() monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) - response = await start_shadow_eval(_start_request(), ADMIN) + response = await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN) assert response.status == "running" - assert response.max_turns == 200 assert response.judged_count is None - sweep_sql, sweep_key = prisma.db.execute_raw.call_args.args + sweep_sql = prisma.db.execute_raw.call_args.args[0] assert "stopped_at IS NULL" in sweep_sql - assert "ends_at <= NOW()" in sweep_sql - assert ">= j.max_turns" in sweep_sql - assert sweep_key == "key-hash" + assert "j.ends_at <= NOW()" in sweep_sql + assert ">= k.max_turns" in sweep_sql + assert "a.api_key_id = k.api_key_id" in sweep_sql + assert "k.api_key_id = ANY($1::text[])" in sweep_sql + assert prisma.db.litellm_shadowevaljob.create.call_args.kwargs["include"] == {"keys": True} + assert prisma.db.execute_raw.call_args.args[1] == ["key-hash", "key-hash-2"] create_data = prisma.db.litellm_shadowevaljob.create.call_args.kwargs["data"] - assert create_data["api_key_id"] == "key-hash" + assert create_data["keys"]["create"] == [ + {"api_key_id": "key-hash", "max_turns": 200}, + {"api_key_id": "key-hash-2", "max_turns": 200}, + ] assert create_data["created_by"] == "admin" - assert "status" not in create_data + assert "status" not in create_data and "api_key_id" not in create_data @pytest.mark.asyncio @@ -591,7 +623,8 @@ async def test_start_shadow_eval_creates_job_and_frees_expired_or_exhausted_ones (ADMIN, {"router_name": "not-a-router"}, None, 400), (ADMIN, {"judge_model": "not/a real model!"}, None, 400), (ADMIN, {"judge_model": "my-router"}, None, 400), - (ADMIN, {}, "active", 409), + (ADMIN, {}, ("key-hash",), 409), + (ADMIN, {"api_key_ids": ("key-hash", "key-hash-2")}, ("key-hash-2",), 409), (ADMIN, {"direction": "reverse", "baseline_model": "my-router"}, None, 400), (ADMIN, {"direction": "reverse", "baseline_model": "not/a real model!"}, None, 400), (ADMIN, {"direction": "reverse", "baseline_model": "openai/gpt-4o", "router_name": "not-a-router"}, None, 400), @@ -603,6 +636,7 @@ async def test_start_shadow_eval_creates_job_and_frees_expired_or_exhausted_ones "unresolvable-judge", "router-as-judge", "already-active", + "one-of-several-keys-already-active", "router-as-baseline", "unresolvable-baseline", "reverse-still-needs-an-auto-router", @@ -613,7 +647,7 @@ async def test_start_shadow_eval_rejections( ): import litellm.proxy.proxy_server as proxy_server - prisma = _shadow_prisma(active_job=_job_record() if active else None) + prisma = _shadow_prisma(claimed=[_claim(key) for key in active or ()]) monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) @@ -644,11 +678,8 @@ async def test_start_shadow_eval_reverse_records_its_arms_and_holds_its_own_slot the slot must not block a reverse one. The second reverse start still 409s.""" import litellm.proxy.proxy_server as proxy_server - prisma = _shadow_prisma() - active = {"forward": _job_record()} - prisma.db.litellm_shadowevaljob.find_first = AsyncMock( - side_effect=lambda where, **_: active.get(str(where.get("direction"))) - ) + claimed = [_claim("key-hash", job_id="job-1", direction="forward")] + prisma = _shadow_prisma(claimed=claimed) prisma.db.litellm_shadowevaljob.create = AsyncMock( return_value=_job_record(direction="reverse", baseline_model="openai/gpt-4o") ) @@ -663,7 +694,7 @@ async def test_start_shadow_eval_reverse_records_its_arms_and_holds_its_own_slot assert create_data["direction"] == "reverse" assert create_data["baseline_model"] == "openai/gpt-4o" - active["reverse"] = _job_record(id="job-2", direction="reverse") + claimed.append(_claim("key-hash", job_id="job-2", direction="reverse")) with pytest.raises(HTTPException) as exc: await start_shadow_eval(reverse, ADMIN) assert exc.value.status_code == 409 @@ -685,19 +716,62 @@ async def test_start_shadow_eval_forward_leaves_the_baseline_column_empty(monkey @pytest.mark.asyncio -async def test_start_shadow_eval_rejects_a_key_this_proxy_does_not_know(monkeypatch: pytest.MonkeyPatch): - """A typo'd api_key_id would otherwise create a job no traffic can ever match.""" +async def test_start_shadow_eval_rejects_a_key_already_active_in_another_job(monkeypatch: pytest.MonkeyPatch): + """One active job per key holds across jobs, so a key busy elsewhere blocks the whole + start rather than being silently dropped from it, and the 409 names which key and where.""" import litellm.proxy.proxy_server as proxy_server - prisma = _shadow_prisma() - prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + prisma = _shadow_prisma(claimed=[_claim("key-hash-2", job_id="job-7")]) monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) with pytest.raises(HTTPException) as exc: - await start_shadow_eval(_start_request(), ADMIN) + await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN) + assert exc.value.status_code == 409 + assert "key-hash-2 (job job-7)" in exc.value.detail + prisma.db.litellm_shadowevaljob.create.assert_not_called() + + +@pytest.mark.asyncio +async def test_start_shadow_eval_reuses_a_key_whose_previous_job_already_stopped(monkeypatch: pytest.MonkeyPatch): + """The claim is held by unstopped key rows only, matching the partial unique index. A read + that forgets that would strand every key that has ever finished a job.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(claimed=[_claim("key-hash", job_id="job-7", stopped_at=datetime.now(timezone.utc))]) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + job = await start_shadow_eval(_start_request(), ADMIN) + + assert job.job_id == "job-1" + prisma.db.litellm_shadowevaljob.create.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_start_shadow_eval_rejects_a_key_this_proxy_does_not_know(monkeypatch: pytest.MonkeyPatch): + """A typo'd api_key_id would otherwise create a job no traffic can ever match. Every + unknown key is named at once, so a caller passing several fixes them in one round.""" + import litellm.proxy.proxy_server as proxy_server + + prisma = _shadow_prisma(known_keys=("key-hash",)) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + + with pytest.raises(HTTPException) as exc: + await start_shadow_eval(_start_request(api_key_ids=("key-hash", "typo-a", "typo-b")), ADMIN) assert exc.value.status_code == 400 - assert "not a key on this proxy" in exc.value.detail + assert "typo-a, typo-b" in exc.value.detail + assert "key-hash," not in exc.value.detail + prisma.db.litellm_shadowevaljob.create.assert_not_called() + + +def test_start_shadow_eval_request_dedupes_and_requires_a_key(): + """A key named twice would collide with itself on the one-active-per-key index, and a + job scoping no key samples nothing.""" + assert _start_request(api_key_ids=("a", "b", "a")).api_key_ids == ("a", "b") + with pytest.raises(ValueError): + _start_request(api_key_ids=()) @pytest.mark.asyncio @@ -725,11 +799,15 @@ async def test_get_shadow_eval_job_derives_counts_spend_and_stratified_results(m {"grp": "SIMPLE", "turn_count": 8, "real_wins": 2, "shadow_wins": 4, "ties": 2, "avg_confidence": 0.8}, {"grp": "REASONING", "turn_count": 2, "real_wins": 2, "shadow_wins": 0, "ties": 0, "avg_confidence": 0.9}, ] - prisma = _shadow_prisma(agg_rows=tier_rows) - prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=_job_record()) - prisma.db.litellm_shadowevalattempt.find_first = AsyncMock( - return_value=MagicMock(error="judge call failed: boom") + key_rows = [ + {"grp": "key-hash", "turn_count": 6, "real_wins": 1, "shadow_wins": 4, "ties": 1, "avg_confidence": 0.7}, + {"grp": "key-hash-2", "turn_count": 4, "real_wins": 3, "shadow_wins": 0, "ties": 1, "avg_confidence": 0.6}, + ] + prisma = _shadow_prisma(agg_rows=tier_rows, by_key_rows=key_rows) + prisma.db.litellm_shadowevaljob.find_unique = AsyncMock( + return_value=_job_record(keys=(_key_row(), _key_row("key-hash-2", max_turns=50))) ) + prisma.db.litellm_shadowevalattempt.find_first = AsyncMock(return_value=MagicMock(error="judge call failed: boom")) monkeypatch.setattr(proxy_server, "prisma_client", prisma) response = await get_shadow_eval_job("job-1", VIEWER) @@ -744,6 +822,10 @@ async def test_get_shadow_eval_job_derives_counts_spend_and_stratified_results(m assert response.results.by_tier[0].shadow_win_rate_pct == 50.0 assert response.results.overall_shadow_win_rate_pct == 40.0 assert response.results.overall_tie_rate_pct == 20.0 + assert [(key.api_key_id, key.max_turns) for key in response.keys] == [("key-hash", 200), ("key-hash-2", 50)] + assert [(s.group, s.turn_count) for s in response.results.by_key] == [("key-hash", 6), ("key-hash-2", 4)] + assert response.results.by_key[0].shadow_win_rate_pct == 66.7 + assert prisma.db.litellm_shadowevaljob.find_unique.await_args.kwargs["include"] == {"keys": True} @pytest.mark.asyncio @@ -765,44 +847,80 @@ async def test_get_shadow_eval_job_404s_and_gates_on_role(monkeypatch: pytest.Mo async def test_list_shadow_eval_jobs_returns_derived_status_without_aggregates(monkeypatch: pytest.MonkeyPatch): import litellm.proxy.proxy_server as proxy_server + stamp = datetime.now(timezone.utc) prisma = _shadow_prisma() prisma.db.litellm_shadowevaljob.find_many = AsyncMock( return_value=[ _job_record(), _job_record(id="job-2", ends_at=datetime.now(timezone.utc) - timedelta(days=1)), - _job_record(id="job-3", stopped_at=datetime.now(timezone.utc)), + _job_record(id="job-3", keys=(_key_row(stopped_at=stamp),)), ] ) monkeypatch.setattr(proxy_server, "prisma_client", prisma) - jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50) + jobs = await list_shadow_eval_jobs(VIEWER, api_key_id="key-hash", limit=50) assert [job.status for job in jobs] == ["running", "completed", "stopped"] - swept = ShadowEvalJobResponse.model_validate( + assert all(job.judged_count is None and job.results is None for job in jobs) + assert prisma.db.query_raw.await_count == 0 + assert prisma.db.litellm_shadowevaljob.find_many.await_args.kwargs["where"] == { + "keys": {"some": {"api_key_id": "key-hash"}} + } + assert prisma.db.litellm_shadowevaljob.find_many.await_args.kwargs["include"] == {"keys": True} + + +@pytest.mark.parametrize( + ("stopped_flags", "days_left", "expected"), + [ + ((False, False), 7, "running"), + ((True, False), 7, "running"), + ((True, True), 7, "stopped"), + ((True, True), -1, "completed"), + ((False, False), -1, "completed"), + ], +) +def test_job_status_runs_until_every_key_stops_and_completed_outranks_stopped( + stopped_flags: tuple[bool, ...], days_left: int, expected: str +): + stamp = datetime.now(timezone.utc) + job = ShadowEvalJobResponse.model_validate( _job_record( - id="job-4", - ends_at=datetime.now(timezone.utc) - timedelta(days=1), - stopped_at=datetime.now(timezone.utc), + ends_at=datetime.now(timezone.utc) + timedelta(days=days_left), + keys=tuple( + _key_row(f"key-{index}", stopped_at=stamp if stopped else None) + for index, stopped in enumerate(stopped_flags) + ), ), from_attributes=True, ) - assert swept.status == "completed" - assert all(job.judged_count is None and job.results is None for job in jobs) - assert prisma.db.query_raw.await_count == 0 + + assert job.status == expected @pytest.mark.asyncio -async def test_stop_shadow_eval_sets_stopped_at_and_rejects_non_running(monkeypatch: pytest.MonkeyPatch): +async def test_stop_shadow_eval_stops_every_unstopped_key_row_and_rejects_non_running( + monkeypatch: pytest.MonkeyPatch, +): import litellm.proxy.proxy_server as proxy_server + stamp = datetime.now(timezone.utc) prisma = _shadow_prisma() - prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=_job_record()) + prisma.db.litellm_shadowevaljob.find_unique = AsyncMock( + side_effect=[ + _job_record(keys=(_key_row(), _key_row("key-hash-2"))), + _job_record(keys=(_key_row(stopped_at=stamp), _key_row("key-hash-2", stopped_at=stamp))), + ] + ) monkeypatch.setattr(proxy_server, "prisma_client", prisma) stopped = await stop_shadow_eval_job("job-1", ADMIN) assert stopped.status == "stopped" - update = prisma.db.litellm_shadowevaljob.update.call_args.kwargs + assert [key.stopped_at for key in stopped.keys] == [stamp, stamp] + update = prisma.db.litellm_shadowevaljobkey.update_many.await_args.kwargs + assert update["where"] == {"job_id": "job-1", "stopped_at": None} assert set(update["data"]) == {"stopped_at"} + assert prisma.db.litellm_shadowevaljob.find_unique.await_args.kwargs["include"] == {"keys": True} + assert not prisma.db.litellm_shadowevaljob.update.called prisma.db.litellm_shadowevaljob.find_unique = AsyncMock( return_value=_job_record(ends_at=datetime.now(timezone.utc) - timedelta(days=1)) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx index e342ee33f25..85fa80baf45 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render, screen } from "@testing-library/react"; +import { fireEvent, render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; @@ -72,12 +72,12 @@ import { const job = (overrides: Partial = {}): ShadowEvalJob => ({ job_id: "job-1", status: "running", - router_name: "claude-auto", direction: "forward", + router_name: "claude-auto", baseline_model: null, judge_model: "anthropic/claude-sonnet-5", shadow_percentage: 10, - max_turns: 200, + keys: [{ api_key_id: "hashed-key-abc", max_turns: 200, stopped_at: null }], judged_count: 42, error_count: 1, judge_spend: 3.21, @@ -110,13 +110,21 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ avg_judge_confidence: 0.8, }, ], + by_key: [ + { + group: "hashed-key-abc", + turn_count: 42, + real_win_rate_pct: 30.0, + shadow_win_rate_pct: 45.0, + tie_rate_pct: 25.0, + avg_judge_confidence: 0.8, + }, + ], overall_shadow_win_rate_pct: 48.0, overall_tie_rate_pct: 22.0, }, created_at: "2026-08-07T00:00:00Z", ends_at: "2026-09-07T00:00:00Z", - stopped_at: null, - api_key_id: "hashed-key-abc", last_error: null, ...overrides, }); @@ -197,8 +205,8 @@ describe("ShadowEvalSection", () => { it("gives every active job its own card with a stop button, with the form still offered", () => { mockHooks({ jobs: [ - job({ job_id: "job-a", status: "running", api_key_id: "key-a" }), - job({ job_id: "job-b", status: "running", api_key_id: "key-b" }), + job({ job_id: "job-a", status: "running", keys: [{ api_key_id: "key-a", max_turns: 200, stopped_at: null }] }), + job({ job_id: "job-b", status: "running", keys: [{ api_key_id: "key-b", max_turns: 200, stopped_at: null }] }), ], }); render(); @@ -219,7 +227,7 @@ describe("ShadowEvalSection", () => { render(); expect(screen.queryByText("Start a shadow eval")).not.toBeInTheDocument(); expect(screen.queryByRole("button", { name: "Stop" })).not.toBeInTheDocument(); - expect(screen.getByText("running")).toBeInTheDocument(); + expect(screen.getAllByText("running").filter((badge) => badge.closest("td") === null)).toHaveLength(1); }); it("never labels a collapsed previous eval as empty from a countless list row", () => { @@ -334,21 +342,91 @@ describe("ShadowEvalSection", () => { expect(screen.getByText("Start a shadow eval")).toBeInTheDocument(); }); + it("breaks results down per key, so one key exhausting its own budget is visible while a sibling runs on", () => { + mockHooks({ + jobs: [ + job({ + judged_count: 205, + keys: [ + { api_key_id: "hash-spent", max_turns: 200, stopped_at: "2026-08-08T00:00:00Z" }, + { api_key_id: "hash-hungry", max_turns: 500, stopped_at: null }, + ], + results: { + by_tier: [], + by_current_model: [], + by_key: [ + { + group: "hash-spent", + turn_count: 200, + real_win_rate_pct: 20.0, + shadow_win_rate_pct: 60.0, + tie_rate_pct: 20.0, + avg_judge_confidence: 0.9, + }, + ], + overall_shadow_win_rate_pct: 60.0, + overall_tie_rate_pct: 20.0, + }, + }), + ], + }); + render(); + + const spent = screen.getByText("hash-spent").closest("tr"); + const hungry = screen.getByText("hash-hungry").closest("tr"); + if (!spent || !hungry) throw new Error("expected a table row per scoped key"); + + expect(within(spent).getByText("stopped")).toBeInTheDocument(); + expect(within(spent).getByText("200 / 200")).toBeInTheDocument(); + expect(within(spent).getByText("60.0%")).toBeInTheDocument(); + + expect(within(hungry).getByText("running")).toBeInTheDocument(); + expect(within(hungry).getByText("0 / 500")).toBeInTheDocument(); + expect(within(hungry).getByText("No verdicts yet")).toBeInTheDocument(); + + expect(screen.getByText(/205 of 700 turns judged/)).toBeInTheDocument(); + }); + + it("reads every key as completed once the job's window closes, whatever its own stop state", () => { + mockHooks({ + jobs: [ + job({ + status: "completed", + keys: [ + { api_key_id: "hash-spent", max_turns: 200, stopped_at: "2026-08-08T00:00:00Z" }, + { api_key_id: "hash-hungry", max_turns: 500, stopped_at: null }, + ], + }), + ], + }); + render(); + + const hungry = screen.getByText("hash-hungry").closest("tr"); + if (!hungry) throw new Error("expected a table row per scoped key"); + expect(within(hungry).getByText("completed")).toBeInTheDocument(); + expect(within(hungry).queryByText("running")).not.toBeInTheDocument(); + }); + it("renders nothing for non-admins when the proxy answers 403", () => { mockHooks({ error: new ApiError("forbidden", 403, {}) }); const { container } = render(); expect(container).toBeEmptyDOMElement(); }); - it("keeps the start button disabled until key, router, and judge model are picked, then submits them", async () => { + it("keeps the start button disabled until key, router, and judge model are picked, then submits every picked key", async () => { const user = userEvent.setup(); const { start } = mockHooks({}); render(); expect(screen.getByText("Start shadow eval")).toBeDisabled(); - await user.click(screen.getByPlaceholderText("Search keys by alias")); - await user.click(await screen.findByText("prod-alpha")); + const keyInput = screen.getByPlaceholderText("Search keys by alias"); + await user.click(keyInput); + const keyList = await screen.findByTestId("paginated-multi-select-list"); + await user.click(within(keyList).getByText("prod-alpha")); + await user.click(keyInput); + await user.click(within(keyList).getByText("staging-beta")); + await user.click(screen.getByPlaceholderText("Select an auto-router")); await user.click(await screen.findByText("gpt-auto")); @@ -359,7 +437,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByText("Start shadow eval")); const expectedBody = { - api_key_id: "hash-alpha", + api_key_ids: ["hash-alpha", "hash-beta"], router_name: "gpt-auto", direction: "forward", shadow_percentage: 10, @@ -394,7 +472,7 @@ describe("ShadowEvalSection", () => { await user.click(screen.getByText("Start shadow eval")); const expectedBody = { - api_key_id: "hash-alpha", + api_key_ids: ["hash-alpha"], router_name: "gpt-auto", direction: "reverse", baseline_model: "prod-claude", @@ -416,7 +494,7 @@ describe("ShadowEvalSection", () => { expect(screen.getByText("52.0%")).toBeInTheDocument(); expect(screen.getByText(/Router won 30.0%/)).toBeInTheDocument(); expect(screen.getByText(/Baseline won 48.0%/)).toBeInTheDocument(); - expect(screen.getAllByText("Baseline wins")).toHaveLength(2); + expect(screen.getAllByText("Baseline wins")).toHaveLength(3); expect(screen.getByText("Router pick")).toBeInTheDocument(); expect(screen.queryByText(/Current model/)).not.toBeInTheDocument(); expect(screen.queryByText("Compared against")).not.toBeInTheDocument(); @@ -444,5 +522,6 @@ describe("ShadowEvalSection", () => { expect(await screen.findByText("SIMPLE")).toBeInTheDocument(); expect(screen.getByText("REASONING")).toBeInTheDocument(); + expect(screen.getByText("42 / 200")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx index df2989d3990..a01030a12e0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx @@ -6,7 +6,7 @@ import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; -import { PaginatedSearchSelect } from "@/components/shared/PaginatedSearchSelect"; +import { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect"; import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -24,7 +24,9 @@ import { useStartShadowEval, useStopShadowEval, type ShadowEvalJob, + type ShadowEvalJobKey, type ShadowEvalSlice, + type StartShadowEvalRequest, } from "./useShadowEval"; const pct = (value: number): string => `${value.toFixed(1)}%`; @@ -164,6 +166,68 @@ const VerdictBar: React.FC<{ direction: ShadowEvalDirection; results: NonNullabl ); }; +const totalBudget = (job: ShadowEvalJob): number => job.keys.reduce((sum, key) => sum + key.max_turns, 0); + +const keyStatus = (job: ShadowEvalJob, key: ShadowEvalJobKey): string => { + if (job.status === "completed") return "completed"; + return key.stopped_at ? "stopped" : "running"; +}; + +const KeyTable: React.FC<{ job: ShadowEvalJob }> = ({ job }) => { + const slices = new Map((job.results?.by_key ?? []).map((slice) => [slice.group, slice])); + return ( + + + + Key + Status + {["Judged of budget", "Router wins", `${otherArmLabel(job.direction)} wins`, "Ties", "Judge confidence"].map( + (label) => ( + + {label} + + ), + )} + + + + {job.keys.map((key) => { + const slice = slices.get(key.api_key_id); + return ( + + + {key.api_key_id} + + + + + + {(slice?.turn_count ?? 0).toLocaleString()} / {key.max_turns.toLocaleString()} + + {slice ? ( + <> + + {pct(routerWinRate(job.direction, slice))} + + + {pct(otherArmWinRate(job.direction, slice))} + + {pct(slice.tie_rate_pct)} + {slice.avg_judge_confidence.toFixed(2)} + + ) : ( + + No verdicts yet + + )} + + ); + })} + +
+ ); +}; + const emptyResultsText = (job: ShadowEvalJob, resultsError: boolean): string => { if (resultsError) return "Results could not be loaded. Retrying."; if (isActive(job)) return "Collecting verdicts. Results appear as sampled requests are judged."; @@ -173,31 +237,41 @@ const emptyResultsText = (job: ShadowEvalJob, resultsError: boolean): string => const ResultsBody: React.FC<{ job: ShadowEvalJob; resultsError?: boolean }> = ({ job, resultsError = false }) => { const results = job.results; - if (!results || (results.by_tier.length === 0 && results.by_current_model.length === 0)) { - return

{emptyResultsText(job, resultsError)}

; - } return ( <> -
-

- Router matched or beat {job.direction === "reverse" ? "the baseline" : "your current model"} -

-

{pct(routerMatchedOrBeatPct(job.direction, results))}

-

of {(job.judged_count ?? 0).toLocaleString()} judged responses

+ {results && (results.by_tier.length > 0 || results.by_current_model.length > 0) ? ( + <> +
+

+ Router matched or beat {job.direction === "reverse" ? "the baseline" : "your current model"} +

+

+ {pct(routerMatchedOrBeatPct(job.direction, results))} +

+

+ of {(job.judged_count ?? 0).toLocaleString()} judged responses +

+
+ + {results.by_current_model.length > 0 && ( + + )} + {results.by_tier.length > 0 && ( +
0 ? "border-t" : ""}> + +
+ )} + + ) : ( +

{emptyResultsText(job, resultsError)}

+ )} +
+
- - {results.by_current_model.length > 0 && ( - - )} - {results.by_tier.length > 0 && ( -
0 ? "border-t" : ""}> - -
- )} ); }; @@ -219,7 +293,7 @@ const JobResults: React.FC<{

{jobHeadline(job)}

- {(job.judged_count ?? 0).toLocaleString()} of {job.max_turns.toLocaleString()} turns judged ·{" "} + {(job.judged_count ?? 0).toLocaleString()} of {totalBudget(job).toLocaleString()} turns judged ·{" "} {(job.error_count ?? 0).toLocaleString()} errored · {usd(job.judge_spend ?? 0)} judge spend {active && remaining ? ` · ${remaining}` : ""}

@@ -294,9 +368,9 @@ const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[ const START_FORM_DESCRIPTION: Record = { forward: - "Duplicates a sampled slice of the key's traffic through the auto-router and has an LLM judge compare both answers blind. The router's answers are never served to users; judge calls bill to the shadowed key.", + "Duplicates a sampled slice of each selected key's traffic through the auto-router and has an LLM judge compare both answers blind. The router's answers are never served to users; judge calls bill to the key whose traffic was sampled.", reverse: - "Duplicates a sampled slice of the traffic the auto-router already serves against a fixed baseline model and has an LLM judge compare both answers blind. The baseline's answers are never served to users; judge calls bill to the shadowed key.", + "Duplicates a sampled slice of the traffic the auto-router already serves for each selected key against a fixed baseline model and has an LLM judge compare both answers blind. The baseline's answers are never served to users; judge calls bill to the key whose traffic was sampled.", }; const DURATION_OPTIONS = [ @@ -321,7 +395,7 @@ const Field: React.FC<{ label: string; htmlFor?: string; className?: string; chi
); -const KeySelect: React.FC<{ value: string; onChange: (token: string) => void }> = ({ value, onChange }) => { +const KeySelect: React.FC<{ value: string[]; onChange: (tokens: string[]) => void }> = ({ value, onChange }) => { const [search, setSearch] = useState(""); const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteKeys(50, { selectedKeyAlias: search || null, @@ -338,7 +412,7 @@ const KeySelect: React.FC<{ value: string; onChange: (token: string) => void }> [data], ); return ( - void }> const StartForm: React.FC = () => { const { accessToken } = useAuthorized(); - const [apiKeyId, setApiKeyId] = useState(""); + const [apiKeyIds, setApiKeyIds] = useState([]); const [routerName, setRouterName] = useState(""); const [direction, setDirection] = useState("forward"); const [baselineModel, setBaselineModel] = useState(""); @@ -381,14 +455,13 @@ const StartForm: React.FC = () => { const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; const parsedMaxTurns = Number.parseInt(maxTurns, 10); const maxTurnsValid = parsedMaxTurns >= 1 && parsedMaxTurns <= 2000; - const filled = - [apiKeyId, routerName, judgeModel].every((field) => field !== "") && - (direction === "forward" || baselineModel !== ""); + const baselinePicked = direction === "forward" || baselineModel !== ""; + const filled = apiKeyIds.length > 0 && [routerName, judgeModel].every((field) => field !== "") && baselinePicked; const boundsValid = percentageValid && maxTurnsValid; const valid = Boolean(accessToken) && filled && boundsValid; const handleStart = () => { - const startBody = { - api_key_id: apiKeyId, + const startBody: StartShadowEvalRequest = { + api_key_ids: apiKeyIds, router_name: routerName, direction, ...(direction === "reverse" ? { baseline_model: baselineModel } : {}), @@ -425,8 +498,8 @@ const StartForm: React.FC = () => { - - + + { value={maxTurns} onChange={(e) => setMaxTurns(e.target.value)} /> - turns judged, max + turns judged per key, max
{maxTurns.trim() !== "" && !maxTurnsValid && (

Enter a value from 1 to 2000

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts index 027003df46f..defbad570a5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useShadowEval.ts @@ -7,6 +7,7 @@ import { $api, fetchClient } from "@/lib/http/api"; import type { components } from "@/lib/http/schema"; export type ShadowEvalJob = components["schemas"]["ShadowEvalJobResponse"]; +export type ShadowEvalJobKey = components["schemas"]["ShadowEvalJobKeyResponse"]; export type ShadowEvalSlice = components["schemas"]["ShadowEvalSlice"]; export type StartShadowEvalRequest = components["schemas"]["StartShadowEvalRequest"]; diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.test.tsx new file mode 100644 index 00000000000..c918016fa19 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.test.tsx @@ -0,0 +1,193 @@ +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { useState } from "react"; +import { describe, expect, it, vi } from "vitest"; + +import { PaginatedMultiSelect } from "./PaginatedMultiSelect"; +import type { SearchSelectOption } from "./SearchSelect"; + +const OPTIONS: SearchSelectOption[] = [ + { label: "alias-alpha", value: "alias-alpha" }, + { label: "alias-beta", value: "alias-beta" }, + { label: "gamma-key", value: "gamma-key" }, +]; + +function renderSelect(overrides: Partial> = {}) { + const props: React.ComponentProps = { + options: OPTIONS, + onValueChange: vi.fn(), + onSearchChange: vi.fn(), + onLoadMore: vi.fn(), + ...overrides, + }; + render(); + return props; +} + +function setListMetrics(list: HTMLElement, metrics: { scrollTop: number; clientHeight: number; scrollHeight: number }) { + Object.defineProperty(list, "scrollTop", { value: metrics.scrollTop, configurable: true }); + Object.defineProperty(list, "clientHeight", { value: metrics.clientHeight, configurable: true }); + Object.defineProperty(list, "scrollHeight", { value: metrics.scrollHeight, configurable: true }); +} + +function chipRemoveButton(label: string): HTMLElement { + const chip = screen.getByText(label).closest('[data-slot="combobox-chip"]'); + if (chip === null) throw new Error(`no chip found for ${label}`); + const button = chip.querySelector('[data-slot="combobox-chip-remove"]'); + if (button === null) throw new Error(`no remove control found on chip for ${label}`); + return button as HTMLElement; +} + +describe("PaginatedMultiSelect", () => { + it("reports the cleared query upstream after a selection, so the next open is not still filtered", async () => { + const user = userEvent.setup(); + const onSearchChange = vi.fn(); + + function Controlled() { + const [value, setValue] = useState([]); + return ( + + ); + } + render(); + + const input = screen.getByRole("combobox"); + await user.click(input); + await user.type(input, "alias-a"); + await waitFor(() => expect(onSearchChange).toHaveBeenCalledWith("alias-a"), { timeout: 2000 }); + + const list = await screen.findByTestId("paginated-multi-select-list"); + await user.click(within(list).getByText("alias-alpha")); + + expect(input).toHaveValue(""); + await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith(""), { timeout: 2000 }); + }); + + it("selects multiple values and reports them cumulatively", async () => { + const user = userEvent.setup(); + const onValueChange = vi.fn(); + + function Controlled() { + const [value, setValue] = useState([]); + return ( + { + setValue(next); + onValueChange(next); + }} + onSearchChange={vi.fn()} + onLoadMore={vi.fn()} + /> + ); + } + render(); + + const input = screen.getByRole("combobox"); + await user.click(input); + const list = await screen.findByTestId("paginated-multi-select-list"); + await user.click(within(list).getByText("alias-alpha")); + + await user.click(input); + await user.click(within(list).getByText("gamma-key")); + + expect(onValueChange).toHaveBeenLastCalledWith(["alias-alpha", "gamma-key"]); + const chips = document.querySelector('[data-slot="combobox-chips"]') as HTMLElement; + expect(within(chips).getByText("alias-alpha")).toBeInTheDocument(); + expect(within(chips).getByText("gamma-key")).toBeInTheDocument(); + }); + + it("deselects one value via the chip remove control and keeps the rest", async () => { + const user = userEvent.setup(); + const onValueChange = vi.fn(); + + function Controlled() { + const [value, setValue] = useState(["alias-alpha", "alias-beta"]); + return ( + { + setValue(next); + onValueChange(next); + }} + onSearchChange={vi.fn()} + onLoadMore={vi.fn()} + /> + ); + } + render(); + + await user.click(chipRemoveButton("alias-alpha")); + + expect(onValueChange).toHaveBeenCalledWith(["alias-beta"]); + expect(screen.queryByText("alias-alpha")).not.toBeInTheDocument(); + expect(screen.getByText("alias-beta")).toBeInTheDocument(); + }); + + it("keeps a selected chip visible after the options page no longer contains it", () => { + const { rerender } = render( + , + ); + + expect(screen.getByText("ghost-key")).toBeInTheDocument(); + + rerender( + , + ); + + expect(screen.getByText("ghost-key")).toBeInTheDocument(); + }); + + it("does not request the next page on scroll when there is no next page", async () => { + const user = userEvent.setup(); + const onLoadMore = vi.fn(); + renderSelect({ onLoadMore, hasNextPage: false }); + + const input = screen.getByRole("combobox"); + await user.click(input); + const list = await screen.findByTestId("paginated-multi-select-list"); + + setListMetrics(list, { scrollTop: 900, clientHeight: 100, scrollHeight: 1000 }); + fireEvent.scroll(list); + + expect(onLoadMore).not.toHaveBeenCalled(); + }); + + it("requests the next page once scrolled past the threshold when a next page exists", async () => { + const user = userEvent.setup(); + const onLoadMore = vi.fn(); + renderSelect({ onLoadMore, hasNextPage: true }); + + const input = screen.getByRole("combobox"); + await user.click(input); + const list = await screen.findByTestId("paginated-multi-select-list"); + + setListMetrics(list, { scrollTop: 0, clientHeight: 100, scrollHeight: 1000 }); + fireEvent.scroll(list); + expect(onLoadMore).not.toHaveBeenCalled(); + + setListMetrics(list, { scrollTop: 850, clientHeight: 100, scrollHeight: 1000 }); + fireEvent.scroll(list); + expect(onLoadMore).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.tsx new file mode 100644 index 00000000000..72ca8986a33 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/PaginatedMultiSelect.tsx @@ -0,0 +1,140 @@ +"use client"; + +import { Loader2 } from "lucide-react"; +import { useMemo, useState } from "react"; + +import { + Combobox, + ComboboxChip, + ComboboxChips, + ComboboxChipsInput, + ComboboxContent, + ComboboxEmpty, + ComboboxItem, + ComboboxList, + ComboboxValue, +} from "@/components/ui/combobox"; + +import type { SearchSelectOption } from "./SearchSelect"; +import { usePaginatedCombobox } from "./usePaginatedCombobox"; + +interface PaginatedMultiSelectProps { + options: SearchSelectOption[]; + value?: string[]; + onValueChange: (value: string[]) => void; + onSearchChange: (query: string) => void; + onLoadMore: () => void; + hasNextPage?: boolean; + isLoading?: boolean; + isFetchingNextPage?: boolean; + placeholder?: string; + emptyText?: string; + errorText?: string; + loadingText?: string; + disabled?: boolean; + className?: string; + inputId?: string; + "aria-invalid"?: true | undefined; + "aria-describedby"?: string; +} + +export function PaginatedMultiSelect({ + options, + value = [], + onValueChange, + onSearchChange, + onLoadMore, + hasNextPage = false, + isLoading = false, + isFetchingNextPage = false, + placeholder = "Search…", + emptyText = "No results", + errorText, + loadingText = "Loading…", + disabled = false, + className, + inputId, + "aria-invalid": ariaInvalid, + "aria-describedby": ariaDescribedBy, +}: PaginatedMultiSelectProps) { + const [query, setQuery] = useState(""); + + const selected = useMemo( + () => + value.map( + (selectedValue) => + options.find((option) => option.value === selectedValue) ?? { label: selectedValue, value: selectedValue }, + ), + [options, value], + ); + + const items = useMemo(() => { + const missing = selected.filter((option) => !options.some((o) => o.value === option.value)); + return missing.length === 0 ? options : [...missing, ...options]; + }, [options, selected]); + + const pagination = { onSearchChange, onLoadMore, hasNextPage, isFetchingNextPage }; + const { handleSearchInput, handleScroll } = usePaginatedCombobox(pagination); + + const handleInputValueChange = (next: string, reason: string) => { + setQuery(next); + handleSearchInput(next, reason); + }; + + return ( + onValueChange(next.map((option) => option.value))} + inputValue={query} + onInputValueChange={(next, eventDetails) => handleInputValueChange(next, eventDetails.reason)} + isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value} + itemToStringLabel={(item: SearchSelectOption) => item.label} + filter={null} + disabled={disabled} + > + + + {(selectedItems: SearchSelectOption[]) => + selectedItems.map((option) => ( + + {option.label} + + )) + } + + + + + + {errorText ?? (isLoading ? loadingText : emptyText)} + + + {(item: SearchSelectOption) => ( + + + {item.label} + {item.sublabel != null && item.sublabel !== "" && ( + {item.sublabel} + )} + + + )} + + {isFetchingNextPage && ( +
+ +
+ )} + + + ); +} diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx index 47dc429bcc5..00041f94521 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx @@ -1,8 +1,7 @@ "use client"; -import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { Loader2 } from "lucide-react"; -import { useMemo, type UIEvent } from "react"; +import { useMemo } from "react"; import { Combobox, @@ -12,13 +11,9 @@ import { ComboboxItem, ComboboxList, } from "@/components/ui/combobox"; -import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import type { SearchSelectOption } from "./SearchSelect"; - -const SCROLL_THRESHOLD = 0.8; - -const SEARCH_REASONS: ReadonlySet = new Set(["input-change", "input-clear", "clear-press"]); +import { usePaginatedCombobox } from "./usePaginatedCombobox"; interface PaginatedSearchSelectProps { options: SearchSelectOption[]; @@ -70,28 +65,15 @@ export function PaginatedSearchSelect({ return [selected, ...options]; }, [options, selected]); - const debouncedSearch = useDebouncedCallback(onSearchChange, { wait: DEBOUNCE_WAIT_MS }); - - const handleInputValueChange = (next: string, reason: string) => { - if (!SEARCH_REASONS.has(reason)) return; - debouncedSearch(next); - }; - - const handleScroll = (event: UIEvent) => { - const target = event.currentTarget; - if (target.scrollHeight === 0) return; - const ratio = (target.scrollTop + target.clientHeight) / target.scrollHeight; - if (ratio >= SCROLL_THRESHOLD && hasNextPage && !isFetchingNextPage) { - onLoadMore(); - } - }; + const pagination = { onSearchChange, onLoadMore, hasNextPage, isFetchingNextPage }; + const { handleSearchInput, handleScroll } = usePaginatedCombobox(pagination); return ( onValueChange(item?.value ?? "")} - onInputValueChange={(next, eventDetails) => handleInputValueChange(next, eventDetails.reason)} + onInputValueChange={(next, eventDetails) => handleSearchInput(next, eventDetails.reason)} isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value} itemToStringLabel={(item: SearchSelectOption) => item.label} filter={null} diff --git a/ui/litellm-dashboard/src/components/shared/usePaginatedCombobox.test.ts b/ui/litellm-dashboard/src/components/shared/usePaginatedCombobox.test.ts new file mode 100644 index 00000000000..0d22a871597 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/usePaginatedCombobox.test.ts @@ -0,0 +1,98 @@ +import { act, renderHook } from "@testing-library/react"; +import type { UIEvent } from "react"; +import { describe, expect, it, vi } from "vitest"; + +import { usePaginatedCombobox } from "./usePaginatedCombobox"; + +const scrollEvent = (scrollTop: number, clientHeight: number, scrollHeight: number) => + ({ currentTarget: { scrollTop, clientHeight, scrollHeight } }) as UIEvent; + +const setup = (overrides: Partial[0]> = {}) => { + const onLoadMore = vi.fn(); + const onSearchChange = vi.fn(); + const options = { onSearchChange, onLoadMore, hasNextPage: true, isFetchingNextPage: false, ...overrides }; + const { result } = renderHook(() => usePaginatedCombobox(options)); + return { result, onLoadMore, onSearchChange }; +}; + +describe("usePaginatedCombobox", () => { + it("loads the next page only once the list is scrolled most of the way down", () => { + const { result, onLoadMore } = setup(); + + result.current.handleScroll(scrollEvent(0, 100, 1000)); + expect(onLoadMore).not.toHaveBeenCalled(); + + result.current.handleScroll(scrollEvent(690, 100, 1000)); + expect(onLoadMore).not.toHaveBeenCalled(); + + result.current.handleScroll(scrollEvent(700, 100, 1000)); + expect(onLoadMore).toHaveBeenCalledTimes(1); + }); + + it("does not ask for a page that does not exist or is already in flight", () => { + const exhausted = setup({ hasNextPage: false }); + exhausted.result.current.handleScroll(scrollEvent(900, 100, 1000)); + expect(exhausted.onLoadMore).not.toHaveBeenCalled(); + + const inFlight = setup({ isFetchingNextPage: true }); + inFlight.result.current.handleScroll(scrollEvent(900, 100, 1000)); + expect(inFlight.onLoadMore).not.toHaveBeenCalled(); + }); + + it("ignores a scroll on an unmeasured list, whose ratio divides by a zero scroll height", () => { + const { result, onLoadMore } = setup(); + result.current.handleScroll(scrollEvent(0, 100, 0)); + expect(onLoadMore).not.toHaveBeenCalled(); + }); + + it("searches on typing and clearing, but not on reasons that leave the query unchanged", async () => { + vi.useFakeTimers(); + try { + const { result, onSearchChange } = setup(); + + result.current.handleSearchInput("alpha", "item-press"); + result.current.handleSearchInput("alpha", "list-navigation"); + await act(async () => { + vi.advanceTimersByTime(1000); + }); + expect(onSearchChange).not.toHaveBeenCalled(); + + result.current.handleSearchInput("alpha", "input-change"); + await act(async () => { + vi.advanceTimersByTime(1000); + }); + expect(onSearchChange).toHaveBeenCalledWith("alpha"); + + result.current.handleSearchInput("", "input-clear"); + await act(async () => { + vi.advanceTimersByTime(1000); + }); + expect(onSearchChange).toHaveBeenLastCalledWith(""); + } finally { + vi.useRealTimers(); + } + }); + + it("debounces typing so one search leaves for a burst of keystrokes", async () => { + vi.useFakeTimers(); + try { + const { result, onSearchChange } = setup(); + + result.current.handleSearchInput("a", "input-change"); + result.current.handleSearchInput("al", "input-change"); + result.current.handleSearchInput("alp", "input-change"); + await act(async () => { + vi.advanceTimersByTime(299); + }); + expect(onSearchChange).not.toHaveBeenCalled(); + + await act(async () => { + vi.advanceTimersByTime(1); + }); + expect(onSearchChange).toHaveBeenCalledTimes(1); + expect(onSearchChange).toHaveBeenCalledWith("alp"); + } finally { + vi.useRealTimers(); + } + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/usePaginatedCombobox.ts b/ui/litellm-dashboard/src/components/shared/usePaginatedCombobox.ts new file mode 100644 index 00000000000..d77b6a9acf1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/usePaginatedCombobox.ts @@ -0,0 +1,47 @@ +"use client"; + +import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; +import type { UIEvent } from "react"; + +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; + +const SCROLL_THRESHOLD = 0.8; + +const SEARCH_REASONS: ReadonlySet = new Set(["input-change", "input-clear", "clear-press"]); + +interface UsePaginatedComboboxOptions { + onSearchChange: (query: string) => void; + onLoadMore: () => void; + hasNextPage: boolean; + isFetchingNextPage: boolean; +} + +interface PaginatedComboboxHandlers { + handleSearchInput: (next: string, reason: string) => void; + handleScroll: (event: UIEvent) => void; +} + +export function usePaginatedCombobox({ + onSearchChange, + onLoadMore, + hasNextPage, + isFetchingNextPage, +}: UsePaginatedComboboxOptions): PaginatedComboboxHandlers { + const debouncedSearch = useDebouncedCallback(onSearchChange, { wait: DEBOUNCE_WAIT_MS }); + + const handleSearchInput = (next: string, reason: string) => { + if (!SEARCH_REASONS.has(reason)) return; + debouncedSearch(next); + }; + + const handleScroll = (event: UIEvent) => { + const target = event.currentTarget; + if (target.scrollHeight === 0) return; + const ratio = (target.scrollTop + target.clientHeight) / target.scrollHeight; + if (ratio >= SCROLL_THRESHOLD && hasNextPage && !isFetchingNextPage) { + onLoadMore(); + } + }; + + return { handleSearchInput, handleScroll }; +} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cf37709c377..5696765bdd4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -838,20 +838,21 @@ export interface paths { put?: never; /** * Start Shadow Eval - * @description Start a shadow eval: duplicate a sampled slice of a key's live traffic against a second - * arm, judge the two responses blind, and stratify win rates by tier and by the model that - * served the real arm. + * @description Start a shadow eval: duplicate a sampled slice of one or more keys' live traffic against + * a second arm, judge the two responses blind, and stratify win rates by tier, by the model + * that served the real arm, and by key. * - * A forward job answers whether the key should adopt router_name: it samples the requests + * A forward job answers whether the keys should adopt router_name: it samples the requests * the router did not serve and duplicates them through it. A reverse job answers whether a * key already on the router still gains from it: it samples the requests the router did * serve and duplicates them against baseline_model. A key can hold one active job per * direction, so both questions can run at once. * - * Shadow responses are never served to users. The job samples until it has judged - * max_turns turns, reaches the end of its window, or is stopped; sampling changes - * propagate to pods within about 10 seconds. Shadow and judge calls bill to the - * shadowed key but are excluded from request counts and auto-router adoption metrics. + * Shadow responses are never served to users. Each key samples until it has judged + * max_turns turns of its own traffic, the job's window ends, or the job is stopped, so a + * busy key running out of budget does not end sampling for the others; sampling changes + * propagate to pods within about 10 seconds. Shadow and judge calls bill to the shadowed + * key but are excluded from request counts and auto-router adoption metrics. */ post: operations["start_shadow_eval_auto_router_shadow_eval_start_post"]; delete?: never; @@ -891,7 +892,9 @@ export interface paths { put?: never; /** * Stop Shadow Eval Job - * @description Stop an active shadow eval job. Attempts are kept; sampling halts within ~10s. + * @description Stop an active shadow eval job, every key it scopes at once. Attempts are kept; + * sampling halts within ~10s. Keys that already stopped on their own budget keep the + * stopped_at they earned, so one write records the outcome for the whole job. */ post: operations["stop_shadow_eval_job_auto_router_shadow_eval__job_id__stop_post"]; delete?: never; @@ -32744,19 +32747,35 @@ export interface components { /** Timeout */ timeout?: number | null; }; + /** + * ShadowEvalJobKeyResponse + * @description One key a job shadows, with its own budget and stop state. + */ + ShadowEvalJobKeyResponse: { + /** + * Api Key Id + * @description The hashed virtual key whose traffic this row scopes + */ + api_key_id: string; + /** + * Max Turns + * @description This key's own sample budget, independent of its siblings' + */ + max_turns: number; + /** + * Stopped At + * @description When this key stopped sampling, whether its own budget ran out, the window closed, or an operator stopped the job. The key reads completed whenever the job does, otherwise stopped once this is set and running until then + */ + stopped_at?: string | null; + }; /** * ShadowEvalJobResponse * @description A shadow-eval job. Validates directly from the prisma record (job_id reads the - * row's id); status is derived from stopped_at and ends_at, never stored, so no writer - * anywhere can produce an inconsistent one. Aggregate fields are populated by the + * row's id); status is derived from the keys' stopped_at and ends_at, never stored, so no + * writer anywhere can produce an inconsistent one. Aggregate fields are populated by the * detail endpoint only and stay None on list responses. */ ShadowEvalJobResponse: { - /** - * Api Key Id - * @description The hashed virtual key whose traffic this job evaluates, and only that key's - */ - api_key_id: string; /** Baseline Model */ baseline_model?: string | null; /** @@ -32794,13 +32813,16 @@ export interface components { * @description Verdicts recorded; detail endpoint only */ judged_count?: number | null; + /** + * Keys + * @description The keys whose traffic this job evaluates, and only those keys', each with its own budget + */ + keys: components["schemas"]["ShadowEvalJobKeyResponse"][]; /** * Last Error * @description Most recent attempt error; detail endpoint only */ last_error?: string | null; - /** Max Turns */ - max_turns: number; /** @description Stratified verdicts; detail endpoint only */ results?: components["schemas"]["ShadowEvalResult"] | null; /** Router Name */ @@ -32809,13 +32831,12 @@ export interface components { shadow_percentage: number; /** * Status - * @description A job whose window has passed reads completed even if a later sweep stamped - * stopped_at; stopped means sampling ended before the window did. + * @description A job whose window has passed reads completed even if a sweep stamped its keys + * first; stopped means every key ended sampling before the window did. One key + * exhausting its own budget leaves the job running while any sibling still samples. * @enum {string} */ readonly status: "running" | "completed" | "stopped"; - /** Stopped At */ - stopped_at?: string | null; }; /** * ShadowEvalResult @@ -32824,9 +32845,15 @@ export interface components { ShadowEvalResult: { /** * By Current Model - * @description Sliced by the model that served the real arm: the key's incumbent models in forward mode, and in reverse the models the router itself picked + * @description Sliced by the model that served the real arm: the keys' incumbent models in forward mode, and in reverse the models the router itself picked */ by_current_model: components["schemas"]["ShadowEvalSlice"][]; + /** + * By Key + * @description One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job scopes but has not yet judged a turn for are absent rather than reported as zero + * @default [] + */ + by_key: components["schemas"]["ShadowEvalSlice"][]; /** By Tier */ by_tier: components["schemas"]["ShadowEvalSlice"][]; /** Overall Shadow Win Rate Pct */ @@ -32836,8 +32863,8 @@ export interface components { }; /** * ShadowEvalSlice - * @description Judge outcomes for one slice of a job's verdicts (a router tier, or one of the - * models that served the real arm). + * @description Judge outcomes for one slice of a job's verdicts (a router tier, one of the models + * that served the real arm, or one of the keys the job is scoped to). */ ShadowEvalSlice: { /** Avg Judge Confidence */ @@ -33034,14 +33061,14 @@ export interface components { }; /** * StartShadowEvalRequest - * @description Start duplicating a key's traffic for blind comparison against an auto-router. + * @description Start duplicating one or more keys' traffic for blind comparison against an auto-router. */ StartShadowEvalRequest: { /** - * Api Key Id - * @description The hashed virtual key whose traffic will be shadowed. Shadow evaluation runs ONLY on this key's traffic; requests made with any other key are not sampled. + * Api Key Ids + * @description The hashed virtual keys whose traffic will be shadowed. Shadow evaluation runs ONLY on these keys' traffic; requests made with any other key are not sampled. Each key gets its own max_turns budget, so one key exhausting its budget does not end sampling for the others. */ - api_key_id: string; + api_key_ids: string[]; /** * Baseline Model * @description Required when direction is reverse and rejected otherwise: the fixed model the router's own responses are judged against. Must be a plain model rather than another auto-router @@ -33068,7 +33095,7 @@ export interface components { judge_model: string; /** * Max Turns - * @description Sample budget: the job judges at most this many turns, then completes. This is also the spend bound; expected judge cost is roughly max_turns times one judge call + * @description Per-key sample budget: the job judges at most this many turns of EACH scoped key's traffic, so a job over N keys judges at most N times max_turns turns. This is also the spend bound; expected judge cost is roughly that turn ceiling times one judge call * @default 200 */ max_turns: number; @@ -37370,7 +37397,7 @@ export interface operations { list_shadow_eval_jobs_auto_router_shadow_eval_get: { parameters: { query?: { - /** @description Filter to jobs shadowing this key */ + /** @description Filter to jobs that shadow this key, alone or alongside others */ api_key_id?: string | null; /** @description Newest jobs to return */ limit?: number;