mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix: attribute shadow eval spend to the shadowed key's budget
Shadow and judge calls were fired with no caller identity on their metadata. The proxy's cost callback requires user_api_key/_team_id/... to log spend and apply budget checks, and silently drops the entry without them — so an admin-enabled eval billed real provider spend that landed on no key, no team, and no budget counter, invisible to every limit the shadowed key is normally subject to. Extract the identity-forwarding rules the auto-router classifier already implements into a shared litellm/litellm_core_utils/internal_call_metadata.py: forward the caller's identity subset, strip the parent's budget reservation (top-level and the copy nested in user_api_key_auth) so a sub-call can't finalize a reservation that belongs to the parent, and stamp the sub-call's origin. The classifier now uses this module instead of its own copy. Wire both shadow eval call sites (_call_router_shadow, _call_judge) through it, and add a per-job spend cap (job cost_actual >= max(3x the quoted estimate, $0.50)) so a bad estimate or a traffic spike can't turn a quoted eval into a runaway bill; a capped job self-completes and is evicted from cache so it can't be resurrected by an in-flight request. Also surface api_key_id/team_id on GetShadowEvalJobResponse so an admin running several jobs can tell which key's traffic a given win rate belongs to, and regenerate the dashboard's OpenAPI types for the new fields. Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Abhimanyu Kapur
parent
0b94603ca8
commit
28878e877b
9 changed files with 540 additions and 46 deletions
|
|
@ -16,7 +16,11 @@ Flow per sampled request (all in a detached background task, zero added latency)
|
|||
3. Write a ``LiteLLM_ShadowEvalVerdict`` row and bump the job's counters.
|
||||
|
||||
Shadow and judge calls carry ``shadow_eval_internal`` metadata so this logger
|
||||
ignores its own traffic and cannot recurse.
|
||||
ignores its own traffic and cannot recurse. They also carry the shadowed key's
|
||||
identity metadata, so the provider spend they incur is attributed to that key
|
||||
(and its team/org/user) and counts against every budget that key is subject to,
|
||||
including the global proxy budget. A job additionally stops itself once its judge
|
||||
spend reaches a multiple of the estimate quoted when it started.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
|
@ -26,6 +30,7 @@ import random
|
|||
import re
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
|
|
@ -34,6 +39,8 @@ from pydantic import BaseModel, TypeAdapter
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.internal_call_metadata import sanitized_forwardable_call_metadata
|
||||
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
|
@ -65,6 +72,14 @@ JUDGE_MAX_OUTPUT_TOKENS: Final = 500
|
|||
|
||||
_SEEN_FLUSH_INTERVAL_SECONDS: Final = 10.0
|
||||
|
||||
# A job stops sampling once its judge spend reaches this multiple of the estimate shown
|
||||
# at start. The headroom absorbs an estimate that undershot the real traffic mix without
|
||||
# letting the job run away; the floor keeps a cent-sized estimate from stopping a job on
|
||||
# its first verdict. cost_actual is read from the job cache, so overshoot is bounded by
|
||||
# _JOB_CACHE_TTL_SECONDS worth of judge calls.
|
||||
_SPEND_CAP_MULTIPLIER: Final = 1.5
|
||||
_SPEND_CAP_FLOOR_USD: Final = 1.0
|
||||
|
||||
_EMPTY_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
PAIRWISE_JUDGE_SYSTEM_PROMPT: Final = """You are an impartial quality judge comparing two responses to the same conversation.
|
||||
|
|
@ -156,6 +171,21 @@ class ActiveShadowEvalJob:
|
|||
shadow_percentage: float
|
||||
judge_model: str
|
||||
status: str
|
||||
cost_estimate: float | None = None
|
||||
cost_actual: float = 0.0
|
||||
|
||||
|
||||
def _job_is_over_spend_cap(job: ActiveShadowEvalJob) -> bool:
|
||||
"""Whether the job has spent past what its start-time estimate justifies.
|
||||
|
||||
Budgets bound what the *key* may spend; this bounds what a single eval may spend
|
||||
even under a generous budget, so a bad estimate or a traffic spike cannot quietly
|
||||
turn a "$3 eval" into a much larger bill. A job with no estimate (older rows,
|
||||
unpriced judge model) has nothing to be a multiple of, so it is uncapped.
|
||||
"""
|
||||
if job.cost_estimate is None or job.cost_estimate <= 0.0:
|
||||
return False
|
||||
return job.cost_actual >= max(job.cost_estimate * _SPEND_CAP_MULTIPLIER, _SPEND_CAP_FLOOR_USD)
|
||||
|
||||
|
||||
_JobCache: TypeAlias = "dict[str, tuple[float, ActiveShadowEvalJob | None]]" # mutable-ok: TTL cache
|
||||
|
|
@ -208,6 +238,9 @@ class ShadowEvalLogger(CustomLogger):
|
|||
job: Final = await self._get_active_job(api_key_hash)
|
||||
if job is None:
|
||||
return
|
||||
if _job_is_over_spend_cap(job):
|
||||
await self._stop_job_over_spend_cap(job)
|
||||
return
|
||||
request_id: Final = payload.get("id") or ""
|
||||
if not request_id:
|
||||
return
|
||||
|
|
@ -238,6 +271,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
model_parameters=MappingProxyType(
|
||||
dict(payload.get("model_parameters") or {}) # mutable-ok: frozen snapshot
|
||||
),
|
||||
parent_metadata=MappingProxyType(dict(request_metadata)), # mutable-ok: frozen snapshot
|
||||
)
|
||||
)
|
||||
task.add_done_callback(lambda _: setattr(self, "_inflight_shadow_tasks", self._inflight_shadow_tasks - 1))
|
||||
|
|
@ -269,6 +303,8 @@ class ShadowEvalLogger(CustomLogger):
|
|||
shadow_percentage=float(record.shadow_percentage),
|
||||
judge_model=str(record.judge_model),
|
||||
status=str(record.status),
|
||||
cost_estimate=float(record.cost_estimate) if record.cost_estimate is not None else None,
|
||||
cost_actual=float(record.cost_actual or 0.0),
|
||||
)
|
||||
if record is not None
|
||||
else None
|
||||
|
|
@ -279,6 +315,38 @@ class ShadowEvalLogger(CustomLogger):
|
|||
self._job_cache[api_key_hash] = (now, job)
|
||||
return job
|
||||
|
||||
async def _stop_job_over_spend_cap(self, job: ActiveShadowEvalJob) -> None:
|
||||
"""Flip a runaway job to completed, keeping the verdicts it already produced.
|
||||
|
||||
Guarded on the job still being pending/running so two pods breaching the cap on
|
||||
the same job cannot resurrect one an admin stopped in between.
|
||||
"""
|
||||
prisma: Final = self._prisma_provider()
|
||||
if prisma is None:
|
||||
return
|
||||
self._job_cache = { # mutable-ok: TTL cache
|
||||
k: v for k, v in self._job_cache.items() if v[1] is None or v[1].id != job.id
|
||||
}
|
||||
verbose_logger.info(
|
||||
"shadow_eval: stopping job %s, spend $%.4f reached the cap for its $%.4f estimate",
|
||||
job.id,
|
||||
job.cost_actual,
|
||||
job.cost_estimate or 0.0,
|
||||
)
|
||||
try:
|
||||
await prisma.db.litellm_shadowevaljob.update_many(
|
||||
where={ # mutable-ok: Prisma filter
|
||||
"id": job.id,
|
||||
"status": {"in": ["pending", "running"]}, # mutable-ok: Prisma filter
|
||||
},
|
||||
data={ # mutable-ok: Prisma payload
|
||||
"status": "completed",
|
||||
"completed_at": datetime.now(timezone.utc),
|
||||
},
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # logging hooks must never fail the request
|
||||
verbose_logger.debug("shadow_eval: failed to stop over-cap job %s: %s", job.id, e)
|
||||
|
||||
async def _flush_seen_counts(self) -> None:
|
||||
prisma: Final = self._prisma_provider()
|
||||
if prisma is None:
|
||||
|
|
@ -304,6 +372,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
response_obj: object,
|
||||
real_model: str,
|
||||
model_parameters: Mapping[str, object],
|
||||
parent_metadata: Mapping[str, object],
|
||||
) -> None:
|
||||
"""Detached background task: shadow call -> blind judge -> verdict row."""
|
||||
prisma: Final = self._prisma_provider()
|
||||
|
|
@ -312,7 +381,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
if not real_text or not messages:
|
||||
return
|
||||
|
||||
shadow: Final = await self._call_router_shadow(job.router_name, messages, model_parameters)
|
||||
shadow: Final = await self._call_router_shadow(job.router_name, messages, model_parameters, parent_metadata)
|
||||
if shadow is None:
|
||||
await self._bump_failed(job.id)
|
||||
return
|
||||
|
|
@ -323,6 +392,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
messages=messages,
|
||||
real_text=real_text,
|
||||
shadow_text=shadow_text,
|
||||
parent_metadata=parent_metadata,
|
||||
)
|
||||
if verdict is None:
|
||||
await self._bump_failed(job.id)
|
||||
|
|
@ -373,16 +443,23 @@ class ShadowEvalLogger(CustomLogger):
|
|||
verbose_logger.debug("shadow_eval: failed_count increment failed: %s", e)
|
||||
|
||||
async def _call_router_shadow(
|
||||
self, router_name: str, messages: Sequence[Mapping[str, object]], model_parameters: Mapping[str, object]
|
||||
self,
|
||||
router_name: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
model_parameters: Mapping[str, object],
|
||||
parent_metadata: Mapping[str, object],
|
||||
) -> tuple[str, str, str | None, int | None] | None:
|
||||
"""Send the prompt through the auto-router; return (text, model, tier, completion_tokens)."""
|
||||
router: Final = self._router_provider()
|
||||
if router is None:
|
||||
verbose_logger.debug("shadow_eval: no router available")
|
||||
return None
|
||||
# The router's pre-routing hook writes its routing decision into this
|
||||
# metadata dict; read it back after the call for tier attribution.
|
||||
shadow_metadata: Final[dict[str, object]] = { # mutable-ok: router writes back
|
||||
# Carries the shadowed key's identity so this call's real provider spend is
|
||||
# attributed to, and budget-checked against, the key whose traffic it copies.
|
||||
# The router's pre-routing hook also writes its routing decision into this dict;
|
||||
# read it back after the call for tier attribution.
|
||||
attribution: Final = sanitized_forwardable_call_metadata(parent_metadata, SHADOW_EVAL_ROUTER_CALL_ORIGIN)
|
||||
shadow_metadata: Final[dict[str, object]] = attribution | { # mutable-ok: router writes back
|
||||
SHADOW_EVAL_INTERNAL_MARKER: True
|
||||
}
|
||||
shadow_params: Final = { # mutable-ok: splatted as kwargs
|
||||
|
|
@ -417,6 +494,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
messages: Sequence[Mapping[str, object]],
|
||||
real_text: str,
|
||||
shadow_text: str,
|
||||
parent_metadata: Mapping[str, object],
|
||||
) -> tuple[str, float, str, float] | None:
|
||||
"""Blind pairwise judge. Returns (preference, confidence, reasoning, cost)."""
|
||||
real_is_a: Final = random.random() < 0.5
|
||||
|
|
@ -434,6 +512,9 @@ class ShadowEvalLogger(CustomLogger):
|
|||
f"Response B:\n{response_b[:_MAX_JUDGE_CHARS]}\n\n"
|
||||
"Which response is better?"
|
||||
)
|
||||
judge_metadata: Final = sanitized_forwardable_call_metadata(parent_metadata, SHADOW_EVAL_JUDGE_CALL_ORIGIN) | {
|
||||
SHADOW_EVAL_INTERNAL_MARKER: True
|
||||
}
|
||||
try:
|
||||
response: Final = await litellm.acompletion(
|
||||
model=judge_model,
|
||||
|
|
@ -443,7 +524,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
],
|
||||
temperature=0,
|
||||
max_tokens=JUDGE_MAX_OUTPUT_TOKENS,
|
||||
metadata={SHADOW_EVAL_INTERNAL_MARKER: True}, # mutable-ok: SDK metadata
|
||||
metadata=judge_metadata,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # judge outages are a counted failure, not a crash
|
||||
verbose_logger.debug("shadow_eval: judge call failed: %s", e)
|
||||
|
|
|
|||
107
litellm/litellm_core_utils/internal_call_metadata.py
Normal file
107
litellm/litellm_core_utils/internal_call_metadata.py
Normal file
|
|
@ -0,0 +1,107 @@
|
|||
"""Metadata a request forwards to the internal LLM sub-calls it triggers.
|
||||
|
||||
Internal features (the auto-router's classifier and its keyword embeddings, shadow
|
||||
eval's shadow + judge calls) bill real provider spend that nobody typed a prompt for.
|
||||
That spend must land on the same key/team/org/user as the request that caused it, so
|
||||
the sub-call has to carry the caller's identity metadata: the proxy's cost callback
|
||||
reads ``user_api_key`` / ``user_api_key_team_id`` / ... straight off the call's
|
||||
metadata, and drops the spend log entirely when they are missing.
|
||||
|
||||
Two things must never be forwarded as-is:
|
||||
|
||||
* ``user_api_key_budget_reservation`` (and the same reservation nested inside
|
||||
``user_api_key_auth``) belongs to the parent completion, not to the sub-call. If a
|
||||
sub-call's cost callback sees it, that callback finalizes the reservation and the
|
||||
parent's own callback then skips incrementing the key/team budget counters, losing
|
||||
the parent's spend. ``user_api_key_auth`` itself is kept, sanitized, because model
|
||||
access-group filtering needs it.
|
||||
* The origin of the sub-call. ``INTERNAL_CALL_ORIGIN_METADATA_KEY`` stamps which
|
||||
feature made the call so a spend-log row can say it is not traffic the caller sent.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.types.utils import InternalCallOrigin
|
||||
|
||||
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
|
||||
|
||||
_USER_API_KEY_AUTH_KEY: Final = "user_api_key_auth"
|
||||
|
||||
# The caller-identity subset a detached sub-call needs to be attributed and
|
||||
# budget-checked exactly like the request that spawned it. Everything else on the
|
||||
# parent's metadata (routing decisions, guardrail state, the standard logging object,
|
||||
# the proxy request body) describes the parent call and would be a lie on a sub-call
|
||||
# that runs after it has already returned.
|
||||
FORWARDABLE_IDENTITY_METADATA_KEYS: Final = frozenset(
|
||||
{
|
||||
"user_api_key",
|
||||
"user_api_key_hash",
|
||||
"user_api_key_alias",
|
||||
"user_api_key_team_id",
|
||||
"user_api_key_org_id",
|
||||
"user_api_key_user_id",
|
||||
"user_api_key_end_user_id",
|
||||
_USER_API_KEY_AUTH_KEY,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def sanitize_user_api_key_auth(auth: object) -> object:
|
||||
"""Copy of the auth object with its budget reservation removed.
|
||||
|
||||
The proxy's cost callback falls back to reading the reservation from inside the
|
||||
auth object when the top-level key is absent, so forwarding it unsanitized
|
||||
re-creates the double-finalization that stripping the top-level key prevents.
|
||||
"""
|
||||
if isinstance(auth, dict):
|
||||
return {k: v for k, v in auth.items() if k != "budget_reservation"} # mutable-ok: SDK metadata value
|
||||
reservation: Final[object] = getattr(auth, "budget_reservation", None)
|
||||
model_copy: Final[object] = getattr(auth, "model_copy", None)
|
||||
if reservation is not None and callable(model_copy):
|
||||
return model_copy(update={"budget_reservation": None}) # mutable-ok: pydantic update payload
|
||||
return auth
|
||||
|
||||
|
||||
def _sanitized(parent_metadata: Mapping[str, object]) -> dict[str, object]: # mutable-ok: SDK metadata kwarg
|
||||
return { # mutable-ok: SDK metadata kwarg
|
||||
k: sanitize_user_api_key_auth(v) if k == _USER_API_KEY_AUTH_KEY else v
|
||||
for k, v in parent_metadata.items()
|
||||
if k not in BUDGET_RESERVATION_METADATA_KEYS
|
||||
}
|
||||
|
||||
|
||||
def forwarded_internal_call_metadata(
|
||||
parent_metadata: Mapping[str, object] | None,
|
||||
call_origin: InternalCallOrigin,
|
||||
) -> dict[str, object]: # mutable-ok: SDK metadata kwarg
|
||||
"""Parent metadata, minus its budget reservation, stamped with the sub-call's origin.
|
||||
|
||||
For sub-calls made *inside* the parent request (the auto-router classifier and its
|
||||
embeddings), where the parent's full context still describes the call being made.
|
||||
A parent with no metadata at all had nothing to attribute in the first place, so
|
||||
the sub-call is left unstamped rather than carrying a lone origin marker.
|
||||
"""
|
||||
if not parent_metadata:
|
||||
return {} # mutable-ok: SDK metadata kwarg
|
||||
return _sanitized(parent_metadata) | { # mutable-ok: SDK metadata kwarg
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin
|
||||
}
|
||||
|
||||
|
||||
def sanitized_forwardable_call_metadata(
|
||||
parent_metadata: Mapping[str, object],
|
||||
call_origin: InternalCallOrigin,
|
||||
) -> dict[str, object]: # mutable-ok: SDK metadata kwarg
|
||||
"""Just the caller's identity, stamped with the sub-call's origin.
|
||||
|
||||
For sub-calls detached from the parent request (shadow eval), which outlive it and
|
||||
must not inherit per-request state such as its routing decision or logging payload.
|
||||
The origin stamp is unconditional here: these calls are always internal, whether or
|
||||
not there was an identity to forward.
|
||||
"""
|
||||
identity: Final = {k: v for k, v in parent_metadata.items() if k in FORWARDABLE_IDENTITY_METADATA_KEYS}
|
||||
return _sanitized(identity) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} # mutable-ok: SDK metadata kwarg
|
||||
|
|
@ -598,6 +598,8 @@ class _ShadowEvalJobRow(BaseModel):
|
|||
id: str
|
||||
status: ShadowEvalStatus
|
||||
router_name: str
|
||||
api_key_id: str
|
||||
team_id: str | None = None
|
||||
shadow_percentage: float
|
||||
request_count: int
|
||||
completed_count: int
|
||||
|
|
@ -617,6 +619,8 @@ def _job_to_response(record: object, results: ShadowEvalResult | None) -> GetSha
|
|||
job_id=row.id,
|
||||
status=row.status,
|
||||
router_name=row.router_name,
|
||||
api_key_id=row.api_key_id,
|
||||
team_id=row.team_id,
|
||||
shadow_percentage=row.shadow_percentage,
|
||||
request_count=row.request_count,
|
||||
completed_count=row.completed_count,
|
||||
|
|
|
|||
|
|
@ -26,8 +26,9 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
|
|||
from pydantic import BaseModel, create_model
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, RETURN_RAW_MODEL_NAME_METADATA_KEY
|
||||
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata
|
||||
from litellm.llms.base_llm.base_utils import type_to_response_format_param
|
||||
from litellm.types.utils import (
|
||||
AUTOROUTER_CLASSIFIER_CALL_ORIGIN,
|
||||
|
|
@ -172,38 +173,8 @@ def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str]
|
|||
return [*base_keywords, *deduped_custom.values()]
|
||||
|
||||
|
||||
# Metadata keys that carry only the parent request's budget reservation state. These
|
||||
# must not reach internal sub-calls (classifier, embedding): the reservation belongs to
|
||||
# the routed completion being decided on, not to the sub-call itself, and forwarding it
|
||||
# would let the sub-call's cost callback finalize the reservation, causing the routed
|
||||
# completion's callback to skip incrementing key/team budget counters.
|
||||
#
|
||||
# Note: user_api_key_auth itself is intentionally kept; it is required by
|
||||
# _filter_deployments_by_model_access_groups to scope embedding/classifier model
|
||||
# selection to the caller's authorized access groups. It is forwarded as a sanitized
|
||||
# copy with its budget_reservation sub-field removed, because the proxy cost callback
|
||||
# (_get_budget_reservation_from_metadata) falls back to reading the reservation from
|
||||
# inside the auth object when the top-level key is absent; forwarding it unsanitized
|
||||
# would re-create the exact double-finalization this stripping exists to prevent.
|
||||
_BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
|
||||
|
||||
|
||||
def _sanitize_user_api_key_auth(auth: Any) -> Any:
|
||||
if isinstance(auth, dict):
|
||||
return {k: v for k, v in auth.items() if k != "budget_reservation"}
|
||||
if getattr(auth, "budget_reservation", None) is not None and hasattr(auth, "model_copy"):
|
||||
return auth.model_copy(update={"budget_reservation": None})
|
||||
return auth
|
||||
|
||||
|
||||
def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]:
|
||||
if not metadata:
|
||||
return {}
|
||||
return {
|
||||
k: _sanitize_user_api_key_auth(v) if k == "user_api_key_auth" else v
|
||||
for k, v in metadata.items()
|
||||
if k not in _BUDGET_RESERVATION_METADATA_KEYS
|
||||
} | {INTERNAL_CALL_ORIGIN_METADATA_KEY: AUTOROUTER_CLASSIFIER_CALL_ORIGIN}
|
||||
def _classifier_call_metadata(metadata: Mapping[str, Any] | None) -> dict[str, Any]:
|
||||
return forwarded_internal_call_metadata(metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN)
|
||||
|
||||
|
||||
def _parent_session_kwargs(request_kwargs: Mapping[str, Any] | None) -> Mapping[str, Any]:
|
||||
|
|
|
|||
|
|
@ -158,7 +158,12 @@ DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5"
|
|||
class StartShadowEvalRequest(BaseModel):
|
||||
"""Start shadowing a deployment's traffic through an auto-router for comparison."""
|
||||
|
||||
api_key_id: str = Field(description="The hashed virtual key whose traffic will be shadowed")
|
||||
api_key_id: str = Field(
|
||||
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."
|
||||
)
|
||||
)
|
||||
router_name: str = Field(description="The auto-router config to shadow requests through")
|
||||
shadow_percentage: float = Field(
|
||||
ge=0.1,
|
||||
|
|
@ -213,6 +218,8 @@ class GetShadowEvalJobResponse(BaseModel):
|
|||
job_id: str
|
||||
status: ShadowEvalStatus
|
||||
router_name: str
|
||||
api_key_id: str = Field(description="The hashed virtual key whose traffic this job evaluates, and only that key's")
|
||||
team_id: str | None = Field(default=None, description="Team the shadowed key belongs to, when it has one")
|
||||
shadow_percentage: float
|
||||
request_count: int = Field(description="Total requests observed on the shadowed key since the job started")
|
||||
completed_count: int = Field(description="Verdicts written so far")
|
||||
|
|
|
|||
|
|
@ -2781,11 +2781,13 @@ RoutingDecisionCause = Literal[
|
|||
]
|
||||
|
||||
|
||||
InternalCallOrigin = Literal["autorouter_classifier"]
|
||||
InternalCallOrigin = Literal["autorouter_classifier", "shadow_eval_router", "shadow_eval_judge"]
|
||||
"""Which internal litellm feature originated a billed sub-call, so a spend log row
|
||||
records that it is not traffic the caller sent."""
|
||||
|
||||
AUTOROUTER_CLASSIFIER_CALL_ORIGIN: Final[InternalCallOrigin] = "autorouter_classifier"
|
||||
SHADOW_EVAL_ROUTER_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_router"
|
||||
SHADOW_EVAL_JUDGE_CALL_ORIGIN: Final[InternalCallOrigin] = "shadow_eval_judge"
|
||||
|
||||
|
||||
class StandardLoggingRoutingDecision(TypedDict, total=False):
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.integrations.shadow_eval_logger import (
|
||||
JUDGE_MAX_OUTPUT_TOKENS,
|
||||
SHADOW_EVAL_INTERNAL_MARKER,
|
||||
|
|
@ -15,6 +16,7 @@ from litellm.integrations.shadow_eval_logger import (
|
|||
_sample_hits,
|
||||
_unmask_preference,
|
||||
)
|
||||
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
|
||||
|
||||
|
||||
class TestSampling:
|
||||
|
|
@ -97,7 +99,7 @@ class TestJudgeOutputBudget:
|
|||
|
||||
logger = ShadowEvalLogger(router_provider=lambda: MagicMock(), prisma_provider=lambda: MagicMock())
|
||||
verdict = await logger._call_judge(
|
||||
"gpt-4o-mini", [{"role": "user", "content": "hi"}], "real text", "shadow text"
|
||||
"gpt-4o-mini", [{"role": "user", "content": "hi"}], "real text", "shadow text", {}
|
||||
)
|
||||
|
||||
assert verdict is not None
|
||||
|
|
@ -188,6 +190,7 @@ class TestCallRouterShadowForwardsParameters:
|
|||
"claude-auto",
|
||||
[{"role": "user", "content": "hi"}],
|
||||
{"temperature": 0.2, "tools": [{"type": "function"}], "max_tokens": 500},
|
||||
{},
|
||||
)
|
||||
|
||||
_, kwargs = router.acompletion.call_args
|
||||
|
|
@ -204,13 +207,16 @@ class TestCallRouterShadowForwardsParameters:
|
|||
"claude-auto",
|
||||
[{"role": "user", "content": "hi"}],
|
||||
{"stream": True, "metadata": {"user_api_key_hash": "leaked"}, "temperature": 0.5},
|
||||
{},
|
||||
)
|
||||
|
||||
_, kwargs = router.acompletion.call_args
|
||||
assert "stream" not in kwargs
|
||||
assert kwargs["temperature"] == 0.5
|
||||
# metadata must stay the logger's own internal marker, not the caller's
|
||||
assert kwargs["metadata"] == {SHADOW_EVAL_INTERNAL_MARKER: True}
|
||||
# model_parameters carries a stale copy of the parent's metadata; the shadow call
|
||||
# builds its own from the parent metadata argument instead.
|
||||
assert "leaked" not in str(kwargs["metadata"])
|
||||
assert kwargs["metadata"][SHADOW_EVAL_INTERNAL_MARKER] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -290,6 +296,7 @@ class TestStoppedJobCannotBeReactivated:
|
|||
response_obj={"choices": [{"message": {"content": "real text"}}]},
|
||||
real_model="gpt-4o",
|
||||
model_parameters={},
|
||||
parent_metadata={},
|
||||
)
|
||||
|
||||
prisma.db.litellm_shadowevaljob.update_many.assert_awaited_once()
|
||||
|
|
@ -326,12 +333,274 @@ class TestVerdictWriteAccumulatesCost:
|
|||
response_obj={"choices": [{"message": {"content": "real text"}}]},
|
||||
real_model="gpt-4o",
|
||||
model_parameters={},
|
||||
parent_metadata={},
|
||||
)
|
||||
|
||||
_, call_kwargs = prisma.db.litellm_shadowevaljob.update_many.call_args
|
||||
assert call_kwargs["data"]["cost_actual"] == {"increment": 0.05}
|
||||
|
||||
|
||||
_PARENT_METADATA = {
|
||||
"user_api_key": "hashed-key",
|
||||
"user_api_key_alias": "team-prod-key",
|
||||
"user_api_key_team_id": "team-1",
|
||||
"user_api_key_org_id": "org-1",
|
||||
"user_api_key_user_id": "user-1",
|
||||
"user_api_key_end_user_id": "end-user-1",
|
||||
"user_api_key_budget_reservation": {"reserved_cost": 2.5, "finalized": False},
|
||||
"user_api_key_auth": {"models": ["gpt-4o"], "budget_reservation": {"reserved_cost": 2.5}},
|
||||
"routing_decision": {"tier": "SIMPLE"},
|
||||
"standard_logging_object": {"id": "parent-req"},
|
||||
}
|
||||
|
||||
_IDENTITY_KEYS = (
|
||||
"user_api_key",
|
||||
"user_api_key_alias",
|
||||
"user_api_key_team_id",
|
||||
"user_api_key_org_id",
|
||||
"user_api_key_user_id",
|
||||
"user_api_key_end_user_id",
|
||||
)
|
||||
|
||||
|
||||
def _assert_attributed(metadata, expected_origin):
|
||||
"""Every assertion the proxy's cost callback depends on to bill an internal sub-call."""
|
||||
for key in _IDENTITY_KEYS:
|
||||
assert metadata[key] == _PARENT_METADATA[key], f"{key} must reach the sub-call for spend attribution"
|
||||
assert metadata[INTERNAL_CALL_ORIGIN_METADATA_KEY] == expected_origin
|
||||
# The parent completion owns this reservation. A sub-call carrying it would let the
|
||||
# sub-call's cost callback finalize it, so the parent's own callback skips
|
||||
# incrementing the key/team counters and the parent's spend is lost.
|
||||
assert "user_api_key_budget_reservation" not in metadata
|
||||
assert "budget_reservation" not in metadata["user_api_key_auth"]
|
||||
assert metadata["user_api_key_auth"]["models"] == ["gpt-4o"]
|
||||
# Per-request state of a request that already finished says nothing true about this one.
|
||||
assert "routing_decision" not in metadata
|
||||
assert "standard_logging_object" not in metadata
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestSubCallsAreAttributedToTheShadowedKey:
|
||||
"""Shadow and judge calls bill real provider spend nobody typed a prompt for.
|
||||
|
||||
Without the caller's identity on their metadata, _PROXY_track_cost_callback drops
|
||||
the spend log and no budget counter moves, so an admin-enabled eval spends against
|
||||
a customer's key invisibly and past every configured limit.
|
||||
"""
|
||||
|
||||
async def test_judge_call_carries_the_key_identity_without_its_reservation(self, monkeypatch: pytest.MonkeyPatch):
|
||||
import litellm as litellm_module
|
||||
|
||||
acompletion = AsyncMock(
|
||||
return_value={
|
||||
"choices": [{"message": {"content": '{"preference": "A", "confidence": 0.8, "reasoning": "clearer"}'}}]
|
||||
}
|
||||
)
|
||||
monkeypatch.setattr(litellm_module, "acompletion", acompletion)
|
||||
logger = ShadowEvalLogger(router_provider=lambda: MagicMock(), prisma_provider=lambda: MagicMock())
|
||||
|
||||
verdict = await logger._call_judge(
|
||||
"gpt-4o-mini", [{"role": "user", "content": "hi"}], "real text", "shadow text", _PARENT_METADATA
|
||||
)
|
||||
|
||||
assert verdict is not None
|
||||
metadata = acompletion.call_args.kwargs["metadata"]
|
||||
_assert_attributed(metadata, SHADOW_EVAL_JUDGE_CALL_ORIGIN)
|
||||
assert metadata[SHADOW_EVAL_INTERNAL_MARKER] is True
|
||||
|
||||
async def test_shadow_router_call_carries_the_key_identity_and_the_recursion_guard(self):
|
||||
router = MagicMock()
|
||||
router.acompletion = AsyncMock(return_value={"choices": [{"message": {"content": "shadow reply"}}]})
|
||||
logger = ShadowEvalLogger(router_provider=lambda: router, prisma_provider=lambda: MagicMock())
|
||||
|
||||
result = await logger._call_router_shadow(
|
||||
"claude-auto", [{"role": "user", "content": "hi"}], {}, _PARENT_METADATA
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
metadata = router.acompletion.call_args.kwargs["metadata"]
|
||||
_assert_attributed(metadata, SHADOW_EVAL_ROUTER_CALL_ORIGIN)
|
||||
# Without the marker the shadow call's own success event would shadow itself.
|
||||
assert metadata[SHADOW_EVAL_INTERNAL_MARKER] is True
|
||||
|
||||
async def test_shadow_metadata_stays_writable_for_the_routing_decision_read_back(self):
|
||||
"""Tier attribution reads routing_decision back out of the dict the router was given."""
|
||||
router = MagicMock()
|
||||
|
||||
async def acompletion(**kwargs):
|
||||
kwargs["metadata"]["routing_decision"] = {"tier_label": "COMPLEX", "routed_model": "o1"}
|
||||
return {"choices": [{"message": {"content": "shadow reply"}}]}
|
||||
|
||||
router.acompletion = acompletion
|
||||
logger = ShadowEvalLogger(router_provider=lambda: router, prisma_provider=lambda: MagicMock())
|
||||
|
||||
result = await logger._call_router_shadow(
|
||||
"claude-auto", [{"role": "user", "content": "hi"}], {}, _PARENT_METADATA
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result[2] == "COMPLEX"
|
||||
|
||||
async def test_parent_metadata_reaches_both_legs_from_the_success_hook(self):
|
||||
"""The hook is the only place the parent's metadata exists; a break here is invisible
|
||||
to unit tests of the two legs in isolation."""
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_shadowevalverdict.create = AsyncMock()
|
||||
prisma.db.litellm_shadowevaljob.update_many = AsyncMock()
|
||||
job = ActiveShadowEvalJob(id="j1", router_name="r", shadow_percentage=100.0, judge_model="m", status="running")
|
||||
logger, _, _ = _logger_with_mocks(job)
|
||||
logger._prisma_provider = lambda: prisma
|
||||
logger._call_router_shadow = AsyncMock(return_value=("shadow text", "shadow-model", "SIMPLE", 10))
|
||||
logger._call_judge = AsyncMock(return_value=("real", 0.9, "clearer", 0.01))
|
||||
|
||||
await logger.async_log_success_event(
|
||||
{
|
||||
"standard_logging_object": {
|
||||
"id": "req-1",
|
||||
"model": "gpt-4o",
|
||||
"call_type": "acompletion",
|
||||
"metadata": {"user_api_key_hash": "key-hash"},
|
||||
},
|
||||
"litellm_params": {"metadata": dict(_PARENT_METADATA)},
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
},
|
||||
{"choices": [{"message": {"content": "real text"}}]},
|
||||
None,
|
||||
None,
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
assert logger._call_router_shadow.await_args.args[3]["user_api_key"] == "hashed-key"
|
||||
assert logger._call_judge.await_args.kwargs["parent_metadata"]["user_api_key"] == "hashed-key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestPerJobSpendCap:
|
||||
"""Budgets bound the key; the cap bounds a single eval, so a bad estimate or a
|
||||
traffic spike cannot turn a quoted eval into a much larger bill."""
|
||||
|
||||
@staticmethod
|
||||
def _success_kwargs():
|
||||
return {
|
||||
"standard_logging_object": {
|
||||
"id": "req-1",
|
||||
"model": "gpt-4o",
|
||||
"call_type": "acompletion",
|
||||
"metadata": {"user_api_key_hash": "key-hash"},
|
||||
},
|
||||
"litellm_params": {"metadata": {}},
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
|
||||
async def test_job_over_the_cap_stops_sampling_and_completes_the_job(self):
|
||||
job = ActiveShadowEvalJob(
|
||||
id="j1",
|
||||
router_name="r",
|
||||
shadow_percentage=100.0,
|
||||
judge_model="m",
|
||||
status="running",
|
||||
cost_estimate=10.0,
|
||||
cost_actual=15.0,
|
||||
)
|
||||
logger, prisma, router = _logger_with_mocks(job)
|
||||
prisma.db.litellm_shadowevaljob.update_many = AsyncMock()
|
||||
logger._run_shadow_eval = AsyncMock()
|
||||
|
||||
await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
logger._run_shadow_eval.assert_not_awaited()
|
||||
router.acompletion.assert_not_called()
|
||||
prisma.db.litellm_shadowevaljob.update_many.assert_awaited_once()
|
||||
call_kwargs = prisma.db.litellm_shadowevaljob.update_many.call_args.kwargs
|
||||
assert call_kwargs["where"]["id"] == "j1"
|
||||
# Guarded so a pod breaching the cap cannot resurrect a job an admin already stopped.
|
||||
assert set(call_kwargs["where"]["status"]["in"]) == {"pending", "running"}
|
||||
assert call_kwargs["data"]["status"] == "completed"
|
||||
assert call_kwargs["data"]["completed_at"] is not None
|
||||
|
||||
async def test_job_just_under_the_cap_keeps_evaluating(self):
|
||||
job = ActiveShadowEvalJob(
|
||||
id="j1",
|
||||
router_name="r",
|
||||
shadow_percentage=100.0,
|
||||
judge_model="m",
|
||||
status="running",
|
||||
cost_estimate=10.0,
|
||||
cost_actual=14.99,
|
||||
)
|
||||
logger, prisma, _ = _logger_with_mocks(job)
|
||||
prisma.db.litellm_shadowevaljob.update_many = AsyncMock()
|
||||
logger._run_shadow_eval = AsyncMock()
|
||||
|
||||
await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
logger._run_shadow_eval.assert_awaited_once()
|
||||
prisma.db.litellm_shadowevaljob.update_many.assert_not_awaited()
|
||||
|
||||
async def test_job_without_an_estimate_is_uncapped(self):
|
||||
"""A missing estimate is no multiple to compare against; treating it as 0 would
|
||||
stop every such job on its first request instead."""
|
||||
job = ActiveShadowEvalJob(
|
||||
id="j1",
|
||||
router_name="r",
|
||||
shadow_percentage=100.0,
|
||||
judge_model="m",
|
||||
status="running",
|
||||
cost_estimate=None,
|
||||
cost_actual=500.0,
|
||||
)
|
||||
logger, prisma, _ = _logger_with_mocks(job)
|
||||
prisma.db.litellm_shadowevaljob.update_many = AsyncMock()
|
||||
logger._run_shadow_eval = AsyncMock()
|
||||
|
||||
await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
logger._run_shadow_eval.assert_awaited_once()
|
||||
prisma.db.litellm_shadowevaljob.update_many.assert_not_awaited()
|
||||
|
||||
async def test_a_cent_sized_estimate_is_not_stopped_by_its_first_verdict(self):
|
||||
"""1.5x a $0.01 estimate is $0.015; without the floor a single judge call ends the job."""
|
||||
job = ActiveShadowEvalJob(
|
||||
id="j1",
|
||||
router_name="r",
|
||||
shadow_percentage=100.0,
|
||||
judge_model="m",
|
||||
status="running",
|
||||
cost_estimate=0.01,
|
||||
cost_actual=0.08,
|
||||
)
|
||||
logger, prisma, _ = _logger_with_mocks(job)
|
||||
prisma.db.litellm_shadowevaljob.update_many = AsyncMock()
|
||||
logger._run_shadow_eval = AsyncMock()
|
||||
|
||||
await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None)
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
logger._run_shadow_eval.assert_awaited_once()
|
||||
|
||||
async def test_stopping_evicts_the_cached_job_so_later_requests_do_not_rewrite_it(self):
|
||||
job = ActiveShadowEvalJob(
|
||||
id="j1",
|
||||
router_name="r",
|
||||
shadow_percentage=100.0,
|
||||
judge_model="m",
|
||||
status="running",
|
||||
cost_estimate=10.0,
|
||||
cost_actual=15.0,
|
||||
)
|
||||
logger, prisma, _ = _logger_with_mocks(job)
|
||||
prisma.db.litellm_shadowevaljob.update_many = AsyncMock()
|
||||
prisma.db.litellm_shadowevaljob.find_first = AsyncMock(return_value=None)
|
||||
|
||||
await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None)
|
||||
await logger.async_log_success_event(self._success_kwargs(), MagicMock(), None, None)
|
||||
|
||||
assert prisma.db.litellm_shadowevaljob.update_many.await_count == 1
|
||||
|
||||
|
||||
class TestExtractResponseText:
|
||||
def test_dict_response(self):
|
||||
resp = {"choices": [{"message": {"content": "hello"}}]}
|
||||
|
|
|
|||
|
|
@ -466,6 +466,49 @@ class TestAutoRouterBenchmarks:
|
|||
assert response.groups[0].tier_turns == expected
|
||||
|
||||
|
||||
class TestShadowEvalJobResponseNamesWhatIsBeingEvaluated:
|
||||
"""A shadow eval only samples one key's traffic. An admin running several jobs cannot
|
||||
tell which key a win rate belongs to unless the job response says so."""
|
||||
|
||||
@staticmethod
|
||||
def _record(**overrides: object):
|
||||
from datetime import datetime, timezone
|
||||
|
||||
fields = {
|
||||
"id": "job-1",
|
||||
"status": "running",
|
||||
"router_name": "claude-auto",
|
||||
"api_key_id": "hashed-key-abc",
|
||||
"team_id": "team-7",
|
||||
"shadow_percentage": 10.0,
|
||||
"request_count": 100,
|
||||
"completed_count": 9,
|
||||
"failed_count": 1,
|
||||
"cost_estimate": 3.0,
|
||||
"cost_actual": 0.42,
|
||||
"created_at": datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
"completed_at": None,
|
||||
**overrides,
|
||||
}
|
||||
return type("Row", (), fields)()
|
||||
|
||||
def test_response_reports_the_shadowed_key_and_its_team(self):
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import _job_to_response
|
||||
|
||||
response = _job_to_response(self._record(), None)
|
||||
|
||||
assert response.api_key_id == "hashed-key-abc"
|
||||
assert response.team_id == "team-7"
|
||||
|
||||
def test_a_keyless_team_still_produces_a_response(self):
|
||||
from litellm.proxy.management_endpoints.auto_router_endpoints import _job_to_response
|
||||
|
||||
response = _job_to_response(self._record(team_id=None), None)
|
||||
|
||||
assert response.api_key_id == "hashed-key-abc"
|
||||
assert response.team_id is None
|
||||
|
||||
|
||||
class TestJudgeCostEstimate:
|
||||
"""The upfront estimate must price the judge's real output budget, not a stale hardcoded one."""
|
||||
|
||||
|
|
|
|||
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
12
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -25329,6 +25329,11 @@ export interface components {
|
|||
* @description Status and, once available, results of a shadow-eval job.
|
||||
*/
|
||||
GetShadowEvalJobResponse: {
|
||||
/**
|
||||
* Api Key Id
|
||||
* @description The hashed virtual key whose traffic this job evaluates, and only that key's
|
||||
*/
|
||||
api_key_id: string;
|
||||
/** Completed At */
|
||||
completed_at?: string | null;
|
||||
/**
|
||||
|
|
@ -25369,6 +25374,11 @@ export interface components {
|
|||
* @enum {string}
|
||||
*/
|
||||
status: "pending" | "running" | "completed" | "failed";
|
||||
/**
|
||||
* Team Id
|
||||
* @description Team the shadowed key belongs to, when it has one
|
||||
*/
|
||||
team_id?: string | null;
|
||||
};
|
||||
/**
|
||||
* GetTeamMemberPermissionsResponse
|
||||
|
|
@ -32768,7 +32778,7 @@ export interface components {
|
|||
StartShadowEvalRequest: {
|
||||
/**
|
||||
* Api Key Id
|
||||
* @description The hashed virtual key whose traffic will be shadowed
|
||||
* @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_id: string;
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue