mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_decrease_anys_opus5_r2
# Conflicts: # basedpyright-code-budget.json # type-discipline-budget.json
This commit is contained in:
commit
3a750cbf92
23 changed files with 2884 additions and 478 deletions
|
|
@ -99,7 +99,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44368
|
||||
"limit": 44364
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
|
|
|
|||
|
|
@ -0,0 +1,21 @@
|
|||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (
|
||||
SELECT 1 FROM information_schema.columns
|
||||
WHERE table_name = 'LiteLLM_ShadowEvalJob' AND column_name = 'api_key_id'
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" RENAME COLUMN "api_key_id" TO "target_id";
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "target_type" TEXT NOT NULL DEFAULT 'key';
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_one_active_per_key_direction";
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_one_active_per_target_direction"
|
||||
ON "LiteLLM_ShadowEvalJob"("target_type", "target_id", "direction") WHERE "stopped_at" IS NULL;
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_api_key_id_idx";
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_target_type_target_id_idx"
|
||||
ON "LiteLLM_ShadowEvalJob"("target_type", "target_id");
|
||||
|
|
@ -1529,14 +1529,15 @@ model LiteLLM_AutoRouterSession {
|
|||
model LiteLLM_ShadowEvalJob {
|
||||
id String @id @default(cuid())
|
||||
group_id String // legs of one job share this; the API's job id
|
||||
api_key_id String // hashed virtual key whose traffic this leg shadows
|
||||
target_type String @default("key") // key | team | user
|
||||
target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows
|
||||
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-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise
|
||||
max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
max_budget Float? // per-target USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
ends_at DateTime
|
||||
|
|
@ -1544,7 +1545,7 @@ model LiteLLM_ShadowEvalJob {
|
|||
stopped_by String? // operator who stopped it early; null when it ended on its own
|
||||
|
||||
@@index([group_id])
|
||||
@@index([api_key_id])
|
||||
@@index([target_type, target_id])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -484,6 +484,22 @@ FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80))
|
|||
#### Logging callback constants ####
|
||||
REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM"
|
||||
MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50))
|
||||
# Backpressure + lifetime bounds for the /v1/messages streaming relay (see
|
||||
# BaseAnthropicMessagesStreamingIterator.async_sse_wrapper). The relay queue is
|
||||
# bounded so a slow client throttles the upstream pump instead of letting it
|
||||
# buffer the whole response in memory; the detached-drain cap bounds how many
|
||||
# post-disconnect drains may run concurrently so client behavior can't create
|
||||
# unbounded worker state.
|
||||
ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int(
|
||||
os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024")
|
||||
)
|
||||
# Setting this to 0 disables detached draining entirely: every post-disconnect
|
||||
# pump bills whatever partial output it has already collected and aborts the
|
||||
# upstream stream immediately, instead of continuing to drain for the real
|
||||
# terminal usage.
|
||||
ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int(
|
||||
os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100")
|
||||
)
|
||||
LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0
|
||||
LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000))
|
||||
LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0))
|
||||
|
|
|
|||
|
|
@ -592,7 +592,12 @@ _JOBS_CACHE_KEY: Final = "shadow_eval:active_jobs"
|
|||
|
||||
|
||||
class ShadowEvalLogger(CustomLogger):
|
||||
"""Fires blind pairwise shadow evaluations for keys with an active shadow-eval job."""
|
||||
"""Fires blind pairwise shadow evaluations for targets with an active shadow-eval job.
|
||||
|
||||
A job targets a virtual key, a team, or a user; a request qualifies for a job when
|
||||
any of its resolved identities (key hash, team id, user id) matches the job's
|
||||
target, so team and user jobs cover JWT-authenticated traffic, which carries no
|
||||
key hash at all."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -617,10 +622,10 @@ class ShadowEvalLogger(CustomLogger):
|
|||
# generation; the refill absorbs written rows and resets.
|
||||
self._job_starts: dict[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."""
|
||||
async def _active_jobs(self) -> Mapping[tuple[str, str], tuple[ActiveShadowEvalJob, ...]]:
|
||||
"""Active jobs by (target_type, target_id), cache-first. A target 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."""
|
||||
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
|
||||
|
|
@ -652,10 +657,10 @@ class ShadowEvalLogger(CustomLogger):
|
|||
)
|
||||
for row in grouped or []
|
||||
}
|
||||
by_key: Final = tuple(
|
||||
by_target: Final = tuple(
|
||||
sorted(
|
||||
(
|
||||
(str(record.api_key_id), job)
|
||||
((str(record.target_type), str(record.target_id)), job)
|
||||
for record in records or []
|
||||
if (job := _as_active_job(record, *attempt_stats.get(str(record.id), (0, 0.0)))) is not None
|
||||
),
|
||||
|
|
@ -663,7 +668,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
)
|
||||
)
|
||||
jobs: Final = MappingProxyType(
|
||||
{key: tuple(job for _, job in group) for key, group in groupby(by_key, key=itemgetter(0))}
|
||||
{target: tuple(job for _, job in group) for target, group in groupby(by_target, key=itemgetter(0))}
|
||||
)
|
||||
await self._jobs_cache.async_set_cache(_JOBS_CACHE_KEY, jobs)
|
||||
self._job_starts = {} # rebind-ok: new generation, counts absorbed into the fill
|
||||
|
|
@ -720,8 +725,18 @@ class ShadowEvalLogger(CustomLogger):
|
|||
if should_redact_message_logging(dict(kwargs)): # mutable-ok: predicate takes a plain dict
|
||||
return
|
||||
metadata: Final = payload.get("metadata") or _EMPTY_METADATA
|
||||
api_key_hash: Final = metadata.get("user_api_key_hash")
|
||||
if not api_key_hash:
|
||||
# Each identity the request resolved to is a candidate target; JWT-auth
|
||||
# requests carry no key hash but do carry a team and user.
|
||||
targets: Final = tuple(
|
||||
(target_type, str(value))
|
||||
for target_type, value in (
|
||||
("key", metadata.get("user_api_key_hash")),
|
||||
("team", metadata.get("user_api_key_team_id")),
|
||||
("user", metadata.get("user_api_key_user_id")),
|
||||
)
|
||||
if value
|
||||
)
|
||||
if not targets:
|
||||
return
|
||||
request_id: Final = payload.get("id") or ""
|
||||
if not request_id:
|
||||
|
|
@ -731,8 +746,11 @@ class ShadowEvalLogger(CustomLogger):
|
|||
return # only surfaces this table can normalize are comparable; unknown types fail closed
|
||||
if ops.wire_params and _request_mutating_guardrail_ran(request_metadata):
|
||||
return # the wire-body snapshot predates the rewrite; replaying it would resurrect stripped content
|
||||
active_jobs: Final = await self._active_jobs()
|
||||
eligible: Final = self._sampled_jobs(
|
||||
(await self._active_jobs()).get(str(api_key_hash), ()), request_metadata, request_id
|
||||
tuple(job for target in targets for job in active_jobs.get(target, ())),
|
||||
request_metadata,
|
||||
request_id,
|
||||
)
|
||||
if not eligible:
|
||||
return
|
||||
|
|
@ -1056,7 +1074,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
)
|
||||
|
||||
|
||||
_EMPTY_JOBS: Final[Mapping[str, tuple[ActiveShadowEvalJob, ...]]] = MappingProxyType({})
|
||||
_EMPTY_JOBS: Final[Mapping[tuple[str, str], tuple[ActiveShadowEvalJob, ...]]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _default_prisma_provider() -> "PrismaClient | None":
|
||||
|
|
|
|||
|
|
@ -8,6 +8,10 @@ import httpx
|
|||
from pydantic import TypeAdapter
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from litellm.constants import (
|
||||
ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS,
|
||||
ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
|
@ -21,6 +25,9 @@ from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
|
|||
|
||||
GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging()
|
||||
|
||||
_UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks
|
||||
_DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains
|
||||
|
||||
INCOMPLETE_STREAM_ERROR_MESSAGE: Final = (
|
||||
"Provider stream ended before emitting a message_stop event; "
|
||||
"the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated."
|
||||
|
|
@ -133,6 +140,34 @@ def _is_terminal_stream_chunk(chunk: object) -> bool:
|
|||
return _is_message_stop_chunk(chunk) or _is_provider_error_chunk(chunk)
|
||||
|
||||
|
||||
def _try_claim_detached_drain_slot() -> bool:
|
||||
"""Claim a detached-drain slot for the current task, bounding concurrency.
|
||||
|
||||
Returns True if a slot was claimed (the caller may keep draining upstream
|
||||
for billing) or False if the cap is already reached (the caller should stop
|
||||
and bill what it has). Only touched from the event loop, so the check +
|
||||
insert need no lock.
|
||||
"""
|
||||
if len(_DETACHED_STREAM_DRAINS) >= ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS:
|
||||
return False
|
||||
current_task: Final = asyncio.current_task()
|
||||
if current_task is not None:
|
||||
_DETACHED_STREAM_DRAINS.add(current_task)
|
||||
current_task.add_done_callback(_DETACHED_STREAM_DRAINS.discard)
|
||||
return True
|
||||
|
||||
|
||||
def _exception_left_unconsumed(queue: "asyncio.Queue[bytes | None | BaseException]", exc: BaseException) -> bool:
|
||||
"""After client detach the relay never reads the queue again, so drain it here.
|
||||
|
||||
The forwarded exception still sitting in the queue means the relay tore
|
||||
down before re-raising it, so the proxy's failure handling never ran and
|
||||
the caller must salvage spend itself.
|
||||
"""
|
||||
remaining: Final = tuple(queue.get_nowait() for _ in range(queue.qsize()))
|
||||
return any(item is exc for item in remaining)
|
||||
|
||||
|
||||
def _sse_event(event_type: str, payload: Mapping[str, object]) -> bytes:
|
||||
return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
|
@ -414,17 +449,167 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
|
||||
async def async_sse_wrapper(
|
||||
self,
|
||||
completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict],
|
||||
completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]],
|
||||
) -> AsyncIterator[bytes]:
|
||||
"""
|
||||
Generic async SSE wrapper that converts streaming chunks to SSE format
|
||||
and handles logging.
|
||||
|
||||
The upstream read runs in a detached background task (``_pump_upstream``)
|
||||
so that a client disconnect tears down only this client-facing generator,
|
||||
never the upstream drain + billing. The provider (e.g. Bedrock) keeps
|
||||
generating and billing the full response regardless of the client, so
|
||||
draining it to completion is what lets spend tracking see the real
|
||||
terminal ``message_delta`` / ``message_stop`` usage instead of a
|
||||
truncated placeholder count.
|
||||
|
||||
Chunks reach the client through a bounded queue. While the client is
|
||||
connected the pump blocks on a full queue (racing the disconnect
|
||||
signal), so a slow reader throttles the upstream read exactly as the old
|
||||
direct ``yield`` did instead of letting the whole response buffer in
|
||||
memory. Once the client goes away the pump stops enqueueing and only
|
||||
keeps a single ``collected_chunks`` copy for billing, and the number of
|
||||
such post-disconnect drains running at once is capped so client behavior
|
||||
can't create unbounded worker state; over the cap the pump bills what it
|
||||
has rather than draining further. Detached-drain lifetime is otherwise
|
||||
bounded by the upstream stream/read timeout.
|
||||
|
||||
An upstream failure (Bedrock read / decode / chunk-conversion error)
|
||||
that happens while the client is still connected is forwarded through
|
||||
the queue and re-raised here, so the original provider exception (and
|
||||
its status) reaches the proxy's failure handling unchanged rather than
|
||||
being masked by a generic incomplete-stream event.
|
||||
|
||||
This method provides the common logic for both Anthropic and Bedrock implementations.
|
||||
"""
|
||||
collected_chunks: Final = []
|
||||
saw_terminal_event = False
|
||||
queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue(
|
||||
maxsize=ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE
|
||||
)
|
||||
client_detached: Final = asyncio.Event()
|
||||
|
||||
pump_task: Final = asyncio.create_task(self._pump_upstream_to_queue(completion_stream, queue, client_detached))
|
||||
_UPSTREAM_PUMP_TASKS.add(pump_task)
|
||||
pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard)
|
||||
|
||||
reached_end = False # rebind-ok: flipped once the relay consumes the end-of-stream sentinel
|
||||
try:
|
||||
while True:
|
||||
item = await queue.get()
|
||||
if item is None:
|
||||
reached_end = True
|
||||
break
|
||||
if isinstance(item, BaseException):
|
||||
raise item
|
||||
yield item
|
||||
finally:
|
||||
client_detached.set()
|
||||
if not reached_end:
|
||||
self._dispatch_pending_deferred_logging()
|
||||
|
||||
def _dispatch_pending_deferred_logging(self) -> None:
|
||||
"""Fire deferred billing that a torn-down response would otherwise drop.
|
||||
|
||||
When the pump finishes draining while the client is still connected it
|
||||
stores the logging coroutine for ProxyLogging._fire_deferred_stream_logging,
|
||||
which the proxy only fires on a normally completed response: a client
|
||||
disconnect (GeneratorExit / CancelledError) re-raises past it. Without
|
||||
this dispatch that window loses the spend row entirely.
|
||||
"""
|
||||
deferred_cb: Final = getattr(self.litellm_logging_obj, "_on_deferred_stream_complete", None)
|
||||
deferred_args: Final = getattr(self.litellm_logging_obj, "_deferred_stream_complete_args", None)
|
||||
if deferred_cb is None or deferred_args is None:
|
||||
return
|
||||
self.litellm_logging_obj._on_deferred_stream_complete = None
|
||||
self.litellm_logging_obj._deferred_stream_complete_args = None
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=deferred_cb(*deferred_args))
|
||||
|
||||
async def _bill_collected_chunks(
|
||||
self,
|
||||
collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _handle_streaming_logging
|
||||
*,
|
||||
stream_teardown: bool,
|
||||
) -> None:
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
try:
|
||||
await self._handle_streaming_logging(collected_chunks, stream_teardown=stream_teardown)
|
||||
except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump
|
||||
verbose_proxy_logger.warning(
|
||||
"async_sse_wrapper billing failed after %d chunks: %s(%s)",
|
||||
len(collected_chunks),
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _abort_upstream(
|
||||
completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]],
|
||||
) -> None:
|
||||
"""Close the upstream provider stream so it stops generating and billing."""
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
try:
|
||||
await aclose_if_supported(completion_stream)
|
||||
except Exception as exc: # noqa: BLE001 # abort is best-effort; log and continue
|
||||
verbose_proxy_logger.warning(
|
||||
"async_sse_wrapper failed to abort upstream stream: %s(%s)",
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _enqueue_for_client(
|
||||
queue: "asyncio.Queue[bytes | None | BaseException]",
|
||||
client_detached: "asyncio.Event",
|
||||
item: bytes | None | BaseException,
|
||||
) -> bool:
|
||||
"""Deliver one item to the client, applying backpressure.
|
||||
|
||||
Returns True if the item was queued, False if the client disconnected
|
||||
before there was room (the item is then dropped, since a gone client
|
||||
can't receive it). Never blocks once the client has detached.
|
||||
"""
|
||||
if client_detached.is_set():
|
||||
return False
|
||||
try:
|
||||
queue.put_nowait(item)
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
else:
|
||||
return True
|
||||
put_task: Final = asyncio.ensure_future(queue.put(item))
|
||||
detached_task: Final = asyncio.ensure_future(client_detached.wait())
|
||||
try:
|
||||
await asyncio.wait(frozenset((put_task, detached_task)), return_when=asyncio.FIRST_COMPLETED)
|
||||
finally:
|
||||
if not detached_task.done():
|
||||
detached_task.cancel()
|
||||
if put_task.done() and not put_task.cancelled():
|
||||
return True
|
||||
put_task.cancel()
|
||||
return False
|
||||
|
||||
async def _pump_upstream_to_queue(
|
||||
self,
|
||||
completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | Mapping[str, object]],
|
||||
queue: "asyncio.Queue[bytes | None | BaseException]",
|
||||
client_detached: "asyncio.Event",
|
||||
) -> None:
|
||||
"""Drain the whole upstream into ``queue`` (backpressured) and bill once.
|
||||
|
||||
Runs detached so a client disconnect can't interrupt the upstream read;
|
||||
see ``async_sse_wrapper`` for the full rationale. On a completed drain
|
||||
the success billing (or deferred park) happens before the end-of-stream
|
||||
sentinel is enqueued: the relay can only tear down after consuming the
|
||||
sentinel, so its teardown can never outrun the park and get mistaken
|
||||
for a client disconnect, and a sentinel the client never consumes falls
|
||||
back to dispatching the parked billing here.
|
||||
"""
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
collected_chunks: Final[list[bytes]] = [] # mutable-ok: SSE billing buffer appended to across the drain
|
||||
saw_terminal_event = False # rebind-ok: accumulates across the upstream loop
|
||||
draining_detached = False # rebind-ok: set once this pump claims a detached-drain slot
|
||||
try:
|
||||
async for chunk in completion_stream:
|
||||
if self.completion_start_time is None:
|
||||
|
|
@ -432,17 +617,62 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk)
|
||||
encoded_chunk = self._convert_chunk_to_sse_format(chunk)
|
||||
collected_chunks.append(encoded_chunk)
|
||||
yield encoded_chunk
|
||||
except (GeneratorExit, asyncio.CancelledError):
|
||||
# A client disconnect tears the generator down at the yield, so the
|
||||
# post-loop logging below never runs and the tokens already streamed
|
||||
# (and billed by the provider) would never reach spend tracking. See LIT-5839.
|
||||
if collected_chunks:
|
||||
await self._handle_streaming_logging(collected_chunks, stream_teardown=True)
|
||||
raise
|
||||
if not client_detached.is_set():
|
||||
await self._enqueue_for_client(queue, client_detached, encoded_chunk)
|
||||
continue
|
||||
if not draining_detached:
|
||||
if not _try_claim_detached_drain_slot():
|
||||
verbose_proxy_logger.warning(
|
||||
"async_sse_wrapper: detached-drain cap (%d) reached; billing %d partial "
|
||||
"chunks and aborting the upstream stream to stop provider billing",
|
||||
ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS,
|
||||
len(collected_chunks),
|
||||
)
|
||||
await self._bill_collected_chunks(collected_chunks, stream_teardown=True)
|
||||
await self._abort_upstream(completion_stream)
|
||||
return
|
||||
draining_detached = True
|
||||
except Exception as exc: # noqa: BLE001 # upstream errors are handled/forwarded by _handle_pump_upstream_error
|
||||
await self._handle_pump_upstream_error(queue, client_detached, collected_chunks, exc)
|
||||
return
|
||||
|
||||
if not saw_terminal_event:
|
||||
yield _incomplete_stream_error_sse_event()
|
||||
if client_detached.is_set():
|
||||
await self._bill_collected_chunks(collected_chunks, stream_teardown=True)
|
||||
return
|
||||
if not saw_terminal_event and not await self._enqueue_for_client(
|
||||
queue, client_detached, _incomplete_stream_error_sse_event()
|
||||
):
|
||||
await self._bill_collected_chunks(collected_chunks, stream_teardown=True)
|
||||
return
|
||||
await self._bill_collected_chunks(collected_chunks, stream_teardown=False)
|
||||
if not await self._enqueue_for_client(queue, client_detached, None):
|
||||
self._dispatch_pending_deferred_logging()
|
||||
|
||||
# Handle logging after all chunks are processed
|
||||
await self._handle_streaming_logging(collected_chunks)
|
||||
async def _handle_pump_upstream_error(
|
||||
self,
|
||||
queue: "asyncio.Queue[bytes | None | BaseException]",
|
||||
client_detached: "asyncio.Event",
|
||||
collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks
|
||||
exc: BaseException,
|
||||
) -> None:
|
||||
"""Forward a provider error to a still-connected client, else salvage partial spend.
|
||||
|
||||
Handing the original exception to the client-facing generator lets it
|
||||
re-raise so the proxy's failure handling keeps the provider status and
|
||||
owns logging (no success-bill). If the client already went away, or
|
||||
disconnects before ever consuming the queued exception, no failure hook
|
||||
runs, so bill the partial instead of dropping the request.
|
||||
"""
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc):
|
||||
await client_detached.wait()
|
||||
if not _exception_left_unconsumed(queue, exc):
|
||||
return
|
||||
verbose_proxy_logger.warning(
|
||||
"async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)",
|
||||
len(collected_chunks),
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
await self._bill_collected_chunks(collected_chunks, stream_teardown=True)
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ POST /auto_router/validate_complexity_router_config - Dry-run the complexity-rou
|
|||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from itertools import groupby
|
||||
from itertools import chain, groupby
|
||||
from operator import attrgetter
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Protocol
|
||||
|
|
@ -58,10 +58,11 @@ from litellm.types.management_endpoints.auto_router_endpoints import (
|
|||
ComplexityRouterConfigValidationResponse,
|
||||
RequestComplexityRouterConfig,
|
||||
ShadowEvalDirection,
|
||||
ShadowEvalJobKeyResponse,
|
||||
ShadowEvalJobResponse,
|
||||
ShadowEvalJobTargetResponse,
|
||||
ShadowEvalResult,
|
||||
ShadowEvalSlice,
|
||||
ShadowEvalTargetType,
|
||||
StartShadowEvalRequest,
|
||||
)
|
||||
|
||||
|
|
@ -104,10 +105,43 @@ class _VerificationTokenTable(Protocol):
|
|||
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_VerificationTokenRow]: ...
|
||||
|
||||
|
||||
class _TeamRow(Protocol):
|
||||
@property
|
||||
def team_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def team_alias(self) -> str | None: ...
|
||||
|
||||
|
||||
class _TeamRowsTable(Protocol):
|
||||
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_TeamRow]: ...
|
||||
|
||||
|
||||
class _UserRow(Protocol):
|
||||
@property
|
||||
def user_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def user_email(self) -> str | None: ...
|
||||
|
||||
|
||||
class _UserRowsTable(Protocol):
|
||||
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_UserRow]: ...
|
||||
|
||||
|
||||
class _ShadowEvalJobRow(Protocol):
|
||||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def group_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def target_type(self) -> str: ...
|
||||
|
||||
@property
|
||||
def target_id(self) -> str: ...
|
||||
|
||||
|
||||
class _ShadowEvalJobTable(Protocol):
|
||||
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_ShadowEvalJobRow]: ...
|
||||
|
|
@ -138,6 +172,14 @@ def _verification_tokens(prisma_client: "PrismaClient") -> _VerificationTokenTab
|
|||
return prisma_client.db.litellm_verificationtoken
|
||||
|
||||
|
||||
def _team_rows(prisma_client: "PrismaClient") -> _TeamRowsTable:
|
||||
return prisma_client.db.litellm_teamtable
|
||||
|
||||
|
||||
def _user_rows(prisma_client: "PrismaClient") -> _UserRowsTable:
|
||||
return prisma_client.db.litellm_usertable
|
||||
|
||||
|
||||
def _shadow_eval_jobs(prisma_client: "PrismaClient") -> _ShadowEvalJobTable:
|
||||
return prisma_client.db.litellm_shadowevaljob
|
||||
|
||||
|
|
@ -836,7 +878,7 @@ def _validate_judge_is_not_a_candidate(
|
|||
|
||||
|
||||
def _is_unique_violation(error: Exception) -> bool:
|
||||
"""Whether a Prisma create failed on a unique index. One active job per key and
|
||||
"""Whether a Prisma create failed on a unique index. One active job per target and
|
||||
direction lives in a partial unique index (raw SQL in the migration; schema.prisma
|
||||
cannot express partial indexes), so the read-then-create check above it is advisory:
|
||||
two concurrent starts pass the read, and the loser must surface as the same 409
|
||||
|
|
@ -885,7 +927,7 @@ _ATTEMPT_AGG_BY_LEG_SQL: Final = "SELECT job_id AS grp," + _ATTEMPT_AGG_SELECT
|
|||
# direction, and mid-deploy rows from old pods price as judge-only until the deploy ends).
|
||||
_SWEEP_FINISHED_JOBS_SQL: Final = """
|
||||
UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = (NOW() AT TIME ZONE 'utc')
|
||||
WHERE j.api_key_id = ANY($1::text[]) AND j.stopped_at IS NULL
|
||||
WHERE j.target_type = $2 AND j.target_id = ANY($1::text[]) AND j.stopped_at IS NULL
|
||||
AND (
|
||||
j.ends_at <= (NOW() AT TIME ZONE 'utc')
|
||||
OR (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = j.id) >= j.max_turns
|
||||
|
|
@ -966,10 +1008,10 @@ WHERE group_id IN (
|
|||
)
|
||||
"""
|
||||
|
||||
_LIST_LEGS_BY_KEY_SQL: Final = """
|
||||
_LIST_LEGS_BY_TARGET_SQL: Final = """
|
||||
SELECT * FROM "LiteLLM_ShadowEvalJob"
|
||||
WHERE group_id IN (
|
||||
SELECT group_id FROM "LiteLLM_ShadowEvalJob" WHERE api_key_id = $2
|
||||
SELECT group_id FROM "LiteLLM_ShadowEvalJob" WHERE target_type = $2 AND target_id = $3
|
||||
GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1::int
|
||||
)
|
||||
"""
|
||||
|
|
@ -1007,15 +1049,16 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]:
|
|||
|
||||
class _LegRow(BaseModel):
|
||||
"""One LiteLLM_ShadowEvalJob row, validated off the untyped prisma record. A row is
|
||||
one key's leg of a job; the legs of a job share group_id and identical config, written
|
||||
together by one create_many. The API's job id is the group id, so leg ids never leave
|
||||
the server (attempts reference them internally)."""
|
||||
one target's leg of a job; the legs of a job share group_id and identical config,
|
||||
written together by one create_many. The API's job id is the group id, so leg ids
|
||||
never leave the server (attempts reference them internally)."""
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
id: str
|
||||
group_id: str
|
||||
api_key_id: str
|
||||
target_type: ShadowEvalTargetType
|
||||
target_id: str
|
||||
router_name: str
|
||||
direction: ShadowEvalDirection
|
||||
baseline_model: str | None = None
|
||||
|
|
@ -1068,16 +1111,17 @@ def _group_response(
|
|||
first: Final = legs[0]
|
||||
return ShadowEvalJobResponse(
|
||||
job_id=group_id,
|
||||
keys=tuple(
|
||||
ShadowEvalJobKeyResponse(
|
||||
api_key_id=leg.api_key_id,
|
||||
targets=tuple(
|
||||
ShadowEvalJobTargetResponse(
|
||||
target_type=leg.target_type,
|
||||
target_id=leg.target_id,
|
||||
max_turns=leg.max_turns,
|
||||
max_budget=leg.max_budget,
|
||||
stopped_at=leg.stopped_at,
|
||||
attempt_count=stats.attempt_count if (stats := attempt_counts.get(leg.id)) else 0,
|
||||
spend=round(stats.spend, 6) if stats else 0.0,
|
||||
)
|
||||
for leg in sorted(legs, key=lambda leg: leg.api_key_id)
|
||||
for leg in sorted(legs, key=lambda leg: (leg.target_type, leg.target_id))
|
||||
),
|
||||
router_name=first.router_name,
|
||||
direction=first.direction,
|
||||
|
|
@ -1090,34 +1134,85 @@ def _group_response(
|
|||
)
|
||||
|
||||
|
||||
_NO_KEY_LABELS: Final[tuple[str | None, str | None]] = (None, None)
|
||||
_NO_TARGET_LABELS: Final[tuple[str | None, str | None]] = (None, None)
|
||||
|
||||
|
||||
async def _with_key_labels(
|
||||
def _target_labels(
|
||||
key_rows: Sequence[_VerificationTokenRow],
|
||||
team_rows: Sequence[_TeamRow],
|
||||
user_rows: Sequence[_UserRow],
|
||||
) -> Mapping[tuple[str, str], tuple[str | None, str | None]]:
|
||||
"""Display labels by (target_type, target_id): a key's (alias, masked name), a
|
||||
team's (alias, None), a user's (email, None)."""
|
||||
return MappingProxyType(
|
||||
{ # mutable-ok: MappingProxyType needs a dict to wrap
|
||||
key: value
|
||||
for key, value in chain(
|
||||
((("key", row.token), (row.key_alias, row.key_name)) for row in key_rows),
|
||||
((("team", row.team_id), (row.team_alias, None)) for row in team_rows),
|
||||
((("user", row.user_id), (row.user_email, None)) for row in user_rows),
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _target_ids_of(responses: Sequence[ShadowEvalJobResponse], target_type: ShadowEvalTargetType) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
sorted(
|
||||
frozenset(
|
||||
target.target_id
|
||||
for response in responses
|
||||
for target in response.targets
|
||||
if target.target_type == target_type
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _with_target_labels(
|
||||
prisma_client: "PrismaClient", responses: Sequence[ShadowEvalJobResponse]
|
||||
) -> tuple[ShadowEvalJobResponse, ...]:
|
||||
"""Resolve every scoped key's hash to its alias and masked name in one batched read,
|
||||
so the UI can say whose traffic a job shadows. Deleted keys resolve to None."""
|
||||
"""Resolve every scoped target's id to a display label in one batched read per kind,
|
||||
so the UI can say whose traffic a job shadows: a key's alias and masked name, a
|
||||
team's alias, a user's email. Deleted targets resolve to None."""
|
||||
if not responses:
|
||||
return ()
|
||||
tokens: Final = sorted(frozenset(key.api_key_id for response in responses for key in response.keys))
|
||||
key_rows: Final = await _verification_tokens(prisma_client).find_many(
|
||||
where={"token": {"in": tokens}} # mutable-ok: Prisma filter
|
||||
tokens: Final = _target_ids_of(responses, "key")
|
||||
team_ids: Final = _target_ids_of(responses, "team")
|
||||
user_ids: Final = _target_ids_of(responses, "user")
|
||||
key_rows: Final = (
|
||||
await _verification_tokens(prisma_client).find_many(
|
||||
where={"token": {"in": list(tokens)}} # mutable-ok: Prisma filter
|
||||
)
|
||||
if tokens
|
||||
else ()
|
||||
)
|
||||
labels: Final[Mapping[str, tuple[str | None, str | None]]] = {
|
||||
row.token: (row.key_alias, row.key_name) for row in key_rows or ()
|
||||
}
|
||||
team_rows: Final = (
|
||||
await _team_rows(prisma_client).find_many(
|
||||
where={"team_id": {"in": list(team_ids)}} # mutable-ok: Prisma filter
|
||||
)
|
||||
if team_ids
|
||||
else ()
|
||||
)
|
||||
user_rows: Final = (
|
||||
await _user_rows(prisma_client).find_many(
|
||||
where={"user_id": {"in": list(user_ids)}} # mutable-ok: Prisma filter
|
||||
)
|
||||
if user_ids
|
||||
else ()
|
||||
)
|
||||
labels: Final = _target_labels(key_rows or (), team_rows or (), user_rows or ())
|
||||
return tuple(
|
||||
response.model_copy(
|
||||
update={ # mutable-ok: pydantic update payload
|
||||
"keys": tuple(
|
||||
key.model_copy(
|
||||
"targets": tuple(
|
||||
target.model_copy(
|
||||
update={ # mutable-ok: pydantic update payload
|
||||
"key_alias": labels.get(key.api_key_id, _NO_KEY_LABELS)[0],
|
||||
"key_name": labels.get(key.api_key_id, _NO_KEY_LABELS)[1],
|
||||
"target_alias": labels.get((target.target_type, target.target_id), _NO_TARGET_LABELS)[0],
|
||||
"key_name": labels.get((target.target_type, target.target_id), _NO_TARGET_LABELS)[1],
|
||||
}
|
||||
)
|
||||
for key in response.keys
|
||||
for target in response.targets
|
||||
)
|
||||
}
|
||||
)
|
||||
|
|
@ -1125,29 +1220,37 @@ async def _with_key_labels(
|
|||
)
|
||||
|
||||
|
||||
async def _shadow_eval_results(prisma_client: "PrismaClient", legs: Sequence[_LegRow]) -> ShadowEvalResult | None:
|
||||
"""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."""
|
||||
async def _shadow_eval_results(
|
||||
prisma_client: "PrismaClient", legs: Sequence[_LegRow]
|
||||
) -> tuple[ShadowEvalResult | None, Mapping[tuple[str, str], ShadowEvalSlice]]:
|
||||
"""One job's stratified verdicts, plus each target's own slice keyed by the
|
||||
(target_type, target_id) pair so a key, team, and user sharing an id can never
|
||||
collapse into one entry. 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 targets use today would the router beat" forward, and "for the
|
||||
turns the router sent to X, did X beat the baseline" in reverse; the per-target
|
||||
slices answer "which target's traffic does the router suit". Reads are bounded by
|
||||
the job's own attempts (<= the sum of its targets' max_turns) via the job_id index."""
|
||||
leg_ids: Final = [leg.id for leg in legs] # mutable-ok: query param
|
||||
by_tier: Final = _ATTEMPT_AGG_ROWS.validate_python(
|
||||
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_TIER_SQL, leg_ids) or ()
|
||||
)
|
||||
if not by_tier:
|
||||
return None
|
||||
return None, MappingProxyType({})
|
||||
by_model: Final = _ATTEMPT_AGG_ROWS.validate_python(
|
||||
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_MODEL_SQL, leg_ids) or ()
|
||||
)
|
||||
key_by_leg: Final = MappingProxyType({leg.id: leg.api_key_id for leg in legs})
|
||||
target_by_leg: Final = MappingProxyType({leg.id: (leg.target_type, leg.target_id) for leg in legs})
|
||||
by_leg: Final = _ATTEMPT_AGG_ROWS.validate_python(
|
||||
await _query_raw(prisma_client, _ATTEMPT_AGG_BY_LEG_SQL, leg_ids) or ()
|
||||
)
|
||||
by_key: Final = tuple(
|
||||
row.model_copy(update={"grp": key_by_leg[row.grp]}) # mutable-ok: pydantic update payload
|
||||
for row in by_leg
|
||||
verdicts_by_target: Final[Mapping[tuple[str, str], ShadowEvalSlice]] = MappingProxyType(
|
||||
{ # mutable-ok: MappingProxyType needs a dict to wrap
|
||||
target_by_leg[slice.group]: slice.model_copy(
|
||||
update={"group": target_by_leg[slice.group][1]} # mutable-ok: pydantic update payload
|
||||
)
|
||||
for slice in _slices(by_leg)
|
||||
}
|
||||
)
|
||||
total_turns: Final = sum(r.turn_count for r in by_tier)
|
||||
funnel_rows: Final = await _query_raw(prisma_client, _FUNNEL_TOTALS_SQL, leg_ids)
|
||||
|
|
@ -1155,10 +1258,9 @@ async def _shadow_eval_results(prisma_client: "PrismaClient", legs: Sequence[_Le
|
|||
# Coverage only when EVERY leg has a funnel row: a partial seed (one leg's insert
|
||||
# failed) must read as unknown, not as job-level counts missing a leg's traffic.
|
||||
funnel: Final = counted if counted is not None and counted.legs_with_rows == len(leg_ids) else None
|
||||
return ShadowEvalResult(
|
||||
result: Final = 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),
|
||||
sampled_real_spend=sum(r.real_spend for r in by_tier),
|
||||
|
|
@ -1168,6 +1270,7 @@ async def _shadow_eval_results(prisma_client: "PrismaClient", legs: Sequence[_Le
|
|||
shed_count=funnel.shed if funnel is not None else None,
|
||||
withheld_count=funnel.withheld if funnel is not None else None,
|
||||
)
|
||||
return result, verdicts_by_target
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
@ -1182,22 +1285,29 @@ 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 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.
|
||||
Start a shadow eval: duplicate a sampled slice of one or more targets' 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 target.
|
||||
|
||||
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.
|
||||
A target is a virtual key, a team, or a user. Team and user targets match on the
|
||||
identity every request resolves to at auth time, so they cover JWT-authenticated
|
||||
traffic, which presents no virtual key; a user target samples that user's traffic
|
||||
across all their teams, whether it arrives on a JWT or a key they own.
|
||||
|
||||
Shadow responses are never served to users. Each key samples until its recorded eval
|
||||
spend, the shadow and judge calls' own cost, reaches max_budget dollars, the job's
|
||||
window ends, or the job is stopped, so one 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.
|
||||
A forward job answers whether the targets should adopt router_name: it samples the
|
||||
requests the router did not serve and duplicates them through it. A reverse job
|
||||
answers whether a target already on the router still gains from it: it samples the
|
||||
requests the router did serve and duplicates them against baseline_model. A target
|
||||
can hold one active job per direction, so both questions can run at once, and a
|
||||
request matching several jobs' targets (say its key and its team) is sampled by
|
||||
each, separately budgeted.
|
||||
|
||||
Shadow responses are never served to users. Each target samples until its recorded
|
||||
eval spend, the shadow and judge calls' own cost, reaches max_budget dollars, the
|
||||
job's window ends, or the job is stopped, so one target 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 sampled request's own identity but are
|
||||
excluded from request counts and auto-router adoption metrics.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router, prisma_client
|
||||
|
||||
|
|
@ -1206,35 +1316,88 @@ async def start_shadow_eval(
|
|||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
if llm_router is None or not _is_configured_pre_routing_strategy(llm_router, data.router_name):
|
||||
raise HTTPException(status_code=400, detail=f"'{data.router_name}' is not a configured auto-router")
|
||||
token_rows: Final = await _verification_tokens(prisma_client).find_many(
|
||||
where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter
|
||||
)
|
||||
unknown: Final = tuple(sorted(frozenset(data.api_key_ids) - frozenset(row.token for row in token_rows or ())))
|
||||
if unknown:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
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"
|
||||
),
|
||||
token_rows: Final = (
|
||||
await _verification_tokens(prisma_client).find_many(
|
||||
where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter
|
||||
)
|
||||
if data.api_key_ids
|
||||
else ()
|
||||
)
|
||||
team_rows: Final = (
|
||||
await _team_rows(prisma_client).find_many(
|
||||
where={"team_id": {"in": list(data.team_ids)}} # mutable-ok: Prisma filter
|
||||
)
|
||||
if data.team_ids
|
||||
else ()
|
||||
)
|
||||
user_rows: Final = (
|
||||
await _user_rows(prisma_client).find_many(
|
||||
where={"user_id": {"in": list(data.user_ids)}} # mutable-ok: Prisma filter
|
||||
)
|
||||
if data.user_ids
|
||||
else ()
|
||||
)
|
||||
unknown_keys: Final = sorted(frozenset(data.api_key_ids) - frozenset(row.token for row in token_rows or ()))
|
||||
unknown_teams: Final = sorted(frozenset(data.team_ids) - frozenset(row.team_id for row in team_rows or ()))
|
||||
unknown_users: Final = sorted(frozenset(data.user_ids) - frozenset(row.user_id for row in user_rows or ()))
|
||||
unknown_parts: Final = tuple(
|
||||
part
|
||||
for part in (
|
||||
(
|
||||
f"api_key_ids not on this proxy: {', '.join(unknown_keys)}; pass each key's token hash, "
|
||||
"the value the key list and key info endpoints report"
|
||||
)
|
||||
if unknown_keys
|
||||
else None,
|
||||
f"team_ids not on this proxy: {', '.join(unknown_teams)}" if unknown_teams else None,
|
||||
f"user_ids not on this proxy: {', '.join(unknown_users)}" if unknown_users else None,
|
||||
)
|
||||
if part is not None
|
||||
)
|
||||
if unknown_parts:
|
||||
raise HTTPException(status_code=400, detail=". ".join(unknown_parts))
|
||||
|
||||
# Every model check below runs once per team the job samples for, since that is the
|
||||
# identity the shadow and judge calls carry and therefore what the router selects on.
|
||||
team_ids: Final = tuple(dict.fromkeys(row.team_id for row in token_rows or ()))
|
||||
# A user target's traffic can span teams, so it validates unscoped (None); each
|
||||
# sampled attempt still resolves the judge under its own request's team at eval time.
|
||||
team_ids: Final = tuple(
|
||||
dict.fromkeys(
|
||||
(
|
||||
*(row.team_id for row in token_rows or ()),
|
||||
*data.team_ids,
|
||||
*((None,) if data.user_ids else ()),
|
||||
)
|
||||
)
|
||||
)
|
||||
_validate_plain_model(llm_router, data.judge_model, "judge_model", team_ids)
|
||||
if data.baseline_model is not None:
|
||||
_validate_plain_model(llm_router, data.baseline_model, "baseline_model", team_ids)
|
||||
_validate_judge_is_not_a_candidate(llm_router, data, team_ids)
|
||||
|
||||
requested_targets: Final[tuple[tuple[ShadowEvalTargetType, str], ...]] = (
|
||||
*(("key", key) for key in data.api_key_ids),
|
||||
*(("team", team) for team in data.team_ids),
|
||||
*(("user", user) for user in data.user_ids),
|
||||
)
|
||||
requested_by_type: Final[tuple[tuple[ShadowEvalTargetType, tuple[str, ...]], ...]] = tuple(
|
||||
(target_type, ids)
|
||||
for target_type, ids in (("key", data.api_key_ids), ("team", data.team_ids), ("user", data.user_ids))
|
||||
if ids
|
||||
)
|
||||
# A job whose window passed or whose budget ran out stopped sampling on its own,
|
||||
# but its legs still hold their slots in the per-key, per-direction partial unique index
|
||||
# until stamped; free them so a new eval can start. Sweeping both directions is deliberate.
|
||||
requested: Final = list(data.api_key_ids) # mutable-ok: query param
|
||||
await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, requested)
|
||||
# but its legs still hold their slots in the per-target, per-direction partial unique
|
||||
# index until stamped; free them so a new eval can start. Sweeping both directions is
|
||||
# deliberate. Sweep and claim filter on exact (target_type, id) pairs so a team id
|
||||
# that happens to equal a key hash never matches the other kind's slot.
|
||||
for target_type, ids in requested_by_type:
|
||||
await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, list(ids), target_type) # mutable-ok: query param
|
||||
claimed: Final = await _shadow_eval_jobs(prisma_client).find_many(
|
||||
where={ # mutable-ok: Prisma filter
|
||||
"api_key_id": {"in": requested}, # mutable-ok: Prisma filter
|
||||
"OR": [ # mutable-ok: Prisma filter
|
||||
{"target_type": target_type, "target_id": {"in": list(ids)}} # mutable-ok: Prisma filter
|
||||
for target_type, ids in requested_by_type
|
||||
],
|
||||
"direction": data.direction,
|
||||
"stopped_at": None,
|
||||
},
|
||||
|
|
@ -1244,7 +1407,7 @@ async def start_shadow_eval(
|
|||
status_code=409,
|
||||
detail=(
|
||||
f"Already in an active {data.direction} shadow eval job: "
|
||||
+ ", ".join(sorted(f"{row.api_key_id} (job {row.group_id})" for row in claimed))
|
||||
+ ", ".join(sorted(f"{row.target_type} {row.target_id} (job {row.group_id})" for row in claimed))
|
||||
+ ". Stop it first."
|
||||
),
|
||||
)
|
||||
|
|
@ -1268,10 +1431,16 @@ async def start_shadow_eval(
|
|||
# Leg ids are minted here rather than by the DB default so the funnel seed below
|
||||
# writes from the same values with no read-back, which a lagging read replica
|
||||
# (DATABASE_URL_READ_REPLICA) could otherwise return empty.
|
||||
leg_ids: Final = tuple(str(uuid4()) for _ in data.api_key_ids)
|
||||
leg_ids: Final = tuple(str(uuid4()) for _ in requested_targets)
|
||||
await _shadow_eval_jobs(prisma_client).create_many(
|
||||
data=[ # mutable-ok: Prisma payload
|
||||
{**shared_config, "id": leg_id, "api_key_id": key} for leg_id, key in zip(leg_ids, data.api_key_ids)
|
||||
{ # mutable-ok: Prisma payload
|
||||
**shared_config,
|
||||
"id": leg_id,
|
||||
"target_type": target_type,
|
||||
"target_id": target_id,
|
||||
} # mutable-ok: Prisma payload
|
||||
for leg_id, (target_type, target_id) in zip(leg_ids, requested_targets)
|
||||
]
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -1280,7 +1449,8 @@ async def start_shadow_eval(
|
|||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=(
|
||||
f"A requested key was claimed by another {data.direction} shadow eval job concurrently. Stop it first."
|
||||
f"A requested target was claimed by another {data.direction} shadow eval job concurrently. "
|
||||
"Stop it first."
|
||||
),
|
||||
) from e
|
||||
# Seed a zero funnel row per leg NOW: a fully covered job never skips a request, so
|
||||
|
|
@ -1293,18 +1463,19 @@ async def start_shadow_eval(
|
|||
)
|
||||
except Exception as seed_err: # noqa: BLE001 # coverage is advisory; the job must still start
|
||||
verbose_proxy_logger.error("shadow_eval: funnel seed failed for job %s: %s", group_id, seed_err)
|
||||
labels: Final = MappingProxyType({row.token: row for row in token_rows})
|
||||
labels: Final = _target_labels(token_rows or (), team_rows or (), user_rows or ())
|
||||
return ShadowEvalJobResponse(
|
||||
job_id=group_id,
|
||||
keys=tuple(
|
||||
ShadowEvalJobKeyResponse(
|
||||
api_key_id=api_key_id,
|
||||
targets=tuple(
|
||||
ShadowEvalJobTargetResponse(
|
||||
target_type=target_type,
|
||||
target_id=target_id,
|
||||
max_turns=SHADOW_EVAL_TURN_VALVE,
|
||||
max_budget=data.max_budget,
|
||||
key_alias=labels[api_key_id].key_alias,
|
||||
key_name=labels[api_key_id].key_name,
|
||||
target_alias=labels.get((target_type, target_id), _NO_TARGET_LABELS)[0],
|
||||
key_name=labels.get((target_type, target_id), _NO_TARGET_LABELS)[1],
|
||||
)
|
||||
for api_key_id in sorted(data.api_key_ids)
|
||||
for target_type, target_id in sorted(requested_targets)
|
||||
),
|
||||
router_name=data.router_name,
|
||||
direction=data.direction,
|
||||
|
|
@ -1324,22 +1495,29 @@ 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 that shadow this key, alone or alongside others")
|
||||
target_type: Annotated[
|
||||
ShadowEvalTargetType | None, Query(description="Kind of target to filter on; requires target_id")
|
||||
] = None,
|
||||
target_id: Annotated[
|
||||
str | None, Query(description="Filter to jobs that shadow this target, 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, each key with its attempt count so status is
|
||||
accurate. Judged counts, spend, and results ride the detail endpoint only."""
|
||||
"""List shadow eval jobs, newest first, each target with its attempt count so status
|
||||
is accurate. Judged counts, spend, and results ride the detail endpoint only."""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
_require_admin_viewer(user_api_key_dict, "view shadow evals")
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
|
||||
filter_type: Final = target_type if isinstance(target_type, str) else None
|
||||
filter_id: Final = target_id if isinstance(target_id, str) else None
|
||||
if (filter_type is None) != (filter_id is None):
|
||||
raise HTTPException(status_code=400, detail="target_type and target_id filter together; pass both or neither")
|
||||
legs: Final = _LEG_ROWS.validate_python(
|
||||
(
|
||||
await _query_raw(prisma_client, _LIST_LEGS_BY_KEY_SQL, limit, api_key_id)
|
||||
if api_key_id
|
||||
await _query_raw(prisma_client, _LIST_LEGS_BY_TARGET_SQL, limit, filter_type, filter_id)
|
||||
if filter_type and filter_id
|
||||
else await _query_raw(prisma_client, _LIST_LEGS_SQL, limit)
|
||||
)
|
||||
or ()
|
||||
|
|
@ -1354,7 +1532,7 @@ async def list_shadow_eval_jobs(
|
|||
by_group, key=lambda group_id: max(leg.created_at for leg in by_group[group_id]), reverse=True
|
||||
)
|
||||
counts: Final = await _leg_attempt_counts(prisma_client, legs)
|
||||
return await _with_key_labels(
|
||||
return await _with_target_labels(
|
||||
prisma_client, tuple(_group_response(group_id, by_group[group_id], counts) for group_id in newest_first)
|
||||
)
|
||||
|
||||
|
|
@ -1391,16 +1569,25 @@ async def get_shadow_eval_job(
|
|||
where={"job_id": {"in": leg_ids}, "outcome": "error"}, # mutable-ok: Prisma filter
|
||||
order={"created_at": "desc"}, # mutable-ok: Prisma order
|
||||
)
|
||||
labeled: Final = await _with_key_labels(
|
||||
labeled: Final = await _with_target_labels(
|
||||
prisma_client, (_group_response(job_id, legs, await _leg_attempt_counts(prisma_client, legs)),)
|
||||
)
|
||||
results, verdicts_by_target = await _shadow_eval_results(prisma_client, legs)
|
||||
return labeled[0].model_copy(
|
||||
update={ # mutable-ok: pydantic update payload
|
||||
"judged_count": totals[0].judged_count if totals else 0,
|
||||
"error_count": totals[0].error_count if totals else 0,
|
||||
"judge_spend": round(totals[0].judge_spend, 6) if totals else 0.0,
|
||||
"last_error": latest_error.error if latest_error else None,
|
||||
"results": await _shadow_eval_results(prisma_client, legs),
|
||||
"results": results,
|
||||
"targets": tuple(
|
||||
target.model_copy(
|
||||
update={ # mutable-ok: pydantic update payload
|
||||
"verdicts": verdicts_by_target.get((target.target_type, target.target_id))
|
||||
}
|
||||
)
|
||||
for target in labeled[0].targets
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -1415,8 +1602,8 @@ 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, every key it scopes at once. Attempts are kept;
|
||||
sampling halts within ~10s. Keys that already stopped on their own budget keep the
|
||||
"""Stop an active shadow eval job, every target it scopes at once. Attempts are kept;
|
||||
sampling halts within ~10s. Targets that already stopped on their own budget keep the
|
||||
stopped_at they earned. The statement is the whole state machine: it claims the job
|
||||
only while a leg still samples inside the window with no stop recorded, so a racing
|
||||
operator, a same-instant budget spend, and a repeat stop all read the same 400 with
|
||||
|
|
@ -1443,5 +1630,5 @@ async def stop_shadow_eval_job(
|
|||
current: Final = _group_response(job_id, legs, counts)
|
||||
if claimed == 0:
|
||||
raise HTTPException(status_code=400, detail=f"Job {job_id} is already {current.status}")
|
||||
labeled: Final = await _with_key_labels(prisma_client, (current,))
|
||||
labeled: Final = await _with_target_labels(prisma_client, (current,))
|
||||
return labeled[0]
|
||||
|
|
|
|||
|
|
@ -1529,14 +1529,15 @@ model LiteLLM_AutoRouterSession {
|
|||
model LiteLLM_ShadowEvalJob {
|
||||
id String @id @default(cuid())
|
||||
group_id String // legs of one job share this; the API's job id
|
||||
api_key_id String // hashed virtual key whose traffic this leg shadows
|
||||
target_type String @default("key") // key | team | user
|
||||
target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows
|
||||
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-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise
|
||||
max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
max_budget Float? // per-target USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
ends_at DateTime
|
||||
|
|
@ -1544,7 +1545,7 @@ model LiteLLM_ShadowEvalJob {
|
|||
stopped_by String? // operator who stopped it early; null when it ended on its own
|
||||
|
||||
@@index([group_id])
|
||||
@@index([api_key_id])
|
||||
@@index([target_type, target_id])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -733,11 +733,20 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo
|
|||
of the three: an agent names the conversation on its first turn, so the cheapest tier would be
|
||||
the pin every session starts with, and the real work that follows would run there for the whole
|
||||
TTL. It describes what that one call is, never what the session's traffic looks like.
|
||||
|
||||
A context-window escalation describes the prompt's size, not the session's complexity, and
|
||||
size shrinks again the moment the client compacts: pinning the escalated tier would hold the
|
||||
session on the big-window model long after the oversized context that forced it is gone. The
|
||||
gate re-fires per request, so leaving these unpinned costs nothing but the classifier call.
|
||||
"""
|
||||
return decision is None or decision.get("cause") not in (
|
||||
"default_model_fallback",
|
||||
"plan_mode",
|
||||
"housekeeping",
|
||||
return decision is None or (
|
||||
decision.get("cause")
|
||||
not in (
|
||||
"default_model_fallback",
|
||||
"plan_mode",
|
||||
"housekeeping",
|
||||
)
|
||||
and not decision.get("context_escalated")
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -786,6 +795,39 @@ class ClassificationOutcome(NamedTuple):
|
|||
classifier_cost: float | None = None
|
||||
|
||||
|
||||
def _allowed(models: tuple[str, ...], fit_filter: frozenset[str] | None) -> tuple[str, ...]:
|
||||
return models if fit_filter is None else tuple(model for model in models if model in fit_filter)
|
||||
|
||||
|
||||
def _apply_context_placement(
|
||||
tier: ComplexityTier | str, signals: tuple[str, ...], placement: _ContextWindowPlacement | None
|
||||
) -> tuple[ComplexityTier | str, tuple[str, ...], ComplexityTier | str | None]:
|
||||
"""(final tier, signals, original tier when the gate escalated, else None)."""
|
||||
if placement is None:
|
||||
return tier, signals, None
|
||||
if _tier_name(placement.tier) == _tier_name(tier):
|
||||
return placement.tier, signals, None
|
||||
return placement.tier, (*signals, "context_escalation"), tier
|
||||
|
||||
|
||||
def _window_can_hold(window: int | None, needed: int, buffer: float) -> bool:
|
||||
return window is None or needed <= int(window * buffer)
|
||||
|
||||
|
||||
def _group_provably_fits(facts: tuple[int | None, bool], needed: int, buffer: float) -> bool:
|
||||
window, has_unknown = facts
|
||||
return window is not None and not has_unknown and needed <= int(window * buffer)
|
||||
|
||||
|
||||
class _ContextWindowPlacement(NamedTuple):
|
||||
"""Where the context-window gate placed the request: the placement tier, the subset of its
|
||||
pool the pick may use, and every configured group not provably misfit (the adaptive filter)."""
|
||||
|
||||
tier: ComplexityTier | str
|
||||
allowed_models: tuple[str, ...]
|
||||
holdable_models: frozenset[str]
|
||||
|
||||
|
||||
class _SessionAffinityPin(NamedTuple):
|
||||
model: str
|
||||
tier: ComplexityTier | None
|
||||
|
|
@ -1222,6 +1264,7 @@ class ComplexityRouter(CustomLogger):
|
|||
classifier_cost: float | None = None,
|
||||
conversation_continuing: bool = True,
|
||||
tier_litellm_params: Mapping[str, object] | None = None,
|
||||
context_escalation_original_tier: ComplexityTier | str | None = None,
|
||||
) -> StandardLoggingRoutingDecision:
|
||||
"""Assemble the per-request provenance record for this router's decision.
|
||||
|
||||
|
|
@ -1271,6 +1314,12 @@ class ComplexityRouter(CustomLogger):
|
|||
decision["classifier_model"] = classifier_model
|
||||
if classifier_cost is not None:
|
||||
decision["classifier_cost"] = classifier_cost
|
||||
if context_escalation_original_tier is not None:
|
||||
# The pair travels together: the flag says the gate moved the request off its
|
||||
# decided tier on prompt size, and the original tier names where the decision
|
||||
# (classifier, keyword rule, or session pin) had placed it before physics did.
|
||||
decision["context_escalated"] = True
|
||||
decision["context_escalation_original_tier"] = _tier_name(context_escalation_original_tier)
|
||||
if tier_litellm_params:
|
||||
masked_tier_litellm_params: Final = mask_credentials_in_payload(tier_litellm_params)
|
||||
if isinstance(masked_tier_litellm_params, Mapping):
|
||||
|
|
@ -1671,7 +1720,7 @@ class ComplexityRouter(CustomLogger):
|
|||
return entry.litellm_params if entry is not None else MappingProxyType({})
|
||||
|
||||
@staticmethod
|
||||
def _pick_from_tier_value(model: str | list[str], tier_key: str) -> str:
|
||||
def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str:
|
||||
if isinstance(model, str):
|
||||
return model
|
||||
if not model:
|
||||
|
|
@ -1687,15 +1736,21 @@ class ComplexityRouter(CustomLogger):
|
|||
raw_messages: list[dict[str, Any]] | None,
|
||||
resolved_messages: list[dict[str, Any]] | None,
|
||||
request_kwargs: dict,
|
||||
allowed_models: tuple[str, ...] | None = None,
|
||||
) -> str:
|
||||
if not self.config.plugins:
|
||||
if allowed_models is not None:
|
||||
return self._pick_from_tier_value(allowed_models, _tier_name(tier))
|
||||
return self.get_model_for_tier(tier)
|
||||
|
||||
from litellm.types.router import RoutingContext
|
||||
|
||||
tier_key: Final = _tier_name(tier)
|
||||
metadata_key: Final = get_metadata_variable_name_from_kwargs(request_kwargs)
|
||||
pool: Final = tuple(self._tier_pools().get(tier_key, ()))
|
||||
full_pool: Final = tuple(self._tier_pools().get(tier_key, ()))
|
||||
pool: Final = (
|
||||
tuple(model for model in full_pool if model in allowed_models) if allowed_models is not None else full_pool
|
||||
)
|
||||
if not pool:
|
||||
# Nothing for the plugins to filter. Falling through would raise the
|
||||
# plugin-filtering error below and send the operator hunting for a policy
|
||||
|
|
@ -1789,6 +1844,7 @@ class ComplexityRouter(CustomLogger):
|
|||
request_kwargs: dict[str, Any] | None = None,
|
||||
hard_floor: ComplexityTier | str | None = None,
|
||||
hard_ceiling: ComplexityTier | str | None = None,
|
||||
fit_filter: frozenset[str] | None = None,
|
||||
) -> str:
|
||||
"""hard_floor excludes every candidate whose tiers all sit below it, turning this pick's
|
||||
soft floors (a distance penalty a high-scoring cheap model can outweigh) into a hard
|
||||
|
|
@ -1801,7 +1857,10 @@ class ComplexityRouter(CustomLogger):
|
|||
tier because that is all it is worth, so a bandit trading cost for quality has nothing to
|
||||
win and must not reach above it. Without it the distance penalty is the only thing holding
|
||||
the tier, and a deployment that lowers tier_distance_penalty silently gets the expensive
|
||||
model back while the routing decision still reads as the cheapest tier."""
|
||||
model back while the routing decision still reads as the cheapest tier.
|
||||
|
||||
fit_filter excludes candidates the context-window gate proved cannot hold the prompt,
|
||||
in every phase including cold start and the tier fallbacks."""
|
||||
from litellm.router_strategy.adaptive_router.bandit import (
|
||||
normalized_cost,
|
||||
thompson_sample,
|
||||
|
|
@ -1812,12 +1871,12 @@ class ComplexityRouter(CustomLogger):
|
|||
if adaptive is None or not isinstance(classified_tier, ComplexityTier):
|
||||
# Custom tier names have no severity index; adaptive is rejected alongside
|
||||
# tier_definitions, so this guard is the contract for any future caller.
|
||||
return self.get_model_for_tier(classified_tier)
|
||||
return self._fitting_tier_fallback(classified_tier, fit_filter)
|
||||
|
||||
request_type: Final = classify_prompt(user_message)
|
||||
classified_idx: Final = TIER_SEVERITY_ORDER.index(classified_tier)
|
||||
pools: Final = self._tier_pools()
|
||||
classified_candidates: Final = tuple(pools.get(_tier_name(classified_tier), ()))
|
||||
classified_candidates: Final = _allowed(tuple(pools.get(_tier_name(classified_tier), ())), fit_filter)
|
||||
cold_start_candidates: Final = tuple(
|
||||
model for model in classified_candidates if adaptive._cells[(request_type, model)].total_samples == 0
|
||||
)
|
||||
|
|
@ -1847,9 +1906,9 @@ class ComplexityRouter(CustomLogger):
|
|||
if self.config.adaptive_eligible == "classified_tier":
|
||||
candidates = list(classified_candidates)
|
||||
if not candidates:
|
||||
return self.get_model_for_tier(classified_tier)
|
||||
return self._fitting_tier_fallback(classified_tier, fit_filter)
|
||||
else:
|
||||
candidates = list(adaptive.config.available_models)
|
||||
candidates = list(_allowed(tuple(adaptive.config.available_models), fit_filter))
|
||||
|
||||
all_costs: Final = [adaptive.model_to_cost.get(m, 0.0) for m in candidates]
|
||||
quality_weight: Final = self.config.adaptive_weights.quality
|
||||
|
|
@ -1896,7 +1955,7 @@ class ComplexityRouter(CustomLogger):
|
|||
best_score = score
|
||||
best_model = model
|
||||
if best_model is None:
|
||||
return self.get_model_for_tier(classified_tier)
|
||||
return self._fitting_tier_fallback(classified_tier, fit_filter)
|
||||
if request_kwargs is not None:
|
||||
metadata = request_kwargs.setdefault("metadata", {})
|
||||
if isinstance(metadata, dict):
|
||||
|
|
@ -1913,6 +1972,12 @@ class ComplexityRouter(CustomLogger):
|
|||
}
|
||||
return best_model
|
||||
|
||||
def _fitting_tier_fallback(self, classified_tier: ComplexityTier | str, fit_filter: frozenset[str] | None) -> str:
|
||||
fitting: Final = _allowed(tuple(self._tier_pools().get(_tier_name(classified_tier), ())), fit_filter)
|
||||
if fit_filter is not None and fitting:
|
||||
return self._pick_from_tier_value(fitting, _tier_name(classified_tier))
|
||||
return self.get_model_for_tier(classified_tier)
|
||||
|
||||
def _resolve_plan_mode_floor(self) -> ComplexityTier | str | None:
|
||||
"""The configured floor as an active tier: the built-in enum member, or the defined
|
||||
name itself for a custom tier set; None when the feature is off."""
|
||||
|
|
@ -1983,6 +2048,163 @@ class ComplexityRouter(CustomLogger):
|
|||
return None
|
||||
return name if self.config.has_custom_tiers else ComplexityTier(name)
|
||||
|
||||
def _deployment_window(self, group: str, deployment: Mapping[str, object]) -> int | None:
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
|
||||
|
||||
deployment_model_info: Final = deployment.get("model_info")
|
||||
declared: Final = (
|
||||
deployment_model_info.get("max_input_tokens") if isinstance(deployment_model_info, Mapping) else None
|
||||
)
|
||||
if isinstance(declared, int):
|
||||
return declared
|
||||
litellm_params: Final = deployment.get("litellm_params")
|
||||
params: Final = litellm_params if isinstance(litellm_params, Mapping) else EMPTY_MAPPING
|
||||
provider_override: Final = params.get("custom_llm_provider")
|
||||
# get_router_model_info resolves the provider, and get_llm_provider runs the OAuth device
|
||||
# flow for github_copilot/chatgpt, so a metadata question must never reach it for those.
|
||||
if declared_authenticating_provider(
|
||||
str(params.get("model") or ""), provider_override if isinstance(provider_override, str) else None
|
||||
):
|
||||
return None
|
||||
try:
|
||||
model_info: Final = self.litellm_router_instance.get_router_model_info(
|
||||
deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts
|
||||
received_model_name=group,
|
||||
)
|
||||
window: Final = model_info.get("max_input_tokens")
|
||||
except Exception: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others
|
||||
return None
|
||||
return window if isinstance(window, int) else None
|
||||
|
||||
def _group_window_facts(self, group: str) -> tuple[int | None, bool]:
|
||||
"""(smallest declared context window across the group's deployments, whether any deployment
|
||||
declares none). The core router picks a deployment within the group without a fit check, so
|
||||
the group is only as safe as its smallest member."""
|
||||
list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None)
|
||||
deployments: Final = list_models(model_name=group) if callable(list_models) else None
|
||||
if not isinstance(deployments, list) or not deployments:
|
||||
return (None, True)
|
||||
windows: Final = tuple(
|
||||
window for deployment in deployments if (window := self._deployment_window(group, deployment)) is not None
|
||||
)
|
||||
return (min(windows) if windows else None, len(windows) < len(deployments))
|
||||
|
||||
@staticmethod
|
||||
def _out_of_band_request_text(request_kwargs: Mapping[str, object]) -> str:
|
||||
"""Prompt content the resolved message list never carries: the Responses API's
|
||||
`instructions`, the /v1/messages top-level `system` block, and tool definitions.
|
||||
A coding agent's context is dominated by these."""
|
||||
import json
|
||||
|
||||
instructions: Final = request_kwargs.get("instructions")
|
||||
proxy_request: Final = request_kwargs.get("proxy_server_request")
|
||||
body: Final = proxy_request.get("body") if isinstance(proxy_request, Mapping) else None
|
||||
system: Final = body.get("system") if isinstance(body, Mapping) else None
|
||||
tools: Final = (
|
||||
body.get("tools") if isinstance(body, Mapping) and body.get("tools") else request_kwargs.get("tools")
|
||||
)
|
||||
tools_text = ""
|
||||
if tools:
|
||||
try:
|
||||
tools_text = json.dumps(tools, default=str)
|
||||
except (TypeError, ValueError):
|
||||
tools_text = str(tools)
|
||||
return (
|
||||
(instructions if isinstance(instructions, str) else "")
|
||||
+ (str(system) if system is not None else "")
|
||||
+ tools_text
|
||||
)
|
||||
|
||||
def _request_byte_upper_bound(
|
||||
self, resolved_messages: Sequence[Mapping[str, object]] | None, request_kwargs: Mapping[str, object]
|
||||
) -> int:
|
||||
"""UTF-8 byte length of all prompt content. BPE emits at least one byte per token in every
|
||||
script, so the token count never exceeds this and 'bytes fit' soundly skips counting."""
|
||||
content_bytes: Final = sum(len(str(m.get("content") or "").encode()) for m in resolved_messages or ())
|
||||
return content_bytes + len(self._out_of_band_request_text(request_kwargs).encode())
|
||||
|
||||
async def _counted_request_tokens(
|
||||
self, resolved_messages: Sequence[Mapping[str, object]], request_kwargs: Mapping[str, object]
|
||||
) -> int | None:
|
||||
"""Real-tokenizer count of the resolved messages plus the out-of-band carriers, off the
|
||||
event loop; None when counting fails, and the gate then leaves the placement alone."""
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
|
||||
out_of_band: Final = self._out_of_band_request_text(request_kwargs)
|
||||
try:
|
||||
counted: Final = await asyncify(litellm.token_counter)(
|
||||
messages=cast(list, resolved_messages) # cast-ok: token_counter only iterates the sequence
|
||||
)
|
||||
return counted + (await asyncify(litellm.token_counter)(text=out_of_band) if out_of_band else 0)
|
||||
except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request
|
||||
verbose_router_logger.debug("ComplexityRouter: context-window token count failed. Got - %s", e)
|
||||
return None
|
||||
|
||||
async def _context_window_placement(
|
||||
self,
|
||||
tier: ComplexityTier | str,
|
||||
resolved_messages: Sequence[Mapping[str, object]] | None,
|
||||
request_kwargs: Mapping[str, object],
|
||||
pool_override: tuple[str, ...] | None = None,
|
||||
) -> _ContextWindowPlacement | None:
|
||||
"""Correct a decided placement whose models provably cannot hold the prompt, or None
|
||||
(the placement stands). Only a real tokenizer count ever moves a request, escalation
|
||||
lands only on groups whose every deployment declares a fitting window, and a group
|
||||
with no resolvable window is never moved on faith in either direction."""
|
||||
if not self.config.enable_context_window_escalation or not resolved_messages:
|
||||
return None
|
||||
pools: Final = self._tier_pools()
|
||||
pool: Final = pool_override if pool_override is not None else tuple(pools.get(_tier_name(tier), ()))
|
||||
if not pool:
|
||||
return None
|
||||
facts: Final = MappingProxyType({group: self._group_window_facts(group) for group in pool})
|
||||
known_windows: Final = tuple(window for window, _ in facts.values() if window is not None)
|
||||
if not known_windows:
|
||||
return None
|
||||
buffer: Final = self.config.context_window_escalation_buffer
|
||||
if self._request_byte_upper_bound(resolved_messages, request_kwargs) <= int(min(known_windows) * buffer):
|
||||
return None
|
||||
needed: Final = await self._counted_request_tokens(resolved_messages, request_kwargs)
|
||||
if needed is None:
|
||||
return None
|
||||
return self._placement_for_tokens(tier=tier, pool=pool, pools=pools, facts=facts, needed=needed)
|
||||
|
||||
def _placement_for_tokens(
|
||||
self,
|
||||
*,
|
||||
tier: ComplexityTier | str,
|
||||
pool: tuple[str, ...],
|
||||
pools: Mapping[str, list[str]],
|
||||
facts: Mapping[str, tuple[int | None, bool]],
|
||||
needed: int,
|
||||
) -> _ContextWindowPlacement | None:
|
||||
buffer: Final = self.config.context_window_escalation_buffer
|
||||
in_tier: Final = tuple(group for group in pool if _window_can_hold(facts[group][0], needed, buffer))
|
||||
if in_tier and len(in_tier) == len(pool):
|
||||
return None
|
||||
holdable: Final = frozenset(
|
||||
group
|
||||
for tier_pool in pools.values()
|
||||
for group in tier_pool
|
||||
if _window_can_hold(self._group_window_facts(group)[0], needed, buffer)
|
||||
)
|
||||
if in_tier:
|
||||
return _ContextWindowPlacement(tier=tier, allowed_models=in_tier, holdable_models=holdable)
|
||||
for name in self.config.tier_names()[self._active_tier_severity(tier) + 1 :]:
|
||||
proven = tuple(
|
||||
group
|
||||
for group in pools.get(name, ())
|
||||
if _group_provably_fits(self._group_window_facts(group), needed, buffer)
|
||||
)
|
||||
if proven:
|
||||
return _ContextWindowPlacement(
|
||||
tier=name if self.config.has_custom_tiers else ComplexityTier(name),
|
||||
allowed_models=proven,
|
||||
holdable_models=holdable,
|
||||
)
|
||||
return None
|
||||
|
||||
def _apply_plan_mode_floor(self, tier: ComplexityTier | str) -> ComplexityTier | str:
|
||||
"""The higher of the decided tier and the plan-mode floor; identity when the floor is unset."""
|
||||
floor: Final = self._resolve_plan_mode_floor()
|
||||
|
|
@ -2381,6 +2603,26 @@ class ComplexityRouter(CustomLogger):
|
|||
session_model: Final = routed_model
|
||||
if plan_floored and pinned_tier is not None:
|
||||
routed_model = self.get_model_for_tier(self._apply_plan_mode_floor(pinned_tier))
|
||||
pin_source_tier: Final = self._tier_for_model(routed_model)
|
||||
pin_placement: Final = (
|
||||
await self._context_window_placement(
|
||||
pin_source_tier, resolved_messages, request_kwargs, pool_override=(routed_model,)
|
||||
)
|
||||
if pin_source_tier is not None
|
||||
else None
|
||||
)
|
||||
pin_context_original_tier: Final = (
|
||||
pin_source_tier
|
||||
if pin_placement is not None
|
||||
and pin_source_tier is not None
|
||||
and _tier_name(pin_placement.tier) != _tier_name(pin_source_tier)
|
||||
else None
|
||||
)
|
||||
if pin_placement is not None and pin_context_original_tier is not None:
|
||||
# The stored pin below keeps the session's own model on purpose.
|
||||
routed_model = self._pick_from_tier_value(
|
||||
pin_placement.allowed_models, _tier_name(pin_placement.tier)
|
||||
)
|
||||
# Refresh the TTL on every hit so an active session doesn't lose its
|
||||
# pin mid-conversation just because it outlives the original write.
|
||||
await self.litellm_router_instance.cache.async_set_cache(
|
||||
|
|
@ -2405,7 +2647,11 @@ class ComplexityRouter(CustomLogger):
|
|||
verbose_router_logger.info(
|
||||
"ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model
|
||||
)
|
||||
routed_pin_tier: Final = self._tier_for_model(routed_model) if plan_floored else resolved_pin_tier
|
||||
routed_pin_tier: Final = (
|
||||
pin_placement.tier
|
||||
if pin_placement is not None and pin_context_original_tier is not None
|
||||
else (self._tier_for_model(routed_model) if plan_floored else resolved_pin_tier)
|
||||
)
|
||||
session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model)
|
||||
has_original_messages: Final = messages is not None and len(messages) > 0
|
||||
return self._with_session_deployment_affinity(
|
||||
|
|
@ -2422,6 +2668,7 @@ class ComplexityRouter(CustomLogger):
|
|||
escalated=escalated,
|
||||
conversation_continuing=conversation_continuing,
|
||||
tier_litellm_params=session_tier_litellm_params,
|
||||
context_escalation_original_tier=pin_context_original_tier,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
|
@ -2616,6 +2863,8 @@ class ComplexityRouter(CustomLogger):
|
|||
plan_floored: Final = tier != pre_floor_tier
|
||||
if plan_floored:
|
||||
signals = (*signals, "plan_mode_floor")
|
||||
context_placement: Final = await self._context_window_placement(tier, resolved_messages, request_kwargs)
|
||||
tier, signals, context_original_tier = _apply_context_placement(tier, signals, context_placement)
|
||||
score_repr: Final = f"{score:.3f}" if score is not None else "n/a"
|
||||
fallback_model: Final = self.config.default_model if not self.config.plugins else None
|
||||
# A sentinel-carrying request skips the failure exit below, whether or not the floor
|
||||
|
|
@ -2662,8 +2911,15 @@ class ComplexityRouter(CustomLogger):
|
|||
# the cheapest tier would then contradict the floor and bound the pick below the tier
|
||||
# the decision reports.
|
||||
housekeeping_ceiling: Final = tier if outcome.cause == "housekeeping" else None
|
||||
# A context-escalated tier becomes the hard floor: a floor the bandit can slide
|
||||
# under is not a floor.
|
||||
routed_model = self._soft_floor_pick(
|
||||
tier, user_message, request_kwargs, hard_floor=plan_floor, hard_ceiling=housekeeping_ceiling
|
||||
tier,
|
||||
user_message,
|
||||
request_kwargs,
|
||||
hard_floor=tier if context_original_tier is not None else plan_floor,
|
||||
hard_ceiling=housekeeping_ceiling,
|
||||
fit_filter=context_placement.holdable_models if context_placement is not None else None,
|
||||
)
|
||||
adaptive: Final = self._ensure_adaptive_router()
|
||||
if adaptive is not None:
|
||||
|
|
@ -2680,7 +2936,13 @@ class ComplexityRouter(CustomLogger):
|
|||
routed_model,
|
||||
)
|
||||
else:
|
||||
routed_model = await self._pick_model_for_tier(tier, messages, resolved_messages, request_kwargs)
|
||||
routed_model = await self._pick_model_for_tier(
|
||||
tier,
|
||||
messages,
|
||||
resolved_messages,
|
||||
request_kwargs,
|
||||
allowed_models=context_placement.allowed_models if context_placement is not None else None,
|
||||
)
|
||||
verbose_router_logger.info(
|
||||
"ComplexityRouter: routing decision cause=%s, tier=%s, score=%s, signals=%s, routed_model=%s",
|
||||
outcome.cause,
|
||||
|
|
@ -2733,5 +2995,6 @@ class ComplexityRouter(CustomLogger):
|
|||
classifier_model=classifier_model,
|
||||
classifier_cost=outcome.classifier_cost,
|
||||
tier_litellm_params=tier_litellm_params,
|
||||
context_escalation_original_tier=context_original_tier,
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -823,6 +823,32 @@ class ComplexityRouterConfig(BaseModel):
|
|||
),
|
||||
)
|
||||
|
||||
enable_context_window_escalation: bool = Field(
|
||||
default=True,
|
||||
description=(
|
||||
"Escalate a request off a tier whose models provably cannot hold its prompt, before "
|
||||
"dispatch. The classifier scores complexity and never prompt size, so a long agentic "
|
||||
"session whose newest ask is trivial lands on a small-window tier and the provider "
|
||||
"rejects it with a context-window 400 that nothing retries. When every model of the "
|
||||
"decided tier has a declared window smaller than the estimated prompt, the request "
|
||||
"moves to the lowest configured tier with a model whose declared window fits; when "
|
||||
"only some of the tier's models fit, the pick is restricted to those and the tier "
|
||||
"keeps the request. Models with no resolvable window are never escalated away from "
|
||||
"and never escalated onto. Set false to dispatch on complexity alone, as before."
|
||||
),
|
||||
)
|
||||
context_window_escalation_buffer: float = Field(
|
||||
default=0.95,
|
||||
gt=0,
|
||||
le=1,
|
||||
description=(
|
||||
"Fraction of a model's declared context window the estimated prompt must fit within. "
|
||||
"The token count is an estimate, so fitting against the full window would dispatch "
|
||||
"prompts that the provider's own tokenizer then rejects; 0.95 leaves room for that "
|
||||
"drift plus the response tokens."
|
||||
),
|
||||
)
|
||||
|
||||
# Semantic (embedding) matching for keyword_tier_rules instead of literal text matching
|
||||
semantic_keyword_matching: bool = Field(
|
||||
default=False,
|
||||
|
|
|
|||
|
|
@ -245,6 +245,8 @@ ShadowEvalStatus: TypeAlias = Literal["running", "completed", "stopped"]
|
|||
|
||||
ShadowEvalDirection: TypeAlias = Literal["forward", "reverse"]
|
||||
|
||||
ShadowEvalTargetType: TypeAlias = Literal["key", "team", "user"]
|
||||
|
||||
DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5"
|
||||
|
||||
# Sample-count ceiling written on every new job: a zero-cost error loop (a shadow arm that
|
||||
|
|
@ -253,16 +255,37 @@ SHADOW_EVAL_TURN_VALVE: Final[int] = 10_000
|
|||
|
||||
|
||||
class StartShadowEvalRequest(BaseModel):
|
||||
"""Start duplicating one or more keys' traffic for blind comparison against an auto-router."""
|
||||
"""Start duplicating one or more targets' traffic for blind comparison against an auto-router.
|
||||
|
||||
A target is a virtual key, a team, or a user; each becomes its own leg with its own
|
||||
budget and stop state. Team and user targets match on the identity every request
|
||||
carries after auth (user_api_key_team_id / user_api_key_user_id), so they cover
|
||||
JWT-authenticated traffic, which presents no virtual key at all."""
|
||||
|
||||
api_key_ids: tuple[str, ...] = Field(
|
||||
min_length=1,
|
||||
default=(),
|
||||
max_length=100,
|
||||
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 carries its own "
|
||||
"max_budget spend budget, so one key exhausting its budget leaves the others sampling. At most 100 "
|
||||
"keys per job, which also bounds every read the job's endpoints make."
|
||||
"Hashed virtual keys whose traffic will be shadowed. Combined with team_ids and user_ids the job "
|
||||
"needs at least one target and at most 100, which also bounds every read the job's endpoints make. "
|
||||
"Each target carries its own max_budget spend budget, so one exhausting its budget leaves the "
|
||||
"others sampling."
|
||||
),
|
||||
)
|
||||
team_ids: tuple[str, ...] = Field(
|
||||
default=(),
|
||||
max_length=100,
|
||||
description=(
|
||||
"Teams whose traffic will be shadowed, matched on the team every authenticated request resolves "
|
||||
"to, so a team's JWT-auth and virtual-key traffic are both sampled"
|
||||
),
|
||||
)
|
||||
user_ids: tuple[str, ...] = Field(
|
||||
default=(),
|
||||
max_length=100,
|
||||
description=(
|
||||
"Users whose traffic will be shadowed, matched on the user every authenticated request resolves "
|
||||
"to across all their teams: JWT requests carrying their subject claim and virtual keys they own"
|
||||
),
|
||||
)
|
||||
router_name: str = Field(description="The auto-router under evaluation, in either direction")
|
||||
|
|
@ -285,7 +308,7 @@ class StartShadowEvalRequest(BaseModel):
|
|||
shadow_percentage: float = Field(
|
||||
ge=0.1,
|
||||
le=100.0,
|
||||
description="Percentage of the key's requests to duplicate through the router",
|
||||
description="Percentage of each target's requests to duplicate through the router",
|
||||
)
|
||||
judge_model: str = Field(
|
||||
default=DEFAULT_SHADOW_EVAL_JUDGE_MODEL,
|
||||
|
|
@ -306,9 +329,9 @@ class StartShadowEvalRequest(BaseModel):
|
|||
ge=0.01,
|
||||
le=10_000,
|
||||
description=(
|
||||
"Per-key USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with "
|
||||
"the same figures the spend pipeline bills. EACH scoped key samples until its recorded eval "
|
||||
"spend reaches this, so a job over N keys spends at most about N times max_budget; in-flight "
|
||||
"Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with "
|
||||
"the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval "
|
||||
"spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight "
|
||||
"samples can overshoot the cap by one sampling cache window"
|
||||
),
|
||||
)
|
||||
|
|
@ -319,7 +342,7 @@ class StartShadowEvalRequest(BaseModel):
|
|||
"""Pydantic ignores unknown fields, so a caller still sending max_turns would
|
||||
silently run on the default dollar budget instead of the bound they asked for."""
|
||||
if isinstance(values, Mapping) and "max_turns" in values:
|
||||
raise ValueError("max_turns was replaced by max_budget, the per-key USD cap on the eval's own spend")
|
||||
raise ValueError("max_turns was replaced by max_budget, the per-target USD cap on the eval's own spend")
|
||||
return values
|
||||
|
||||
@field_validator("shadow_percentage")
|
||||
|
|
@ -327,12 +350,21 @@ class StartShadowEvalRequest(BaseModel):
|
|||
def _round_percentage(cls, value: float) -> float:
|
||||
return round(value, 2)
|
||||
|
||||
@field_validator("api_key_ids")
|
||||
@field_validator("api_key_ids", "team_ids", "user_ids")
|
||||
@classmethod
|
||||
def _dedupe_keys(cls, value: tuple[str, ...]) -> tuple[str, ...]:
|
||||
"""A key named twice would collide with itself on the one-active-per-(key, direction) index."""
|
||||
def _dedupe_targets(cls, value: tuple[str, ...]) -> tuple[str, ...]:
|
||||
"""A target named twice would collide with itself on the one-active-per-(target, direction) index."""
|
||||
return tuple(dict.fromkeys(value))
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _at_least_one_target_at_most_hundred(self) -> "StartShadowEvalRequest":
|
||||
total: Final = len(self.api_key_ids) + len(self.team_ids) + len(self.user_ids)
|
||||
if total < 1:
|
||||
raise ValueError("at least one target is required: pass api_key_ids, team_ids, or user_ids")
|
||||
if total > 100:
|
||||
raise ValueError("at most 100 targets per job across api_key_ids, team_ids, and user_ids")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _baseline_model_matches_direction(self) -> "StartShadowEvalRequest":
|
||||
if self.direction == "reverse" and self.baseline_model is None:
|
||||
|
|
@ -343,8 +375,9 @@ 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 scoped target (embedded on that target's
|
||||
own entry, so slices never need re-joining to a target by id)."""
|
||||
|
||||
group: str
|
||||
turn_count: int
|
||||
|
|
@ -395,12 +428,6 @@ class ShadowEvalResult(BaseModel):
|
|||
"and in reverse the models the router itself picked"
|
||||
)
|
||||
)
|
||||
by_key: tuple[ShadowEvalSlice, ...] = Field(
|
||||
description=(
|
||||
"One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job "
|
||||
"scopes but has not judged a turn for yet are absent rather than reported as zero"
|
||||
),
|
||||
)
|
||||
overall_shadow_win_rate_pct: float
|
||||
overall_tie_rate_pct: float
|
||||
sampled_real_spend: float = Field(
|
||||
|
|
@ -436,27 +463,28 @@ class ShadowEvalResult(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
class ShadowEvalJobKeyResponse(BaseModel):
|
||||
"""One key a job shadows, with its own budget and stop state."""
|
||||
class ShadowEvalJobTargetResponse(BaseModel):
|
||||
"""One target a job shadows (a key, team, or user), with its own budget and stop state."""
|
||||
|
||||
api_key_id: str = Field(description="The hashed virtual key whose traffic this entry scopes")
|
||||
target_type: ShadowEvalTargetType = Field(description="What kind of entity this entry scopes")
|
||||
target_id: str = Field(description="The hashed virtual key, team id, or user id whose traffic this entry scopes")
|
||||
max_turns: int = Field(
|
||||
description=(
|
||||
"This key's sample-count ceiling: the whole budget for jobs created before max_budget "
|
||||
"This target's sample-count ceiling: the whole budget for jobs created before max_budget "
|
||||
"existed, and the error-loop safety valve otherwise"
|
||||
)
|
||||
)
|
||||
max_budget: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"This key's own USD budget for the eval's shadow and judge spend, independent of its "
|
||||
"This target's own USD budget for the eval's shadow and judge spend, independent of its "
|
||||
"siblings'; None on jobs created before spend budgets existed, which max_turns alone bounds"
|
||||
),
|
||||
)
|
||||
stopped_at: datetime | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"When this key's slot was stamped free, whether its own budget ran out, the window closed, "
|
||||
"When this target's slot was stamped free, whether its own budget ran out, the window closed, "
|
||||
"or an operator stopped the job; status is derived, so a spent budget reads completed even "
|
||||
"while this is still unset"
|
||||
),
|
||||
|
|
@ -464,45 +492,53 @@ class ShadowEvalJobKeyResponse(BaseModel):
|
|||
attempt_count: int | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"This key's sampled attempts so far, judged and errored alike, the same count the sampler "
|
||||
"This target's sampled attempts so far, judged and errored alike, the same count the sampler "
|
||||
"budgets against max_turns; populated on list and detail responses. Frozen at stopped_at "
|
||||
"once the key is stamped, so in-flight attempts landing after a stop never reclassify it"
|
||||
"once the target is stamped, so in-flight attempts landing after a stop never reclassify it"
|
||||
),
|
||||
)
|
||||
spend: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"This key's recorded shadow plus judge spend in USD, the same figure the sampler budgets "
|
||||
"This target's recorded shadow plus judge spend in USD, the same figure the sampler budgets "
|
||||
"against max_budget; populated on list and detail responses and frozen at stopped_at "
|
||||
"exactly like attempt_count"
|
||||
),
|
||||
)
|
||||
|
||||
verdicts: "ShadowEvalSlice | None" = Field(
|
||||
default=None,
|
||||
description="This target's own judged-verdict slice; detail endpoint only, None until a turn is judged",
|
||||
)
|
||||
|
||||
@property
|
||||
def budget_spent(self) -> bool:
|
||||
over_spend: Final = self.max_budget is not None and self.spend is not None and self.spend >= self.max_budget
|
||||
return over_spend or (self.attempt_count is not None and self.attempt_count >= self.max_turns)
|
||||
|
||||
key_alias: str | None = Field(
|
||||
target_alias: str | None = Field(
|
||||
default=None,
|
||||
description="Alias of the shadowed key, resolved from the key row at read time; None when unset or deleted",
|
||||
description=(
|
||||
"Display label resolved from the target's own row at read time: the key's alias, the team's "
|
||||
"alias, or the user's email; None when unset or deleted"
|
||||
),
|
||||
)
|
||||
key_name: str | None = Field(
|
||||
default=None,
|
||||
description="Masked display name (sk-...) of the shadowed key, resolved at read time like key_alias",
|
||||
description="Masked display name (sk-...) for key targets, resolved at read time; None for teams and users",
|
||||
)
|
||||
|
||||
|
||||
class ShadowEvalJobResponse(BaseModel):
|
||||
"""A shadow-eval job over one or more keys, each with its own budget and stop state;
|
||||
status is derived from stopped_by, the keys' stop and budget state, and ends_at,
|
||||
"""A shadow-eval job over one or more targets, each with its own budget and stop state;
|
||||
status is derived from stopped_by, the targets' stop and budget state, 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."""
|
||||
|
||||
job_id: str
|
||||
keys: tuple[ShadowEvalJobKeyResponse, ...] = Field(
|
||||
targets: tuple[ShadowEvalJobTargetResponse, ...] = Field(
|
||||
min_length=1,
|
||||
description="The keys whose traffic this job evaluates, and only those keys', each with its own budget",
|
||||
description="The targets whose traffic this job evaluates, and only theirs, each with its own budget",
|
||||
)
|
||||
router_name: str
|
||||
direction: ShadowEvalDirection = "forward"
|
||||
|
|
@ -531,8 +567,8 @@ class ShadowEvalJobResponse(BaseModel):
|
|||
def status(self) -> ShadowEvalStatus:
|
||||
"""Three recorded facts, no history-guessing: a stop is stopped_by (the migration
|
||||
backfills it for every job that displayed stopped when the column arrived, so the
|
||||
pre-column population is closed), completion is the window passing or every key
|
||||
spending its budget, and anything else is running. The all-keys-stamped fallback
|
||||
pre-column population is closed), completion is the window passing or every target
|
||||
spending its budget, and anything else is running. The all-targets-stamped fallback
|
||||
covers only stops written by pre-column pods during a rolling deploy."""
|
||||
if self.stopped_by is not None:
|
||||
return "stopped"
|
||||
|
|
@ -540,8 +576,8 @@ class ShadowEvalJobResponse(BaseModel):
|
|||
self.ends_at if self.ends_at.tzinfo else self.ends_at.replace(tzinfo=timezone.utc)
|
||||
):
|
||||
return "completed"
|
||||
if all(key.budget_spent for key in self.keys):
|
||||
if all(target.budget_spent for target in self.targets):
|
||||
return "completed"
|
||||
if all(key.stopped_at is not None for key in self.keys):
|
||||
if all(target.stopped_at is not None for target in self.targets):
|
||||
return "stopped"
|
||||
return "running"
|
||||
|
|
|
|||
|
|
@ -2886,6 +2886,8 @@ class StandardLoggingRoutingDecision(TypedDict, total=False):
|
|||
classifier_model: str
|
||||
classifier_cost: float
|
||||
escalated: bool
|
||||
context_escalated: bool # writable-ok: Pydantic warns on ReadOnly TypedDict fields
|
||||
context_escalation_original_tier: str # writable-ok: Pydantic warns on ReadOnly TypedDict fields
|
||||
tier_boundaries: StandardLoggingRoutingDecisionTierBoundaries
|
||||
reasoning_override_min_score: float # writable-ok: Pydantic warns on ReadOnly TypedDict fields
|
||||
conversation_continuing: bool
|
||||
|
|
@ -2912,6 +2914,8 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset(
|
|||
"classifier_model",
|
||||
"classifier_cost",
|
||||
"escalated",
|
||||
"context_escalated",
|
||||
"context_escalation_original_tier",
|
||||
"tier_boundaries",
|
||||
"reasoning_override_min_score",
|
||||
"conversation_continuing",
|
||||
|
|
|
|||
|
|
@ -1529,14 +1529,15 @@ model LiteLLM_AutoRouterSession {
|
|||
model LiteLLM_ShadowEvalJob {
|
||||
id String @id @default(cuid())
|
||||
group_id String // legs of one job share this; the API's job id
|
||||
api_key_id String // hashed virtual key whose traffic this leg shadows
|
||||
target_type String @default("key") // key | team | user
|
||||
target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows
|
||||
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-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise
|
||||
max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
max_budget Float? // per-target USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
ends_at DateTime
|
||||
|
|
@ -1544,7 +1545,7 @@ model LiteLLM_ShadowEvalJob {
|
|||
stopped_by String? // operator who stopped it early; null when it ended on its own
|
||||
|
||||
@@index([group_id])
|
||||
@@index([api_key_id])
|
||||
@@index([target_type, target_id])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -33,6 +33,13 @@ EXCLUDED_ROLLOUT_FLAGS = {
|
|||
"LITELLM_RUST",
|
||||
}
|
||||
|
||||
# Internal infrastructure tuning parameters for streaming/queue management
|
||||
# These are advanced settings with sensible defaults that most users should not modify
|
||||
EXCLUDED_INTERNAL_TUNING_VARS = {
|
||||
"ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS",
|
||||
"ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE",
|
||||
}
|
||||
|
||||
EXCLUDED_TERMINAL_VARS = {
|
||||
"TERM",
|
||||
"TERM_PROGRAM",
|
||||
|
|
@ -50,7 +57,9 @@ EXCLUDED_TERMINAL_VARS = {
|
|||
"ALACRITTY_SOCKET",
|
||||
}
|
||||
|
||||
EXCLUDED_KEYS = frozenset(EXCLUDED_TERMINAL_VARS | EXCLUDED_GUARD_ONLY_VARS | EXCLUDED_ROLLOUT_FLAGS)
|
||||
EXCLUDED_KEYS = frozenset(
|
||||
EXCLUDED_TERMINAL_VARS | EXCLUDED_GUARD_ONLY_VARS | EXCLUDED_ROLLOUT_FLAGS | EXCLUDED_INTERNAL_TUNING_VARS
|
||||
)
|
||||
|
||||
# Directories to skip (dependencies, venvs, caches) - only scan litellm source
|
||||
SKIP_DIRS = {
|
||||
|
|
|
|||
|
|
@ -58,11 +58,12 @@ def _prisma(jobs=(), attempt_counts=(), attempt_costs=()) -> MagicMock:
|
|||
return prisma
|
||||
|
||||
|
||||
def _job_record(job: ActiveShadowEvalJob, api_key_id="key-hash") -> MagicMock:
|
||||
def _job_record(job: ActiveShadowEvalJob, target_type="key", target_id="key-hash") -> MagicMock:
|
||||
record = MagicMock()
|
||||
for field, value in dict(
|
||||
id=job.id,
|
||||
api_key_id=api_key_id,
|
||||
target_type=target_type,
|
||||
target_id=target_id,
|
||||
router_name=job.router_name,
|
||||
direction=job.direction,
|
||||
baseline_model=job.baseline_model,
|
||||
|
|
@ -123,7 +124,7 @@ def _spend_counter(store=None):
|
|||
return counter, read, write
|
||||
|
||||
|
||||
def _logger(router=None, prisma=None, jobs=(), counter_store=None) -> ShadowEvalLogger:
|
||||
def _logger(router=None, prisma=None, jobs=(), counter_store=None, jobs_by_target=None) -> ShadowEvalLogger:
|
||||
cache = InMemoryCache(max_size_in_memory=4, default_ttl=60)
|
||||
counter, read, write = _spend_counter(counter_store)
|
||||
funnel_events = []
|
||||
|
|
@ -137,8 +138,9 @@ def _logger(router=None, prisma=None, jobs=(), counter_store=None) -> ShadowEval
|
|||
)
|
||||
logger._test_counter = counter
|
||||
logger._test_funnel = funnel_events
|
||||
if jobs:
|
||||
cache.set_cache("shadow_eval:active_jobs", {"key-hash": tuple(jobs)})
|
||||
seeded = jobs_by_target if jobs_by_target is not None else ({("key", "key-hash"): tuple(jobs)} if jobs else None)
|
||||
if seeded is not None:
|
||||
cache.set_cache("shadow_eval:active_jobs", seeded)
|
||||
return logger
|
||||
|
||||
|
||||
|
|
@ -837,6 +839,86 @@ class TestSuccessHookSkipChain:
|
|||
prisma.db.litellm_shadowevalattempt.create.assert_not_called()
|
||||
|
||||
|
||||
JWT_IDENTITY = {"user_api_key_hash": None, "user_api_key_team_id": "team-eng", "user_api_key_user_id": "dev-alice"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestTargetMatching:
|
||||
"""A request qualifies for a job through ANY of its resolved identities: key hash,
|
||||
team id, or user id. Team and user jobs must therefore sample JWT-authenticated
|
||||
traffic, which carries no key hash at all."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"target,sampled",
|
||||
[
|
||||
(("team", "team-eng"), True),
|
||||
(("user", "dev-alice"), True),
|
||||
(("key", "some-key"), False),
|
||||
],
|
||||
ids=["team-job-samples-jwt-traffic", "user-job-samples-jwt-traffic", "key-jobs-never-match-keyless-traffic"],
|
||||
)
|
||||
async def test_jwt_shaped_traffic_matches_team_and_user_jobs_but_no_key_job(self, target, sampled):
|
||||
prisma = _prisma()
|
||||
router = _router()
|
||||
logger = _logger(router=router, prisma=prisma, jobs_by_target={target: (_job(),)})
|
||||
hook_kwargs = _success_kwargs()
|
||||
hook_kwargs["standard_logging_object"]["metadata"] = dict(JWT_IDENTITY)
|
||||
|
||||
await logger.async_log_success_event(hook_kwargs, RESPONSE, None, None)
|
||||
await _drain(logger)
|
||||
|
||||
if sampled:
|
||||
prisma.db.litellm_shadowevalattempt.create.assert_awaited_once()
|
||||
assert prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"]["job_id"] == "job-1"
|
||||
else:
|
||||
router.acompletion.assert_not_called()
|
||||
prisma.db.litellm_shadowevalattempt.create.assert_not_called()
|
||||
|
||||
async def test_an_event_with_no_identity_early_returns_without_a_cache_read(self):
|
||||
prisma = _prisma()
|
||||
router = _router()
|
||||
cache = MagicMock(spec=InMemoryCache)
|
||||
cache.async_get_cache = AsyncMock()
|
||||
logger = ShadowEvalLogger(
|
||||
router_provider=lambda: router,
|
||||
prisma_provider=lambda: prisma,
|
||||
jobs_cache=cache,
|
||||
)
|
||||
hook_kwargs = _success_kwargs()
|
||||
hook_kwargs["standard_logging_object"]["metadata"] = {}
|
||||
|
||||
await logger.async_log_success_event(hook_kwargs, RESPONSE, None, None)
|
||||
|
||||
cache.async_get_cache.assert_not_awaited()
|
||||
router.acompletion.assert_not_called()
|
||||
prisma.db.litellm_shadowevalattempt.create.assert_not_called()
|
||||
|
||||
async def test_an_event_matching_a_key_job_and_a_team_job_fires_both(self):
|
||||
"""A request's key and its team can each hold a job; the two are separately
|
||||
budgeted experiments, so both fire and each counts its own start."""
|
||||
prisma = _prisma()
|
||||
logger = _logger(
|
||||
router=_router(),
|
||||
prisma=prisma,
|
||||
jobs_by_target={
|
||||
("key", "key-hash"): (_job(id="key-job"),),
|
||||
("team", "team-eng"): (_job(id="team-job"),),
|
||||
},
|
||||
)
|
||||
hook_kwargs = _success_kwargs()
|
||||
hook_kwargs["standard_logging_object"]["metadata"] = {
|
||||
"user_api_key_hash": "key-hash",
|
||||
"user_api_key_team_id": "team-eng",
|
||||
}
|
||||
|
||||
await logger.async_log_success_event(hook_kwargs, RESPONSE, None, None)
|
||||
await _drain(logger)
|
||||
|
||||
rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.call_args_list]
|
||||
assert sorted(row["job_id"] for row in rows) == ["key-job", "team-job"]
|
||||
assert logger._job_starts == {"key-job": 1, "team-job": 1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestActiveJobsCache:
|
||||
async def test_cache_miss_reads_db_once_then_serves_from_cache(self):
|
||||
|
|
@ -851,8 +933,8 @@ class TestActiveJobsCache:
|
|||
first = await logger._active_jobs()
|
||||
second = await logger._active_jobs()
|
||||
|
||||
assert [job.id for job in first["key-hash"]] == ["job-1"]
|
||||
assert second["key-hash"][0].attempts == 7
|
||||
assert [job.id for job in first[("key", "key-hash")]] == ["job-1"]
|
||||
assert second[("key", "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 where["stopped_at"] is None
|
||||
|
|
@ -899,8 +981,8 @@ class TestActiveJobsCache:
|
|||
jobs = await logger._active_jobs()
|
||||
|
||||
assert logger._job_starts == {}
|
||||
assert jobs["key-hash"][0].attempts == 7
|
||||
assert jobs["key-hash"][0].spend == 0.05
|
||||
assert jobs[("key", "key-hash")][0].attempts == 7
|
||||
assert jobs[("key", "key-hash")][0].spend == 0.05
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1249,13 +1331,14 @@ class TestActiveJobsFailClosed:
|
|||
jobs_cache=InMemoryCache(max_size_in_memory=4, default_ttl=60),
|
||||
)
|
||||
|
||||
assert [job.id for job in (await logger._active_jobs())["key-hash"]] == ["job-ok"]
|
||||
assert [job.id for job in (await logger._active_jobs())[("key", "key-hash")]] == ["job-ok"]
|
||||
|
||||
async def test_both_of_a_key_s_jobs_survive_the_lookup(self):
|
||||
async def test_every_targets_jobs_survive_the_lookup_keyed_by_type_and_id(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"),
|
||||
_job_record(_job(id="job-other"), target_id="other-key"),
|
||||
_job_record(_job(id="job-team"), target_type="team", target_id="team-eng"),
|
||||
]
|
||||
prisma = _prisma(jobs=records, attempt_counts=[("job-reverse", 3)])
|
||||
logger = ShadowEvalLogger(
|
||||
|
|
@ -1266,9 +1349,11 @@ class TestActiveJobsFailClosed:
|
|||
|
||||
jobs = await logger._active_jobs()
|
||||
|
||||
assert sorted(job.id for job in jobs["key-hash"]) == ["job-forward", "job-reverse"]
|
||||
assert [job.id for job in jobs["other-key"]] == ["job-other"]
|
||||
assert {job.id: job.attempts for job in jobs["key-hash"]}["job-reverse"] == 3
|
||||
assert sorted(job.id for job in jobs[("key", "key-hash")]) == ["job-forward", "job-reverse"]
|
||||
assert [job.id for job in jobs[("key", "other-key")]] == ["job-other"]
|
||||
assert [job.id for job in jobs[("team", "team-eng")]] == ["job-team"]
|
||||
assert ("team-eng",) not in jobs and "team-eng" not in jobs
|
||||
assert {job.id: job.attempts for job in jobs[("key", "key-hash")]}["job-reverse"] == 3
|
||||
|
||||
|
||||
def _failing_router():
|
||||
|
|
|
|||
|
|
@ -6,9 +6,7 @@ import pytest
|
|||
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages import (
|
||||
streaming_iterator as streaming_iterator_module,
|
||||
)
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages import streaming_iterator as streaming_iterator_module
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
INCOMPLETE_STREAM_ERROR_MESSAGE,
|
||||
AnthropicMessagesStreamHiddenParams,
|
||||
|
|
@ -338,47 +336,6 @@ async def _events_then_hang(events):
|
|||
await asyncio.Event().wait()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_logs_partial_chunks_on_client_disconnect():
|
||||
"""
|
||||
Regression test for LIT-5839: a client disconnect tears the generator
|
||||
down with GeneratorExit at the yield, which used to skip the post-loop
|
||||
logging dispatch entirely, so the partial output tokens the provider
|
||||
already generated (and billed) never reached spend tracking.
|
||||
"""
|
||||
iterator = _RecordingLoggingIterator(
|
||||
litellm_logging_obj=_make_logging_obj("test_disconnect_logs_partial_chunks"),
|
||||
request_body={},
|
||||
)
|
||||
wrapped = iterator.async_sse_wrapper(_events_then_hang(TRUNCATED_TOOL_USE_EVENTS))
|
||||
streamed = [await wrapped.__anext__() for _ in range(len(TRUNCATED_TOOL_USE_EVENTS))]
|
||||
assert iterator.logging_call_count == 0
|
||||
|
||||
await wrapped.aclose()
|
||||
|
||||
assert iterator.logging_call_count == 1
|
||||
assert iterator.logged_chunks == streamed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_logs_partial_chunks_on_cancellation():
|
||||
iterator = _RecordingLoggingIterator(
|
||||
litellm_logging_obj=_make_logging_obj("test_cancellation_logs_partial_chunks"),
|
||||
request_body={},
|
||||
)
|
||||
wrapped = iterator.async_sse_wrapper(_events_then_hang(TRUNCATED_TOOL_USE_EVENTS))
|
||||
streamed = [await wrapped.__anext__() for _ in range(len(TRUNCATED_TOOL_USE_EVENTS))]
|
||||
|
||||
consume_task = asyncio.ensure_future(wrapped.__anext__())
|
||||
await asyncio.sleep(0.01)
|
||||
consume_task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await consume_task
|
||||
|
||||
assert iterator.logging_call_count == 1
|
||||
assert iterator.logged_chunks == streamed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_skips_logging_on_disconnect_before_first_chunk():
|
||||
iterator = _RecordingLoggingIterator(
|
||||
|
|
@ -408,6 +365,561 @@ def test_incomplete_stream_error_sse_event_is_valid_anthropic_error():
|
|||
assert event.endswith("\n\n")
|
||||
|
||||
|
||||
_STREAM_PREFIX = (
|
||||
{"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}},
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "The Roman"}},
|
||||
)
|
||||
_STREAM_TAIL = (
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " Empire ..."}},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 64}},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
|
||||
|
||||
def _output_tokens_from_logged_chunks(chunks: list[bytes]) -> int | None:
|
||||
"""Read the last output_tokens the billing path would see from the SSE bytes."""
|
||||
latest: int | None = None
|
||||
for raw in chunks:
|
||||
for line in raw.decode().splitlines():
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
data = json.loads(line[len("data:"):].strip())
|
||||
usage = data.get("usage") if isinstance(data, dict) else None
|
||||
if isinstance(usage, dict) and usage.get("output_tokens") is not None:
|
||||
latest = usage["output_tokens"]
|
||||
return latest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_bills_full_stream_after_client_disconnect():
|
||||
"""
|
||||
Regression: on a client disconnect mid-stream the upstream provider keeps
|
||||
generating (and billing) the full response. The wrapper must keep draining
|
||||
that upstream to its terminal ``message_delta`` and bill the real
|
||||
output_tokens (64), not the partial count the client drained before leaving
|
||||
(the message_start placeholder, 1).
|
||||
|
||||
A ``tail_gated`` event holds back the stream tail until the client has
|
||||
disconnected, so the tail can only be captured by a drain that survives the
|
||||
client teardown - exactly the path the previous implementation dropped.
|
||||
"""
|
||||
tail_gated = asyncio.Event()
|
||||
upstream_fully_drained = asyncio.Event()
|
||||
|
||||
async def _gated_stream():
|
||||
for event in _STREAM_PREFIX:
|
||||
yield event
|
||||
await tail_gated.wait()
|
||||
for event in _STREAM_TAIL:
|
||||
yield event
|
||||
upstream_fully_drained.set()
|
||||
|
||||
iterator = _RecordingLoggingIterator(
|
||||
litellm_logging_obj=_make_logging_obj("test_bills_full_stream_after_disconnect"),
|
||||
request_body={},
|
||||
)
|
||||
|
||||
gen = iterator.async_sse_wrapper(_gated_stream())
|
||||
|
||||
client_chunks = []
|
||||
async for chunk in gen:
|
||||
client_chunks.append(chunk)
|
||||
if len(client_chunks) == len(_STREAM_PREFIX):
|
||||
break
|
||||
await gen.aclose()
|
||||
|
||||
tail_gated.set()
|
||||
await asyncio.wait_for(upstream_fully_drained.wait(), timeout=5)
|
||||
for _ in range(100):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(client_chunks) == len(_STREAM_PREFIX)
|
||||
|
||||
assert iterator.logged_chunks, "pump never billed after client disconnect"
|
||||
assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64
|
||||
assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks)
|
||||
assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_bills_full_stream_when_client_reads_all():
|
||||
"""Happy path: when the client drains the whole stream, billing still sees
|
||||
the terminal output_tokens (64) and the client gets every chunk."""
|
||||
tail_gated = asyncio.Event()
|
||||
tail_gated.set() # no gating; full stream flows immediately
|
||||
|
||||
async def _full_stream():
|
||||
for event in (*_STREAM_PREFIX, *_STREAM_TAIL):
|
||||
yield event
|
||||
|
||||
iterator = _RecordingLoggingIterator(
|
||||
litellm_logging_obj=_make_logging_obj("test_bills_full_stream_happy_path"),
|
||||
request_body={},
|
||||
)
|
||||
client_chunks = [chunk async for chunk in iterator.async_sse_wrapper(_full_stream())]
|
||||
|
||||
for _ in range(100):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(client_chunks) == len(_STREAM_PREFIX) + len(_STREAM_TAIL)
|
||||
assert _output_tokens_from_logged_chunks(iterator.logged_chunks) == 64
|
||||
assert not any(c.startswith(b"event: error\n") for c in iterator.logged_chunks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_dispatches_deferred_logging_when_client_disconnects_mid_tail():
|
||||
"""
|
||||
Regression: when the pump finishes draining while the client is still
|
||||
connected, ``_handle_streaming_logging`` defers billing for the proxy's
|
||||
post-response hook (``ProxyLogging._fire_deferred_stream_logging``), which
|
||||
only fires on a normally completed response. If the client then disconnects
|
||||
before consuming the queued tail, the response generator tears down via
|
||||
GeneratorExit and that hook never runs. The relay teardown must dispatch
|
||||
the stored deferred billing itself, or the request logs no spend at all.
|
||||
"""
|
||||
dispatched = []
|
||||
deferred_fired = asyncio.Event()
|
||||
|
||||
def _deferred_stream_complete(logging_coroutine):
|
||||
dispatched.append(logging_coroutine)
|
||||
|
||||
async def _consume():
|
||||
logging_coroutine.close()
|
||||
deferred_fired.set()
|
||||
|
||||
return _consume()
|
||||
|
||||
logging_obj = _make_logging_obj("test_deferred_dispatch_on_disconnect_mid_tail")
|
||||
logging_obj._on_deferred_stream_complete = _deferred_stream_complete
|
||||
iterator = BaseAnthropicMessagesStreamingIterator(litellm_logging_obj=logging_obj, request_body={})
|
||||
|
||||
async def _full_stream():
|
||||
for event in (*_STREAM_PREFIX, *_STREAM_TAIL):
|
||||
yield event
|
||||
|
||||
gen = iterator.async_sse_wrapper(_full_stream())
|
||||
client_chunks = []
|
||||
async for chunk in gen:
|
||||
client_chunks.append(chunk)
|
||||
if len(client_chunks) == len(_STREAM_PREFIX):
|
||||
break
|
||||
|
||||
for _ in range(100):
|
||||
if getattr(logging_obj, "_deferred_stream_complete_args", None) is not None:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert getattr(logging_obj, "_deferred_stream_complete_args", None) is not None, "pump never deferred billing"
|
||||
|
||||
await gen.aclose()
|
||||
|
||||
assert len(dispatched) == 1, "relay teardown did not dispatch the deferred billing"
|
||||
assert logging_obj._on_deferred_stream_complete is None
|
||||
assert logging_obj._deferred_stream_complete_args is None
|
||||
await asyncio.wait_for(deferred_fired.wait(), timeout=5)
|
||||
|
||||
|
||||
class _ProviderStreamError(Exception):
|
||||
"""Stand-in for a provider-specific streaming failure carrying a status code."""
|
||||
|
||||
def __init__(self, message: str, status_code: int):
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_reraises_upstream_error_to_connected_client():
|
||||
"""
|
||||
Regression: an upstream failure (Bedrock read / decode / chunk-conversion)
|
||||
before message_stop must propagate the ORIGINAL provider exception to a
|
||||
still-connected client, so the proxy's failure handling keeps the
|
||||
provider-specific status. The pump must not swallow it into a generic
|
||||
api_error event + normal termination.
|
||||
"""
|
||||
|
||||
async def _failing_stream():
|
||||
yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}}
|
||||
raise _ProviderStreamError("bedrock stream blew up", status_code=529)
|
||||
|
||||
iterator = _RecordingLoggingIterator(
|
||||
litellm_logging_obj=_make_logging_obj("test_reraises_upstream_error"),
|
||||
request_body={},
|
||||
)
|
||||
|
||||
received = []
|
||||
|
||||
async def _drain():
|
||||
async for chunk in iterator.async_sse_wrapper(_failing_stream()):
|
||||
received.append(chunk)
|
||||
|
||||
with pytest.raises(_ProviderStreamError) as excinfo:
|
||||
await _drain()
|
||||
|
||||
assert excinfo.value.status_code == 529
|
||||
assert received
|
||||
assert not any(c.startswith(b"event: error\n") for c in received)
|
||||
assert iterator.logged_chunks == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_disconnect():
|
||||
"""
|
||||
When the upstream errors AFTER the client has already disconnected there is
|
||||
no live client to re-raise to and no failure hook will run, so the pump
|
||||
salvages partial spend from what it collected instead of dropping the row.
|
||||
"""
|
||||
tail_gated = asyncio.Event()
|
||||
|
||||
async def _gated_failing_stream():
|
||||
yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}}
|
||||
await tail_gated.wait()
|
||||
raise _ProviderStreamError("late failure", status_code=500)
|
||||
|
||||
iterator = _RecordingLoggingIterator(
|
||||
litellm_logging_obj=_make_logging_obj("test_salvage_partial_on_late_error"),
|
||||
request_body={},
|
||||
)
|
||||
|
||||
gen = iterator.async_sse_wrapper(_gated_failing_stream())
|
||||
received = [await gen.__anext__(), await gen.__anext__()]
|
||||
await gen.aclose() # client disconnects before the upstream error
|
||||
|
||||
tail_gated.set() # let the upstream raise now, after disconnect
|
||||
for _ in range(100):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(received) == 2
|
||||
assert iterator.logged_chunks == received
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_salvages_spend_when_queued_error_is_never_consumed():
|
||||
"""
|
||||
When the upstream errors while the client is still connected, the pump
|
||||
forwards the exception through the queue expecting the relay to re-raise it
|
||||
into the proxy's failure handling. If the client disconnects before
|
||||
consuming that queued exception, the handoff never happens and no failure
|
||||
hook runs, so the pump must notice the unconsumed exception at teardown and
|
||||
salvage partial spend instead of dropping the row entirely.
|
||||
"""
|
||||
upstream_errored = asyncio.Event()
|
||||
|
||||
async def _failing_stream():
|
||||
yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 52, "output_tokens": 1}}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}}
|
||||
upstream_errored.set()
|
||||
raise _ProviderStreamError("mid-stream failure", status_code=500)
|
||||
|
||||
iterator = _RecordingLoggingIterator(
|
||||
litellm_logging_obj=_make_logging_obj("test_salvage_on_unconsumed_queued_error"),
|
||||
request_body={},
|
||||
)
|
||||
|
||||
gen = iterator.async_sse_wrapper(_failing_stream())
|
||||
received = [await gen.__anext__(), await gen.__anext__()]
|
||||
await upstream_errored.wait() # exception is now queued behind the consumed chunks
|
||||
await gen.aclose() # client disconnects without ever consuming the queued exception
|
||||
|
||||
for _ in range(100):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert iterator.logging_call_count == 1
|
||||
assert iterator.logged_chunks == received
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_applies_backpressure_to_slow_client(monkeypatch):
|
||||
"""
|
||||
Regression: the relay queue is bounded, so a slow client throttles the
|
||||
upstream read instead of letting the pump buffer the whole response in
|
||||
memory. With a tiny queue and a client that reads a single chunk, the pump
|
||||
must stall after producing only a queue's worth of chunks ahead, not race
|
||||
to the end of a large stream.
|
||||
"""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2)
|
||||
|
||||
total = 200
|
||||
produced = 0
|
||||
|
||||
async def _fast_stream():
|
||||
nonlocal produced
|
||||
for i in range(total):
|
||||
produced += 1
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}}
|
||||
|
||||
iterator = _make_iterator("test_backpressure_slow_client")
|
||||
gen = iterator.async_sse_wrapper(_fast_stream())
|
||||
try:
|
||||
await gen.__anext__()
|
||||
for _ in range(500):
|
||||
await asyncio.sleep(0)
|
||||
assert produced <= 2 + 3, f"pump ran ahead unthrottled: produced {produced} of {total}"
|
||||
assert produced < total
|
||||
finally:
|
||||
await gen.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(monkeypatch):
|
||||
"""
|
||||
Regression: when the concurrent detached-drain cap is already reached, a
|
||||
pump whose client has disconnected must bill what it collected instead of
|
||||
continuing to drain (and accumulating) the rest of a large upstream stream,
|
||||
so slow/abandoned clients can't pin unbounded worker state.
|
||||
|
||||
The cap slot set is pre-occupied so the single slot is unavailable when this
|
||||
pump reaches its first post-disconnect chunk; that isolates the cap decision
|
||||
from multi-pump scheduling races.
|
||||
"""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1)
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4)
|
||||
|
||||
async def _hold_slot():
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
holder = asyncio.ensure_future(_hold_slot())
|
||||
streaming_iterator_module._DETACHED_STREAM_DRAINS.add(holder)
|
||||
tail_reached = False
|
||||
|
||||
async def _long_stream():
|
||||
nonlocal tail_reached
|
||||
yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}}
|
||||
for i in range(100):
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}}
|
||||
tail_reached = True
|
||||
yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}}
|
||||
yield {"type": "message_stop"}
|
||||
|
||||
iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drain_cap_full"), request_body={})
|
||||
try:
|
||||
gen = iterator.async_sse_wrapper(_long_stream())
|
||||
await gen.__anext__() # message_start
|
||||
await gen.__anext__() # first delta
|
||||
await gen.aclose() # client disconnects; 100+ chunks remain upstream
|
||||
|
||||
for _ in range(200):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert iterator.logged_chunks, "capped pump never billed"
|
||||
assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE
|
||||
assert len(iterator.logged_chunks) < 100
|
||||
assert not any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks)
|
||||
assert tail_reached is False, "pump kept draining past the cap instead of stopping"
|
||||
finally:
|
||||
holder.cancel()
|
||||
streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_bills_partial_when_detached_drains_disabled(monkeypatch):
|
||||
"""
|
||||
Regression: ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS=0 must disable
|
||||
detached draining entirely, not just shrink the cap. With no slots ever
|
||||
available, the very first post-disconnect chunk must fall back to partial
|
||||
spend logging instead of hanging on a cap that's unreachable.
|
||||
"""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 0)
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4)
|
||||
|
||||
tail_reached = False
|
||||
|
||||
async def _long_stream():
|
||||
nonlocal tail_reached
|
||||
yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}}
|
||||
for i in range(100):
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}}
|
||||
tail_reached = True
|
||||
yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}}
|
||||
yield {"type": "message_stop"}
|
||||
|
||||
iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drains_disabled"), request_body={})
|
||||
gen = iterator.async_sse_wrapper(_long_stream())
|
||||
await gen.__anext__() # message_start
|
||||
await gen.__anext__() # first delta
|
||||
await gen.aclose() # client disconnects; 100+ chunks remain upstream
|
||||
|
||||
for _ in range(200):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert iterator.logged_chunks, "pump never billed with detached drains disabled"
|
||||
assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE
|
||||
assert len(iterator.logged_chunks) < 100
|
||||
assert not any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks)
|
||||
assert tail_reached is False, "pump kept draining despite detached drains being disabled"
|
||||
assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached(monkeypatch):
|
||||
"""
|
||||
Regression: when the cap is full and a disconnected pump bails, it must call
|
||||
aclose on the upstream stream so the provider stops generating and billing,
|
||||
not continue running the stream while we record only the partial prefix.
|
||||
"""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1)
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4)
|
||||
|
||||
async def _hold_slot():
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
holder = asyncio.ensure_future(_hold_slot())
|
||||
streaming_iterator_module._DETACHED_STREAM_DRAINS.add(holder)
|
||||
|
||||
class _AbortableStream:
|
||||
def __init__(self):
|
||||
self.aclose_called = False
|
||||
self._remaining = iter(
|
||||
(
|
||||
{"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}},
|
||||
)
|
||||
+ tuple(
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}}
|
||||
for i in range(50)
|
||||
)
|
||||
)
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
try:
|
||||
return next(self._remaining)
|
||||
except StopIteration:
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def aclose(self):
|
||||
self.aclose_called = True
|
||||
|
||||
stream = _AbortableStream()
|
||||
iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("abort_upstream_at_cap"), request_body={})
|
||||
try:
|
||||
gen = iterator.async_sse_wrapper(stream)
|
||||
await gen.__anext__()
|
||||
await gen.__anext__()
|
||||
await gen.aclose()
|
||||
|
||||
for _ in range(200):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert iterator.logged_chunks, "capped pump never billed"
|
||||
assert stream.aclose_called, "upstream aclose was not called when the detached-drain cap was reached"
|
||||
finally:
|
||||
holder.cancel()
|
||||
streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_abort_upstream_logs_warning_when_aclose_raises(caplog):
|
||||
"""_abort_upstream must swallow and log any exception from aclose()."""
|
||||
import logging
|
||||
|
||||
class _ExplodingStream:
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def aclose(self):
|
||||
raise RuntimeError("aclose exploded")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
await BaseAnthropicMessagesStreamingIterator._abort_upstream(_ExplodingStream())
|
||||
|
||||
assert any("abort" in r.message and "RuntimeError" in r.message for r in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enqueue_for_client_returns_false_when_already_detached():
|
||||
"""_enqueue_for_client must return False immediately (without touching the queue)
|
||||
when client_detached is already set before the call."""
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
BaseAnthropicMessagesStreamingIterator,
|
||||
)
|
||||
|
||||
queue: asyncio.Queue[bytes | None | BaseException] = asyncio.Queue(maxsize=1)
|
||||
client_detached = asyncio.Event()
|
||||
client_detached.set()
|
||||
|
||||
result = await BaseAnthropicMessagesStreamingIterator._enqueue_for_client(queue, client_detached, b"chunk")
|
||||
assert result is False
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enqueue_for_client_returns_false_when_client_detaches_while_queue_full():
|
||||
"""_enqueue_for_client must return False (and cancel the put) when the queue
|
||||
is full and client_detached fires before space becomes available."""
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
BaseAnthropicMessagesStreamingIterator,
|
||||
)
|
||||
|
||||
queue: asyncio.Queue[bytes | None | BaseException] = asyncio.Queue(maxsize=1)
|
||||
queue.put_nowait(b"already-full")
|
||||
|
||||
client_detached = asyncio.Event()
|
||||
|
||||
async def _set_detached_soon():
|
||||
await asyncio.sleep(0.01)
|
||||
client_detached.set()
|
||||
|
||||
asyncio.create_task(_set_detached_soon())
|
||||
result = await BaseAnthropicMessagesStreamingIterator._enqueue_for_client(queue, client_detached, b"new-chunk")
|
||||
assert result is False
|
||||
assert queue.qsize() == 1
|
||||
assert queue.get_nowait() == b"already-full"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch):
|
||||
"""Complement to the cap test: with a slot free, a disconnected pump drains
|
||||
the full upstream and bills the terminal usage, and releases its slot after."""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1)
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4)
|
||||
|
||||
async def _stream():
|
||||
yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}}
|
||||
for i in range(20):
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"m{i}"}}
|
||||
yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}}
|
||||
yield {"type": "message_stop"}
|
||||
|
||||
iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drain_cap_free"), request_body={})
|
||||
gen = iterator.async_sse_wrapper(_stream())
|
||||
await gen.__anext__()
|
||||
await gen.__anext__()
|
||||
await gen.aclose()
|
||||
|
||||
for _ in range(300):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks)
|
||||
assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0
|
||||
|
||||
|
||||
def _decode_sse_events(events: tuple[bytes, ...]) -> list[tuple[str, dict]]:
|
||||
decoded = []
|
||||
for event in events:
|
||||
|
|
@ -599,20 +1111,35 @@ async def test_normal_end_with_deferred_dispatch_armed_parks_logging_coroutine(m
|
|||
@pytest.mark.asyncio
|
||||
async def test_client_disconnect_enqueues_immediately_even_when_deferred_dispatch_armed(monkeypatch):
|
||||
"""
|
||||
On client disconnect the guardrail end-of-stream scan never runs, so
|
||||
deferral would strand the spend log; the teardown path must keep
|
||||
enqueueing immediately (LIT-5839) even when the deferred callback is armed.
|
||||
Regression: on client disconnect the guardrail end-of-stream scan never
|
||||
runs, so deferral would strand the spend log. The detached pump's
|
||||
post-disconnect bill must bypass the deferred-dispatch park and enqueue
|
||||
immediately (LIT-5839) even when the deferred callback is armed (LIT-6409).
|
||||
"""
|
||||
worker = _RecordingLoggingWorker()
|
||||
monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker)
|
||||
iterator = _make_iterator("test_disconnect_enqueues_when_armed")
|
||||
iterator.litellm_logging_obj._on_deferred_stream_complete = _noop_deferred_dispatch
|
||||
|
||||
wrapped = iterator.async_sse_wrapper(_events_then_hang(TRUNCATED_TOOL_USE_EVENTS))
|
||||
tail_gated = asyncio.Event()
|
||||
|
||||
async def _gated_stream():
|
||||
for event in TRUNCATED_TOOL_USE_EVENTS:
|
||||
yield event
|
||||
await tail_gated.wait()
|
||||
yield {"type": "message_stop"}
|
||||
|
||||
wrapped = iterator.async_sse_wrapper(_gated_stream())
|
||||
for _ in range(len(TRUNCATED_TOOL_USE_EVENTS)):
|
||||
await wrapped.__anext__()
|
||||
await wrapped.aclose()
|
||||
|
||||
tail_gated.set()
|
||||
for _ in range(100):
|
||||
if worker.enqueued:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(worker.enqueued) == 1
|
||||
assert getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is None
|
||||
worker.close_enqueued()
|
||||
|
|
@ -629,3 +1156,123 @@ async def test_normal_end_without_deferred_dispatch_enqueues_immediately(monkeyp
|
|||
assert len(worker.enqueued) == 1
|
||||
assert getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is None
|
||||
worker.close_enqueued()
|
||||
|
||||
|
||||
def _backpressured_wrapper(iterator, upstream_exhausted: asyncio.Event):
|
||||
async def _stream():
|
||||
try:
|
||||
for event in COMPLETE_STREAM_EVENTS:
|
||||
yield event
|
||||
finally:
|
||||
upstream_exhausted.set()
|
||||
|
||||
return iterator.async_sse_wrapper(_stream())
|
||||
|
||||
|
||||
async def _drain_with_pauses_until_upstream_exhausted(gen, upstream_exhausted: asyncio.Event) -> list:
|
||||
received = []
|
||||
while not upstream_exhausted.is_set():
|
||||
received.append(await gen.__anext__())
|
||||
for _ in range(25):
|
||||
await asyncio.sleep(0)
|
||||
assert len(received) <= len(COMPLETE_STREAM_EVENTS)
|
||||
return received
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_normal_end_parks_deferred_logging_even_when_sentinel_enqueue_backpressured(monkeypatch):
|
||||
"""
|
||||
Regression: with a full relay queue at end of stream, the pump suspends
|
||||
while enqueueing the end-of-stream sentinel, and a client that then drains
|
||||
the whole tail tears the relay down (setting ``client_detached``) before
|
||||
the pump resumes. That teardown is a normally completed response, not a
|
||||
disconnect: billing must still park for the proxy's post-response hook
|
||||
(preserving post_call decoration such as guardrail_information) instead of
|
||||
enqueueing immediately through the teardown path.
|
||||
"""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2)
|
||||
worker = _RecordingLoggingWorker()
|
||||
monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker)
|
||||
|
||||
dispatched = []
|
||||
|
||||
async def _deferred_stream_complete(logging_coroutine):
|
||||
dispatched.append(logging_coroutine)
|
||||
logging_coroutine.close()
|
||||
|
||||
iterator = _make_iterator("test_sentinel_backpressure_normal_end")
|
||||
iterator.litellm_logging_obj._on_deferred_stream_complete = _deferred_stream_complete
|
||||
|
||||
upstream_exhausted = asyncio.Event()
|
||||
gen = _backpressured_wrapper(iterator, upstream_exhausted)
|
||||
received = await _drain_with_pauses_until_upstream_exhausted(gen, upstream_exhausted)
|
||||
|
||||
while True:
|
||||
try:
|
||||
received.append(await gen.__anext__())
|
||||
except StopAsyncIteration:
|
||||
break
|
||||
|
||||
for _ in range(100):
|
||||
if worker.enqueued or getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None):
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(received) == len(COMPLETE_STREAM_EVENTS)
|
||||
assert worker.enqueued == [], "fully delivered stream billed through the teardown path"
|
||||
assert dispatched == []
|
||||
parked = getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None)
|
||||
assert parked is not None, "pump never parked deferred billing"
|
||||
parked[0].close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_consumed(monkeypatch):
|
||||
"""
|
||||
Regression: when the pump has parked deferred billing but its end-of-stream
|
||||
sentinel never fits in the full relay queue (the client disconnects without
|
||||
draining the tail), the proxy's post-response hook never fires. Exactly one
|
||||
of the relay teardown or the pump's fallback must dispatch the parked
|
||||
billing, or the request logs no spend at all.
|
||||
"""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2)
|
||||
worker = _RecordingLoggingWorker()
|
||||
monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker)
|
||||
|
||||
dispatched = []
|
||||
deferred_fired = asyncio.Event()
|
||||
|
||||
def _deferred_stream_complete(logging_coroutine):
|
||||
dispatched.append(logging_coroutine)
|
||||
|
||||
async def _consume():
|
||||
logging_coroutine.close()
|
||||
deferred_fired.set()
|
||||
|
||||
return _consume()
|
||||
|
||||
iterator = _make_iterator("test_sentinel_never_consumed_dispatch")
|
||||
iterator.litellm_logging_obj._on_deferred_stream_complete = _deferred_stream_complete
|
||||
|
||||
upstream_exhausted = asyncio.Event()
|
||||
gen = _backpressured_wrapper(iterator, upstream_exhausted)
|
||||
await _drain_with_pauses_until_upstream_exhausted(gen, upstream_exhausted)
|
||||
|
||||
for _ in range(100):
|
||||
if getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is not None:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
await gen.aclose()
|
||||
|
||||
for _ in range(100):
|
||||
if dispatched:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert len(dispatched) == 1, "parked billing was never dispatched"
|
||||
assert getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is None
|
||||
assert getattr(iterator.litellm_logging_obj, "_on_deferred_stream_complete", None) is None
|
||||
assert len(worker.enqueued) == 1, "teardown billing enqueued alongside the deferred dispatch"
|
||||
await worker.enqueued[0]
|
||||
assert deferred_fired.is_set()
|
||||
|
|
|
|||
|
|
@ -3066,17 +3066,21 @@ def test_bedrock_invoke_messages_allows_converted_websearch_function_tool():
|
|||
async def test_bedrock_sse_wrapper_dispatches_logging_on_client_disconnect():
|
||||
"""
|
||||
Regression test for LIT-5839: closing the outer bedrock_sse_wrapper
|
||||
mid-stream (what the proxy does on a client disconnect) must close the
|
||||
inner async_sse_wrapper deterministically so the partial-stream logging
|
||||
fires. `completion_start_time` is only stamped on the logging object by
|
||||
that dispatch, so it observing a value proves the whole chain ran.
|
||||
mid-stream (what the proxy does on a client disconnect) must not lose the
|
||||
stream's spend logging. Since the detached-pump relay, the upstream read
|
||||
survives the disconnect and billing fires once the provider stream ends,
|
||||
so the dispatch is awaited after releasing the upstream instead of being
|
||||
observed synchronously at aclose(). `completion_start_time` is only
|
||||
stamped on the logging object by that dispatch, so it observing a value
|
||||
proves the whole chain ran.
|
||||
"""
|
||||
cfg = AmazonAnthropicClaudeMessagesConfig()
|
||||
release_upstream = asyncio.Event()
|
||||
|
||||
async def _hanging_stream():
|
||||
async def _gated_stream():
|
||||
yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 25, "output_tokens": 1}}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}}
|
||||
await asyncio.Event().wait()
|
||||
await release_upstream.wait()
|
||||
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="bedrock/invoke/anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
|
|
@ -3087,11 +3091,16 @@ async def test_bedrock_sse_wrapper_dispatches_logging_on_client_disconnect():
|
|||
litellm_call_id="test_bedrock_sse_wrapper_disconnect_logging",
|
||||
function_id="test_bedrock_sse_wrapper_disconnect_logging",
|
||||
)
|
||||
wrapped = cfg.bedrock_sse_wrapper(_hanging_stream(), litellm_logging_obj=logging_obj, request_body={})
|
||||
wrapped = cfg.bedrock_sse_wrapper(_gated_stream(), litellm_logging_obj=logging_obj, request_body={})
|
||||
await wrapped.__anext__()
|
||||
await wrapped.__anext__()
|
||||
assert logging_obj.completion_start_time is None
|
||||
|
||||
await wrapped.aclose()
|
||||
release_upstream.set()
|
||||
|
||||
for _ in range(500):
|
||||
if logging_obj.completion_start_time is not None:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert logging_obj.completion_start_time is not None
|
||||
|
|
|
|||
|
|
@ -878,7 +878,8 @@ def _leg_record(**overrides: object) -> MagicMock:
|
|||
defaults = {
|
||||
"id": "leg-1",
|
||||
"group_id": "job-1",
|
||||
"api_key_id": "key-hash",
|
||||
"target_type": "key",
|
||||
"target_id": "key-hash",
|
||||
"router_name": "my-router",
|
||||
"direction": "forward",
|
||||
"baseline_model": None,
|
||||
|
|
@ -912,8 +913,28 @@ def _key_record(
|
|||
return record
|
||||
|
||||
|
||||
def _team_record(team_id: str, team_alias: str | None) -> MagicMock:
|
||||
record = MagicMock(spec=["team_id", "team_alias"])
|
||||
record.team_id = team_id
|
||||
record.team_alias = team_alias
|
||||
return record
|
||||
|
||||
|
||||
def _user_record(user_id: str, user_email: str | None) -> MagicMock:
|
||||
record = MagicMock(spec=["user_id", "user_email"])
|
||||
record.user_id = user_id
|
||||
record.user_email = user_email
|
||||
return record
|
||||
|
||||
|
||||
def _shadow_prisma(
|
||||
legs=(), agg_rows=None, by_leg_rows=None, known_keys=("key-hash", "key-hash-2"), key_teams=None
|
||||
legs=(),
|
||||
agg_rows=None,
|
||||
by_leg_rows=None,
|
||||
known_keys=("key-hash", "key-hash-2"),
|
||||
key_teams=None,
|
||||
known_teams=None,
|
||||
known_users=None,
|
||||
) -> MagicMock:
|
||||
"""The job-table fake honours the filters it is handed, so a read that forgets
|
||||
stopped_at sees rows the partial index would have released, one that forgets
|
||||
|
|
@ -921,6 +942,8 @@ def _shadow_prisma(
|
|||
group read that matched on a leg id would come back empty."""
|
||||
prisma = MagicMock()
|
||||
teams: Final = key_teams or {}
|
||||
team_aliases: Final = known_teams or {}
|
||||
user_emails: Final = known_users or {}
|
||||
|
||||
async def find_tokens(*, where):
|
||||
"""Honours the token filter, like the job-table fake below: the endpoint derives the
|
||||
|
|
@ -931,6 +954,17 @@ def _shadow_prisma(
|
|||
|
||||
prisma.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=find_tokens)
|
||||
|
||||
async def find_teams(*, where):
|
||||
requested = where["team_id"]["in"]
|
||||
return [_team_record(t, alias) for t, alias in team_aliases.items() if t in requested]
|
||||
|
||||
async def find_users(*, where):
|
||||
requested = where["user_id"]["in"]
|
||||
return [_user_record(u, email) for u, email in user_emails.items() if u in requested]
|
||||
|
||||
prisma.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_teams)
|
||||
prisma.db.litellm_usertable.find_many = AsyncMock(side_effect=find_users)
|
||||
|
||||
async def execute_raw(sql: str, *params: object):
|
||||
if "SET stopped_by" in sql:
|
||||
group = [row for row in stored if row.group_id == params[0]]
|
||||
|
|
@ -959,9 +993,19 @@ def _shadow_prisma(
|
|||
async def find_many_legs(where=None, **_: object):
|
||||
current = list(stored)
|
||||
w = dict(where or {})
|
||||
if "api_key_id" in w:
|
||||
wanted = w["api_key_id"]["in"] if isinstance(w["api_key_id"], dict) else [w["api_key_id"]]
|
||||
current = [row for row in current if row.api_key_id in wanted]
|
||||
if "OR" in w:
|
||||
pairs = [
|
||||
(
|
||||
branch["target_type"],
|
||||
branch["target_id"]["in"] if isinstance(branch["target_id"], dict) else [branch["target_id"]],
|
||||
)
|
||||
for branch in w["OR"]
|
||||
]
|
||||
current = [
|
||||
row
|
||||
for row in current
|
||||
if any(row.target_type == target_type and row.target_id in ids for target_type, ids in pairs)
|
||||
]
|
||||
if "direction" in w:
|
||||
current = [row for row in current if row.direction == w["direction"]]
|
||||
if "stopped_at" in w:
|
||||
|
|
@ -983,7 +1027,8 @@ def _shadow_prisma(
|
|||
fields = (
|
||||
"id",
|
||||
"group_id",
|
||||
"api_key_id",
|
||||
"target_type",
|
||||
"target_id",
|
||||
"router_name",
|
||||
"direction",
|
||||
"baseline_model",
|
||||
|
|
@ -1009,7 +1054,11 @@ def _shadow_prisma(
|
|||
if "AS attempt_count" in sql:
|
||||
return prisma.attempt_rows
|
||||
if "GROUP BY group_id" in sql:
|
||||
scoped = [row for row in stored if "api_key_id = $2" not in sql or row.api_key_id == params[1]]
|
||||
scoped = [
|
||||
row
|
||||
for row in stored
|
||||
if "target_type = $2" not in sql or (row.target_type == params[1] and row.target_id == params[2])
|
||||
]
|
||||
keep = set(newest_groups(scoped, params[0]))
|
||||
return [leg_dict(row) for row in stored if row.group_id in keep]
|
||||
if "FILTER (WHERE outcome != 'error')::int AS judged_count" in sql:
|
||||
|
|
@ -1058,7 +1107,7 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp
|
|||
|
||||
response = await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN)
|
||||
|
||||
sweep_sql, sweep_keys = prisma.db.execute_raw.call_args.args
|
||||
sweep_sql, sweep_ids, sweep_type = prisma.db.execute_raw.call_args.args
|
||||
assert "stopped_at IS NULL" in sweep_sql
|
||||
assert "j.ends_at <= (NOW() AT TIME ZONE 'utc')" in sweep_sql
|
||||
assert "SET stopped_at = (NOW() AT TIME ZONE 'utc')" in sweep_sql
|
||||
|
|
@ -1066,12 +1115,13 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp
|
|||
assert "j.max_budget IS NOT NULL" in sweep_sql
|
||||
assert ">= j.max_budget" in sweep_sql
|
||||
assert "SUM(a.judge_cost + a.shadow_cost + a.shadow_classifier_cost)" in sweep_sql
|
||||
assert "j.api_key_id = ANY($1::text[])" in sweep_sql
|
||||
assert sweep_keys == ["key-hash", "key-hash-2"]
|
||||
assert "j.target_type = $2 AND j.target_id = ANY($1::text[])" in sweep_sql
|
||||
assert sweep_ids == ["key-hash", "key-hash-2"]
|
||||
assert sweep_type == "key"
|
||||
prisma.db.litellm_shadowevaljob.create_many.assert_awaited_once()
|
||||
rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"]
|
||||
assert [row["api_key_id"] for row in rows] == ["key-hash", "key-hash-2"]
|
||||
assert len({frozenset((k, v) for k, v in row.items() if k not in ("api_key_id", "id")) for row in rows}) == 1
|
||||
assert [(row["target_type"], row["target_id"]) for row in rows] == [("key", "key-hash"), ("key", "key-hash-2")]
|
||||
assert len({frozenset((k, v) for k, v in row.items() if k not in ("target_id", "id")) for row in rows}) == 1
|
||||
assert len({row["id"] for row in rows}) == len(rows)
|
||||
assert len({row["group_id"] for row in rows}) == 1
|
||||
assert all(row["max_turns"] == SHADOW_EVAL_TURN_VALVE and row["created_by"] == "admin" for row in rows)
|
||||
|
|
@ -1080,11 +1130,12 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp
|
|||
assert response.job_id == rows[0]["group_id"]
|
||||
assert response.status == "running"
|
||||
assert response.judged_count is None
|
||||
assert [(key.api_key_id, key.max_budget, key.key_alias) for key in response.keys] == [
|
||||
assert [(target.target_id, target.max_budget, target.target_alias) for target in response.targets] == [
|
||||
("key-hash", 5.0, "prod-alpha"),
|
||||
("key-hash-2", 5.0, "prod-alpha"),
|
||||
]
|
||||
assert all(key.max_turns == SHADOW_EVAL_TURN_VALVE for key in response.keys)
|
||||
assert all(target.target_type == "key" for target in response.targets)
|
||||
assert all(target.max_turns == SHADOW_EVAL_TURN_VALVE for target in response.targets)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1232,7 +1283,7 @@ async def test_start_shadow_eval_rejections(
|
|||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
_configure_anthropic_sdk_judge(monkeypatch)
|
||||
prisma = _shadow_prisma(legs=[_leg_record(id=f"leg-{key}", group_id="job-7", api_key_id=key) for key in claimed])
|
||||
prisma = _shadow_prisma(legs=[_leg_record(id=f"leg-{key}", group_id="job-7", target_id=key) for key in claimed])
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
|
||||
|
||||
|
|
@ -1315,14 +1366,14 @@ async def test_start_shadow_eval_names_the_busy_key_and_its_job(monkeypatch: pyt
|
|||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
_configure_anthropic_sdk_judge(monkeypatch)
|
||||
prisma = _shadow_prisma(legs=[_leg_record(id="leg-b", group_id="job-7", api_key_id="key-hash-2")])
|
||||
prisma = _shadow_prisma(legs=[_leg_record(id="leg-b", group_id="job-7", target_id="key-hash-2")])
|
||||
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", "key-hash-2")), ADMIN)
|
||||
assert exc.value.status_code == 409
|
||||
assert "key-hash-2 (job job-7)" in exc.value.detail
|
||||
assert "key key-hash-2 (job job-7)" in exc.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1441,6 +1492,180 @@ def test_start_shadow_eval_request_dedupes_and_bounds_the_key_set():
|
|||
_start_request(api_key_ids=tuple(f"k{i}" for i in range(101)))
|
||||
|
||||
|
||||
def test_start_request_bounds_the_combined_target_count_across_types():
|
||||
"""The 1..100 bound counts keys, teams, and users together, so a caller cannot dodge
|
||||
it by spreading targets over the three fields, and a request naming no target of any
|
||||
type samples nothing and is rejected."""
|
||||
with pytest.raises(ValidationError, match="at least one target"):
|
||||
_start_request(api_key_ids=(), team_ids=(), user_ids=())
|
||||
with pytest.raises(ValidationError, match="at most 100 targets"):
|
||||
_start_request(api_key_ids=tuple(f"k{i}" for i in range(60)), team_ids=tuple(f"t{i}" for i in range(41)))
|
||||
mixed = _start_request(api_key_ids=tuple(f"k{i}" for i in range(60)), team_ids=tuple(f"t{i}" for i in range(40)))
|
||||
assert len(mixed.api_key_ids) + len(mixed.team_ids) == 100
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"overrides,prisma_kwargs,expected_target",
|
||||
[
|
||||
(
|
||||
{"api_key_ids": (), "team_ids": ("team-eng",)},
|
||||
{"known_teams": {"team-eng": "Engineering"}},
|
||||
("team", "team-eng", "Engineering"),
|
||||
),
|
||||
(
|
||||
{"api_key_ids": (), "user_ids": ("dev-alice",)},
|
||||
{"known_users": {"dev-alice": "alice@example.com"}},
|
||||
("user", "dev-alice", "alice@example.com"),
|
||||
),
|
||||
],
|
||||
ids=["team-target-labeled-by-team-alias", "user-target-labeled-by-user-email"],
|
||||
)
|
||||
async def test_start_shadow_eval_creates_typed_legs_for_team_and_user_targets(
|
||||
monkeypatch: pytest.MonkeyPatch, overrides, prisma_kwargs, expected_target
|
||||
):
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
_configure_anthropic_sdk_judge(monkeypatch)
|
||||
prisma = _shadow_prisma(**prisma_kwargs)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
|
||||
|
||||
response = await start_shadow_eval(_start_request(**overrides), ADMIN)
|
||||
|
||||
target_type, target_id, target_alias = expected_target
|
||||
rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"]
|
||||
assert [(row["target_type"], row["target_id"]) for row in rows] == [(target_type, target_id)]
|
||||
assert response.status == "running"
|
||||
target = response.targets[0]
|
||||
assert (target.target_type, target.target_id, target.target_alias, target.key_name) == (
|
||||
target_type,
|
||||
target_id,
|
||||
target_alias,
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"overrides,prisma_kwargs,expected_detail",
|
||||
[
|
||||
(
|
||||
{"api_key_ids": (), "team_ids": ("team-eng", "team-ghost")},
|
||||
{"known_teams": {"team-eng": "Engineering"}},
|
||||
"team_ids not on this proxy: team-ghost",
|
||||
),
|
||||
(
|
||||
{"api_key_ids": (), "user_ids": ("dev-alice", "dev-ghost")},
|
||||
{"known_users": {"dev-alice": "alice@example.com"}},
|
||||
"user_ids not on this proxy: dev-ghost",
|
||||
),
|
||||
],
|
||||
ids=["unknown-team", "unknown-user"],
|
||||
)
|
||||
async def test_start_shadow_eval_rejects_teams_and_users_this_proxy_does_not_know(
|
||||
monkeypatch: pytest.MonkeyPatch, overrides, prisma_kwargs, expected_detail
|
||||
):
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
_configure_anthropic_sdk_judge(monkeypatch)
|
||||
prisma = _shadow_prisma(**prisma_kwargs)
|
||||
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(**overrides), ADMIN)
|
||||
assert exc.value.status_code == 400
|
||||
assert expected_detail in exc.value.detail
|
||||
prisma.db.litellm_shadowevaljob.create_many.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_shadow_eval_mixed_targets_create_both_legs_and_sweep_once_per_type(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
_configure_anthropic_sdk_judge(monkeypatch)
|
||||
prisma = _shadow_prisma(known_teams={"team-eng": "Engineering"})
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
|
||||
|
||||
response = await start_shadow_eval(_start_request(team_ids=("team-eng",)), ADMIN)
|
||||
|
||||
rows = prisma.db.litellm_shadowevaljob.create_many.call_args.kwargs["data"]
|
||||
assert [(row["target_type"], row["target_id"]) for row in rows] == [("key", "key-hash"), ("team", "team-eng")]
|
||||
assert len({row["group_id"] for row in rows}) == 1
|
||||
sweeps = [
|
||||
call.args
|
||||
for call in prisma.db.execute_raw.await_args_list
|
||||
if "SET stopped_at = (NOW() AT TIME ZONE 'utc')" in call.args[0]
|
||||
]
|
||||
assert [(ids, target_type) for _, ids, target_type in sweeps] == [(["key-hash"], "key"), (["team-eng"], "team")]
|
||||
assert [(t.target_type, t.target_id, t.target_alias) for t in response.targets] == [
|
||||
("key", "key-hash", "prod-alpha"),
|
||||
("team", "team-eng", "Engineering"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_shadow_eval_names_the_busy_team_target(monkeypatch: pytest.MonkeyPatch):
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
_configure_anthropic_sdk_judge(monkeypatch)
|
||||
prisma = _shadow_prisma(
|
||||
legs=[_leg_record(id="leg-t", group_id="job-7", target_type="team", target_id="team-eng")],
|
||||
known_teams={"team-eng": "Engineering"},
|
||||
)
|
||||
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=(), team_ids=("team-eng",)), ADMIN)
|
||||
assert exc.value.status_code == 409
|
||||
assert "team team-eng (job job-7)" in exc.value.detail
|
||||
prisma.db.litellm_shadowevaljob.create_many.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_shadow_eval_claim_matches_exact_target_pairs_not_bare_ids(monkeypatch: pytest.MonkeyPatch):
|
||||
"""A key whose hash happens to spell a team's id must not hold the team's slot: the
|
||||
claim matches (target_type, target_id) pairs, never ids across kinds."""
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
_configure_anthropic_sdk_judge(monkeypatch)
|
||||
prisma = _shadow_prisma(
|
||||
legs=[_leg_record(id="leg-k", group_id="job-7", target_type="key", target_id="team-eng")],
|
||||
known_teams={"team-eng": "Engineering"},
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
|
||||
|
||||
response = await start_shadow_eval(_start_request(api_key_ids=(), team_ids=("team-eng",)), ADMIN)
|
||||
|
||||
assert response.status == "running"
|
||||
prisma.db.litellm_shadowevaljob.create_many.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_shadow_eval_jobs_rejects_a_lone_filter_half(monkeypatch: pytest.MonkeyPatch):
|
||||
"""target_type and target_id only mean anything together: a bare id could name a key
|
||||
or a team, and a bare type filters nothing."""
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
prisma = _shadow_prisma(legs=[_leg_record()])
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
|
||||
with pytest.raises(HTTPException) as id_only:
|
||||
await list_shadow_eval_jobs(VIEWER, target_type=None, target_id="key-hash", limit=50)
|
||||
assert id_only.value.status_code == 400
|
||||
|
||||
with pytest.raises(HTTPException) as type_only:
|
||||
await list_shadow_eval_jobs(VIEWER, target_type="key", target_id=None, limit=50)
|
||||
assert type_only.value.status_code == 400
|
||||
prisma.db.query_raw.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_shadow_eval_concurrent_unique_violation_is_a_409(monkeypatch: pytest.MonkeyPatch):
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
|
@ -1530,7 +1755,7 @@ async def test_get_shadow_eval_job_pools_counts_and_slices_results_per_key(monke
|
|||
},
|
||||
]
|
||||
prisma = _shadow_prisma(
|
||||
legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2", max_turns=50)],
|
||||
legs=[_leg_record(), _leg_record(id="leg-2", target_id="key-hash-2", max_turns=50)],
|
||||
agg_rows=tier_rows,
|
||||
by_leg_rows=leg_rows,
|
||||
)
|
||||
|
|
@ -1549,8 +1774,10 @@ async def test_get_shadow_eval_job_pools_counts_and_slices_results_per_key(monke
|
|||
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 [(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
|
||||
verdicts_by_target = {(t.target_type, t.target_id): t.verdicts for t in response.targets}
|
||||
assert verdicts_by_target[("key", "key-hash")].turn_count == 6
|
||||
assert verdicts_by_target[("key", "key-hash")].shadow_win_rate_pct == 66.7
|
||||
assert verdicts_by_target[("key", "key-hash-2")].turn_count == 4
|
||||
agg_sql = next(call.args[0] for call in prisma.db.query_raw.await_args_list if "real_spend" in call.args[0])
|
||||
assert agg_sql.count("FILTER (WHERE real_cost IS NOT NULL AND NOT real_cache_hit)") == 2
|
||||
assert response.results.by_tier[0].real_spend == 0.08
|
||||
|
|
@ -1561,7 +1788,10 @@ async def test_get_shadow_eval_job_pools_counts_and_slices_results_per_key(monke
|
|||
assert response.results.not_sampled_count is None
|
||||
assert response.results.unjudgeable_count is None
|
||||
assert response.results.shed_count is None
|
||||
assert [(key.api_key_id, key.max_turns) for key in response.keys] == [("key-hash", 200), ("key-hash-2", 50)]
|
||||
assert [(target.target_id, target.max_turns) for target in response.targets] == [
|
||||
("key-hash", 200),
|
||||
("key-hash-2", 50),
|
||||
]
|
||||
totals_args = [call.args for call in prisma.db.query_raw.await_args_list if "judged_count" in call.args[0]]
|
||||
assert totals_args == [(totals_args[0][0], ["leg-1", "leg-2"])]
|
||||
error_where = prisma.db.litellm_shadowevalattempt.find_first.call_args.kwargs["where"]
|
||||
|
|
@ -1595,7 +1825,7 @@ async def test_list_shadow_eval_jobs_collapses_legs_into_jobs_newest_first(monke
|
|||
_leg_record(created_at=datetime(2026, 8, 13, tzinfo=timezone.utc)),
|
||||
_leg_record(
|
||||
id="leg-2",
|
||||
api_key_id="key-hash-2",
|
||||
target_id="key-hash-2",
|
||||
stopped_at=stamp,
|
||||
created_at=datetime(2026, 8, 13, tzinfo=timezone.utc),
|
||||
),
|
||||
|
|
@ -1615,14 +1845,14 @@ async def test_list_shadow_eval_jobs_collapses_legs_into_jobs_newest_first(monke
|
|||
)
|
||||
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, target_type=None, target_id=None, limit=50)
|
||||
|
||||
assert [(job.job_id, job.status) for job in jobs] == [
|
||||
("job-1", "running"),
|
||||
("job-2", "stopped"),
|
||||
("job-3", "completed"),
|
||||
]
|
||||
assert [key.api_key_id for key in jobs[0].keys] == ["key-hash", "key-hash-2"]
|
||||
assert [target.target_id for target in jobs[0].targets] == ["key-hash", "key-hash-2"]
|
||||
assert all(job.judged_count is None and job.results is None for job in jobs)
|
||||
legs_sql, legs_limit = prisma.db.query_raw.await_args_list[0].args
|
||||
assert "GROUP BY group_id ORDER BY MAX(created_at) DESC LIMIT $1::int" in legs_sql
|
||||
|
|
@ -1648,17 +1878,20 @@ async def test_list_shadow_eval_jobs_filters_to_jobs_containing_the_key(monkeypa
|
|||
prisma = _shadow_prisma(
|
||||
legs=[
|
||||
_leg_record(),
|
||||
_leg_record(id="leg-2", api_key_id="key-hash-2"),
|
||||
_leg_record(id="leg-3", group_id="job-2", api_key_id="key-hash-2"),
|
||||
_leg_record(id="leg-2", target_id="key-hash-2"),
|
||||
_leg_record(id="leg-3", group_id="job-2", target_id="key-hash-2"),
|
||||
_leg_record(id="leg-4", group_id="job-3"),
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
|
||||
jobs = await list_shadow_eval_jobs(VIEWER, api_key_id="key-hash-2", limit=50)
|
||||
jobs = await list_shadow_eval_jobs(VIEWER, target_type="key", target_id="key-hash-2", limit=50)
|
||||
|
||||
assert [job.job_id for job in jobs] == ["job-1", "job-2"]
|
||||
assert [key.api_key_id for key in jobs[0].keys] == ["key-hash", "key-hash-2"]
|
||||
assert [target.target_id for target in jobs[0].targets] == ["key-hash", "key-hash-2"]
|
||||
legs_sql, *legs_params = prisma.db.query_raw.await_args_list[0].args
|
||||
assert "WHERE target_type = $2 AND target_id = $3" in legs_sql
|
||||
assert legs_params == [50, "key", "key-hash-2"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -1682,7 +1915,7 @@ async def test_job_status_runs_until_every_key_stops_and_completed_outranks_stop
|
|||
legs=[
|
||||
_leg_record(
|
||||
id=f"leg-{index}",
|
||||
api_key_id=f"key-{index}",
|
||||
target_id=f"key-{index}",
|
||||
stopped_at=stamp if stopped else None,
|
||||
ends_at=datetime.now(timezone.utc) + timedelta(days=days_left),
|
||||
)
|
||||
|
|
@ -1691,7 +1924,7 @@ async def test_job_status_runs_until_every_key_stops_and_completed_outranks_stop
|
|||
)
|
||||
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, target_type=None, target_id=None, limit=50)
|
||||
|
||||
assert [job.status for job in jobs] == [expected]
|
||||
|
||||
|
|
@ -1707,9 +1940,9 @@ async def test_list_reads_completed_once_every_key_spends_its_budget(monkeypatch
|
|||
prisma = _shadow_prisma(
|
||||
legs=[
|
||||
_leg_record(max_turns=5),
|
||||
_leg_record(id="leg-2", api_key_id="key-hash-2", max_turns=5),
|
||||
_leg_record(id="leg-3", group_id="job-2", api_key_id="key-hash", max_turns=5),
|
||||
_leg_record(id="leg-4", group_id="job-2", api_key_id="key-hash-2", max_turns=5),
|
||||
_leg_record(id="leg-2", target_id="key-hash-2", max_turns=5),
|
||||
_leg_record(id="leg-3", group_id="job-2", target_id="key-hash", max_turns=5),
|
||||
_leg_record(id="leg-4", group_id="job-2", target_id="key-hash-2", max_turns=5),
|
||||
]
|
||||
)
|
||||
prisma.attempt_rows = [
|
||||
|
|
@ -1720,13 +1953,13 @@ async def test_list_reads_completed_once_every_key_spends_its_budget(monkeypatch
|
|||
]
|
||||
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, target_type=None, target_id=None, limit=50)
|
||||
|
||||
by_id = {job.job_id: job for job in jobs}
|
||||
assert by_id["job-1"].status == "completed"
|
||||
assert all(key.stopped_at is None for key in by_id["job-1"].keys)
|
||||
assert all(target.stopped_at is None for target in by_id["job-1"].targets)
|
||||
assert by_id["job-2"].status == "running"
|
||||
assert {key.api_key_id: key.attempt_count for key in by_id["job-2"].keys} == {"key-hash": 5, "key-hash-2": 3}
|
||||
assert {t.target_id: t.attempt_count for t in by_id["job-2"].targets} == {"key-hash": 5, "key-hash-2": 3}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1740,7 +1973,7 @@ async def test_recorded_operator_stop_outranks_budget_arithmetic(monkeypatch: py
|
|||
prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 6, "spend": 0.0}]
|
||||
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, target_type=None, target_id=None, limit=50)
|
||||
assert jobs[0].status == "stopped"
|
||||
assert jobs[0].stopped_by == "admin"
|
||||
|
||||
|
|
@ -1760,7 +1993,7 @@ async def test_backfilled_legacy_stop_never_reads_as_completion(monkeypatch: pyt
|
|||
prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 6, "spend": 0.0}]
|
||||
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, target_type=None, target_id=None, limit=50)
|
||||
assert jobs[0].status == "stopped"
|
||||
|
||||
|
||||
|
|
@ -1808,6 +2041,56 @@ def test_max_budget_migration_is_additive_and_leaves_legacy_rows_null():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.asyncio
|
||||
async def test_verdicts_keep_same_id_targets_of_different_kinds_distinct(monkeypatch):
|
||||
"""A team and a user can legitimately share an id; their slices must not merge."""
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
leg_rows = [
|
||||
{
|
||||
"grp": "leg-1",
|
||||
"turn_count": 6,
|
||||
"real_wins": 2,
|
||||
"shadow_wins": 4,
|
||||
"ties": 0,
|
||||
"avg_confidence": 0.8,
|
||||
"real_spend": 0.02,
|
||||
"shadow_spend": 0.01,
|
||||
"cache_hit_turns": 0,
|
||||
},
|
||||
{
|
||||
"grp": "leg-2",
|
||||
"turn_count": 4,
|
||||
"real_wins": 3,
|
||||
"shadow_wins": 0,
|
||||
"ties": 1,
|
||||
"avg_confidence": 0.6,
|
||||
"real_spend": 0.05,
|
||||
"shadow_spend": 0.04,
|
||||
"cache_hit_turns": 1,
|
||||
},
|
||||
]
|
||||
prisma = _shadow_prisma(
|
||||
legs=[
|
||||
_leg_record(target_type="team", target_id="dev-alice"),
|
||||
_leg_record(id="leg-2", target_type="user", target_id="dev-alice"),
|
||||
],
|
||||
agg_rows=leg_rows[:1],
|
||||
by_leg_rows=leg_rows,
|
||||
known_teams={"dev-alice": "alias"},
|
||||
known_users={"dev-alice": "alice@example.com"},
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
|
||||
response = await get_shadow_eval_job("job-1", VIEWER)
|
||||
|
||||
verdicts_by_target = {(t.target_type, t.target_id): t.verdicts for t in response.targets}
|
||||
assert verdicts_by_target[("team", "dev-alice")].turn_count == 6
|
||||
assert verdicts_by_target[("team", "dev-alice")].shadow_win_rate_pct == 66.7
|
||||
assert verdicts_by_target[("user", "dev-alice")].turn_count == 4
|
||||
assert verdicts_by_target[("user", "dev-alice")].shadow_win_rate_pct == 0.0
|
||||
|
||||
|
||||
async def test_stop_rejects_a_job_that_already_spent_its_budget(monkeypatch: pytest.MonkeyPatch):
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
|
|
@ -1832,9 +2115,9 @@ async def test_list_reads_completed_once_every_key_spends_its_dollar_budget(monk
|
|||
prisma = _shadow_prisma(
|
||||
legs=[
|
||||
_leg_record(max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0),
|
||||
_leg_record(id="leg-2", api_key_id="key-hash-2", max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0),
|
||||
_leg_record(id="leg-2", target_id="key-hash-2", max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0),
|
||||
_leg_record(
|
||||
id="leg-3", group_id="job-2", api_key_id="key-hash", max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0
|
||||
id="leg-3", group_id="job-2", target_id="key-hash", max_turns=SHADOW_EVAL_TURN_VALVE, max_budget=1.0
|
||||
),
|
||||
]
|
||||
)
|
||||
|
|
@ -1845,13 +2128,13 @@ async def test_list_reads_completed_once_every_key_spends_its_dollar_budget(monk
|
|||
]
|
||||
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, target_type=None, target_id=None, limit=50)
|
||||
|
||||
by_id = {job.job_id: job for job in jobs}
|
||||
assert by_id["job-1"].status == "completed"
|
||||
assert by_id["job-2"].status == "running"
|
||||
assert {key.api_key_id: key.spend for key in by_id["job-1"].keys} == {"key-hash": 1.0, "key-hash-2": 1.25}
|
||||
assert all(key.max_budget == 1.0 for key in by_id["job-1"].keys)
|
||||
assert {t.target_id: t.spend for t in by_id["job-1"].targets} == {"key-hash": 1.0, "key-hash-2": 1.25}
|
||||
assert all(target.max_budget == 1.0 for target in by_id["job-1"].targets)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1880,11 +2163,11 @@ async def test_legacy_jobs_without_a_dollar_budget_stay_turn_gated(monkeypatch:
|
|||
prisma.attempt_rows = [{"job_id": "leg-1", "attempt_count": 40, "spend": 250.0}]
|
||||
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, target_type=None, target_id=None, limit=50)
|
||||
|
||||
assert jobs[0].status == "running"
|
||||
assert jobs[0].keys[0].max_budget is None
|
||||
assert jobs[0].keys[0].spend == 250.0
|
||||
assert jobs[0].targets[0].max_budget is None
|
||||
assert jobs[0].targets[0].spend == 250.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1892,14 +2175,14 @@ async def test_shadow_eval_responses_name_every_shadowed_key(monkeypatch: pytest
|
|||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
prisma = _shadow_prisma(
|
||||
legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="deleted-key-hash")],
|
||||
legs=[_leg_record(), _leg_record(id="leg-2", target_id="deleted-key-hash")],
|
||||
known_keys=("key-hash", "key-hash-2"),
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
|
||||
jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50)
|
||||
assert [(key.key_alias, key.key_name) for key in jobs[0].keys] == [
|
||||
jobs = await list_shadow_eval_jobs(VIEWER, target_type=None, target_id=None, limit=50)
|
||||
assert [(target.target_alias, target.key_name) for target in jobs[0].targets] == [
|
||||
(None, None),
|
||||
("prod-alpha", "sk-...lpha"),
|
||||
]
|
||||
|
|
@ -1907,7 +2190,7 @@ async def test_shadow_eval_responses_name_every_shadowed_key(monkeypatch: pytest
|
|||
assert batched_where == {"token": {"in": ["deleted-key-hash", "key-hash"]}}
|
||||
|
||||
detail = await get_shadow_eval_job("job-1", VIEWER)
|
||||
assert [key.key_alias for key in detail.keys] == [None, "prod-alpha"]
|
||||
assert [target.target_alias for target in detail.targets] == [None, "prod-alpha"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1919,7 +2202,7 @@ async def test_stop_shadow_eval_stops_every_unstopped_leg_and_rejects_non_runnin
|
|||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
earned = datetime.now(timezone.utc) - timedelta(hours=1)
|
||||
prisma = _shadow_prisma(legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2", stopped_at=earned)])
|
||||
prisma = _shadow_prisma(legs=[_leg_record(), _leg_record(id="leg-2", target_id="key-hash-2", stopped_at=earned)])
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
|
||||
stopped = await stop_shadow_eval_job("job-1", ADMIN)
|
||||
|
|
@ -1938,9 +2221,9 @@ async def test_stop_shadow_eval_stops_every_unstopped_leg_and_rejects_non_runnin
|
|||
assert datetime.fromisoformat(stop_stamp).tzinfo is None
|
||||
assert prisma.db.execute_raw.await_count == 1
|
||||
prisma.db.litellm_shadowevaljob.update_many.assert_not_called()
|
||||
by_key = {key.api_key_id: key.stopped_at for key in stopped.keys}
|
||||
assert by_key["key-hash-2"] == earned
|
||||
assert by_key["key-hash"] is not None and by_key["key-hash"] != earned
|
||||
by_target = {target.target_id: target.stopped_at for target in stopped.targets}
|
||||
assert by_target["key-hash-2"] == earned
|
||||
assert by_target["key-hash"] is not None and by_target["key-hash"] != earned
|
||||
|
||||
done_leg = _leg_record(ends_at=datetime.now(timezone.utc) - timedelta(days=1))
|
||||
prisma_done = _shadow_prisma(legs=[done_leg])
|
||||
|
|
@ -2331,7 +2614,7 @@ async def test_get_shadow_eval_job_sums_funnel_rows_across_legs(monkeypatch: pyt
|
|||
},
|
||||
]
|
||||
prisma = _shadow_prisma(
|
||||
legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2")],
|
||||
legs=[_leg_record(), _leg_record(id="leg-2", target_id="key-hash-2")],
|
||||
agg_rows=tier_rows,
|
||||
)
|
||||
prisma.funnel_rows = [{"legs_with_rows": 2, "not_sampled": 30, "unjudgeable": 5, "shed": 2, "withheld": 3}]
|
||||
|
|
@ -2366,7 +2649,7 @@ async def test_partially_seeded_funnel_reads_as_unknown_coverage(monkeypatch: py
|
|||
},
|
||||
]
|
||||
prisma = _shadow_prisma(
|
||||
legs=[_leg_record(), _leg_record(id="leg-2", api_key_id="key-hash-2")],
|
||||
legs=[_leg_record(), _leg_record(id="leg-2", target_id="key-hash-2")],
|
||||
agg_rows=tier_rows,
|
||||
)
|
||||
prisma.funnel_rows = [{"legs_with_rows": 1, "not_sampled": 30, "unjudgeable": 5, "shed": 2, "withheld": 0}]
|
||||
|
|
|
|||
|
|
@ -10021,3 +10021,399 @@ class TestHeuristicFirst:
|
|||
)
|
||||
outcome = await router.aclassify(NO_SIGNAL_PROMPT)
|
||||
assert outcome.cause == "default_model_fallback"
|
||||
|
||||
|
||||
def _windowed_router(*deployments: tuple) -> Router:
|
||||
"""Real Router; each deployment is (group, provider_model, declared window or None).
|
||||
None means no declared override on a model the cost map does not know: unresolvable."""
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": group,
|
||||
"litellm_params": {"model": provider_model, "mock_response": "ok"},
|
||||
**({"model_info": {"max_input_tokens": window}} if window is not None else {}),
|
||||
}
|
||||
for group, provider_model, window in deployments
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
_SMALL = ("small-model", "openai/gpt-3.5-turbo", 16385)
|
||||
_BIG = ("big-model", "openai/gpt-4o-mini", 200000)
|
||||
|
||||
# A long agentic session whose newest ask is trivial: low-density filler the heuristic scores
|
||||
# SIMPLE, sized well past a 16,385-token window so the fit check must move it.
|
||||
_CONTEXT_FILLER = "The meeting notes were saved to the shared folder for later review this week. " * 2000
|
||||
_OVERSIZED_TURNS = [
|
||||
{"role": "user", "content": "Here is everything discussed so far. " + _CONTEXT_FILLER},
|
||||
{"role": "assistant", "content": "Noted, I have read all of it."},
|
||||
{"role": "user", "content": "ok continue"},
|
||||
]
|
||||
# ~40k CJK chars: chars/4 says ~10k tokens, the real tokenizer says several times that. A
|
||||
# character-based shortcut would skip counting and dispatch this to a 16k window.
|
||||
_CJK_TURNS = [
|
||||
{"role": "user", "content": "会议记录已经保存到共享文件夹里,供大家本周晚些时候查阅和讨论使用。" * 1300},
|
||||
{"role": "user", "content": "ok continue"},
|
||||
]
|
||||
|
||||
|
||||
def _tier_config(**overrides) -> Dict:
|
||||
return {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}, **overrides}
|
||||
|
||||
|
||||
class TestContextWindowEscalation:
|
||||
"""A tier decided on complexity alone must still hold the prompt, or the provider 400s.
|
||||
|
||||
The classifier never weighs prompt size (token count is a 0.10-weight scoring dimension,
|
||||
below every tier boundary), so a long session ending in a trivial ask lands on the
|
||||
smallest tier and dies upstream with no retry. The gate checks fit pre-dispatch, against
|
||||
windows resolved through the real Router deployment chain.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_oversized_simple_prompt_escalates_to_the_lowest_tier_that_fits(self):
|
||||
"""The LIT-6503 regression: SIMPLE verdict, 17k-token prompt, 16,385-token tier model.
|
||||
|
||||
Unfixed, this dispatched to the small model and the provider rejected it with a
|
||||
context-window 400 that neither the retry layer nor tier-keyed fallbacks catch.
|
||||
"""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=_windowed_router(_SMALL, _BIG),
|
||||
complexity_router_config=_tier_config(),
|
||||
)
|
||||
|
||||
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == "big-model"
|
||||
assert result.routing_decision["context_escalated"] is True
|
||||
assert result.routing_decision["context_escalation_original_tier"] == "SIMPLE"
|
||||
assert result.routing_decision["tier"] == "COMPLEX"
|
||||
assert "context_escalation" in result.routing_decision["signals"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_prompt_that_fits_routes_exactly_as_before(self):
|
||||
"""The gate must be invisible for normal traffic: same model, no escalation facts."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=_windowed_router(_SMALL, _BIG),
|
||||
complexity_router_config=_tier_config(),
|
||||
)
|
||||
|
||||
result = await router.async_pre_routing_hook(
|
||||
model="test-router", request_kwargs={}, messages=[{"role": "user", "content": "ok continue"}]
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == "small-model"
|
||||
assert "context_escalated" not in result.routing_decision
|
||||
assert "context_escalation_original_tier" not in result.routing_decision
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_pick_prefers_a_fitting_group_inside_the_decided_tier(self):
|
||||
"""A tier holding both a small and a large group keeps the request and picks the one
|
||||
that fits, which is cheaper than escalating and preserves the classifier's decision."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG),
|
||||
complexity_router_config={"tiers": {"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}},
|
||||
)
|
||||
|
||||
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == "mid-model"
|
||||
assert result.routing_decision["tier"] == "SIMPLE"
|
||||
assert "context_escalated" not in result.routing_decision
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_group_is_only_as_safe_as_its_smallest_deployment(self):
|
||||
"""One group name can front deployments with different windows, and the core router
|
||||
picks among them with no fit check, so retaining the group on its largest member
|
||||
turns the pick into a coin flip against a 400. The gate judges the group by its
|
||||
smallest resolvable window and escalates past it."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "mixed-pool",
|
||||
"litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"},
|
||||
"model_info": {"max_input_tokens": 16385},
|
||||
},
|
||||
{
|
||||
"model_name": "mixed-pool",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"},
|
||||
"model_info": {"max_input_tokens": 200000},
|
||||
},
|
||||
{
|
||||
"model_name": "big-model",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"},
|
||||
"model_info": {"max_input_tokens": 200000},
|
||||
},
|
||||
]
|
||||
),
|
||||
complexity_router_config={"tiers": {"SIMPLE": "mixed-pool", "COMPLEX": "big-model"}},
|
||||
)
|
||||
|
||||
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == "big-model"
|
||||
assert result.routing_decision["context_escalated"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_dense_text_cannot_slip_past_the_counting_shortcut(self):
|
||||
"""CJK text runs several tokens per four characters, so a chars/4 shortcut would skip
|
||||
the real count and dispatch an oversized prompt. The skip is gated on the UTF-8 byte
|
||||
length, which the token count can never exceed."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=_windowed_router(_SMALL, _BIG),
|
||||
complexity_router_config=_tier_config(),
|
||||
)
|
||||
|
||||
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_CJK_TURNS)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == "big-model"
|
||||
assert result.routing_decision["context_escalated"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"deployments,tiers,expected_model",
|
||||
[
|
||||
(
|
||||
(("small-model", "openai/unmapped-model-under-test", None), _BIG),
|
||||
{"SIMPLE": "small-model", "COMPLEX": "big-model"},
|
||||
"small-model",
|
||||
),
|
||||
(
|
||||
(_SMALL, ("mid-model", "openai/another-unmapped-model", None), _BIG),
|
||||
{"SIMPLE": "small-model", "MEDIUM": "mid-model", "COMPLEX": "big-model"},
|
||||
"big-model",
|
||||
),
|
||||
((_SMALL,), {"SIMPLE": "small-model"}, "small-model"),
|
||||
],
|
||||
ids=["unknown-window-stays", "unproven-target-skipped", "nothing-fits-stays"],
|
||||
)
|
||||
async def test_unknown_windows_are_never_acted_on(self, deployments, tiers, expected_model):
|
||||
"""No faith in either direction: a model with no resolvable window is never escalated
|
||||
away from (its misfit is unprovable) and never escalated onto (its fit is unprovable);
|
||||
when nothing provably fits, the classified tier stands and the client owns overflow."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=_windowed_router(*deployments),
|
||||
complexity_router_config={"tiers": tiers},
|
||||
)
|
||||
|
||||
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == expected_model
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_disabled_gate_dispatches_on_complexity_alone(self):
|
||||
"""The escape hatch: enable_context_window_escalation false restores today's behavior."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=_windowed_router(_SMALL, _BIG),
|
||||
complexity_router_config=_tier_config(enable_context_window_escalation=False),
|
||||
)
|
||||
|
||||
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == "small-model"
|
||||
assert "context_escalated" not in result.routing_decision
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_out_of_band_system_and_tools_count_against_the_window(self):
|
||||
"""The Claude Code shape that live-testing caught: a tiny ask riding a top-level
|
||||
`system` block and tool definitions that together dwarf the message list. None of
|
||||
that reaches resolved messages on /v1/messages, so a gate reading only messages
|
||||
dispatches a provably oversized request and the provider 400s anyway."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=_windowed_router(_SMALL, _BIG),
|
||||
complexity_router_config=_tier_config(),
|
||||
)
|
||||
|
||||
result = await router.async_pre_routing_hook(
|
||||
model="test-router",
|
||||
request_kwargs={
|
||||
"proxy_server_request": {
|
||||
"body": {
|
||||
"system": _CONTEXT_FILLER,
|
||||
"tools": [{"name": f"tool_{i}", "description": _CONTEXT_FILLER[:500]} for i in range(20)],
|
||||
}
|
||||
}
|
||||
},
|
||||
messages=[{"role": "user", "content": "reply with exactly: rig check ok"}],
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == "big-model"
|
||||
assert result.routing_decision["context_escalated"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_escalated_first_turn_never_becomes_the_session_pin(self):
|
||||
"""Escalation describes the prompt's size, not the session: once the client compacts,
|
||||
the next turn fits again, so pinning the big-window tier would hold the whole session
|
||||
on it for the TTL. The escalated turn routes big, and the next fitting turn classifies
|
||||
fresh instead of inheriting a pin."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=_windowed_router(_SMALL, _BIG),
|
||||
complexity_router_config=_tier_config(session_affinity=True),
|
||||
)
|
||||
session_kwargs = lambda: {"metadata": {"session_id": "s-1", "user_api_key_hash": "k-1"}} # noqa: E731
|
||||
|
||||
first = await router.async_pre_routing_hook(
|
||||
model="test-router", request_kwargs=session_kwargs(), messages=_OVERSIZED_TURNS
|
||||
)
|
||||
second = await router.async_pre_routing_hook(
|
||||
model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}]
|
||||
)
|
||||
|
||||
assert first is not None and first.model == "big-model"
|
||||
assert second is not None and second.model == "small-model"
|
||||
assert second.routing_decision["cause"] != "session_affinity_pin"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_pinned_session_escalates_per_request_and_keeps_its_pin(self):
|
||||
"""The pin fast path skips classification, not physics: an oversized turn on a session
|
||||
pinned to the small tier is served by the fitting tier, while the stored pin keeps the
|
||||
session's own model so the first turn that fits again routes exactly as pinned."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=_windowed_router(_SMALL, _BIG),
|
||||
complexity_router_config=_tier_config(session_affinity=True),
|
||||
)
|
||||
session_kwargs = lambda: {"metadata": {"session_id": "s-2", "user_api_key_hash": "k-2"}} # noqa: E731
|
||||
|
||||
pinned = await router.async_pre_routing_hook(
|
||||
model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}]
|
||||
)
|
||||
oversized = await router.async_pre_routing_hook(
|
||||
model="test-router", request_kwargs=session_kwargs(), messages=_OVERSIZED_TURNS
|
||||
)
|
||||
back_to_small = await router.async_pre_routing_hook(
|
||||
model="test-router", request_kwargs=session_kwargs(), messages=[{"role": "user", "content": "ok continue"}]
|
||||
)
|
||||
|
||||
assert pinned is not None and pinned.model == "small-model"
|
||||
assert oversized is not None and oversized.model == "big-model"
|
||||
assert oversized.routing_decision["cause"] == "session_affinity_pin"
|
||||
assert oversized.routing_decision["context_escalated"] is True
|
||||
assert oversized.routing_decision["context_escalation_original_tier"] == "SIMPLE"
|
||||
assert back_to_small is not None and back_to_small.model == "small-model"
|
||||
assert back_to_small.routing_decision["cause"] == "session_affinity_pin"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_adaptive_cold_start_never_samples_a_model_that_cannot_hold_the_prompt(self):
|
||||
"""The bandit's exploration is still bounded by physics: with the whole classified tier
|
||||
unobserved, cold start samples only among models whose window holds the prompt."""
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "small-model",
|
||||
"litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"},
|
||||
"model_info": {"max_input_tokens": 16385},
|
||||
},
|
||||
{
|
||||
"model_name": "mid-model",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"},
|
||||
"model_info": {"max_input_tokens": 200000},
|
||||
},
|
||||
]
|
||||
),
|
||||
complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["small-model", "mid-model"]}},
|
||||
)
|
||||
|
||||
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == "mid-model"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_gate_never_resolves_an_authenticating_provider(self, monkeypatch, tmp_path):
|
||||
"""Resolving github_copilot runs its OAuth device flow, so a window question must adopt
|
||||
the declaration instead of resolving: the copilot group reads as unknown-window and the
|
||||
request stays put, with zero copilot resolutions recorded."""
|
||||
import json
|
||||
import time
|
||||
|
||||
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
|
||||
(tmp_path / "api-key.json").write_text(json.dumps({"token": "tid=test", "expires_at": int(time.time()) + 3600}))
|
||||
router = ComplexityRouter(
|
||||
model_name="test-router",
|
||||
litellm_router_instance=Router(
|
||||
model_list=[
|
||||
{"model_name": "cop-pool", "litellm_params": {"model": "github_copilot/gpt-4o"}},
|
||||
{
|
||||
"model_name": "big-model",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"},
|
||||
"model_info": {"max_input_tokens": 200000},
|
||||
},
|
||||
]
|
||||
),
|
||||
complexity_router_config={"tiers": {"SIMPLE": "cop-pool", "COMPLEX": "big-model"}},
|
||||
)
|
||||
real_get_llm_provider = litellm.get_llm_provider
|
||||
copilot_resolutions: List = []
|
||||
|
||||
def _guarded(*args, **kwargs):
|
||||
target = str(kwargs.get("model") or (args[0] if args else "")) + str(kwargs.get("custom_llm_provider") or "")
|
||||
if "github_copilot" in target:
|
||||
copilot_resolutions.append(target)
|
||||
raise RuntimeError("the gate must not resolve an authenticating provider")
|
||||
return real_get_llm_provider(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(litellm, "get_llm_provider", _guarded)
|
||||
|
||||
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
|
||||
|
||||
assert result is not None
|
||||
assert result.model == "cop-pool"
|
||||
assert copilot_resolutions == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_full_routing_path_serves_the_escalated_deployment(self):
|
||||
"""End to end through Router.async_get_available_deployment: the auto-router alias with
|
||||
an oversized prompt resolves to the big tier's deployment, and a small prompt to the
|
||||
small tier's, with no mocking anywhere in the resolution chain."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}},
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "small-model",
|
||||
"litellm_params": {"model": "openai/gpt-3.5-turbo", "mock_response": "ok"},
|
||||
"model_info": {"max_input_tokens": 16385},
|
||||
},
|
||||
{
|
||||
"model_name": "big-model",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"},
|
||||
"model_info": {"max_input_tokens": 200000},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
oversized = await router.async_get_available_deployment(
|
||||
model="smart-router", request_kwargs={}, messages=_OVERSIZED_TURNS
|
||||
)
|
||||
small = await router.async_get_available_deployment(
|
||||
model="smart-router", request_kwargs={}, messages=[{"role": "user", "content": "ok continue"}]
|
||||
)
|
||||
|
||||
assert oversized["model_name"] == "big-model"
|
||||
assert small["model_name"] == "small-model"
|
||||
|
|
|
|||
|
|
@ -39,6 +39,35 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({
|
|||
})),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
|
||||
useInfiniteTeams: vi.fn(() => ({
|
||||
data: { pages: [{ teams: [{ team_id: "team-eng", team_alias: "engineering" }], page: 1, total_pages: 1 }] },
|
||||
isLoading: false,
|
||||
fetchNextPage: vi.fn(),
|
||||
hasNextPage: false,
|
||||
isFetchingNextPage: false,
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/users/useUsers", () => ({
|
||||
useInfiniteUsers: vi.fn(() => ({
|
||||
data: {
|
||||
pages: [
|
||||
{
|
||||
users: [{ user_id: "dev-alice", user_alias: null, user_email: "alice@example.com" }],
|
||||
page: 1,
|
||||
total_pages: 1,
|
||||
},
|
||||
],
|
||||
},
|
||||
isPending: false,
|
||||
isError: false,
|
||||
fetchNextPage: vi.fn(),
|
||||
hasNextPage: false,
|
||||
isFetchingNextPage: false,
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({
|
||||
useAutoRouters: vi.fn(() => ({
|
||||
data: [
|
||||
|
|
@ -60,7 +89,7 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({
|
|||
})),
|
||||
}));
|
||||
|
||||
import ShadowEvalSection, { shadowedKeyLabel } from "./ShadowEvalSection";
|
||||
import ShadowEvalSection, { shadowedTargetLabel } from "./ShadowEvalSection";
|
||||
import {
|
||||
useShadowEvalJob,
|
||||
useShadowEvalJobs,
|
||||
|
|
@ -77,14 +106,15 @@ const job = (overrides: Partial<ShadowEvalJob> = {}): ShadowEvalJob => ({
|
|||
baseline_model: null,
|
||||
judge_model: "anthropic/claude-sonnet-5",
|
||||
shadow_percentage: 10,
|
||||
keys: [
|
||||
targets: [
|
||||
{
|
||||
api_key_id: "hashed-key-abc",
|
||||
target_type: "key",
|
||||
target_id: "hashed-key-abc",
|
||||
max_turns: 10000,
|
||||
max_budget: 10,
|
||||
spend: 3.21,
|
||||
stopped_at: null,
|
||||
key_alias: "prod-alpha",
|
||||
target_alias: "prod-alpha",
|
||||
key_name: "sk-...alpha",
|
||||
},
|
||||
],
|
||||
|
|
@ -129,7 +159,6 @@ const job = (overrides: Partial<ShadowEvalJob> = {}): ShadowEvalJob => ({
|
|||
cache_hit_turns: 2,
|
||||
},
|
||||
],
|
||||
by_key: [],
|
||||
overall_shadow_win_rate_pct: 48.0,
|
||||
overall_tie_rate_pct: 22.0,
|
||||
sampled_real_spend: 0.6,
|
||||
|
|
@ -144,17 +173,18 @@ const job = (overrides: Partial<ShadowEvalJob> = {}): ShadowEvalJob => ({
|
|||
...overrides,
|
||||
});
|
||||
|
||||
const keyEntry = (
|
||||
api_key_id: string,
|
||||
overrides: Partial<ShadowEvalJob["keys"][number]> = {},
|
||||
): ShadowEvalJob["keys"][number] => ({
|
||||
api_key_id,
|
||||
const targetEntry = (
|
||||
target_id: string,
|
||||
overrides: Partial<ShadowEvalJob["targets"][number]> = {},
|
||||
): ShadowEvalJob["targets"][number] => ({
|
||||
target_type: "key",
|
||||
target_id,
|
||||
max_turns: 10000,
|
||||
max_budget: 10,
|
||||
spend: 0,
|
||||
stopped_at: null,
|
||||
attempt_count: null,
|
||||
key_alias: null,
|
||||
target_alias: null,
|
||||
key_name: null,
|
||||
...overrides,
|
||||
});
|
||||
|
|
@ -235,8 +265,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", keys: [keyEntry("key-a")] }),
|
||||
job({ job_id: "job-b", status: "running", keys: [keyEntry("key-b")] }),
|
||||
job({ job_id: "job-a", status: "running", targets: [targetEntry("key-a")] }),
|
||||
job({ job_id: "job-b", status: "running", targets: [targetEntry("key-b")] }),
|
||||
],
|
||||
});
|
||||
render(<ShadowEvalSection />);
|
||||
|
|
@ -347,7 +377,7 @@ describe("ShadowEvalSection", () => {
|
|||
});
|
||||
|
||||
it("shows spend without a budget cap for a job from before spend budgets existed", () => {
|
||||
const j = job({ keys: [keyEntry("hashed-key-abc", { max_budget: null, spend: 3.21 })] });
|
||||
const j = job({ targets: [targetEntry("hashed-key-abc", { max_budget: null, spend: 3.21 })] });
|
||||
mockHooks({ jobs: [j], detailsById: { "job-1": j } });
|
||||
render(<ShadowEvalSection />);
|
||||
expect(screen.getByText(/\$3\.21 eval spend/)).toBeInTheDocument();
|
||||
|
|
@ -417,6 +447,38 @@ describe("ShadowEvalSection", () => {
|
|||
|
||||
const expectedBody = {
|
||||
api_key_ids: ["hash-alpha", "hash-beta"],
|
||||
team_ids: [],
|
||||
user_ids: [],
|
||||
router_name: "gpt-auto",
|
||||
direction: "forward",
|
||||
shadow_percentage: 10,
|
||||
duration_days: 7,
|
||||
max_budget: 10,
|
||||
judge_model: "anthropic/claude-sonnet-5",
|
||||
};
|
||||
expect(start.mutate).toHaveBeenCalledWith(expectedBody);
|
||||
});
|
||||
|
||||
it("submits a team-only job with team_ids and no keys", async () => {
|
||||
const user = userEvent.setup();
|
||||
const { start } = mockHooks({});
|
||||
render(<ShadowEvalSection />);
|
||||
|
||||
expect(screen.getByText("Start shadow eval")).toBeDisabled();
|
||||
|
||||
await user.click(screen.getByPlaceholderText("Search teams by alias"));
|
||||
const teamList = await screen.findByTestId("paginated-multi-select-list");
|
||||
await user.click(within(teamList).getByText("engineering"));
|
||||
await user.click(screen.getByPlaceholderText("Select an auto-router"));
|
||||
await user.click(await screen.findByText("gpt-auto"));
|
||||
await user.click(screen.getByPlaceholderText("Select a judge model"));
|
||||
await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ }));
|
||||
await user.click(screen.getByText("Start shadow eval"));
|
||||
|
||||
const expectedBody = {
|
||||
api_key_ids: [],
|
||||
team_ids: ["team-eng"],
|
||||
user_ids: [],
|
||||
router_name: "gpt-auto",
|
||||
direction: "forward",
|
||||
shadow_percentage: 10,
|
||||
|
|
@ -453,6 +515,8 @@ describe("ShadowEvalSection", () => {
|
|||
|
||||
const expectedBody = {
|
||||
api_key_ids: ["hash-alpha"],
|
||||
team_ids: [],
|
||||
user_ids: [],
|
||||
router_name: "gpt-auto",
|
||||
direction: "reverse",
|
||||
baseline_model: "prod-claude",
|
||||
|
|
@ -485,9 +549,13 @@ describe("ShadowEvalSection", () => {
|
|||
});
|
||||
|
||||
it("labels the shadowed key by alias, then masked name, then truncated hash", () => {
|
||||
expect(shadowedKeyLabel(job().keys[0])).toBe("prod-alpha");
|
||||
expect(shadowedKeyLabel(keyEntry("hashed-key-abc", { key_name: "sk-...alpha" }))).toBe("sk-...alpha");
|
||||
expect(shadowedKeyLabel(keyEntry("hashed-key-abc"))).toBe("hashed-key…");
|
||||
expect(shadowedTargetLabel(job().targets[0])).toBe("prod-alpha");
|
||||
expect(shadowedTargetLabel(targetEntry("hashed-key-abc", { key_name: "sk-...alpha" }))).toBe("sk-...alpha");
|
||||
expect(shadowedTargetLabel(targetEntry("hashed-key-abc"))).toBe("hashed-key…");
|
||||
expect(shadowedTargetLabel(targetEntry("team-eng", { target_type: "team" }))).toBe("team-eng");
|
||||
expect(shadowedTargetLabel(targetEntry("team-eng", { target_type: "team", target_alias: "engineering" }))).toBe(
|
||||
"engineering",
|
||||
);
|
||||
});
|
||||
|
||||
it("breaks results down per key, so one key exhausting its own budget is visible while a sibling runs on", () => {
|
||||
|
|
@ -495,15 +563,12 @@ describe("ShadowEvalSection", () => {
|
|||
jobs: [
|
||||
job({
|
||||
judged_count: 205,
|
||||
keys: [
|
||||
keyEntry("hash-spent", { max_budget: 2, spend: 1.5, stopped_at: "2026-08-08T00:00:00Z" }),
|
||||
keyEntry("hash-hungry", { max_budget: 5, spend: 0.2 }),
|
||||
],
|
||||
results: {
|
||||
by_tier: [],
|
||||
by_current_model: [],
|
||||
by_key: [
|
||||
{
|
||||
targets: [
|
||||
targetEntry("hash-spent", {
|
||||
max_budget: 2,
|
||||
spend: 1.5,
|
||||
stopped_at: "2026-08-08T00:00:00Z",
|
||||
verdicts: {
|
||||
group: "hash-spent",
|
||||
turn_count: 200,
|
||||
real_win_rate_pct: 20.0,
|
||||
|
|
@ -514,7 +579,12 @@ describe("ShadowEvalSection", () => {
|
|||
shadow_spend: 0.5,
|
||||
cache_hit_turns: 0,
|
||||
},
|
||||
],
|
||||
}),
|
||||
targetEntry("hash-hungry", { max_budget: 5, spend: 0.2 }),
|
||||
],
|
||||
results: {
|
||||
by_tier: [],
|
||||
by_current_model: [],
|
||||
overall_shadow_win_rate_pct: 60.0,
|
||||
overall_tie_rate_pct: 20.0,
|
||||
sampled_real_spend: 0.9,
|
||||
|
|
@ -539,7 +609,7 @@ describe("ShadowEvalSection", () => {
|
|||
|
||||
expect(screen.getByText(/205 turns judged/)).toBeInTheDocument();
|
||||
expect(screen.getByText(/Shadowing 10% of/)).toBeInTheDocument();
|
||||
expect(screen.getByText("2 keys")).toBeInTheDocument();
|
||||
expect(screen.getByText("2 targets")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("reads a key that spent its budget as completed even before the sweep stamps it", () => {
|
||||
|
|
@ -547,9 +617,9 @@ describe("ShadowEvalSection", () => {
|
|||
mockHooks({
|
||||
jobs: [
|
||||
job({
|
||||
keys: [
|
||||
keyEntry("hash-spent", { max_budget: 2, spend: 2, attempt_count: 40 }),
|
||||
keyEntry("hash-hungry", legacyTurnBudgetLeg),
|
||||
targets: [
|
||||
targetEntry("hash-spent", { max_budget: 2, spend: 2, attempt_count: 40 }),
|
||||
targetEntry("hash-hungry", legacyTurnBudgetLeg),
|
||||
],
|
||||
}),
|
||||
],
|
||||
|
|
@ -571,9 +641,9 @@ describe("ShadowEvalSection", () => {
|
|||
job({
|
||||
judged_count: 0,
|
||||
results: null,
|
||||
keys: [
|
||||
keyEntry("hash-spent", { max_budget: 0.5, spend: 0.5, attempt_count: 2 }),
|
||||
keyEntry("hash-hungry", { max_budget: 5, spend: 0.01, attempt_count: 1 }),
|
||||
targets: [
|
||||
targetEntry("hash-spent", { max_budget: 0.5, spend: 0.5, attempt_count: 2 }),
|
||||
targetEntry("hash-hungry", { max_budget: 5, spend: 0.01, attempt_count: 1 }),
|
||||
],
|
||||
}),
|
||||
],
|
||||
|
|
@ -594,9 +664,9 @@ describe("ShadowEvalSection", () => {
|
|||
jobs: [
|
||||
job({
|
||||
status: "completed",
|
||||
keys: [
|
||||
keyEntry("hash-spent", { max_turns: 200, stopped_at: "2026-08-08T00:00:00Z" }),
|
||||
keyEntry("hash-hungry", { max_turns: 500 }),
|
||||
targets: [
|
||||
targetEntry("hash-spent", { max_turns: 200, stopped_at: "2026-08-08T00:00:00Z" }),
|
||||
targetEntry("hash-hungry", { max_turns: 500 }),
|
||||
],
|
||||
}),
|
||||
],
|
||||
|
|
|
|||
|
|
@ -3,10 +3,13 @@
|
|||
import React, { useMemo, useState } from "react";
|
||||
|
||||
import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
|
||||
import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers";
|
||||
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 { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect";
|
||||
import TeamMultiSelect from "@/components/common_components/team_multi_select";
|
||||
import { userOptionLabel } from "@/components/common_components/UserDropdown";
|
||||
import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
|
|
@ -27,7 +30,7 @@ import {
|
|||
useStartShadowEval,
|
||||
useStopShadowEval,
|
||||
type ShadowEvalJob,
|
||||
type ShadowEvalJobKey,
|
||||
type ShadowEvalJobTarget,
|
||||
type ShadowEvalSlice,
|
||||
} from "./useShadowEval";
|
||||
|
||||
|
|
@ -66,29 +69,31 @@ const routerMatchedOrBeatPct = (
|
|||
? 100 - results.overall_shadow_win_rate_pct
|
||||
: results.overall_shadow_win_rate_pct + results.overall_tie_rate_pct;
|
||||
|
||||
export const shadowedKeyLabel = (key: ShadowEvalJobKey): string =>
|
||||
key.key_alias || key.key_name || `${key.api_key_id.slice(0, 10)}…`;
|
||||
export const shadowedTargetLabel = (target: ShadowEvalJobTarget): string =>
|
||||
target.target_alias ||
|
||||
target.key_name ||
|
||||
(target.target_type === "key" ? `${target.target_id.slice(0, 10)}…` : target.target_id);
|
||||
|
||||
const shadowedKeysLabel = (job: ShadowEvalJob): string =>
|
||||
job.keys.length === 1 ? shadowedKeyLabel(job.keys[0]) : `${job.keys.length} keys`;
|
||||
const shadowedTargetsLabel = (job: ShadowEvalJob): string =>
|
||||
job.targets.length === 1 ? shadowedTargetLabel(job.targets[0]) : `${job.targets.length} targets`;
|
||||
|
||||
const totalBudget = (job: ShadowEvalJob): number | null =>
|
||||
job.keys.reduce<number | null>(
|
||||
(sum, key) => (sum === null || key.max_budget == null ? null : sum + key.max_budget),
|
||||
job.targets.reduce<number | null>(
|
||||
(sum, target) => (sum === null || target.max_budget == null ? null : sum + target.max_budget),
|
||||
0,
|
||||
);
|
||||
|
||||
const totalSpend = (job: ShadowEvalJob): number => job.keys.reduce((sum, key) => sum + (key.spend ?? 0), 0);
|
||||
const totalSpend = (job: ShadowEvalJob): number => job.targets.reduce((sum, target) => sum + (target.spend ?? 0), 0);
|
||||
|
||||
const keySpent = (key: ShadowEvalJobKey): boolean => {
|
||||
const spendBudgetReached = key.max_budget != null && key.spend != null && key.spend >= key.max_budget;
|
||||
const turnValveReached = key.attempt_count != null && key.attempt_count >= key.max_turns;
|
||||
const targetSpent = (target: ShadowEvalJobTarget): boolean => {
|
||||
const spendBudgetReached = target.max_budget != null && target.spend != null && target.spend >= target.max_budget;
|
||||
const turnValveReached = target.attempt_count != null && target.attempt_count >= target.max_turns;
|
||||
return spendBudgetReached || turnValveReached;
|
||||
};
|
||||
|
||||
const keyStatus = (job: ShadowEvalJob, key: ShadowEvalJobKey): string => {
|
||||
if (job.status === "completed" || (key.stopped_at == null && keySpent(key))) return "completed";
|
||||
return key.stopped_at != null ? "stopped" : "running";
|
||||
const targetStatus = (job: ShadowEvalJob, target: ShadowEvalJobTarget): string => {
|
||||
if (job.status === "completed" || (target.stopped_at == null && targetSpent(target))) return "completed";
|
||||
return target.stopped_at != null ? "stopped" : "running";
|
||||
};
|
||||
|
||||
const jobHeadline = (job: ShadowEvalJob): React.ReactNode =>
|
||||
|
|
@ -96,12 +101,12 @@ const jobHeadline = (job: ShadowEvalJob): React.ReactNode =>
|
|||
<>
|
||||
Comparing <span className="font-mono text-xs">{job.router_name}</span> to{" "}
|
||||
<span className="font-mono text-xs">{job.baseline_model}</span> on {job.shadow_percentage}% of{" "}
|
||||
<span className="font-mono text-xs">{shadowedKeysLabel(job)}</span> traffic
|
||||
<span className="font-mono text-xs">{shadowedTargetsLabel(job)}</span> traffic
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
Shadowing {job.shadow_percentage}% of <span className="font-mono text-xs">{shadowedKeysLabel(job)}</span> traffic
|
||||
via <span className="font-mono text-xs">{job.router_name}</span>
|
||||
Shadowing {job.shadow_percentage}% of <span className="font-mono text-xs">{shadowedTargetsLabel(job)}</span>{" "}
|
||||
traffic via <span className="font-mono text-xs">{job.router_name}</span>
|
||||
</>
|
||||
);
|
||||
|
||||
|
|
@ -255,13 +260,12 @@ const VerdictBar: React.FC<{ direction: ShadowEvalDirection; results: NonNullabl
|
|||
);
|
||||
};
|
||||
|
||||
const KeyTable: React.FC<{ job: ShadowEvalJob }> = ({ job }) => {
|
||||
const slices = new Map((job.results?.by_key ?? []).map((slice) => [slice.group, slice]));
|
||||
const TargetTable: React.FC<{ job: ShadowEvalJob }> = ({ job }) => {
|
||||
return (
|
||||
<Table>
|
||||
<TableHeader>
|
||||
<TableRow>
|
||||
<TableHead>Key</TableHead>
|
||||
<TableHead>Target</TableHead>
|
||||
<TableHead>Status</TableHead>
|
||||
{["Budget used", "Router wins", `${otherArmLabel(job.direction)} wins`].map((label) => (
|
||||
<TableHead key={label} className="text-right">
|
||||
|
|
@ -271,18 +275,23 @@ const KeyTable: React.FC<{ job: ShadowEvalJob }> = ({ job }) => {
|
|||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
{job.keys.map((key) => {
|
||||
const slice = slices.get(key.api_key_id);
|
||||
{job.targets.map((target) => {
|
||||
const slice = target.verdicts;
|
||||
return (
|
||||
<TableRow key={key.api_key_id}>
|
||||
<TableCell className="font-medium text-foreground">{shadowedKeyLabel(key)}</TableCell>
|
||||
<TableRow key={`${target.target_type}:${target.target_id}`}>
|
||||
<TableCell className="font-medium text-foreground">
|
||||
{shadowedTargetLabel(target)}
|
||||
{target.target_type !== "key" && (
|
||||
<span className="ml-2 text-xs font-normal text-muted-foreground">{target.target_type}</span>
|
||||
)}
|
||||
</TableCell>
|
||||
<TableCell>
|
||||
<StatusBadge status={keyStatus(job, key)} />
|
||||
<StatusBadge status={targetStatus(job, target)} />
|
||||
</TableCell>
|
||||
<TableCell className="text-right tabular-nums">
|
||||
{key.max_budget != null
|
||||
? `${usd(key.spend ?? 0)} / ${usd(key.max_budget)}`
|
||||
: `${(key.attempt_count ?? slice?.turn_count ?? 0).toLocaleString()} / ${key.max_turns.toLocaleString()} turns`}
|
||||
{target.max_budget != null
|
||||
? `${usd(target.spend ?? 0)} / ${usd(target.max_budget)}`
|
||||
: `${(target.attempt_count ?? slice?.turn_count ?? 0).toLocaleString()} / ${target.max_turns.toLocaleString()} turns`}
|
||||
</TableCell>
|
||||
{slice ? (
|
||||
<>
|
||||
|
|
@ -318,9 +327,9 @@ const ResultsBody: React.FC<{ job: ShadowEvalJob; resultsError?: boolean }> = ({
|
|||
const hasVerdicts = results != null && (results.by_tier.length > 0 || results.by_current_model.length > 0);
|
||||
return (
|
||||
<>
|
||||
{job.keys.length > 1 && (
|
||||
{job.targets.length > 1 && (
|
||||
<div className="border-b">
|
||||
<KeyTable job={job} />
|
||||
<TargetTable job={job} />
|
||||
</div>
|
||||
)}
|
||||
{/* results == null re-stated for TS narrowing; hasVerdicts alone cannot narrow it */}
|
||||
|
|
@ -454,9 +463,9 @@ const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[
|
|||
|
||||
const START_FORM_DESCRIPTION: Record<ShadowEvalDirection, string> = {
|
||||
forward:
|
||||
"Duplicates a sampled slice of the selected keys' traffic through the auto-router and has an LLM judge compare both answers blind. Each key gets its own spend budget. The router's answers are never served to users; judge calls bill to the shadowed key.",
|
||||
"Duplicates a sampled slice of the selected targets' traffic (keys, teams, or users) through the auto-router and has an LLM judge compare both answers blind. Each target gets its own spend budget. The router's answers are never served to users; judge calls bill to the sampled traffic's own identity.",
|
||||
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. Each key gets its own spend budget. 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 against a fixed baseline model and has an LLM judge compare both answers blind. Each target gets its own spend budget. The baseline's answers are never served to users; judge calls bill to the sampled traffic's own identity.",
|
||||
};
|
||||
|
||||
const DURATION_OPTIONS = [
|
||||
|
|
@ -515,9 +524,46 @@ const KeySelect: React.FC<{ value: string[]; onChange: (tokens: string[]) => voi
|
|||
);
|
||||
};
|
||||
|
||||
const UserSelect: React.FC<{ value: string[]; onChange: (ids: string[]) => void }> = ({ value, onChange }) => {
|
||||
const [search, setSearch] = useState("");
|
||||
const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteUsers(
|
||||
50,
|
||||
search || undefined,
|
||||
);
|
||||
const options = useMemo<SearchSelectOption[]>(
|
||||
() =>
|
||||
Array.from(
|
||||
new Map(
|
||||
(data?.pages ?? [])
|
||||
.flatMap((page) => page.users)
|
||||
.map((user) => [user.user_id, { label: userOptionLabel(user), value: user.user_id }] as const),
|
||||
).values(),
|
||||
),
|
||||
[data],
|
||||
);
|
||||
return (
|
||||
<PaginatedMultiSelect
|
||||
inputId="shadow-eval-user"
|
||||
options={options}
|
||||
value={value}
|
||||
onValueChange={onChange}
|
||||
onSearchChange={setSearch}
|
||||
onLoadMore={() => void fetchNextPage()}
|
||||
hasNextPage={hasNextPage}
|
||||
isFetchingNextPage={isFetchingNextPage}
|
||||
isLoading={isPending}
|
||||
placeholder="Search users by email"
|
||||
emptyText="No matching users"
|
||||
errorText={isError ? "Users could not be loaded. Refresh the page to retry." : undefined}
|
||||
/>
|
||||
);
|
||||
};
|
||||
|
||||
const StartForm: React.FC = () => {
|
||||
const { accessToken } = useAuthorized();
|
||||
const [apiKeyIds, setApiKeyIds] = useState<string[]>([]);
|
||||
const [teamIds, setTeamIds] = useState<string[]>([]);
|
||||
const [userIds, setUserIds] = useState<string[]>([]);
|
||||
const [routerName, setRouterName] = useState("");
|
||||
const [direction, setDirection] = useState<ShadowEvalDirection>("forward");
|
||||
const [baselineModel, setBaselineModel] = useState("");
|
||||
|
|
@ -542,12 +588,15 @@ const StartForm: React.FC = () => {
|
|||
const parsedMaxBudget = Number.parseFloat(maxBudget);
|
||||
const maxBudgetValid = parsedMaxBudget >= 0.01 && parsedMaxBudget <= 10000;
|
||||
const baselinePicked = direction === "forward" || baselineModel !== "";
|
||||
const filled = apiKeyIds.length > 0 && [routerName, judgeModel].every((field) => field !== "") && baselinePicked;
|
||||
const targetsPicked = apiKeyIds.length + teamIds.length + userIds.length > 0;
|
||||
const filled = targetsPicked && [routerName, judgeModel].every((field) => field !== "") && baselinePicked;
|
||||
const boundsValid = percentageValid && maxBudgetValid;
|
||||
const valid = Boolean(accessToken) && filled && boundsValid;
|
||||
const handleStart = () => {
|
||||
const startBody = {
|
||||
api_key_ids: apiKeyIds,
|
||||
team_ids: teamIds,
|
||||
user_ids: userIds,
|
||||
router_name: routerName,
|
||||
direction,
|
||||
...(direction === "reverse" ? { baseline_model: baselineModel } : {}),
|
||||
|
|
@ -587,6 +636,12 @@ const StartForm: React.FC = () => {
|
|||
<Field label="Keys to shadow" htmlFor="shadow-eval-key">
|
||||
<KeySelect value={apiKeyIds} onChange={setApiKeyIds} />
|
||||
</Field>
|
||||
<Field label="Teams to shadow">
|
||||
<TeamMultiSelect value={teamIds} onChange={setTeamIds} placeholder="Search teams by alias" />
|
||||
</Field>
|
||||
<Field label="Users to shadow" htmlFor="shadow-eval-user">
|
||||
<UserSelect value={userIds} onChange={setUserIds} />
|
||||
</Field>
|
||||
<Field label="Auto-router">
|
||||
<SearchSelect
|
||||
options={routerOptions}
|
||||
|
|
@ -642,7 +697,7 @@ const StartForm: React.FC = () => {
|
|||
value={maxBudget}
|
||||
onChange={(e) => setMaxBudget(e.target.value)}
|
||||
/>
|
||||
<span className="text-sm text-muted-foreground">max shadow + judge spend, per key</span>
|
||||
<span className="text-sm text-muted-foreground">max shadow + judge spend, per target</span>
|
||||
</div>
|
||||
{maxBudget.trim() !== "" && !maxBudgetValid && (
|
||||
<p className="text-xs text-destructive">Enter a value from 0.01 to 10000</p>
|
||||
|
|
@ -774,8 +829,9 @@ const ShadowEvalSection: React.FC = () => {
|
|||
<div className="flex flex-wrap items-baseline gap-2">
|
||||
<h2 className="text-xl font-semibold text-foreground">Shadow eval</h2>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
Blind-judge the auto-router on your real traffic: against the models a key uses today before switching, or
|
||||
against a fixed baseline after it has switched.
|
||||
Blind-judge the auto-router on the real traffic of a key, team, or user (teams and users cover
|
||||
JWT-authenticated traffic): against the models they use today before switching, or against a fixed baseline
|
||||
after they have switched.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +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 ShadowEvalJobTarget = components["schemas"]["ShadowEvalJobTargetResponse"];
|
||||
export type ShadowEvalSlice = components["schemas"]["ShadowEvalSlice"];
|
||||
export type StartShadowEvalRequest = components["schemas"]["StartShadowEvalRequest"];
|
||||
|
||||
|
|
|
|||
219
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
219
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -1234,8 +1234,8 @@ export interface paths {
|
|||
};
|
||||
/**
|
||||
* List Shadow Eval Jobs
|
||||
* @description List shadow eval jobs, newest first, each key with its attempt count so status is
|
||||
* accurate. Judged counts, spend, and results ride the detail endpoint only.
|
||||
* @description List shadow eval jobs, newest first, each target with its attempt count so status
|
||||
* is accurate. Judged counts, spend, and results ride the detail endpoint only.
|
||||
*/
|
||||
get: operations["list_shadow_eval_jobs_auto_router_shadow_eval_get"];
|
||||
put?: never;
|
||||
|
|
@ -1257,22 +1257,29 @@ export interface paths {
|
|||
put?: never;
|
||||
/**
|
||||
* Start Shadow Eval
|
||||
* @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.
|
||||
* @description Start a shadow eval: duplicate a sampled slice of one or more targets' 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 target.
|
||||
*
|
||||
* 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.
|
||||
* A target is a virtual key, a team, or a user. Team and user targets match on the
|
||||
* identity every request resolves to at auth time, so they cover JWT-authenticated
|
||||
* traffic, which presents no virtual key; a user target samples that user's traffic
|
||||
* across all their teams, whether it arrives on a JWT or a key they own.
|
||||
*
|
||||
* Shadow responses are never served to users. Each key samples until its recorded eval
|
||||
* spend, the shadow and judge calls' own cost, reaches max_budget dollars, the job's
|
||||
* window ends, or the job is stopped, so one 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.
|
||||
* A forward job answers whether the targets should adopt router_name: it samples the
|
||||
* requests the router did not serve and duplicates them through it. A reverse job
|
||||
* answers whether a target already on the router still gains from it: it samples the
|
||||
* requests the router did serve and duplicates them against baseline_model. A target
|
||||
* can hold one active job per direction, so both questions can run at once, and a
|
||||
* request matching several jobs' targets (say its key and its team) is sampled by
|
||||
* each, separately budgeted.
|
||||
*
|
||||
* Shadow responses are never served to users. Each target samples until its recorded
|
||||
* eval spend, the shadow and judge calls' own cost, reaches max_budget dollars, the
|
||||
* job's window ends, or the job is stopped, so one target 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 sampled request's own identity but are
|
||||
* excluded from request counts and auto-router adoption metrics.
|
||||
*/
|
||||
post: operations["start_shadow_eval_auto_router_shadow_eval_start_post"];
|
||||
delete?: never;
|
||||
|
|
@ -1312,8 +1319,8 @@ export interface paths {
|
|||
put?: never;
|
||||
/**
|
||||
* Stop Shadow Eval Job
|
||||
* @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
|
||||
* @description Stop an active shadow eval job, every target it scopes at once. Attempts are kept;
|
||||
* sampling halts within ~10s. Targets that already stopped on their own budget keep the
|
||||
* stopped_at they earned. The statement is the whole state machine: it claims the job
|
||||
* only while a leg still samples inside the window with no stop recorded, so a racing
|
||||
* operator, a same-instant budget spend, and a repeat stop all read the same 400 with
|
||||
|
|
@ -34385,6 +34392,12 @@ export interface components {
|
|||
* @description Keywords indicating code-related content
|
||||
*/
|
||||
code_keywords?: string[] | null;
|
||||
/**
|
||||
* Context Window Escalation Buffer
|
||||
* @description Fraction of a model's declared context window the estimated prompt must fit within. The token count is an estimate, so fitting against the full window would dispatch prompts that the provider's own tokenizer then rejects; 0.95 leaves room for that drift plus the response tokens.
|
||||
* @default 0.95
|
||||
*/
|
||||
context_window_escalation_buffer: number;
|
||||
/**
|
||||
* Custom Technical Keywords
|
||||
* @description Domain-specific technical keywords appended to the effective base list (technical_keywords if set, otherwise DEFAULT_TECHNICAL_KEYWORDS). Order is preserved; duplicates are removed case-insensitively against the base list and within this list.
|
||||
|
|
@ -34413,6 +34426,12 @@ export interface components {
|
|||
* @description Embedding model (LiteLLM model name) used when semantic_keyword_matching is enabled
|
||||
*/
|
||||
embedding_model?: string | null;
|
||||
/**
|
||||
* Enable Context Window Escalation
|
||||
* @description Escalate a request off a tier whose models provably cannot hold its prompt, before dispatch. The classifier scores complexity and never prompt size, so a long agentic session whose newest ask is trivial lands on a small-window tier and the provider rejects it with a context-window 400 that nothing retries. When every model of the decided tier has a declared window smaller than the estimated prompt, the request moves to the lowest configured tier with a model whose declared window fits; when only some of the tier's models fit, the pick is restricted to those and the tier keeps the request. Models with no resolvable window are never escalated away from and never escalated onto. Set false to dispatch on complexity alone, as before.
|
||||
* @default true
|
||||
*/
|
||||
enable_context_window_escalation: boolean;
|
||||
/**
|
||||
* Escalation Keywords
|
||||
* @description Case-sensitive phrases a user can include to force a bump to the next-higher complexity tier when they aren't satisfied with results (they can force a stronger model, but not choose which one). Defaults to ['LITELLM ESCALATE'] when unset; set to an empty list to disable.
|
||||
|
|
@ -35250,56 +35269,10 @@ 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 entry scopes
|
||||
*/
|
||||
api_key_id: string;
|
||||
/**
|
||||
* Attempt Count
|
||||
* @description This key's sampled attempts so far, judged and errored alike, the same count the sampler budgets against max_turns; populated on list and detail responses. Frozen at stopped_at once the key is stamped, so in-flight attempts landing after a stop never reclassify it
|
||||
*/
|
||||
attempt_count?: number | null;
|
||||
/**
|
||||
* Key Alias
|
||||
* @description Alias of the shadowed key, resolved from the key row at read time; None when unset or deleted
|
||||
*/
|
||||
key_alias?: string | null;
|
||||
/**
|
||||
* Key Name
|
||||
* @description Masked display name (sk-...) of the shadowed key, resolved at read time like key_alias
|
||||
*/
|
||||
key_name?: string | null;
|
||||
/**
|
||||
* Max Budget
|
||||
* @description This key's own USD budget for the eval's shadow and judge spend, independent of its siblings'; None on jobs created before spend budgets existed, which max_turns alone bounds
|
||||
*/
|
||||
max_budget?: number | null;
|
||||
/**
|
||||
* Max Turns
|
||||
* @description This key's sample-count ceiling: the whole budget for jobs created before max_budget existed, and the error-loop safety valve otherwise
|
||||
*/
|
||||
max_turns: number;
|
||||
/**
|
||||
* Spend
|
||||
* @description This key's recorded shadow plus judge spend in USD, the same figure the sampler budgets against max_budget; populated on list and detail responses and frozen at stopped_at exactly like attempt_count
|
||||
*/
|
||||
spend?: number | null;
|
||||
/**
|
||||
* Stopped At
|
||||
* @description When this key's slot was stamped free, whether its own budget ran out, the window closed, or an operator stopped the job; status is derived, so a spent budget reads completed even while this is still unset
|
||||
*/
|
||||
stopped_at?: string | null;
|
||||
};
|
||||
/**
|
||||
* ShadowEvalJobResponse
|
||||
* @description A shadow-eval job over one or more keys, each with its own budget and stop state;
|
||||
* status is derived from stopped_by, the keys' stop and budget state, and ends_at,
|
||||
* @description A shadow-eval job over one or more targets, each with its own budget and stop state;
|
||||
* status is derived from stopped_by, the targets' stop and budget state, 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.
|
||||
*/
|
||||
|
|
@ -35341,11 +35314,6 @@ 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
|
||||
|
|
@ -35361,8 +35329,8 @@ export interface components {
|
|||
* Status
|
||||
* @description Three recorded facts, no history-guessing: a stop is stopped_by (the migration
|
||||
* backfills it for every job that displayed stopped when the column arrived, so the
|
||||
* pre-column population is closed), completion is the window passing or every key
|
||||
* spending its budget, and anything else is running. The all-keys-stamped fallback
|
||||
* pre-column population is closed), completion is the window passing or every target
|
||||
* spending its budget, and anything else is running. The all-targets-stamped fallback
|
||||
* covers only stops written by pre-column pods during a rolling deploy.
|
||||
* @enum {string}
|
||||
*/
|
||||
|
|
@ -35372,6 +35340,65 @@ export interface components {
|
|||
* @description The operator who stopped the job early, recorded by the stop endpoint; 'unknown' backfilled by migration for jobs that displayed stopped when the column arrived; None when the job ended on its own. Its presence is what makes a job read stopped rather than completed
|
||||
*/
|
||||
stopped_by?: string | null;
|
||||
/**
|
||||
* Targets
|
||||
* @description The targets whose traffic this job evaluates, and only theirs, each with its own budget
|
||||
*/
|
||||
targets: components["schemas"]["ShadowEvalJobTargetResponse"][];
|
||||
};
|
||||
/**
|
||||
* ShadowEvalJobTargetResponse
|
||||
* @description One target a job shadows (a key, team, or user), with its own budget and stop state.
|
||||
*/
|
||||
ShadowEvalJobTargetResponse: {
|
||||
/**
|
||||
* Attempt Count
|
||||
* @description This target's sampled attempts so far, judged and errored alike, the same count the sampler budgets against max_turns; populated on list and detail responses. Frozen at stopped_at once the target is stamped, so in-flight attempts landing after a stop never reclassify it
|
||||
*/
|
||||
attempt_count?: number | null;
|
||||
/**
|
||||
* Key Name
|
||||
* @description Masked display name (sk-...) for key targets, resolved at read time; None for teams and users
|
||||
*/
|
||||
key_name?: string | null;
|
||||
/**
|
||||
* Max Budget
|
||||
* @description This target's own USD budget for the eval's shadow and judge spend, independent of its siblings'; None on jobs created before spend budgets existed, which max_turns alone bounds
|
||||
*/
|
||||
max_budget?: number | null;
|
||||
/**
|
||||
* Max Turns
|
||||
* @description This target's sample-count ceiling: the whole budget for jobs created before max_budget existed, and the error-loop safety valve otherwise
|
||||
*/
|
||||
max_turns: number;
|
||||
/**
|
||||
* Spend
|
||||
* @description This target's recorded shadow plus judge spend in USD, the same figure the sampler budgets against max_budget; populated on list and detail responses and frozen at stopped_at exactly like attempt_count
|
||||
*/
|
||||
spend?: number | null;
|
||||
/**
|
||||
* Stopped At
|
||||
* @description When this target's slot was stamped free, whether its own budget ran out, the window closed, or an operator stopped the job; status is derived, so a spent budget reads completed even while this is still unset
|
||||
*/
|
||||
stopped_at?: string | null;
|
||||
/**
|
||||
* Target Alias
|
||||
* @description Display label resolved from the target's own row at read time: the key's alias, the team's alias, or the user's email; None when unset or deleted
|
||||
*/
|
||||
target_alias?: string | null;
|
||||
/**
|
||||
* Target Id
|
||||
* @description The hashed virtual key, team id, or user id whose traffic this entry scopes
|
||||
*/
|
||||
target_id: string;
|
||||
/**
|
||||
* Target Type
|
||||
* @description What kind of entity this entry scopes
|
||||
* @enum {string}
|
||||
*/
|
||||
target_type: "key" | "team" | "user";
|
||||
/** @description This target's own judged-verdict slice; detail endpoint only, None until a turn is judged */
|
||||
verdicts?: components["schemas"]["ShadowEvalSlice"] | null;
|
||||
};
|
||||
/**
|
||||
* ShadowEvalResult
|
||||
|
|
@ -35383,11 +35410,6 @@ export interface components {
|
|||
* @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 judged a turn for yet are absent rather than reported as zero
|
||||
*/
|
||||
by_key: components["schemas"]["ShadowEvalSlice"][];
|
||||
/** By Tier */
|
||||
by_tier: components["schemas"]["ShadowEvalSlice"][];
|
||||
/**
|
||||
|
|
@ -35429,8 +35451,9 @@ 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 scoped target (embedded on that target's
|
||||
* own entry, so slices never need re-joining to a target by id).
|
||||
*/
|
||||
ShadowEvalSlice: {
|
||||
/** Avg Judge Confidence */
|
||||
|
|
@ -35602,6 +35625,10 @@ export interface components {
|
|||
classifier_cost?: number;
|
||||
/** Classifier Model */
|
||||
classifier_model?: string;
|
||||
/** Context Escalated */
|
||||
context_escalated?: boolean;
|
||||
/** Context Escalation Original Tier */
|
||||
context_escalation_original_tier?: string;
|
||||
/** Conversation Continuing */
|
||||
conversation_continuing?: boolean;
|
||||
/** Escalated */
|
||||
|
|
@ -35656,12 +35683,18 @@ export interface components {
|
|||
};
|
||||
/**
|
||||
* StartShadowEvalRequest
|
||||
* @description Start duplicating one or more keys' traffic for blind comparison against an auto-router.
|
||||
* @description Start duplicating one or more targets' traffic for blind comparison against an auto-router.
|
||||
*
|
||||
* A target is a virtual key, a team, or a user; each becomes its own leg with its own
|
||||
* budget and stop state. Team and user targets match on the identity every request
|
||||
* carries after auth (user_api_key_team_id / user_api_key_user_id), so they cover
|
||||
* JWT-authenticated traffic, which presents no virtual key at all.
|
||||
*/
|
||||
StartShadowEvalRequest: {
|
||||
/**
|
||||
* 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 carries its own max_budget spend budget, so one key exhausting its budget leaves the others sampling. At most 100 keys per job, which also bounds every read the job's endpoints make.
|
||||
* @description Hashed virtual keys whose traffic will be shadowed. Combined with team_ids and user_ids the job needs at least one target and at most 100, which also bounds every read the job's endpoints make. Each target carries its own max_budget spend budget, so one exhausting its budget leaves the others sampling.
|
||||
* @default []
|
||||
*/
|
||||
api_key_ids: string[];
|
||||
/**
|
||||
|
|
@ -35690,7 +35723,7 @@ export interface components {
|
|||
judge_model: string;
|
||||
/**
|
||||
* Max Budget
|
||||
* @description Per-key USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with the same figures the spend pipeline bills. EACH scoped key samples until its recorded eval spend reaches this, so a job over N keys spends at most about N times max_budget; in-flight samples can overshoot the cap by one sampling cache window
|
||||
* @description Per-target USD budget for the eval's own overhead, the shadow-arm and judge calls, priced with the same figures the spend pipeline bills. EACH scoped target samples until its recorded eval spend reaches this, so a job over N targets spends at most about N times max_budget; in-flight samples can overshoot the cap by one sampling cache window
|
||||
* @default 10
|
||||
*/
|
||||
max_budget: number;
|
||||
|
|
@ -35701,9 +35734,21 @@ export interface components {
|
|||
router_name: string;
|
||||
/**
|
||||
* Shadow Percentage
|
||||
* @description Percentage of the key's requests to duplicate through the router
|
||||
* @description Percentage of each target's requests to duplicate through the router
|
||||
*/
|
||||
shadow_percentage: number;
|
||||
/**
|
||||
* Team Ids
|
||||
* @description Teams whose traffic will be shadowed, matched on the team every authenticated request resolves to, so a team's JWT-auth and virtual-key traffic are both sampled
|
||||
* @default []
|
||||
*/
|
||||
team_ids: string[];
|
||||
/**
|
||||
* User Ids
|
||||
* @description Users whose traffic will be shadowed, matched on the user every authenticated request resolves to across all their teams: JWT requests carrying their subject claim and virtual keys they own
|
||||
* @default []
|
||||
*/
|
||||
user_ids: string[];
|
||||
};
|
||||
/**
|
||||
* SuccessfulKeyUpdate
|
||||
|
|
@ -40665,8 +40710,10 @@ export interface operations {
|
|||
list_shadow_eval_jobs_auto_router_shadow_eval_get: {
|
||||
parameters: {
|
||||
query?: {
|
||||
/** @description Filter to jobs that shadow this key, alone or alongside others */
|
||||
api_key_id?: string | null;
|
||||
/** @description Kind of target to filter on; requires target_id */
|
||||
target_type?: ("key" | "team" | "user") | null;
|
||||
/** @description Filter to jobs that shadow this target, alone or alongside others */
|
||||
target_id?: string | null;
|
||||
/** @description Newest jobs to return */
|
||||
limit?: number;
|
||||
};
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue