fix(proxy): finish LDAP revocations safely and skip malformed groups

This commit is contained in:
Mrinal Chanshetty 2026-11-05 09:02:04 -08:00
commit ed91537b36
729 changed files with 108896 additions and 11582 deletions

View file

@ -419,6 +419,7 @@ jobs:
- shard: router
test-path: >-
tests/unit/router_rust
tests/unit/test_router
python-version: "3.12"
workers: 4

View file

@ -0,0 +1,60 @@
# Manual skill security scans
This standalone adapter runs NVIDIA SkillSpector on a separate machine. It is not installed in the LiteLLM proxy and does not require a shared Helm or Compose deployment
Start it with Python 3.12 or newer and uv. The script pins SkillSpector and its HTTP dependencies
```sh
export SKILL_SCANNER_API_KEY="$(openssl rand -hex 32)"
uv run server.py
```
For another host, place the loopback service behind an HTTPS reverse proxy that forwards the Authorization header and enforces a 65 MiB request limit. Run it under an unprivileged account with a writable temporary directory and no production files or credentials. The scanner host needs Git and network access to requested sources and optional model providers. Apply network egress restrictions suitable for those sources
Configure the same scanner secret in the gateway's environment, then add this to LiteLLM's configuration
```yaml
general_settings:
skill_scanner:
endpoint: https://scanner.example.com/scan
api_key_env: SKILL_SCANNER_API_KEY
use_llm: false
timeout_seconds: 180
```
The endpoint must be reachable from the gateway. A customer's laptop works while it is reachable through their network, VPN or authenticated HTTPS tunnel. `localhost` refers to the gateway machine, not the admin's browser. HTTP is accepted only for loopback development
Open Skills, Submissions, and Review as a proxy admin. Click **Scan for threats**, inspect the report, then independently approve or reject the proposal. Scanning is optional and never changes approval permissions or blocks publication. The report shows findings, locations, remediation, risk score, scanner version and whether semantic analysis actually ran. Incomplete reports do not establish safety
Uploaded ZIPs are sent privately to the scanner with their exact stored bytes. Metadata-only edits reuse the matching stored archive. If that archive is missing or changed, submit a new ZIP before scanning. The displayed SHA-256 identifies that archive, and the scan examines its files rather than proposed catalog metadata. Remote URL reports describe the source at scan time and do not freeze its contents. GitHub branch/tag refs and explicit-ref subdirectories are supported. Other ref/subdirectory combinations require a ZIP so the adapter does not silently scan the wrong revision
The latest scan result is retained with the submission. Running scans have a deadline. If a gateway restarts or the scanner becomes unavailable, the attempt expires and an admin can retry. Retries replace the latest report. Approving or rejecting during a scan does not cancel it or change the decision
Set `use_llm: true` to request semantic analysis and configure the provider on the scanner machine using [SkillSpector's provider instructions](https://github.com/NVIDIA/skillspector#llm-analysis). Provider credentials stay on that machine. Static analysis can also call OSV for dependency checks
## HTTP contract
The gateway sends an authenticated `POST` to the configured endpoint. The body contains form fields `source` (JSON source descriptor), `use_llm` (boolean), and `expires_at` (timezone-aware ISO timestamp), plus `file` (ZIP) when bytes are available. The service returns HTTP 200 with this JSON shape
```json
{
"scanner_version": "2.12.0",
"complete": true,
"scan_mode": "static-only",
"risk_score": 70,
"findings": [{
"rule_id": "example-rule",
"severity": "HIGH",
"message": "Potential data exfiltration",
"file": "SKILL.md",
"start_line": 8,
"explanation": "Instructions request sending local files to an external host",
"remediation": "Remove the external upload instruction"
}],
"limitations": []
}
```
Reports are limited to 1 MiB and 500 findings. Scanner uploads allow 64 MiB because the gateway can expand compressed skill uploads into stored compatibility ZIPs. The gateway rejects malformed reports, refuses redirects and verifies TLS. Only the configured scanner credential is sent. A timeout, HTTP error or incomplete report remains visible to the reviewer
The bundled adapter runs at most two child scans concurrently. Each child and its descendants are terminated at the gateway deadline, capped at 570 seconds. Keep the gateway and scanner clocks synchronized. Direct adapter clients may omit `expires_at` to use the cap. Child processes never execute submitted skill scripts. Raise the gateway deadline up to 600 seconds when needed for semantic analysis and configure the reverse proxy timeout to match

View file

@ -0,0 +1,172 @@
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "fastapi==0.162.0", "uvicorn==0.54.0", "python-multipart==0.0.33",
# "skillspector @ git+https://github.com/NVIDIA/skillspector.git@62c2fe8622f580ac292d4078c3b71f790194dbbd",
# ]
# ///
import asyncio
import hmac
import json
import os
import signal
import sys
import tempfile
from contextlib import suppress
from datetime import datetime, timezone
from pathlib import Path
from typing import Annotated, Final
from urllib.parse import quote, urlsplit
from fastapi import Depends, FastAPI, File, Form, HTTPException, UploadFile
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import AwareDatetime, TypeAdapter, ValidationError
MAX_UPLOAD: Final = 64 * 1024 * 1024
MAX_REPORT: Final = 1024 * 1024
SCAN_TIMEOUT: Final = 570
SOURCE: Final = TypeAdapter(dict[str, str])
SLOTS: Final = asyncio.Semaphore(2)
app: Final = FastAPI(title="SkillSpector adapter", docs_url=None, redoc_url=None, openapi_url=None)
def authenticate(credentials: Annotated[HTTPAuthorizationCredentials, Depends(HTTPBearer())]) -> None:
token: Final = os.environ.get("SKILL_SCANNER_API_KEY", "")
if not token or not hmac.compare_digest(credentials.credentials, token):
raise HTTPException(401, "Invalid scanner credentials")
def source_target(source: dict[str, str]) -> str:
target: Final = (
"https://github.com/" + source.get("repo", "") if source.get("source") == "github" else source.get("url", "")
)
parsed: Final = urlsplit(target)
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
raise HTTPException(400, "Remote sources require an HTTPS URL without credentials or a fragment")
ref: Final = source.get("ref")
subdir: Final = source.get("path")
if ref or subdir:
if parsed.hostname != "github.com" or len(parsed.path.strip("/").split("/")) != 2 or not ref:
raise HTTPException(400, "Upload a ZIP to scan this source's exact ref and subdirectory")
suffix: Final = "/" + quote(subdir, safe="/") if subdir else ""
return target.removesuffix(".git").rstrip("/") + "/tree/" + quote(ref, safe="/") + suffix
return target
async def worker(target: str, output: Path, use_llm: bool) -> None:
from skillspector.mcp_server import run_scan
result: Final = await run_scan(target, use_llm=use_llm, allow_local_targets=not target.startswith("https://"))
coverage: Final = result.get("analysis_completeness", {})
findings: Final = result["findings"]
complete: Final = (
result.get("execution_successful") is True
and coverage.get("is_complete") is True
and (not use_llm or (result.get("llm_available") is True and result.get("llm_used") is True))
and len(findings) <= 500
)
limitations: Final = tuple(str(value)[:2000] for value in coverage.get("limitations", ()))
report: Final = {
"scanner_version": result["version"],
"complete": complete,
"scan_mode": result["scan_mode"],
"risk_score": result["risk_score"],
"findings": [
{
"rule_id": finding["id"],
"severity": finding["severity"],
"message": finding.get("pattern") or finding.get("finding") or "Security finding",
"file": finding["location"]["file"],
"start_line": finding["location"].get("start_line"),
"explanation": finding.get("explanation"),
"remediation": finding.get("remediation"),
}
for finding in findings[:500]
],
"limitations": (
*limitations,
*(("Only the first 500 findings are displayed.",) if len(findings) > 500 else ()),
),
}
encoded: Final = json.dumps(report).encode()
if len(encoded) > MAX_REPORT:
raise ValueError("Scanner report exceeded the response limit")
output.write_bytes(encoded)
async def run_worker(target: str, output: Path, use_llm: bool, expires_at: datetime | None = None) -> dict[str, object]:
timeout: Final = (
SCAN_TIMEOUT
if expires_at is None
else min(SCAN_TIMEOUT, (expires_at - datetime.now(timezone.utc)).total_seconds())
)
if timeout <= 0:
raise HTTPException(504, "SkillSpector scan deadline expired")
process: Final = await asyncio.create_subprocess_exec(
sys.executable,
str(Path(__file__).resolve()),
"--scan",
target,
str(output),
str(use_llm).lower(),
stdout=asyncio.subprocess.DEVNULL,
stderr=asyncio.subprocess.DEVNULL,
start_new_session=True,
env={
**os.environ,
"HOME": str(output.parent),
"TMPDIR": str(output.parent),
"GIT_CONFIG_GLOBAL": os.devnull,
"GIT_CONFIG_NOSYSTEM": "1",
"GIT_TERMINAL_PROMPT": "0",
"GIT_ALLOW_PROTOCOL": "https",
},
)
try:
code: Final = await asyncio.wait_for(process.wait(), timeout)
if code != 0 or not output.is_file() or output.stat().st_size > MAX_REPORT:
raise HTTPException(502, "SkillSpector could not produce a report for this source")
return TypeAdapter(dict[str, object]).validate_json(output.read_bytes())
except TimeoutError:
raise HTTPException(504, "SkillSpector scan timed out") from None
finally:
with suppress(ProcessLookupError):
os.killpg(process.pid, signal.SIGKILL)
await process.wait()
@app.post("/scan", dependencies=[Depends(authenticate)])
async def scan(
source: Annotated[str, Form(max_length=16384)],
use_llm: Annotated[bool, Form()] = False,
expires_at: Annotated[AwareDatetime | None, Form()] = None,
file: Annotated[UploadFile | None, File()] = None,
) -> dict[str, object]:
if SLOTS.locked():
raise HTTPException(429, "Scanner is busy; try again shortly")
try:
parsed: Final = SOURCE.validate_json(source)
except ValidationError:
raise HTTPException(400, "Invalid source metadata") from None
async with SLOTS:
with tempfile.TemporaryDirectory(prefix="skill-scan-") as directory:
root: Final = Path(directory)
if file is not None:
content: Final = await file.read(MAX_UPLOAD + 1)
if len(content) > MAX_UPLOAD:
raise HTTPException(413, "Upload exceeds 64 MiB")
archive: Final = root / "skill.zip"
archive.write_bytes(content)
return await run_worker(str(archive), root / "report.json", use_llm, expires_at)
return await run_worker(source_target(parsed), root / "report.json", use_llm, expires_at)
if __name__ == "__main__":
if len(sys.argv) > 1 and sys.argv[1] == "--scan":
asyncio.run(worker(sys.argv[2], Path(sys.argv[3]), sys.argv[4] == "true"))
else:
import uvicorn
if not os.environ.get("SKILL_SCANNER_API_KEY"):
raise SystemExit("Set SKILL_SCANNER_API_KEY before starting the scanner")
uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("PORT", "8099")))

View file

@ -50,8 +50,8 @@ LDAP is authoritative for the membership of every matched team, including manual
A group no longer returned by LDAP leaves its team and memberships unchanged. A successfully read group with no members removes all memberships from its corresponding team. No team or user is ever deleted, and no user is provisioned by this sync.
Failed or incomplete group searches, referrals, duplicate group IDs, invalid responses and LDAP operation exceptions abort the directory snapshot without applying it. A member DN whose entry no longer exists skips only its group, so other groups still synchronize and that team keeps its current members until the DN is removed or restored. Unreadable entries still abort the snapshot. Other member validation failures, such as ambiguous identities or nested groups, also skip their affected group. A membership-write failure for one valid group does not prevent other valid groups from synchronizing. A user disappearing during the membership write causes that group's transaction to make no changes. Repeated successful runs with the same roster do not rewrite it.
Failed or incomplete group searches, referrals, duplicate group IDs, invalid responses and LDAP operation exceptions abort the directory snapshot without applying it. A member DN whose entry no longer exists skips only its group, so other groups still synchronize and that team keeps its current members until the DN is removed or restored. Unreadable entries still abort the snapshot. Groups with missing, blank or multivalued IDs skip individually. Other member validation failures, such as ambiguous identities or nested groups, also skip their affected group. A membership-write failure for one valid group does not prevent other valid groups from synchronizing. A user disappearing during the membership write causes that group's transaction to make no changes. Repeated successful runs with the same roster do not rewrite it.
The LDAP identity must have permission to read every selected group, its complete membership attributes, and the user identity attributes. LDAP can omit attributes hidden by access controls just as it omits empty attributes, so the sync cannot distinguish every permission restriction from an empty group. Verify the bind account and dry-run output before enabling writes.
Runs have a timeout shorter than their shared Redis lease. A timeout cancels further writes and signals the LDAP reader to stop between entries and lookups. LDAP socket timeouts and server search time limits use `timeout_seconds` as an upper bound and are capped further for short sync intervals, and the next interval retries from current state. Cancellation does not undo groups already committed earlier in the run. Stopping the sync (`enabled: false` or removing the block and restarting) leaves current teams and memberships in place.
Runs have a timeout shorter than their shared Redis lease. A timeout cancels further writes and signals the LDAP reader to stop between entries and lookups. LDAP socket timeouts and server search time limits use `timeout_seconds` as an upper bound and are capped further for short sync intervals, and the next interval retries from current state. Cancellation does not undo groups already committed earlier in the run. Cache cleanup for a committed membership change finishes before cancellation is reported, so removed keys are not left cached by the job timeout. This cleanup can extend past the run deadline but performs no further membership writes. Stopping the sync (`enabled: false` or removing the block and restarting) leaves current teams and memberships in place.

View file

@ -1,9 +1,9 @@
import enum
from typing import Dict, List, Optional
from pydantic import Field
from litellm.proxy._types import WebhookEvent
from litellm.types.integrations.model_deprecation_email import MODEL_DEPRECATION_EMAIL_EVENT
from litellm.types.llms.base import LiteLLMBaseModel
@ -26,7 +26,7 @@ class SendKeyCreatedEmailEvent(WebhookEvent):
class SendKeyRotatedEmailEvent(WebhookEvent):
virtual_key: str
key_alias: Optional[str] = None
key_alias: str | None = None
"""
The virtual key that was rotated
this will be sk-123xxx, since we will be emailing this to the user to start using the new key
@ -34,23 +34,33 @@ class SendKeyRotatedEmailEvent(WebhookEvent):
class EmailEvent(str, enum.Enum):
model_deprecations = MODEL_DEPRECATION_EMAIL_EVENT
virtual_key_created = "Virtual Key Created"
new_user_invitation = "New User Invitation"
virtual_key_rotated = "Virtual Key Rotated"
soft_budget_crossed = "Soft Budget Crossed"
max_budget_alert = "Max Budget Alert"
class EmailEventSettings(LiteLLMBaseModel):
event: EmailEvent
enabled: bool
class EmailEventSettingsUpdateRequest(LiteLLMBaseModel):
settings: List[EmailEventSettings]
settings: list[EmailEventSettings]
class EmailEventSettingsResponse(LiteLLMBaseModel):
settings: List[EmailEventSettings]
settings: list[EmailEventSettings]
class DefaultEmailSettings(LiteLLMBaseModel):
"""Default settings for email events"""
settings: Dict[EmailEvent, bool] = Field(
settings: dict[EmailEvent, bool] = Field(
default_factory=lambda: {
EmailEvent.model_deprecations: False, # Explicit opt-in for owner retirement reminders.
EmailEvent.virtual_key_created: True, # On by default
EmailEvent.new_user_invitation: True, # On by default
EmailEvent.virtual_key_rotated: True, # On by default
@ -58,10 +68,12 @@ class DefaultEmailSettings(LiteLLMBaseModel):
EmailEvent.max_budget_alert: True, # On by default
}
)
def to_dict(self) -> Dict[str, bool]:
def to_dict(self) -> dict[str, bool]:
"""Convert to dictionary with string keys for storage"""
return {event.value: enabled for event, enabled in self.settings.items()}
@classmethod
def get_defaults(cls) -> Dict[str, bool]:
def get_defaults(cls) -> dict[str, bool]:
"""Get the default settings as a dictionary with string keys"""
return cls().to_dict()
return cls().to_dict()

View file

@ -0,0 +1,19 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogDelivery" (
"delivery_id" TEXT PRIMARY KEY DEFAULT gen_random_uuid()::text,
"request_id" TEXT NOT NULL,
"state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending',
"attempts" INTEGER NOT NULL DEFAULT 0,
"next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"applied_at" TIMESTAMP(3),
"last_error" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"accounting_payload" JSONB NOT NULL,
"is_batch" BOOLEAN NOT NULL
);
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_next_attempt_at_idx"
ON "LiteLLM_SpendLogDelivery" ("state", "next_attempt_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_state_applied_at_idx"
ON "LiteLLM_SpendLogDelivery" ("state", "applied_at");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogDelivery_request_id_idx"
ON "LiteLLM_SpendLogDelivery" ("request_id");

View file

@ -0,0 +1,5 @@
SET LOCAL lock_timeout = '10s';
DROP TABLE IF EXISTS "LiteLLM_AccountingDelivery";
DROP TABLE IF EXISTS "LiteLLM_AccountingEvent";
DROP TYPE IF EXISTS "AccountingDeliveryTarget";
DROP TYPE IF EXISTS "AccountingEventKind";

View file

@ -0,0 +1,16 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion" (
"id" TEXT NOT NULL,
"plugin_id" TEXT NOT NULL,
"version" TEXT,
"description" TEXT,
"manifest_json" TEXT,
"published_at" TIMESTAMP(3),
"published_by" TEXT,
"archive_available" BOOLEAN NOT NULL DEFAULT false,
"content" BYTEA,
CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_pkey" PRIMARY KEY ("id"),
CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_plugin_id_fkey" FOREIGN KEY ("plugin_id") REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE
);
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_version_key" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "version");
CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_published_at_idx" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "published_at");

View file

@ -0,0 +1,27 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_SkillSubmission" (
"id" TEXT NOT NULL,
"plugin_name" TEXT NOT NULL,
"operation" TEXT NOT NULL,
"status" TEXT NOT NULL DEFAULT 'pending_review',
"proposal_json" TEXT NOT NULL,
"archive_content" BYTEA,
"submitted_by" TEXT NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"base_plugin_id" TEXT,
"base_updated_at" TIMESTAMP(3),
"reviewed_by" TEXT,
"reviewed_at" TIMESTAMP(3),
"rejection_reason" TEXT,
CONSTRAINT "LiteLLM_SkillSubmission_pkey" PRIMARY KEY ("id")
);
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_created_at_id_idx"
ON "LiteLLM_SkillSubmission"("submitted_by", "created_at" DESC, "id" DESC);
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_status_created_at_id_idx"
ON "LiteLLM_SkillSubmission"("submitted_by", "status", "created_at" DESC, "id" DESC);
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_status_created_at_id_idx"
ON "LiteLLM_SkillSubmission"("status", "created_at" DESC, "id" DESC);
CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_created_at_id_idx"
ON "LiteLLM_SkillSubmission"("created_at" DESC, "id" DESC);

View file

@ -0,0 +1,3 @@
SET LOCAL lock_timeout = '10s';
ALTER TABLE "LiteLLM_SkillSubmission" ADD COLUMN IF NOT EXISTS "scan_json" TEXT;

View file

@ -0,0 +1,37 @@
SET LOCAL lock_timeout = '10s';
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_attribute
WHERE attrelid = to_regclass('"LiteLLM_ClaudeCodePluginTable"')
AND attname = 'anonymous_download_count' AND NOT attisdropped
) THEN
ALTER TABLE "LiteLLM_ClaudeCodePluginTable"
ADD COLUMN IF NOT EXISTS "anonymous_download_count" BIGINT NOT NULL DEFAULT 0;
END IF;
IF to_regclass('"LiteLLM_ClaudeCodePluginDownload"') IS NULL THEN
CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload" (
"id" BIGSERIAL NOT NULL,
"plugin_id" TEXT NOT NULL,
"downloaded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"version" TEXT,
"digest" TEXT NOT NULL,
"authenticated" BOOLEAN NOT NULL,
"user_id" TEXT,
"user_email" TEXT,
"key_alias" TEXT,
"team_id" TEXT,
"team_alias" TEXT,
CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_pkey" PRIMARY KEY ("id"),
CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_plugin_id_fkey" FOREIGN KEY ("plugin_id")
REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE
);
END IF;
IF to_regclass('"LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx"') IS NULL THEN
CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx"
ON "LiteLLM_ClaudeCodePluginDownload"("plugin_id", "id");
END IF;
END $$;

View file

@ -1616,10 +1616,69 @@ model LiteLLM_ClaudeCodePluginTable {
created_by String?
uploaded_archive LiteLLM_ClaudeCodePluginUpload?
versions LiteLLM_ClaudeCodePluginVersion[]
downloads LiteLLM_ClaudeCodePluginDownload[]
anonymous_download_count BigInt @default(0)
@@map("LiteLLM_ClaudeCodePluginTable")
}
model LiteLLM_ClaudeCodePluginVersion {
id String @id @default(uuid())
plugin_id String
version String?
description String?
manifest_json String?
published_at DateTime?
published_by String?
archive_available Boolean @default(false)
content Bytes?
plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade)
@@unique([plugin_id, version])
@@index([plugin_id, published_at])
}
model LiteLLM_SkillSubmission {
id String @id @default(uuid())
plugin_name String
operation String
status String @default("pending_review")
proposal_json String
archive_content Bytes?
submitted_by String
created_at DateTime @default(now())
base_plugin_id String?
base_updated_at DateTime?
reviewed_by String?
reviewed_at DateTime?
rejection_reason String?
scan_json String?
@@index([submitted_by, created_at(sort: Desc), id(sort: Desc)])
@@index([submitted_by, status, created_at(sort: Desc), id(sort: Desc)])
@@index([status, created_at(sort: Desc), id(sort: Desc)])
@@index([created_at(sort: Desc), id(sort: Desc)])
}
model LiteLLM_ClaudeCodePluginDownload {
id BigInt @id @default(autoincrement())
plugin_id String
downloaded_at DateTime @default(now())
version String?
digest String
authenticated Boolean
user_id String?
user_email String?
key_alias String?
team_id String?
team_alias String?
plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade)
@@index([plugin_id, id])
}
model LiteLLM_ClaudeCodePluginUpload {
plugin_name String @id
content Bytes
@ -2071,74 +2130,27 @@ model LiteLLM_TelemetryReport {
@@index([created_at])
}
enum AccountingEventKind {
inbound_usage
outbound_usage
cost
}
enum AccountingDeliveryState {
pending
applied
parked
}
enum AccountingDeliveryTarget {
LiteLLM_VerificationToken
LiteLLM_UserTable
LiteLLM_TeamTable
LiteLLM_TeamMembership
LiteLLM_OrganizationTable
LiteLLM_OrganizationMembership
LiteLLM_EndUserTable
LiteLLM_AgentsTable
LiteLLM_TagTable
LiteLLM_ProjectTable
LiteLLM_ModelAccessGroupBudgetTable
LiteLLM_BudgetTable
LiteLLM_BudgetWindowSpend
LiteLLM_DailyUserSpend
LiteLLM_DailyTeamSpend
LiteLLM_DailyOrganizationSpend
LiteLLM_DailyEndUserSpend
LiteLLM_DailyAgentSpend
LiteLLM_DailyTagSpend
LiteLLM_DailyModelUsage
LiteLLM_DailyToolSpend
LiteLLM_AutoRouterDailySpend
}
model LiteLLM_SpendLogDelivery {
delivery_id String @id @default(dbgenerated("(gen_random_uuid())::text"))
request_id String
state AccountingDeliveryState @default(pending)
accounting_payload Json
is_batch Boolean
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
created_at DateTime @default(now())
// Durable accounting ledger. One row per usage or cost event; payload holds the
// private event envelope versioned by schema_version. Usage rows retain request-scoped
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
model LiteLLM_AccountingEvent {
event_id String @id
kind AccountingEventKind
source_event_id String?
schema_version Int
payload_hash String
occurred_at DateTime
recorded_at DateTime @default(now())
payload Json
@@index([recorded_at])
@@index([kind, source_event_id])
}
// One row per (event, target table, version): pending work for the projector and,
// once applied in the same transaction as the target write, the receipt that stops a retry.
model LiteLLM_AccountingDelivery {
event_id String
target AccountingDeliveryTarget
version Int
state AccountingDeliveryState @default(pending)
attempts Int @default(0)
next_attempt_at DateTime @default(now())
applied_at DateTime?
last_error String?
@@id([event_id, target, version])
@@index([target, state, next_attempt_at])
@@index([state, next_attempt_at])
@@index([state, applied_at])
@@index([request_id])
}
model LiteLLM_SubscriptionAccountTable {

View file

@ -4288,6 +4288,7 @@ dependencies = [
name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"async-trait",
"aws-sdk-secretsmanager",
"bytes",
"criterion",
@ -4312,6 +4313,7 @@ dependencies = [
"litellm-llms",
"litellm-llms-types",
"litellm-model-catalog",
"litellm-router-core",
"litellm-router-types",
"litellm-secret-value",
"litellm-secrets",
@ -4393,6 +4395,24 @@ dependencies = [
[[package]]
name = "litellm-router-core"
version = "0.1.0"
dependencies = [
"arc-swap",
"async-trait",
"jiff",
"litellm-cache",
"litellm-router-types",
"parking_lot",
"rand 0.8.7",
"rstest",
"serde",
"serde_json",
"serde_with",
"sha2 0.10.9",
"strum",
"thiserror 2.0.19",
"tokio",
"tracing",
]
[[package]]
name = "litellm-router-types"

View file

@ -13,6 +13,7 @@ litellm-config = { path = "crates/config" }
litellm-error = { path = "crates/error" }
litellm-router = { path = "crates/router" }
litellm-router-types = { path = "crates/router-types" }
litellm-router-core = { path = "crates/router-core" }
litellm-tracing = { path = "crates/tracing" }
litellm-spend-clickhouse = { path = "crates/spend-clickhouse" }
litellm-storage-clickhouse = { path = "crates/storage-clickhouse" }
@ -76,6 +77,8 @@ litellm-cost = { path = "crates/cost" }
litellm-model-catalog = { path = "crates/model-catalog" }
litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" }
arc-swap = "1.9.2"
async-trait = "0.1.91"
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
tracing = "0.1"
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] }

View file

@ -25,6 +25,7 @@ litellm-tracing.workspace = true
litellm-spend-clickhouse.workspace = true
litellm-storage-clickhouse.workspace = true
litellm-host.workspace = true
async-trait.workspace = true
bytes.workspace = true
futures-util.workspace = true
litellm-cache.workspace = true
@ -47,6 +48,7 @@ litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "has
litellm-secrets-types.workspace = true
litellm-llms-types.workspace = true
litellm-router-types.workspace = true
litellm-router-core.workspace = true
litellm-host-python.workspace = true
litellm-model-catalog.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }

View file

@ -12,6 +12,7 @@ mod marshal;
mod model_registry;
mod preflight;
mod python_settings;
mod router;
mod routes;
mod secrets;
mod tokenizer;
@ -36,6 +37,8 @@ mod _native {
#[pymodule_export]
use crate::model_registry::bundled_model_cost_map;
#[pymodule_export]
use crate::router::{RouterCall, RouterCoreHandle};
#[pymodule_export]
use crate::routes::audio_transcription::{atranscription, transcription};
#[pymodule_export]
use crate::routes::chat_completions::{acompletion, completion};
@ -115,6 +118,8 @@ mod tests {
"bundled_model_cost_map",
"TokenCounter",
"Tokenizer",
"RouterCoreHandle",
"RouterCall",
"gil_stats",
"process_state_started",
"reserve_process_for_forking",

View file

@ -0,0 +1,408 @@
//! The documents `RouterCoreHandle` reads: the router's static and dynamic config, a call's
//! request and an `apply` change set. A config key router-core does not read declines the
//! handle, so Python keeps the router on its own core.
use std::collections::BTreeMap;
use litellm_router_core::select::PassThrough;
use litellm_router_core::{
ChangeSet, ConfigError, DeploymentSpec, OutcomeCallbacks, RouterConfig, RouterDynamicConfig,
RouterRequest,
config::{Fallbacks, RetryPolicy},
host::UncountedRetries,
mock::MockFailure,
registry::IndexOverrides,
};
use serde::{Deserialize, de::DeserializeOwned};
use serde_json::{Map, Value};
/// The `RouterConfig` keys router-core reads.
const STATIC_KEYS: &[&str] = &[
"model_list",
"routing_strategy",
"routing_groups",
"routing_group_members",
"plugins",
"cooldown_time",
"disable_cooldowns",
"allowed_fails_policy",
"retry_policy",
"model_group_retry_policy",
"retry_after",
"timeout",
"request_timeout",
"drop_params",
"enable_pre_call_checks",
"routing_plugins",
"callback_filters",
"enable_health_check_routing",
"health_check_staleness_threshold",
"background_health_check_model_groups",
"enable_weighted_failover",
"fallback_on_unsupported_params",
"provider_default_deployment_ids",
];
/// The `RouterDynamicConfig` keys router-core reads.
const DYNAMIC_KEYS: &[&str] = &[
"version",
"model_alias_map",
"model_group_alias",
"fallbacks",
"default_fallbacks",
"context_window_fallbacks",
"content_policy_fallbacks",
"num_retries",
"max_fallbacks",
"allowed_fails",
"allowed_fails_is_default",
"fallback_keys",
"default_deployment",
"default_deployment_raises",
];
/// Router arguments that do not change routing, accepted in either config and dropped.
const IGNORED_KEYS: &[&str] = &[
"set_verbose",
"debug_level",
"alerting_config",
"client_ttl",
];
/// Fallback lists Python treats as unset when empty (`context_window_fallbacks or None`).
const EMPTY_MEANS_UNSET: &[&str] = &["context_window_fallbacks", "content_policy_fallbacks"];
/// Which document a failure is about.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)]
pub(crate) enum Document {
#[strum(serialize = "router config")]
Static,
#[strum(serialize = "dynamic router config")]
Dynamic,
#[strum(serialize = "router request")]
Request,
#[strum(serialize = "change set")]
Changes,
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum Error {
#[error("the {document} must be a JSON object")]
NotObject { document: Document },
#[error("router-core does not read `{key}` in the {document}")]
UnknownKey { document: Document, key: String },
#[error("the {document} does not have the shape router-core reads")]
Shape {
document: Document,
#[source]
source: serde_json::Error,
},
#[error("the request names no model")]
MissingModel,
#[error("router-core rejected the router config")]
Core(#[from] ConfigError),
}
impl Error {
/// Whether router-core cannot serve this config, so Python keeps its own core.
pub(crate) fn declined(&self) -> bool {
match self {
Self::UnknownKey { .. } | Self::Core(_) => true,
Self::Shape { document, .. } => {
matches!(document, Document::Static | Document::Dynamic)
}
Self::NotObject { .. } | Self::MissingModel => false,
}
}
}
fn object(document: Document, value: Value) -> Result<Map<String, Value>, Error> {
match value {
Value::Object(object) => Ok(object),
_ => Err(Error::NotObject { document }),
}
}
fn shaped<T: DeserializeOwned>(document: Document, object: Map<String, Value>) -> Result<T, Error> {
serde_json::from_value(Value::Object(object))
.map_err(|source| Error::Shape { document, source })
}
/// Drop unset and ignored keys, then reject any key outside `known`.
fn routing_keys(
document: Document,
value: Value,
known: &[&str],
) -> Result<Map<String, Value>, Error> {
let mut object = object(document, value)?;
object.retain(|key, value| {
let empty_unset = EMPTY_MEANS_UNSET.contains(&key.as_str())
&& value.as_array().is_some_and(Vec::is_empty);
!(value.is_null() || empty_unset || IGNORED_KEYS.contains(&key.as_str()))
});
if let Some(key) = object.keys().find(|key| !known.contains(&key.as_str())) {
return Err(Error::UnknownKey {
document,
key: key.clone(),
});
}
Ok(object)
}
pub(crate) fn parse_router_config(value: Value) -> Result<RouterConfig, Error> {
let object = routing_keys(Document::Static, value, STATIC_KEYS)?;
shaped(Document::Static, object)
}
pub(crate) fn parse_dynamic_config(value: Value) -> Result<RouterDynamicConfig, Error> {
let object = routing_keys(Document::Dynamic, value, DYNAMIC_KEYS)?;
shaped(Document::Dynamic, object)
}
/// A call's request as Python spells it.
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct RequestSpec {
model: Option<String>,
team_id: Option<String>,
listing_team_id: Option<String>,
proxy_admin: Option<bool>,
allowed_models: Option<Vec<String>>,
admin_model: Option<String>,
specific_deployment: Option<bool>,
num_retries: Option<u32>,
uncounted_retries: Option<UncountedRetries>,
disable_fallbacks: Option<bool>,
router_weights: Option<Value>,
pinned_deployment_id: Option<String>,
routing_strategy: Option<String>,
filter_tags: Option<bool>,
/// `listed` or `healthy`: which pass-through pick the request makes.
pass_through: Option<String>,
target_order: Option<i64>,
excluded_ids: Option<Vec<String>>,
retry_skipped_ids: Option<Vec<String>>,
pre_routes: Option<bool>,
has_image: Option<bool>,
endpoint: Option<String>,
web_search: Option<bool>,
route_handler: Option<String>,
decisions_media: Option<bool>,
wildcard_request: Option<String>,
unrouted: Option<bool>,
passthrough_on_no_deployment: Option<bool>,
skip_weighted_failover: Option<bool>,
same_group_fallbacks_only: Option<bool>,
model_group_retry_policy: Option<BTreeMap<String, RetryPolicy>>,
retry_policy_raises: Option<Vec<String>>,
#[serde(default, deserialize_with = "present")]
fallbacks: Option<Option<Fallbacks>>,
#[serde(default, deserialize_with = "present")]
context_window_fallbacks: Option<Option<Fallbacks>>,
#[serde(default, deserialize_with = "present")]
content_policy_fallbacks: Option<Option<Fallbacks>>,
max_fallbacks: Option<u32>,
mock_failure: Option<MockFailure>,
cooldown_time: Option<f64>,
health_check_probe: Option<bool>,
outcome_callbacks: Option<OutcomeCallbacks>,
body: Option<Value>,
}
/// A key the request passed, even as `None`: `Some(None)`.
fn present<'de, D, T>(deserializer: D) -> Result<Option<Option<T>>, D::Error>
where
D: serde::Deserializer<'de>,
T: Deserialize<'de>,
{
Option::<T>::deserialize(deserializer).map(Some)
}
/// Read a request; `model` replaces the request's own when given.
pub(crate) fn parse_request(value: Value, model: Option<&str>) -> Result<RouterRequest, Error> {
let spec: RequestSpec = shaped(Document::Request, object(Document::Request, value)?)?;
let model = model
.map(str::to_owned)
.or(spec.model)
.ok_or(Error::MissingModel)?;
Ok(RouterRequest {
model,
body: spec.body.unwrap_or_default(),
team_id: spec.team_id,
listing_team_id: spec.listing_team_id,
proxy_admin: spec.proxy_admin.unwrap_or(false),
allowed_models: spec.allowed_models.unwrap_or_default(),
admin_model: spec.admin_model,
specific_deployment: spec.specific_deployment.unwrap_or(false),
num_retries: spec.num_retries,
uncounted_retries: spec.uncounted_retries,
disable_fallbacks: spec.disable_fallbacks.unwrap_or(false),
router_weights: spec.router_weights,
pinned_deployment_id: spec.pinned_deployment_id,
routing_strategy: spec.routing_strategy,
filter_tags: spec.filter_tags.unwrap_or(false),
pass_through: match spec.pass_through.as_deref() {
None => None,
Some("healthy") => Some(PassThrough::Healthy),
Some(_) => Some(PassThrough::Listed),
},
target_order: spec.target_order,
excluded_ids: spec.excluded_ids.unwrap_or_default(),
retry_skipped_ids: spec.retry_skipped_ids.unwrap_or_default(),
pre_routes: spec.pre_routes.unwrap_or(false),
has_image: spec.has_image.unwrap_or(false),
endpoint: spec.endpoint,
web_search: spec.web_search.unwrap_or(false),
route_handler: spec.route_handler,
decisions_media: spec.decisions_media.unwrap_or(false),
wildcard_request: spec.wildcard_request,
unrouted: spec.unrouted.unwrap_or(false),
passthrough_on_no_deployment: spec.passthrough_on_no_deployment.unwrap_or(false),
skip_weighted_failover: spec.skip_weighted_failover.unwrap_or(false),
same_group_fallbacks_only: spec.same_group_fallbacks_only.unwrap_or(false),
model_group_retry_policy: spec.model_group_retry_policy,
retry_policy_raises: spec.retry_policy_raises.unwrap_or_default(),
fallbacks: spec.fallbacks,
context_window_fallbacks: spec.context_window_fallbacks,
content_policy_fallbacks: spec.content_policy_fallbacks,
max_fallbacks: spec.max_fallbacks,
mock_failure: spec.mock_failure,
cooldown_time: spec.cooldown_time,
health_check_probe: spec.health_check_probe.unwrap_or(false),
outcome_callbacks: spec.outcome_callbacks.unwrap_or_default(),
payload: None,
})
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ChangesSpec {
/// Every deployment, replacing the current list before the upserts and deletes.
#[serde(default)]
replace: Option<Vec<DeploymentSpec>>,
/// Deployments added at the end, beside any that share their id.
#[serde(default)]
append: Vec<DeploymentSpec>,
#[serde(default)]
upsert: Vec<DeploymentSpec>,
#[serde(default)]
delete: Vec<String>,
/// A new dynamic document (the same keys `RouterCoreHandle(static, dynamic)` takes).
#[serde(default)]
settings: Option<Value>,
/// The registry indexes the host wrote (`IndexOverrides`); `{}` drops them.
#[serde(default)]
index: Option<IndexOverrides>,
/// The deployment at the id map's position for this id, removed with every index shifted.
#[serde(default)]
remove: Option<String>,
/// A write to the indexes the host hands out (`HostIndexes`), field by field.
#[serde(default)]
host_index: Option<litellm_router_core::host_index::IndexWrite>,
}
pub(crate) fn parse_changes(value: Value) -> Result<ChangeSet, Error> {
let spec: ChangesSpec = shaped(Document::Changes, object(Document::Changes, value)?)?;
Ok(ChangeSet {
replace: spec.replace,
appends: spec.append,
upserts: spec.upsert,
deletes: spec.delete,
settings: spec.settings.map(parse_dynamic_config).transpose()?,
index: spec.index,
remove: spec.remove,
host_index: spec.host_index,
})
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::unknown_key(json!({"routing_strategy_args": {"ttl": 1}}), "routing_strategy_args")]
#[case::dynamic_key_in_static(json!({"num_retries": 2}), "num_retries")]
fn an_unread_static_key_declines_naming_the_key(#[case] json: Value, #[case] key: &str) {
let error = parse_router_config(json).unwrap_err();
assert!(error.declined());
assert!(matches!(&error, Error::UnknownKey { key: named, .. } if named == key));
}
#[rstest]
#[case::null_unknown(json!({"routing_strategy_args": null}))]
#[case::ignored(json!({"set_verbose": true, "debug_level": "INFO"}))]
#[case::pre_call_checks_off(json!({"enable_pre_call_checks": false}))]
#[case::pre_call_checks_on(json!({"enable_pre_call_checks": true}))]
fn unset_and_ignored_keys_are_accepted(#[case] json: Value) {
parse_router_config(json).unwrap();
}
#[rstest]
#[case::not_object(json!([]), false)]
#[case::bad_shape(json!({"retry_after": "soon"}), true)]
fn a_config_that_is_not_an_object_is_a_value_error_and_a_wrong_shape_declines(
#[case] json: Value,
#[case] declined: bool,
) {
assert_eq!(parse_router_config(json).unwrap_err().declined(), declined);
}
#[rstest]
fn python_router_defaults_read_as_unset() {
let dynamic = parse_dynamic_config(json!({
"num_retries": null, "fallbacks": [], "context_window_fallbacks": [],
"content_policy_fallbacks": [], "model_group_alias": {"gpt": "gpt-4"},
}))
.unwrap();
assert_eq!(dynamic.num_retries, None);
assert!(dynamic.context_window_fallbacks.is_none());
assert!(dynamic.content_policy_fallbacks.is_none());
assert_eq!(dynamic.model_alias_map["gpt"], "gpt-4");
}
#[rstest]
#[case::own_model(json!({"model": "gpt", "num_retries": 1}), None, "gpt")]
#[case::override_model(json!({"model": "gpt"}), Some("other"), "other")]
fn a_request_reads_its_model(
#[case] json: Value,
#[case] model: Option<&str>,
#[case] expected: &str,
) {
assert_eq!(parse_request(json, model).unwrap().model, expected);
}
#[rstest]
#[case::negative("negative", UncountedRetries::Negative)]
#[case::not_positive("not_positive", UncountedRetries::NotPositive)]
#[case::raises("raises", UncountedRetries::Raises)]
fn a_request_reads_a_num_retries_that_is_no_count(
#[case] spelled: &str,
#[case] expected: UncountedRetries,
) {
let request =
parse_request(json!({"model": "gpt", "uncounted_retries": spelled}), None).unwrap();
assert_eq!(request.uncounted_retries, Some(expected));
assert_eq!(request.num_retries, None);
}
#[rstest]
#[case::unknown_field(json!({"model": "gpt", "temperature": 1}))]
#[case::no_model(json!({"team_id": "t"}))]
fn a_bad_request_is_a_value_error(#[case] json: Value) {
assert!(!parse_request(json, None).unwrap_err().declined());
}
#[rstest]
fn a_change_set_reads_upserts_and_deletes() {
let changes = parse_changes(json!({
"upsert": [{"model_name": "m", "litellm_params": {"model": "openai/x"}}],
"delete": ["old"],
}))
.unwrap();
assert_eq!(changes.upserts.len(), 1);
assert_eq!(changes.deletes, ["old"]);
}
}

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,55 @@
//! The retry decisions `RouterCoreHandle` exposes as members: `get_num_retries_from_retry_policy`
//! and `should_retry_this_error`, decided by router-core for a failure Python describes the way it
//! describes a failed attempt.
use litellm_router_core::{
AttemptError, Router,
retry::{RetryContext, resolve_retry_policy, should_retry as core_should_retry},
};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
fn failure(outcome: &Bound<'_, PyDict>) -> PyResult<AttemptError> {
match super::attempt_outcome(outcome)? {
Err(failure) => Ok(failure),
Ok(_) => Err(PyValueError::new_err("a failure needs `ok: False`")),
}
}
pub(crate) fn policy_retries(
router: &Router,
outcome: &Bound<'_, PyDict>,
model_group: Option<&str>,
) -> PyResult<Option<u32>> {
let failure = failure(outcome)?;
let settings = router.settings();
let policy = match model_group {
Some(group) => resolve_retry_policy(
settings.retry_policy.as_ref(),
group,
Some(&settings.model_group_retry_policy),
),
None => settings.retry_policy.as_ref(),
};
Ok(policy.and_then(|policy| policy.retries_for(&failure)))
}
pub(crate) fn allowed_fails(router: &Router, outcome: &Bound<'_, PyDict>) -> PyResult<Option<u32>> {
Ok(router.settings().allowed_fails_for(&failure(outcome)?))
}
pub(crate) fn should_retry(
outcome: &Bound<'_, PyDict>,
(healthy, all): (usize, usize),
(context_window, content_policy, regular): (bool, bool, bool),
) -> PyResult<bool> {
Ok(core_should_retry(
&failure(outcome)?,
RetryContext {
healthy,
all,
has_context_window_fallbacks: context_window,
has_content_policy_fallbacks: content_policy,
has_regular_fallbacks: regular,
},
))
}

View file

@ -0,0 +1,94 @@
//! The stream-recovery entry points of `RouterCoreHandle`: what may still take over a stream that
//! failed before content, the same-group stream retry, and the fallback walk re-entered from a
//! stream. The host's stream wrappers call them; the calls are driven like `start`.
use litellm_router_core::{
FailureKind, Op, Router,
call::{StreamOrigin, StreamPolicy},
};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use super::{RouterCall, attempt_outcome, config, config_error_to_pyerr, document, host::Finish};
/// A stream recovery a host wrapper asks for.
#[derive(Clone, Copy)]
pub(super) enum Recovery {
/// `_astream_retry_same_group`.
Retry,
/// The fallback walk re-entered from the stream wrapper.
Fallback,
}
fn origin(value: &Bound<'_, PyAny>) -> PyResult<StreamOrigin> {
serde_json::from_value(document(value)?)
.map_err(|error| PyValueError::new_err(format!("a stream origin: {error}")))
}
fn kind(kind: Option<&str>) -> PyResult<Option<FailureKind>> {
kind.map(|kind| {
kind.parse()
.map_err(|_| PyValueError::new_err(format!("`{kind}` is not a router failure kind")))
})
.transpose()
}
/// `{"can_retry", "can_fall_back", "refusal_fallback", "retry_budget": int | None,
/// "policy_applies": bool}`.
pub(super) fn policy<'py>(
py: Python<'py>,
router: &Router,
request: &Bound<'py, PyAny>,
stream_origin: &Bound<'py, PyAny>,
failure_kind: Option<&str>,
) -> PyResult<Bound<'py, PyDict>> {
let request = config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?;
let StreamPolicy {
can_retry,
can_fall_back,
refusal_fallback,
retry_budget,
} = router.stream_policy(&request, &origin(stream_origin)?, kind(failure_kind)?);
let dict = PyDict::new(py);
dict.set_item("can_retry", can_retry)?;
dict.set_item("can_fall_back", can_fall_back)?;
dict.set_item("refusal_fallback", refusal_fallback)?;
dict.set_item("retry_budget", retry_budget.map(|(budget, _)| budget))?;
dict.set_item(
"policy_applies",
retry_budget.is_some_and(|(_, applies)| applies),
)?;
Ok(dict)
}
/// A recovery call for a stream from `stream_origin` that failed with `failure`.
pub(super) fn recover(
router: &Router,
recovery: Recovery,
op: &str,
request: &Bound<'_, PyAny>,
stream_origin: &Bound<'_, PyAny>,
failure: &Bound<'_, PyDict>,
) -> PyResult<RouterCall> {
let op = Op::new(op);
let request = config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?;
let stream_origin = origin(stream_origin)?;
let Err(failure) = attempt_outcome(failure)? else {
return Err(PyValueError::new_err("a stream failure needs `ok: False`"));
};
let router = router.clone();
Ok(RouterCall::new(async move {
let response = match recovery {
Recovery::Retry => {
router
.retry_stream(op, request, stream_origin, failure)
.await
}
Recovery::Fallback => {
router
.fallback_stream(op, request, stream_origin, failure)
.await
}
};
response.map(Finish::Call)
}))
}

View file

@ -7,3 +7,21 @@ repository.workspace = true
description = "The router core behind litellm.Router: registry, settings, health, selection and the call loop"
[dependencies]
arc-swap.workspace = true
async-trait.workspace = true
litellm-cache.workspace = true
litellm-router-types.workspace = true
jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-zoneinfo", "tzdb-bundle-always"] }
parking_lot.workspace = true
rand.workspace = true
serde.workspace = true
serde_with.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
sha2.workspace = true
strum.workspace = true
thiserror.workspace = true
tracing.workspace = true
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

View file

@ -0,0 +1,236 @@
//! Who may reach a deployment: model access windows (`litellm_core_utils/access_windows.py`,
//! a deployment reserved for some teams during a daily wall-clock window) and model access
//! groups (`_filter_deployments_by_model_access_groups`, a caller allowed only through access
//! groups is kept to its groups' deployments).
use std::{collections::BTreeSet, sync::Arc};
use jiff::{Timestamp, tz::TimeZone};
use serde_json::Value;
use crate::registry::Deployment;
/// One `model_info.access_windows` entry (`ModelAccessWindow`), validated when the deployment
/// was added.
#[derive(Clone, Debug, PartialEq)]
pub struct AccessWindow {
/// Microseconds after midnight.
pub start: u64,
pub end: u64,
pub timezone: String,
pub team_ids: Vec<String>,
}
const MICROS_PER_SECOND: u64 = 1_000_000;
/// A `datetime.time` as pydantic reads it: `HH:MM[:SS[.ffffff]]`, or seconds after midnight.
fn wall_clock(value: &Value) -> Option<u64> {
match value {
Value::Number(seconds) => {
let seconds = seconds.as_f64()?;
(0.0..86_400.0)
.contains(&seconds)
.then(|| (seconds * MICROS_PER_SECOND as f64).round() as u64)
}
Value::String(text) => {
let mut parts = text.splitn(3, ':');
let hour: u64 = parts.next()?.parse().ok()?;
let minute: u64 = parts.next()?.parse().ok()?;
let (second, micros) = match parts.next() {
None => (0, 0),
Some(rest) => {
let (whole, fraction) = rest.split_once('.').unwrap_or((rest, ""));
let digits: String = fraction.chars().take(6).collect();
let micros = if digits.is_empty() {
0
} else {
format!("{digits:0<6}").parse().ok()?
};
(whole.parse().ok()?, micros)
}
};
(hour < 24 && minute < 60 && second < 60)
.then(|| ((hour * 60 + minute) * 60 + second) * MICROS_PER_SECOND + micros)
}
_ => None,
}
}
impl AccessWindow {
fn parse(value: &Value) -> Option<Self> {
let window = value.as_object()?;
Some(Self {
start: wall_clock(window.get("start")?)?,
end: wall_clock(window.get("end")?)?,
timezone: window.get("timezone")?.as_str()?.to_owned(),
team_ids: window
.get("team_ids")?
.as_array()?
.iter()
.filter_map(Value::as_str)
.map(str::to_owned)
.collect(),
})
}
/// `is_window_active`: the window's wall clock contains `now`, wrapping past midnight when
/// it ends before it starts.
pub fn is_active(&self, now: Timestamp) -> bool {
let Ok(zone) = TimeZone::get(&self.timezone) else {
return false;
};
let local = now.to_zoned(zone).time();
let micros = ((u64::from(local.hour().unsigned_abs()) * 60
+ u64::from(local.minute().unsigned_abs()))
* 60
+ u64::from(local.second().unsigned_abs()))
* MICROS_PER_SECOND
+ u64::from(local.subsec_nanosecond().unsigned_abs()) / 1_000;
if self.start < self.end {
self.start <= micros && micros < self.end
} else {
micros >= self.start || micros < self.end
}
}
/// The window's end as `{end:%H:%M}`.
pub fn end_hhmm(&self) -> String {
let minutes = self.end / MICROS_PER_SECOND / 60;
format!("{:02}:{:02}", minutes / 60, minutes % 60)
}
}
/// A deployment's access windows.
pub fn access_windows(deployment: &Deployment) -> Vec<AccessWindow> {
deployment
.model_info
.raw
.get("access_windows")
.and_then(Value::as_array)
.map(|windows| windows.iter().filter_map(AccessWindow::parse).collect())
.unwrap_or_default()
}
/// `_reservation_blocking_window`: the first active window, unless the caller's team is named
/// by one of the active windows.
fn blocking_window(
deployment: &Deployment,
team_id: Option<&str>,
now: Timestamp,
) -> Option<AccessWindow> {
let active: Vec<AccessWindow> = access_windows(deployment)
.into_iter()
.filter(|window| window.is_active(now))
.collect();
if let Some(team) = team_id
&& active
.iter()
.any(|window| window.team_ids.iter().any(|id| id == team))
{
return None;
}
active.into_iter().next()
}
/// `filter_reserved_deployments`: the deployments no active window withholds from the caller,
/// and the first window that withheld one.
pub fn filter_reserved(
deployments: Vec<Arc<Deployment>>,
team_id: Option<&str>,
now: Timestamp,
) -> (Vec<Arc<Deployment>>, Option<AccessWindow>) {
let mut blocking = None;
let kept = deployments
.into_iter()
.filter(
|deployment| match blocking_window(deployment, team_id, now) {
None => true,
Some(window) => {
blocking.get_or_insert(window);
false
}
},
)
.collect();
(kept, blocking)
}
/// `_filter_deployments_by_model_access_groups`: a caller whose allowed models (`models` plus
/// `team_models`) reach `model` only through access groups is kept to the deployments in those
/// groups. `model_groups` are the access groups the rows listed for `model` carry.
pub fn filter_by_access_groups(
model: &str,
deployments: Vec<Arc<Deployment>>,
allowed_models: &[String],
model_groups: impl FnOnce() -> BTreeSet<String>,
) -> Vec<Arc<Deployment>> {
if deployments.is_empty() || allowed_models.is_empty() {
return deployments;
}
let allowed = |name: &str| allowed_models.iter().any(|m| m == name);
if allowed(model) || allowed("*") || allowed("all-proxy-models") {
return deployments;
}
let groups: BTreeSet<String> = model_groups()
.into_iter()
.filter(|group| allowed(group))
.collect();
if groups.is_empty() {
return deployments;
}
deployments
.into_iter()
.filter(|d| {
d.model_info
.access_groups
.iter()
.any(|g| groups.contains(g))
})
.collect()
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::minutes(json!("09:30"), Some(34_200_000_000))]
#[case::seconds(json!("09:30:15"), Some(34_215_000_000))]
#[case::fraction(json!("00:00:00.5"), Some(500_000))]
#[case::number(json!(60), Some(60_000_000))]
#[case::bad(json!("25:00"), None)]
fn wall_clock_reads_pydantic_times(#[case] raw: Value, #[case] expected: Option<u64>) {
assert_eq!(wall_clock(&raw), expected);
}
#[rstest]
#[case::inside("2026-01-01T10:00:00Z", true)]
#[case::at_end("2026-01-01T17:00:00Z", false)]
#[case::before("2026-01-01T08:59:59Z", false)]
fn a_window_is_active_inside_its_wall_clock(#[case] now: &str, #[case] expected: bool) {
let window = AccessWindow {
start: wall_clock(&json!("09:00")).unwrap(),
end: wall_clock(&json!("17:00")).unwrap(),
timezone: "UTC".to_owned(),
team_ids: vec!["t".to_owned()],
};
assert_eq!(window.is_active(now.parse().unwrap()), expected);
assert_eq!(window.end_hhmm(), "17:00");
}
#[test]
fn an_overnight_window_wraps_past_midnight() {
let window = AccessWindow {
start: wall_clock(&json!("22:00")).unwrap(),
end: wall_clock(&json!("02:00")).unwrap(),
timezone: "UTC".to_owned(),
team_ids: vec!["t".to_owned()],
};
assert!(window.is_active("2026-01-01T23:00:00Z".parse().unwrap()));
assert!(window.is_active("2026-01-01T01:00:00Z".parse().unwrap()));
assert!(!window.is_active("2026-01-01T12:00:00Z".parse().unwrap()));
}
}

View file

@ -0,0 +1,386 @@
//! The deployment build: how configured entries become the registry's deployments.
//!
//! The core owns the loop (`set_model_list`, `_create_deployment` and `_add_deployment`): which
//! values are secret references, the id of a deployment that declares none, one deployment per
//! `organization`, and the order everything happens in. What needs the embedding process is a
//! host step, asked for one deployment at a time so an error surfaces exactly where the Python
//! router raises it:
//! - [`Host::secrets`]: the `os.environ/` references the entry holds;
//! - [`Host::validate`]: the host's checks, and the entry as it normalizes it;
//! - [`Host::provider`]: the provider that serves it;
//! - [`Host::added`]: the host's side effects, and the entry as the registry stores it.
use serde_json::{Map, Value};
use crate::{
error::{ConfigError, Error},
host::{BuildAsk, BuildPath, Built, Host, Provided, SecretAt, SecretRef},
registry::{DeploymentSpec, deployment_id},
};
/// The `litellm_params` keys whose `os.environ/` value is a secret pointer the identity source
/// dereferences at use time, so the build leaves it as written. The router's copy of
/// `_WIF_POINTER_FIELDS` (litellm/types/router.py), which the proxy's config loader also reads.
pub const SECRET_POINTER_PARAMS: &[&str] = &[
"anthropic_issuer_signing_key_ref",
"anthropic_keycloak_client_secret_ref",
];
const SECRET_PREFIX: &str = "os.environ/";
/// How a build's deployments join the registry.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)]
#[strum(serialize_all = "snake_case")]
pub enum Install {
/// They become every deployment (`set_model_list`).
Replace,
/// They are added at the end (a deployment retried once its provider resolves).
Append,
/// Each replaces the deployment with its id, if any, and is added at the end
/// (`add_deployment`, an upsert).
Upsert,
}
/// One deployment the build accepted, as the registry stores it.
#[derive(Clone, Debug, PartialEq)]
pub struct BuiltDeployment {
/// The entry index in the list the build was given.
pub index: usize,
pub spec: DeploymentSpec,
}
/// Build `entries` (model_list entries) through `host`, in order. `resolve_secrets` resolves
/// top-level `os.environ/` references first, as `set_model_list` does; `add_deployment` and
/// upserts take their params as given. `install` takes each deployment as soon as it is
/// accepted, so it joins the registry before the next one is built, as the Python router adds
/// each in turn.
pub async fn build(
host: &dyn Host,
path: BuildPath,
entries: Vec<Map<String, Value>>,
resolve_secrets: bool,
mut install: impl FnMut(&BuiltDeployment) -> Result<(), Error> + Send,
) -> Result<Vec<BuiltDeployment>, Error> {
let declared: Vec<Option<String>> = entries.iter().map(declared_id).collect();
let mut built = Vec::new();
for (index, mut entry) in entries.into_iter().enumerate() {
let model_name = text(entry.get("model_name"));
let mut params = match entry.remove("litellm_params") {
Some(Value::Object(params)) => params,
_ => Map::new(),
};
if resolve_secrets {
let refs = param_refs(&params);
if !refs.is_empty() {
let values = host.secrets(index, &refs).await?;
for (secret, value) in refs.into_iter().zip(values) {
if let SecretAt::Param(key) = secret.at {
params.insert(key, value);
}
}
}
}
let mut info = match entry.remove("model_info") {
Some(Value::Object(info)) => info,
_ => Map::new(),
};
let declared_id = declared[index].clone();
if !info.contains_key("id") {
info.insert(
"id".to_owned(),
Value::String(deployment_id(&model_name, &params)),
);
}
// `set_model_list` adds one deployment per organization of a list, each with the
// organization set on the same params (so every one keeps the list's id).
let organizations = match (path, params.get("organization")) {
(BuildPath::Create, Some(Value::Array(organizations))) => organizations.clone(),
_ => vec![Value::Null],
};
let duplicate_id = path == BuildPath::Create
&& declared_id.as_ref().is_some_and(|id| {
declared
.iter()
.filter(|other| other.as_ref() == Some(id))
.count()
> 1
});
for organization in organizations {
if !organization.is_null() {
params.insert("organization".to_owned(), organization);
}
let mut deployment = entry.clone();
deployment.insert("model_name".to_owned(), Value::String(model_name.clone()));
deployment.insert("litellm_params".to_owned(), Value::Object(params.clone()));
deployment.insert("model_info".to_owned(), Value::Object(info.clone()));
if let Some(spec) = build_one(
host,
path,
index,
&deployment,
declared_id.as_deref(),
duplicate_id,
)
.await?
{
let deployment = BuiltDeployment { index, spec };
install(&deployment)?;
built.push(deployment);
}
}
}
Ok(built)
}
/// One deployment's steps: validate, provider, the data-source references, added.
async fn build_one(
host: &dyn Host,
path: BuildPath,
index: usize,
entry: &Map<String, Value>,
declared_id: Option<&str>,
duplicate_id: bool,
) -> Result<Option<DeploymentSpec>, Error> {
let ask = |entry| BuildAsk {
path,
index,
entry,
declared_id,
duplicate_id,
};
let validated = match host.validate(&ask(entry)).await? {
Built::Entry(validated) => validated,
Built::Skip | Built::Invalid => return Ok(None),
};
let provider = match host.provider(&ask(&validated)).await? {
Provided::Provider(provider) => provider,
Provided::Invalid => return Ok(None),
};
let refs = data_source_refs(&validated);
if !refs.is_empty() {
let _ = host.secrets(index, &refs).await?;
}
let stored = match host.added(&ask(&validated), provider.as_deref()).await? {
Built::Entry(stored) => stored,
Built::Skip | Built::Invalid => return Ok(None),
};
let model_name = text(stored.get("model_name"));
let mut spec: DeploymentSpec = serde_json::from_value(Value::Object(stored))
.map_err(|source| ConfigError::InvalidDeployment { model_name, source })?;
spec.provider = provider;
Ok(Some(spec))
}
/// The top-level `litellm_params` values `set_model_list` resolves.
fn param_refs(params: &Map<String, Value>) -> Vec<SecretRef> {
params
.iter()
.filter(|(key, _)| !SECRET_POINTER_PARAMS.contains(&key.as_str()))
.filter_map(|(key, value)| {
let name = value
.as_str()
.filter(|name| name.starts_with(SECRET_PREFIX))?;
Some(SecretRef {
name: name.to_owned(),
at: SecretAt::Param(key.clone()),
})
})
.collect()
}
/// The `dataSources[].parameters.endpoint` / `key` references `_add_deployment` reads from the
/// environment.
fn data_source_refs(entry: &Map<String, Value>) -> Vec<SecretRef> {
let sources = entry
.get("litellm_params")
.and_then(|params| params.get("dataSources"))
.and_then(Value::as_array);
let mut refs = Vec::new();
for (source, data_source) in sources.into_iter().flatten().enumerate() {
let Some(parameters) = data_source.get("parameters").and_then(Value::as_object) else {
continue;
};
for key in ["endpoint", "key"] {
if let Some(name) = parameters
.get(key)
.and_then(Value::as_str)
.filter(|name| name.starts_with(SECRET_PREFIX))
{
refs.push(SecretRef {
name: name.to_owned(),
at: SecretAt::DataSource {
source,
key: key.to_owned(),
},
});
}
}
}
refs
}
/// The `model_info.id` an entry declares, as Python's `str()` spells it.
fn declared_id(entry: &Map<String, Value>) -> Option<String> {
match entry.get("model_info")?.get("id")? {
Value::Null => None,
Value::String(id) => Some(id.clone()),
other => Some(other.to_string()),
}
}
/// A name as Python's `str()` spells it.
fn text(value: Option<&Value>) -> String {
match value {
Some(Value::String(text)) => text.clone(),
Some(Value::Null) | None => "None".to_owned(),
Some(other) => other.to_string(),
}
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use serde_json::json;
use super::*;
use crate::{
cache::{InMemoryRouterCache, RouterCache},
error::AttemptError,
host::{Attempt, Output, RouterEvent, RouterRequest, RoutingContext, SelectStage},
registry::Deployment,
};
/// Records every build step and resolves `os.environ/X` to `secret-X`.
#[derive(Default)]
struct Recording {
steps: Mutex<Vec<String>>,
}
#[async_trait]
impl Host for Recording {
async fn invoke(&self, _attempt: &Attempt) -> Result<Output, AttemptError> {
unreachable!("the build runs no attempt")
}
async fn select(
&self,
_stage: &SelectStage,
_ctx: &RoutingContext<'_>,
candidates: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, Error> {
Ok(candidates)
}
async fn allow_fallback(
&self,
_request: &RouterRequest,
_target: &str,
) -> Result<bool, crate::error::AttemptError> {
Ok(true)
}
async fn sleep(&self, _duration: std::time::Duration) {}
fn observe(&self, _event: &RouterEvent) {}
fn cache(&self) -> Arc<dyn RouterCache> {
Arc::new(InMemoryRouterCache::default())
}
async fn secrets(&self, _index: usize, refs: &[SecretRef]) -> Result<Vec<Value>, Error> {
self.steps.lock().unwrap().push(format!("secrets {refs:?}"));
Ok(refs
.iter()
.map(|r| json!(format!("secret-{}", &r.name[SECRET_PREFIX.len()..])))
.collect())
}
async fn validate(&self, ask: &BuildAsk<'_>) -> Result<Built, Error> {
self.steps.lock().unwrap().push(format!(
"validate {} {}",
ask.index, ask.entry["litellm_params"]
));
Ok(Built::Entry(ask.entry.clone()))
}
async fn provider(&self, ask: &BuildAsk<'_>) -> Result<Provided, Error> {
self.steps
.lock()
.unwrap()
.push(format!("provider {}", ask.index));
Ok(Provided::Provider(Some("azure".to_owned())))
}
async fn added(&self, ask: &BuildAsk<'_>, provider: Option<&str>) -> Result<Built, Error> {
self.steps
.lock()
.unwrap()
.push(format!("added {} {provider:?}", ask.index));
Ok(Built::Entry(ask.entry.clone()))
}
}
fn entries(value: Value) -> Vec<Map<String, Value>> {
serde_json::from_value(value).unwrap()
}
#[tokio::test]
async fn create_resolves_secrets_then_builds_one_deployment_per_organization() {
let host = Recording::default();
let built = build(
&host,
BuildPath::Create,
entries(json!([{
"model_name": "g",
"litellm_params": {
"model": "azure/x",
"api_key": "os.environ/KEY",
"anthropic_issuer_signing_key_ref": "os.environ/POINTER",
"organization": ["a", "b"],
},
}])),
true,
|_| Ok(()),
)
.await
.unwrap();
let steps = host.steps.lock().unwrap().clone();
assert_eq!(steps.len(), 7, "{steps:?}");
assert!(steps[0].starts_with("secrets") && steps[0].contains("KEY"));
assert!(!steps[0].contains("POINTER"));
assert!(steps[1].contains(r#""organization":"a""#));
assert!(steps[1].contains("secret-KEY"));
assert_eq!(steps[3], "added 0 Some(\"azure\")");
assert!(steps[4].contains(r#""organization":"b""#));
assert_eq!(built.len(), 2);
assert_eq!(built[0].spec.provider.as_deref(), Some("azure"));
assert_eq!(
built[0].spec.model_info["id"],
built[1].spec.model_info["id"]
);
let deployment = Deployment::from_spec(built[0].spec.clone()).unwrap();
assert_eq!(deployment.provider(), "azure");
}
#[tokio::test]
async fn add_leaves_params_as_given_and_asks_for_data_source_references() {
let host = Recording::default();
let built = build(
&host,
BuildPath::Add,
entries(json!([{
"model_name": "g",
"litellm_params": {
"model": "azure/x",
"api_key": "os.environ/KEY",
"organization": ["a", "b"],
"dataSources": [{"parameters": {"endpoint": "os.environ/END", "key": "plain"}}],
},
"model_info": {"id": "given"},
}])),
false,
|_| Ok(()),
)
.await
.unwrap();
let steps = host.steps.lock().unwrap().clone();
assert_eq!(steps.len(), 4, "{steps:?}");
assert!(steps[0].contains("os.environ/KEY"));
assert!(steps[2].contains("DataSource { source: 0, key: \"endpoint\" }"));
assert_eq!(built.len(), 1);
assert_eq!(built[0].spec.model_info["id"], "given");
}
}

View file

@ -0,0 +1,159 @@
//! [`RouterCache`]: the DualCache-style store all routing state lives in, under the keys,
//! encodings and TTLs the Python router uses, so Python and Rust workers share it. Method names
//! follow `litellm-cache`'s `BaseCache`, `BatchCache`, `CounterCache` and `DeleteCache`.
use std::{
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
use async_trait::async_trait;
use litellm_cache::Error;
use parking_lot::Mutex;
use serde_json::Value;
#[async_trait]
pub trait RouterCache: Send + Sync {
async fn async_get_cache(&self, key: &str) -> Result<Option<Value>, Error>;
async fn async_set_cache(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<(), Error>;
/// Add `amount` to the number at `key` (0 when missing) and return the new value. The TTL
/// applies when the key is created, as Redis `INCRBYFLOAT` + `EXPIRE NX` does.
async fn async_increment(
&self,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, Error>;
async fn async_batch_get_cache(&self, keys: &[String]) -> Result<Vec<Option<Value>>, Error>;
async fn async_delete_cache(&self, key: &str) -> Result<(), Error>;
/// Read this worker's in-memory tier only (`local_only=True`).
async fn async_get_local(&self, key: &str) -> Result<Option<Value>, Error> {
self.async_get_cache(key).await
}
/// Add to a counter in this worker's in-memory tier only (`local_only=True`).
async fn async_increment_local(
&self,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
self.async_increment(key, amount, ttl).await
}
/// `_increment_deployment_usage`'s write: add each `(key, amount)` in one pipeline. With a
/// shared tier, this worker's in-memory tier then holds the shared totals, so its reads see
/// what other workers counted up to now.
async fn async_increment_usage(
&self,
increments: &[(String, f64)],
ttl: Option<Duration>,
) -> Result<(), Error> {
for (key, amount) in increments {
self.async_increment(key, *amount, ttl).await?;
}
Ok(())
}
}
#[derive(Default)]
struct Entry {
value: Value,
expires: Option<Instant>,
}
/// An in-process [`RouterCache`], for tests and for state Python keeps `local_only`.
#[derive(Clone, Default)]
pub struct InMemoryRouterCache {
entries: Arc<Mutex<HashMap<String, Entry>>>,
}
impl InMemoryRouterCache {
fn live(entries: &mut HashMap<String, Entry>, key: &str) -> Option<Value> {
let expired = entries
.get(key)
.is_some_and(|entry| entry.expires.is_some_and(|at| at <= Instant::now()));
if expired {
entries.remove(key);
}
entries.get(key).map(|entry| entry.value.clone())
}
pub fn get(&self, key: &str) -> Option<Value> {
Self::live(&mut self.entries.lock(), key)
}
pub fn increment(&self, key: &str, amount: f64, ttl: Option<Duration>) -> f64 {
let mut entries = self.entries.lock();
let current = Self::live(&mut entries, key).and_then(|value| value.as_f64());
let next = current.unwrap_or(0.0) + amount;
let value = serde_json::Number::from_f64(next)
.map(|number| {
if next.fract() == 0.0 && next.abs() < i64::MAX as f64 {
Value::from(next as i64)
} else {
Value::Number(number)
}
})
.unwrap_or(Value::Null);
let expires = match (current, entries.get(key)) {
(Some(_), Some(entry)) => entry.expires,
_ => ttl.map(|ttl| Instant::now() + ttl),
};
entries.insert(key.to_owned(), Entry { value, expires });
next
}
}
#[async_trait]
impl RouterCache for InMemoryRouterCache {
async fn async_get_cache(&self, key: &str) -> Result<Option<Value>, Error> {
Ok(self.get(key))
}
async fn async_set_cache(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<(), Error> {
let expires = ttl.map(|ttl| Instant::now() + ttl);
self.entries
.lock()
.insert(key.to_owned(), Entry { value, expires });
Ok(())
}
async fn async_increment(
&self,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
Ok(self.increment(key, amount, ttl))
}
async fn async_batch_get_cache(&self, keys: &[String]) -> Result<Vec<Option<Value>>, Error> {
let mut entries = self.entries.lock();
Ok(keys
.iter()
.map(|key| Self::live(&mut entries, key))
.collect())
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
self.entries.lock().remove(key);
Ok(())
}
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,456 @@
//! Recovering a stream that failed after its attempt returned it. The host wraps every streamed
//! response; when the stream fails before the caller received content, the wrapper asks the core
//! what may still take over and hands the failure back here. Mirrors the decisions of the Python
//! router's stream wrappers: `_stream_can_fall_back`, `_refusal_fallback_available`,
//! `_anthropic_messages_stream_can_retry`, `_stream_retry_budget`, `_astream_retry_same_group`,
//! and the `async_function_with_fallbacks_common_utils` re-entry of the mid-stream fallback.
use std::sync::Arc;
use serde::Deserialize;
use super::{Failed, FallbackChain, FallbackHop, Lists, Target, Tried, group_members};
use crate::{
Router, RouterStateSnapshot,
config::Op,
error::{AttemptError, Error, FailureKind},
fallback::FallbackTarget,
health::should_retry_status,
host::{RouterEvent, RouterRequest, UncountedRetries},
select::{Candidates, SelectInputs},
};
/// Where a streamed response came from: the attempt that returned it.
#[derive(Clone, Debug, Default, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct StreamOrigin {
/// The name the attempt served: the requested group, or the fallback target.
pub group: String,
/// The `order` level an order fallback targeted.
pub order: Option<i64>,
/// `fallback_depth` of the attempt.
pub depth: u32,
pub level: u32,
/// The deployment the attempt ran on.
pub deployment_id: Option<String>,
/// Same-group stream retries this request already spent (`attempted_retries_for_request`).
pub attempted_retries: u32,
/// The budget the first stream retry committed to (`committed_retry_budget_for_request`).
pub committed_budget: Option<u32>,
}
/// What may still take over a stream that fails before content.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct StreamPolicy {
/// `_anthropic_messages_stream_can_retry`.
pub can_retry: bool,
/// `_stream_can_fall_back`.
pub can_fall_back: bool,
/// `_refusal_fallback_available`.
pub refusal_fallback: bool,
/// `_stream_retry_budget` for the failure kind asked about: the budget, and whether a retry
/// policy governs the retry.
pub retry_budget: Option<(u32, bool)>,
}
impl StreamOrigin {
fn target(&self) -> Target {
Target {
group: self.group.clone(),
order: self.order,
params: None,
excluded: Arc::from([]),
failover_excluded: Arc::from([]),
}
}
/// `fallback_lookup_groups`: the served group, then the group the request was for.
fn lookup<'a>(&'a self, request: &'a RouterRequest) -> Vec<&'a str> {
let mut lookup = vec![self.group.as_str()];
if request.model != self.group {
lookup.push(request.model.as_str());
}
lookup
}
/// `has_unattempted_fallback_target`. The first attempt carries no attempted targets, so any
/// resolved chain counts; a fallback hop has already tried its own group and the original one.
fn has_unattempted(&self, request: &RouterRequest, chain: Option<Vec<String>>) -> bool {
let Some(chain) = chain else {
return false;
};
if self.depth == 0 {
return true;
}
chain
.iter()
.any(|target| *target != self.group && *target != request.model)
}
}
/// The group names a chain targets (a dict target counts by its `model`).
fn names(chain: Option<Vec<FallbackTarget>>) -> Option<Vec<String>> {
chain.map(|targets| {
targets
.iter()
.filter_map(|target| target.model().map(str::to_owned))
.collect()
})
}
/// The `content_policy_fallbacks` chain for the lookup groups: the request's list where it
/// passed one, else the router's (`_has_content_policy_fallback`).
fn content_policy_targets(lists: &Lists<'_>, lookup: &[&str]) -> Option<Option<Vec<String>>> {
lists
.content_policy
.map(|list| names(list.lookup_groups(lookup, lists.keys)))
}
/// The router's `default_fallbacks`, which Python folds into its `fallbacks` as `{"*": [...]}`: a
/// request's own list replaces them.
fn folded_defaults<'a>(snapshot: &'a RouterStateSnapshot, request: &RouterRequest) -> &'a [String] {
match request.fallbacks {
Some(_) => &[],
None => &snapshot.settings.dynamic.default_fallbacks,
}
}
/// The regular fallback targets for the lookup groups, then the folded `default_fallbacks`.
fn regular_targets(
snapshot: &RouterStateSnapshot,
request: &RouterRequest,
lists: &Lists<'_>,
lookup: &[&str],
) -> Option<Vec<String>> {
let defaults = folded_defaults(snapshot, request);
names(
lists
.fallbacks
.and_then(|list| list.lookup_groups(lookup, lists.keys)),
)
.or_else(|| (!defaults.is_empty()).then(|| defaults.to_vec()))
}
/// `_has_default_fallbacks`: the router's list, whatever the request passed.
fn has_default_fallbacks(snapshot: &RouterStateSnapshot) -> bool {
let dynamic = &snapshot.settings.dynamic;
dynamic.fallbacks.has_star() || !dynamic.default_fallbacks.is_empty()
}
impl Router {
/// What may still take over a stream from `origin` that fails before content, and the retry
/// budget for a failure of `kind`.
pub fn stream_policy(
&self,
request: &RouterRequest,
origin: &StreamOrigin,
kind: Option<FailureKind>,
) -> StreamPolicy {
let snapshot = self.core.snapshot.load_full();
StreamPolicy {
can_retry: stream_can_retry(&snapshot, request, origin),
can_fall_back: stream_can_fall_back(&snapshot, request, origin),
refusal_fallback: refusal_fallback_available(&snapshot, request, origin),
retry_budget: kind.map(|kind| stream_retry_budget(&snapshot, request, origin, kind)),
}
}
/// The fallback chain for a stream from `origin` that failed with `failure` before content:
/// `async_function_with_fallbacks_common_utils` re-entered from the stream wrapper, so the
/// group's own retry loop does not run again.
pub async fn fallback_stream(
&self,
op: Op,
request: RouterRequest,
origin: StreamOrigin,
failure: AttemptError,
) -> Result<super::RouterResponse, Error> {
self.live()?;
let snapshot = self.core.snapshot.load_full();
let request = Arc::new(request);
let mut chain = FallbackChain::default();
if request.model != origin.group {
chain
.attempted
.push(Tried::Target(FallbackTarget::Name(request.model.clone())));
}
self.core
.fall_back(
&snapshot,
&op,
&request,
origin.target(),
FallbackHop {
depth: origin.depth,
level: origin.level,
scope: 0,
},
&mut chain,
Failed::from(Error::Attempt(failure)),
)
.await
}
/// `_astream_retry_same_group`: re-run the attempt in the stream's own group until a retry
/// returns a response or the budget runs out. The failed result carries the newest failure.
pub async fn retry_stream(
&self,
op: Op,
request: RouterRequest,
origin: StreamOrigin,
failure: AttemptError,
) -> Result<super::RouterResponse, Error> {
self.live()?;
let snapshot = self.core.snapshot.load_full();
let request = Arc::new(request);
let target = origin.target();
let at = FallbackHop {
depth: origin.depth,
level: origin.level,
scope: 0,
};
let counting = SelectInputs {
count_only: true,
..SelectInputs::default()
};
let counted = |c: &Candidates, n: usize| if c.kind.is_single() { 0 } else { n };
let (healthy, all) = self
.core
.candidates(&snapshot, &op, &request, &target.group, &counting)
.await
.map(|c| (counted(&c, c.healthy.len()), counted(&c, c.all.len())))
.unwrap_or((0, 0));
let (budget, policy_applies) =
stream_retry_budget(&snapshot, &request, &origin, failure.kind);
let inputs = SelectInputs {
target_order: target.order,
requested: Some(target.group.clone()),
fallback_hop: at.level > 0,
..SelectInputs::default()
};
let mut last = failure;
for attempt in origin.attempted_retries..budget {
if last.agentic_followup {
break;
}
if !policy_applies
&& !self
.core
.should_retry(&snapshot, &request, &last, healthy, all)
{
break;
}
let delay = self.core.time_to_sleep(
&snapshot,
last.retry_after,
budget - attempt,
budget,
healthy,
all,
);
self.core.host.sleep(delay).await;
self.core.host.observe(&RouterEvent::Retry {
model_group: target.group.clone(),
retry: attempt + 1,
delay,
});
match self
.core
.attempt(
&snapshot,
&op,
&request,
&target,
&inputs,
(attempt + 1, budget),
(at, &[]),
)
.await
{
Ok(mut response) => {
response.attempted_retries = attempt + 1;
response.max_retries = Some(budget);
return Ok(response);
}
Err(failed) => match *failed.error {
// `_stream_fallback_error_for_raised`: a status the pre-stream retry rule
// would not retry ends the retries with that failure.
Error::Attempt(failure) => {
let stop = failure.status.is_some_and(|s| !should_retry_status(s));
last = failure;
if stop {
break;
}
}
other => return Err(other),
},
}
}
Err(Error::Attempt(last))
}
}
/// The request's `num_retries`, then the routed deployment's, then the router's
/// (`_anthropic_messages_plain_retry_budget`).
fn plain_retry_budget(
snapshot: &RouterStateSnapshot,
request: &RouterRequest,
origin: &StreamOrigin,
) -> u32 {
if let Some(retries) = request.num_retries {
return retries;
}
if request.uncounted_retries == Some(UncountedRetries::Negative) {
return 0;
}
let deployment = origin
.deployment_id
.as_deref()
.and_then(|id| snapshot.registry.deployment(id));
let configured = deployment.and_then(|d| match d.raw_params.get("num_retries") {
Some(serde_json::Value::Number(n)) => n.as_u64().and_then(|n| u32::try_from(n).ok()),
Some(serde_json::Value::String(s)) if s.chars().all(|c| c.is_ascii_digit()) => {
s.parse().ok()
}
_ => None,
});
configured.unwrap_or(snapshot.settings.num_retries)
}
/// The retry policy for the request's model group (`_anthropic_messages_resolved_retry_policy`),
/// unless the request opted out with `num_retries=0`. A policy that does not read as one governs
/// nothing here, with no fall through to the router's `retry_policy`.
fn retry_policy<'a>(
snapshot: &'a RouterStateSnapshot,
request: &'a RouterRequest,
origin: &StreamOrigin,
) -> Option<&'a crate::config::RetryPolicy> {
if request.num_retries == Some(0) || request.retry_policy_raises.contains(&origin.group) {
return None;
}
let settings = &snapshot.settings;
crate::retry::resolve_retry_policy(
settings.retry_policy.as_ref(),
&origin.group,
Some(
request
.model_group_retry_policy
.as_ref()
.unwrap_or(&settings.model_group_retry_policy),
),
)
}
fn stream_retry_budget(
snapshot: &RouterStateSnapshot,
request: &RouterRequest,
origin: &StreamOrigin,
kind: FailureKind,
) -> (u32, bool) {
let failure = AttemptError::new(kind, None);
let policy_retries =
retry_policy(snapshot, request, origin).and_then(|p| p.retries_for(&failure));
if origin.attempted_retries > 0
&& let Some(committed) = origin.committed_budget
{
return (committed, policy_retries.is_some());
}
match policy_retries {
Some(retries) => (retries, true),
None => (plain_retry_budget(snapshot, request, origin), false),
}
}
fn stream_can_retry(
snapshot: &RouterStateSnapshot,
request: &RouterRequest,
origin: &StreamOrigin,
) -> bool {
let attempted = origin.attempted_retries;
if attempted > 0
&& let Some(committed) = origin.committed_budget
{
return committed > attempted;
}
let plain = plain_retry_budget(snapshot, request, origin);
// `_retry_policy_ceiling`: every field the policy dumps.
let ceiling = retry_policy(snapshot, request, origin).map_or(plain, |policy| {
let granted = [
policy.BadRequestErrorRetries,
policy.AuthenticationErrorRetries,
policy.TimeoutErrorRetries,
policy.RateLimitErrorRetries,
policy.ContentPolicyViolationErrorRetries,
policy.InternalServerErrorRetries,
policy.ServiceUnavailableErrorRetries,
policy.NotFoundErrorRetries,
policy.DefaultRetries,
];
plain.max(granted.into_iter().flatten().max().unwrap_or(0))
});
ceiling > attempted
}
fn stream_can_fall_back(
snapshot: &RouterStateSnapshot,
request: &RouterRequest,
origin: &StreamOrigin,
) -> bool {
if request.disable_fallbacks {
return false;
}
if snapshot.settings.enable_weighted_failover {
return true;
}
// `_stream_order_levels`, the levels the fallback walk sees for the group.
let listing_team = request.listing_team_id.as_deref();
let mut orders: Vec<i64> = group_members(snapshot, &origin.group, listing_team)
.iter()
.filter(|d| d.listed_for_team(listing_team))
.filter_map(|d| d.order())
.collect();
orders.sort_unstable();
orders.dedup();
if orders.len() > 1 {
let current = origin.order.unwrap_or(orders[0]);
if orders.iter().any(|order| *order > current) {
return true;
}
}
let lists = Lists::of(snapshot, request);
let lookup = origin.lookup(request);
if let Some(chain) = content_policy_targets(&lists, &lookup)
&& origin.has_unattempted(request, chain)
{
return true;
}
if !lists.has_regular() && folded_defaults(snapshot, request).is_empty() {
return false;
}
if lists
.fallbacks
.is_some_and(|list| list.is_client_side(lists.keys))
{
return true;
}
origin.has_unattempted(request, regular_targets(snapshot, request, &lists, &lookup))
}
fn refusal_fallback_available(
snapshot: &RouterStateSnapshot,
request: &RouterRequest,
origin: &StreamOrigin,
) -> bool {
if request.disable_fallbacks {
return false;
}
let lists = Lists::of(snapshot, request);
let lookup = origin.lookup(request);
if let Some(chain) = content_policy_targets(&lists, &lookup) {
return origin.has_unattempted(request, chain);
}
if has_default_fallbacks(snapshot) {
return true;
}
origin.has_unattempted(request, regular_targets(snapshot, request, &lists, &lookup))
}

View file

@ -0,0 +1,355 @@
//! What the router is told: [`RouterConfig`], fixed for the router's life, and
//! [`RouterDynamicConfig`], a versioned snapshot read once per call. The router never reads
//! process globals.
use std::{collections::BTreeMap, time::Duration};
use serde::Deserialize;
use crate::{
error::AttemptError,
registry::{DeploymentSpec, RoutingGroupSpec},
};
/// `DEFAULT_COOLDOWN_TIME_SECONDS` in `litellm/constants.py`.
pub const DEFAULT_COOLDOWN_TIME: Duration = Duration::from_secs(5);
/// `ROUTER_MAX_FALLBACKS` in `litellm/constants.py`.
pub const DEFAULT_MAX_FALLBACKS: u32 = 5;
/// The router operation a call runs: the Python router's endpoint method name (`acompletion`,
/// `aembedding`, ...). The core runs every op alike; what differs per endpoint is the host's.
#[derive(Clone, Debug, PartialEq, Eq, Hash, Deserialize)]
#[serde(transparent)]
pub struct Op(String);
impl Op {
pub fn new(name: impl Into<String>) -> Self {
Self(name.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl From<&str> for Op {
fn from(name: &str) -> Self {
Self::new(name)
}
}
impl std::fmt::Display for Op {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
/// `routing_strategy`. Only simple-shuffle is native; the rest go through `Host::select`.
#[derive(
Clone,
Debug,
Default,
PartialEq,
Eq,
strum::EnumString,
strum::Display,
serde_with::DeserializeFromStr,
)]
pub enum RoutingStrategy {
#[default]
#[strum(serialize = "simple-shuffle")]
SimpleShuffle,
#[strum(serialize = "least-busy")]
LeastBusy,
#[strum(serialize = "usage-based-routing")]
UsageBased,
#[strum(serialize = "usage-based-routing-v2")]
UsageBasedV2,
#[strum(serialize = "latency-based-routing")]
LatencyBased,
#[strum(serialize = "cost-based-routing")]
CostBased,
/// `routing_strategy=None` passed explicitly (sent as `"None"`): stored as None, so no
/// strategy arm matches and every pick finds no deployment.
#[strum(serialize = "None")]
Unset,
#[strum(default, transparent)]
Other(String),
}
impl RoutingStrategy {
pub fn is_native(&self) -> bool {
matches!(self, Self::SimpleShuffle | Self::Unset)
}
}
/// `PluginKind` in `litellm/router_api.py`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, strum::Display)]
pub enum PluginKind {
#[serde(rename = "auto")]
#[strum(serialize = "auto")]
Auto,
#[serde(rename = "complexity")]
#[strum(serialize = "complexity")]
Complexity,
#[serde(rename = "quality")]
#[strum(serialize = "quality")]
Quality,
#[serde(rename = "adaptive")]
#[strum(serialize = "adaptive")]
Adaptive,
#[serde(rename = "routing")]
#[strum(serialize = "routing")]
Routing,
#[serde(rename = "budget_limiter")]
#[strum(serialize = "budget_limiter")]
BudgetLimiter,
#[serde(rename = "prompt_caching_pins")]
#[strum(serialize = "prompt_caching_pins")]
PromptCachingPins,
}
/// One registered plugin. Its behaviour lives in the host until it is ported.
#[derive(Clone, Debug, PartialEq, Deserialize)]
pub struct PluginSpec {
pub kind: PluginKind,
pub name: String,
#[serde(default)]
pub config: serde_json::Value,
}
/// `RetryPolicy` in `litellm/types/router.py`: retries per exception type.
#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
#[allow(non_snake_case)]
pub struct RetryPolicy {
pub BadRequestErrorRetries: Option<u32>,
pub AuthenticationErrorRetries: Option<u32>,
pub TimeoutErrorRetries: Option<u32>,
pub RateLimitErrorRetries: Option<u32>,
pub ContentPolicyViolationErrorRetries: Option<u32>,
pub InternalServerErrorRetries: Option<u32>,
pub ServiceUnavailableErrorRetries: Option<u32>,
pub NotFoundErrorRetries: Option<u32>,
pub DefaultRetries: Option<u32>,
}
/// `AllowedFailsPolicy` in `litellm/types/router.py`: allowed fails per minute per exception type.
#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
#[allow(non_snake_case)]
pub struct AllowedFailsPolicy {
pub BadRequestErrorAllowedFails: Option<u32>,
pub AuthenticationErrorAllowedFails: Option<u32>,
pub TimeoutErrorAllowedFails: Option<u32>,
pub RateLimitErrorAllowedFails: Option<u32>,
pub ContentPolicyViolationErrorAllowedFails: Option<u32>,
pub InternalServerErrorAllowedFails: Option<u32>,
pub ServiceUnavailableErrorAllowedFails: Option<u32>,
pub BadGatewayErrorAllowedFails: Option<u32>,
pub NotFoundErrorAllowedFails: Option<u32>,
}
impl AllowedFailsPolicy {
/// `_core_get_allowed_fails_from_policy`: the first set field whose exception class the
/// failure is an instance of, in Python's order.
pub fn allowed_fails_for(&self, failure: &AttemptError) -> Option<u32> {
[
("AuthenticationError", self.AuthenticationErrorAllowedFails),
("Timeout", self.TimeoutErrorAllowedFails),
("RateLimitError", self.RateLimitErrorAllowedFails),
(
"ContentPolicyViolationError",
self.ContentPolicyViolationErrorAllowedFails,
),
("BadRequestError", self.BadRequestErrorAllowedFails),
("InternalServerError", self.InternalServerErrorAllowedFails),
(
"ServiceUnavailableError",
self.ServiceUnavailableErrorAllowedFails,
),
("BadGatewayError", self.BadGatewayErrorAllowedFails),
("NotFoundError", self.NotFoundErrorAllowedFails),
]
.into_iter()
.find_map(|(class, allowed)| allowed.filter(|_| failure.is_a(class)))
}
}
pub use crate::fallback::{FallbackKeys, Fallbacks};
/// Static config, fixed for the router's life.
#[derive(Clone, Debug, Default, Deserialize)]
#[serde(default)]
pub struct RouterConfig {
pub model_list: Vec<DeploymentSpec>,
pub routing_strategy: RoutingStrategy,
pub routing_groups: Vec<RoutingGroupSpec>,
/// `_model_to_group` as Python last built it (member model group to routing group), set
/// only when a raw `routing_groups` write replaced the groups since: member calls keep
/// routing by it, as the Python router's stale map does.
pub routing_group_members: Option<BTreeMap<String, String>>,
pub plugins: Vec<PluginSpec>,
/// Router default `cooldown_time`; `None` is `DEFAULT_COOLDOWN_TIME_SECONDS`.
pub cooldown_time: Option<f64>,
pub disable_cooldowns: bool,
pub allowed_fails_policy: Option<AllowedFailsPolicy>,
pub retry_policy: Option<RetryPolicy>,
/// `model_group_retry_policy`: a group's own policy, which wins over `retry_policy`.
pub model_group_retry_policy: BTreeMap<String, RetryPolicy>,
/// `retry_after`: the minimum backoff between retries, in seconds.
pub retry_after: f64,
pub timeout: Option<f64>,
pub request_timeout: Option<f64>,
pub drop_params: Option<bool>,
pub enable_pre_call_checks: bool,
/// `Router(plugins=[...])` is set: the host runs the routing-plugin pipeline.
pub routing_plugins: bool,
/// The host's `CustomLogger` deployment filters run on every async candidate list
/// (`async_callback_filter_deployments`).
pub callback_filters: bool,
pub enable_health_check_routing: bool,
/// How old (seconds) a background health check result may be and still count
/// (`health_check_staleness_threshold`, else the default interval times its multiplier).
pub health_check_staleness_threshold: Option<f64>,
/// `background_health_check_model_groups`: when set, only these groups are health-filtered.
pub background_health_check_model_groups: Option<Vec<String>>,
/// `enable_weighted_failover`: a failed deployment's group is re-picked without it before
/// any cross-group fallback (simple-shuffle only).
pub enable_weighted_failover: bool,
/// `fallback_on_unsupported_params`; `None` is Python's default, `True`.
pub fallback_on_unsupported_params: Option<bool>,
/// Deployment ids that are a provider's default and never cool down
/// (`provider_default_deployment_ids`).
pub provider_default_deployment_ids: Vec<String>,
/// The operations this router must serve: `Host::supports` must accept each one for every
/// configured provider, or the router fails to build.
pub served_ops: Vec<Op>,
}
/// Dynamic config: a versioned snapshot read once per call.
#[derive(Clone, Debug, Deserialize)]
#[serde(default)]
pub struct RouterDynamicConfig {
pub version: u64,
/// `model_group_alias` as Python spells it: each alias maps to a group name, or to
/// `{"model": <group>, "hidden": <bool>}`.
#[serde(alias = "model_group_alias", deserialize_with = "alias_targets")]
pub model_alias_map: BTreeMap<String, String>,
pub fallbacks: Fallbacks,
pub default_fallbacks: Vec<String>,
pub context_window_fallbacks: Option<Fallbacks>,
pub content_policy_fallbacks: Option<Fallbacks>,
pub num_retries: Option<u32>,
pub max_fallbacks: u32,
/// `allowed_fails`, as the router copied it when it was built.
pub allowed_fails: Option<u32>,
/// `allowed_fails` is only the process default: it was not passed to the router and still
/// equals `litellm.allowed_fails`. Anything else switches cooldowns from the failure-ratio
/// rule to the allowed-fails counter (`_is_allowed_fails_set_on_router`).
pub allowed_fails_is_default: bool,
/// Other spellings a fallback key matches.
pub fallback_keys: FallbackKeys,
/// `default_deployment`: what a name no deployment, alias, routing group, team route or
/// wildcard serves is routed to, as a copy whose `litellm_params.model` is the name
/// (`_try_early_resolve_deployments_for_model_not_in_names`).
pub default_deployment: Option<DeploymentSpec>,
/// The `default_deployment` is a value Python cannot copy for a name: resolving to it
/// raises the copy's error ([`crate::Error::DefaultDeploymentRaises`]).
pub default_deployment_raises: bool,
}
/// `model_group_alias` as `resolve_model_group_alias` reads it: an alias maps to a non-empty
/// group name, or to `{"model": <non-empty group name>, ...}`; any other value rewrites nothing.
fn alias_targets<'de, D>(deserializer: D) -> Result<BTreeMap<String, String>, D::Error>
where
D: serde::Deserializer<'de>,
{
let aliases = BTreeMap::<String, serde_json::Value>::deserialize(deserializer)?;
Ok(aliases
.into_iter()
.filter_map(|(alias, target)| {
let group = match &target {
serde_json::Value::Object(entry) => entry.get("model")?,
other => other,
};
group
.as_str()
.filter(|group| !group.is_empty())
.map(|group| (alias, group.to_owned()))
})
.collect())
}
impl Default for RouterDynamicConfig {
fn default() -> Self {
Self {
version: 0,
model_alias_map: BTreeMap::new(),
fallbacks: Fallbacks::default(),
default_fallbacks: Vec::new(),
context_window_fallbacks: None,
content_policy_fallbacks: None,
num_retries: None,
max_fallbacks: DEFAULT_MAX_FALLBACKS,
allowed_fails: None,
allowed_fails_is_default: false,
fallback_keys: FallbackKeys::default(),
default_deployment: None,
default_deployment_raises: false,
}
}
}
impl RouterDynamicConfig {
/// The `default_deployment` as resolution serves it.
pub fn default_ref(&self) -> Option<crate::resolve::DefaultDeployment<'_>> {
if self.default_deployment_raises {
return Some(crate::resolve::DefaultDeployment::Raises);
}
self.default_deployment
.as_ref()
.map(crate::resolve::DefaultDeployment::Spec)
}
/// The regular fallback targets for the lookup groups, then `default_fallbacks`, which
/// Python folds into `fallbacks` as `{"*": [...]}`: what `_has_fallback_group` reads.
pub fn fallback_targets(&self, groups: &[&str]) -> Option<Vec<String>> {
if groups
.iter()
.any(|group| self.fallbacks.has_keyed_chain(group, &self.fallback_keys))
{
return Some(Vec::new());
}
(!self.default_fallbacks.is_empty()).then(|| self.default_fallbacks.clone())
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::group_name(json!({"model_group_alias": {"gpt": "gpt-4"}}))]
#[case::hidden_entry(json!({"model_group_alias": {"gpt": {"model": "gpt-4", "hidden": true}}}))]
#[case::native_name(json!({"model_alias_map": {"gpt": "gpt-4"}}))]
fn model_group_alias_reads_both_python_spellings(#[case] wire: serde_json::Value) {
let dynamic: RouterDynamicConfig = serde_json::from_value(wire).unwrap();
assert_eq!(
dynamic.model_alias_map,
BTreeMap::from([("gpt".to_owned(), "gpt-4".to_owned())])
);
}
#[rstest]
#[case::native("simple-shuffle", RoutingStrategy::SimpleShuffle)]
#[case::unknown("my-strategy", RoutingStrategy::Other("my-strategy".to_owned()))]
fn routing_strategy_parses_its_wire_name(
#[case] wire: &str,
#[case] expected: RoutingStrategy,
) {
assert_eq!(wire.parse::<RoutingStrategy>().unwrap(), expected);
assert_eq!(expected.to_string(), wire);
}
}

View file

@ -0,0 +1,240 @@
//! Which deployments may serve a request's endpoint: a port of
//! `litellm/router_utils/endpoint_eligibility.py`. A deployment that sets
//! `enforce_supported_endpoints` and lists none of the chat-shaped endpoints serves none of them;
//! the host names the endpoint for the calls `endpoint_selection` wraps.
use std::{collections::BTreeSet, sync::Arc};
use serde_json::{Map, Value};
use crate::{error::Error, registry::Deployment, resolve::Caller};
/// `CHAT_ENDPOINTS`: `_endpoint_path` of each value of `CHAT_ENDPOINT_BY_CALL`.
const CHAT_ENDPOINTS: [&str; 8] = [
"/chat/completions",
"/completions",
"/messages",
"/responses",
"/generatecontent",
"/streamgeneratecontent",
"/interactions",
"/responses/compact",
];
/// `_endpoint_path`.
fn endpoint_path(endpoint: &str) -> String {
let rooted = format!("/{}", endpoint.trim().trim_matches('/').to_lowercase());
let path = rooted.strip_prefix("/openai").unwrap_or(&rooted);
let path = path.strip_prefix("/v1beta").unwrap_or(path);
let path = path.strip_prefix("/v1").unwrap_or(path);
match path.rsplit_once(':') {
Some((_, action)) => format!("/{action}"),
None => path.to_owned(),
}
}
/// The flag as pydantic's lax `bool | None` reads it: `None` for a value it refuses, which fails
/// the whole declaration.
fn enforce_flag(value: Option<&Value>) -> Option<bool> {
match value {
None | Some(Value::Null) => Some(false),
Some(Value::Bool(flag)) => Some(*flag),
Some(Value::Number(number)) => match number.as_f64() {
Some(0.0) => Some(false),
Some(1.0) => Some(true),
_ => None,
},
Some(Value::String(text)) => match text.to_lowercase().as_str() {
"0" | "off" | "f" | "false" | "n" | "no" => Some(false),
"1" | "on" | "t" | "true" | "y" | "yes" => Some(true),
_ => None,
},
Some(_) => None,
}
}
/// `_enforced_endpoints`: the declared endpoints of a deployment that enforces them. One endpoint
/// may be written as a string; a declaration that does not validate enforces nothing.
fn enforced_endpoints(model_info: &Map<String, Value>) -> Option<Vec<String>> {
if !enforce_flag(model_info.get("enforce_supported_endpoints"))? {
return None;
}
match model_info.get("supported_endpoints")? {
Value::String(endpoint) => Some(vec![endpoint.clone()]),
Value::Array(endpoints) => endpoints
.iter()
.map(|endpoint| endpoint.as_str().map(str::to_owned))
.collect(),
_ => None,
}
}
/// `non_chat_endpoints`: what a deployment serves when it serves no chat-shaped endpoint, else
/// `None`.
fn non_chat_endpoints(deployment: &Deployment) -> Option<BTreeSet<String>> {
let declared = enforced_endpoints(&deployment.model_info.raw)?;
let chat = declared
.iter()
.any(|endpoint| CHAT_ENDPOINTS.contains(&endpoint_path(endpoint).as_str()));
(!chat).then(|| declared.into_iter().collect())
}
/// `eligible_deployments`: the deployments that may serve the caller's endpoint. When none may,
/// the refusal names what they do serve.
pub(crate) fn eligible_deployments(
model: &str,
caller: Caller<'_>,
deployments: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, Error> {
let Some(endpoint) = caller.endpoint.filter(|_| !deployments.is_empty()) else {
return Ok(deployments);
};
let declared: Vec<Option<BTreeSet<String>>> = deployments
.iter()
.map(|deployment| non_chat_endpoints(deployment))
.collect();
let eligible: Vec<Arc<Deployment>> = deployments
.into_iter()
.zip(&declared)
.filter(|(_, endpoints)| endpoints.is_none())
.map(|(deployment, _)| deployment)
.collect();
if !eligible.is_empty() {
return Ok(eligible);
}
let supported: BTreeSet<String> = declared.into_iter().flatten().flatten().collect();
let supported = if supported.is_empty() {
"none".to_owned()
} else {
Vec::from_iter(supported).join(", ")
};
let requested = if caller.requested.is_empty() {
model
} else {
caller.requested
};
Err(Error::EndpointNotSupported {
model: model.to_owned(),
message: format!(
"{endpoint} is not supported by model={requested}. Supported endpoints: {supported}"
),
})
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
use crate::registry::DeploymentSpec;
fn deployment(id: &str, info: Value) -> Arc<Deployment> {
let mut model_info = info;
model_info["id"] = json!(id);
let spec: DeploymentSpec = serde_json::from_value(json!({
"model_name": "g", "litellm_params": {"model": "openai/a"}, "model_info": model_info,
}))
.unwrap();
Arc::new(Deployment::from_spec(spec).unwrap())
}
fn chat<'a>() -> Caller<'a> {
Caller {
endpoint: Some("/v1/chat/completions"),
requested: "g",
..Caller::default()
}
}
#[rstest]
#[case::v1_prefix("/v1/chat/completions", "/chat/completions")]
#[case::openai_prefix("/openai/v1/responses/", "/responses")]
#[case::bare_gemini_action("generateContent", "/generatecontent")]
#[case::beta_prefix("/v1beta/interactions", "/interactions")]
#[case::gemini_model_path(
"/v1beta/models/gemini:streamGenerateContent",
"/streamgeneratecontent"
)]
fn endpoint_paths_normalize(#[case] endpoint: &str, #[case] expected: &str) {
assert_eq!(endpoint_path(endpoint), expected);
}
#[rstest]
#[case::embeddings_only(
json!({"enforce_supported_endpoints": true, "supported_endpoints": ["/v1/embeddings"]}),
false
)]
#[case::one_endpoint_as_a_string(
json!({"enforce_supported_endpoints": true, "supported_endpoints": "/v1/embeddings"}),
false
)]
#[case::string_flag(
json!({"enforce_supported_endpoints": "yes", "supported_endpoints": ["/v1/embeddings"]}),
false
)]
#[case::an_empty_list_serves_nothing(
json!({"enforce_supported_endpoints": true, "supported_endpoints": []}),
false
)]
#[case::any_chat_shaped_endpoint_serves_every_one(
json!({"enforce_supported_endpoints": true, "supported_endpoints": ["/v1/responses"]}),
true
)]
#[case::not_enforced(
json!({"enforce_supported_endpoints": false, "supported_endpoints": ["/v1/embeddings"]}),
true
)]
#[case::unset_flag(json!({"supported_endpoints": ["/v1/embeddings"]}), true)]
#[case::enforced_without_a_list(json!({"enforce_supported_endpoints": true}), true)]
#[case::an_invalid_flag_enforces_nothing(
json!({"enforce_supported_endpoints": "maybe", "supported_endpoints": ["/v1/embeddings"]}),
true
)]
#[case::an_invalid_list_enforces_nothing(
json!({"enforce_supported_endpoints": true, "supported_endpoints": [1]}),
true
)]
fn a_declaration_keeps_or_drops_a_deployment(#[case] info: Value, #[case] kept: bool) {
let open = deployment("open", json!({}));
let declared = deployment("declared", info);
let eligible = eligible_deployments("g", chat(), vec![declared, open]).unwrap();
assert_eq!(eligible.len(), if kept { 2 } else { 1 });
}
#[test]
fn no_endpoint_keeps_everything() {
let declared = deployment(
"d",
json!({"enforce_supported_endpoints": true, "supported_endpoints": []}),
);
let eligible = eligible_deployments("g", Caller::default(), vec![declared]).unwrap();
assert_eq!(eligible.len(), 1);
}
#[rstest]
#[case::lists_what_they_serve(
json!(["/v1/rerank", "/v1/embeddings"]),
"/v1/chat/completions is not supported by model=g. \
Supported endpoints: /v1/embeddings, /v1/rerank"
)]
#[case::nothing_declared(
json!([]),
"/v1/chat/completions is not supported by model=g. Supported endpoints: none"
)]
fn every_deployment_refusing_names_the_endpoint(
#[case] endpoints: Value,
#[case] expected: &str,
) {
let declared = deployment(
"d",
json!({"enforce_supported_endpoints": true, "supported_endpoints": endpoints}),
);
let refused = eligible_deployments("resolved", chat(), vec![declared]);
assert!(matches!(
refused,
Err(Error::EndpointNotSupported { model, message })
if model == "resolved" && message == expected
));
}
}

View file

@ -0,0 +1,473 @@
//! The router's errors: [`ConfigError`] when a router is built or changed, [`Error`] when a call
//! or a selection fails, and [`AttemptError`], the typed facts of one failed attempt.
use std::{any::Any, fmt, sync::Arc, time::Duration};
use crate::config::{Op, RoutingStrategy};
/// What a host cannot serve, found when the router is built rather than partway through a call.
#[derive(Clone, Debug, PartialEq, thiserror::Error)]
pub enum Unsupported {
#[error("routing strategy `{0}` needs a host that can select")]
Strategy(RoutingStrategy),
#[error("routing plugin `{0}` needs a host that can select")]
Plugin(String),
#[error("the host cannot run operation `{op}` on provider `{provider}`")]
Op { op: Op, provider: String },
}
#[derive(Debug, thiserror::Error)]
pub enum ConfigError {
#[error("the router config asks for something this host cannot serve")]
Unsupported(#[from] Unsupported),
#[error("deployment `{model_name}` has invalid litellm_params")]
InvalidDeployment {
model_name: String,
#[source]
source: serde_json::Error,
},
#[error("deployment id `{id}` is configured more than once")]
DuplicateId { id: String },
/// A removal the host's indexes could not finish ([`crate::host_index::RemovalError`]).
#[error("{0}")]
Removal(String),
}
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("the router config is invalid")]
Config(#[from] ConfigError),
#[error("no deployment serves model `{model}`")]
UnknownModel { model: String },
#[error("no deployment of `{model}` is available")]
NoHealthyDeployments { model: String },
#[error("the attempt failed")]
Attempt(#[from] AttemptError),
#[error("the routing cache failed")]
Cache(#[from] litellm_cache::Error),
#[error("the router was retired")]
Retired,
/// A blocked deployment named by its id (`model` is its litellm model).
#[error("deployment model `{model}` is paused")]
Paused { model: String },
/// A pass-through pick named a deployment id without `use_in_pass_through`.
#[error("deployment `{id}` does not support pass-through")]
NotPassThrough { id: String },
/// A pass-through pick of a group with healthy deployments, none of them pass-through.
#[error("model `{model}` has no pass-through deployment")]
NoPassThroughDeployments { model: String },
/// Every deployment `model` resolved to is a pre-routing strategy marker (`auto_router/...`).
#[error("every deployment of `{model}` is a strategy router marker")]
OnlyMarkers { model: String },
/// Every deployment `model` resolved to refuses the request's endpoint
/// (`eligible_deployments`); `message` names what they serve.
#[error("{message}")]
EndpointNotSupported { model: String, message: String },
/// Every deployment `model` resolved to is split into PTU shares the caller holds none of.
#[error("deployment `{model}` is reserved for the teams holding a PTU share of it")]
PtuReserved { model: String },
/// A proxy admin without a team named a public model name several teams use.
#[error("model name `{model}` matches deployments from multiple teams")]
MultipleTeams { model: String },
/// Every deployment `model` resolved to is reserved for other teams by an active access
/// window, the first of which ends at `until` (`HH:MM`) in `timezone`.
#[error("deployment `{model}` is reserved for another team until {until} {timezone}")]
Reserved {
model: String,
until: String,
timezone: String,
},
/// The pre-call checks found no deployment whose context window fits the prompt; `detail`
/// names each one dropped for it.
#[error("no deployment of `{model}` has a context window large enough for this call")]
ContextWindowExceeded { model: String, detail: String },
/// The pre-call checks found no deployment of `model` that takes the requested `params`.
#[error("no deployment of `{model}` supports {params:?}")]
UnsupportedParams { model: String, params: Vec<String> },
/// The routing plugins left none of the healthy deployments' models.
#[error("no deployments left after routing-plugin filtering: {models:?}")]
RoutingPluginsLeftNothing { models: Vec<String> },
/// The request's `_router_weights` failed `validate_router_weights` when simple-shuffle picked.
#[error("the request's router weights are invalid")]
InvalidRouterWeights,
/// A deployment's `litellm_params` `metric` (weight, rpm or tpm) is not a number.
#[error("deployment {metric} must be numeric")]
NonNumericWeight { metric: &'static str },
/// A router with routing plugins was asked for a pick on the sync path, which never runs them.
#[error("routing plugins only run on the async selection path")]
RoutingPluginsNeedAsync,
/// The stale routing-group member map (`_model_to_group`, kept as it was when a raw
/// `routing_groups` write replaced the groups) names a group that is no longer configured:
/// Python's `_get_routing_context` raises `KeyError(group)` for every pick of the name.
#[error("routing group `{group}` is no longer configured")]
MissingRoutingGroup { group: String },
/// The router's `default_deployment` is not a value Python can copy for `model`
/// (`_try_early_resolve_deployments_for_model_not_in_names` raises the copy's error).
#[error("the default deployment cannot be served for `{model}`")]
DefaultDeploymentRaises { model: String },
/// A written model-name or team index names a position past the end of the model list:
/// Python's lookup of `model` raises `IndexError`.
#[error("a written index of `{model}` names a position past the model list")]
IndexOutOfRange { model: String },
/// A model-info query raised where the Python router raises: `exception` names the Python
/// exception class (`KeyError`, `TypeError`, `AttributeError`, `ValueError`) and `detail` its
/// argument.
#[error("the model-info query raised {exception}({detail:?})")]
ModelInfo {
exception: &'static str,
detail: String,
},
}
impl Error {
/// The failed attempt behind this error, when a deployment was called.
pub fn attempt(&self) -> Option<&AttemptError> {
match self {
Self::Attempt(attempt) => Some(attempt),
Self::Config(_)
| Self::UnknownModel { .. }
| Self::NoHealthyDeployments { .. }
| Self::Cache(_)
| Self::Retired
| Self::Paused { .. }
| Self::NotPassThrough { .. }
| Self::NoPassThroughDeployments { .. }
| Self::OnlyMarkers { .. }
| Self::EndpointNotSupported { .. }
| Self::PtuReserved { .. }
| Self::MultipleTeams { .. }
| Self::Reserved { .. }
| Self::ContextWindowExceeded { .. }
| Self::UnsupportedParams { .. }
| Self::RoutingPluginsLeftNothing { .. }
| Self::InvalidRouterWeights
| Self::NonNumericWeight { .. }
| Self::RoutingPluginsNeedAsync
| Self::MissingRoutingGroup { .. }
| Self::DefaultDeploymentRaises { .. }
| Self::IndexOutOfRange { .. }
| Self::ModelInfo { .. } => None,
}
}
/// The Python exception a pick that failed with this error raises, as the retry loop
/// decides by it (its class names, most specific first, and its status): `None` for an
/// attempt's failure and for errors no pick raises.
pub fn pick_failure(&self) -> Option<AttemptError> {
const BAD_REQUEST: &[&str] = &["BadRequestError"];
let (classes, status): (&[&str], Option<u16>) = match self {
Self::UnknownModel { .. }
| Self::OnlyMarkers { .. }
| Self::PtuReserved { .. }
| Self::NotPassThrough { .. }
| Self::NoPassThroughDeployments { .. }
| Self::MultipleTeams { .. }
| Self::Reserved { .. } => (BAD_REQUEST, Some(400)),
Self::EndpointNotSupported { .. } => (
&["EndpointNotSupportedByModelError", "BadRequestError"],
Some(400),
),
Self::ContextWindowExceeded { .. } => (
&["ContextWindowExceededError", "BadRequestError"],
Some(400),
),
Self::UnsupportedParams { .. } => {
(&["UnsupportedParamsError", "BadRequestError"], Some(400))
}
Self::Paused { .. } => (&["ServiceUnavailableError"], Some(503)),
Self::NoHealthyDeployments { .. } => (&["RouterRateLimitError", "ValueError"], None),
Self::RoutingPluginsLeftNothing { .. } | Self::RoutingPluginsNeedAsync => {
(&["ValueError"], None)
}
Self::InvalidRouterWeights => (&["ValidationError", "ValueError"], None),
Self::NonNumericWeight { .. } => (&["TypeError"], None),
Self::MissingRoutingGroup { .. } => (&["KeyError", "LookupError"], None),
Self::IndexOutOfRange { .. } => (&["IndexError", "LookupError"], None),
Self::DefaultDeploymentRaises { .. } => (&["Exception"], None),
Self::Attempt(_)
| Self::Config(_)
| Self::Cache(_)
| Self::Retired
| Self::ModelInfo { .. } => return None,
};
let mut failure = AttemptError::host_raised(classes);
failure.status = status;
Some(failure)
}
/// A copy of a pick's error, for the events that name it; `None` for any other error.
pub fn pick_copy(&self) -> Option<Self> {
Some(match self {
Self::UnknownModel { model } => Self::UnknownModel {
model: model.clone(),
},
Self::NoHealthyDeployments { model } => Self::NoHealthyDeployments {
model: model.clone(),
},
Self::Paused { model } => Self::Paused {
model: model.clone(),
},
Self::NotPassThrough { id } => Self::NotPassThrough { id: id.clone() },
Self::NoPassThroughDeployments { model } => Self::NoPassThroughDeployments {
model: model.clone(),
},
Self::OnlyMarkers { model } => Self::OnlyMarkers {
model: model.clone(),
},
Self::EndpointNotSupported { model, message } => Self::EndpointNotSupported {
model: model.clone(),
message: message.clone(),
},
Self::PtuReserved { model } => Self::PtuReserved {
model: model.clone(),
},
Self::MultipleTeams { model } => Self::MultipleTeams {
model: model.clone(),
},
Self::Reserved {
model,
until,
timezone,
} => Self::Reserved {
model: model.clone(),
until: until.clone(),
timezone: timezone.clone(),
},
Self::ContextWindowExceeded { model, detail } => Self::ContextWindowExceeded {
model: model.clone(),
detail: detail.clone(),
},
Self::UnsupportedParams { model, params } => Self::UnsupportedParams {
model: model.clone(),
params: params.clone(),
},
Self::RoutingPluginsLeftNothing { models } => Self::RoutingPluginsLeftNothing {
models: models.clone(),
},
Self::InvalidRouterWeights => Self::InvalidRouterWeights,
Self::NonNumericWeight { metric } => Self::NonNumericWeight { metric },
Self::RoutingPluginsNeedAsync => Self::RoutingPluginsNeedAsync,
Self::MissingRoutingGroup { group } => Self::MissingRoutingGroup {
group: group.clone(),
},
Self::DefaultDeploymentRaises { model } => Self::DefaultDeploymentRaises {
model: model.clone(),
},
Self::IndexOutOfRange { model } => Self::IndexOutOfRange {
model: model.clone(),
},
Self::Attempt(_)
| Self::Config(_)
| Self::Cache(_)
| Self::Retired
| Self::ModelInfo { .. } => return None,
})
}
}
/// The class of a failed attempt, as litellm's exception mapping names it.
#[derive(
Clone,
Copy,
Debug,
PartialEq,
Eq,
Hash,
strum::Display,
strum::EnumString,
strum::IntoStaticStr,
strum::VariantArray,
)]
pub enum FailureKind {
#[strum(serialize = "RateLimitError")]
RateLimit,
#[strum(serialize = "ContextWindowExceededError")]
ContextWindowExceeded,
#[strum(serialize = "ContentPolicyViolationError")]
ContentPolicyViolation,
#[strum(serialize = "AuthenticationError")]
Authentication,
#[strum(serialize = "PermissionDeniedError")]
PermissionDenied,
#[strum(serialize = "NotFoundError")]
NotFound,
#[strum(serialize = "BadRequestError")]
BadRequest,
#[strum(serialize = "Timeout")]
Timeout,
#[strum(serialize = "APIConnectionError")]
ApiConnection,
#[strum(serialize = "InternalServerError")]
InternalServer,
#[strum(serialize = "ServiceUnavailableError")]
ServiceUnavailable,
#[strum(serialize = "BadGatewayError")]
BadGateway,
/// A guardrail blocked the request: never retried, never a fallback.
#[strum(serialize = "GuardrailIntervention")]
GuardrailIntervention,
/// litellm rejected the request before calling the deployment (`MalformedRequestFieldError`,
/// `UnsupportedEndpointError`): never cools the deployment down.
#[strum(serialize = "Declined")]
Declined,
/// The caller's own stored credential was rejected, not the deployment's
/// (`CallerCredentialAuthenticationError`): never cools the deployment down.
#[strum(serialize = "CallerCredentialAuthenticationError")]
CallerCredentialAuthentication,
/// A per-user credential exchange was rate limited (`CallerCredentialRateLimitError`): never
/// cools the deployment down.
#[strum(serialize = "CallerCredentialRateLimitError")]
CallerCredentialRateLimit,
#[strum(serialize = "APIError")]
Other,
}
/// An opaque value the host hands through the router unchanged, such as the original Python
/// exception or response object.
#[derive(Clone)]
pub struct Payload(pub Arc<dyn Any + Send + Sync>);
impl fmt::Debug for Payload {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("Payload(..)")
}
}
/// One failed attempt, as typed facts the retry, fallback and cooldown rules decide by.
#[derive(Clone, Debug, thiserror::Error)]
#[error("the deployment call failed with {kind}")]
pub struct AttemptError {
pub kind: FailureKind,
pub status: Option<u16>,
/// The provider's `retry-after` header (`_get_retry_after_from_exception_header`).
pub retry_after: Option<Duration>,
/// A deployment-level `num_retries` the failure carries (Python's `exception.num_retries`).
pub num_retries: Option<u32>,
/// The deployment the exception itself names (`_deployment_ids_to_skip_on_retry` reads its
/// `retry_skip_deployment_id`, then `failed_deployment_id`): an exception raised again keeps
/// the id its first failure stamped, whichever deployment raised it now.
pub failed_deployment_id: Option<String>,
/// The provider's error text, kept only for the cooldown record Python writes
/// (`exception_received`, masked past 50 characters).
pub detail: Option<Arc<str>>,
/// The exception's class names, most specific first (its `__mro__`), for the rules Python
/// decides by `isinstance`. Empty when the host reports none; the kind then stands in.
pub classes: Vec<Arc<str>>,
/// The host already recorded what this failure does to the deployment's health (a routing
/// pre-call check's failure, which Python's attempt handles itself: a rate limit cools the
/// deployment down, anything else records nothing), or no failure callback runs for it (it
/// was raised outside the litellm call, or a shared logging object already logged a failure),
/// so the core records nothing.
pub recorded: bool,
/// The exception names its provider (`llm_provider` is set), which marks a provider's own
/// 400 for `fallback_on_unsupported_params`.
pub names_provider: bool,
/// An agentic follow-up call failed after the first call was served
/// (`is_agentic_followup_failure`): never retried, since a retry reruns every paid round. The
/// failure keeps its own class, which decides its fallbacks and cooldown as for any failure.
pub agentic_followup: bool,
/// One of `deployment_callback_on_failure`'s early exits holds: the failure says nothing about
/// the deployment, so it is neither counted nor cooled down (a failed request still counts
/// against the minute's requests, which `async_deployment_callback_on_failure` does apart).
pub outside_deployment: bool,
/// One of `_should_run_cooldown_logic`'s per-caller exits holds (`_is_per_user_provider_request`
/// or a per-user Copilot deployment): the failure is counted but never cools the deployment.
pub caller_scoped: bool,
/// The exception has no `status_code` (Python's status `""`): `_is_cooldown_required` holds
/// no cooldown for it, unlike a status the host could not read, which cools down as a 500.
pub status_absent: bool,
/// The failure counts against no per-minute requests: `async_deployment_callback_on_failure`
/// does not run for a sync call, and returns at once for a batch retrieve.
pub uncounted_request: bool,
pub payload: Option<Payload>,
}
impl AttemptError {
pub fn new(kind: FailureKind, status: Option<u16>) -> Self {
Self {
kind,
status,
retry_after: None,
num_retries: None,
failed_deployment_id: None,
detail: None,
classes: Vec::new(),
recorded: false,
names_provider: false,
agentic_followup: false,
outside_deployment: false,
caller_scoped: false,
status_absent: false,
uncounted_request: false,
payload: None,
}
}
/// The exception's class names, most specific first: the host's, else the kind's.
pub fn class_names(&self) -> Vec<&str> {
if self.classes.is_empty() {
return crate::retry::kind_classes(self.kind).to_vec();
}
self.classes.iter().map(AsRef::as_ref).collect()
}
/// `isinstance(error, <class named name>)`.
pub fn is_a(&self, name: &str) -> bool {
self.class_names().contains(&name)
}
#[must_use]
pub fn with_retry_after(self, retry_after: Duration) -> Self {
Self {
retry_after: Some(retry_after),
..self
}
}
/// A plain Python exception the host raised outside any attempt (a `KeyError` from a stale
/// routing-group lookup, an error a retry policy lookup raised): no status, no provider,
/// nothing recorded against a deployment. `classes` are its class names, most specific first.
pub fn host_raised(classes: &[&str]) -> Self {
Self {
classes: classes.iter().map(|class| (*class).into()).collect(),
recorded: true,
..Self::new(FailureKind::Other, None)
}
}
/// The status Python's cooldown rules read: the provider's, else 500.
pub fn cooldown_status(&self) -> u16 {
self.status.unwrap_or(500)
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::rate_limit("RateLimitError", FailureKind::RateLimit)]
#[case::context_window("ContextWindowExceededError", FailureKind::ContextWindowExceeded)]
#[case::caller_credential_auth(
"CallerCredentialAuthenticationError",
FailureKind::CallerCredentialAuthentication
)]
#[case::caller_credential_rate_limit(
"CallerCredentialRateLimitError",
FailureKind::CallerCredentialRateLimit
)]
#[case::other("APIError", FailureKind::Other)]
fn failure_kind_round_trips_its_exception_name(#[case] wire: &str, #[case] kind: FailureKind) {
assert_eq!(wire.parse::<FailureKind>().unwrap(), kind);
assert_eq!(<&'static str>::from(kind), wire);
}
#[rstest]
#[case::empty("")]
#[case::lowercase("ratelimiterror")]
fn failure_kind_rejects_unknown_names(#[case] wire: &str) {
assert!(wire.parse::<FailureKind>().is_err());
}
}

View file

@ -0,0 +1,330 @@
//! Fallback lists as Python spells them, and how a failed group's chain is found in one.
//! Mirrors `get_fallback_model_group`, `get_fallback_model_group_for_lookup_groups`,
//! `check_non_standard_fallback_format` and `fallback_attempt_key` in
//! `router_utils/fallback_event_handlers.py`.
//!
//! A list holds keyed entries (`{"gpt-4": ["gpt-3.5", {"model": "x", "temperature": 0}]}`),
//! bare names (`"claude"`, a chain for every group) and, in a client-side list, dict targets
//! (`{"model": "x", "messages": [...]}`), which override request params for their attempt.
use serde::Deserialize;
use serde_json::{Map, Value};
use std::collections::BTreeMap;
/// One item of a fallback list, kept as Python holds it.
#[derive(Clone, Debug, PartialEq)]
enum Entry {
Name(String),
Dict(Map<String, Value>),
/// Anything else: Python skips it when it looks a chain up.
Other,
}
/// A fallback list, in order. `"*"` keys the default chain.
#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
#[serde(from = "Vec<Value>")]
pub struct Fallbacks(Vec<Entry>);
impl From<Vec<Value>> for Fallbacks {
fn from(entries: Vec<Value>) -> Self {
Self(
entries
.into_iter()
.map(|entry| match entry {
Value::String(name) => Entry::Name(name),
Value::Object(dict) => Entry::Dict(dict),
_ => Entry::Other,
})
.collect(),
)
}
}
/// One target of a chain: a group name, or a dict target that also overrides request params.
#[derive(Clone, Debug, PartialEq)]
pub enum FallbackTarget {
Name(String),
Params(Map<String, Value>),
}
impl FallbackTarget {
fn of(value: &Value) -> Option<Self> {
match value {
Value::String(name) => Some(Self::Name(name.clone())),
Value::Object(dict) => Some(Self::Params(dict.clone())),
_ => None,
}
}
/// `_get_fallback_target_model_group`: the group the target names, if any.
pub fn model(&self) -> Option<&str> {
match self {
Self::Name(name) => Some(name),
Self::Params(dict) => dict.get("model").and_then(Value::as_str),
}
}
/// The `model` of a dict target that sets a non-string one: the hop routes that value, which
/// no group answers (`_get_fallback_target_model_group` finds no group to check, but
/// `kwargs.update(target)` still rebinds the request's model).
pub fn unroutable_model(&self) -> Option<&Value> {
match self {
Self::Params(dict) => dict.get("model").filter(|model| !model.is_string()),
Self::Name(_) => None,
}
}
/// `fallback_attempt_key`: a bare name and `{"model": name}` are the same attempt; a dict
/// carrying anything else is its own.
pub fn identity(&self) -> FallbackTarget {
match self {
Self::Params(dict) if dict.len() == 1 => match dict.get("model") {
Some(Value::String(name)) => Self::Name(name.clone()),
_ => self.clone(),
},
_ => self.clone(),
}
}
/// The params a dict target merges into the request (`kwargs.update(target)`).
pub fn params(&self) -> Option<&Map<String, Value>> {
match self {
Self::Params(dict) => Some(dict),
Self::Name(_) => None,
}
}
}
/// What fallback keys match besides the exact group name: Python's `provider_list`, for
/// `"openai/gpt-4o"` to use a `"gpt-4o"` entry, the provider-prefixed spelling of each bare
/// group name (`_provider_prefixed_model_group`), for `"gpt-4o"` to use an `"openai/gpt-4o"`
/// entry, and the `LiteLLMParamsTypedDict` keys that mark a client-side dict target.
#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
#[serde(default)]
pub struct FallbackKeys {
pub providers: Vec<String>,
pub provider_prefixed: BTreeMap<String, String>,
pub param_keys: Vec<String>,
}
impl FallbackKeys {
/// `_check_stripped_model_group`.
fn is_stripped(&self, group: &str, key: &str) -> bool {
self.providers.iter().any(|provider| {
let prefix = format!("{provider}/");
group.starts_with(&prefix) && group.replace(&prefix, "") == key
})
}
}
/// A chain found for a group, and whether it is the `"*"` chain itself.
struct Found {
targets: Vec<FallbackTarget>,
generic: bool,
}
/// The targets of a keyed entry's value: a list of names and dicts. A bare string is walked
/// character by character, as Python iterates it.
fn chain_of(value: &Value) -> Vec<FallbackTarget> {
match value {
Value::Array(items) => items.iter().filter_map(FallbackTarget::of).collect(),
Value::String(text) => text
.chars()
.map(|c| FallbackTarget::Name(c.to_string()))
.collect(),
_ => Vec::new(),
}
}
impl Fallbacks {
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
/// `check_non_standard_fallback_format`: a client-side list of names, or of dicts that set
/// a litellm param, is one chain for every group.
pub fn is_client_side(&self, keys: &FallbackKeys) -> bool {
if self.0.is_empty() {
return false;
}
if self.0.iter().all(|entry| matches!(entry, Entry::Name(_))) {
return true;
}
let dicts: Vec<_> = self
.0
.iter()
.filter_map(|entry| match entry {
Entry::Dict(dict) => Some(dict),
_ => None,
})
.collect();
dicts.len() == self.0.len()
&& dicts.iter().any(|dict| {
keys.param_keys
.iter()
.any(|key| dict.get(key).is_some_and(|value| !value.is_array()))
})
}
/// The whole list as one chain (a client-side list).
pub fn as_chain(&self) -> Vec<FallbackTarget> {
self.0
.iter()
.filter_map(|entry| match entry {
Entry::Name(name) => Some(FallbackTarget::Name(name.clone())),
Entry::Dict(dict) => Some(FallbackTarget::Params(dict.clone())),
Entry::Other => None,
})
.collect()
}
/// `_get_first_default_fallback`: the first target of the first `"*"` entry.
pub fn default_target(&self) -> Option<&str> {
self.0.iter().find_map(|entry| match entry {
Entry::Dict(dict) => match dict.get("*") {
Some(Value::Array(targets)) => targets.first().and_then(Value::as_str),
_ => None,
},
_ => None,
})
}
/// `_has_default_fallbacks`: some entry keys `"*"`.
pub fn has_star(&self) -> bool {
self.0
.iter()
.any(|entry| matches!(entry, Entry::Dict(dict) if dict.contains_key("*")))
}
/// `get_fallback_model_group`.
fn lookup(&self, group: &str, keys: &FallbackKeys) -> Option<Found> {
let prefixed = keys.provider_prefixed.get(group).map(String::as_str);
let mut generic = None;
let mut stripped = None;
let mut found = None;
for entry in &self.0 {
match entry {
Entry::Dict(dict) => {
let Some((key, value)) = dict.iter().next() else {
continue;
};
if key == group {
found = Some(chain_of(value));
break;
} else if Some(key.as_str()) == prefixed || keys.is_stripped(group, key) {
stripped = Some(chain_of(value));
} else if key == "*" {
generic = Some(value);
}
}
Entry::Name(name) => found = Some(vec![FallbackTarget::Name(name.clone())]),
Entry::Other => {}
}
}
found
.or(stripped)
.map(|targets| Found {
targets,
generic: false,
})
.or_else(|| {
generic.map(|value| Found {
targets: chain_of(value),
generic: true,
})
})
}
/// `get_fallback_model_group_for_lookup_groups`: the first lookup group with its own chain
/// wins; `"*"` applies only after every group missed.
pub fn lookup_groups(
&self,
groups: &[&str],
keys: &FallbackKeys,
) -> Option<Vec<FallbackTarget>> {
let results: Vec<Found> = groups
.iter()
.filter_map(|group| self.lookup(group, keys))
.collect();
let index = results
.iter()
.position(|found| !found.generic)
.or_else(|| (!results.is_empty()).then_some(0))?;
results.into_iter().nth(index).map(|found| found.targets)
}
/// `_has_fallback_group`'s view: only keyed entries whose targets are names or dicts count.
pub fn has_keyed_chain(&self, group: &str, keys: &FallbackKeys) -> bool {
let keyed = Fallbacks(
self.0
.iter()
.filter(|entry| matches!(entry, Entry::Dict(_)))
.cloned()
.collect(),
);
keyed
.lookup(group, keys)
.is_some_and(|found| !found.targets.is_empty())
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::an_int(json!({"model": 5}), true)]
#[case::none(json!({"model": null}), true)]
#[case::a_string(json!({"model": "b"}), false)]
#[case::no_model(json!({"temperature": 0}), false)]
fn only_a_set_non_string_model_is_unroutable(#[case] target: Value, #[case] unroutable: bool) {
let target = FallbackTarget::of(&target).unwrap();
assert_eq!(target.unroutable_model().is_some(), unroutable);
}
fn names(targets: Option<Vec<FallbackTarget>>) -> Option<Vec<String>> {
targets.map(|targets| {
targets
.iter()
.map(|t| t.model().unwrap_or_default().to_owned())
.collect()
})
}
#[rstest]
#[case::own_entry(json!([{"a": ["b"]}, {"*": ["d"]}]), "a", Some(vec!["b"]))]
#[case::star_entry(json!([{"a": ["b"]}, {"*": ["d"]}]), "x", Some(vec!["d"]))]
#[case::string_is_every_groups_chain(json!(["b", {"x": ["c"]}]), "a", Some(vec!["b"]))]
#[case::keyed_entry_wins(json!(["c", {"a": ["b"]}]), "a", Some(vec!["b"]))]
#[case::string_value_is_iterated(json!([{"a": "bc"}]), "a", Some(vec!["b", "c"]))]
#[case::none(json!([{"x": ["c"]}]), "a", None)]
fn chains_follow_python(
#[case] list: Value,
#[case] group: &str,
#[case] expected: Option<Vec<&str>>,
) {
let fallbacks: Fallbacks = serde_json::from_value(list).unwrap();
assert_eq!(
names(fallbacks.lookup_groups(&[group], &FallbackKeys::default())),
expected.map(|e| e.into_iter().map(str::to_owned).collect())
);
}
#[rstest]
#[case::names(json!(["a", "b"]), true)]
#[case::param_dicts(json!([{"model": "a", "temperature": 1}]), true)]
#[case::keyed(json!([{"a": ["b"]}]), false)]
#[case::mixed(json!(["a", {"a": ["b"]}]), false)]
fn client_side_lists_follow_python(#[case] list: Value, #[case] expected: bool) {
let fallbacks: Fallbacks = serde_json::from_value(list).unwrap();
let keys = FallbackKeys {
param_keys: vec!["model".to_owned(), "temperature".to_owned()],
..FallbackKeys::default()
};
assert_eq!(fallbacks.is_client_side(&keys), expected);
}
}

View file

@ -0,0 +1,646 @@
//! Model group info and the configured-value reads: what the router reports a model group, or a
//! concrete deployment behind a listed name, can do.
//!
//! The rules are the Python router core's (`_set_model_group_info`, `get_model_listing_info`,
//! `get_configured_token_limits`, `get_configured_mode`, `get_configured_display_name`,
//! `get_configured_service_tiers`). The host picks the deployments (the registry lookups) and
//! answers what needs its objects as [`Host::model_info`] lookups asked while the query runs: a
//! deployment's validated params, its merged model info and provider, its capabilities, and the
//! group's first entry as its model validates it. The core decides which deployments count and
//! how their values combine.
use serde_json::{Map, Number, Value};
use crate::{
error::Error,
host::{Host, ModelInfoLookup},
model_info::{
discovered, is_opaque, lookup, malformed, python_float, python_int, raised, truthy,
without_digit_underscores,
},
};
/// One `_set_model_group_info` query.
#[derive(Clone, Debug, Default, serde::Deserialize)]
#[serde(default)]
pub struct GroupInfoQuery {
/// The model group the deployments serve (an alias already resolved).
pub model_group: String,
/// The name the group is reported as.
pub user_facing_model_group_name: String,
/// The group is a team's public model name the team's deployments serve.
pub serves_team_deployments: bool,
/// A wildcard pattern routes `model_group` (asked only when an entry needs it).
pub pattern_routes: bool,
/// The first string `model_info.description` among the deployments.
pub description: Value,
/// The deployments the group routes to, as stored.
pub deployments: Vec<Value>,
}
/// The group info: `None` when no deployment matched, else the fields the merge assigned after
/// the host built the group from its first deployment ([`ModelInfoLookup::GroupSeed`]).
#[derive(Clone, Debug, Default, PartialEq)]
pub struct GroupInfo {
pub assignments: Map<String, Value>,
}
/// The type name Python reports for a value, in its error messages.
pub(crate) fn type_name(value: &Value) -> &'static str {
match value {
Value::Null => "NoneType",
Value::Bool(_) => "bool",
Value::Number(number) if number.is_f64() => "float",
Value::Number(_) => "int",
Value::String(_) => "str",
Value::Array(_) => "list",
Value::Object(_) if is_opaque(value) => "object",
Value::Object(_) => "dict",
}
}
/// A number as Python's arithmetic reads it (a bool is 0 or 1), or `None`.
fn arithmetic(value: &Value) -> Option<Number> {
match value {
Value::Bool(flag) => Some(Number::from(u8::from(*flag))),
Value::Number(number) => Some(number.clone()),
_ => None,
}
}
/// `total += value`, with `total` starting at `0`; adding a non-number raises `TypeError`.
fn add(total: Option<Value>, value: &Value) -> Result<Value, Error> {
let total = total.unwrap_or_else(|| Value::Number(Number::from(0)));
let (Some(left), Some(right)) = (arithmetic(&total), arithmetic(value)) else {
return Err(raised(
"TypeError",
format!(
"unsupported operand type(s) for +=: '{}' and '{}'",
type_name(&total),
type_name(value)
),
));
};
if let (Some(left), Some(right)) = (left.as_i64(), right.as_i64())
&& let Some(sum) = left.checked_add(right)
{
return Ok(Value::Number(Number::from(sum)));
}
let sum = left.as_f64().unwrap_or(0.0) + right.as_f64().unwrap_or(0.0);
Ok(Number::from_f64(sum).map_or(Value::Null, Value::Number))
}
/// `left > right`, for numbers or strings; any other pair raises `TypeError`.
fn greater(left: &Value, right: &Value) -> Result<bool, Error> {
if let (Some(left), Some(right)) = (arithmetic(left), arithmetic(right)) {
return Ok(left.as_f64().unwrap_or(0.0) > right.as_f64().unwrap_or(0.0));
}
if let (Value::String(left), Value::String(right)) = (left, right) {
return Ok(left > right);
}
Err(raised(
"TypeError",
format!(
"'>' not supported between instances of '{}' and '{}'",
type_name(left),
type_name(right)
),
))
}
/// `_cost_value_as_float`: `float(value)`, or `None` when it is unset or not a number.
fn cost_as_float(value: Option<&Value>) -> Option<f64> {
match value? {
Value::Bool(flag) => Some(f64::from(u8::from(*flag))),
Value::Number(number) => number.as_f64(),
Value::String(text) => {
without_digit_underscores(text.trim()).and_then(|digits| python_float(&digits))
}
_ => None,
}
}
fn float(value: Option<f64>) -> Value {
value
.and_then(Number::from_f64)
.map_or(Value::Null, Value::Number)
}
/// `value.get(key)` on a dict; a value that is no dict raises `AttributeError`, as `.get` does.
fn get<'a>(value: &'a Value, key: &str) -> Result<Option<&'a Value>, Error> {
match value {
Value::Object(map) if !is_opaque(value) => Ok(map.get(key)),
other => Err(raised(
"AttributeError",
format!("'{}' object has no attribute 'get'", type_name(other)),
)),
}
}
/// `model.get(key)`, then `litellm_params.get(key)`, then `model_info.get(key)`: the first that is
/// not `None`.
fn first_set(
model: &Value,
params: &Value,
info: &Value,
key: &str,
) -> Result<Option<Value>, Error> {
for layer in [model, params, info] {
if let Some(value) = get(layer, key)?.filter(|value| !value.is_null()) {
return Ok(Some(value.clone()));
}
}
Ok(None)
}
fn object(answer: Value, kind: &str) -> Result<Map<String, Value>, Error> {
match answer {
Value::Object(map) => Ok(map),
_ => Err(malformed(kind)),
}
}
/// The group being merged: the host's validated first entry, and what the merge assigned since.
struct Merged {
current: Map<String, Value>,
assignments: Map<String, Value>,
}
impl Merged {
fn get(&self, field: &str) -> &Value {
self.current.get(field).unwrap_or(&Value::Null)
}
fn set(&mut self, field: &str, value: Value) {
self.current.insert(field.to_owned(), value.clone());
self.assignments.insert(field.to_owned(), value);
}
/// `if info.get(field) is not None and info[field] > group.field (or group.field is None)`.
fn widen(&mut self, field: &str, info: &Map<String, Value>) -> Result<(), Error> {
let Some(value) = info.get(field).filter(|value| !value.is_null()) else {
return Ok(());
};
let current = self.get(field);
if current.is_null() || greater(value, current)? {
self.set(field, value.clone());
}
Ok(())
}
/// The cost `field` widened to `float(info[field])` when that exceeds `group.field or 0.0`.
fn widen_cost(&mut self, field: &str, info: &Map<String, Value>) {
let Some(cost) = cost_as_float(info.get(field)) else {
return;
};
let current = self.get(field);
let floor = if truthy(current) {
current.as_f64().unwrap_or(0.0)
} else {
0.0
};
if current.is_null() || cost > floor {
self.set(field, float(Some(cost)));
}
}
/// `if info.get(field) is True: group.field = True`.
fn any(&mut self, field: &str, info: &Map<String, Value>) {
if info.get(field) == Some(&Value::Bool(true)) {
self.set(field, Value::Bool(true));
}
}
}
/// `tuple(dict.fromkeys((*current, *added))) or None`.
fn union(current: &Value, added: &[Value]) -> Value {
let mut kept: Vec<Value> = Vec::new();
let existing = current.as_array().map(Vec::as_slice).unwrap_or_default();
for value in existing.iter().chain(added) {
if !kept.contains(value) {
kept.push(value.clone());
}
}
if kept.is_empty() {
Value::Null
} else {
Value::Array(kept)
}
}
/// `_set_model_group_info`: the group's info merged across the deployments that serve it.
pub async fn model_group_info(
host: &dyn Host,
query: &GroupInfoQuery,
) -> Result<Option<GroupInfo>, Error> {
let mut merged: Option<Merged> = None;
let mut totals: [(&str, Option<Value>); 4] =
[("tpm", None), ("rpm", None), ("itpm", None), ("otpm", None)];
let mut clientside_auth_params = Value::Null;
let mut efforts_initialized = false;
let mut efforts_unknown = false;
let empty = Value::Object(Map::new());
for (index, model) in query.deployments.iter().enumerate() {
let named = get(model, "model_name")?;
let is_match = query.serves_team_deployments
|| named.is_some_and(|name| name.as_str() == Some(query.model_group.as_str()))
|| (named.is_some() && query.pattern_routes);
if !is_match {
continue;
}
let params = object(
lookup(host, ModelInfoLookup::GroupParams { index }).await?,
"group_params",
)?;
clientside_auth_params = params
.get("configurable_clientside_auth_params")
.cloned()
.unwrap_or(Value::Null);
let model_params = get(model, "litellm_params")?.unwrap_or(&empty);
let model_info_dict = get(model, "model_info")?.unwrap_or(&empty);
let mut deployment = [None, None, None, None];
for (slot, (metric, _)) in deployment.iter_mut().zip(totals.iter()) {
*slot = first_set(model, model_params, model_info_dict, metric)?;
}
let facts = object(
lookup(
host,
ModelInfoLookup::GroupDeployment {
index,
model: params.get("model").cloned().unwrap_or(Value::Null),
},
)
.await?,
"group_deployment",
)?;
let mapped = facts.get("mapped") == Some(&Value::Bool(true));
let litellm_model = facts.get("litellm_model").cloned().unwrap_or(Value::Null);
let llm_provider = facts.get("llm_provider").cloned().unwrap_or(Value::Null);
let model_info = match facts.get("model_info") {
Some(Value::Object(info)) if !is_opaque(&Value::Object(info.clone())) => info.clone(),
Some(Value::Null) | None => {
// No model info: what the deployment's own (database) model_info says.
let db_info = get(model, "model_info")?.unwrap_or(&empty);
let mut info = Map::new();
info.insert("key".to_owned(), Value::String(query.model_group.clone()));
for key in ["max_tokens", "max_input_tokens", "max_output_tokens"] {
info.insert(key.to_owned(), Value::Null);
}
for key in ["input_cost_per_token", "output_cost_per_token"] {
info.insert(key.to_owned(), float(cost_as_float(get(db_info, key)?)));
}
info.insert("litellm_provider".to_owned(), llm_provider.clone());
info.insert(
"mode".to_owned(),
match get(db_info, "mode")? {
Some(mode) => mode.clone(),
None => Value::String("chat".to_owned()),
},
);
info.insert(
"supported_openai_params".to_owned(),
match facts.get("supported_openai_params") {
Some(Value::Null) | None => Value::Array(Vec::new()),
Some(params) => params.clone(),
},
);
info.insert("supports_system_messages".to_owned(), Value::Null);
info
}
Some(_) => return Err(malformed("group_deployment")),
};
match merged.as_mut() {
None => {
let mut seed = Map::new();
seed.insert(
"model_group".to_owned(),
Value::String(query.user_facing_model_group_name.clone()),
);
seed.insert(
"providers".to_owned(),
Value::Array(vec![llm_provider.clone()]),
);
seed.insert("description".to_owned(), query.description.clone());
seed.extend(model_info.clone());
seed.insert(
"supports_continuation".to_owned(),
Value::Bool(
model_info.get("supports_continuation") == Some(&Value::Bool(true)),
),
);
seed.insert("supports_fast_mode".to_owned(), Value::Bool(true));
seed.insert("supported_reasoning_efforts".to_owned(), Value::Null);
seed.insert("supported_endpoints".to_owned(), Value::Null);
let current = object(
lookup(host, ModelInfoLookup::GroupSeed { info: seed }).await?,
"group_seed",
)?;
merged = Some(Merged {
current,
assignments: Map::new(),
});
}
Some(group) => {
let providers = group
.get("providers")
.as_array()
.cloned()
.unwrap_or_default();
if !providers.contains(&llm_provider) {
let mut providers = providers;
providers.push(llm_provider.clone());
group.set("providers", Value::Array(providers));
}
group.widen("max_input_tokens", &model_info)?;
group.widen("max_output_tokens", &model_info)?;
group.widen_cost("input_cost_per_token", &model_info);
group.widen_cost("output_cost_per_token", &model_info);
for field in [
"supports_parallel_function_calling",
"supports_vision",
"supports_function_calling",
"supports_web_search",
"supports_url_context",
"supports_continuation",
"supports_reasoning",
] {
group.any(field, &model_info);
}
if let Some(params) = model_info
.get("supported_openai_params")
.filter(|params| !params.is_null())
{
group.set("supported_openai_params", params.clone());
}
for (slot, metric) in deployment.iter_mut().zip(["tpm", "rpm"]) {
if slot.is_none()
&& let Some(value) = model_info.get(metric).filter(|v| !v.is_null())
{
*slot = Some(value.clone());
}
}
}
}
let Some(group) = merged.as_mut() else {
unreachable!("the group was just built");
};
let capabilities = object(
lookup(
host,
ModelInfoLookup::GroupCapabilities {
index,
model_info: Value::Object(model_info),
litellm_model,
llm_provider,
mapped,
},
)
.await?,
"group_capabilities",
)?;
let endpoints = capabilities
.get("supported_endpoints")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let endpoints = union(group.get("supported_endpoints"), &endpoints);
group.set("supported_endpoints", endpoints);
let fast_mode = truthy(group.get("supports_fast_mode"))
&& capabilities.get("supports_fast_mode") == Some(&Value::Bool(true));
group.set("supports_fast_mode", Value::Bool(fast_mode));
match capabilities
.get("reasoning_efforts")
.cloned()
.unwrap_or(Value::Null)
{
Value::Null => {
efforts_unknown = true;
group.set("supported_reasoning_efforts", Value::Null);
}
resolved if !efforts_initialized => {
efforts_initialized = true;
if !efforts_unknown {
group.set("supported_reasoning_efforts", resolved);
}
}
resolved if !efforts_unknown => {
let intersected = lookup(
host,
ModelInfoLookup::ReasoningIntersect {
current: group.get("supported_reasoning_efforts").clone(),
resolved,
},
)
.await?;
group.set("supported_reasoning_efforts", intersected);
}
_ => {}
}
for ((_, total), value) in totals.iter_mut().zip(deployment) {
if let Some(value) = value {
*total = Some(add(total.take(), &value)?);
}
}
}
let Some(mut group) = merged else {
return Ok(None);
};
for (metric, total) in totals {
if let Some(total) = total {
group.set(metric, total);
}
}
if !clientside_auth_params.is_null() {
group.set(
"configurable_clientside_auth_params",
clientside_auth_params,
);
}
Ok(Some(GroupInfo {
assignments: group.assignments,
}))
}
/// `coerce_token_limit`: a token limit as an int, `None` when it is unset or not a usable number
/// (a bool never is).
pub fn coerce_token_limit(value: Option<&Value>) -> Option<i64> {
match value? {
Value::Number(number) => number.as_i64().or_else(|| {
number
.as_f64()
.filter(|n| n.is_finite())
.map(|n| n.trunc() as i64)
}),
Value::String(text) => without_digit_underscores(text.trim())
.and_then(|digits| python_int(&digits))
.and_then(|number| number.as_i64()),
_ => None,
}
}
/// The deployment's discovered limits under its configured `model_info` values that are set.
async fn effective_model_info(
host: &dyn Host,
id: Value,
configured: &Map<String, Value>,
) -> Result<Map<String, Value>, Error> {
let mut info = discovered(host, id).await?;
info.extend(
configured
.iter()
.filter(|(_, value)| !value.is_null())
.map(|(key, value)| (key.clone(), value.clone())),
);
Ok(info)
}
/// `get_configured_token_limits`: the (max input, max output) tokens a concrete deployment
/// configures, or its refresh discovered, from its `model_info` (`exclude_none` dump) and id.
pub async fn configured_token_limits(
host: &dyn Host,
id: Value,
configured: &Map<String, Value>,
) -> Result<(Option<i64>, Option<i64>), Error> {
let info = effective_model_info(host, id, configured).await?;
Ok((
coerce_token_limit(info.get("max_input_tokens")),
coerce_token_limit(info.get("max_output_tokens")),
))
}
/// `get_configured_mode` / `get_configured_display_name`: a configured string that is not blank.
pub fn configured_text(value: Option<&Value>) -> Option<String> {
value
.and_then(Value::as_str)
.filter(|text| !text.trim().is_empty())
.map(str::to_owned)
}
/// `get_configured_service_tiers`: each deployment's `model_info.service_tiers`, unvalidated.
pub fn configured_service_tiers(model_infos: &[Map<String, Value>]) -> Vec<Value> {
model_infos
.iter()
.map(|info| info.get("service_tiers").cloned().unwrap_or(Value::Null))
.collect()
}
/// What a listed name's deployments contribute to its `/v1/models` entry.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ListingInfo {
pub cost_map_keys: Vec<String>,
pub max_input_tokens: Option<i64>,
pub max_output_tokens: Option<i64>,
}
fn named_text(value: Option<&Value>) -> Option<&str> {
value
.and_then(Value::as_str)
.filter(|text| !text.is_empty())
}
/// `deployment_pricing_models(deployment)[0]`: `model_info.base_model`, then
/// `litellm_params.base_model`, then `litellm_params.model`; malformed blocks count as unset.
fn pricing_key(deployment: &Value) -> Option<String> {
let block = |key: &str| {
deployment
.get(key)
.and_then(Value::as_object)
.filter(|_| !is_opaque(deployment.get(key).unwrap_or(&Value::Null)))
};
let params = block("litellm_params");
let info = block("model_info");
let base_model = info
.and_then(|info| info.get("base_model"))
.filter(|base| truthy(base))
.or_else(|| params.and_then(|params| params.get("base_model")));
named_text(base_model)
.or_else(|| named_text(params.and_then(|params| params.get("model"))))
.map(str::to_owned)
}
/// `get_model_listing_info` over the deployments a listed name names: the cost-map keys of their
/// models, deduplicated, and the widest token limits they configure or discovered.
pub async fn model_listing_info(
host: &dyn Host,
deployments: &[Value],
) -> Result<ListingInfo, Error> {
let mut listing = ListingInfo::default();
let empty = Value::Object(Map::new());
for deployment in deployments {
let info = match get(deployment, "model_info")? {
Some(info) if truthy(info) => info,
_ => &empty,
};
let id = get(info, "id")?.cloned().unwrap_or(Value::Null);
let configured = match info {
Value::Object(map) => map.clone(),
_ => Map::new(),
};
let effective = effective_model_info(host, id, &configured).await?;
for (field, widest) in [
("max_input_tokens", &mut listing.max_input_tokens),
("max_output_tokens", &mut listing.max_output_tokens),
] {
if let Some(limit) = coerce_token_limit(effective.get(field)) {
*widest = Some(widest.map_or(limit, |current| current.max(limit)));
}
}
}
for key in deployments.iter().filter_map(pricing_key) {
if !listing.cost_map_keys.contains(&key) {
listing.cost_map_keys.push(key);
}
}
Ok(listing)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn coerce_token_limit_reads_python_ints() {
assert_eq!(coerce_token_limit(Some(&json!(128))), Some(128));
assert_eq!(coerce_token_limit(Some(&json!(12.9))), Some(12));
assert_eq!(coerce_token_limit(Some(&json!(" 1_000 "))), Some(1000));
assert_eq!(coerce_token_limit(Some(&json!("128,000"))), None);
assert_eq!(coerce_token_limit(Some(&json!(true))), None);
assert_eq!(coerce_token_limit(None), None);
}
#[test]
fn totals_add_like_python() {
assert_eq!(add(None, &json!(5)).unwrap(), json!(5));
assert_eq!(add(Some(json!(5)), &json!(true)).unwrap(), json!(6));
assert_eq!(add(Some(json!(5)), &json!(0.5)).unwrap(), json!(5.5));
assert!(add(Some(json!(5)), &json!("1")).is_err());
}
#[test]
fn union_keeps_first_seen_order_and_empty_reads_as_none() {
assert_eq!(
union(&json!(["a", "b"]), &[json!("b"), json!("c")]),
json!(["a", "b", "c"])
);
assert_eq!(union(&Value::Null, &[]), Value::Null);
}
#[test]
fn pricing_key_prefers_the_base_model() {
let deployment = json!({
"litellm_params": {"model": "azure/x", "base_model": "params-base"},
"model_info": {"base_model": "info-base"},
});
assert_eq!(pricing_key(&deployment).as_deref(), Some("info-base"));
let deployment = json!({"litellm_params": {"model": "azure/x"}, "model_info": "junk"});
assert_eq!(pricing_key(&deployment).as_deref(), Some("azure/x"));
}
}

View file

@ -0,0 +1,960 @@
//! [`Health`]: cooldowns, allowed-fails counters and per-minute usage, kept in the shared
//! [`RouterCache`] under the Python router's keys, encodings and TTLs:
//!
//! | state | key | value | TTL |
//! |---|---|---|---|
//! | cooldown | `deployment:{id}:cooldown` | `{exception_received, status_code, timestamp, cooldown_time}` | cooldown time |
//! | allowed fails | `deployment:{id}:allowed_fails[:{suffix}]` | int counter | cooldown time |
//! | usage | `global_router:{id}:{model}:tpm:{HH-MM}` and `:rpm:` | int counter | 60s |
//! | outcomes (local only) | `{id}:successes`, `{id}:fails` | int counter | 60s |
//!
//! From `router_utils/cooldown_cache.py`, `router_utils/cooldown_handlers.py`,
//! `router_utils/router_callbacks/track_deployment_metrics.py`, `RouterCacheEnum` in
//! `types/router.py` and `Router._increment_usage` / `deployment_callback_on_failure` in
//! `router.py`.
use std::{
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use serde::{Deserialize, Serialize};
use tracing::debug;
use crate::{
cache::RouterCache,
error::{AttemptError, Error, FailureKind},
outcome::{OutcomeCallbacks, OutcomeWrite},
registry::{Deployment, Registry},
settings::Settings,
};
/// `RoutingArgs.ttl`: per-minute usage keys expire after a minute.
pub const USAGE_TTL: Duration = Duration::from_secs(60);
/// `DEFAULT_FAILURE_THRESHOLD_PERCENT`.
pub const FAILURE_THRESHOLD_PERCENT: f64 = 0.5;
/// `DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS`.
pub const FAILURE_THRESHOLD_MINIMUM_REQUESTS: u64 = 5;
/// `SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD`.
pub const SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD: u64 = 1000;
/// `exception_masker` in `CooldownCache`: the visible prefix of `exception_received`.
const EXCEPTION_VISIBLE_PREFIX: usize = 50;
pub trait Clock: Send + Sync {
fn now(&self) -> SystemTime;
}
pub struct SystemClock;
impl Clock for SystemClock {
fn now(&self) -> SystemTime {
SystemTime::now()
}
}
fn unix_seconds(time: SystemTime) -> f64 {
time.duration_since(UNIX_EPOCH)
.map(|elapsed| elapsed.as_secs_f64())
.unwrap_or_default()
}
/// `get_utc_datetime().strftime("%H-%M")`.
pub fn current_minute(time: SystemTime) -> String {
let seconds = time
.duration_since(UNIX_EPOCH)
.map(|elapsed| elapsed.as_secs())
.unwrap_or_default()
% 86_400;
format!("{:02}-{:02}", seconds / 3600, (seconds % 3600) / 60)
}
pub fn cooldown_key(id: &str) -> String {
format!("deployment:{id}:cooldown")
}
pub fn allowed_fails_key(id: &str, suffix: Option<&str>) -> String {
match suffix {
Some(suffix) => format!("deployment:{id}:allowed_fails:{suffix}"),
None => format!("deployment:{id}:allowed_fails"),
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::EnumString)]
pub enum UsageMetric {
#[strum(serialize = "tpm")]
Tpm,
#[strum(serialize = "rpm")]
Rpm,
/// Input tokens per minute, which the io-token rate limit check counts.
#[strum(serialize = "itpm")]
Itpm,
/// Output tokens per minute.
#[strum(serialize = "otpm")]
Otpm,
}
/// `RouterCacheEnum.TPM` / `RPM` / `ITPM` / `OTPM`: `model` is the deployment's
/// `litellm_params.model`.
pub fn usage_key(id: &str, model: &str, metric: UsageMetric, minute: &str) -> String {
format!("global_router:{id}:{model}:{metric}:{minute}")
}
/// `CooldownCacheValue`.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CooldownValue {
pub exception_received: String,
pub status_code: String,
pub timestamp: f64,
/// Whole seconds are written as an int, as Python's default `cooldown_time` is, so the
/// "Try again in N seconds" text reads the same.
#[serde(serialize_with = "whole_as_int")]
pub cooldown_time: f64,
}
fn whole_as_int<S: serde::Serializer>(value: &f64, serializer: S) -> Result<S::Ok, S::Error> {
if value.fract() == 0.0 && value.abs() < i64::MAX as f64 {
serializer.serialize_i64(*value as i64)
} else {
serializer.serialize_f64(*value)
}
}
/// `litellm.should_retry`.
pub fn should_retry_status(status: u16) -> bool {
matches!(status, 408 | 409 | 429) || status >= 500
}
/// `_is_cooldown_required`: a 4xx cools down only for 429, 401, 402, 404 and 408, every other
/// status does, and a connection error or an empty status never does.
pub fn is_cooldown_required(failure: &AttemptError) -> bool {
if failure.kind == FailureKind::ApiConnection || failure.status_absent {
return false;
}
match failure.cooldown_status() {
status @ 400..500 => matches!(status, 429 | 401 | 402 | 408 | 404),
_ => true,
}
}
/// `SensitiveDataMasker(visible_prefix=50, visible_suffix=0, mask_short_values=False)`.
fn mask_exception(text: &str) -> String {
let length = text.chars().count();
if length <= EXCEPTION_VISIBLE_PREFIX {
return text.to_owned();
}
let mut masked: String = text.chars().take(EXCEPTION_VISIBLE_PREFIX).collect();
masked.extend(std::iter::repeat_n('*', length - EXCEPTION_VISIBLE_PREFIX));
masked
}
/// `_EXCEPTION_POLICY_FIELDS`: a deployment policy's fields in the order Python matches them
/// (a content-policy violation is a bad request, so it comes first).
const DEPLOYMENT_POLICY_FIELDS: [(&str, &str); 9] = [
(
"ContentPolicyViolationError",
"ContentPolicyViolationErrorAllowedFails",
),
("BadRequestError", "BadRequestErrorAllowedFails"),
("AuthenticationError", "AuthenticationErrorAllowedFails"),
("Timeout", "TimeoutErrorAllowedFails"),
("RateLimitError", "RateLimitErrorAllowedFails"),
("InternalServerError", "InternalServerErrorAllowedFails"),
(
"ServiceUnavailableError",
"ServiceUnavailableErrorAllowedFails",
),
("BadGatewayError", "BadGatewayErrorAllowedFails"),
("NotFoundError", "NotFoundErrorAllowedFails"),
];
/// `_resolve_allowed_fails_from_policy` over a deployment's `model_info.allowed_fails_policy`:
/// the first field set for a class the failure is an instance of.
pub fn deployment_allowed_fails(
policy: Option<&serde_json::Map<String, serde_json::Value>>,
failure: &AttemptError,
) -> Option<f64> {
let policy = policy?;
DEPLOYMENT_POLICY_FIELDS
.iter()
.filter(|(class, _)| failure.is_a(class))
.find_map(|(_, field)| policy.get(*field).and_then(serde_json::Value::as_f64))
}
/// What the cooldown rules need to know about the deployment's group.
#[derive(Clone, Copy, Debug, Default)]
pub struct GroupShape {
/// The deployment is the only one in its model group.
pub single_deployment: bool,
/// The requested group has a fallback configured.
pub has_fallback_group: bool,
}
/// `_core_team_model_has_alternatives`: another routable deployment serves the deployment's team
/// public name.
pub(crate) fn team_model_has_alternatives(registry: &Registry, deployment: &Deployment) -> bool {
let info = &deployment.model_info;
let (Some(team_id), Some(public_name)) = (&info.team_id, &info.team_public_model_name) else {
return false;
};
registry
.team_named(team_id, public_name)
.is_some_and(|siblings| {
siblings
.iter()
.filter(|sibling| !sibling.model_info.blocked)
.count()
> 1
})
}
/// `_core_routing_group_has_alternatives`: `model_group` names a callable routing group whose
/// members span more than one deployment.
pub(crate) fn routing_group_has_alternatives(registry: &Registry, model_group: &str) -> bool {
let resolved = registry.alias_of(model_group).unwrap_or(model_group);
registry.routing_group_named(resolved).is_some_and(|group| {
group
.models
.iter()
.map(|member| registry.named(member).map_or(0, |members| members.len()))
.sum::<usize>()
> 1
})
}
/// An outcome the router did not see itself (`Health.report`).
#[derive(Clone, Debug)]
pub enum Outcome {
Success { total_tokens: u64 },
Failure(AttemptError),
}
#[derive(Clone)]
pub struct Health {
cache: Arc<dyn RouterCache>,
clock: Arc<dyn Clock>,
}
impl Health {
pub fn new(cache: Arc<dyn RouterCache>, clock: Arc<dyn Clock>) -> Self {
Self { cache, clock }
}
pub fn cache(&self) -> &Arc<dyn RouterCache> {
&self.cache
}
pub(crate) fn now_time(&self) -> SystemTime {
self.clock.now()
}
fn now(&self) -> f64 {
unix_seconds(self.clock.now())
}
/// The ids among `ids` with an unexpired cooldown (`async_get_active_cooldowns`).
pub async fn active_cooldowns(
&self,
ids: &[String],
) -> Result<Vec<(String, CooldownValue)>, Error> {
let keys: Vec<String> = ids.iter().map(|id| cooldown_key(id)).collect();
let values = self.cache.async_batch_get_cache(&keys).await?;
let now = self.now();
Ok(ids
.iter()
.zip(values)
.filter_map(|(id, value)| {
let value: CooldownValue = serde_json::from_value(value?).ok()?;
(value.timestamp + value.cooldown_time - now > 0.0).then(|| (id.clone(), value))
})
.collect())
}
/// `add_deployment_to_cooldown`.
pub async fn add_cooldown(
&self,
id: &str,
failure: &AttemptError,
cooldown: Duration,
) -> Result<(), Error> {
let value = CooldownValue {
exception_received: mask_exception(failure.detail.as_deref().unwrap_or_default()),
status_code: failure.cooldown_status().to_string(),
timestamp: self.now(),
cooldown_time: cooldown.as_secs_f64(),
};
let value = serde_json::to_value(value).unwrap_or_default();
self.cache
.async_set_cache(&cooldown_key(id), value, Some(cooldown))
.await?;
Ok(())
}
async fn local_count(&self, key: &str) -> Result<u64, Error> {
Ok(self
.cache
.async_get_local(key)
.await?
.and_then(|value| value.as_f64())
.map_or(0, |count| count as u64))
}
/// The per-minute (successes, fails) this worker saw (`{id}:successes`, `{id}:fails`, kept
/// `local_only`).
pub async fn outcomes_this_minute(&self, id: &str) -> Result<(u64, u64), Error> {
Ok((
self.local_count(&format!("{id}:successes")).await?,
self.local_count(&format!("{id}:fails")).await?,
))
}
/// `increment_deployment_successes_for_current_minute`: returns the key it counted under.
pub async fn count_success(&self, id: &str) -> Result<String, Error> {
let key = format!("{id}:successes");
self.cache
.async_increment_local(&key, 1.0, Some(USAGE_TTL))
.await?;
Ok(key)
}
/// `increment_deployment_failures_for_current_minute`.
pub async fn count_failure(&self, id: &str) -> Result<(), Error> {
self.cache
.async_increment_local(&format!("{id}:fails"), 1.0, Some(USAGE_TTL))
.await?;
Ok(())
}
/// Cooldown duration precedence from `deployment_callback_on_failure`: the deployment's
/// `cooldown_time`, then the provider's retry-after, then the router default.
pub fn cooldown_duration(
settings: &Settings,
deployment: &Deployment,
failure: &AttemptError,
) -> Duration {
Self::cooldown_for_call(settings, deployment, None, failure)
}
/// `cooldown_duration` for a call whose own params carry `cooldown_time` (a request's
/// value wins over the deployment's `litellm_params`, not over its `model_info`).
pub fn cooldown_for_call(
settings: &Settings,
deployment: &Deployment,
call_cooldown: Option<f64>,
failure: &AttemptError,
) -> Duration {
let configured = deployment
.model_info
.cooldown_time
.or(call_cooldown)
.or_else(|| {
deployment
.raw_params
.get("cooldown_time")
.and_then(|v| v.as_f64())
})
.filter(|seconds| *seconds >= 0.0)
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok());
configured
.or(failure.retry_after)
.unwrap_or(settings.cooldown_time)
}
/// The deployment's own `cooldown_time`: `model_info` first, then `litellm_params`
/// (`_first_present(model_info, litellm_params, key="cooldown_time")`). It is both the
/// cooldown and the TTL of the deployment's `allowed_fails:generic` counter.
fn configured_cooldown(deployment: &Deployment) -> Option<Duration> {
deployment
.model_info
.cooldown_time
.or_else(|| {
deployment
.raw_params
.get("cooldown_time")
.and_then(|v| v.as_f64())
})
.filter(|seconds| *seconds >= 0.0)
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
}
/// Record a failed attempt and cool the deployment down when the rules say so. Returns the
/// cooldown applied, if any.
pub async fn record_failure(
&self,
settings: &Settings,
deployment: &Deployment,
shape: GroupShape,
call_cooldown: Option<f64>,
failure: &AttemptError,
) -> Result<Option<Duration>, Error> {
if failure.outside_deployment {
return Ok(None);
}
self.count_failure(&deployment.id).await?;
let cooldown = Self::cooldown_for_call(settings, deployment, call_cooldown, failure);
self.cool_down_after(settings, deployment, shape, failure, cooldown)
.await
}
/// `set_cooldown_deployments` for a failure already counted: cool the deployment down for
/// `cooldown` when the rules say so. Returns the cooldown applied, if any.
pub async fn cool_down_after(
&self,
settings: &Settings,
deployment: &Deployment,
shape: GroupShape,
failure: &AttemptError,
cooldown: Duration,
) -> Result<Option<Duration>, Error> {
if !Self::should_run_cooldown_logic(settings, deployment, failure, cooldown) {
return Ok(None);
}
if !self
.should_cooldown(settings, deployment, shape, failure)
.await?
{
return Ok(None);
}
debug!(deployment_id = %deployment.id, cooldown_seconds = cooldown.as_secs_f64(), "cooling deployment down");
self.add_cooldown(&deployment.id, failure, cooldown).await?;
Ok(Some(cooldown))
}
/// `_should_run_cooldown_logic`.
fn should_run_cooldown_logic(
settings: &Settings,
deployment: &Deployment,
failure: &AttemptError,
cooldown: Duration,
) -> bool {
!matches!(
failure.kind,
FailureKind::Declined
| FailureKind::CallerCredentialAuthentication
| FailureKind::CallerCredentialRateLimit
) && !failure.caller_scoped
&& !cooldown.is_zero()
&& !settings.disable_cooldowns
&& (is_cooldown_required(failure)
|| deployment_allowed_fails(
deployment.model_info.allowed_fails_policy.as_ref(),
failure,
)
.is_some())
&& !settings
.provider_default_deployment_ids
.contains(&deployment.id)
}
/// `_should_cooldown_deployment`.
async fn should_cooldown(
&self,
settings: &Settings,
deployment: &Deployment,
shape: GroupShape,
failure: &AttemptError,
) -> Result<bool, Error> {
let status = failure.cooldown_status();
let policy_fails = settings.allowed_fails_for(failure);
let lone = shape.single_deployment && !shape.has_fallback_group;
let deployment_policy = deployment.model_info.allowed_fails_policy.as_ref();
let deployment_fails = deployment_allowed_fails(deployment_policy, failure);
if shape.single_deployment
&& status == 402
&& deployment_fails.is_none()
&& policy_fails.is_none()
{
return Ok(false);
}
let deployment_allowed = deployment.model_info.allowed_fails;
if deployment_policy.is_some() || deployment_allowed.is_some() {
// `_should_cooldown_based_on_deployment_policy`: a deployment-wide `allowed_fails`
// keeps the lone-deployment safety net; a policy entry naming the error does not.
if deployment_fails.is_none()
&& lone
&& (deployment_allowed.is_some() || policy_fails.is_none())
{
return Ok(false);
}
let (allowed, suffix) = match deployment_fails {
Some(allowed) => (
Some(allowed),
failure.class_names().first().map(|name| (*name).to_owned()),
),
None => (
deployment_allowed.map(f64::from),
deployment_allowed.map(|_| "generic".to_owned()),
),
};
let allowed = allowed
.or_else(|| policy_fails.map(f64::from))
.unwrap_or_else(|| f64::from(settings.effective_allowed_fails()));
let ttl =
Self::configured_cooldown(deployment).unwrap_or_else(|| settings.fails_window());
return self
.exceeds_allowed_fails(&deployment.id, suffix.as_deref(), allowed, ttl)
.await;
}
if settings.allowed_fails_policy.is_none() && !settings.allowed_fails_set() {
let (successes, fails) = self.outcomes_this_minute(&deployment.id).await?;
let total = successes + fails;
let percent_fails = if total > 0 {
fails as f64 / total as f64
} else {
0.0
};
return Ok((status == 429 && !shape.single_deployment)
|| (percent_fails == 1.0 && total >= SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD)
|| (percent_fails > FAILURE_THRESHOLD_PERCENT
&& total >= FAILURE_THRESHOLD_MINIMUM_REQUESTS
&& !shape.single_deployment)
|| !should_retry_status(status));
}
if lone && policy_fails.is_none() {
return Ok(false);
}
let allowed = policy_fails.unwrap_or_else(|| settings.effective_allowed_fails());
self.exceeds_allowed_fails(
&deployment.id,
None,
f64::from(allowed),
settings.fails_window(),
)
.await
}
/// `should_cooldown_based_on_allowed_fails_policy`: count this failure, then compare.
async fn exceeds_allowed_fails(
&self,
id: &str,
suffix: Option<&str>,
allowed: f64,
ttl: Duration,
) -> Result<bool, Error> {
let fails = self
.cache
.async_increment(&allowed_fails_key(id, suffix), 1.0, Some(ttl))
.await?;
Ok(fails > allowed)
}
/// Count a success, and the tokens and request against the deployment's per-minute usage
/// when `tracks` says it has limits to track. The count is the success callback's; the usage
/// is `make_call`'s, which runs whatever the callbacks do.
pub async fn record_success(
&self,
callbacks: OutcomeCallbacks,
deployment: &Deployment,
total_tokens: u64,
tracks: bool,
) -> Result<(), Error> {
callbacks
.record(OutcomeWrite::Success, self.count_success(&deployment.id))
.await?;
if !tracks {
return Ok(());
}
callbacks
.record(
OutcomeWrite::AttemptUsage,
self.add_usage(deployment, None, total_tokens as f64, 1.0),
)
.await
}
/// `_increment_deployment_usage`, for a deployment that tracks usage: tokens and requests
/// under `model` (metadata's `deployment`, which a wildcard route fills; else the
/// deployment's own), unless there is nothing to add.
pub async fn add_usage(
&self,
deployment: &Deployment,
model: Option<&str>,
total_tokens: f64,
requests: f64,
) -> Result<(), Error> {
if total_tokens <= 0.0 && requests <= 0.0 {
return Ok(());
}
let minute = current_minute(self.clock.now());
let model = model.unwrap_or(&deployment.params.model);
self.cache
.async_increment_usage(
&[
(
usage_key(&deployment.id, model, UsageMetric::Tpm, &minute),
total_tokens,
),
(
usage_key(&deployment.id, model, UsageMetric::Rpm, &minute),
requests,
),
],
Some(USAGE_TTL),
)
.await?;
Ok(())
}
/// `async_deployment_callback_on_failure`: a failed request counts against the deployment's
/// per-minute requests under `model` (see [`Self::add_usage`]), whether or not it tracks usage.
pub async fn count_failed_request(
&self,
deployment: &Deployment,
model: Option<&str>,
) -> Result<(), Error> {
let minute = current_minute(self.clock.now());
self.cache
.async_increment(
&usage_key(
&deployment.id,
model.unwrap_or(&deployment.params.model),
UsageMetric::Rpm,
&minute,
),
1.0,
Some(USAGE_TTL),
)
.await?;
Ok(())
}
/// This minute's (tpm, rpm) for a deployment.
pub async fn usage(&self, deployment: &Deployment) -> Result<(u64, u64), Error> {
let minute = current_minute(self.clock.now());
let model = &deployment.params.model;
let keys = [UsageMetric::Tpm, UsageMetric::Rpm]
.map(|metric| usage_key(&deployment.id, model, metric, &minute));
let values = self.cache.async_batch_get_cache(&keys).await?;
let count = |index: usize| {
values
.get(index)
.cloned()
.flatten()
.and_then(|value| value.as_f64())
.map_or(0, |count| count as u64)
};
Ok((count(0), count(1)))
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use rstest::rstest;
use super::*;
use crate::{
cache::InMemoryRouterCache,
config::{RouterConfig, RouterDynamicConfig},
registry::{DeploymentSpec, RoutingGroupSpec},
};
#[rstest]
#[case::rate_limit(FailureKind::RateLimit, Some(429), true)]
#[case::auth(FailureKind::Authentication, Some(401), true)]
#[case::payment(FailureKind::Other, Some(402), true)]
#[case::not_found(FailureKind::NotFound, Some(404), true)]
#[case::timeout(FailureKind::Timeout, Some(408), true)]
#[case::bad_request(FailureKind::BadRequest, Some(400), false)]
#[case::forbidden(FailureKind::PermissionDenied, Some(403), false)]
#[case::server(FailureKind::InternalServer, Some(500), true)]
#[case::unreadable_status(FailureKind::Other, None, true)]
#[case::connection(FailureKind::ApiConnection, Some(500), false)]
fn cooldown_required_follows_python(
#[case] kind: FailureKind,
#[case] status: Option<u16>,
#[case] expected: bool,
) {
assert_eq!(
is_cooldown_required(&AttemptError::new(kind, status)),
expected
);
}
/// An exception with no `status_code` is Python's status `""`, which no rule cools down.
#[rstest]
fn an_empty_status_requires_no_cooldown() {
let mut failure = AttemptError::new(FailureKind::Other, None);
failure.status_absent = true;
assert!(!is_cooldown_required(&failure));
}
#[rstest]
#[case::midnight(0, "00-00")]
#[case::afternoon(13 * 3600 + 7 * 60 + 59, "13-07")]
fn current_minute_is_utc_hours_and_minutes(#[case] seconds: u64, #[case] expected: &str) {
let time = UNIX_EPOCH + Duration::from_secs(86_400 * 3 + seconds);
assert_eq!(current_minute(time), expected);
}
#[rstest]
#[case::content_policy_first(&["ContentPolicyViolationError", "BadRequestError"], Some(1.0))]
#[case::bad_request(&["BadRequestError"], Some(2.0))]
#[case::unnamed(&["Timeout"], None)]
fn a_deployment_policy_matches_in_python_order(
#[case] classes: &[&str],
#[case] expected: Option<f64>,
) {
let policy = serde_json::json!({
"BadRequestErrorAllowedFails": 2,
"ContentPolicyViolationErrorAllowedFails": 1,
});
let mut failure = AttemptError::new(FailureKind::BadRequest, Some(400));
failure.classes = classes.iter().map(|class| (*class).into()).collect();
assert_eq!(
deployment_allowed_fails(policy.as_object(), &failure),
expected
);
}
fn deployment(id: &str, info: serde_json::Value) -> DeploymentSpec {
let mut model_info = serde_json::json!({"id": id});
model_info
.as_object_mut()
.unwrap()
.extend(info.as_object().unwrap().clone());
serde_json::from_value(serde_json::json!({
"model_name": id,
"litellm_params": {"model": format!("openai/{id}")},
"model_info": model_info,
}))
.unwrap()
}
fn health() -> Health {
Health::new(
Arc::new(InMemoryRouterCache::default()),
Arc::new(SystemClock),
)
}
fn settings(dynamic: RouterDynamicConfig) -> Settings {
Settings::new(&RouterConfig::default(), Arc::new(dynamic))
}
#[rstest]
#[case::held(false, false, 1, true)]
#[case::caller_scoped(false, true, 1, false)]
#[case::outside_the_deployment(true, false, 0, false)]
#[tokio::test]
async fn the_failure_exits_decide_what_is_counted_and_cooled(
#[case] outside_deployment: bool,
#[case] caller_scoped: bool,
#[case] counted: u64,
#[case] cooled: bool,
) {
let health = health();
let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap();
let mut failure = AttemptError::new(FailureKind::RateLimit, Some(429));
failure.outside_deployment = outside_deployment;
failure.caller_scoped = caller_scoped;
let cooldown = health
.record_failure(
&settings(RouterDynamicConfig::default()),
&deployment,
GroupShape::default(),
None,
&failure,
)
.await
.unwrap();
assert_eq!(cooldown.is_some(), cooled);
assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, counted);
}
#[rstest]
#[case::authentication(
FailureKind::CallerCredentialAuthentication,
401,
"AuthenticationError"
)]
#[case::rate_limit(FailureKind::CallerCredentialRateLimit, 429, "RateLimitError")]
#[tokio::test]
async fn a_caller_credential_failure_counts_without_cooling(
#[case] kind: FailureKind,
#[case] status: u16,
#[case] base_class: &str,
) {
let health = health();
let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap();
let failure = AttemptError::new(kind, Some(status));
let cooldown = health
.record_failure(
&settings(RouterDynamicConfig::default()),
&deployment,
GroupShape::default(),
None,
&failure,
)
.await
.unwrap();
assert_eq!(cooldown, None);
assert!(failure.is_a(base_class));
assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, 1);
}
/// `litellm.allowed_fails` written after the build: the copied 3 starts counting instead of
/// the failure-ratio rule, which cools a 429 down at once.
#[rstest]
#[case::still_the_process_default(true, true)]
#[case::the_process_default_moved(false, false)]
#[tokio::test]
async fn the_copied_allowed_fails_counts_once_it_differs_from_the_process_default(
#[case] allowed_fails_is_default: bool,
#[case] cooled: bool,
) {
let settings = settings(RouterDynamicConfig {
allowed_fails: Some(3),
allowed_fails_is_default,
..RouterDynamicConfig::default()
});
let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap();
let failure = AttemptError::new(FailureKind::RateLimit, Some(429));
let cooldown = health()
.record_failure(
&settings,
&deployment,
GroupShape::default(),
None,
&failure,
)
.await
.unwrap();
assert_eq!(cooldown.is_some(), cooled);
}
fn team(id: &str, team_id: &str, blocked: bool) -> DeploymentSpec {
deployment(
id,
serde_json::json!({"team_id": team_id, "team_public_model_name": "tp", "blocked": blocked}),
)
}
#[rstest]
#[case::a_sibling(team("b", "t1", false), true)]
#[case::a_blocked_sibling(team("b", "t1", true), false)]
#[case::another_teams_deployment(team("b", "t2", false), false)]
#[case::no_team(deployment("b", serde_json::json!({})), false)]
fn a_team_deployment_has_alternatives_under_its_public_name(
#[case] other: DeploymentSpec,
#[case] expected: bool,
) {
let registry = Registry::build(
vec![team("a", "t1", false), other],
BTreeMap::new(),
Vec::new(),
)
.unwrap();
let first = registry.deployment("a").unwrap();
assert_eq!(team_model_has_alternatives(&registry, &first), expected);
}
#[rstest]
#[case::spans_two_deployments(&["a", "b"], "rg", true)]
#[case::one_deployment(&["a"], "rg", false)]
#[case::through_its_alias(&["a", "b"], "alias", true)]
#[case::a_member_called_directly(&["a", "b"], "a", false)]
fn a_routing_group_has_alternatives_past_one_deployment(
#[case] members: &[&str],
#[case] requested: &str,
#[case] expected: bool,
) {
let group: RoutingGroupSpec = serde_json::from_value(serde_json::json!({
"group_name": "rg", "models": members, "routing_strategy": "simple-shuffle",
}))
.unwrap();
let registry = Registry::build(
vec![
deployment("a", serde_json::json!({})),
deployment("b", serde_json::json!({})),
],
BTreeMap::from([("alias".to_owned(), "rg".to_owned())]),
vec![group],
)
.unwrap();
assert_eq!(
routing_group_has_alternatives(&registry, requested),
expected
);
}
/// An empty status is counted and cools nothing down, unless the deployment's own policy
/// names the error.
#[rstest]
#[case::no_policy(serde_json::json!({}), false)]
#[case::policy_names_the_error(
serde_json::json!({"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}),
true
)]
#[tokio::test]
async fn an_empty_status_on_the_attempt_path_needs_the_deployments_policy_to_cool(
#[case] info: serde_json::Value,
#[case] cooled: bool,
) {
let health = health();
let deployment = Deployment::from_spec(deployment("a", info)).unwrap();
let mut failure = AttemptError::new(FailureKind::BadRequest, None);
failure.status_absent = true;
let cooldown = health
.record_failure(
&settings(RouterDynamicConfig::default()),
&deployment,
GroupShape::default(),
None,
&failure,
)
.await
.unwrap();
assert_eq!(cooldown.is_some(), cooled);
assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, 1);
}
#[rstest]
#[case::the_deployments_own_model(None, "openai/a")]
#[case::a_filled_wildcard(Some("openai/gpt-4o"), "openai/gpt-4o")]
#[tokio::test]
async fn usage_is_keyed_by_the_model_served(#[case] model: Option<&str>, #[case] keyed: &str) {
let cache = InMemoryRouterCache::default();
let clock = Arc::new(SystemClock);
let health = Health::new(Arc::new(cache.clone()), clock.clone());
let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap();
health
.add_usage(&deployment, model, 15.0, 1.0)
.await
.unwrap();
let minute = current_minute(clock.now());
assert_eq!(
(
cache.get(&usage_key("a", keyed, UsageMetric::Tpm, &minute)),
cache.get(&usage_key("a", keyed, UsageMetric::Rpm, &minute)),
),
(Some(serde_json::json!(15)), Some(serde_json::json!(1))),
);
}
#[rstest]
fn keys_match_python() {
assert_eq!(
(
cooldown_key("abc"),
allowed_fails_key("abc", Some("generic")),
usage_key("abc", "openai/gpt-4o", UsageMetric::Tpm, "09-05"),
),
(
"deployment:abc:cooldown".to_owned(),
"deployment:abc:allowed_fails:generic".to_owned(),
"global_router:abc:openai/gpt-4o:tpm:09-05".to_owned(),
)
);
}
}

View file

@ -0,0 +1,126 @@
//! Health-check-driven routing: the background health checks' results, kept in the shared
//! [`RouterCache`](crate::cache::RouterCache) under one key as Python's `DeploymentHealthCache`
//! writes them, and the filter `_async_filter_health_check_unhealthy_deployments` /
//! `_filter_health_check_unhealthy_deployments` run before the cooldown filter.
use std::{collections::BTreeSet, sync::Arc, time::UNIX_EPOCH};
use serde_json::Value;
use crate::{health::Health, registry::Deployment, settings::Settings};
/// `DeploymentHealthCache.CACHE_KEY`.
pub const HEALTH_STATE_KEY: &str = "litellm:health_check:deployment_health_state";
/// `_extract_unhealthy_ids`: the ids marked unhealthy whose result is younger than `staleness`.
pub fn unhealthy_ids(raw: Option<&Value>, now: f64, staleness: f64) -> BTreeSet<String> {
let Some(Value::Object(states)) = raw else {
return BTreeSet::new();
};
states
.iter()
.filter(|(_, state)| {
let Value::Object(state) = state else {
return false;
};
// `not state.get("is_healthy", True)`: any falsy value is unhealthy.
let unhealthy = match state.get("is_healthy") {
None => false,
Some(Value::Null) => true,
Some(Value::Bool(value)) => !value,
Some(Value::Number(value)) => value.as_f64() == Some(0.0),
Some(Value::String(value)) => value.is_empty(),
Some(Value::Array(value)) => value.is_empty(),
Some(Value::Object(value)) => value.is_empty(),
};
let timestamp = state
.get("timestamp")
.and_then(Value::as_f64)
.unwrap_or(0.0);
unhealthy && now - timestamp < staleness
})
.map(|(id, _)| id.clone())
.collect()
}
/// What the filter does when it would leave nothing.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum WhenAllUnhealthy {
/// The safety net: every candidate is kept.
KeepAll,
/// A health check probe sees no deployment.
Nothing,
}
impl Health {
/// The deployments the background health checks have not marked unhealthy, as the
/// router's health-check routing settings say.
pub async fn health_check_filter(
&self,
settings: &Settings,
deployments: Vec<Arc<Deployment>>,
when_all_unhealthy: WhenAllUnhealthy,
) -> Vec<Arc<Deployment>> {
if !settings.enable_health_check_routing {
return deployments;
}
let scoped = settings.background_health_check_model_groups.as_deref();
// With an allowed_fails_policy cooldowns are the only exclusion, unless the filter is
// scoped to listed groups.
if settings.allowed_fails_policy.is_some() && scoped.is_none() {
return deployments;
}
// A read that fails counts as no health state.
let raw = self
.cache()
.async_get_cache(HEALTH_STATE_KEY)
.await
.ok()
.flatten();
let now = self
.now_time()
.duration_since(UNIX_EPOCH)
.map(|elapsed| elapsed.as_secs_f64())
.unwrap_or_default();
let unhealthy = unhealthy_ids(
raw.as_ref(),
now,
settings.health_check_staleness.as_secs_f64(),
);
if unhealthy.is_empty() {
return deployments;
}
let filtered: Vec<Arc<Deployment>> = deployments
.iter()
.filter(|d| {
!unhealthy.contains(&d.id)
|| scoped.is_some_and(|groups| !groups.contains(&d.model_name))
})
.cloned()
.collect();
match (filtered.is_empty(), when_all_unhealthy) {
(false, _) => filtered,
(true, WhenAllUnhealthy::KeepAll) => deployments,
(true, WhenAllUnhealthy::Nothing) => Vec::new(),
}
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::fresh_unhealthy(json!({"a": {"is_healthy": false, "timestamp": 95.0}}), vec!["a"])]
#[case::stale(json!({"a": {"is_healthy": false, "timestamp": 50.0}}), vec![])]
#[case::healthy(json!({"a": {"is_healthy": true, "timestamp": 95.0}}), vec![])]
#[case::no_flag(json!({"a": {"timestamp": 95.0}}), vec![])]
#[case::not_a_dict(json!({"a": 1}), vec![])]
fn unhealthy_ids_follow_python(#[case] raw: Value, #[case] expected: Vec<&str>) {
let ids = unhealthy_ids(Some(&raw), 100.0, 10.0);
assert_eq!(ids.iter().map(String::as_str).collect::<Vec<_>>(), expected);
}
}

View file

@ -0,0 +1,726 @@
//! [`Host`]: the router core's only boundary. Everything the core needs from the embedding
//! process goes through it, and every attempt runs through [`Host::invoke`]: the core never runs
//! an operation itself. `PythonHost` serves it in a Python process (attempts call `litellm.<op>`,
//! and the deployment build asks it for secrets, validation, providers and side effects).
use std::{collections::BTreeMap, sync::Arc, time::Duration};
use async_trait::async_trait;
use serde::Deserialize;
use serde_json::{Map, Value};
use crate::{
cache::RouterCache,
config::{Op, RetryPolicy, RoutingStrategy},
error::{AttemptError, Error, FailureKind, Payload},
fallback::Fallbacks,
health::Health,
mock::MockFailure,
outcome::OutcomeCallbacks,
precall::PreCallFacts,
prerouting::{PreRoute, PreRouted},
registry::{Deployment, Registry},
settings::Settings,
strategy::StrategyScope,
};
/// One call into the router.
#[derive(Clone, Debug, Default)]
pub struct RouterRequest {
/// The requested name: a model group, alias, routing group, deployment id or wildcard match.
pub model: String,
/// The request params, for the plugins and host steps that read them.
pub body: Value,
/// `user_api_key_team_id`, from the first metadata bucket that holds one
/// (`get_request_team_id`).
pub team_id: Option<String>,
/// The `metadata` bucket's `user_api_key_team_id` alone: the team a failed level lists the
/// group's deployments for (`async_function_with_fallbacks_common_utils`,
/// `_stream_order_levels`). `None` lists every team's.
pub listing_team_id: Option<String>,
/// The caller's `user_role` is `proxy_admin`.
pub proxy_admin: bool,
/// The caller's `models` plus `team_models` (`user_api_key_auth`), which narrow a caller
/// allowed only through access groups to those groups' deployments.
pub allowed_models: Vec<String>,
/// For a proxy admin, the model the request names (`model`, else the metadata
/// `model_group`), which keeps one team's deployments listed under it (`filter_team_based_models`).
pub admin_model: Option<String>,
pub specific_deployment: bool,
/// A per-request `num_retries`, which wins over the router's.
pub num_retries: Option<u32>,
/// A per-request `num_retries` that is no count, which replaces the router's as well.
pub uncounted_retries: Option<UncountedRetries>,
pub disable_fallbacks: bool,
/// `_router_weights` as the request passed it: model group to deployment id to weight,
/// validated when simple-shuffle picks ([`crate::weights::validate`]).
pub router_weights: Option<Value>,
/// The deployment a call on an existing resource (a video id) names: picked whenever the
/// filters keep it (`_narrow_to_video_deployment`).
pub pinned_deployment_id: Option<String>,
/// A per-request `routing_strategy` override, as Python normalizes it.
pub routing_strategy: Option<String>,
/// Whether tag routing may apply to this request (router, request or model setting), so the
/// host's tag stage runs.
pub filter_tags: bool,
/// A pass-through pick: only `use_in_pass_through` deployments may serve it.
pub pass_through: Option<crate::select::PassThrough>,
/// `_target_order`: the order level to pick from instead of the lowest.
pub target_order: Option<i64>,
/// `_excluded_deployment_ids`: never picked, even when that leaves nothing.
pub excluded_ids: Vec<String>,
/// `_retry_skipped_deployment_ids`: skipped unless that leaves nothing.
pub retry_skipped_ids: Vec<String>,
/// The endpoint `endpoint_selection` names around the call, which deployments that declare
/// `enforce_supported_endpoints` may refuse ([`crate::eligibility`]).
pub endpoint: Option<String>,
/// A web search request on an async pick: `filter_web_search_deployments` drops the
/// deployments that declare no support. The sync pick never runs that filter.
pub web_search: bool,
/// The handler name of a call only some providers serve (`_ROUTE_CAPABILITY_RESOLVERS`):
/// the deployments whose provider cannot serve it are dropped unless that leaves nothing.
pub route_handler: Option<String>,
/// A Decisions call carrying audio or videos: the deployments that refuse them are skipped
/// unless that leaves nothing (`_without_decisions_deployments_refusing_media`).
pub decisions_media: bool,
/// `f"{provider}/{model}"` for a lookup that cannot ask the host
/// (`_core_get_routable_deployments`); a pick asks [`Host::provider_request`].
pub wildcard_request: Option<String>,
/// The request carries an image (`_request_carries_image`), so deployments that declare no
/// vision are skipped while any other remains.
pub has_image: bool,
/// An async pick: the host's pre-routing step runs before each pick resolves its name.
pub pre_routes: bool,
/// Which of the router's outcome callbacks run for this request.
pub outcome_callbacks: OutcomeCallbacks,
/// The requested name is no model group (a search tool name): nothing is selected, every
/// attempt is handed to the host as is, and no outcome is recorded for a deployment
/// (`_asearch_with_fallbacks`, whose helper picks among the router's search tools).
pub unrouted: bool,
/// `passthrough_on_no_deployment`: a pick that fails hands the attempt to the host under the
/// requested name, inside the level's retries and fallbacks, with no outcome recorded for a
/// deployment (`_ageneric_api_call_with_fallbacks_helper`).
pub passthrough_on_no_deployment: bool,
/// The endpoint's attempt stamps no `failed_deployment_id` on its error (file and batch
/// creation, whose result belongs to the deployment that served it), so weighted failover
/// never runs for it.
pub skip_weighted_failover: bool,
/// The request creates or names a resource only one provider holds (a file upload, a batch
/// on an uploaded file): it falls back within its own model group only
/// (`run_async_fallback`'s `same_model_group_only`).
pub same_group_fallbacks_only: bool,
/// Fallback lists the request passed (`Some(None)` when it passed `None`), which replace
/// the router's.
pub fallbacks: Option<Option<Fallbacks>>,
pub context_window_fallbacks: Option<Option<Fallbacks>>,
pub content_policy_fallbacks: Option<Option<Fallbacks>>,
/// A per-request `max_fallbacks`, which replaces the router's.
pub max_fallbacks: Option<u32>,
/// A per-request `cooldown_time`: the call's own params carry it, so it cools a failed
/// deployment down for that long unless its `model_info` sets one.
pub cooldown_time: Option<f64>,
/// A `mock_testing_*` flag the request set.
pub mock_failure: Option<MockFailure>,
/// A per-request `model_group_retry_policy`, which replaces the router's.
pub model_group_retry_policy: Option<BTreeMap<String, RetryPolicy>>,
/// The model groups whose `model_group_retry_policy` lookup raises (the policy is no mapping
/// of policies, or the group's entry is not one): the first failure of a hop serving one ends
/// the hop with the host's error instead ([`RouterEvent::RetryPolicyRaised`]), no retry.
pub retry_policy_raises: Vec<String>,
/// `health_check_probe`: a health check's own pick, which sees no deployment when every
/// one is marked unhealthy (the safety net keeps them all for any other request).
pub health_check_probe: bool,
/// What the host needs to run the call itself, such as the Python kwargs.
pub payload: Option<Payload>,
}
/// A per-request `num_retries` that is no count of retries. `async_function_with_retries` only
/// reads it once a hop's first failure was not taken over by a retry policy: a call that succeeds
/// first time never does.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UncountedRetries {
/// A negative int: `num_retries > 0` is false, so the failure stands, and a stream's plain
/// budget is spent (`isinstance(num_retries, int)`).
Negative,
/// Any other value `num_retries > 0` is false for (a negative float, a NaN): the failure
/// stands, and a stream falls through to the deployment's budget, then the router's.
NotPositive,
/// `num_retries > 0` or `range(num_retries)` raises (a string, a positive float): the hop ends
/// with that error ([`RouterEvent::RetryCountRaised`]).
Raises,
}
/// One `os.environ/` reference the deployment build found ([`crate::build`]).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SecretRef {
/// The value as written, `os.environ/<name>`.
pub name: String,
/// Where the value sits in the deployment.
pub at: SecretAt,
}
/// Where a [`SecretRef`] sits in a deployment.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SecretAt {
/// A top-level `litellm_params` value: resolved through the host's secret manager.
Param(String),
/// `litellm_params.dataSources[source].parameters[key]` (`endpoint` or `key`): read from the
/// environment alone, empty when unset.
DataSource { source: usize, key: String },
}
/// Which of the router's adds a build runs.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)]
#[strum(serialize_all = "snake_case")]
pub enum BuildPath {
/// `set_model_list` (and the constructor's list): config entries, one deployment per
/// organization, ignored when invalid if the router ignores invalid deployments.
Create,
/// `add_deployment` (and upserts): one validated deployment.
Add,
}
/// What the build asks the host about one deployment.
#[derive(Clone, Debug)]
pub struct BuildAsk<'a> {
pub path: BuildPath,
/// The entry's position in the list the build was given.
pub index: usize,
/// The entry as the build holds it: `model_name`, `litellm_params` (secrets resolved, the
/// organization of this deployment set), `model_info` (with its `id`) and any other keys.
pub entry: &'a Map<String, Value>,
/// The `model_info.id` the entry declared, before the build assigned one.
pub declared_id: Option<&'a str>,
/// Another entry of the same list declares the same id.
pub duplicate_id: bool,
}
/// The host's answer about one deployment.
#[derive(Clone, Debug, PartialEq)]
pub enum Built {
Entry(Map<String, Value>),
/// Left out without an error (`model_info.supported_environments` excludes this one, or an
/// `add_deployment` id that already exists).
Skip,
/// The host failed the deployment and the router ignores invalid deployments.
Invalid,
}
/// The host's provider answer.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Provided {
Provider(Option<String>),
Invalid,
}
/// What the model-info merge ([`crate::model_info`]) asks the host while it answers one query.
/// Asked at query time, so a price reload or a registration the host made since the router was
/// built is what the next query reads.
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
#[serde(tag = "lookup", rename_all = "snake_case")]
pub enum ModelInfoLookup {
/// The process cost map's row for `model` (`get_model_info(model=...)`), the host choosing
/// the keys it tries. Answered with the row, or `None`. A lookup that raises propagates,
/// unless `tolerate` is set: then the host answers `{"failed": True}`. `exact` also asks
/// whether the row came from a fallback-generalization rule (`{"row": ..., "generalized":
/// bool}`).
ModelInfo {
model: Value,
tolerate: bool,
exact: bool,
},
/// The raw cost-map entry under `key` (`model_cost.get(key)`): answered with
/// `{"entry": value}`, or `None` when the key is absent. `process` reads the process global
/// (`litellm.model_cost`) rather than the host's live settings, where the Python router does.
CostEntry { key: String, process: bool },
/// `get_llm_provider` for the queried deployment's `model` and validated params: answered
/// with `[model, provider]`.
Provider { model: String },
/// The deployments a wildcard pattern routes `name` to: answered with `[{id, model}]`, or
/// `None` when no pattern matches.
PatternModels { name: String },
/// The base model a bedrock model id names (`get_bedrock_base_model`).
BedrockBase { model_id: String },
/// The limits the background model-info refresh discovered for deployment `id`.
Discovered { id: Value },
/// The queried deployment's own `model_info`, read when the merge reaches it (a malformed one
/// raises there).
UserModelInfo,
/// A log record, emitted where the Python router logs it.
Log { level: LogLevel, message: String },
/// Group info: deployment `index`'s validated params (`LiteLLM_Params(**litellm_params)`, which
/// raises on bad params): answered with `{"model", "configurable_clientside_auth_params"}`.
GroupParams { index: usize },
/// Group info: what the host knows of deployment `index` serving `model`: answered with
/// `{"model_info"}` (its merged model info, `None` when that fails), `"mapped"`
/// (`deployment_is_catalog_mapped`), `"litellm_model"` and `"llm_provider"`, and, when it has
/// no model info, `"supported_openai_params"`.
GroupDeployment { index: usize, model: Value },
/// Group info: deployment `index`'s capabilities under `model_info`: answered with
/// `{"supported_endpoints"}` (empty when the team may not use it), `"supports_fast_mode"` and
/// `"reasoning_efforts"` (`resolve_supported_reasoning_efforts`, or `None`).
GroupCapabilities {
index: usize,
model_info: Value,
litellm_model: Value,
llm_provider: Value,
mapped: bool,
},
/// Group info: the group's first deployment as the host builds it (`ModelGroupInfo(**info)`,
/// whose validation may raise): answered with the fields the merge compares, as validated.
GroupSeed { info: Map<String, Value> },
/// `intersect_supported_reasoning_efforts(current, resolved)`, in the host's advertisement order.
ReasoningIntersect { current: Value, resolved: Value },
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum LogLevel {
Debug,
Error,
}
/// One attempt against one deployment.
#[derive(Clone, Debug)]
pub struct Attempt {
pub op: Op,
pub request: Arc<RouterRequest>,
pub deployment: Arc<Deployment>,
/// The name this attempt resolved: the requested name, or the fallback target.
pub name: String,
/// The model group being served, after fallbacks.
pub model_group: String,
/// The `order` level an order fallback targets in the same group.
pub target_order: Option<i64>,
/// Every deployment `name` resolved to for this caller, the picked one included.
pub candidate_ids: Vec<String>,
/// 0 for the first attempt in this group, then the retry number.
pub retry: u32,
pub fallback_depth: u32,
/// How many fallbacks deep the chain is: 0 for the requested name.
pub fallback_level: u32,
pub timeout: Option<Duration>,
/// Request params the fallback chain's dict targets merged in (`kwargs.update(target)`).
pub params: Option<Arc<Map<String, Value>>>,
/// The chain level whose request metadata this attempt carries (see [`RouterEvent::Fallback`]).
pub scope: u32,
/// The retry budget the level's loop stamps on the request before this attempt.
pub max_retries: u32,
/// How many deployments `get_model_list` lists for the level's group (`model_group_size`).
pub model_group_size: usize,
}
/// What a successful attempt returns.
#[derive(Clone, Debug)]
pub struct Output {
pub value: Payload,
/// Prompt plus completion tokens, counted against the deployment's per-minute usage.
pub total_tokens: u64,
/// A response the provider blocked, which the router may raise to fall back.
pub blocked: Option<Blocked>,
/// A stream the host reports the outcome of when it ends ([`crate::Router::report_success`],
/// [`crate::Router::report_failure`]), or a realtime session, whose success the host counts
/// only when it ended cleanly: the attempt counts its request and `total_tokens` now, and no
/// success.
pub deferred: bool,
/// Whether the attempt counts against the deployment's per-minute usage, when the host knows
/// better than the deployment's own limits ([`Deployment::tracks_usage`]): a limit only the
/// model catalog sets (`_deployment_tracks_usage`), or a realtime session, whose `None`
/// response counts nothing (`increment_deployment_usage_for_response`).
pub tracks_usage: Option<bool>,
}
/// How a provider blocked a response the call otherwise served.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)]
#[strum(serialize_all = "snake_case")]
pub enum Blocked {
/// A chat response whose first choice was content-filtered (or that has no choices).
ContentFilter,
/// An Anthropic /v1/messages safeguard refusal carrying `stop_details`.
Refusal,
}
/// What a plugin or strategy sees: never the router itself.
pub struct RoutingContext<'a> {
pub registry: &'a Registry,
pub settings: &'a Settings,
pub health: &'a Health,
pub request: &'a RouterRequest,
pub op: &'a Op,
pub model_group: &'a str,
/// The name the request resolved from (Python's `model`), which strategies key state by.
pub name: &'a str,
/// Whose selector a strategy stage reads.
pub scope: StrategyScope,
/// The deployments the name resolved to for this caller, before health filters
/// (`request_deployments`), for the tag stage.
pub pool: &'a [Arc<Deployment>],
/// `_excluded_deployment_ids` a weighted failover hop put on the request, which the callback
/// filters read (`deployment_ids_failed_in_request`).
pub excluded_ids: &'a [String],
/// `_retry_skipped_deployment_ids` the retry loop put on the request, as `excluded_ids`.
pub skipped_ids: &'a [String],
/// `_target_order` as the filters read it: the level an order fallback targets, else the one
/// the request named.
pub target_order: Option<i64>,
}
/// Which selection step the host is asked to run.
#[derive(Clone, Debug, PartialEq)]
pub enum SelectStage {
/// Pre-call checks and sub-routers registered as plugins: return the deployments that remain.
Plugins,
/// A strategy not yet ported: return the candidates with the pick first.
Strategy(RoutingStrategy),
/// Tag routing (`get_deployments_for_tag`), not yet ported: return the deployments that remain.
Tags,
/// `_deployment_declares_no_vision`: return the candidates that declare no vision support.
TextOnly,
/// `_deployment_serves_route`: return the candidates whose provider cannot serve the call
/// `handler` names (a Decisions call on a provider with no Decisions config).
RouteIncapable { handler: String },
/// `_decisions_deployment_refuses_media`: return the candidates that refuse the audio or
/// videos a Decisions call carries.
MediaRefusing,
}
/// Routing events, delivered through [`Host::observe`].
#[derive(Clone, Debug, PartialEq)]
pub enum RouterEvent {
Selected {
model_group: String,
deployment_id: String,
/// How long the pick took (`_acompletion`'s `async_get_available_deployment` service event).
pick: Duration,
},
Succeeded {
deployment_id: String,
},
AttemptFailed {
deployment_id: String,
kind: FailureKind,
status: Option<u16>,
},
CooldownSet {
deployment_id: String,
cooldown: Duration,
},
Retry {
model_group: String,
retry: u32,
delay: Duration,
},
Fallback {
from: String,
to: String,
/// The order level of `to` a same-group hop walks to.
order: Option<i64>,
depth: u32,
/// The fallback level `to` is tried at: 1 for the request's own fallbacks.
level: u32,
/// The new level's request metadata scope, copied from `parent`'s at the hop.
scope: u32,
parent: u32,
/// The `max_fallbacks` the chain reads, stamped on the hop's request.
max_fallbacks: u32,
/// The non-string `model` of the dict target this hop routes
/// ([`crate::fallback::FallbackTarget::unroutable_model`]): the hop fails on it before any
/// attempt, and the host raises what Python's level raises for that value.
unroutable: Option<Value>,
},
/// A hop of the fallback chain returned: `log_success_fallback_event` or
/// `log_failure_fallback_event` for the level at `parent`.
FallbackDone {
from: String,
to: String,
level: u32,
scope: u32,
parent: u32,
succeeded: bool,
},
/// The router raised a `ContentPolicyViolationError` for a blocked response, since a
/// fallback would serve it (`_should_raise_content_policy_error`,
/// `_should_raise_anthropic_refusal_error`).
Blocked {
level: u32,
model: String,
blocked: Blocked,
},
/// `log_retry`: the router recorded a failure in `scope`'s request metadata
/// (`previous_models`, `request_retry_count`) before retrying or falling back.
RetryLogged {
scope: u32,
level: u32,
/// The group the failed level served (`kwargs["model"]`).
model: String,
/// The retry number stamped before the failed try; `None` keeps the scope's own.
attempted_retries: Option<u32>,
cause: RetryCause,
},
/// `async_function_with_retries` opened a level's retry loop, before any pick: `scope`'s
/// request metadata tracks no retry yet, `max_retries` of them, and the group's size.
RetriesOpened {
scope: u32,
max_retries: u32,
model_group_size: usize,
},
/// The level's first failure settled its overrides: the budget `scope` then tracks.
RetryBudget {
scope: u32,
max_retries: u32,
},
/// A pick of the async selection or the sync one failed on `error`: the async one's `except`
/// logs every refusal to the request's logging object (`_async_get_available_deployment`).
PickFailed {
level: u32,
model: String,
error: PickError,
},
/// A level of the fallback chain serving `model` ended on an error a pick raised (no
/// deployment for the caller, none available, a stale routing group, ...): Python raises
/// that error for the level, whatever its earlier attempts raised.
LevelFailed {
level: u32,
model: String,
error: PickError,
/// Raised by the level's fallback handling, while the level's own error was handled
/// (so it chains to that error), rather than by the level's retry loop.
handling: bool,
},
/// The router raised a `mock_testing_*` failure for `model` instead of calling it.
MockFailed {
mock: MockFailure,
level: u32,
model: String,
num_retries: Option<u32>,
},
/// The first failure of a hop serving `model` looked up a `model_group_retry_policy` that
/// raises for it: the hop ends with the error the lookup raised (the host raises it, from the
/// failure `cause` names), without retrying, and falls back as from any plain exception.
RetryPolicyRaised {
level: u32,
model: String,
cause: RetryCause,
},
/// The first failure of a hop serving `model` read a request `num_retries` that raises
/// ([`UncountedRetries::Raises`]): the hop ends with that error (the host raises it, from the
/// failure `cause` names), and falls back as from any plain exception.
RetryCountRaised {
level: u32,
model: String,
cause: RetryCause,
},
/// The level serving `model` reached the fallback handling that appends the context window
/// and content policy hints to its error, before any hop runs
/// (`async_function_with_fallbacks_common_utils`). The order path never gets there.
FallingBack {
level: u32,
model: String,
},
/// The order fallback chain of the level at `level` served nothing, so the level raises from
/// `async_function_with_fallbacks_common_utils`' order path, with no outcome message: what the
/// chain's last hop raised when `last_hop`, else its own error (a context window error, or a
/// chain no hop of which ran). Weighted failover raising what it does not catch
/// (`_maybe_run_weighted_failover`) ends the level the same way.
OrderChainFailed {
level: u32,
last_hop: bool,
},
/// The regular fallback handling of the level at `level` served nothing, where
/// `async_function_with_fallbacks_common_utils` reaches its `except` and logs the failure.
/// `chain` when the level had targets to walk: the error it logs is what the last hop raised
/// when `last_hop`, else its own. With no `chain`, the host decides: only the host holds
/// whether a level with no match had a fallback list at all, which Python does not log.
/// The order path and weighted failover are outside that `try` and never report this.
FallbacksFailed {
level: u32,
chain: bool,
last_hop: bool,
},
/// A group's retry loop ran every retry and the last one failed: Python stamps
/// `num_retries` and `max_retries` on that error.
RetriesExhausted {
model_group: String,
attempted: u32,
max: u32,
/// The level the retry loop served.
level: u32,
/// The last failure was a pick's error, which the level then raises
/// ([`RouterEvent::LevelFailed`]), rather than an attempt's.
picked: bool,
},
}
/// A pick's error, as the events that name it carry it.
#[derive(Clone, Debug)]
pub struct PickError(pub Arc<Error>);
impl PickError {
/// The copy of `error` events carry, when a pick raised it.
pub fn of(error: &Error) -> Option<Self> {
error.pick_copy().map(|copy| Self(Arc::new(copy)))
}
}
impl PartialEq for PickError {
fn eq(&self, other: &Self) -> bool {
format!("{:?}", self.0) == format!("{:?}", other.0)
}
}
/// Which error a [`RouterEvent::RetryLogged`] records.
#[derive(Clone, Debug, PartialEq)]
pub enum RetryCause {
/// The last attempt's failure.
Attempt,
/// The error the level ended on (`run_async_fallback`'s `original_exception`).
Level,
/// A pick raised this error: the group was unknown or had none available, the stale
/// routing-group member map named a group no longer configured, ...
Pick(PickError),
/// Any other error no attempt raised.
Other,
}
impl RetryCause {
/// The cause that names `error`.
pub fn of(error: &Error) -> Self {
match error {
Error::Attempt(_) => Self::Attempt,
_ => PickError::of(error).map_or(Self::Other, Self::Pick),
}
}
}
#[async_trait]
pub trait Host: Send + Sync {
/// Run one attempt against the picked deployment. The core records its outcome (cooldowns,
/// allowed-fails counters, usage) itself; the host only runs it.
async fn invoke(&self, attempt: &Attempt) -> Result<Output, AttemptError>;
/// Run a selection stage the core does not implement.
async fn select(
&self,
stage: &SelectStage,
ctx: &RoutingContext<'_>,
candidates: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, Error>;
/// Pre-routing before a pick resolves `ask.name`: the model group to resolve instead, a
/// custom strategy's deployment to serve, or `None` to resolve the name itself. Asked only
/// for picks that pre-route (`pre_routes`).
async fn pre_route(&self, ask: &PreRoute<'_>) -> Result<Option<PreRouted>, Error> {
let _ = ask;
Ok(None)
}
/// Access and budget checks before a cross-group fallback target.
/// `Err` when a check raised: the chain ends there, outside any target's `try`.
async fn allow_fallback(
&self,
request: &RouterRequest,
target: &str,
) -> Result<bool, AttemptError>;
/// Whether `run_async_fallback` refuses a chain for one of its dict `targets`, a server-owned
/// federation field or an `oidc/` api key. The chain is walked by `level`.
async fn refuses_fallbacks(&self, _level: u32, _targets: &[Map<String, Value>]) -> bool {
false
}
async fn sleep(&self, duration: Duration);
fn observe(&self, event: &RouterEvent);
/// The shared store routing state lives in.
fn cache(&self) -> Arc<dyn RouterCache>;
/// Whether [`Self::invoke`] can run `op` on `provider`. Checked for every served op and
/// deployment when the router is built or changed, so an unsupported pair fails the config,
/// never a request. `PythonHost` runs every op.
fn supports(&self, _op: &Op, _provider: &str) -> bool {
true
}
/// Whether [`Self::select`] can run strategies and plugins.
fn can_select(&self) -> bool {
true
}
/// Run the routing plugins for `ctx.name` when `run` (the async pick's pre-routing hook), and
/// report the candidate models they left on the request, if any.
async fn routing_plugin_models(
&self,
_ctx: &RoutingContext<'_>,
_run: bool,
) -> Result<Option<Vec<String>>, Error> {
Ok(None)
}
/// The facts the `enable_pre_call_checks` filter decides by (token count, context windows,
/// param support), for `candidates` and the deployments in `ctx.pool`. A host without them
/// reports no input, so nothing is checked.
async fn pre_call_facts(
&self,
_ctx: &RoutingContext<'_>,
_candidates: &[Arc<Deployment>],
) -> Result<PreCallFacts, Error> {
Ok(PreCallFacts::default())
}
/// `f"{provider}/{model}"` for `ctx.name`, which `PatternMatchRouter.get_pattern` tries when
/// no pattern matches the name itself: only the host can infer a provider. `None` from a
/// host that infers none.
async fn provider_request(&self, _ctx: &RoutingContext<'_>) -> Result<Option<String>, Error> {
Ok(None)
}
/// `async_callback_filter_deployments`: the host's routing callbacks narrow `candidates` in
/// turn. A filter that raises fails the try with [`Error::Attempt`], which the call retries and
/// falls back from like any attempt failure. A host without filters keeps every candidate.
async fn callback_filters(
&self,
_ctx: &RoutingContext<'_>,
candidates: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, Error> {
Ok(candidates)
}
/// The values of `refs`, one per ref, in order: what the deployment build substitutes for
/// each `os.environ/` reference it found in entry `index` ([`crate::build`]). Asked once per
/// deployment, as the build reaches it. A host without secrets leaves every reference as
/// written.
async fn secrets(&self, index: usize, refs: &[SecretRef]) -> Result<Vec<Value>, Error> {
let _ = index;
Ok(refs.iter().map(|r| Value::String(r.name.clone())).collect())
}
/// Check one deployment before it is added, and answer the entry as the host normalized it,
/// a skip (the deployment is left out), or [`Built::Invalid`] (it failed and the router
/// ignores invalid deployments). A host without checks accepts the entry as given.
async fn validate(&self, ask: &BuildAsk<'_>) -> Result<Built, Error> {
Ok(Built::Entry(ask.entry.clone()))
}
/// The provider that serves one validated deployment (`get_llm_provider`), `None` for one
/// served without a provider (a prompt-management model), or [`Provided::Invalid`]. A host
/// without provider resolution answers `None`, so [`Deployment::provider`] reads the model.
async fn provider(&self, ask: &BuildAsk<'_>) -> Result<Provided, Error> {
let _ = ask;
Ok(Provided::Provider(None))
}
/// Run the host's side effects for an accepted deployment and answer the entry as the
/// registry stores it. A host without side effects stores the entry it validated.
async fn added(&self, ask: &BuildAsk<'_>, provider: Option<&str>) -> Result<Built, Error> {
let _ = provider;
Ok(Built::Entry(ask.entry.clone()))
}
/// Answer one model-info lookup ([`ModelInfoLookup`]). A host without a cost map answers
/// `None` to everything, so no deployment has model info.
async fn model_info(&self, lookup: &ModelInfoLookup) -> Result<Value, Error> {
let _ = lookup;
Ok(Value::Null)
}
}

View file

@ -0,0 +1,727 @@
//! [`HostIndexes`]: the registry indexes the Python router hands out, kept the way it keeps them.
//!
//! The Python router keeps its indexes by history rather than deriving them from the deployment
//! list: `model_names` is the set of names `set_model_list` was given (invalid entries included),
//! `deployment_names` only ever grows, assigning `model_list` reindexes nothing, a removal shifts
//! the positions every index holds (a caller-written one included), and the wildcard pattern
//! routers keep the copy of a deployment they were given. The registry keeps one of these per
//! version and applies those rules on its own writes; the host reads them back as the live
//! objects `registry().index()`, `names()` and `patterns()` return, and hands a caller's write to
//! any of them back as a changeset ([`IndexWrite`]).
//!
//! Values are JSON so whatever a caller wrote reads back as written, junk included.
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
/// One wildcard pattern router (`PatternMatchRouter`): regex to the deployments it routes to,
/// most specific pattern first.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct PatternRouter {
pub patterns: Vec<(String, Vec<Value>)>,
}
/// The indexes, as the host reads and writes them.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(default)]
pub struct HostIndexes {
/// `model_id_to_deployment_index_map`: id to position, in insertion order.
pub id_map: Map<String, Value>,
/// `model_name_to_deployment_indices`: name to its positions.
pub model_name: Map<String, Value>,
/// `team_model_to_deployment_indices`: `[team_id, public name, positions]`.
pub team_model: Vec<(String, String, Value)>,
/// `model_names`: a set, its members in the order they joined.
pub model_names: Vec<Value>,
pub deployment_names: Vec<Value>,
pub team_public_model_names: Vec<String>,
pub provider_default_ids: Vec<String>,
/// `pattern_router`.
pub patterns: PatternRouter,
/// `team_pattern_routers`: team id to its router.
pub team_patterns: Vec<(String, PatternRouter)>,
}
/// A host write to the indexes: each field it names replaces the one held.
#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
#[serde(default)]
pub struct IndexWrite {
pub id_map: Option<Map<String, Value>>,
pub model_name: Option<Map<String, Value>>,
pub team_model: Option<Vec<(String, String, Value)>>,
pub model_names: Option<Vec<Value>>,
pub deployment_names: Option<Vec<Value>>,
pub team_public_model_names: Option<Vec<String>>,
pub provider_default_ids: Option<Vec<String>>,
pub patterns: Option<PatternRouter>,
pub team_patterns: Option<Vec<(String, PatternRouter)>>,
}
/// A removal the Python router could not finish: an index held a position it cannot compare
/// (`_without_indices` raised `TypeError`). The indexes keep what it shifted before that.
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("'>' not supported between instances of '{found}' and 'int'")]
pub struct RemovalError {
pub found: &'static str,
}
impl HostIndexes {
/// Apply a host write.
pub fn write(&mut self, write: IndexWrite) {
let IndexWrite {
id_map,
model_name,
team_model,
model_names,
deployment_names,
team_public_model_names,
provider_default_ids,
patterns,
team_patterns,
} = write;
if let Some(value) = id_map {
self.id_map = value;
}
if let Some(value) = model_name {
self.model_name = value;
}
if let Some(value) = team_model {
self.team_model = value;
}
if let Some(value) = model_names {
self.model_names = unique(value);
}
if let Some(value) = deployment_names {
self.deployment_names = value;
}
if let Some(value) = team_public_model_names {
self.team_public_model_names = value;
}
if let Some(value) = provider_default_ids {
self.provider_default_ids = value;
}
if let Some(value) = patterns {
self.patterns = value;
}
if let Some(value) = team_patterns {
self.team_patterns = value;
}
}
/// `set_model_list` starting over: the id, name and team indexes and the team routers go;
/// the deployment names, the wildcard router and the provider default ids stay.
pub fn reset(&mut self) {
self.id_map.clear();
self.model_name.clear();
self.team_model.clear();
self.team_patterns.clear();
self.team_public_model_names.clear();
}
/// `set_model_list`'s `model_names`: the names it was given.
pub fn set_model_names(&mut self, names: Vec<String>) {
self.model_names = unique(names.into_iter().map(Value::String).collect());
}
/// A deployment joined the list at `position` (`_add_deployment` then
/// `_add_model_to_list_and_index_map`); `add` is `add_deployment`, which also names its group.
pub fn added(&mut self, entry: &Map<String, Value>, position: usize, add: bool) {
let params = entry.get("litellm_params").and_then(Value::as_object);
let info = entry.get("model_info").and_then(Value::as_object);
let model_name = entry.get("model_name").and_then(Value::as_str);
let id = info
.and_then(|info| info.get("id"))
.filter(|id| !id.is_null());
self.deployment_names.push(
params
.and_then(|params| params.get("model"))
.cloned()
.unwrap_or(Value::Null),
);
if let Some(name) = model_name.filter(|name| name.contains('*')) {
self.patterns.add(name, Value::Object(entry.clone()));
if let Some(id) = id.and_then(Value::as_str).filter(|id| !id.is_empty()) {
self.provider_default_ids.push(id.to_owned());
}
}
let team_id = info
.and_then(|info| info.get("team_id"))
.filter(|v| !v.is_null());
let public = info
.and_then(|info| info.get("team_public_model_name"))
.filter(|v| !v.is_null());
if let (Some(team_id), Some(Value::String(public))) = (team_id, public)
&& public.contains('*')
{
let team = python_str(team_id);
let router = match self.team_patterns.iter_mut().find(|(t, _)| *t == team) {
Some((_, router)) => router,
None => {
self.team_patterns.push((team, PatternRouter::default()));
&mut self.team_patterns.last_mut().expect("just pushed").1
}
};
router.add(public, Value::Object(entry.clone()));
}
let position = Value::from(position);
if let Some(id) = id {
self.id_map.insert(python_str(id), position.clone());
}
if let Some(name) = model_name.filter(|name| !name.is_empty()) {
push_position(
self.model_name
.entry(name.to_owned())
.or_insert_with(|| Value::Array(Vec::new())),
&position,
);
}
if add {
let name = Value::String(model_name.unwrap_or_default().to_owned());
if !self.model_names.contains(&name) {
self.model_names.push(name);
}
}
// `_update_team_model_index`: a truthy team id and public name.
if let (Some(team_id), Some(public)) = (team_id, public)
&& truthy(team_id)
&& truthy(public)
{
let (team, public) = (python_str(team_id), python_str(public));
if !self.team_public_model_names.contains(&public) {
self.team_public_model_names.push(public.clone());
}
let held = match self
.team_model
.iter_mut()
.find(|(t, p, _)| *t == team && *p == public)
{
Some((_, _, held)) => held,
None => {
self.team_model
.push((team, public, Value::Array(Vec::new())));
&mut self.team_model.last_mut().expect("just pushed").2
}
};
if let Value::Array(positions) = held
&& !positions.contains(&position)
{
positions.push(position);
}
}
}
/// The deployment `id` left the list from `position` (`_update_deployment_indices_after_removal`).
pub fn removed(&mut self, id: &str, position: usize) -> Result<(), RemovalError> {
let mut id_map = Map::new();
for (key, held) in std::mem::take(&mut self.id_map) {
if key == id {
continue;
}
let shifted = shift(&held, position)?.unwrap_or(held);
id_map.insert(key, shifted);
}
self.id_map = id_map;
let names: Vec<String> = self.model_name.keys().cloned().collect();
for name in names {
let kept = without(&self.model_name[&name], position)?;
if kept.is_empty() {
self.model_name.remove(&name);
self.model_names
.retain(|held| held.as_str() != Some(name.as_str()));
} else {
self.model_name.insert(name, Value::Array(kept));
}
}
let mut team_model = Vec::new();
let held = std::mem::take(&mut self.team_model);
let mut rest = held.into_iter();
for (team, public, positions) in rest.by_ref() {
match without(&positions, position) {
Ok(kept) if kept.is_empty() => {}
Ok(kept) => team_model.push((team, public, Value::Array(kept))),
Err(error) => {
team_model.push((team, public, positions));
team_model.extend(rest);
self.team_model = team_model;
return Err(error);
}
}
}
self.team_model = team_model;
let mut publics: Vec<String> = Vec::new();
for (_, public, _) in &self.team_model {
if !publics.contains(public) {
publics.push(public.clone());
}
}
self.team_public_model_names = publics;
self.patterns.remove(id);
for (_, router) in &mut self.team_patterns {
router.remove(id);
}
self.team_patterns
.retain(|(_, router)| !router.patterns.is_empty());
Ok(())
}
/// Whether the id map holds `id` (`_has_model_id`).
pub fn has_id(&self, id: &str) -> bool {
self.id_map.contains_key(id)
}
/// What the id map holds for `id`, as written (`model_id_to_deployment_index_map.get(id)`).
pub fn id_entry(&self, id: &str) -> Option<&Value> {
self.id_map.get(id)
}
/// What the name index holds for `name`, as written.
pub fn named(&self, name: &str) -> Option<&Value> {
self.model_name.get(name)
}
/// What the team index holds for `(team, public)`, as written.
pub fn team_named(&self, team: &str, public: &str) -> Option<&Value> {
self.team_model
.iter()
.find(|(t, p, _)| t == team && p == public)
.map(|(_, _, positions)| positions)
}
/// `team_public_model_names_of`: the public names the team index holds for `team`.
pub fn team_public_names_of(&self, team: &str) -> Vec<String> {
let mut names: Vec<String> = Vec::new();
for (t, public, _) in &self.team_model {
if t == team && !names.contains(public) {
names.push(public.clone());
}
}
names
}
/// The position the id map holds for `id`, when it is a position.
pub fn position_of(&self, id: &str) -> Option<usize> {
self.id_map
.get(id)
.and_then(Value::as_u64)
.and_then(|position| usize::try_from(position).ok())
}
}
impl PatternRouter {
/// `add_pattern`: a new regex joins in specificity order (a stable sort, most specific
/// first); a known one gains the deployment.
pub fn add(&mut self, pattern: &str, deployment: Value) {
let regex = pattern_to_regex(pattern);
if let Some((_, deployments)) = self.patterns.iter_mut().find(|(held, _)| *held == regex) {
deployments.push(deployment);
return;
}
self.patterns.push((regex, vec![deployment]));
self.patterns
.sort_by_key(|(regex, _)| std::cmp::Reverse(regex_specificity(regex)));
}
/// `remove_deployment`: every deployment with this id goes, and so does a pattern left empty.
pub fn remove(&mut self, id: &str) {
for (_, deployments) in &mut self.patterns {
deployments.retain(|deployment| {
deployment
.get("model_info")
.and_then(|info| info.get("id"))
.and_then(Value::as_str)
!= Some(id)
});
}
self.patterns
.retain(|(_, deployments)| !deployments.is_empty());
}
/// `route`: the deployments of the first pattern `request` matches, each a copy whose
/// `litellm_params.model` takes the parts the wildcards matched; `filtered` keeps to the
/// patterns of those names.
pub fn route(&self, request: &str, filtered: Option<&[String]>) -> Option<Vec<Value>> {
let filtered: Option<Vec<String>> =
filtered.map(|names| names.iter().map(|name| pattern_to_regex(name)).collect());
for (regex, deployments) in &self.patterns {
if filtered
.as_ref()
.is_some_and(|names| !names.contains(regex))
{
continue;
}
let Some(groups) = regex_match(regex, request) else {
continue;
};
let mut routed = Vec::with_capacity(deployments.len());
for deployment in deployments {
let mut copy = deployment.clone();
let model = deployment
.get("litellm_params")
.and_then(|params| params.get("model"));
let Some(Value::String(model)) = model else {
// `deployment["litellm_params"]["model"]` raising ends the route with None.
return None;
};
let named = deployment_model_name(&groups, request, model);
if model.contains('*')
&& let Some(params) =
deployment.get("litellm_params").and_then(Value::as_object)
&& let Some(info) = copy.get_mut("model_info").and_then(Value::as_object_mut)
{
for flag in cache_flags_left_to_the_catalog(params) {
info.remove(*flag);
}
}
if let Some(params) = copy
.get_mut("litellm_params")
.and_then(Value::as_object_mut)
{
params.insert("model".to_owned(), Value::String(named));
}
routed.push(copy);
}
return Some(routed);
}
None
}
/// `get_models_by_pattern`'s pick: the first (request, pattern) pair that matches, trying
/// `model` against every pattern, then `provider/model`.
pub fn models_by_pattern(&self, requests: &[String]) -> Option<Vec<String>> {
for request in requests {
for (regex, deployments) in &self.patterns {
if let Some(groups) = regex_match(regex, request) {
return Some(
deployments
.iter()
.filter_map(|deployment| {
deployment
.get("litellm_params")
.and_then(|params| params.get("model"))
.and_then(Value::as_str)
.map(|model| deployment_model_name(&groups, request, model))
})
.collect(),
);
}
}
}
None
}
}
/// `_cache_flags_left_to_the_catalog`: the `model_info` flags a wildcard deployment does not speak
/// for on the models it matches. An azure_ai wildcard's `supports_prompt_cache_breakpoint` does.
/// `inferred_provider` names `azure_ai` only for a model carrying that prefix.
pub(crate) fn cache_flags_left_to_the_catalog(
params: &Map<String, Value>,
) -> &'static [&'static str] {
let provider = params
.get("custom_llm_provider")
.filter(|provider| truthy(provider))
.and_then(Value::as_str)
.or_else(|| {
params
.get("model")
.and_then(Value::as_str)
.and_then(|model| model.split_once('/'))
.map(|(prefix, _)| prefix)
});
if provider == Some("azure_ai") {
&["supports_prompt_caching"]
} else {
&[
"supports_prompt_cache_breakpoint",
"supports_prompt_caching",
]
}
}
/// `re.escape(pattern).replace(r"\*", "(.*)")`.
pub fn pattern_to_regex(pattern: &str) -> String {
let mut regex = String::new();
for c in pattern.chars() {
match c {
'*' => regex.push_str("(.*)"),
'(' | ')' | '[' | ']' | '{' | '}' | '?' | '+' | '-' | '|' | '^' | '$' | '\\' | '.'
| '&' | '~' | '#' | ' ' | '\t' | '\n' | '\r' | '\x0b' | '\x0c' => {
regex.push('\\');
regex.push(c);
}
_ => regex.push(c),
}
}
regex
}
/// `PatternUtils.calculate_pattern_specificity`: the regex's length, then how much regex syntax
/// it has.
pub fn regex_specificity(regex: &str) -> (usize, usize) {
let complexity = regex
.chars()
.filter(|c| matches!(c, '*' | '+' | '?' | '\\' | '^' | '$' | '|' | '(' | ')'))
.count();
(regex.chars().count(), complexity)
}
/// A pattern regex read back as its literal runs between wildcards, or None when it holds regex
/// syntax `pattern_to_regex` never writes (a caller wrote it).
fn literal_runs(regex: &str) -> Option<Vec<String>> {
let mut runs = vec![String::new()];
let mut chars = regex.chars().peekable();
while let Some(c) = chars.next() {
match c {
'\\' => runs.last_mut()?.push(chars.next()?),
'(' => {
let rest: String = chars.by_ref().take(3).collect();
if rest != ".*)" {
return None;
}
runs.push(String::new());
}
'.' | '*' | '+' | '?' | '[' | ']' | '{' | '}' | '|' | '^' | '$' | ')' => return None,
other => runs.last_mut()?.push(other),
}
}
Some(runs)
}
/// `re.match(regex, request)` for a pattern regex: anchored at the start only, each `(.*)`
/// greedy. The groups it captured, or None when it does not match.
pub fn regex_match(regex: &str, request: &str) -> Option<Vec<String>> {
let runs = literal_runs(regex)?;
let rest = request.strip_prefix(runs[0].as_str())?;
let mut groups = Vec::new();
greedy(&runs[1..], rest, &mut groups).then_some(groups)
}
/// Match `runs` (each a literal that follows a group) against `rest`, longest group first.
fn greedy(runs: &[String], rest: &str, groups: &mut Vec<String>) -> bool {
let Some((literal, later)) = runs.split_first() else {
return true;
};
let mut ends: Vec<usize> = rest.char_indices().map(|(at, _)| at).collect();
ends.push(rest.len());
for &end in ends.iter().rev() {
let (group, after) = rest.split_at(end);
let Some(after) = after.strip_prefix(literal.as_str()) else {
continue;
};
groups.push(group.to_owned());
if greedy(later, after, groups) {
return true;
}
groups.pop();
}
false
}
/// `set_deployment_model_name`: the deployment's model with each wildcard replaced by the group
/// matched for it, the request itself when more groups matched than it has wildcards.
pub fn deployment_model_name(groups: &[String], request: &str, model: &str) -> String {
if !model.contains('*') {
return model.to_owned();
}
if groups.len() > model.matches('*').count() {
return request.to_owned();
}
let mut named = model.to_owned();
for group in groups {
named = named.replacen('*', group, 1);
}
named
}
fn unique(values: Vec<Value>) -> Vec<Value> {
let mut kept: Vec<Value> = Vec::with_capacity(values.len());
for value in values {
if !kept.contains(&value) {
kept.push(value);
}
}
kept
}
fn push_position(held: &mut Value, position: &Value) {
if let Value::Array(positions) = held {
positions.push(position.clone());
}
}
/// Python's truthiness for a JSON value.
pub(crate) fn truthy(value: &Value) -> bool {
match value {
Value::Null => false,
Value::Bool(flag) => *flag,
Value::Number(number) => number.as_f64() != Some(0.0),
Value::String(text) => !text.is_empty(),
Value::Array(items) => !items.is_empty(),
Value::Object(items) => !items.is_empty(),
}
}
/// Python's `str()` for an id or team id.
pub(crate) fn python_str(value: &Value) -> String {
match value {
Value::String(text) => text.clone(),
Value::Bool(true) => "True".to_owned(),
Value::Bool(false) => "False".to_owned(),
Value::Null => "None".to_owned(),
other => other.to_string(),
}
}
fn type_name(value: &Value) -> &'static str {
match value {
Value::Null => "NoneType",
Value::Bool(_) => "bool",
Value::Number(number) if number.is_f64() => "float",
Value::Number(_) => "int",
Value::String(_) => "str",
Value::Array(_) => "list",
Value::Object(_) => "dict",
}
}
/// `idx - 1 if idx > removal_idx else idx`: None keeps `held` as is.
fn shift(held: &Value, position: usize) -> Result<Option<Value>, RemovalError> {
let removal = position as f64;
let number = match held {
Value::Bool(flag) => f64::from(u8::from(*flag)),
Value::Number(number) => number.as_f64().unwrap_or_default(),
other => {
return Err(RemovalError {
found: type_name(other),
});
}
};
if number <= removal {
return Ok(None);
}
Ok(Some(match held {
Value::Number(n) if n.is_i64() || n.is_u64() => Value::from(n.as_i64().unwrap_or(0) - 1),
_ => Value::from(number - 1.0),
}))
}
/// `_without_indices`: the positions without `position`, the later ones shifted down.
fn without(held: &Value, position: usize) -> Result<Vec<Value>, RemovalError> {
let Value::Array(positions) = held else {
return Ok(Vec::new());
};
let removal = Value::from(position);
let mut kept = Vec::with_capacity(positions.len());
for index in positions {
if *index == removal {
continue;
}
kept.push(shift(index, position)?.unwrap_or_else(|| index.clone()));
}
Ok(kept)
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
fn entry(name: &str, id: &str, model: &str) -> Map<String, Value> {
serde_json::from_value(json!({
"model_name": name,
"litellm_params": {"model": model},
"model_info": {"id": id},
}))
.unwrap()
}
#[test]
fn removal_shifts_every_index_and_drops_emptied_names() {
let mut indexes = HostIndexes::default();
indexes.added(&entry("gpt", "a", "openai/a"), 0, false);
indexes.added(&entry("gpt", "b", "openai/b"), 1, false);
indexes.added(&entry("claude", "c", "anthropic/c"), 2, false);
indexes.set_model_names(vec!["gpt".into(), "claude".into(), "broken".into()]);
indexes.removed("a", 0).unwrap();
assert_eq!(
indexes.id_map,
json!({"b": 0, "c": 1}).as_object().unwrap().clone()
);
assert_eq!(indexes.model_name["gpt"], json!([0]));
indexes.removed("c", 1).unwrap();
assert!(!indexes.model_name.contains_key("claude"));
assert_eq!(indexes.model_names, vec![json!("gpt"), json!("broken")]);
assert_eq!(indexes.deployment_names.len(), 3);
}
#[test]
fn a_pattern_routes_with_the_matched_parts() {
let mut router = PatternRouter::default();
router.add(
"llmengine/*",
json!({"litellm_params": {"model": "openai/*"}, "model_info": {"id": "w"}}),
);
router.add(
"*meta.llama3*",
json!({"litellm_params": {"model": "bedrock/meta.llama3*"}, "model_info": {"id": "m"}}),
);
let routed = router.route("llmengine/foo", None).unwrap();
assert_eq!(routed[0]["litellm_params"]["model"], "openai/foo");
let routed = router.route("hello-world-meta.llama3-70b", None).unwrap();
assert_eq!(
routed[0]["litellm_params"]["model"],
"hello-world-meta.llama3-70b"
);
assert_eq!(
regex_match(r"a(.*)b(.*)", "axbyb").unwrap(),
vec!["xby", ""]
);
router.remove("w");
assert!(router.route("llmengine/foo", None).is_none());
}
#[rstest]
#[case::any_wildcard(
json!({"model": "openai/*"}),
json!({"id": "w", "keep": 1}),
)]
#[case::azure_ai_keeps_the_breakpoint_flag(
json!({"model": "azure_ai/*"}),
json!({"id": "w", "keep": 1, "supports_prompt_cache_breakpoint": true}),
)]
#[case::azure_ai_named_by_custom_llm_provider(
json!({"model": "*", "custom_llm_provider": "azure_ai"}),
json!({"id": "w", "keep": 1, "supports_prompt_cache_breakpoint": true}),
)]
#[case::a_fixed_model_speaks_for_itself(
json!({"model": "openai/gpt-4o"}),
json!({
"id": "w",
"keep": 1,
"supports_prompt_cache_breakpoint": true,
"supports_prompt_caching": true,
}),
)]
fn a_wildcard_leaves_the_cache_flags_to_the_catalog(
#[case] params: Value,
#[case] expected: Value,
) {
let mut router = PatternRouter::default();
router.add(
"any/*",
json!({
"litellm_params": params,
"model_info": {
"id": "w",
"keep": 1,
"supports_prompt_cache_breakpoint": true,
"supports_prompt_caching": true,
},
}),
);
let routed = router.route("any/gpt-4o", None).unwrap();
assert_eq!(routed[0]["model_info"], expected);
}
}

View file

@ -43,21 +43,706 @@
//! - `call(op, request)`: select, run the attempt, and retry or fall back until it succeeds or the
//! policy is exhausted.
//!
//! # Execution
//! Each attempt resolves an `Executor` after a deployment is picked. A Rust-native operation runs
//! its `inference-*` route machine in-process; anything else goes through `Host::invoke`. A
//! started attempt is never replayed on a different executor.
//! # Hosts: the only boundary
//! The core reaches outside itself only through [`Host`], and every attempt runs through
//! `Host::invoke`: the core picks the deployment, the host runs the call, the core records the
//! outcome. There is no second, "native" attempt path in the core. Which code runs an attempt is
//! the host's business:
//! - `PythonHost` (inside a Python process) builds the request for the picked deployment and calls
//! `litellm.<op>`. It never calls back into a Python router.
//!
//! # Hosts
//! `Host` is what the embedding process provides: `invoke` for non-native operations, `select`
//! for strategies and sub-routers not yet ported, `allow_fallback` for access and budget checks,
//! `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that holds all
//! routing state under the same keys, encodings and TTLs the Python router uses. `PythonHost`
//! serves them from the existing Python code when the router runs inside Python; `NativeHost`
//! serves them in the Rust server, where an operation without a Rust implementation makes the
//! config fail at build time, never partway through a request.
//! An op or provider a host cannot run fails `Host::supports` when the router is built or
//! changed, never partway through a request.
//!
//! The core also builds its deployments through the host ([`build`]): for each configured entry
//! it finds the secret references and asks `secrets`, then `validate`, `provider` and `added`.
//!
//! Besides `invoke` and `supports`, a host provides `select` for strategies and plugins not yet
//! ported, `allow_fallback` for access and budget checks, `refuses_fallbacks` for a chain's dict
//! targets, `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that
//! holds all routing state under the same keys, encodings and TTLs the Python router uses. The
//! core writes cooldowns, allowed-fails counters and usage itself, through that cache.
//!
//! # Plugins
//! Routing strategies, pre-call checks and sub-routers (auto, complexity, quality, adaptive)
//! register through `register_plugin` and see a `RoutingContext` (registry, settings, health and
//! the current request), never the router itself.
//!
//! # Modules
//! - [`registry`]: deployments and their indices, `deployment_id`, `ChangeSet`.
//! - [`config`] / [`settings`]: static and dynamic config, and the frozen `Settings` view.
//! - [`cache`] / [`health`]: the shared routing cache, cooldowns and usage.
//! - [`select`]: resolve, filter chain, simple-shuffle.
//! - [`call`]: retries and fallbacks.
//! - [`eligibility`]: the endpoints a deployment refuses to serve.
//! - [`precall`]: the `enable_pre_call_checks` filter.
//! - [`model_info`]: model info merged from the cost map, asked of the host at query time.
//! - [`host`]: the `Host` boundary.
pub mod access;
pub mod build;
pub mod cache;
pub mod call;
pub mod config;
pub mod eligibility;
pub mod error;
pub mod fallback;
pub mod group_info;
pub mod health;
pub mod health_routing;
pub mod host;
pub mod host_index;
pub mod lookup;
pub mod mock;
pub mod model_info;
pub mod outcome;
pub mod precall;
pub mod prerouting;
pub mod ptu;
pub mod registry;
pub mod resolve;
pub mod retry;
pub mod routing_groups;
pub mod routing_plugins;
pub mod select;
pub mod settings;
pub mod strategy;
pub mod weights;
use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
};
use arc_swap::ArcSwap;
use parking_lot::Mutex;
pub use crate::{
cache::{InMemoryRouterCache, RouterCache},
call::RouterResponse,
config::{Op, RouterConfig, RouterDynamicConfig, RoutingStrategy},
error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported},
health::{Clock, Health, Outcome, SystemClock},
host::{
Attempt, Blocked, BuildAsk, BuildPath, Built, Host, Output, Provided, RetryCause,
RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage,
},
outcome::{OutcomeCallbacks, OutcomeWrite},
prerouting::{PreRoute, PreRouted},
registry::{ChangeSet, Deployment, DeploymentSpec, Registry, RegistryVersion},
select::{Randomness, SeededRandomness, ThreadRandomness},
settings::Settings,
};
use crate::{
health::{GroupShape, routing_group_has_alternatives, team_model_has_alternatives},
registry::DeploymentFilter,
select::SelectInputs,
};
/// The registry and settings one call reads, swapped atomically by `apply`.
pub(crate) struct RouterStateSnapshot {
pub registry: Arc<Registry>,
pub settings: Arc<Settings>,
}
impl RouterStateSnapshot {
/// `_get_routing_context` for a pick of `name`: unless the request overrides the strategy,
/// a stale member map naming a group no longer configured raises `KeyError(group)`.
pub(crate) fn missing_routing_group(
&self,
request: &RouterRequest,
name: &str,
) -> Result<(), Error> {
if strategy::request_override(request.routing_strategy.as_deref()).is_some() {
return Ok(());
}
match self.registry.missing_member_group(name) {
Some(group) => Err(Error::MissingRoutingGroup {
group: group.to_owned(),
}),
None => Ok(()),
}
}
}
pub(crate) struct Core {
host: Arc<dyn Host>,
rng: Arc<dyn Randomness>,
health: Health,
config: RouterConfig,
snapshot: ArcSwap<RouterStateSnapshot>,
writer: Mutex<()>,
retired: AtomicBool,
}
impl Core {
/// Reject what this host cannot serve, before any request.
fn validate(&self, registry: &Registry, settings: &Settings) -> Result<(), ConfigError> {
let can_select = self.host.can_select();
let strategies = std::iter::once(&settings.routing_strategy).chain(
registry
.routing_groups()
.iter()
.map(|group| &group.routing_strategy),
);
for strategy in strategies {
if !strategy.is_native() && !can_select {
return Err(Unsupported::Strategy(strategy.clone()).into());
}
}
if let Some(plugin) = settings.plugins.first()
&& !can_select
{
return Err(Unsupported::Plugin(plugin.name.clone()).into());
}
for deployment in registry.deployments(&DeploymentFilter::default()) {
self.validate_deployment(&deployment)?;
}
Ok(())
}
/// Reject a deployment whose provider this host cannot run a served op on.
fn validate_deployment(&self, deployment: &Deployment) -> Result<(), ConfigError> {
for op in &self.config.served_ops {
if !self.host.supports(op, deployment.provider()) {
return Err(Unsupported::Op {
op: op.clone(),
provider: deployment.provider().to_owned(),
}
.into());
}
}
Ok(())
}
}
/// How a reported failure is recorded.
#[derive(Clone, Debug, Default)]
pub struct FailureReport {
/// Cool down for the deployment's own `cooldown_time` (then the provider's retry-after),
/// as the failure callback does, rather than the router default.
pub deployment_cooldown: bool,
/// The `cooldown_time` the failed call's own params carried (`model_info`, then
/// `litellm_params`), which wins over the provider's retry-after for a callback report.
pub cooldown_time: Option<f64>,
/// Count the failure only: its status is one no cooldown rule reads (an empty status).
pub count_only: bool,
/// The model group the request was for, whose fallbacks keep a lone deployment cooling.
pub requested_group: Option<String>,
/// Cool down without counting the failure (`set_cooldown_deployments` alone).
pub uncounted: bool,
/// Also count the failed request against the deployment's per-minute requests
/// (`async_deployment_callback_on_failure`).
pub counts_request: bool,
/// The model the request's usage is keyed by (metadata's `deployment`, which a wildcard route
/// fills), when it is not the deployment's own.
pub usage_model: Option<String>,
/// The failure callbacks of a logged failure; a report no callback makes keeps the default.
pub callbacks: OutcomeCallbacks,
}
/// How a reported success is recorded.
#[derive(Clone, Debug, Default)]
pub struct SuccessReport {
/// See [`host::Output::tracks_usage`].
pub tracks_usage: Option<bool>,
/// See [`FailureReport::usage_model`].
pub usage_model: Option<String>,
/// See [`FailureReport::callbacks`].
pub callbacks: OutcomeCallbacks,
}
/// Builds a [`Router`], with test seams.
pub struct RouterBuilder {
config: RouterConfig,
dynamic: RouterDynamicConfig,
host: Arc<dyn Host>,
rng: Arc<dyn Randomness>,
clock: Arc<dyn Clock>,
}
impl RouterBuilder {
#[must_use]
pub fn randomness(self, rng: Arc<dyn Randomness>) -> Self {
Self { rng, ..self }
}
#[must_use]
pub fn clock(self, clock: Arc<dyn Clock>) -> Self {
Self { clock, ..self }
}
pub fn build(self) -> Result<Router, ConfigError> {
let registry = Registry::build(
self.config.model_list.clone(),
self.dynamic.model_alias_map.clone(),
self.config.routing_groups.clone(),
)?
.with_member_groups(self.config.routing_group_members.clone());
let settings = Settings::new(&self.config, Arc::new(self.dynamic));
let core = Core {
health: Health::new(self.host.cache(), self.clock),
host: self.host,
rng: self.rng,
config: self.config,
snapshot: ArcSwap::from_pointee(RouterStateSnapshot {
registry: Arc::new(Registry::default()),
settings: Arc::new(settings.clone()),
}),
writer: Mutex::new(()),
retired: AtomicBool::new(false),
};
core.validate(&registry, &settings)?;
core.snapshot.store(Arc::new(RouterStateSnapshot {
registry: Arc::new(registry),
settings: Arc::new(settings),
}));
Ok(Router {
core: Arc::new(core),
})
}
}
/// The router: the entry points `litellm.Router` delegates to.
#[derive(Clone)]
pub struct Router {
core: Arc<Core>,
}
impl Router {
pub fn new(
config: RouterConfig,
dynamic: RouterDynamicConfig,
host: Arc<dyn Host>,
) -> Result<Self, ConfigError> {
Self::builder(config, dynamic, host).build()
}
pub fn builder(
config: RouterConfig,
dynamic: RouterDynamicConfig,
host: Arc<dyn Host>,
) -> RouterBuilder {
RouterBuilder {
config,
dynamic,
host,
rng: Arc::new(ThreadRandomness),
clock: Arc::new(SystemClock),
}
}
fn live(&self) -> Result<(), Error> {
if self.core.retired.load(Ordering::Acquire) {
return Err(Error::Retired);
}
Ok(())
}
/// Select, run the attempt, and retry or fall back until it succeeds or the policy is
/// exhausted.
pub async fn call(&self, op: Op, request: RouterRequest) -> Result<RouterResponse, Error> {
self.live()?;
self.core.call(op, Arc::new(request)).await
}
/// Resolve, filter and pick a deployment without calling it.
pub async fn select(&self, op: &Op, request: &RouterRequest) -> Result<Arc<Deployment>, Error> {
self.live()?;
let snapshot = self.core.snapshot.load_full();
routing_plugins::refuse_sync_pick(
snapshot.settings.routing_plugins,
&snapshot.settings.routing_strategy,
request,
)?;
let routed = self
.core
.pre_routed(PreRoute {
op,
request,
name: &request.model,
fallback_depth: 0,
retry: 0,
})
.await?;
let name = match routed {
PreRouted::Model(name) => name,
PreRouted::Deployment(deployment) => return Ok(deployment),
};
snapshot.missing_routing_group(request, &name)?;
let candidates = self
.core
.candidates(&snapshot, op, request, &name, &SelectInputs::default())
.await?;
self.core.pick(&snapshot, op, request, &candidates).await
}
/// The deployments a pick for `request` would choose among: resolved, team- and
/// cooldown-filtered, through the plugins and the lowest `order`.
pub async fn healthy(
&self,
op: &Op,
request: &RouterRequest,
) -> Result<select::Candidates, Error> {
self.live()?;
let snapshot = self.core.snapshot.load_full();
let candidates = self
.core
.candidates(
&snapshot,
op,
request,
&request.model,
&SelectInputs::default(),
)
.await?;
// `_core_async_get_healthy_deployments` raises when nothing is left.
if candidates.healthy.is_empty() {
return Err(Error::NoHealthyDeployments {
model: candidates.failed_model(&snapshot.registry),
});
}
Ok(candidates)
}
pub fn registry(&self) -> Arc<Registry> {
self.core.snapshot.load().registry.clone()
}
pub fn settings(&self) -> Arc<Settings> {
self.core.snapshot.load().settings.clone()
}
pub fn health(&self) -> &Health {
&self.core.health
}
/// Record an outcome the router did not see itself, such as a health check. Returns whether
/// the deployment was cooled down.
pub async fn report(&self, id: &str, outcome: Outcome) -> Result<bool, Error> {
match outcome {
Outcome::Success { total_tokens } => self
.report_success(id, total_tokens as f64, 1.0, SuccessReport::default())
.await
.map(|_| false),
Outcome::Failure(failure) => self
.report_failure(id, &failure, FailureReport::default())
.await
.map(|cooldown| cooldown.is_some()),
}
}
/// A success a logging callback saw (`deployment_callback_on_success`): counted for a
/// configured deployment, with its tokens and requests when it tracks usage. Returns the
/// success counter's key, or `None` for an unknown deployment.
pub async fn report_success(
&self,
id: &str,
total_tokens: f64,
requests: f64,
report: SuccessReport,
) -> Result<Option<String>, Error> {
report
.callbacks
.record(
OutcomeWrite::Success,
self.reported_success(id, total_tokens, requests, &report),
)
.await
}
async fn reported_success(
&self,
id: &str,
total_tokens: f64,
requests: f64,
report: &SuccessReport,
) -> Result<Option<String>, Error> {
let snapshot = self.core.snapshot.load_full();
let Some(deployment) = snapshot.registry.deployment(id) else {
return Ok(None);
};
let key = self.core.health.count_success(id).await?;
if report
.tracks_usage
.unwrap_or_else(|| deployment.tracks_usage())
{
self.core
.health
.add_usage(
&deployment,
report.usage_model.as_deref(),
total_tokens,
requests,
)
.await?;
}
Ok(Some(key))
}
/// `sync_deployment_callback_on_success`: count a success against `id`, configured or not,
/// and return the counter's key.
pub async fn count_success(
&self,
id: &str,
callbacks: OutcomeCallbacks,
) -> Result<Option<String>, Error> {
callbacks
.record(OutcomeWrite::Success, async {
self.core.health.count_success(id).await.map(Some)
})
.await
}
/// A failure the router did not see itself: counted against `id` (known or not), then
/// cooled down by the rules (`record_failure`, `deployment_callback_on_failure`). Returns the
/// cooldown applied, if any.
pub async fn report_failure(
&self,
id: &str,
failure: &AttemptError,
report: FailureReport,
) -> Result<Option<std::time::Duration>, Error> {
report
.callbacks
.record(
OutcomeWrite::Failure,
self.reported_failure(id, failure, &report),
)
.await
}
async fn reported_failure(
&self,
id: &str,
failure: &AttemptError,
report: &FailureReport,
) -> Result<Option<std::time::Duration>, Error> {
if !report.uncounted && !failure.outside_deployment {
self.core.health.count_failure(id).await?;
}
let snapshot = self.core.snapshot.load_full();
let Some(deployment) = snapshot.registry.deployment(id) else {
return Ok(None);
};
if report.counts_request {
self.core
.health
.count_failed_request(&deployment, report.usage_model.as_deref())
.await?;
}
if failure.outside_deployment {
return Ok(None);
}
// An empty status is no cooldown (`_is_cooldown_required`), unless the deployment's
// own policy names the error.
if report.count_only
&& crate::health::deployment_allowed_fails(
deployment.model_info.allowed_fails_policy.as_ref(),
failure,
)
.is_none()
{
return Ok(None);
}
let cooldown = if report.deployment_cooldown {
report
.cooldown_time
.filter(|seconds| *seconds >= 0.0)
.and_then(|seconds| std::time::Duration::try_from_secs_f64(seconds).ok())
.or(failure.retry_after)
.unwrap_or(snapshot.settings.cooldown_time)
} else {
snapshot.settings.cooldown_time
};
let shape = GroupShape {
single_deployment: snapshot
.registry
.group(&deployment.model_name)
.is_some_and(|group| group.len() == 1)
&& !report
.requested_group
.as_deref()
.is_some_and(|group| routing_group_has_alternatives(&snapshot.registry, group))
&& !team_model_has_alternatives(&snapshot.registry, &deployment),
has_fallback_group: snapshot
.settings
.dynamic
.fallback_targets(&[report
.requested_group
.as_deref()
.unwrap_or(&deployment.model_name)])
.is_some(),
};
self.core
.health
.cool_down_after(&snapshot.settings, &deployment, shape, failure, cooldown)
.await
}
/// Apply one change set atomically and return the new registry version.
pub fn apply(&self, changes: ChangeSet) -> Result<RegistryVersion, ConfigError> {
let _writer = self.core.writer.lock();
let current = self.core.snapshot.load_full();
let aliases = changes
.settings
.as_ref()
.map(|dynamic| dynamic.model_alias_map.clone());
let registry = current.registry.apply(&changes, aliases)?;
let settings = match changes.settings {
Some(dynamic) => Settings::new(&self.core.config, Arc::new(dynamic)),
None => Settings::clone(&current.settings),
};
self.core.validate(&registry, &settings)?;
let version = registry.version();
self.core.snapshot.store(Arc::new(RouterStateSnapshot {
registry: Arc::new(registry),
settings: Arc::new(settings),
}));
Ok(version)
}
/// `get_router_model_info` for one deployment, its cost-map reads asked of the host now
/// ([`model_info::router_model_info`]).
pub async fn router_model_info(
&self,
query: &model_info::RouterModelInfoQuery,
) -> Result<Option<model_info::MergedModelInfo>, Error> {
model_info::router_model_info(self.core.host.as_ref(), query).await
}
/// `_set_model_group_info` for the deployments `query` names ([`group_info::model_group_info`]).
pub async fn model_group_info(
&self,
query: &group_info::GroupInfoQuery,
) -> Result<Option<group_info::GroupInfo>, Error> {
group_info::model_group_info(self.core.host.as_ref(), query).await
}
/// `get_configured_token_limits` for deployment `id` ([`group_info::configured_token_limits`]).
pub async fn configured_token_limits(
&self,
id: serde_json::Value,
configured: &serde_json::Map<String, serde_json::Value>,
) -> Result<(Option<i64>, Option<i64>), Error> {
group_info::configured_token_limits(self.core.host.as_ref(), id, configured).await
}
/// `get_model_listing_info` over a listed name's deployments ([`group_info::model_listing_info`]).
pub async fn model_listing_info(
&self,
deployments: &[serde_json::Value],
) -> Result<group_info::ListingInfo, Error> {
group_info::model_listing_info(self.core.host.as_ref(), deployments).await
}
/// `get_deployment_model_info` for deployment `model_id` serving `model_name`
/// ([`model_info::deployment_model_info`]).
pub async fn deployment_model_info(
&self,
model_id: &str,
model_name: &serde_json::Value,
) -> Result<Option<serde_json::Map<String, serde_json::Value>>, Error> {
model_info::deployment_model_info(self.core.host.as_ref(), model_id, model_name).await
}
/// Build `entries` through the host ([`build::build`]), installing each deployment it
/// accepts as soon as it does, as `install` says: a replace empties the registry first (and
/// takes `set_model_list`'s `model_names` once every entry is built), an append or upsert
/// adds at the end. When a step fails, what was installed before it stays, as the Python
/// router keeps the deployments it added before the one that raised.
///
/// The build installs into a registry of its own, indexed in place, and serves it once it
/// ends (or fails), so n deployments build in O(n). A write that lands while it runs (a step
/// answer that applies a changeset) is kept: what the build installed is added over it.
pub async fn build(
&self,
path: BuildPath,
entries: Vec<serde_json::Map<String, serde_json::Value>>,
resolve_secrets: bool,
install: build::Install,
) -> Result<Vec<build::BuiltDeployment>, Error> {
let model_names: Vec<String> = entries
.iter()
.map(|entry| match entry.get("model_name") {
Some(serde_json::Value::String(name)) => name.clone(),
None | Some(serde_json::Value::Null) => "None".to_owned(),
Some(other) => other.to_string(),
})
.collect();
let replace = install == build::Install::Replace;
let base = self.core.snapshot.load_full().registry.clone();
let mut working = if replace {
base.emptied()
} else {
Registry::clone(&base)
};
let mut installed: Vec<DeploymentSpec> = Vec::new();
let add = path == BuildPath::Add;
let core = &self.core;
let built = build::build(
core.host.as_ref(),
path,
entries,
resolve_secrets,
|deployment| {
working.push_built(deployment.spec.clone(), add)?;
installed.push(deployment.spec.clone());
core.validate_deployment(working.last().expect("just installed"))
.map_err(Error::from)
},
)
.await;
if built.is_ok() && replace && path == BuildPath::Create {
working.set_model_names(model_names);
}
self.serve_built(&base, working, installed, add)?;
built
}
/// Serve the registry a build made from `base`. When another write replaced `base` while
/// the build ran, the build's deployments are added over that write instead.
fn serve_built(
&self,
base: &Arc<Registry>,
built: Registry,
installed: Vec<DeploymentSpec>,
add: bool,
) -> Result<(), ConfigError> {
let _writer = self.core.writer.lock();
let current = self.core.snapshot.load_full();
let registry = if Arc::ptr_eq(&current.registry, base) {
built
} else {
let mut next = Registry::clone(&current.registry);
for spec in installed {
next.push_built(spec, add)?;
}
next
};
self.core.snapshot.store(Arc::new(RouterStateSnapshot {
registry: Arc::new(registry),
settings: current.settings.clone(),
}));
Ok(())
}
/// Keep the providers `from`'s host resolved for the deployments this router serves too (a
/// router rebuilt over the same deployments for new static settings).
pub fn adopt_providers(&self, from: &Self) {
let _writer = self.core.writer.lock();
let current = self.core.snapshot.load_full();
let registry = current
.registry
.with_providers_of(&from.core.snapshot.load().registry);
self.core.snapshot.store(Arc::new(RouterStateSnapshot {
registry: Arc::new(registry),
settings: current.settings.clone(),
}));
}
/// Stop serving: every later call and selection fails with [`Error::Retired`].
pub fn retire(&self) {
self.core.retired.store(true, Ordering::Release);
}
}

View file

@ -0,0 +1,883 @@
//! The registry lookups the host serves as `registry()` queries, answered over the deployments and
//! the host indexes ([`crate::host_index::HostIndexes`]) the way the Python router answers them.
//!
//! A lookup answers with [`Row`]s: positions in the deployment list (each with how the row renames
//! the deployment), wildcard routes, or a position the Python list could not be read at. The host
//! turns a position into the entry object it stores, so identity and in-place reads stay its own;
//! a [`Row::Raw`] is handed back for the host to index its list by, which raises exactly where the
//! Python router's own read raised (an index a caller wrote junk into).
use std::sync::Arc;
use serde_json::{Map, Value, json};
use crate::{
host_index::{HostIndexes, python_str, truthy},
registry::{Deployment, Registry},
};
/// One row of a lookup's answer.
#[derive(Clone, Debug, PartialEq)]
pub enum Row {
/// The deployment at `position`, under `name` when the row renames it (an alias or a routing
/// group), as a routing group's member `member` of `group` when `group` is set (its row drops
/// the member's `access_groups` and takes the group's priority as its `order`).
At {
position: usize,
name: Option<String>,
group: Option<(String, String)>,
},
/// A wildcard route: a copy of a pattern router's deployment, under `name`.
Pattern { entry: Value, name: String },
/// A position the list cannot be read at: the host indexes its list by it, and raises.
Raw(Value),
/// A position container that is not a list: the host iterates it as its list's positions.
RawPositions(Value),
}
impl Row {
/// The row as the host reads it.
pub fn wire(&self) -> Value {
match self {
Self::At {
position,
name,
group,
} => {
let mut row = Map::new();
row.insert("at".to_owned(), json!(position));
if let Some(name) = name {
row.insert("name".to_owned(), json!(name));
}
if let Some((group, member)) = group {
row.insert("group".to_owned(), json!(group));
row.insert("member".to_owned(), json!(member));
}
Value::Object(row)
}
Self::Pattern { entry, name } => json!({"pattern": entry, "name": name}),
Self::Raw(value) => json!({"raw": value}),
Self::RawPositions(value) => json!({"raw_positions": value}),
}
}
}
/// An alias `model_list` lists: its name and the group it targets (`hidden` aliases already left
/// out by the host when the whole list is asked for).
#[derive(Clone, Debug, PartialEq, serde::Deserialize)]
pub struct AliasItem {
pub alias: String,
pub target: String,
}
/// A routing group callable by its name (the host leaves out one a deployment or an alias
/// shadows): its name and member model groups.
#[derive(Clone, Debug, PartialEq, serde::Deserialize)]
pub struct CallableGroup {
pub name: String,
pub members: Vec<String>,
}
/// What `get_model_list` composes beyond the deployments: the aliases and callable routing groups
/// it lists, and the wildcard requests a name no deployment serves is routed by.
#[derive(Clone, Debug, Default, PartialEq, serde::Deserialize)]
#[serde(default)]
pub struct ListScope {
pub aliases: Vec<AliasItem>,
pub groups: Vec<CallableGroup>,
/// `get_deployments_by_pattern`'s requests: the name, then the name under its provider.
pub wildcard: Vec<String>,
/// Only one part of the list: `"aliases"` (`_core_get_model_list_from_model_alias`) or
/// `"groups"` (`_core_get_model_list_from_routing_groups`).
pub only: Option<String>,
}
/// A row read for a lookup's filters: the deployment's fields as the row shows them.
struct View<'a> {
name: &'a str,
info: &'a Map<String, Value>,
params: &'a Map<String, Value>,
strips_access_groups: bool,
}
/// The lookups over one registry version.
pub struct Lookup<'a> {
deployments: &'a [Arc<Deployment>],
host: &'a HostIndexes,
}
impl<'a> Lookup<'a> {
pub fn new(registry: &'a Registry) -> Self {
Self {
deployments: registry.deployments_from(0),
host: registry.host_indexes(),
}
}
/// `list[value]` as Python reads it: an int (or bool) in range, negative from the end.
fn position(&self, value: &Value) -> Option<usize> {
let len = i64::try_from(self.deployments.len()).ok()?;
let index = match value {
Value::Bool(flag) => i64::from(*flag),
Value::Number(number) => number.as_i64()?,
_ => return None,
};
let index = if index < 0 { len + index } else { index };
(0..len)
.contains(&index)
.then(|| usize::try_from(index).ok())
.flatten()
}
/// Every position `held` (what an index holds) names, read before any is filtered, as the
/// Python router reads them into a tuple first: `Err` with the row that raises.
fn positions(&self, held: &Value) -> Result<Vec<usize>, Row> {
let items = match held {
Value::Array(items) => items,
// Iterating an empty string or dict reads no position; anything else raises.
Value::String(text) if text.is_empty() => return Ok(Vec::new()),
Value::Object(map) if map.is_empty() => return Ok(Vec::new()),
_ => return Err(Row::RawPositions(held.clone())),
};
items
.iter()
.map(|item| self.position(item).ok_or_else(|| Row::Raw(item.clone())))
.collect()
}
fn view(&self, position: usize) -> View<'a> {
let deployment = &self.deployments[position];
View {
name: &deployment.model_name,
info: &deployment.model_info.raw,
params: &deployment.raw_params,
strips_access_groups: false,
}
}
fn row_view<'r>(&self, row: &'r Row) -> Option<View<'r>>
where
'a: 'r,
{
match row {
Row::At {
position,
name,
group,
} => {
let base = self.view(*position);
Some(View {
name: name.as_deref().unwrap_or(base.name),
strips_access_groups: group.is_some(),
..base
})
}
Row::Pattern { entry, name } => Some(View {
name,
info: entry
.get("model_info")
.and_then(Value::as_object)
.unwrap_or(empty()),
params: entry
.get("litellm_params")
.and_then(Value::as_object)
.unwrap_or(empty()),
strips_access_groups: false,
}),
Row::Raw(_) | Row::RawPositions(_) => None,
}
}
// -- deployments by name -------------------------------------------------------------------
/// `_core_get_all_deployments`: the team's deployments under public name `model_name`, else
/// the group's, else (for a team) every deployment the name includes; each under `rename`.
pub fn all(&self, model_name: &str, team_id: Option<&str>, rename: Option<&str>) -> Vec<Row> {
let rows = |positions: Vec<usize>| -> Vec<Row> {
positions
.into_iter()
.filter(|position| should_include(model_name, &self.view(*position), team_id))
.map(|position| Row::At {
position,
name: rename.map(str::to_owned),
group: None,
})
.collect()
};
if let Some(team) = team_id
&& let Some(held) = self.host.team_named(team, model_name)
{
match self.positions(held) {
Ok(positions) => {
let found = rows(positions);
if !found.is_empty() {
return found;
}
}
Err(raw) => return vec![raw],
}
}
if let Some(held) = self.host.named(model_name) {
return match self.positions(held) {
Ok(positions) => rows(positions),
Err(raw) => vec![raw],
};
}
if team_id.is_some() {
return rows((0..self.deployments.len()).collect());
}
Vec::new()
}
/// `_team_deployments_across_teams`: every team's deployments under public name `model`, and
/// how many distinct teams own them.
pub fn across_teams(&self, model: &str) -> (Vec<Row>, usize) {
let mut rows = Vec::new();
for (_, public, held) in &self.host.team_model {
if public != model {
continue;
}
match self.positions(held) {
Ok(positions) => rows.extend(positions.into_iter().map(|position| Row::At {
position,
name: None,
group: None,
})),
Err(raw) => return (vec![raw], 0),
}
}
let mut teams: Vec<&Value> = Vec::new();
for row in &rows {
if let Some(view) = self.row_view(row)
&& let Some(team) = view.info.get("team_id").filter(|team| !team.is_null())
&& !teams.contains(&team)
{
teams.push(team);
}
}
let count = teams.len();
(rows, count)
}
/// `_core_config_deployments`: the deployments a config declared (no `db_model: true`).
pub fn config(&self) -> Vec<Row> {
(0..self.deployments.len())
.filter(|position| {
self.view(*position).info.get("db_model") != Some(&Value::Bool(true))
})
.map(|position| Row::At {
position,
name: None,
group: None,
})
.collect()
}
/// `_core_get_model_ids`: the ids of the group's deployments, or (no group) of the id map's
/// entries, without team deployments when `exclude_team_models`. `Err` with the row that raises.
pub fn model_ids(
&self,
model_name: Option<&str>,
exclude_team_models: bool,
) -> Result<Vec<String>, Row> {
let kept = |view: &View<'_>| {
view.info.contains_key("id")
&& !(exclude_team_models && view.info.get("team_id").is_some_and(truthy))
};
if let Some(name) = model_name {
let Some(held) = self.host.named(name) else {
return Ok(Vec::new());
};
return Ok(self
.positions(held)?
.into_iter()
.map(|position| self.view(position))
.filter(kept)
.map(|view| python_str(&view.info["id"]))
.collect());
}
let mut ids = Vec::new();
for (id, held) in &self.host.id_map {
let position = self.position(held).ok_or_else(|| Row::Raw(held.clone()))?;
if kept(&self.view(position)) {
ids.push(id.clone());
}
}
Ok(ids)
}
// -- the model list --------------------------------------------------------------------------
/// `_core_get_model_list`: the name's deployments, its alias rows and routing group rows, its
/// wildcard routes when nothing else serves it, or (no name) every alias and group row and
/// then every deployment.
pub fn model_list(
&self,
model_name: Option<&str>,
team_id: Option<&str>,
scope: &ListScope,
) -> Vec<Row> {
let only = scope.only.as_deref();
let mut rows = Vec::new();
if let Some(name) = model_name
&& only.is_none()
{
rows.extend(self.all(name, team_id, None));
}
let aliases: &[AliasItem] = if only == Some("groups") {
&[]
} else {
&scope.aliases
};
for alias in aliases {
match scope.groups.iter().find(|group| group.name == alias.target) {
Some(group) => {
rows.extend(self.group_rows(group).into_iter().map(|row| match row {
Row::At {
position, group, ..
} => Row::At {
position,
name: Some(alias.alias.clone()),
group,
},
other => other,
}))
}
None => rows.extend(self.all(&alias.target, None, Some(&alias.alias))),
}
}
match (model_name, only) {
(_, Some("aliases")) => {}
(Some(name), _) => {
if let Some(group) = scope.groups.iter().find(|group| group.name == name) {
rows.extend(self.group_rows(group));
}
}
(None, _) => {
for group in &scope.groups {
rows.extend(self.group_rows(group));
}
}
}
if only.is_some() {
return rows;
}
if let Some(name) = model_name
&& rows.is_empty()
{
rows.extend(self.wildcard(name, team_id, &scope.wildcard));
}
if model_name.is_none() {
rows.extend((0..self.deployments.len()).map(|position| Row::At {
position,
name: None,
group: None,
}));
}
rows
}
/// The whole list's rows, or the first row the host's read of the list raises at.
fn readable(&self, rows: Vec<Row>) -> Result<Vec<Row>, Row> {
match rows
.iter()
.find(|row| matches!(row, Row::Raw(_) | Row::RawPositions(_)))
{
Some(raw) => Err(raw.clone()),
None => Ok(rows),
}
}
/// A callable routing group's rows: each member's deployments under the group's name.
fn group_rows(&self, group: &CallableGroup) -> Vec<Row> {
let mut rows = Vec::new();
for member in &group.members {
for row in self.all(member, None, Some(&group.name)) {
rows.push(match row {
Row::At { position, name, .. } => Row::At {
position,
name,
group: Some((group.name.clone(), member.clone())),
},
other => other,
});
}
}
rows
}
/// `_get_wildcard_deployments`: the global router's routes for the requests, then the team's.
fn wildcard(&self, name: &str, team_id: Option<&str>, requests: &[String]) -> Vec<Row> {
let route = |router: &crate::host_index::PatternRouter| -> Vec<Value> {
requests
.iter()
.find_map(|request| {
router
.route(request, None)
.filter(|routed| !routed.is_empty())
})
.unwrap_or_default()
};
let mut routed = route(&self.host.patterns);
if let Some(team) = team_id
&& let Some((_, router)) = self.host.team_patterns.iter().find(|(t, _)| t == team)
{
routed.extend(route(router));
}
routed
.into_iter()
.map(|entry| Row::Pattern {
entry,
name: name.to_owned(),
})
.collect()
}
/// The deployments `team_id`'s own wildcard pattern routes `name` to
/// (`_core_team_pattern_routers`, keyed by `team_public_model_name`).
pub fn team_wildcard(&self, name: &str, team_id: &str) -> Vec<Arc<Deployment>> {
let Some((_, router)) = self.host.team_patterns.iter().find(|(t, _)| t == team_id) else {
return Vec::new();
};
router
.route(name, None)
.unwrap_or_default()
.iter()
.filter_map(|entry| entry.get("model_info")?.get("id").map(python_str))
.filter_map(|id| self.deployments.iter().find(|d| d.id == id).cloned())
.collect()
}
/// `_core_get_model_names`: per row of the whole list, a team deployment's public name for
/// its own team, else the row's name.
pub fn model_names(
&self,
team_id: Option<&str>,
scope: &ListScope,
) -> Result<Vec<String>, Row> {
let rows = self.readable(self.model_list(None, None, scope))?;
let mut names = Vec::new();
for row in &rows {
let Some(view) = self.row_view(row) else {
continue;
};
match view.info.get("team_id").filter(|team| truthy(team)) {
Some(team) => {
if team_id.is_some_and(|id| team == &Value::String(id.to_owned()))
&& let Some(public) = view
.info
.get("team_public_model_name")
.filter(|public| truthy(public))
{
names.push(python_str(public));
}
}
None => names.push(view.name.to_owned()),
}
}
Ok(names)
}
/// `_core_get_model_access_groups`: access group to the names of the rows that declare it,
/// in the order they appear.
pub fn access_groups(
&self,
model_name: Option<&str>,
access_group: Option<&str>,
team_id: Option<&str>,
scope: &ListScope,
) -> Result<Vec<(String, Vec<String>)>, Row> {
let mut groups: Vec<(String, Vec<String>)> = Vec::new();
for row in self.readable(self.model_list(model_name, team_id, scope))? {
let Some(view) = self.row_view(&row) else {
continue;
};
if view.info.is_empty() || view.strips_access_groups {
continue;
}
// `access_groups or []`, iterated: a list's items, a string's characters.
let declared: Vec<String> = match view.info.get("access_groups") {
Some(Value::Array(items)) => items.iter().map(python_str).collect(),
Some(Value::String(text)) => text.chars().map(String::from).collect(),
_ => continue,
};
for group in declared {
if access_group.is_some_and(|wanted| wanted != group) {
continue;
}
match groups.iter_mut().find(|(held, _)| *held == group) {
Some((_, names)) => names.push(view.name.to_owned()),
None => groups.push((group, vec![view.name.to_owned()])),
}
}
}
Ok(groups)
}
/// `_core_resolve_model_name_from_model_id`'s scan of the whole list: whether it lists any
/// row, and the name of the first row whose `litellm_params.model` is `model_id`, or ends in
/// `/model_id` or `:model_id`.
pub fn name_serving(
&self,
model_id: &str,
scope: &ListScope,
) -> Result<(bool, Option<String>), Row> {
let rows = self.readable(self.model_list(None, None, scope))?;
let (slash, colon) = (format!("/{model_id}"), format!(":{model_id}"));
let name = rows.iter().find_map(|row| {
let view = self.row_view(row)?;
let model = view.params.get("model").and_then(Value::as_str)?;
(model == model_id || model.ends_with(&slash) || model.ends_with(&colon))
.then_some(view.name)
.filter(|name| !name.is_empty())
.map(str::to_owned)
});
Ok((!rows.is_empty(), name))
}
// -- teams and credentials -------------------------------------------------------------------
/// `_core_team_model_has_alternatives`: whether the team's public name of deployment `id` has
/// another unpaused deployment of that team.
pub fn team_model_has_alternatives(&self, id: &str) -> Result<bool, Row> {
let Some(held) = self.host.id_entry(id) else {
return Ok(false);
};
let position = self.position(held).ok_or_else(|| Row::Raw(held.clone()))?;
let info = self.view(position).info;
let (Some(Value::String(team)), Some(Value::String(public))) =
(info.get("team_id"), info.get("team_public_model_name"))
else {
return Ok(false);
};
let Some(siblings) = self.host.team_named(team, public) else {
return Ok(false);
};
let unpaused = self
.positions(siblings)?
.into_iter()
.filter(|position| self.view(*position).info.get("blocked") != Some(&Value::Bool(true)))
.count();
Ok(unpaused > 1)
}
/// `_get_model_group_deployment_usable_by_team`: the first deployment of the group the team
/// may use.
pub fn group_deployment_usable_by_team(
&self,
model_group: &str,
team_id: Option<&str>,
) -> Result<Option<usize>, Row> {
// `.get(name) or ()`: nothing held, or a falsy value, reads no position.
let Some(held) = self.host.named(model_group).filter(|held| truthy(held)) else {
return Ok(None);
};
let Value::Array(items) = held else {
return Err(Row::RawPositions(held.clone()));
};
for item in items {
let position = self.position(item).ok_or_else(|| Row::Raw(item.clone()))?;
if usable_by_team(&self.view(position), team_id) {
return Ok(Some(position));
}
}
Ok(None)
}
/// `_get_team_public_name_deployment`: the first deployment of the team's public name.
pub fn team_public_name_deployment(
&self,
model_id: &str,
team_id: &str,
) -> Result<Option<usize>, Row> {
match self.host.team_named(team_id, model_id) {
Some(Value::Array(items)) => match items.first() {
Some(item) => self
.position(item)
.map(Some)
.ok_or_else(|| Row::Raw(item.clone())),
None => Ok(None),
},
Some(held) if truthy(held) => Err(Row::RawPositions(held.clone())),
Some(_) | None => Ok(None),
}
}
// -- routable deployments --------------------------------------------------------------------
/// `_routable_deployments`: the routable group's deployments for the team (those it may use
/// when any), without paused ones.
pub fn routable(&self, model_group: &str, team_id: Option<&str>) -> Vec<Row> {
let named = self.all(model_group, team_id, None);
if named
.iter()
.any(|row| matches!(row, Row::Raw(_) | Row::RawPositions(_)))
{
return named;
}
let usable: Vec<&Row> = named
.iter()
.filter(|row| {
self.row_view(row)
.is_some_and(|view| usable_by_team(&view, team_id))
})
.collect();
let pool: Vec<&Row> = if usable.is_empty() {
named.iter().collect()
} else {
usable
};
pool.into_iter()
.filter(|row| {
self.row_view(row)
.is_some_and(|view| view.info.get("blocked") != Some(&Value::Bool(true)))
})
.cloned()
.collect()
}
/// `_core_get_candidate_model_ids_for_route`'s ids of `rows`, in first-seen order.
pub fn ids_of(&self, rows: &[Row]) -> Result<Vec<String>, Row> {
let mut ids: Vec<String> = Vec::new();
for row in rows {
let view = match row {
Row::Raw(_) | Row::RawPositions(_) => return Err(row.clone()),
_ => self.row_view(row).expect("a readable row"),
};
if let Some(id) = view.info.get("id").filter(|id| !id.is_null()) {
let id = python_str(id);
if !ids.contains(&id) {
ids.push(id);
}
}
}
Ok(ids)
}
/// `_core_get_candidate_model_ids_for_route` once the alias resolved `resolved`: the
/// members' ids when it names a callable routing group (`members`), else, for a name that is
/// no model group, the team's deployments under it or its wildcard routes, else the group's.
pub fn candidate_ids(
&self,
resolved: &str,
team_id: Option<&str>,
members: Option<&[String]>,
requests: &[String],
) -> Result<Vec<String>, Row> {
if let Some(members) = members {
let mut ids: Vec<String> = Vec::new();
for member in members {
for id in self.ids_of(&self.all(member, team_id, None))? {
if !ids.contains(&id) {
ids.push(id);
}
}
}
return Ok(ids);
}
let named = Value::String(resolved.to_owned());
if !self.host.model_names.contains(&named) {
if team_id.is_some() {
let team_rows = self.all(resolved, team_id, None);
if !team_rows.is_empty() {
return self.ids_of(&team_rows);
}
}
let wildcard = self.wildcard_route(resolved, team_id, requests);
if !wildcard.is_empty() {
return self.ids_of(&wildcard);
}
}
self.ids_of(&self.all(resolved, team_id, None))
}
/// `_core_wildcard_route_deployments`: the team router's routes for `model`, else the global
/// router's.
pub fn wildcard_route(
&self,
model: &str,
team_id: Option<&str>,
requests: &[String],
) -> Vec<Row> {
let route = |router: &crate::host_index::PatternRouter| -> Vec<Value> {
requests
.iter()
.find_map(|request| {
router
.route(request, None)
.filter(|routed| !routed.is_empty())
})
.unwrap_or_default()
};
let team_owned = team_id
.and_then(|team| self.host.team_patterns.iter().find(|(t, _)| t == team))
.map(|(_, router)| route(router))
.unwrap_or_default();
let routed = if team_owned.is_empty() {
route(&self.host.patterns)
} else {
team_owned
};
routed
.into_iter()
.map(|entry| Row::Pattern {
entry,
name: model.to_owned(),
})
.collect()
}
}
/// `_core_should_include_deployment`: the team's deployment under its public name, or a
/// deployment of the group the team may use.
fn should_include(model_name: &str, view: &View<'_>, team_id: Option<&str>) -> bool {
let team_value = team_id.map(|team| Value::String(team.to_owned()));
let model_team = view.info.get("team_id");
if let Some(team) = &team_value
&& model_team == Some(team)
&& view.info.get("team_public_model_name") == Some(&Value::String(model_name.to_owned()))
{
return true;
}
if view.name == model_name {
return team_value.is_none()
|| model_team.is_none_or(Value::is_null)
|| model_team == team_value.as_ref();
}
false
}
/// `deployment_usable_by_team`: a deployment no team owns, or one the team owns.
fn usable_by_team(view: &View<'_>, team_id: Option<&str>) -> bool {
match view.info.get("team_id").filter(|team| !team.is_null()) {
None => true,
Some(owner) => team_id.is_some_and(|team| owner == &Value::String(team.to_owned())),
}
}
fn empty() -> &'static Map<String, Value> {
static EMPTY: std::sync::OnceLock<Map<String, Value>> = std::sync::OnceLock::new();
EMPTY.get_or_init(Map::new)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
use crate::registry::DeploymentSpec;
fn registry(entries: Value) -> Registry {
let mut registry = Registry::default();
for entry in entries.as_array().unwrap() {
let spec: DeploymentSpec = serde_json::from_value(entry.clone()).unwrap();
registry.push_built(spec, false).unwrap();
}
registry
}
#[test]
fn a_teams_wildcard_routes_only_the_names_its_public_pattern_matches() {
let registry = registry(json!([
{"model_name": "t1-w", "litellm_params": {"model": "openai/*"},
"model_info": {"id": "tw1", "team_id": "T1", "team_public_model_name": "tw-*"}},
{"model_name": "g", "litellm_params": {"model": "openai/g"}, "model_info": {"id": "gp"}},
]));
let lookup = Lookup::new(&registry);
let ids = |name: &str, team: &str| -> Vec<String> {
let routed = lookup.team_wildcard(name, team);
routed.iter().map(|d| d.id.clone()).collect()
};
assert_eq!(ids("tw-gpt-4o", "T1"), ["tw1"]);
assert!(ids("gpt-4o", "T1").is_empty());
assert!(ids("tw-gpt-4o", "T2").is_empty());
}
fn entry(name: &str, id: &str, info: Value) -> Value {
let mut info = info.as_object().cloned().unwrap_or_default();
info.insert("id".to_owned(), json!(id));
json!({"model_name": name, "litellm_params": {"model": format!("openai/{id}")}, "model_info": info})
}
fn at(position: usize) -> Row {
Row::At {
position,
name: None,
group: None,
}
}
#[test]
fn all_prefers_the_teams_public_name_then_the_group() {
let registry = registry(json!([
entry("gpt", "a", json!({})),
entry(
"team-gpt",
"t",
json!({"team_id": "T", "team_public_model_name": "gpt"})
),
entry("gpt", "b", json!({"team_id": "U"})),
]));
let lookup = Lookup::new(&registry);
assert_eq!(lookup.all("gpt", Some("T"), None), vec![at(1)]);
assert_eq!(lookup.all("gpt", Some("V"), None), vec![at(0)]);
assert_eq!(lookup.all("gpt", None, None), vec![at(0), at(2)]);
assert_eq!(lookup.all("missing", None, None), vec![]);
}
#[test]
fn a_junk_position_is_handed_back_to_raise() {
let mut registry = registry(json!([entry("gpt", "a", json!({}))]));
let mut model_name = registry.host_indexes().model_name.clone();
model_name.insert("gpt".to_owned(), json!([0, "x"]));
let changes = crate::registry::ChangeSet {
host_index: Some(crate::host_index::IndexWrite {
model_name: Some(model_name),
..Default::default()
}),
..Default::default()
};
registry = registry.apply(&changes, None).unwrap();
let lookup = Lookup::new(&registry);
assert_eq!(lookup.all("gpt", None, None), vec![Row::Raw(json!("x"))]);
}
#[test]
fn the_model_list_lists_aliases_and_groups_then_every_deployment() {
let registry = registry(json!([
entry("gpt", "a", json!({"access_groups": ["g1"]})),
entry("claude", "b", json!({})),
]));
let lookup = Lookup::new(&registry);
let scope = ListScope {
aliases: vec![AliasItem {
alias: "fast".to_owned(),
target: "gpt".to_owned(),
}],
groups: vec![CallableGroup {
name: "pool".to_owned(),
members: vec!["gpt".to_owned(), "claude".to_owned()],
}],
wildcard: Vec::new(),
only: None,
};
let rows = lookup.model_list(None, None, &scope);
assert_eq!(rows.len(), 5);
assert_eq!(
rows[0],
Row::At {
position: 0,
name: Some("fast".to_owned()),
group: None
}
);
assert_eq!(
lookup.model_names(None, &scope).unwrap(),
["fast", "pool", "pool", "gpt", "claude"]
);
assert_eq!(
lookup.access_groups(None, None, None, &scope).unwrap(),
vec![("g1".to_owned(), vec!["fast".to_owned(), "gpt".to_owned()])]
);
}
}

View file

@ -0,0 +1,44 @@
//! The request's `mock_testing_*` flags: a failure the router raises itself, without calling a
//! deployment (`_handle_mock_testing_fallbacks` and `_handle_mock_testing_rate_limit_error`).
use serde::Deserialize;
use crate::error::{AttemptError, FailureKind};
/// Which mock failure a request asks for.
#[derive(
Clone, Copy, Debug, PartialEq, Eq, Deserialize, strum::IntoStaticStr, strum::EnumString,
)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum MockFailure {
/// `mock_testing_fallbacks`: an `InternalServerError` before the first level's retries.
Fallbacks,
/// `mock_testing_context_fallbacks`: a `ContextWindowExceededError`, likewise.
ContextWindow,
/// `mock_testing_content_policy_fallbacks`: a `ContentPolicyViolationError`, likewise.
ContentPolicy,
/// `mock_testing_rate_limit_error`: a `RateLimitError` in place of every level's first
/// attempt, which the retry loop handles like a failed attempt.
RateLimit,
}
impl MockFailure {
/// Whether the failure stands in for the first level's retry loop (raised before its retries).
pub fn before_retries(self) -> bool {
!matches!(self, Self::RateLimit)
}
/// The failure as the rules see it.
pub fn failure(self, num_retries: Option<u32>) -> AttemptError {
let (kind, status) = match self {
Self::Fallbacks => (FailureKind::InternalServer, 500),
Self::ContextWindow => (FailureKind::ContextWindowExceeded, 400),
Self::ContentPolicy => (FailureKind::ContentPolicyViolation, 400),
Self::RateLimit => (FailureKind::RateLimit, 429),
};
let mut failure = AttemptError::new(kind, Some(status));
failure.num_retries = num_retries;
failure
}
}

View file

@ -0,0 +1,716 @@
//! Model info: what the router reports a deployment can do and costs, merged from the process
//! cost map, what the background refresh discovered, and the deployment's own `model_info`.
//!
//! The merge rules are the Python router core's (`get_router_model_info`,
//! `get_deployment_model_info`); every cost-map read is a [`Host::model_info`] lookup asked while
//! the query runs, so a price reload or a registration the host made since the router was built is
//! what the next query reads. The host keeps choosing the keys a lookup tries (`get_model_info`);
//! the core chooses which lookups a query makes.
//!
//! Values the host cannot spell as JSON travel as opaque tokens ([`is_opaque`]): the merge moves
//! them between layers without reading them, and the host restores them.
use serde_json::{Map, Number, Value};
use crate::{
error::Error,
host::{Host, LogLevel, ModelInfoLookup},
};
/// The key the host wraps a value it cannot spell as JSON in.
pub const OPAQUE_KEY: &str = "__litellm_python_value__";
/// A value the host passed as an opaque token: the merge moves it, never reads it.
pub fn is_opaque(value: &Value) -> bool {
matches!(value, Value::Object(map) if map.len() == 1 && map.contains_key(OPAQUE_KEY))
}
/// One `get_router_model_info` query, as the host reads it off the deployment.
#[derive(Clone, Debug, Default, serde::Deserialize)]
#[serde(default)]
pub struct RouterModelInfoQuery {
/// `model_info.id`.
pub deployment_id: Value,
/// `model_info.base_model`, then `litellm_params.base_model`: the first that is not `None`
/// names the model whose limits and costs the deployment reports.
pub model_info_base_model: Value,
pub params_base_model: Value,
/// The validated `litellm_params.model`.
pub params_model: String,
/// `base_model` among the validated params' extra fields.
pub extra_base_model: Value,
/// The name the caller asked for, which a wildcard deployment matches.
pub received_model_name: String,
}
/// A merged model info, and the keys whose values are the deployment's own (the last layer).
#[derive(Clone, Debug, PartialEq)]
pub struct MergedModelInfo {
pub info: Map<String, Value>,
pub user_keys: Vec<String>,
}
pub(crate) fn raised(exception: &'static str, detail: impl Into<String>) -> Error {
Error::ModelInfo {
exception,
detail: detail.into(),
}
}
pub(crate) fn malformed(lookup: &str) -> Error {
raised(
"ValueError",
format!("the host answered a `{lookup}` lookup malformed"),
)
}
/// Python truthiness of a JSON value.
pub(crate) fn truthy(value: &Value) -> bool {
match value {
Value::Null => false,
Value::Bool(flag) => *flag,
Value::Number(number) => number.as_f64().is_some_and(|n| n != 0.0),
Value::String(text) => !text.is_empty(),
Value::Array(items) => !items.is_empty(),
Value::Object(map) => is_opaque(value) || !map.is_empty(),
}
}
/// `(value or 0) > 0`, for a value that is a number (anything else reads as not positive).
fn positive_or_zero(value: Option<&Value>) -> bool {
value
.filter(|value| truthy(value))
.and_then(Value::as_f64)
.is_some_and(|n| n > 0.0)
}
pub(crate) async fn lookup(host: &dyn Host, lookup: ModelInfoLookup) -> Result<Value, Error> {
host.model_info(&lookup).await
}
/// `get_model_info(model=...)`: the row, or `None` when the lookup tolerates a failure and failed.
async fn catalog_row(
host: &dyn Host,
model: Value,
tolerate: bool,
) -> Result<Option<Map<String, Value>>, Error> {
match lookup(
host,
ModelInfoLookup::ModelInfo {
model,
tolerate,
exact: false,
},
)
.await?
{
Value::Null => Ok(None),
Value::Object(row) if tolerate && row.get("failed") == Some(&Value::Bool(true)) => Ok(None),
Value::Object(row) => Ok(Some(row)),
_ => Err(malformed("model_info")),
}
}
/// `model_cost.get(key)`: `None` when the key is absent, else the entry as stored.
async fn cost_entry(host: &dyn Host, key: &str, process: bool) -> Result<Option<Value>, Error> {
match lookup(
host,
ModelInfoLookup::CostEntry {
key: key.to_owned(),
process,
},
)
.await?
{
Value::Null => Ok(None),
Value::Object(mut answer) => Ok(Some(answer.remove("entry").unwrap_or(Value::Null))),
_ => Err(malformed("cost_entry")),
}
}
pub(crate) async fn discovered(host: &dyn Host, id: Value) -> Result<Map<String, Value>, Error> {
match lookup(host, ModelInfoLookup::Discovered { id }).await? {
Value::Null => Ok(Map::new()),
Value::Object(limits) => Ok(limits),
_ => Err(malformed("discovered")),
}
}
async fn log(host: &dyn Host, level: LogLevel, message: String) -> Result<(), Error> {
lookup(host, ModelInfoLookup::Log { level, message })
.await
.map(|_| ())
}
/// `_provider_model_id`: the model id a cost-map key or base model names, bedrock base model ids
/// normalized.
async fn provider_model_id(host: &dyn Host, name: &str, provider: &str) -> Result<String, Error> {
let model_id = name.rsplit('/').next().unwrap_or(name).to_owned();
if !provider.starts_with("bedrock") {
return Ok(model_id);
}
match lookup(host, ModelInfoLookup::BedrockBase { model_id }).await? {
Value::String(base) => Ok(base),
_ => Err(malformed("bedrock_base")),
}
}
/// `_exact_base_model_max_input_tokens`: the base model's own window, unless its row came from a
/// fallback-generalization rule or it is unmapped.
async fn exact_base_model_max_input_tokens(
host: &dyn Host,
base_model: &str,
) -> Result<Option<Value>, Error> {
let answer = lookup(
host,
ModelInfoLookup::ModelInfo {
model: Value::String(base_model.to_owned()),
tolerate: true,
exact: true,
},
)
.await?;
let Value::Object(mut answer) = answer else {
return Ok(None);
};
if answer.get("failed") == Some(&Value::Bool(true))
|| answer.get("generalized") == Some(&Value::Bool(true))
{
return Ok(None);
}
let Some(Value::Object(row)) = answer.remove("row") else {
return Ok(None);
};
Ok(row
.get("max_input_tokens")
.filter(|tokens| !tokens.is_null())
.cloned())
}
/// `_base_model_widened_window`: a base model naming the same model as the route can only widen
/// its context window.
async fn base_model_widened_window(
host: &dyn Host,
route: &Map<String, Value>,
base_model: Option<&str>,
provider: &str,
) -> Result<Option<Value>, Error> {
let Some(route_max) = route
.get("max_input_tokens")
.filter(|tokens| !tokens.is_null())
else {
return Ok(None);
};
let Some(base_model) = base_model else {
return Ok(None);
};
let route_key = match route.get("key") {
Some(Value::String(key)) => key.clone(),
Some(_) => return Err(raised("AttributeError", "key")),
None => return Err(raised("KeyError", "key")),
};
let route_id = provider_model_id(host, &route_key, provider).await?;
let base_id = provider_model_id(host, base_model, provider).await?;
let unversioned = base_id
.split_once('@')
.map_or(base_id.as_str(), |(id, _)| id);
if route_id != base_id && route_id != unversioned {
return Ok(None);
}
let Some(base_max) = exact_base_model_max_input_tokens(host, base_model).await? else {
return Ok(None);
};
match (base_max.as_f64(), route_max.as_f64()) {
(Some(base), Some(route)) if base > route => Ok(Some(base_max)),
(Some(_), Some(_)) => Ok(None),
_ => Err(raised(
"TypeError",
"'<=' not supported between the base model's and the route's max_input_tokens",
)),
}
}
/// `_declared_base_model`: a `base_model` that is a string.
fn declared_base_model(value: &Value) -> Option<&str> {
value.as_str()
}
/// `get_router_model_info` after the host validated the deployment's params: the cost-map row of
/// the model the deployment serves, widened by its base model, overlaid with what the background
/// refresh discovered and then with every value the deployment sets itself. `None` when the
/// lookup answers `None`; an unmapped model raises as the lookup raises.
pub async fn router_model_info(
host: &dyn Host,
query: &RouterModelInfoQuery,
) -> Result<Option<MergedModelInfo>, Error> {
let base_model = if query.model_info_base_model.is_null() {
query.params_base_model.clone()
} else {
query.model_info_base_model.clone()
};
let mut model = base_model.clone();
let (resolved_model, provider) = match lookup(
host,
ModelInfoLookup::Provider {
model: query.params_model.clone(),
},
)
.await?
{
Value::Array(pair) => match pair.as_slice() {
[Value::String(model), Value::String(provider), ..] => {
(model.clone(), provider.clone())
}
_ => return Err(malformed("provider")),
},
_ => return Err(malformed("provider")),
};
if provider == "azure" && base_model.is_null() {
// Router init registers every deployment name into the cost map as a zeroed stub, so
// membership alone cannot tell a resolvable name apart; it needs usable limits or costs.
let fallback_key = if resolved_model.starts_with("azure/") {
resolved_model.clone()
} else {
format!("azure/{resolved_model}")
};
let resolves = match cost_entry(host, &fallback_key, false).await? {
Some(Value::Object(entry)) if !is_opaque(&Value::Object(entry.clone())) => {
positive_or_zero(entry.get("max_input_tokens"))
|| positive_or_zero(entry.get("input_cost_per_token"))
}
_ => false,
};
if resolves {
log(
host,
LogLevel::Debug,
format!(
"Azure deployment '{resolved_model}' has no base_model set; using '{fallback_key}' from the model cost map for max tokens, cost tracking, etc."
),
)
.await?;
} else {
log(
host,
LogLevel::Error,
format!(
"Could not identify azure model '{resolved_model}'. Set azure 'base_model' for accurate max tokens, cost tracking, etc.- https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models"
),
)
.await?;
}
} else if provider != "azure" {
model = Value::String(resolved_model.clone());
if resolved_model.contains('*') {
// Only a wildcard deployment asks which model its pattern routed the name to.
if let Value::Array(routed) = lookup(
host,
ModelInfoLookup::PatternModels {
name: query.received_model_name.clone(),
},
)
.await?
&& let Some(routed) = routed
.iter()
.find(|routed| routed.get("id").unwrap_or(&Value::Null) == &query.deployment_id)
{
model = routed.get("model").cloned().unwrap_or(Value::Null);
}
}
}
if model.is_null() {
// An azure deployment without a base_model is looked up by its own model.
model = Value::String(resolved_model);
}
let Value::String(model) = model else {
return Err(raised(
"AttributeError",
format!(
"'{}' object has no attribute 'startswith'",
crate::group_info::type_name(&model)
),
));
};
let model_info_name = if model.starts_with(&format!("{provider}/")) {
model
} else {
format!("{provider}/{model}")
};
let Some(row) = catalog_row(host, Value::String(model_info_name), false).await? else {
return Ok(None);
};
let user_model_info = match lookup(host, ModelInfoLookup::UserModelInfo).await? {
Value::Object(info) => info,
Value::Null => Map::new(),
_ => return Err(malformed("user_model_info")),
};
let declared = user_model_info
.get("base_model")
.and_then(declared_base_model)
.or_else(|| declared_base_model(&query.extra_base_model));
let widened = base_model_widened_window(host, &row, declared, &provider).await?;
let discovered = discovered(host, query.deployment_id.clone()).await?;
let mut info = row;
if let Some(max_input_tokens) = widened {
info.insert("max_input_tokens".to_owned(), max_input_tokens);
}
info.extend(discovered);
let mut user_keys = Vec::new();
for (key, value) in &user_model_info {
if !value.is_null() {
info.insert(key.clone(), value.clone());
user_keys.push(key.clone());
}
}
Ok(Some(MergedModelInfo { info, user_keys }))
}
/// `_convert_stringified_numbers`: a string that reads as a number becomes one (a float when it
/// has an exponent or a point, else an int); any other string stays as written.
fn convert_stringified_number(text: &str) -> Value {
let trimmed = text.trim();
let digits = without_digit_underscores(trimmed);
let lower = text.to_lowercase();
if lower.contains('e') || text.contains('.') {
return digits
.and_then(|digits| python_float(&digits))
.and_then(Number::from_f64)
.map_or_else(|| Value::String(text.to_owned()), Value::Number);
}
digits
.and_then(|digits| python_int(&digits))
.map_or_else(|| Value::String(text.to_owned()), Value::Number)
}
/// The text with the underscores Python's `int()`/`float()` accept between digits removed, or
/// `None` when an underscore sits anywhere else.
pub(crate) fn without_digit_underscores(text: &str) -> Option<String> {
let chars: Vec<char> = text.chars().collect();
let mut out = String::with_capacity(text.len());
for (index, ch) in chars.iter().enumerate() {
if *ch == '_' {
let before = index.checked_sub(1).and_then(|i| chars.get(i));
let after = chars.get(index + 1);
if !(before.is_some_and(char::is_ascii_digit)
&& after.is_some_and(char::is_ascii_digit))
{
return None;
}
continue;
}
out.push(*ch);
}
Some(out)
}
pub(crate) fn python_float(text: &str) -> Option<f64> {
let lower = text.to_ascii_lowercase();
let unsigned = lower.trim_start_matches(['+', '-']);
if matches!(unsigned, "inf" | "infinity" | "nan") {
// Not finite: JSON cannot carry it, so the value stays as written.
return None;
}
let well_formed = text
.chars()
.all(|ch| ch.is_ascii_digit() || matches!(ch, '+' | '-' | '.' | 'e' | 'E'));
if !well_formed {
return None;
}
text.parse::<f64>().ok().filter(|n| n.is_finite())
}
pub(crate) fn python_int(text: &str) -> Option<Number> {
let unsigned = text.strip_prefix(['+', '-']).unwrap_or(text);
if unsigned.is_empty() || !unsigned.chars().all(|ch| ch.is_ascii_digit()) {
return None;
}
let text = text.strip_prefix('+').unwrap_or(text);
text.parse::<i64>()
.map(Number::from)
.ok()
.or_else(|| text.parse::<u64>().map(Number::from).ok())
}
/// `update_dictionary`: `new`'s values over `existing`, skipping `None`s; a string that reads as a
/// number becomes one, and a dict merges one level into the dict it replaces.
pub fn update_dictionary(existing: &mut Map<String, Value>, new: &Map<String, Value>) {
for (key, value) in new {
match value {
Value::Null => {}
Value::String(text) => {
existing.insert(key.clone(), convert_stringified_number(text));
}
Value::Object(nested) if !is_opaque(value) => {
let merged = match existing.get(key) {
Some(current @ Value::Object(current_map)) if !is_opaque(current) => {
let mut merged = current_map.clone();
merged.extend(nested.clone());
merged
}
_ => nested.clone(),
};
existing.insert(key.clone(), Value::Object(merged));
}
_ => {
existing.insert(key.clone(), value.clone());
}
}
}
}
/// `cost_map_omits_token_price`: the cost-map entries under `keys` exist, and none carries a
/// per-token price.
async fn cost_map_omits_token_price(host: &dyn Host, keys: &[Option<&str>]) -> Result<bool, Error> {
let mut entries = 0;
for key in keys.iter().flatten() {
if let Some(entry) = cost_entry(host, key, true).await?
&& !is_opaque(&entry)
&& let Value::Object(map) = entry
{
entries += 1;
if map.contains_key("input_cost_per_token") || map.contains_key("output_cost_per_token")
{
return Ok(false);
}
}
}
Ok(entries > 0)
}
/// `get_deployment_model_info`: deployment `model_id`'s own cost-map entry (with what the refresh
/// discovered, over its base model's row), merged over the row of its litellm model `model_name`.
/// A lookup that fails leaves its layer out, as the Python router ignores it.
pub async fn deployment_model_info(
host: &dyn Host,
model_id: &str,
model_name: &Value,
) -> Result<Option<Map<String, Value>>, Error> {
let mut custom: Option<Map<String, Value>> = None;
if let Some(entry) = cost_entry(host, model_id, false).await? {
let entry = if truthy(&entry) {
match entry {
Value::Object(map) if !is_opaque(&Value::Object(map.clone())) => Some(map),
// Not a mapping: the Python router's TypeError, which it ignores.
_ => None,
}
} else {
Some(Map::new())
};
if let Some(mut entry) = entry {
entry.extend(discovered(host, Value::String(model_id.to_owned())).await?);
custom = Some(entry);
}
}
let builtin = catalog_row(host, model_name.clone(), true).await?;
let mut base_model_key: Option<String> = None;
if let Some(current) = &custom
&& let Some(base_model) = current.get("base_model").filter(|base| !base.is_null())
&& let Some(mut base_row) = catalog_row(host, base_model.clone(), true).await?
{
base_model_key = base_row
.get("key")
.and_then(Value::as_str)
.map(str::to_owned);
// The base model provides defaults; the deployment's own entry overrides them.
update_dictionary(&mut base_row, current);
custom = Some(base_row);
}
let builtin_key = builtin
.as_ref()
.and_then(|row| row.get("key"))
.and_then(Value::as_str)
.map(str::to_owned);
let info = match (custom, builtin) {
(Some(custom), Some(mut builtin)) => {
update_dictionary(&mut builtin, &custom);
builtin
}
(None, Some(builtin)) => builtin,
(Some(custom), None) => custom,
(None, None) => return Ok(None),
};
if cost_map_omits_token_price(
host,
&[
Some(model_id),
builtin_key.as_deref(),
base_model_key.as_deref(),
],
)
.await?
{
let mut info = info;
info.insert("input_cost_per_token".to_owned(), Value::Null);
info.insert("output_cost_per_token".to_owned(), Value::Null);
return Ok(Some(info));
}
Ok(Some(info))
}
#[cfg(test)]
mod tests {
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
};
use async_trait::async_trait;
use serde_json::json;
use super::*;
use crate::{
cache::{InMemoryRouterCache, RouterCache},
error::AttemptError,
host::{Attempt, Output, RouterEvent, RouterRequest, RoutingContext, SelectStage},
registry::Deployment,
};
/// A cost map and a provider, answering lookups the way `PythonHost` does.
#[derive(Default)]
struct Catalog {
rows: BTreeMap<String, Value>,
provider: (String, String),
user_model_info: Value,
lookups: Mutex<Vec<Value>>,
}
#[async_trait]
impl Host for Catalog {
async fn invoke(&self, _attempt: &Attempt) -> Result<Output, AttemptError> {
unreachable!("a model-info query runs no attempt")
}
async fn select(
&self,
_stage: &SelectStage,
_ctx: &RoutingContext<'_>,
candidates: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, Error> {
Ok(candidates)
}
async fn allow_fallback(
&self,
_request: &RouterRequest,
_target: &str,
) -> Result<bool, crate::error::AttemptError> {
Ok(true)
}
async fn sleep(&self, _duration: std::time::Duration) {}
fn observe(&self, _event: &RouterEvent) {}
fn cache(&self) -> Arc<dyn RouterCache> {
Arc::new(InMemoryRouterCache::default())
}
async fn model_info(&self, lookup: &ModelInfoLookup) -> Result<Value, Error> {
self.lookups
.lock()
.unwrap()
.push(serde_json::to_value(lookup).unwrap());
Ok(match lookup {
ModelInfoLookup::ModelInfo {
model, tolerate, ..
} => match model.as_str().and_then(|model| self.rows.get(model)) {
Some(row) => row.clone(),
None if *tolerate => json!({"failed": true}),
None => return Err(raised("ValueError", "unmapped")),
},
ModelInfoLookup::CostEntry { key, .. } => self
.rows
.get(key)
.map_or(Value::Null, |entry| json!({"entry": entry})),
ModelInfoLookup::Provider { .. } => {
json!([self.provider.0, self.provider.1])
}
ModelInfoLookup::UserModelInfo => self.user_model_info.clone(),
_ => Value::Null,
})
}
}
#[test]
fn update_dictionary_converts_numbers_and_merges_nested_dicts_one_level() {
let mut existing = json!({"a": 1, "n": {"x": 1, "y": 2}, "k": "keep"})
.as_object()
.unwrap()
.clone();
let new = json!({"a": None::<i64>, "n": {"y": 3}, "i": "8_192", "f": "3e-07", "s": "abc"})
.as_object()
.unwrap()
.clone();
update_dictionary(&mut existing, &new);
assert_eq!(
Value::Object(existing),
json!({"a": 1, "n": {"x": 1, "y": 3}, "k": "keep", "i": 8192, "f": 3e-07, "s": "abc"})
);
}
#[tokio::test]
async fn router_model_info_overlays_the_deployments_own_values_on_the_row() {
let host = Catalog {
rows: BTreeMap::from([(
"openai/gpt-4o".to_owned(),
json!({"key": "gpt-4o", "max_input_tokens": 100, "input_cost_per_token": 1.0}),
)]),
provider: ("gpt-4o".to_owned(), "openai".to_owned()),
user_model_info: json!({"id": "d", "input_cost_per_token": 2.0, "mode": None::<String>}),
..Catalog::default()
};
let query = RouterModelInfoQuery {
deployment_id: json!("d"),
params_model: "openai/gpt-4o".to_owned(),
received_model_name: "g".to_owned(),
..RouterModelInfoQuery::default()
};
let merged = router_model_info(&host, &query).await.unwrap().unwrap();
assert_eq!(
Value::Object(merged.info),
json!({"key": "gpt-4o", "max_input_tokens": 100, "input_cost_per_token": 2.0, "id": "d"})
);
assert_eq!(merged.user_keys, vec!["id", "input_cost_per_token"]);
}
#[tokio::test]
async fn deployment_model_info_merges_the_entry_over_the_builtin_row() {
let host = Catalog {
rows: BTreeMap::from([
(
"gpt-4o".to_owned(),
json!({"key": "gpt-4o", "max_tokens": 10, "input_cost_per_token": 1.0}),
),
("d".to_owned(), json!({"max_tokens": "20"})),
]),
..Catalog::default()
};
let info = deployment_model_info(&host, "d", &json!("gpt-4o"))
.await
.unwrap()
.unwrap();
assert_eq!(
Value::Object(info),
json!({"key": "gpt-4o", "max_tokens": 20, "input_cost_per_token": 1.0})
);
}
#[tokio::test]
async fn an_unpriced_entry_reports_no_token_prices() {
let host = Catalog {
rows: BTreeMap::from([("d".to_owned(), json!({"max_tokens": 5}))]),
..Catalog::default()
};
let info = deployment_model_info(&host, "d", &json!("unmapped"))
.await
.unwrap()
.unwrap();
assert_eq!(
Value::Object(info),
json!({"max_tokens": 5, "input_cost_per_token": null, "output_cost_per_token": null})
);
}
}

View file

@ -0,0 +1,125 @@
//! Which of a request's outcomes reach the deployment's health. The Python router learns them from
//! logging callbacks it registers on the process lists; `_core_serve_through_request_logging` and
//! `_core_discard` take those off, and `_core_bind_to_request_logging` puts them back on one
//! request's logging object. This core registers none, so the host answers, for each request, what
//! would have run ([`OutcomeCallbacks`]), and every write a callback owns passes through
//! [`OutcomeCallbacks::record`].
//!
//! What the Python router writes outside its callbacks is not gated, and reaches [`crate::Health`]
//! with the default value: `make_call`'s usage, `async_routing_strategy_pre_call_checks`'
//! cooldown, `_trigger_cooldown_for_failed_deployment`, `RouterHealth.record_failure` and a
//! realtime session's success.
use std::future::Future;
use serde::Deserialize;
use crate::error::Error;
/// Which of the router's outcome callbacks litellm's logging runs for a request. The default is a
/// router whose callbacks are on the process lists.
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)]
#[serde(default, deny_unknown_fields)]
pub struct OutcomeCallbacks {
/// `deployment_callback_on_success`, `sync_deployment_callback_on_success` for a sync call.
/// `_core_bind_to_request_logging` binds only the former.
pub success: bool,
/// `deployment_callback_on_failure` and `async_deployment_callback_on_failure`.
pub failure: bool,
/// The attempt runs through `make_call`, which counts its usage whatever the callbacks do.
pub make_call: bool,
}
impl Default for OutcomeCallbacks {
fn default() -> Self {
Self {
success: true,
failure: true,
make_call: true,
}
}
}
/// What an outcome writes.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OutcomeWrite {
/// The success counter, or the usage the success callback adds at the end of a stream.
Success,
/// The tokens and request of an attempt that returned: `make_call`'s write, or the success
/// callback's for an endpoint that skips `make_call`.
AttemptUsage,
/// The failure counter, the allowed-fails counter, the cooldown and the failed request.
Failure,
}
impl OutcomeCallbacks {
/// Run `record` when `write` happens for this request, else record nothing.
pub(crate) async fn record<T: Default>(
self,
write: OutcomeWrite,
record: impl Future<Output = Result<T, Error>>,
) -> Result<T, Error> {
let runs = match write {
OutcomeWrite::Success => self.success,
OutcomeWrite::AttemptUsage => self.success || self.make_call,
OutcomeWrite::Failure => self.failure,
};
if runs { record.await } else { Ok(T::default()) }
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
const REGISTERED: OutcomeCallbacks = OutcomeCallbacks {
success: true,
failure: true,
make_call: true,
};
const UNBOUND: OutcomeCallbacks = OutcomeCallbacks {
success: false,
failure: false,
make_call: true,
};
const UNBOUND_WITHOUT_MAKE_CALL: OutcomeCallbacks = OutcomeCallbacks {
success: false,
failure: false,
make_call: false,
};
const BOUND_SYNC: OutcomeCallbacks = OutcomeCallbacks {
success: false,
failure: true,
make_call: true,
};
#[rstest]
#[case::registered_success(REGISTERED, OutcomeWrite::Success, true)]
#[case::registered_usage(REGISTERED, OutcomeWrite::AttemptUsage, true)]
#[case::registered_failure(REGISTERED, OutcomeWrite::Failure, true)]
#[case::unbound_success(UNBOUND, OutcomeWrite::Success, false)]
#[case::unbound_usage(UNBOUND, OutcomeWrite::AttemptUsage, true)]
#[case::unbound_failure(UNBOUND, OutcomeWrite::Failure, false)]
#[case::no_make_call_usage(UNBOUND_WITHOUT_MAKE_CALL, OutcomeWrite::AttemptUsage, false)]
#[case::bound_sync_success(BOUND_SYNC, OutcomeWrite::Success, false)]
#[case::bound_sync_failure(BOUND_SYNC, OutcomeWrite::Failure, true)]
#[tokio::test]
async fn a_write_runs_when_its_callback_does(
#[case] callbacks: OutcomeCallbacks,
#[case] write: OutcomeWrite,
#[case] runs: bool,
) {
let recorded = callbacks.record(write, async { Ok(true) }).await;
assert_eq!(recorded.ok(), Some(runs));
}
#[test]
fn a_router_on_the_process_lists_is_the_default() {
let parsed: OutcomeCallbacks = serde_json::from_str("{}").expect("an empty document");
assert_eq!(parsed, REGISTERED);
}
}

View file

@ -0,0 +1,300 @@
//! `enable_pre_call_checks` (`_pre_call_checks`): before a pick, drop the deployments whose
//! context window the prompt exceeds, those outside the request's `allowed_model_region`, and
//! those that cannot take a requested param, then decide which error a group with nothing left
//! raises. The facts that need live model data (the prompt's token count, each deployment's
//! context window, which params it takes) come from [`Host::pre_call_facts`]; every decision is
//! made here.
//!
//! [`Host::pre_call_facts`]: crate::host::Host::pre_call_facts
use std::{collections::BTreeMap, sync::Arc};
use serde::Deserialize;
use serde_json::Value;
use crate::{error::Error, registry::Deployment};
/// What the host knows about one deployment for the pre-call checks.
#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
#[serde(default)]
pub struct DeploymentFacts {
/// The deployment's context window (`get_router_model_info`'s `max_input_tokens`), when it
/// declares one and resolves.
pub max_input_tokens: Option<i64>,
/// `_deployment_lacks_params`: it does not support one of the requested params.
pub lacks_params: bool,
/// `_deployment_rejects_params`: it lacks one the request and deployment neither drop nor allow.
pub rejects_params: bool,
}
/// What the host knows about a request for the pre-call checks.
#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
#[serde(default)]
pub struct PreCallFacts {
/// The request carries `messages` or `input`: without either no check runs.
pub has_input: bool,
/// The prompt may be counted against a context window (no compaction is pending).
pub countable: bool,
/// The prompt's token count, when it was counted.
pub input_tokens: Option<i64>,
/// The request's `allowed_model_region`.
pub allowed_model_region: Option<Value>,
/// The requested params the check filters on, sorted (`_pre_call_check_requested_chat_params`).
pub requested_params: Vec<String>,
/// Facts by deployment id, for the candidates and the group's deployments.
pub deployments: BTreeMap<String, DeploymentFacts>,
}
impl PreCallFacts {
fn of(&self, deployment: &Deployment) -> DeploymentFacts {
self.deployments
.get(&deployment.id)
.cloned()
.unwrap_or_default()
}
}
/// The name a context-window error gives a deployment: its base model, else its model.
fn deployment_model(deployment: &Deployment) -> String {
let truthy = |value: Option<&Value>| match value {
Some(Value::String(text)) if !text.is_empty() => Some(text.clone()),
Some(Value::Null | Value::Bool(false)) | None => None,
Some(Value::String(_)) => None,
Some(other) => Some(python_str(other)),
};
truthy(deployment.model_info.raw.get("base_model"))
.or_else(|| truthy(deployment.raw_params.get("base_model")))
.unwrap_or_else(|| deployment.params.model.clone())
}
/// `str(value)` for the scalar values a base model may hold.
fn python_str(value: &Value) -> String {
match value {
Value::Bool(true) => "True".to_owned(),
Value::Number(number) => number.to_string(),
other => other.to_string(),
}
}
/// `is_region_allowed`: the deployment's `region_name` equals the allowed region.
fn in_region(deployment: &Deployment, allowed: &Value) -> bool {
deployment.raw_params.get("region_name") == Some(allowed)
}
/// `_pre_call_checks` over `candidates`; `pool` is the group's deployments before health filters
/// (`request_deployments`), which decide whether nothing left is an unsupported-params error.
pub fn filter(
model: &str,
candidates: Vec<Arc<Deployment>>,
pool: &[Arc<Deployment>],
facts: &PreCallFacts,
) -> Result<Vec<Arc<Deployment>>, Error> {
if !facts.has_input {
return Ok(candidates);
}
let mut invalid = vec![false; candidates.len()];
let mut context_window_error = false;
let mut detail = String::new();
for (index, deployment) in candidates.iter().enumerate() {
let found = facts.of(deployment);
if let (Some(limit), true, Some(tokens)) =
(found.max_input_tokens, facts.countable, facts.input_tokens)
&& tokens > limit
{
invalid[index] = true;
context_window_error = true;
detail.push_str(&format!(
"Model={}, Max Input Tokens={limit}, Got={tokens}",
deployment_model(deployment)
));
continue;
}
if let Some(allowed) = &facts.allowed_model_region
&& !allowed.is_null()
&& !in_region(deployment, allowed)
{
invalid[index] = true;
}
}
let params_checked = !facts.requested_params.is_empty();
let param_candidates: Vec<usize> = (0..candidates.len()).filter(|i| !invalid[*i]).collect();
let lacking: Vec<bool> = param_candidates
.iter()
.map(|i| params_checked && facts.of(&candidates[*i]).lacks_params)
.collect();
let none_takes = !lacking.is_empty() && lacking.iter().all(|lacks| *lacks);
let failing: Vec<bool> = if none_takes {
param_candidates
.iter()
.map(|i| params_checked && facts.of(&candidates[*i]).rejects_params)
.collect()
} else {
lacking
};
for (index, fails) in param_candidates.iter().zip(failing) {
if fails {
invalid[*index] = true;
}
}
let excluded = invalid.iter().filter(|x| **x).count();
if excluded == candidates.len() {
if !context_window_error
&& !pool.is_empty()
&& pool
.iter()
.all(|d| params_checked && facts.of(d).rejects_params)
{
return Err(Error::UnsupportedParams {
model: model.to_owned(),
params: facts.requested_params.clone(),
});
}
if context_window_error {
return Err(Error::ContextWindowExceeded {
model: model.to_owned(),
detail,
});
}
}
Ok(candidates
.into_iter()
.zip(invalid)
.filter_map(|(deployment, invalid)| (!invalid).then_some(deployment))
.collect())
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
use crate::registry::DeploymentSpec;
fn deployment(id: &str, params: Value, model_info: Value) -> Arc<Deployment> {
let mut info = model_info;
info["id"] = json!(id);
let spec: DeploymentSpec = serde_json::from_value(json!({
"model_name": "g", "litellm_params": params, "model_info": info,
}))
.unwrap();
Arc::new(Deployment::from_spec(spec).unwrap())
}
fn facts(entries: &[(&str, DeploymentFacts)]) -> PreCallFacts {
PreCallFacts {
has_input: true,
countable: true,
deployments: entries
.iter()
.map(|(id, f)| ((*id).to_owned(), f.clone()))
.collect(),
..PreCallFacts::default()
}
}
fn window(limit: i64) -> DeploymentFacts {
DeploymentFacts {
max_input_tokens: Some(limit),
..DeploymentFacts::default()
}
}
#[rstest]
fn a_prompt_over_a_context_window_drops_that_deployment() {
let small = deployment("small", json!({"model": "openai/a"}), json!({}));
let big = deployment("big", json!({"model": "openai/b"}), json!({}));
let mut found = facts(&[("small", window(5)), ("big", window(10_000))]);
found.input_tokens = Some(100);
let kept = filter("g", vec![small, big], &[], &found).unwrap();
assert_eq!(
kept.iter().map(|d| d.id.as_str()).collect::<Vec<_>>(),
["big"]
);
}
#[rstest]
fn no_window_left_names_each_deployment_back_to_back() {
let one = deployment("one", json!({"model": "openai/one"}), json!({}));
let two = deployment(
"two",
json!({"model": "openai/two"}),
json!({"base_model": "azure/base"}),
);
let mut found = facts(&[("one", window(5)), ("two", window(6))]);
found.input_tokens = Some(9);
let error = filter("g", vec![one, two], &[], &found).unwrap_err();
let Error::ContextWindowExceeded { model, detail } = error else {
panic!("expected a context window error, got {error:?}");
};
assert_eq!(model, "g");
assert_eq!(
detail,
"Model=openai/one, Max Input Tokens=5, Got=9Model=azure/base, Max Input Tokens=6, Got=9"
);
}
#[rstest]
fn without_input_nothing_is_checked() {
let small = deployment("small", json!({"model": "openai/a"}), json!({}));
let mut found = facts(&[("small", window(5))]);
found.input_tokens = Some(100);
found.has_input = false;
assert_eq!(filter("g", vec![small], &[], &found).unwrap().len(), 1);
}
#[rstest]
#[case::in_region("eu", 1)]
#[case::out_of_region("ap", 0)]
fn the_allowed_region_keeps_only_its_deployments(#[case] region: &str, #[case] kept: usize) {
let eu = deployment(
"eu",
json!({"model": "openai/a", "region_name": "eu"}),
json!({}),
);
let none = deployment("none", json!({"model": "openai/a"}), json!({}));
let mut found = facts(&[]);
found.allowed_model_region = Some(json!(region));
assert_eq!(
filter("g", vec![eu, none], &[], &found).unwrap().len(),
kept
);
}
#[rstest]
fn a_group_where_every_deployment_rejects_the_param_is_unsupported() {
let old = deployment("old", json!({"model": "openai/old"}), json!({}));
let mut found = facts(&[(
"old",
DeploymentFacts {
lacks_params: true,
rejects_params: true,
..DeploymentFacts::default()
},
)]);
found.requested_params = vec!["response_format".to_owned()];
let error = filter("g", vec![old.clone()], &[old], &found).unwrap_err();
assert!(
matches!(error, Error::UnsupportedParams { params, .. } if params == ["response_format"])
);
}
#[rstest]
fn a_deployment_that_takes_the_param_wins_over_one_that_would_drop_it() {
let old = deployment("old", json!({"model": "openai/old"}), json!({}));
let new = deployment("new", json!({"model": "openai/new"}), json!({}));
let mut found = facts(&[(
"old",
DeploymentFacts {
lacks_params: true,
..DeploymentFacts::default()
},
)]);
found.requested_params = vec!["response_format".to_owned()];
let kept = filter("g", vec![old, new], &[], &found).unwrap();
assert_eq!(
kept.iter().map(|d| d.id.as_str()).collect::<Vec<_>>(),
["new"]
);
}
}

View file

@ -0,0 +1,53 @@
//! Pre-routing: before an async pick resolves its name, the host may route the request to another
//! model group (Python's `_core_async_pre_routing_hook`). The complexity, auto, adaptive and
//! quality routers registered from `auto_router/...` marker deployments, and the Claude Code
//! session binding, live host-side; the core asks once per pick, before resolution, and resolves
//! the name the host answers. The attempt still serves the requested name (its `model_group`):
//! only the deployments come from the routed group.
//!
//! A custom routing strategy (`register_plugin`, `set_custom_routing_strategy`, lar1) replaces the
//! whole pick (`_core_async_get_available_deployment` and `get_available_deployment`): the host
//! answers the same step with the deployment document the strategy returned, which may name a
//! deployment the registry does not hold, and the core serves that deployment as is.
use std::sync::Arc;
use crate::{Core, config::Op, error::Error, host::RouterRequest, registry::Deployment};
/// What the host's pre-routing step sees.
#[derive(Clone, Copy, Debug)]
pub struct PreRoute<'a> {
pub op: &'a Op,
pub request: &'a RouterRequest,
/// The name this pick resolves: the requested name or the fallback target.
pub name: &'a str,
/// How many fallbacks deep the chain is (Python's `fallback_depth`): only a first attempt
/// binds a Claude Code session.
pub fallback_depth: u32,
/// 0 for the first attempt in this group, then the retry number.
pub retry: u32,
}
/// The host's answer to a pre-routing step.
#[derive(Clone, Debug)]
pub enum PreRouted {
/// Resolve this model group instead of the requested name.
Model(String),
/// A custom strategy's pick: serve this deployment, without resolving or filtering.
Deployment(Arc<Deployment>),
}
impl Core {
/// What a pick of `ask.name` serves: the host's pre-routed group or custom pick, else the
/// name itself. Only picks that ask for it pre-route (`request.pre_routes`).
pub(crate) async fn pre_routed(&self, ask: PreRoute<'_>) -> Result<PreRouted, Error> {
if !ask.request.pre_routes {
return Ok(PreRouted::Model(ask.name.to_owned()));
}
Ok(self
.host
.pre_route(&ask)
.await?
.unwrap_or_else(|| PreRouted::Model(ask.name.to_owned())))
}
}

View file

@ -0,0 +1,96 @@
//! PTU shares (`litellm_core_utils/ptu_shares.py`): a deployment whose `model_info.ptu_shares`
//! splits it across teams serves only those teams, and the teams holding a share are kept to
//! their share.
use std::{collections::BTreeSet, sync::Arc};
use serde_json::{Map, Value};
use crate::registry::Deployment;
/// `parsed_ptu_shares`: the teams of a non-empty `ptu_shares` map whose every entry is a team id
/// and a positive whole number of PTUs (`2.0` reads as 2, a bool is no number), else None.
pub fn share_teams(model_info: &Map<String, Value>) -> Option<BTreeSet<String>> {
let shares = model_info.get("ptu_shares")?.as_object()?;
if shares.is_empty() {
return None;
}
shares
.iter()
.map(|(team, share)| (!team.is_empty() && whole_ptus(share)).then(|| team.clone()))
.collect()
}
/// `_whole_ptus`: a positive whole number.
fn whole_ptus(share: &Value) -> bool {
match share {
Value::Number(number) => number.as_i64().map_or_else(
|| {
number
.as_u64()
.map(|whole| whole > 0)
.or_else(|| number.as_f64().map(|f| f.fract() == 0.0 && f > 0.0))
.unwrap_or(false)
},
|whole| whole > 0,
),
_ => false,
}
}
/// `filter_ptu_shared_deployments`: the deployments `team_id` holds a share on when it holds one,
/// else the ones split into no shares; a caller without a team holds none. Also whether any
/// deployment was withheld.
pub fn filter_shared(
deployments: Vec<Arc<Deployment>>,
team_id: Option<&str>,
) -> (Vec<Arc<Deployment>>, bool) {
let team = team_id.unwrap_or("");
let checks: Vec<(Arc<Deployment>, Option<BTreeSet<String>>)> = deployments
.into_iter()
.map(|deployment| {
let shares = share_teams(&deployment.model_info.raw);
(deployment, shares)
})
.collect();
let holds_share = checks
.iter()
.any(|(_, shares)| shares.as_ref().is_some_and(|teams| teams.contains(team)));
let total = checks.len();
let kept: Vec<Arc<Deployment>> = checks
.into_iter()
.filter(|(_, shares)| match shares {
Some(teams) => teams.contains(team),
None => !holds_share,
})
.map(|(deployment, _)| deployment)
.collect();
let withheld = kept.len() < total;
(kept, withheld)
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case(json!({"TA": 50}), Some(vec!["TA"]))]
#[case(json!({"TA": 2.0, "TB": 1}), Some(vec!["TA", "TB"]))]
#[case(json!({"TA": 2.5}), None)]
#[case(json!({"TA": true}), None)]
#[case(json!({"TA": 0}), None)]
#[case(json!({}), None)]
#[case(json!({"": 1}), None)]
#[case(json!("TA"), None)]
fn reads_the_share_teams(#[case] shares: Value, #[case] expected: Option<Vec<&str>>) {
let mut model_info = Map::new();
model_info.insert("ptu_shares".to_owned(), shares);
assert_eq!(
share_teams(&model_info),
expected.map(|teams| teams.into_iter().map(str::to_owned).collect())
);
}
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,585 @@
//! A requested name to the deployments that may serve this caller: a port of
//! `_common_checks_available_deployment`. In order: a specific deployment, a deployment id, the
//! alias, a routing group, then for a name no deployment has, a team's public name, every team's
//! public name for a proxy admin, a wildcard; else the group (team-first), kept to the caller's
//! access groups, then deployments whose litellm model is the name, then the first default
//! fallback. Every arm drops strategy markers and deployments that refuse the request's endpoint,
//! and withholds deployments reserved for other teams.
use std::{collections::BTreeSet, sync::Arc};
use jiff::Timestamp;
use serde_json::Value;
use crate::{
access::{filter_by_access_groups, filter_reserved},
eligibility::eligible_deployments,
error::Error,
host_index::{pattern_to_regex, regex_match},
registry::{Deployment, DeploymentSpec, Registry, Resolution, ResolveKind, routed_deployment},
select::is_strategy_marker,
};
/// Who is asking, as the request's auth metadata says.
#[derive(Clone, Copy, Debug, Default)]
pub struct Caller<'a> {
/// `user_api_key_team_id`.
pub team_id: Option<&'a str>,
/// The caller's `user_role` is `proxy_admin`.
pub proxy_admin: bool,
/// The caller's `models` plus `team_models`.
pub allowed_models: &'a [String],
/// The endpoint the pick runs under (`endpoint_selection`), which deployments may refuse.
pub endpoint: Option<&'a str>,
/// The model `endpoint_selection` was given: the name an endpoint refusal reads.
pub requested: &'a str,
/// `f"{provider}/{model}"`, the spelling `PatternMatchRouter.get_pattern` tries when no
/// pattern matches the name itself.
pub wildcard_request: Option<&'a str>,
}
/// `_core_should_include_deployment`: the team's own deployment under its public name, or a
/// deployment named `model_name` the team may use.
fn included(deployment: &Deployment, model_name: &str, team_id: Option<&str>) -> bool {
let info = &deployment.model_info;
if team_id.is_some()
&& info.team_id.as_deref() == team_id
&& info.team_public_model_name.as_deref() == Some(model_name)
{
return true;
}
deployment.model_name == model_name
&& (team_id.is_none() || info.team_id.is_none() || info.team_id.as_deref() == team_id)
}
/// `_core_get_all_deployments(model_name, team_id)`: the team's deployments under that public
/// name when it has any, else the deployments named `model_name` the team may use; a name the
/// model-name index does not hold is looked for in every deployment, for a team.
pub fn all_deployments(
registry: &Registry,
model_name: &str,
team_id: Option<&str>,
) -> Result<Vec<Arc<Deployment>>, Error> {
let keep = |deployments: Vec<Arc<Deployment>>| -> Vec<Arc<Deployment>> {
deployments
.into_iter()
.filter(|d| included(d, model_name, team_id))
.collect()
};
// A written index position past the model list raises `IndexError` where it is read.
let out_of_range = || Error::IndexOutOfRange {
model: model_name.to_owned(),
};
if let Some(team) = team_id {
if registry.team_index_overflows(team, model_name) {
return Err(out_of_range());
}
if let Some(own) = registry.team_named(team, model_name) {
let own = keep(own);
if !own.is_empty() {
return Ok(own);
}
}
}
if registry.name_index_overflows(model_name) {
return Err(out_of_range());
}
Ok(match registry.named(model_name) {
Some(named) => keep(named),
None if team_id.is_some() => keep(registry.all().to_vec()),
None => Vec::new(),
})
}
/// `_team_deployments_across_teams`: every team's deployments under public name `model`, for a
/// proxy admin calling without a team; a name several teams share is refused.
fn across_teams(registry: &Registry, model: &str) -> Result<Vec<Arc<Deployment>>, Error> {
if registry.across_teams_index_overflows(model) {
return Err(Error::IndexOutOfRange {
model: model.to_owned(),
});
}
let deployments = registry.team_named_across_teams(model);
let teams: BTreeSet<&str> = deployments
.iter()
.filter_map(|d| d.model_info.team_id.as_deref())
.collect();
if teams.len() > 1 {
return Err(Error::MultipleTeams {
model: model.to_owned(),
});
}
Ok(deployments)
}
/// The access groups of the rows `_core_get_model_list(model_name, team_id)` lists: the
/// group's deployments, an alias's target deployments, a routing group's rows (which carry no
/// access groups), and only when there are none of those, the wildcard's deployments.
fn model_access_groups(
registry: &Registry,
model: &str,
team_id: Option<&str>,
) -> BTreeSet<String> {
let mut rows = all_deployments(registry, model, team_id).unwrap_or_default();
let mut stripped_rows = 0;
if let Some(target) = registry.alias_of(model) {
match registry.routing_group_named(target) {
Some(group) => stripped_rows += routing_group_rows(registry, group, None).len(),
None => rows.extend(all_deployments(registry, target, None).unwrap_or_default()),
}
}
if let Some(group) = registry.routing_group_named(model) {
stripped_rows += routing_group_rows(registry, group, None).len();
}
if rows.is_empty() && stripped_rows == 0 {
rows = registry
.wildcard_for(model)
.map(|(_, deployments)| deployments)
.unwrap_or_default();
}
rows.iter()
.flat_map(|d| d.model_info.access_groups.iter().cloned())
.collect()
}
fn routing_group_rows(
registry: &Registry,
group: &crate::registry::RoutingGroupSpec,
team_id: Option<&str>,
) -> Vec<Arc<Deployment>> {
crate::routing_groups::members(group, |member| {
all_deployments(registry, member, team_id).unwrap_or_default()
})
}
/// `_filter_reserved_deployments` for deployments resolved as `model`: markers dropped (only
/// markers is refused), those that refuse the caller's endpoint dropped (`_drop_strategy_markers`),
/// then the deployments an active access window withholds; when that leaves nothing, the window
/// is named in the refusal.
fn reserved(
model: &str,
deployments: Vec<Arc<Deployment>>,
caller: Caller<'_>,
now: Timestamp,
) -> Result<Vec<Arc<Deployment>>, Error> {
let any = !deployments.is_empty();
let selectable: Vec<Arc<Deployment>> = deployments
.into_iter()
.filter(|d| !is_strategy_marker(d))
.collect();
if any && selectable.is_empty() {
return Err(Error::OnlyMarkers {
model: model.to_owned(),
});
}
let selectable = eligible_deployments(model, caller, selectable)?;
let (kept, blocking) = filter_reserved(selectable, caller.team_id, now);
match blocking {
Some(window) if kept.is_empty() => Err(Error::Reserved {
model: model.to_owned(),
until: window.end_hhmm(),
timezone: window.timezone,
}),
_ => Ok(kept),
}
}
/// `PatternMatchRouter.route`: copies of a pattern's deployments serving `request`.
fn routed_wildcard(
pattern: &str,
request: &str,
deployments: &[Arc<Deployment>],
) -> Vec<Arc<Deployment>> {
let groups = regex_match(&pattern_to_regex(pattern), request).unwrap_or_default();
deployments
.iter()
.map(|deployment| routed_deployment(deployment, &groups, request))
.collect()
}
/// `get_pattern` for `model`: the most specific pattern matching it as spelled, else its
/// provider-prefixed spelling, with the pattern's deployments serving that request.
fn wildcard_route(
registry: &Registry,
model: &str,
wildcard_request: Option<&str>,
) -> Option<(String, Vec<Arc<Deployment>>)> {
let route = |request: &str| {
registry
.wildcard_for(request)
.map(|(pattern, deployments)| {
let routed = routed_wildcard(&pattern, request, &deployments);
(pattern, routed)
})
};
route(model).or_else(|| route(wildcard_request?))
}
/// The name resolution would look for under its provider (`get_pattern`'s second spelling), which
/// only the host can infer: one that reaches the pattern routers and that no pattern matches as
/// spelled.
pub(crate) fn wildcard_name<'a>(
registry: &'a Registry,
name: &'a str,
specific_deployment: bool,
) -> Option<&'a str> {
let model = registry.alias_of(name).unwrap_or(name);
(registry.has_wildcards()
&& !specific_deployment
&& !registry.has_name(model)
&& (registry.has_name(name) || registry.deployment(name).is_none())
&& registry.routing_group_named(model).is_none()
&& registry.wildcard_for(model).is_none())
.then_some(model)
}
/// The first default fallback a name no deployment serves is routed to.
pub fn default_target(dynamic: &crate::config::RouterDynamicConfig) -> Option<&str> {
dynamic
.fallbacks
.default_target()
.or_else(|| dynamic.default_fallbacks.first().map(String::as_str))
}
/// The router's `default_deployment`, as resolution serves it.
#[derive(Clone, Copy, Debug)]
pub enum DefaultDeployment<'a> {
Spec(&'a DeploymentSpec),
/// A value Python cannot copy for a name: resolving to it raises.
Raises,
}
/// The default deployment as served for `model`: a copy whose `litellm_params.model` is the name.
fn default_for(spec: &DeploymentSpec, model: &str) -> Option<Arc<Deployment>> {
let mut spec = spec.clone();
spec.litellm_params
.insert("model".to_owned(), Value::String(model.to_owned()));
Deployment::from_spec(spec).ok().map(Arc::new)
}
/// `_core_get_routable_deployments`: the deployments resolution would consider for `request`,
/// before any team, cooldown or strategy filter; nothing when it resolves to nothing or refuses.
pub fn routable(
registry: &Registry,
dynamic: &crate::config::RouterDynamicConfig,
request: &crate::host::RouterRequest,
) -> Vec<Arc<Deployment>> {
let caller = Caller {
team_id: request.team_id.as_deref(),
proxy_admin: request.proxy_admin,
allowed_models: &request.allowed_models,
endpoint: request.endpoint.as_deref(),
requested: &request.model,
wildcard_request: request.wildcard_request.as_deref(),
};
match resolve_request(
registry,
&request.model,
false,
caller,
default_target(dynamic),
dynamic.default_ref(),
Timestamp::now(),
) {
Ok(Some(resolution)) => resolution.deployments,
Ok(None) | Err(_) => Vec::new(),
}
}
/// Resolve `name` for `caller`. `None` when nothing serves it (`no_healthy_deployments`).
pub fn resolve_request(
registry: &Registry,
name: &str,
specific_deployment: bool,
caller: Caller<'_>,
default_target: Option<&str>,
default_deployment: Option<DefaultDeployment<'_>>,
now: Timestamp,
) -> Result<Option<Resolution>, Error> {
let team_id = caller.team_id;
let resolved = |kind, group: &str, deployments| {
Ok(Some(Resolution {
kind,
group: group.to_owned(),
deployments,
}))
};
if specific_deployment {
let matching = registry
.all()
.iter()
.filter(|d| d.params.model == name)
.cloned()
.collect();
let kept = reserved(name, matching, caller, now)?;
return resolved(ResolveKind::SpecificDeployment, name, kept);
}
if !registry.has_name(name)
&& let Some(deployment) = registry.deployment(name)
{
let group = deployment.model_name.clone();
let litellm_model = deployment.params.model.clone();
let kept = reserved(&litellm_model, vec![deployment], caller, now)?;
return resolved(ResolveKind::Id, &group, kept);
}
let aliased = registry.alias_of(name);
let model = aliased.unwrap_or(name);
let routing_group = registry
.routing_group_named(model)
.map(|group| routing_group_rows(registry, group, team_id));
if routing_group.is_none() && !registry.has_name(model) {
let early = match team_id {
Some(team) => Some(all_deployments(registry, model, Some(team))?),
None if caller.proxy_admin => Some(across_teams(registry, model)?),
None => None,
}
.filter(|deployments| !deployments.is_empty());
if let Some(deployments) = early {
let kept = reserved(model, deployments, caller, now)?;
return resolved(ResolveKind::Group, model, kept);
}
// `_core_wildcard_route_deployments`: the team's own pattern, else the proxy-wide one.
let team_owned = team_id
.map(|team| crate::lookup::Lookup::new(registry).team_wildcard(model, team))
.filter(|deployments| !deployments.is_empty());
if let Some(deployments) = team_owned {
let routed = deployments.iter().flat_map(|deployment| {
let public = deployment.model_info.team_public_model_name.as_deref();
let pattern = public.unwrap_or(model);
routed_wildcard(pattern, model, std::slice::from_ref(deployment))
});
let kept = reserved(model, routed.collect(), caller, now)?;
return resolved(ResolveKind::Wildcard, model, kept);
}
if let Some((pattern, deployments)) =
wildcard_route(registry, model, caller.wildcard_request)
{
let kept = reserved(model, deployments, caller, now)?;
return resolved(ResolveKind::Wildcard, &pattern, kept);
}
match default_deployment {
Some(DefaultDeployment::Raises) => {
return Err(Error::DefaultDeploymentRaises {
model: model.to_owned(),
});
}
Some(DefaultDeployment::Spec(spec)) => {
if let Some(deployment) = default_for(spec, model) {
let kept = reserved(model, vec![deployment], caller, now)?;
return resolved(ResolveKind::Default, model, kept);
}
}
None => {}
}
}
let (kind, listed) = match routing_group {
Some(deployments) => (ResolveKind::RoutingGroup, deployments),
None if aliased.is_some() => (
ResolveKind::Alias,
all_deployments(registry, model, team_id)?,
),
None => (
ResolveKind::Group,
all_deployments(registry, model, team_id)?,
),
};
let narrow = |model: &str, deployments: Vec<Arc<Deployment>>| {
let allowed = filter_by_access_groups(model, deployments, caller.allowed_models, || {
model_access_groups(registry, model, team_id)
});
reserved(model, allowed, caller, now)
};
let listed_any = !listed.is_empty();
let deployments = narrow(model, listed)?;
if !deployments.is_empty() {
return resolved(kind, model, deployments);
}
let mut access_emptied = listed_any;
if !listed_any {
// The caller may have named a deployment's litellm model instead.
let by_model: Vec<Arc<Deployment>> = registry
.all()
.iter()
.filter(|d| d.params.model == model)
.cloned()
.collect();
let found = !by_model.is_empty();
let deployments = narrow(model, by_model)?;
if !deployments.is_empty() {
return resolved(ResolveKind::Group, model, deployments);
}
access_emptied = found;
}
if let Some(target) = default_target.filter(|target| !access_emptied && !target.is_empty()) {
let deployments = narrow(target, all_deployments(registry, target, team_id)?)?;
if !deployments.is_empty() {
return resolved(ResolveKind::Group, target, deployments);
}
// `_common_checks_available_deployment` re-assigned the name: its error names the target.
return Err(Error::UnknownModel {
model: target.to_owned(),
});
}
Ok(None)
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::{Value, json};
use super::*;
use crate::registry::DeploymentSpec;
fn registry(entries: Value) -> Registry {
let specs: Vec<DeploymentSpec> = serde_json::from_value(entries).unwrap();
Registry::build(specs, BTreeMap::new(), Vec::new()).unwrap()
}
fn ids(resolution: Option<Resolution>) -> Vec<String> {
resolution
.map(|r| r.deployments.iter().map(|d| d.id.clone()).collect())
.unwrap_or_default()
}
#[test]
fn an_access_group_only_caller_is_kept_to_its_groups() {
let registry = registry(json!([
{"model_name": "g", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a", "access_groups": ["A"]}},
{"model_name": "g", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "b", "access_groups": ["B"]}},
]));
let allowed = vec!["A".to_owned()];
let caller = Caller {
allowed_models: &allowed,
..Caller::default()
};
let resolved = resolve_request(&registry, "g", false, caller, None, None, Timestamp::now());
assert_eq!(ids(resolved.unwrap()), ["a"]);
// The litellm model string is not narrowed (python-fixes-to-make.md).
let by_model = resolve_request(
&registry,
"openai/b",
false,
caller,
None,
None,
Timestamp::now(),
);
assert_eq!(ids(by_model.unwrap()), ["b"]);
}
#[test]
fn a_proxy_admin_is_refused_a_public_name_several_teams_share() {
let registry = registry(json!([
{"model_name": "x", "litellm_params": {"model": "openai/x"}, "model_info": {"id": "x", "team_id": "tx", "team_public_model_name": "pub"}},
{"model_name": "y", "litellm_params": {"model": "openai/y"}, "model_info": {"id": "y", "team_id": "ty", "team_public_model_name": "pub"}},
]));
let caller = Caller {
proxy_admin: true,
..Caller::default()
};
let refused = resolve_request(
&registry,
"pub",
false,
caller,
None,
None,
Timestamp::now(),
);
assert!(matches!(refused, Err(Error::MultipleTeams { model }) if model == "pub"));
}
fn wildcard_group() -> Value {
json!([{
"model_name": "openai/*",
"litellm_params": {"model": "openai/*"},
"model_info": {"id": "w", "supports_prompt_cache_breakpoint": true},
}])
}
#[test]
fn a_wildcard_serves_the_matched_model_without_its_cache_flags() {
let registry = registry(wildcard_group());
let resolved = resolve_request(
&registry,
"openai/gpt-4o",
false,
Caller::default(),
None,
None,
Timestamp::now(),
);
let resolution = resolved.unwrap().unwrap();
assert_eq!(resolution.kind, ResolveKind::Wildcard);
assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o");
assert_eq!(
resolution.deployments[0].record()["model_info"],
json!({"id": "w"})
);
}
#[test]
fn an_alias_to_a_wildcard_fills_the_model_from_its_target() {
let specs: Vec<DeploymentSpec> = serde_json::from_value(wildcard_group()).unwrap();
let aliases = BTreeMap::from([("al".to_owned(), "openai/gpt-4o".to_owned())]);
let registry = Registry::build(specs, aliases, Vec::new()).unwrap();
let resolved = resolve_request(
&registry,
"al",
false,
Caller::default(),
None,
None,
Timestamp::now(),
);
assert_eq!(
resolved.unwrap().unwrap().deployments[0].params.model,
"openai/gpt-4o"
);
}
#[test]
fn a_name_no_pattern_matches_is_tried_under_its_provider() {
let registry = registry(wildcard_group());
assert_eq!(wildcard_name(&registry, "gpt-4o", false), Some("gpt-4o"));
assert_eq!(wildcard_name(&registry, "openai/gpt-4o", false), None);
let caller = Caller {
wildcard_request: Some("openai/gpt-4o"),
..Caller::default()
};
let resolved = resolve_request(
&registry,
"gpt-4o",
false,
caller,
None,
None,
Timestamp::now(),
);
let resolution = resolved.unwrap().unwrap();
assert_eq!(resolution.group, "openai/*");
assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o");
}
#[test]
fn a_deployment_named_by_id_refuses_the_endpoint_like_any_other() {
let registry = registry(json!([
{"model_name": "g", "litellm_params": {"model": "openai/a"}, "model_info": {
"id": "a", "enforce_supported_endpoints": true, "supported_endpoints": ["/v1/embeddings"],
}},
]));
let caller = Caller {
endpoint: Some("/v1/chat/completions"),
requested: "a",
..Caller::default()
};
let refused = resolve_request(&registry, "a", false, caller, None, None, Timestamp::now());
assert!(matches!(
refused,
Err(Error::EndpointNotSupported { model, .. }) if model == "openai/a"
));
}
}

View file

@ -0,0 +1,219 @@
//! Retry decisions: how many retries a `RetryPolicy` grants for a failure, and whether a failure
//! is retried in the same group at all. Mirrors `get_num_retries_from_retry_policy` and
//! `resolve_retry_policy` in `litellm_core_utils/get_retry_from_policy.py`, and
//! `Router.should_retry_this_error`.
use std::collections::BTreeMap;
use crate::{
config::RetryPolicy,
error::{AttemptError, FailureKind},
health::should_retry_status,
};
/// `resolve_retry_policy`: the group's own policy when `model_group_retry_policy` has one, else
/// the router's `retry_policy`.
pub fn resolve_retry_policy<'a>(
retry_policy: Option<&'a RetryPolicy>,
model_group: &str,
model_group_retry_policy: Option<&'a BTreeMap<String, RetryPolicy>>,
) -> Option<&'a RetryPolicy> {
model_group_retry_policy
.and_then(|policies| policies.get(model_group))
.or(retry_policy)
}
impl RetryPolicy {
/// `get_num_retries_from_retry_policy`: `DefaultRetries` for an `UnsupportedEndpointError`;
/// else `NotFoundErrorRetries` for any 404 answer, then the most specific exception class
/// the policy names, then `DefaultRetries`.
pub fn retries_for(&self, failure: &AttemptError) -> Option<u32> {
if failure.is_a("UnsupportedEndpointError") {
return self.DefaultRetries;
}
let not_found = (failure.status == Some(404))
.then_some(self.NotFoundErrorRetries)
.flatten();
not_found
.or_else(|| {
failure
.class_names()
.iter()
.find_map(|class| self.field_for(class))
})
.or(self.DefaultRetries)
}
/// `_RETRIES_BY_EXCEPTION_TYPE`: the policy field an exception class reads.
fn field_for(&self, class: &str) -> Option<u32> {
match class {
"AuthenticationError" => self.AuthenticationErrorRetries,
"Timeout" => self.TimeoutErrorRetries,
"RateLimitError" => self.RateLimitErrorRetries,
"ContentPolicyViolationError" => self.ContentPolicyViolationErrorRetries,
"BadRequestError" => self.BadRequestErrorRetries,
"ServiceUnavailableError" => self.ServiceUnavailableErrorRetries,
"InternalServerError" => self.InternalServerErrorRetries,
_ => None,
}
}
}
/// What `should_retry_this_error` knows besides the failure.
#[derive(Clone, Copy, Debug, Default)]
pub struct RetryContext {
pub healthy: usize,
pub all: usize,
/// `context_window_fallbacks is not None`.
pub has_context_window_fallbacks: bool,
/// `content_policy_fallbacks is not None`.
pub has_content_policy_fallbacks: bool,
/// `regular_fallbacks is not None and len(regular_fallbacks) > 0`.
pub has_regular_fallbacks: bool,
}
/// `should_retry_this_error`, as a predicate: `false` where Python raises the error.
pub fn should_retry(failure: &AttemptError, ctx: RetryContext) -> bool {
if failure.is_a("ContextWindowExceededError") && ctx.has_context_window_fallbacks {
return false;
}
if failure.is_a("ContentPolicyViolationError") && ctx.has_content_policy_fallbacks {
return false;
}
if let Some(status) = failure.status
&& !should_retry_status(status)
&& !matches!(status, 401 | 403)
&& !failure.is_a("UnsupportedEndpointError")
{
return false;
}
if failure.is_a("NotFoundError") {
return false;
}
if failure.is_a("RateLimitError") && ctx.healthy == 0 && ctx.has_regular_fallbacks {
return false;
}
if (failure.is_a("AuthenticationError") || failure.is_a("PermissionDeniedError"))
&& ctx.all <= 1
{
return false;
}
ctx.healthy > 0
}
/// `ContentPolicyViolationError`'s classes, most specific first.
pub(crate) fn content_policy_classes() -> Vec<std::sync::Arc<str>> {
[
"ContentPolicyViolationError",
"BadRequestError",
"APIStatusError",
"APIError",
"OpenAIError",
"Exception",
"BaseException",
"object",
]
.into_iter()
.map(Into::into)
.collect()
}
/// The exception classes a failure kind stands for, most specific first, for a host that
/// reports no class names (`litellm/exceptions.py`).
pub(crate) fn kind_classes(kind: FailureKind) -> &'static [&'static str] {
match kind {
FailureKind::RateLimit => &["RateLimitError"],
FailureKind::ContextWindowExceeded => &["ContextWindowExceededError", "BadRequestError"],
FailureKind::ContentPolicyViolation => &["ContentPolicyViolationError", "BadRequestError"],
FailureKind::Authentication => &["AuthenticationError"],
FailureKind::PermissionDenied => &["PermissionDeniedError"],
FailureKind::NotFound => &["NotFoundError"],
FailureKind::BadRequest => &["BadRequestError"],
FailureKind::Timeout => &["Timeout", "APITimeoutError", "APIConnectionError"],
FailureKind::ApiConnection => &["APIConnectionError"],
FailureKind::InternalServer => &["InternalServerError"],
FailureKind::ServiceUnavailable => &["ServiceUnavailableError"],
FailureKind::BadGateway => &["BadGatewayError"],
FailureKind::CallerCredentialAuthentication => {
&["CallerCredentialAuthenticationError", "AuthenticationError"]
}
FailureKind::CallerCredentialRateLimit => {
&["CallerCredentialRateLimitError", "RateLimitError"]
}
FailureKind::GuardrailIntervention | FailureKind::Declined | FailureKind::Other => &[],
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn failure(kind: FailureKind, status: Option<u16>, classes: &[&str]) -> AttemptError {
let mut failure = AttemptError::new(kind, status);
failure.classes = classes.iter().map(|class| (*class).into()).collect();
failure
}
#[rstest]
#[case::own_field(failure(FailureKind::InternalServer, Some(500), &[]), Some(4))]
#[case::default(failure(FailureKind::BadGateway, Some(502), &[]), Some(7))]
#[case::not_found_status(failure(FailureKind::Other, Some(404), &["APIError"]), Some(1))]
#[case::unsupported_endpoint(
failure(FailureKind::Declined, Some(400), &["UnsupportedEndpointError", "BadRequestError"]),
Some(7)
)]
#[case::malformed_reads_bad_request(
failure(FailureKind::Declined, Some(400), &["MalformedRequestFieldError", "BadRequestError"]),
Some(3)
)]
fn retries_follow_the_most_specific_class(
#[case] failure: AttemptError,
#[case] expected: Option<u32>,
) {
let policy = RetryPolicy {
InternalServerErrorRetries: Some(4),
BadRequestErrorRetries: Some(3),
NotFoundErrorRetries: Some(1),
DefaultRetries: Some(7),
..RetryPolicy::default()
};
assert_eq!(policy.retries_for(&failure), expected);
}
#[rstest]
#[case::rate_limit_with_healthy(FailureKind::RateLimit, Some(429), 1, 2, false, true)]
#[case::rate_limit_none_healthy_with_fallbacks(
FailureKind::RateLimit,
Some(429),
0,
2,
true,
false
)]
#[case::bad_request(FailureKind::BadRequest, Some(400), 2, 2, false, false)]
#[case::auth_lone_deployment(FailureKind::Authentication, Some(401), 1, 1, false, false)]
#[case::auth_with_siblings(FailureKind::Authentication, Some(401), 1, 2, false, true)]
#[case::not_found(FailureKind::NotFound, Some(404), 2, 2, false, false)]
#[case::server_error_no_healthy(FailureKind::InternalServer, Some(500), 0, 2, false, false)]
fn should_retry_follows_python(
#[case] kind: FailureKind,
#[case] status: Option<u16>,
#[case] healthy: usize,
#[case] all: usize,
#[case] has_regular_fallbacks: bool,
#[case] expected: bool,
) {
let ctx = RetryContext {
healthy,
all,
has_regular_fallbacks,
..RetryContext::default()
};
assert_eq!(
should_retry(&AttemptError::new(kind, status), ctx),
expected
);
}
}

View file

@ -0,0 +1,48 @@
//! Routing groups: a callable name over the deployments of its member model groups
//! (`_get_routing_group_deployments`). A priority group serves each member's priority as the
//! member deployments' `order` (`apply_routing_group_priority`), so the order filter and order
//! fallbacks walk the priorities.
use std::sync::Arc;
use serde_json::Value;
use crate::registry::{Deployment, RoutingGroupSpec};
/// The member deployments of `group`, in member order, with a priority group's priorities as
/// their `order`.
pub fn members(
group: &RoutingGroupSpec,
deployments_of: impl Fn(&str) -> Vec<Arc<Deployment>>,
) -> Vec<Arc<Deployment>> {
let priorities = group
.model_priorities
.as_ref()
.filter(|_| is_priority(group));
group
.models
.iter()
.flat_map(|member| {
let priority = priorities.and_then(|p| p.get(member)).copied();
deployments_of(member)
.into_iter()
.map(move |deployment| match priority {
Some(order) => with_order(&deployment, order),
None => deployment,
})
})
.collect()
}
/// Whether `group` is a priority group.
pub fn is_priority(group: &RoutingGroupSpec) -> bool {
matches!(&group.routing_strategy, crate::config::RoutingStrategy::Other(name) if name == "priority")
}
fn with_order(deployment: &Arc<Deployment>, order: i64) -> Arc<Deployment> {
let mut prioritized = Deployment::clone(deployment);
prioritized
.raw_params
.insert("order".to_owned(), Value::from(order));
Arc::new(prioritized)
}

View file

@ -0,0 +1,114 @@
//! `Router(plugins=[...])`: the routing-plugin pipeline. The host runs the plugins (live Python
//! objects) in the async pick's pre-routing hook and reports the candidate models they left on
//! the request (`_routing_plugin_candidate_models`); router-core decides which deployments that
//! keeps (`_filter_by_routing_plugin_candidates`) and refuses the paths that never run the
//! pipeline (`_get_available_deployment_unresolved`).
use std::sync::Arc;
use crate::{config::RoutingStrategy, error::Error, host::RouterRequest, registry::Deployment};
/// Whether the router's own strategy has an async selector: the others pick on the sync path
/// even for async callers (`_async_get_available_deployment`).
pub(crate) fn picks_async(strategy: &RoutingStrategy) -> bool {
matches!(
strategy,
RoutingStrategy::UsageBasedV2
| RoutingStrategy::SimpleShuffle
| RoutingStrategy::CostBased
| RoutingStrategy::LatencyBased
| RoutingStrategy::LeastBusy
)
}
/// The sync pick never runs the pipeline, so a router with plugins refuses it rather than
/// letting a policy plugin be bypassed. `request` is a pick or a call, never a listing.
pub(crate) fn refuse_sync_pick(
has_plugins: bool,
strategy: &RoutingStrategy,
request: &RouterRequest,
) -> Result<(), Error> {
if has_plugins
&& request.pass_through.is_none()
&& (!request.pre_routes || !picks_async(strategy))
{
return Err(Error::RoutingPluginsNeedAsync);
}
Ok(())
}
/// `_filter_by_routing_plugin_candidates`: keep the deployments whose litellm model the plugins
/// left; nothing left is an error, never a fallback to the unfiltered pool.
pub(crate) fn narrow(
healthy: Vec<Arc<Deployment>>,
models: Option<&[String]>,
) -> Result<Vec<Arc<Deployment>>, Error> {
let Some(models) = models else {
return Ok(healthy);
};
let kept: Vec<Arc<Deployment>> = healthy
.into_iter()
.filter(|d| models.contains(&d.params.model))
.collect();
if kept.is_empty() {
return Err(Error::RoutingPluginsLeftNothing {
models: models.to_vec(),
});
}
Ok(kept)
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
use crate::registry::DeploymentSpec;
fn deployment(id: &str, model: &str) -> Arc<Deployment> {
let spec: DeploymentSpec = serde_json::from_value(json!({
"model_name": "g", "litellm_params": {"model": model}, "model_info": {"id": id},
}))
.unwrap();
Arc::new(Deployment::from_spec(spec).unwrap())
}
#[rstest]
fn the_plugins_models_narrow_the_pool() {
let pool = vec![deployment("a", "openai/a"), deployment("b", "openai/b")];
let kept = narrow(pool, Some(&["openai/b".to_owned()])).unwrap();
assert_eq!(
kept.iter().map(|d| d.id.as_str()).collect::<Vec<_>>(),
["b"]
);
}
#[rstest]
fn narrowing_to_nothing_is_an_error() {
let pool = vec![deployment("a", "openai/a")];
assert!(matches!(
narrow(pool, Some(&[])),
Err(Error::RoutingPluginsLeftNothing { models }) if models.is_empty()
));
}
#[rstest]
#[case::async_pick(true, RoutingStrategy::SimpleShuffle, false)]
#[case::sync_pick(false, RoutingStrategy::SimpleShuffle, true)]
#[case::v1_strategy(true, RoutingStrategy::UsageBased, true)]
fn only_the_async_pick_runs_plugins(
#[case] pre_routes: bool,
#[case] strategy: RoutingStrategy,
#[case] refused: bool,
) {
let request = RouterRequest {
pre_routes,
..RouterRequest::default()
};
assert_eq!(
refuse_sync_pick(true, &strategy, &request).is_err(),
refused
);
}
}

View file

@ -0,0 +1,795 @@
//! Selection: resolve a name, filter its deployments in the order of
//! `Router._core_async_get_healthy_deployments` (team, web search, cooldown, blocked, plugins,
//! pre-call checks, tags, order, retry skips), then pick one with the strategy. Simple-shuffle is
//! native (`router_strategy/simple_shuffle.py`); other strategies go through `Host::select`.
use std::{collections::BTreeMap, sync::Arc};
use parking_lot::Mutex;
use rand::{Rng, SeedableRng, rngs::StdRng};
use crate::{
Core, RouterStateSnapshot,
config::{Op, RoutingStrategy},
error::Error,
health_routing::WhenAllUnhealthy,
host::{RouterRequest, RoutingContext, SelectStage},
registry::{Deployment, Registry, ResolveKind},
resolve::{Caller, resolve_request},
strategy::{StrategyScope, routing_context},
};
/// Randomness for the shuffle and retry jitter, injectable so tests are deterministic.
pub trait Randomness: Send + Sync {
/// A float in `[0, 1)`, as Python's `random.random()`.
fn next_f64(&self) -> f64;
/// An index below `len`, as Python's `random.choice` draws one.
fn index_below(&self, len: usize) -> usize {
((self.next_f64() * len as f64) as usize).min(len - 1)
}
}
pub struct ThreadRandomness;
impl Randomness for ThreadRandomness {
fn next_f64(&self) -> f64 {
rand::thread_rng().r#gen()
}
}
pub struct SeededRandomness(Mutex<StdRng>);
impl SeededRandomness {
pub fn new(seed: u64) -> Self {
Self(Mutex::new(StdRng::seed_from_u64(seed)))
}
}
impl Randomness for SeededRandomness {
fn next_f64(&self) -> f64 {
self.0.lock().r#gen()
}
}
/// Which pass-through pick a request makes.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PassThrough {
/// `get_available_deployment_for_pass_through` (and `select(kind="pass_through")`): the
/// pass-through deployments are chosen first, then cooldown and blocked filters, and only a
/// priority group filters by order.
Listed,
/// `async_get_available_deployment_for_pass_through`: the full healthy chain, then the
/// pass-through filter.
Healthy,
}
/// A group's deployments before and after filtering.
#[derive(Clone, Debug)]
pub struct Candidates {
/// The name that was resolved (Python's `model`).
pub name: String,
pub model_group: String,
pub kind: ResolveKind,
/// Every deployment the caller may use (after the team filter).
pub all: Vec<Arc<Deployment>>,
pub healthy: Vec<Arc<Deployment>>,
/// The provider-prefixed spelling the name resolved under (`get_pattern`), when the host
/// was asked for one.
pub wildcard_request: Option<String>,
}
impl Candidates {
/// The `model` a no-deployments error names: what `_common_checks_available_deployment`
/// returned, so the alias-resolved name a wildcard serves and not its pattern.
pub(crate) fn failed_model(&self, registry: &Registry) -> String {
let model = match self.kind {
ResolveKind::Wildcard => registry.alias_of(&self.name).unwrap_or(&self.name),
_ => self.model_group.as_str(),
};
model.to_owned()
}
}
/// Per-attempt selection inputs the call loop carries.
#[derive(Clone, Debug, Default)]
pub(crate) struct SelectInputs {
pub target_order: Option<i64>,
pub skipped_ids: Vec<String>,
/// A weighted failover hop's `_excluded_deployment_ids`: never picked.
pub excluded_ids: Arc<[String]>,
/// Stop before the plugins, order and skip filters: the group less cooldowns and blocked
/// deployments, which `_async_get_healthy_deployments` counts when deciding whether to retry.
pub count_only: bool,
/// The level's name, which an endpoint refusal reads (`endpoint_selection(model=...)`): the
/// name before pre-routing replaced it.
pub requested: Option<String>,
/// A fallback hop: `run_async_fallback` pops the request's `_target_order`, so only the
/// hop's own order applies.
pub fallback_hop: bool,
}
/// `simple_shuffle`: weighted by `_router_weights[group]`, else `weight`, else `rpm`, else
/// `tpm`, each normalized by its maximum, else uniform. A metric is read only when every set
/// before it weighed nothing, and a non-numeric one fails the pick then (`_metric_weight`).
pub fn simple_shuffle(
deployments: &[Arc<Deployment>],
model_group: &str,
router_weights: Option<&BTreeMap<String, BTreeMap<String, f64>>>,
rng: &dyn Randomness,
) -> Result<Option<Arc<Deployment>>, Error> {
if deployments.is_empty() {
return Ok(None);
}
let scoped = router_weights
.and_then(|weights| weights.get(model_group))
.filter(|weights| !weights.is_empty())
.map(|weights| {
deployments
.iter()
.map(|d| weights.get(&d.id).copied().unwrap_or(0.0))
.collect::<Vec<_>>()
});
let metric_sets = ["weight", "rpm", "tpm"].into_iter().map(|metric| {
deployments
.iter()
.map(|d| d.metric_weight(metric))
.collect::<Option<Vec<_>>>()
.ok_or(Error::NonNumericWeight { metric })
});
for weights in scoped.into_iter().map(Ok).chain(metric_sets) {
let weights = weights?;
let largest = weights.iter().copied().fold(0.0, f64::max);
if largest <= 0.0 {
continue;
}
let normalized: Vec<f64> = weights.iter().map(|weight| weight / largest).collect();
let total: f64 = normalized.iter().sum();
if total <= 0.0 {
continue;
}
// `random.choices`: bisect_right over the cumulative weights.
let target = rng.next_f64() * total;
let mut cumulative = 0.0;
let index = normalized
.iter()
.position(|weight| {
cumulative += weight;
target < cumulative
})
.unwrap_or(deployments.len() - 1);
return Ok(Some(deployments[index].clone()));
}
Ok(Some(
deployments[rng.index_below(deployments.len())].clone(),
))
}
/// A pre-routing strategy marker (`classify_strategy_router_model`): its model is `auto_router/...`.
pub fn is_strategy_marker(deployment: &Deployment) -> bool {
deployment.params.model.starts_with("auto_router/")
}
/// `litellm_params.use_in_pass_through`, read for truthiness as Python does.
fn uses_pass_through(deployment: &Deployment) -> bool {
match deployment.raw_params.get("use_in_pass_through") {
None | Some(serde_json::Value::Null) => false,
Some(serde_json::Value::Bool(value)) => *value,
Some(serde_json::Value::Number(value)) => value.as_f64().is_some_and(|n| n != 0.0),
Some(serde_json::Value::String(value)) => !value.is_empty(),
Some(serde_json::Value::Array(value)) => !value.is_empty(),
Some(serde_json::Value::Object(value)) => !value.is_empty(),
}
}
/// `filter_team_based_models`: a team-owned deployment serves only its team, except that a proxy
/// admin without a team keeps deployments that are all team-owned and listed under the model the
/// request names, unless they belong to several teams.
fn filter_team_based(
deployments: Vec<Arc<Deployment>>,
request: &RouterRequest,
) -> Result<Vec<Arc<Deployment>>, Error> {
let team_id = request.team_id.as_deref();
if team_id.is_none()
&& request.proxy_admin
&& let Some(requested) = request.admin_model.as_deref()
&& !deployments.is_empty()
&& deployments.iter().all(|d| {
d.model_info.team_id.is_some()
&& (d.model_name == requested
|| d.model_info.team_public_model_name.as_deref() == Some(requested))
})
{
let teams: std::collections::BTreeSet<&str> = deployments
.iter()
.filter_map(|d| d.model_info.team_id.as_deref())
.collect();
if teams.len() > 1 {
return Err(Error::MultipleTeams {
model: requested.to_owned(),
});
}
return Ok(deployments);
}
Ok(deployments
.into_iter()
.filter(|d| d.usable_by_team(team_id))
.collect())
}
impl Core {
fn context<'a>(
&'a self,
snapshot: &'a RouterStateSnapshot,
op: &'a Op,
request: &'a RouterRequest,
model_group: &'a str,
name: &'a str,
scope: StrategyScope,
) -> RoutingContext<'a> {
RoutingContext {
registry: &snapshot.registry,
settings: &snapshot.settings,
health: &self.health,
request,
op,
model_group,
name,
scope,
pool: &[],
excluded_ids: &[],
skipped_ids: &[],
target_order: None,
}
}
/// The eligible deployments for `name`, in filter-chain order.
pub(crate) async fn candidates(
&self,
snapshot: &RouterStateSnapshot,
op: &Op,
request: &RouterRequest,
name: &str,
inputs: &SelectInputs,
) -> Result<Candidates, Error> {
let registry = &snapshot.registry;
// The async pick's pre-routing hook runs the routing plugins before anything resolves;
// every pick but a listed pass-through one reads what they left on the request.
let plugin_models = if snapshot.settings.routing_plugins
&& !inputs.count_only
&& request.pass_through != Some(PassThrough::Listed)
{
let ctx = self.context(snapshot, op, request, name, name, StrategyScope::Default);
self.host
.routing_plugin_models(&ctx, request.pre_routes)
.await?
} else {
None
};
// No deployment for the name: route to the first default fallback.
let default_target = crate::resolve::default_target(&snapshot.settings.dynamic);
// `get_pattern`: a name no pattern matches is tried again under its inferred provider.
let wildcard_name =
crate::resolve::wildcard_name(registry, name, request.specific_deployment);
let wildcard_request = match wildcard_name {
Some(model) => {
let scope = StrategyScope::Default;
let ctx = self.context(snapshot, op, request, model, model, scope);
self.host.provider_request(&ctx).await?
}
None => None,
};
// `endpoint_selection` wraps the attempt, not the count a retry decision reads.
let endpoint = request.endpoint.as_deref().filter(|_| !inputs.count_only);
let caller = Caller {
team_id: request.team_id.as_deref(),
proxy_admin: request.proxy_admin,
allowed_models: &request.allowed_models,
endpoint,
requested: inputs.requested.as_deref().unwrap_or(name),
wildcard_request: wildcard_request.as_deref(),
};
let resolution = resolve_request(
registry,
name,
request.specific_deployment,
caller,
default_target,
snapshot.settings.dynamic.default_ref(),
jiff::Timestamp::now(),
)?
.ok_or_else(|| Error::UnknownModel {
// The name `_common_checks_available_deployment` looked for, past the alias.
model: registry.alias_of(name).unwrap_or(name).to_owned(),
})?;
let first_model = resolution
.deployments
.first()
.map(|d| d.params.model.clone());
let team_id = request.team_id.as_deref();
let usable = filter_team_based(resolution.deployments.clone(), request)?;
// `_filter_reserved_deployments`: PTU shares keep holders to their share and everyone
// else to the deployments split into none.
let (mut all, withheld) = crate::ptu::filter_shared(usable, team_id);
if withheld && all.is_empty() {
return Err(Error::PtuReserved {
model: match resolution.kind {
ResolveKind::Id => first_model.unwrap_or_else(|| name.to_owned()),
ResolveKind::Wildcard | ResolveKind::SpecificDeployment => name.to_owned(),
_ => resolution.group.clone(),
},
});
}
// `filter_web_search_deployments` runs before `request_deployments` is read, so every later
// stage sees the group without the deployments it dropped.
if request.web_search && !inputs.count_only && !resolution.kind.is_single() {
all.retain(|d| d.supports_web_search());
}
// `_without_deployments_that_cannot_serve_route`, right after the web search filter: the
// deployments whose provider cannot serve the call, unless that leaves nothing. Only a pick
// runs it (`_route_capability_pin`): the retry count reads the whole group.
if let Some(handler) = request.route_handler.as_deref()
&& !inputs.count_only
&& !resolution.kind.is_single()
&& !all.is_empty()
{
let ctx = self.context(
snapshot,
op,
request,
&resolution.group,
name,
StrategyScope::Default,
);
let stage = SelectStage::RouteIncapable {
handler: handler.to_owned(),
};
let incapable = self.host.select(&stage, &ctx, all.clone()).await?;
let capable: Vec<Arc<Deployment>> = all
.iter()
.filter(|d| !incapable.iter().any(|i| i.id == d.id))
.cloned()
.collect();
if !capable.is_empty() {
all = capable;
}
}
let listed =
request.pass_through == Some(PassThrough::Listed) && !resolution.kind.is_single();
if listed {
all.retain(|d| uses_pass_through(d));
if all.is_empty() {
// The name `_common_checks_available_deployment` resolved: an alias's target, a
// routing group's name, the alias-resolved name a wildcard serves.
let model = match resolution.kind {
ResolveKind::Wildcard => registry.alias_of(name).unwrap_or(name),
ResolveKind::SpecificDeployment => name,
_ => resolution.group.as_str(),
};
return Err(Error::NoPassThroughDeployments {
model: model.to_owned(),
});
}
}
// A deployment named by its id is served even while it cools down
// (`get_deployment_by_model_id` returns before the cooldown filter).
let settings = &snapshot.settings;
let cooling = if resolution.kind.is_single() {
Vec::new()
} else {
// `async_get_cooldown_deployments` reads every deployment's, as one throttled batch.
self.health
.active_cooldowns(&snapshot.registry.ids())
.await?
};
// Health-check filtering runs before the cooldown filter, except for a deployment named
// by its id and the healthy count a retry decision reads.
let checked = if resolution.kind.is_single() || inputs.count_only {
all.clone()
} else {
let when_all_unhealthy = if request.health_check_probe {
WhenAllUnhealthy::Nothing
} else {
WhenAllUnhealthy::KeepAll
};
self.health
.health_check_filter(settings, all.clone(), when_all_unhealthy)
.await
};
let mut healthy: Vec<Arc<Deployment>> = checked
.iter()
.filter(|d| !cooling.iter().any(|(id, _)| *id == d.id))
.cloned()
.collect();
// The safety net when health-check routing drives cooldowns (an allowed_fails_policy):
// never on a listed pass-through pick or the healthy count.
if healthy.is_empty()
&& !listed
&& !inputs.count_only
&& settings.enable_health_check_routing
&& settings.allowed_fails_policy.is_some()
{
healthy = checked;
}
if resolution.kind.is_single()
&& let Some(deployment) = all.first()
{
// `_core_async_get_healthy_deployments`: a deployment named by id is checked alone.
if deployment.model_info.blocked {
return Err(Error::Paused {
model: deployment.params.model.clone(),
});
}
if request.pass_through.is_some() && !uses_pass_through(deployment) {
return Err(Error::NotPassThrough {
id: deployment.id.clone(),
});
}
}
healthy.retain(|d| !d.model_info.blocked);
// `continuation_deployments`: a mid-stream continuation never lands on a deployment the host
// named as unable to take its prefill.
if let Some(excluded) = request
.body
.get("continuation_excluded_ids")
.and_then(serde_json::Value::as_array)
{
healthy.retain(|d| !excluded.iter().any(|id| id.as_str() == Some(d.id.as_str())));
}
let request_order = request.target_order.filter(|_| !inputs.fallback_hop);
// `async_callback_filter_deployments` runs after the cooldown and blocked filters on every
// async candidate list but a deployment named by id and a listed pass-through pick.
if settings.callback_filters
&& !inputs.count_only
&& !resolution.kind.is_single()
&& request.pass_through != Some(PassThrough::Listed)
{
let mut ctx = self.context(
snapshot,
op,
request,
&resolution.group,
name,
StrategyScope::Default,
);
ctx.excluded_ids = &inputs.excluded_ids;
ctx.skipped_ids = &inputs.skipped_ids;
ctx.target_order = inputs.target_order.or(request_order);
healthy = self.host.callback_filters(&ctx, healthy).await?;
}
if inputs.count_only {
return Ok(Candidates {
name: name.to_owned(),
model_group: resolution.group,
kind: resolution.kind,
all,
healthy,
wildcard_request,
});
}
if listed {
if settings.enable_pre_call_checks && !resolution.kind.is_single() {
healthy = self
.pre_call_checks(
snapshot,
op,
request,
&resolution.group,
name,
&all,
healthy,
)
.await?;
}
// Only a priority group narrows a listed pass-through pick by order.
let priority = resolution.kind == ResolveKind::RoutingGroup
&& registry.routing_groups().iter().any(|g| {
g.group_name == resolution.group && crate::routing_groups::is_priority(g)
});
if priority {
let order = request
.target_order
.or_else(|| healthy.iter().filter_map(|d| d.order()).min());
if let Some(order) = order {
healthy.retain(|d| d.order() == Some(order));
}
}
return Ok(Candidates {
name: name.to_owned(),
model_group: resolution.group,
kind: resolution.kind,
all,
healthy,
wildcard_request,
});
}
if !settings.plugins.is_empty() && !healthy.is_empty() {
let ctx = self.context(
snapshot,
op,
request,
&resolution.group,
name,
StrategyScope::Default,
);
healthy = self
.host
.select(&SelectStage::Plugins, &ctx, healthy)
.await?;
}
if settings.enable_pre_call_checks && !resolution.kind.is_single() {
healthy = self
.pre_call_checks(
snapshot,
op,
request,
&resolution.group,
name,
&all,
healthy,
)
.await?;
}
if request.filter_tags && !healthy.is_empty() {
let mut ctx = self.context(
snapshot,
op,
request,
&resolution.group,
name,
StrategyScope::Default,
);
ctx.pool = &all;
healthy = self.host.select(&SelectStage::Tags, &ctx, healthy).await?;
}
if !resolution.kind.is_single() {
healthy = crate::routing_plugins::narrow(healthy, plugin_models.as_deref())?;
}
// `_without_declared_text_only_deployments`: an image request skips the deployments that
// declare no vision, unless that leaves nothing.
if request.has_image && !resolution.kind.is_single() && !healthy.is_empty() {
let ctx = self.context(
snapshot,
op,
request,
&resolution.group,
name,
StrategyScope::Default,
);
let text_only = self
.host
.select(&SelectStage::TextOnly, &ctx, healthy.clone())
.await?;
let vision: Vec<Arc<Deployment>> = healthy
.iter()
.filter(|d| !text_only.iter().any(|t| t.id == d.id))
.cloned()
.collect();
if !vision.is_empty() {
healthy = vision;
}
}
// `_without_decisions_deployments_refusing_media`: a Decisions call carrying audio or
// videos skips the deployments that refuse them, unless that leaves nothing.
if request.decisions_media && !resolution.kind.is_single() && !healthy.is_empty() {
let ctx = self.context(
snapshot,
op,
request,
&resolution.group,
name,
StrategyScope::Default,
);
let refusing = self
.host
.select(&SelectStage::MediaRefusing, &ctx, healthy.clone())
.await?;
let taking: Vec<Arc<Deployment>> = healthy
.iter()
.filter(|d| !refusing.iter().any(|r| r.id == d.id))
.cloned()
.collect();
if !taking.is_empty() {
healthy = taking;
}
}
if let Some(pinned) = request.pinned_deployment_id.as_deref()
&& healthy.iter().any(|d| d.id == pinned)
{
healthy.retain(|d| d.id == pinned);
}
let order = inputs
.target_order
.or(request_order)
.or_else(|| healthy.iter().filter_map(|d| d.order()).min());
if let Some(order) = order {
healthy.retain(|d| d.order() == Some(order));
}
healthy.retain(|d| {
!request.excluded_ids.contains(&d.id) && !inputs.excluded_ids.contains(&d.id)
});
if !inputs.skipped_ids.is_empty() || !request.retry_skipped_ids.is_empty() {
let kept: Vec<Arc<Deployment>> = healthy
.iter()
.filter(|d| {
!inputs.skipped_ids.contains(&d.id)
&& !request.retry_skipped_ids.contains(&d.id)
})
.cloned()
.collect();
if !kept.is_empty() {
healthy = kept;
}
}
if request.pass_through == Some(PassThrough::Healthy)
&& !resolution.kind.is_single()
&& !healthy.is_empty()
{
healthy.retain(|d| uses_pass_through(d));
if healthy.is_empty() {
return Err(Error::NoPassThroughDeployments {
model: name.to_owned(),
});
}
}
Ok(Candidates {
name: name.to_owned(),
model_group: resolution.group,
kind: resolution.kind,
all,
healthy,
wildcard_request,
})
}
/// `_pre_call_checks` over `healthy`, with the host's facts.
#[allow(clippy::too_many_arguments)]
async fn pre_call_checks(
&self,
snapshot: &RouterStateSnapshot,
op: &Op,
request: &RouterRequest,
model_group: &str,
name: &str,
pool: &[Arc<Deployment>],
healthy: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, Error> {
let mut ctx = self.context(
snapshot,
op,
request,
model_group,
name,
StrategyScope::Default,
);
ctx.pool = pool;
let facts = self.host.pre_call_facts(&ctx, &healthy).await?;
crate::precall::filter(model_group, healthy, pool, &facts)
}
/// Pick one of `candidates.healthy` with the group's strategy.
pub(crate) async fn pick(
&self,
snapshot: &RouterStateSnapshot,
op: &Op,
request: &RouterRequest,
candidates: &Candidates,
) -> Result<Arc<Deployment>, Error> {
let no_deployment = || Error::NoHealthyDeployments {
model: candidates.failed_model(&snapshot.registry),
};
// A deployment named by its id is returned as is, before any strategy runs.
if candidates.kind.is_single() {
return candidates
.healthy
.first()
.cloned()
.ok_or_else(no_deployment);
}
let (strategy, scope) = routing_context(
&snapshot.registry,
&snapshot.settings.routing_strategy,
request.routing_strategy.as_deref(),
&candidates.kind,
&candidates.model_group,
);
if strategy == RoutingStrategy::Unset {
return Err(no_deployment());
}
if strategy == RoutingStrategy::SimpleShuffle {
let weights = crate::weights::validate(request.router_weights.as_ref())
.map_err(|_| Error::InvalidRouterWeights)?;
return simple_shuffle(
&candidates.healthy,
&candidates.model_group,
weights.as_ref(),
self.rng.as_ref(),
)?
.ok_or_else(no_deployment);
}
let ctx = self.context(
snapshot,
op,
request,
&candidates.model_group,
&candidates.name,
scope,
);
self.host
.select(
&SelectStage::Strategy(strategy),
&ctx,
candidates.healthy.clone(),
)
.await?
.into_iter()
.next()
.ok_or_else(no_deployment)
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
use crate::registry::DeploymentSpec;
struct Fixed(f64);
impl Randomness for Fixed {
fn next_f64(&self) -> f64 {
self.0
}
}
fn deployment(id: &str, params: serde_json::Value) -> Arc<Deployment> {
let spec: DeploymentSpec = serde_json::from_value(json!({
"model_name": "g", "litellm_params": params, "model_info": {"id": id},
}))
.unwrap();
Arc::new(Deployment::from_spec(spec).unwrap())
}
#[rstest]
#[case::weight_low_draw(0.1, "a")]
#[case::weight_high_draw(0.9, "b")]
fn shuffle_uses_weight_before_rpm(#[case] draw: f64, #[case] expected: &str) {
let deployments = [
deployment("a", json!({"model": "openai/a", "weight": 1, "rpm": 1000})),
deployment("b", json!({"model": "openai/b", "weight": 3})),
];
let picked = simple_shuffle(&deployments, "g", None, &Fixed(draw))
.unwrap()
.unwrap();
assert_eq!(picked.id, expected);
}
#[rstest]
fn shuffle_prefers_request_weights() {
let deployments = [
deployment("a", json!({"model": "openai/a", "weight": 100})),
deployment("b", json!({"model": "openai/b"})),
];
let weights = BTreeMap::from([("g".to_owned(), BTreeMap::from([("b".to_owned(), 1.0)]))]);
let picked = simple_shuffle(&deployments, "g", Some(&weights), &Fixed(0.99))
.unwrap()
.unwrap();
assert_eq!(picked.id, "b");
}
}

View file

@ -0,0 +1,133 @@
//! [`Settings`]: one typed, frozen view of the configuration a caller may read, built from the
//! static config and the current dynamic snapshot.
use std::{collections::BTreeMap, sync::Arc, time::Duration};
use crate::{
config::{
AllowedFailsPolicy, DEFAULT_COOLDOWN_TIME, PluginSpec, RetryPolicy, RouterConfig,
RouterDynamicConfig, RoutingStrategy,
},
error::AttemptError,
};
#[derive(Clone, Debug)]
pub struct Settings {
pub routing_strategy: RoutingStrategy,
pub plugins: Arc<[PluginSpec]>,
pub cooldown_time: Duration,
pub disable_cooldowns: bool,
pub allowed_fails: Option<u32>,
pub allowed_fails_policy: Option<AllowedFailsPolicy>,
pub retry_policy: Option<RetryPolicy>,
pub model_group_retry_policy: Arc<BTreeMap<String, RetryPolicy>>,
pub num_retries: u32,
pub max_fallbacks: u32,
pub retry_after: Duration,
pub timeout: Option<Duration>,
pub drop_params: Option<bool>,
pub enable_pre_call_checks: bool,
pub routing_plugins: bool,
pub callback_filters: bool,
pub enable_health_check_routing: bool,
pub health_check_staleness: Duration,
pub background_health_check_model_groups: Option<Arc<[String]>>,
pub enable_weighted_failover: bool,
pub fallback_on_unsupported_params: bool,
pub provider_default_deployment_ids: Arc<[String]>,
pub dynamic: Arc<RouterDynamicConfig>,
}
/// `DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER`.
pub const DEFAULT_HEALTH_CHECK_STALENESS: Duration = Duration::from_secs(300 * 2);
fn seconds(value: f64) -> Duration {
Duration::try_from_secs_f64(value).unwrap_or_default()
}
impl Settings {
pub fn new(config: &RouterConfig, dynamic: Arc<RouterDynamicConfig>) -> Self {
Self {
routing_strategy: config.routing_strategy.clone(),
plugins: config.plugins.clone().into(),
cooldown_time: config.cooldown_time.map_or(DEFAULT_COOLDOWN_TIME, seconds),
disable_cooldowns: config.disable_cooldowns,
allowed_fails: dynamic.allowed_fails,
allowed_fails_policy: config.allowed_fails_policy.clone(),
retry_policy: config.retry_policy.clone(),
model_group_retry_policy: Arc::new(config.model_group_retry_policy.clone()),
num_retries: dynamic.num_retries.unwrap_or(0),
max_fallbacks: dynamic.max_fallbacks,
retry_after: seconds(config.retry_after),
timeout: config.request_timeout.or(config.timeout).map(seconds),
drop_params: config.drop_params,
enable_pre_call_checks: config.enable_pre_call_checks,
routing_plugins: config.routing_plugins,
callback_filters: config.callback_filters,
enable_health_check_routing: config.enable_health_check_routing,
health_check_staleness: config
.health_check_staleness_threshold
.map_or(DEFAULT_HEALTH_CHECK_STALENESS, seconds),
background_health_check_model_groups: config
.background_health_check_model_groups
.clone()
.map(Into::into),
enable_weighted_failover: config.enable_weighted_failover,
fallback_on_unsupported_params: config.fallback_on_unsupported_params.unwrap_or(true),
provider_default_deployment_ids: config.provider_default_deployment_ids.clone().into(),
dynamic,
}
}
/// `allowed_fails_for`: the policy's count for this failure.
pub fn allowed_fails_for(&self, failure: &AttemptError) -> Option<u32> {
self.allowed_fails_policy
.as_ref()
.and_then(|policy| policy.allowed_fails_for(failure))
}
/// How long an allowed-fails counter lives: `settings().cooldown_time or
/// DEFAULT_COOLDOWN_TIME_SECONDS`, so a router cooldown of 0 still counts over the default.
pub fn fails_window(&self) -> Duration {
if self.cooldown_time.is_zero() {
DEFAULT_COOLDOWN_TIME
} else {
self.cooldown_time
}
}
/// `_is_allowed_fails_set_on_router`.
pub fn allowed_fails_set(&self) -> bool {
self.allowed_fails.is_some() && !self.dynamic.allowed_fails_is_default
}
/// Python's `allowed_fails` default when set without a value (`litellm.allowed_fails`).
pub fn effective_allowed_fails(&self) -> u32 {
self.allowed_fails.unwrap_or(3)
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::passed_to_the_router(Some(3), false, true)]
#[case::the_process_default(Some(3), true, false)]
#[case::none(None, false, false)]
fn allowed_fails_counts_as_set_unless_it_is_the_process_default(
#[case] allowed_fails: Option<u32>,
#[case] allowed_fails_is_default: bool,
#[case] expected: bool,
) {
let dynamic = RouterDynamicConfig {
allowed_fails,
allowed_fails_is_default,
..RouterDynamicConfig::default()
};
let settings = Settings::new(&RouterConfig::default(), Arc::new(dynamic));
assert_eq!(settings.allowed_fails_set(), expected);
}
}

View file

@ -0,0 +1,110 @@
//! Which routing strategy picks for a request, and whose state it reads: a port of
//! `Router._get_routing_context`. A per-request `routing_strategy` override wins, then the
//! routing group the resolved model belongs to (called as the group, or as a member of a
//! non-priority group), then the router's own strategy.
use crate::{
config::RoutingStrategy,
registry::{Registry, ResolveKind},
};
/// The per-request overrides `_OVERRIDABLE_ROUTING_STRATEGIES` honours: simple-shuffle and every
/// strategy with a default selector. Anything else is ignored.
pub const OVERRIDABLE: [&str; 6] = [
"simple-shuffle",
"least-busy",
"usage-based-routing",
"usage-based-routing-v2",
"latency-based-routing",
"cost-based-routing",
];
/// Whose selector state a strategy pick reads.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum StrategyScope {
/// The router's own strategy (the implicit `"default"` group).
Default,
/// A routing group's own selector.
Group(String),
/// A per-request override: the default selector when it names the router's strategy, else
/// one selector per strategy kept for the router's life.
Override,
}
impl StrategyScope {
/// The scope as the host step spells it.
pub fn wire(&self) -> (&'static str, Option<&str>) {
match self {
Self::Default => ("default", None),
Self::Group(name) => ("group", Some(name)),
Self::Override => ("override", None),
}
}
}
/// A request's `routing_strategy` override, when it is one Python honours.
pub fn request_override(raw: Option<&str>) -> Option<RoutingStrategy> {
let raw = raw?;
OVERRIDABLE
.contains(&raw)
.then(|| raw.parse().unwrap_or_default())
}
/// The strategy and selector scope for a pick of `model_group` (the resolved group, after
/// aliases), reached through a resolution of `kind`.
pub fn routing_context(
registry: &Registry,
default: &RoutingStrategy,
override_strategy: Option<&str>,
kind: &ResolveKind,
model_group: &str,
) -> (RoutingStrategy, StrategyScope) {
if let Some(strategy) = request_override(override_strategy) {
return (strategy, StrategyScope::Override);
}
let group = match kind {
ResolveKind::RoutingGroup => registry
.routing_groups()
.iter()
.find(|group| group.group_name == model_group),
ResolveKind::Alias | ResolveKind::Group => registry.member_routing_group(model_group),
ResolveKind::SpecificDeployment
| ResolveKind::Id
| ResolveKind::Default
| ResolveKind::Wildcard => None,
};
match group {
None => (default.clone(), StrategyScope::Default),
Some(group) if is_priority(&group.routing_strategy) => {
(RoutingStrategy::SimpleShuffle, StrategyScope::Default)
}
Some(group) => (
group.routing_strategy.clone(),
StrategyScope::Group(group.group_name.clone()),
),
}
}
fn is_priority(strategy: &RoutingStrategy) -> bool {
matches!(strategy, RoutingStrategy::Other(name) if name == "priority")
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::honoured(Some("least-busy"), Some(RoutingStrategy::LeastBusy))]
#[case::shuffle(Some("simple-shuffle"), Some(RoutingStrategy::SimpleShuffle))]
#[case::lar1_ignored(Some("lar1"), None)]
#[case::unknown_ignored(Some("bogus"), None)]
#[case::unset(None, None)]
fn only_overridable_strategies_override(
#[case] raw: Option<&str>,
#[case] expected: Option<RoutingStrategy>,
) {
assert_eq!(request_override(raw), expected);
}
}

View file

@ -0,0 +1,81 @@
//! `_router_weights` as `validate_router_weights` reads it (`litellm/types/router_weights.py`):
//! model group to deployment id to a finite weight of at least 0 (an int or a float, never a
//! bool or a string), with non-blank ids and at least one positive weight in every non-empty group.
use std::collections::BTreeMap;
use serde_json::Value;
/// Validated request weights.
pub type RouterWeights = BTreeMap<String, BTreeMap<String, f64>>;
/// The request's weights didn't validate; the host raises the validator's own error.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct InvalidRouterWeights;
fn identifier(text: &str) -> bool {
text.chars().any(|c| !c.is_whitespace())
}
fn weight(value: &Value) -> Option<f64> {
match value {
Value::Number(number) => number.as_f64().filter(|w| w.is_finite() && *w >= 0.0),
_ => None,
}
}
/// `validate_router_weights(raw)`: `None` when the request passed none.
pub fn validate(raw: Option<&Value>) -> Result<Option<RouterWeights>, InvalidRouterWeights> {
let groups = match raw {
None | Some(Value::Null) => return Ok(None),
Some(Value::Object(groups)) => groups,
Some(_) => return Err(InvalidRouterWeights),
};
let mut validated = RouterWeights::new();
for (group, weights) in groups {
let Value::Object(weights) = weights else {
return Err(InvalidRouterWeights);
};
if !identifier(group) {
return Err(InvalidRouterWeights);
}
let mut entries = BTreeMap::new();
for (id, value) in weights {
if !identifier(id) {
return Err(InvalidRouterWeights);
}
entries.insert(id.clone(), weight(value).ok_or(InvalidRouterWeights)?);
}
if !entries.is_empty() && !entries.values().any(|w| *w > 0.0) {
return Err(InvalidRouterWeights);
}
validated.insert(group.clone(), entries);
}
Ok(Some(validated))
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::all_zero(json!({"g": {"a": 0.0, "b": 0.0}}))]
#[case::negative(json!({"g": {"a": -1.0}}))]
#[case::text(json!({"g": {"a": "1"}}))]
#[case::boolean(json!({"g": {"a": true}}))]
#[case::blank_id(json!({"g": {" ": 1.0}}))]
#[case::not_a_mapping(json!(["g"]))]
fn invalid_weights_are_refused(#[case] raw: Value) {
assert_eq!(validate(Some(&raw)), Err(InvalidRouterWeights));
}
#[rstest]
#[case::empty_group(json!({"g": {}}))]
#[case::int_weight(json!({"g": {"a": 1, "b": 0}}))]
fn valid_weights_are_kept(#[case] raw: Value) {
assert!(validate(Some(&raw)).unwrap().is_some());
}
}

File diff suppressed because it is too large Load diff

View file

@ -1,59 +0,0 @@
"""Public accounting contracts. Importing these values does not enable accounting."""
from litellm.accounting.measurements import measurements_for
from litellm.accounting.protocols import CostAggregator
from litellm.accounting.types import (
AdmissionDecision,
AggregateQuery,
AggregateRow,
CostEvent,
DeliveryState,
DeliveryTarget,
Dimension,
DimensionFilter,
DimensionName,
EventKind,
Granularity,
InboundUsageEvent,
Measurement,
Metric,
MetricName,
Multiplier,
NamedCost,
OutboundUsageEvent,
Outcome,
PricedQuantity,
PricingSnapshot,
Rate,
Unpriced,
validate_aggregate_rows,
)
__all__ = [
"AdmissionDecision",
"AggregateQuery",
"AggregateRow",
"CostAggregator",
"CostEvent",
"DeliveryState",
"DeliveryTarget",
"Dimension",
"DimensionFilter",
"DimensionName",
"EventKind",
"Granularity",
"InboundUsageEvent",
"Measurement",
"Metric",
"MetricName",
"Multiplier",
"NamedCost",
"OutboundUsageEvent",
"Outcome",
"PricedQuantity",
"PricingSnapshot",
"Rate",
"Unpriced",
"measurements_for",
"validate_aggregate_rows",
]

View file

@ -0,0 +1,71 @@
from typing import TYPE_CHECKING, Final
from litellm.accounting.delivery_projection import entity_transactions, project_deliveries
from litellm.accounting.projection import Contribution
from litellm.proxy.db.budget_window_spend_writer import commit_window_spend_updates
from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key
from litellm.proxy.db.db_span import db_span
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction
from litellm.proxy.db.model_usage_rollup import queue_model_usage_upsert
from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction
if TYPE_CHECKING:
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.utils import PrismaClient, ProxyLogging
async def apply_deliveries(
transaction: LedgerTransaction,
deliveries: tuple[DeliveryProjection, ...],
*,
prisma_client: "PrismaClient",
writer: "DBSpendUpdateWriter",
proxy_logging_obj: "ProxyLogging",
) -> tuple[Contribution, ...]:
from litellm.proxy.guardrails.usage_tracking import UsageUnitIncrement, UsageUnitKey, upsert_usage_unit_row
contributions: Final = project_deliveries(deliveries)
await writer.apply_accounting_entities(
prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions)
)
for table in sorted(DAILY_SPEND_TABLES.values(), key=lambda item: item.name):
rows = tuple(
{**dict(item.key), **dict(item.values), **dict(item.increments)}
for item in contributions
if item.table == table.name
)
merged = merge_by_conflict_key(table, rows)
for offset in range(0, len(merged), 100):
sql, params = build_bulk_upsert(table, merged[offset : offset + 100])
async with db_span("upsert_daily_spend", table.name):
await transaction.execute_raw(sql, *params)
async with transaction.batch_() as batcher:
for item in contributions:
if item.table == "LiteLLM_DailyModelUsage":
queue_model_usage_upsert(batcher, dict(item.key), dict(item.increments))
for item in contributions:
if item.table == "LiteLLM_DailyGuardrailUsageUnits":
key = dict(item.key)
counters = dict(item.increments)
await upsert_usage_unit_row(
prisma_client,
UsageUnitKey(key["guardrail_id"], key["date"], key["team_id"], key["api_key"], key["usage_unit"]),
UsageUnitIncrement(int(counters["units"]), float(counters["cost"]), int(counters["untracked_units"])),
outer_transaction=transaction,
)
windows: Final = tuple(_window(item) for item in contributions if item.table == "LiteLLM_BudgetWindowSpend")
await commit_window_spend_updates(prisma_client, windows, outer_transaction=transaction)
return contributions
def _window(item: Contribution) -> WindowSpendTransaction:
key: Final = dict(item.key)
values: Final = dict(item.values)
return WindowSpendTransaction(
entity_type=key["entity_type"],
entity_id=key["entity_id"],
window_duration=key["window_duration"],
window_start=values["window_start"] or "",
spend=float(dict(item.increments)["spend"]),
started_at=values["started_at"],
)

View file

@ -0,0 +1,78 @@
from collections.abc import Iterable
from itertools import chain, groupby
from typing import TYPE_CHECKING, Final
from urllib.parse import quote
from litellm.accounting.projection import Contribution, project
from litellm.repositories.accounting_delivery_repository import DeliveryProjection
def _destination(contribution: Contribution) -> tuple[str, tuple[tuple[str, str], ...], str]:
return (
contribution.table,
contribution.key,
(dict(contribution.values).get("window_start") or "")
if contribution.table == "LiteLLM_BudgetWindowSpend"
else "",
)
def _sum(contributions: tuple[Contribution, ...]) -> Contribution:
first: Final = contributions[0]
increments: Final = tuple(
(column, sum(dict(item.increments).get(column, 0) for item in contributions))
for column in dict.fromkeys(
column for column, _ in chain.from_iterable(item.increments for item in contributions)
)
)
starts: Final = tuple(
value
for column, value in chain.from_iterable(item.values for item in contributions)
if column == "started_at" and value
)
values: Final = tuple(
(column, min(starts) if column == "started_at" and starts else value) for column, value in first.values
)
return Contribution(first.table, first.key, increments, values)
def project_deliveries(deliveries: Iterable[DeliveryProjection]) -> tuple[Contribution, ...]:
ordered: Final = sorted(chain.from_iterable(project(delivery.row) for delivery in deliveries), key=_destination)
return tuple(_sum(tuple(group)) for _, group in groupby(ordered, key=_destination))
def entity_transactions(contributions: tuple[Contribution, ...]) -> "DBSpendUpdateTransactions":
return {
"user_list_transactions": dict(_entity_totals(contributions, "LiteLLM_UserTable")),
"end_user_list_transactions": dict(_entity_totals(contributions, "LiteLLM_EndUserTable")),
"key_list_transactions": dict(_entity_totals(contributions, "LiteLLM_VerificationToken")),
"team_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TeamTable")),
"team_member_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TeamMembership")),
"org_list_transactions": dict(_entity_totals(contributions, "LiteLLM_OrganizationTable")),
"org_member_list_transactions": dict(_entity_totals(contributions, "LiteLLM_OrganizationMembership")),
"project_list_transactions": dict(_entity_totals(contributions, "LiteLLM_ProjectTable")),
"tag_list_transactions": dict(_entity_totals(contributions, "LiteLLM_TagTable")),
"agent_list_transactions": dict(_entity_totals(contributions, "LiteLLM_AgentsTable")),
"model_access_group_list_transactions": dict(
_entity_totals(contributions, "LiteLLM_ModelAccessGroupBudgetTable")
),
}
def _entity_totals(contributions: tuple[Contribution, ...], table: str) -> tuple[tuple[str, float], ...]:
return tuple(
(_entity_key(item), float(dict(item.increments)["spend"])) for item in contributions if item.table == table
)
def _entity_key(item: Contribution) -> str:
fields: Final = dict(item.key)
if item.table == "LiteLLM_TeamMembership":
return f"team_id::{fields['team_id']}::user_id::{fields['user_id']}"
if item.table == "LiteLLM_OrganizationMembership":
return f"org_id::{quote(fields['organization_id'], safe='')}::user_id::{quote(fields['user_id'], safe='')}"
return item.key[0][1]
if TYPE_CHECKING:
from litellm.proxy._types import DBSpendUpdateTransactions

View file

@ -1,13 +0,0 @@
from typing import Final
from litellm.accounting.types import CostEvent, InboundUsageEvent, Measurement, OutboundUsageEvent
def measurements_for(event: InboundUsageEvent | OutboundUsageEvent | CostEvent) -> tuple[Measurement, ...]:
"""Return the complete contribution set for one event, ready for one apply call."""
metrics: Final = event.charges if isinstance(event, CostEvent) else event.metrics
occurred_at: Final = event.occurred_at if isinstance(event, CostEvent) else event.started_at
return tuple(
Measurement(source_event_id=event.event_id, occurred_at=occurred_at, dimensions=event.dimensions, metric=metric)
for metric in metrics or ()
)

View file

@ -0,0 +1,30 @@
from typing import Final
from prometheus_client import Gauge, Histogram
from litellm.repositories.accounting_delivery_repository import DeliveryMetrics
_PENDING: Final = Gauge("litellm_accounting_pending", "Pending SpendLogs deliveries", multiprocess_mode="livemax")
_PARKED: Final = Gauge("litellm_accounting_parked", "Parked SpendLogs deliveries", multiprocess_mode="livemax")
_OLDEST: Final = Gauge(
"litellm_accounting_oldest_pending_seconds",
"Age of the oldest pending SpendLogs delivery",
multiprocess_mode="livemax",
)
_BATCH_TIME: Final = Histogram("litellm_accounting_batch_seconds", "SpendLogs delivery worker batch time")
_BATCH_COUNT: Final = Gauge(
"litellm_accounting_batch_deliveries", "Retained batch deliveries", multiprocess_mode="livemax"
)
_TABLE_BYTES: Final = Gauge(
"litellm_accounting_delivery_table_bytes", "Delivery table size in bytes", multiprocess_mode="livemax"
)
def record_deliveries(metrics: DeliveryMetrics, seconds: float) -> None:
_PENDING.set(metrics.pending_count)
_PARKED.set(metrics.parked_count)
_OLDEST.set(max(0.0, metrics.oldest_pending_seconds))
_BATCH_TIME.observe(seconds)
_BATCH_COUNT.set(metrics.batch_count)
_TABLE_BYTES.set(metrics.table_bytes)

View file

@ -1,5 +1,5 @@
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from dataclasses import dataclass, replace
from datetime import datetime, timezone
from itertools import chain
from types import MappingProxyType
@ -24,8 +24,11 @@ from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
_FIELDS: Final = TypeAdapter(dict[str, object])
_STRINGS: Final = TypeAdapter(tuple[str, ...])
_VALUES: Final = TypeAdapter(tuple[object, ...])
_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...])
_TASK_TYPES: Final = frozenset(load_model_insight_tasks())
_MAX_INDEXED_TEXT_BYTES: Final = 2000
_MAX_COUNTER: Final = 2**40
@dataclass(frozen=True, slots=True)
@ -54,7 +57,8 @@ def _tags(row: Mapping[str, object]) -> tuple[str, ...]:
value: Final = row.get("request_tags")
if value is None:
return ()
return _STRINGS.validate_json(value) if isinstance(value, str) else _STRINGS.validate_python(value)
parsed: Final = _VALUES.validate_json(value) if isinstance(value, str) else _VALUES.validate_python(value)
return tuple(tag for tag in parsed if isinstance(tag, str))
def _date(row: Mapping[str, object]) -> str | None:
@ -191,13 +195,31 @@ def _window_contributions(row: Mapping[str, object], metadata: Mapping[str, obje
def project(row: Mapping[str, object]) -> tuple[Contribution, ...]:
metadata: Final = _fields(row.get("metadata"))
return (
contributions: Final = (
*_entity_contributions(row, metadata),
*_daily_contributions(row, metadata),
*_window_contributions(row, metadata),
*_model_contributions(row, metadata),
*_guardrail_contributions(row, metadata),
)
return tuple(_capped(contribution) for contribution in contributions if _indexable(contribution))
def _indexable(contribution: Contribution) -> bool:
texts: Final = (value for _, value in (*contribution.key, *contribution.values) if value is not None)
return all(len(text.encode("utf-8", "surrogatepass")) <= _MAX_INDEXED_TEXT_BYTES for text in texts)
def _capped(contribution: Contribution) -> Contribution:
"""Counters are summed into bigint columns in the charge's transaction, so one capped far below that range
can never overflow a day's total."""
return replace(
contribution,
increments=tuple(
(name, max(-_MAX_COUNTER, min(_MAX_COUNTER, value)) if isinstance(value, int) else value)
for name, value in contribution.increments
),
)
def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, object]) -> tuple[Contribution, ...]:

View file

@ -1,21 +0,0 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import Protocol
from litellm.accounting.types import AggregateQuery, AggregateRow
class CostAggregator(Protocol):
async def apply(self, rows: tuple[Mapping[str, object], ...]) -> None:
"""Apply stored request contributions atomically with their delivery receipts."""
...
async def query(self, query: AggregateQuery) -> tuple[AggregateRow, ...]:
"""Return disjoint [start, end) rows per group using existing daily projections.
Filters combine with AND, values within a filter with IN. DAY emits daily
buckets; None emits whole-interval totals. Reject unsupported metric/scope
combinations. Different groups may overlap, including non-additive tags.
"""
...

View file

@ -2,12 +2,13 @@ from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens
from litellm.proxy.spend_tracking.savings import compute_savings_spend, marks_gateway_injection
from litellm.proxy.spend_tracking.savings import SavingsSpend, compute_savings_spend, marks_gateway_injection
if TYPE_CHECKING:
from litellm.router import Router
@ -66,7 +67,47 @@ def build_accounting_payload(
router: "Callable[[], Router | None]",
) -> Mapping[str, object]:
metadata: Final = _FIELDS.validate_json(payload["metadata"])
savings: Final = compute_savings_spend(
savings: Final = _savings_or_zero(payload, metadata, router)
accounting_metadata: Final = {
**{key: _accounting_metadata_field(key, value) for key, value in metadata.items() if key in _METADATA_FIELDS},
"accounting_entities": entities,
"accounting_savings": savings._asdict(),
"accounting_model_access_groups": model_access_groups,
"accounting_windows": windows,
}
return _FIELDS.validate_python(
safe_json_structure(
{
**{
key: value.isoformat() if isinstance(value, datetime) else value
for key, value in payload.items()
if key in _ACCOUNTING_FIELDS
},
"metadata": accounting_metadata,
},
non_finite_as_null=True,
)
)
def _savings_or_zero(
payload: SpendLogsPayload, metadata: Mapping[str, object], router: "Callable[[], Router | None]"
) -> SavingsSpend:
"""Savings are analytics computed from caller-supplied counts, so one that cannot be computed is stored as zero
rather than losing the charge it rides with."""
try:
return _savings(payload, metadata, router)
except (ArithmeticError, ValueError, TypeError) as error:
verbose_proxy_logger.warning(
"Recording zero savings for request_id=%s, savings could not be computed: %s", payload["request_id"], error
)
return SavingsSpend(compression=0.0, prompt_caching=0.0)
def _savings(
payload: SpendLogsPayload, metadata: Mapping[str, object], router: "Callable[[], Router | None]"
) -> SavingsSpend:
return compute_savings_spend(
model=payload["model"],
custom_llm_provider=payload["custom_llm_provider"],
compression_saved_tokens=extract_compression_saved_tokens(metadata),
@ -80,18 +121,41 @@ def build_accounting_payload(
recorded_autorouter_savings_estimate=_mapping(metadata.get("autorouter_savings_estimate")),
billed_at=payload["endTime"],
)
accounting_metadata: Final = {
**{key: value for key, value in metadata.items() if key in _METADATA_FIELDS},
"accounting_entities": entities,
"accounting_savings": savings._asdict(),
"accounting_model_access_groups": model_access_groups,
"accounting_windows": windows,
_GUARDRAIL_ENTRIES: Final = TypeAdapter(tuple[dict[str, object], ...])
_GUARDRAIL_METADATA_FIELDS: Final = frozenset({"guardrail_information", "standard_logging_guardrail_information"})
_GUARDRAIL_ACCOUNTING_FIELDS: Final = frozenset(
{
"guardrail_id",
"guardrail_name",
"guardrail_provider",
"guardrail_usage",
"guardrail_cost_by_unit",
"guardrail_cost_in_spend",
}
)
_GUARDRAIL_RESPONSE: Final = TypeAdapter(dict[str, object])
def _accounting_guardrail_entries(value: object) -> object:
try:
entries: Final = _GUARDRAIL_ENTRIES.validate_python((value,) if isinstance(value, dict) else value)
except ValidationError:
return None
return [_accounting_guardrail_entry(entry) for entry in entries]
def _accounting_guardrail_entry(entry: Mapping[str, object]) -> Mapping[str, object]:
kept: Final = {key: item for key, item in entry.items() if key in _GUARDRAIL_ACCOUNTING_FIELDS}
response: Final = entry.get("guardrail_response")
if not isinstance(response, dict) or "tokens_saved" not in response:
return kept
return {
**{
key: value.isoformat() if isinstance(value, datetime) else value
for key, value in payload.items()
if key in _ACCOUNTING_FIELDS
},
"metadata": _FIELDS.validate_python(safe_json_structure(accounting_metadata, non_finite_as_null=True)),
**kept,
"guardrail_response": {"tokens_saved": _GUARDRAIL_RESPONSE.validate_python(response)["tokens_saved"]},
}
def _accounting_metadata_field(key: str, value: object) -> object:
return _accounting_guardrail_entries(value) if key in _GUARDRAIL_METADATA_FIELDS else value

View file

@ -1,428 +0,0 @@
"""Accounting values use UTC instants, decimal strings on the wire and USD charges."""
from collections.abc import Mapping
from datetime import datetime, timezone
from decimal import Decimal
from enum import Enum
from typing import Annotated, Final, Literal, TypeAlias
from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field, model_validator
from typing_extensions import Self
from litellm.types.utils import CallTypes
_CONFIG: Final = ConfigDict(frozen=True, extra="forbid", validate_default=True)
Identifier: TypeAlias = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")]
Quantity: TypeAlias = Annotated[Decimal, Field(strict=True, ge=0, allow_inf_nan=False)]
def _utc(value: datetime) -> datetime:
return value.astimezone(timezone.utc)
Instant: TypeAlias = Annotated[AwareDatetime, AfterValidator(_utc)]
class MetricName(str, Enum):
"""Reasoning tokens are included in completion tokens; above-1hr writes are included in cache creation tokens."""
REQUESTS = "requests"
SUCCESSFUL_REQUESTS = "successful_requests"
FAILED_REQUESTS = "failed_requests"
PROMPT_TOKENS = "prompt_tokens"
COMPLETION_TOKENS = "completion_tokens"
REASONING_TOKENS = "reasoning_tokens"
TOTAL_TOKENS = "total_tokens"
CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens"
CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens"
CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR = "cache_creation_input_tokens_above_1hr"
SPEND = "spend"
class DimensionName(str, Enum):
KEY = "key"
USER = "user"
TEAM = "team"
ORGANIZATION = "organization"
END_USER = "end_user"
MODEL = "model"
MODEL_GROUP = "model_group"
DEPLOYMENT = "deployment"
PROVIDER = "provider"
TOOL = "tool"
AGENT = "agent"
PROJECT = "project"
MODEL_ACCESS_GROUP = "model_access_group"
API_BASE = "api_base"
SERVICE_TIER = "service_tier"
TAG = "tag"
class Outcome(str, Enum):
SUCCEEDED = "succeeded"
FAILED = "failed"
CANCELLED = "cancelled"
class Granularity(str, Enum):
DAY = "day"
class EventKind(str, Enum):
INBOUND_USAGE = "inbound_usage"
OUTBOUND_USAGE = "outbound_usage"
COST = "cost"
class DeliveryState(str, Enum):
PENDING = "pending"
APPLIED = "applied"
PARKED = "parked"
class DeliveryTarget(str, Enum):
"""One member per existing table the projector writes; each is applied and receipted independently."""
KEY_SPEND = "LiteLLM_VerificationToken"
USER_SPEND = "LiteLLM_UserTable"
TEAM_SPEND = "LiteLLM_TeamTable"
TEAM_MEMBER_SPEND = "LiteLLM_TeamMembership"
ORGANIZATION_SPEND = "LiteLLM_OrganizationTable"
ORGANIZATION_MEMBER_SPEND = "LiteLLM_OrganizationMembership"
END_USER_SPEND = "LiteLLM_EndUserTable"
AGENT_SPEND = "LiteLLM_AgentsTable"
TAG_SPEND = "LiteLLM_TagTable"
PROJECT_SPEND = "LiteLLM_ProjectTable"
MODEL_ACCESS_GROUP_SPEND = "LiteLLM_ModelAccessGroupBudgetTable"
BUDGET_SPEND = "LiteLLM_BudgetTable"
BUDGET_WINDOW_SPEND = "LiteLLM_BudgetWindowSpend"
DAILY_USER_SPEND = "LiteLLM_DailyUserSpend"
DAILY_TEAM_SPEND = "LiteLLM_DailyTeamSpend"
DAILY_ORGANIZATION_SPEND = "LiteLLM_DailyOrganizationSpend"
DAILY_END_USER_SPEND = "LiteLLM_DailyEndUserSpend"
DAILY_AGENT_SPEND = "LiteLLM_DailyAgentSpend"
DAILY_TAG_SPEND = "LiteLLM_DailyTagSpend"
DAILY_MODEL_USAGE = "LiteLLM_DailyModelUsage"
DAILY_TOOL_SPEND = "LiteLLM_DailyToolSpend"
AUTOROUTER_DAILY_SPEND = "LiteLLM_AutoRouterDailySpend"
_SUBSET_PAIRS: Final = (
(MetricName.REASONING_TOKENS, MetricName.COMPLETION_TOKENS),
(MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, MetricName.CACHE_CREATION_INPUT_TOKENS),
)
def _check_subsets(values: Mapping[str, Decimal]) -> None:
for part, whole in _SUBSET_PAIRS:
if part.value in values and whole.value in values and values[part.value] > values[whole.value]:
raise ValueError(f"{part.value} cannot exceed {whole.value}")
class Metric(BaseModel):
model_config = _CONFIG
name: MetricName
value: Quantity
unit: Literal["count", "USD"] = "count"
@model_validator(mode="after")
def validate_unit(self) -> Self:
if (self.name is MetricName.SPEND) != (self.unit == "USD"):
raise ValueError("SPEND requires USD; usage metrics require count")
if self.unit == "count" and self.value != self.value.to_integral_value():
raise ValueError("Count metrics require integral values")
return self
class Dimension(BaseModel):
model_config = _CONFIG
name: DimensionName
value: Identifier
_REPEATABLE_DIMENSIONS: Final = frozenset({DimensionName.TAG, DimensionName.MODEL_ACCESS_GROUP})
def _dimensions(values: tuple[Dimension, ...]) -> tuple[Dimension, ...]:
if len(frozenset(values)) != len(values):
raise ValueError("Duplicate dimension values")
single_names: Final = tuple(value.name for value in values if value.name not in _REPEATABLE_DIMENSIONS)
if len(frozenset(single_names)) != len(single_names):
raise ValueError("Only tag and model access group dimensions may have multiple values")
return values
Dimensions: TypeAlias = Annotated[tuple[Dimension, ...], AfterValidator(_dimensions)]
_CALLER_METRICS: Final = frozenset((MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS))
def _metrics(values: tuple[Metric, ...]) -> tuple[Metric, ...]:
if len(frozenset(value.name for value in values)) != len(values):
raise ValueError("Duplicate metric names")
return values
Metrics: TypeAlias = Annotated[tuple[Metric, ...], AfterValidator(_metrics)]
def _interval(start: datetime, end: datetime, *, allow_empty: bool = False) -> None:
if end < start or (end == start and not allow_empty):
raise ValueError("End must follow start")
class InboundUsageEvent(BaseModel):
model_config = _CONFIG
event_id: Identifier
request_id: Identifier
call_type: CallTypes
started_at: Instant
ended_at: Instant
outcome: Outcome
dimensions: Dimensions
metrics: Metrics
@model_validator(mode="after")
def validate_caller_activity(self) -> Self:
_interval(self.started_at, self.ended_at, allow_empty=True)
if any(metric.name not in _CALLER_METRICS for metric in self.metrics):
raise ValueError("Inbound events own only caller request metrics")
if not any(metric.name is MetricName.REQUESTS and metric.value == 1 for metric in self.metrics):
raise ValueError("Inbound events count exactly one caller request")
for metric in self.metrics:
if metric.name is MetricName.SUCCESSFUL_REQUESTS and metric.value != int(self.outcome is Outcome.SUCCEEDED):
raise ValueError("Successful request count must match outcome")
if metric.name is MetricName.FAILED_REQUESTS and metric.value != int(self.outcome is not Outcome.SUCCEEDED):
raise ValueError("Failed request count must match outcome")
return self
class OutboundUsageEvent(BaseModel):
model_config = _CONFIG
event_id: Identifier
attempt_id: Identifier
request_id: Identifier | None
retry_of_attempt_id: Identifier | None
call_type: CallTypes
started_at: Instant
ended_at: Instant
outcome: Outcome
dimensions: Dimensions
metrics: Metrics | None
@model_validator(mode="after")
def validate_attempt_usage(self) -> Self:
_interval(self.started_at, self.ended_at, allow_empty=True)
if self.retry_of_attempt_id == self.attempt_id:
raise ValueError("An attempt cannot retry itself")
if self.metrics is not None and any(
metric.name in _CALLER_METRICS or metric.name is MetricName.SPEND for metric in self.metrics
):
raise ValueError("Outbound events own only attempt usage metrics")
if self.metrics is not None:
_check_subsets({metric.name.value: metric.value for metric in self.metrics})
return self
PricingUnit: TypeAlias = Literal["token", "second", "character", "image", "request"]
class Rate(BaseModel):
model_config = _CONFIG
measure: Identifier
unit: PricingUnit
usd_per_unit: Quantity
class PricedQuantity(BaseModel):
"""A pricing input in its billing unit. Only seconds may be fractional."""
model_config = _CONFIG
measure: Identifier
unit: PricingUnit
value: Quantity
@model_validator(mode="after")
def validate_discrete_quantity(self) -> Self:
if self.unit != "second" and self.value != self.value.to_integral_value():
raise ValueError("Discrete priced quantities require integral values")
return self
class Multiplier(BaseModel):
"""A regional or tier factor applied on top of the priced key, e.g. data_residency=eu or vertex_location=us-east5."""
model_config = _CONFIG
name: Identifier
value: Identifier
factor: Annotated[Decimal, Field(strict=True, gt=0, allow_inf_nan=False)]
class NamedCost(BaseModel):
"""A free-form extra charge by name, e.g. azure_model_router_flat_cost."""
model_config = _CONFIG
name: Identifier
amount: Quantity
class PricingSnapshot(BaseModel):
"""Retained inputs and result of authoritative compilation, reproducible after cost-map rates change.
Request-scoped inputs, overrides and effective rates come from the usage row's versioned payload.
Retries reuse the committed charge and snapshot.
Component names match ``CostBreakdown``: reasoning_cost is part of output_cost, not an additional charge.
Component sums and float tolerances belong to reconciliation, not this structural contract.
"""
model_config = _CONFIG
schema_version: Literal[1] = 1
calculator_version: Identifier
model: Identifier
provider: Identifier | None
pricing_source: Literal["cost_map", "deployment", "per_request", "upstream_header"]
priced_model_key: Identifier
multipliers: tuple[Multiplier, ...] = ()
rates: tuple[Rate, ...] = ()
usage: tuple[PricedQuantity, ...] = ()
input_cost: Quantity | None = None
output_cost: Quantity | None = None
cache_read_cost: Quantity | None = None
cache_creation_cost: Quantity | None = None
reasoning_cost: Quantity | None = None
tool_usage_cost: Quantity | None = None
guardrail_cost: Quantity | None = None
additional_costs: tuple[NamedCost, ...] = ()
original_cost: Quantity | None = None
discount_percent: Quantity | None = None
discount_amount: Quantity | None = None
margin_percent: Quantity | None = None
margin_fixed_amount: Quantity | None = None
margin_total_amount: Quantity | None = None
upstream_reported_cost: Quantity | None = None
total_cost: Quantity
@model_validator(mode="after")
def validate_snapshot(self) -> Self:
rate_keys: Final = frozenset((rate.measure, rate.unit) for rate in self.rates)
quantity_keys: Final = frozenset((quantity.measure, quantity.unit) for quantity in self.usage)
if len(rate_keys) != len(self.rates):
raise ValueError("Duplicate rate measures and units")
if len(quantity_keys) != len(self.usage):
raise ValueError("Duplicate priced quantity measures and units")
if not rate_keys.issubset(quantity_keys):
raise ValueError("Every rate requires one priced quantity with the same measure and unit")
if len(frozenset(cost.name for cost in self.additional_costs)) != len(self.additional_costs):
raise ValueError("Duplicate additional cost names")
if len(frozenset(multiplier.name for multiplier in self.multipliers)) != len(self.multipliers):
raise ValueError("Duplicate multiplier names")
_check_subsets({quantity.measure: quantity.value for quantity in self.usage if quantity.unit == "token"})
return self
class CostEvent(BaseModel):
model_config = _CONFIG
event_id: Identifier
usage_event_id: Identifier
occurred_at: Instant
dimensions: Dimensions
charges: Metrics
snapshot: PricingSnapshot
@model_validator(mode="after")
def validate_charge(self) -> Self:
if len(self.charges) != 1 or self.charges[0].name is not MetricName.SPEND:
raise ValueError("A cost event requires one authoritative USD spend metric")
if self.event_id == self.usage_event_id:
raise ValueError("Cost and usage events require distinct identities")
if self.snapshot.total_cost != self.charges[0].value:
raise ValueError("Snapshot total must equal the authoritative spend")
return self
class Unpriced(BaseModel):
model_config = _CONFIG
usage_event_id: Identifier
reason: Identifier
class Measurement(BaseModel):
model_config = _CONFIG
source_event_id: Identifier
occurred_at: Instant
dimensions: Dimensions
metric: Metric
class DimensionFilter(BaseModel):
model_config = _CONFIG
name: DimensionName
values: Annotated[tuple[Identifier, ...], Field(min_length=1)]
class AggregateQuery(BaseModel):
model_config = _CONFIG
metric: MetricName
filters: tuple[DimensionFilter, ...]
group_by: tuple[DimensionName, ...]
start: Instant
end: Instant
granularity: Granularity | None
@model_validator(mode="after")
def validate_query(self) -> Self:
_interval(self.start, self.end)
if len(frozenset(item.name for item in self.filters)) != len(self.filters):
raise ValueError("Use one IN filter per dimension; filters combine with AND")
if len(frozenset(self.group_by)) != len(self.group_by):
raise ValueError("Duplicate group-by dimensions")
if any(
(instant.hour, instant.minute, instant.second, instant.microsecond) != (0, 0, 0, 0)
for instant in (self.start, self.end)
):
raise ValueError("Initial reporting queries require UTC day boundaries")
return self
class AggregateRow(BaseModel):
model_config = _CONFIG
start: Instant
end: Instant
dimensions: Dimensions
value: Quantity
@model_validator(mode="after")
def validate_bounds(self) -> Self:
_interval(self.start, self.end)
return self
def validate_aggregate_rows(rows: tuple[AggregateRow, ...]) -> tuple[AggregateRow, ...]:
"""Validate [start, end) buckets per group; overlapping tag groups are not additive."""
groups: Final = frozenset(frozenset(row.dimensions) for row in rows)
for group in groups:
ordered = sorted((row for row in rows if frozenset(row.dimensions) == group), key=lambda row: row.start)
if any(left.end > right.start for left, right in zip(ordered, ordered[1:])):
raise ValueError("Aggregate buckets overlap within a dimension group")
return rows
class AdmissionDecision(BaseModel):
model_config = _CONFIG
allowed: bool
reason: Identifier | None = None

View file

@ -0,0 +1,102 @@
import os
import time
from typing import (
TYPE_CHECKING,
Final,
Literal,
Protocol,
cast, # noqa: TID251 # generated Prisma transaction signatures cannot express the borrowed transaction protocol
)
from litellm._logging import verbose_proxy_logger
from litellm.accounting.application import apply_deliveries
from litellm.accounting.metrics import record_deliveries
from litellm.accounting.projection import Contribution
from litellm.repositories.accounting_delivery_repository import (
AccountingDeliveryRepository,
DeliveryProjection,
LedgerDatabase,
LedgerTransaction,
)
from litellm.types.integrations.slack_alerting import AlertType
if TYPE_CHECKING:
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.utils import PrismaClient, ProxyLogging
class _AccountingAlerts(Protocol):
async def alerting_handler(
self, message: str, level: Literal["Low", "Medium", "High"], alert_type: AlertType
) -> None: ...
_JOB: Final = "accounting_delivery"
async def flush_accounting_deliveries(
writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging"
) -> None:
acquired: Final = await writer.pod_lock_manager.acquire_lock(cronjob_id=_JOB, allow_reentrant=False)
if acquired is False:
return
applied: tuple[Contribution, ...] = () # rebind-ok: keep committed contributions for cache invalidation
async def apply(transaction: LedgerTransaction, deliveries: tuple[DeliveryProjection, ...]) -> None:
nonlocal applied
applied = await apply_deliveries( # rebind-ok: invalidate caches only after the outer transaction commits
transaction, deliveries, prisma_client=prisma_client, writer=writer, proxy_logging_obj=proxy_logging_obj
)
try:
database: Final = cast( # cast-ok: PrismaClient.tx owns writer routing and yields the generated transaction
LedgerDatabase, prisma_client
)
repository: Final = AccountingDeliveryRepository(database)
started: Final = time.monotonic()
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "10000"))
async def run_batch(*, retry_only: bool) -> bool:
result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only)
if result.applied:
await writer.invalidate_accounting_entities(applied, proxy_logging_obj)
if result.parked:
verbose_proxy_logger.error(
"Accounting delivery parked %d requests after isolated failures", len(result.parked)
)
await _alert(
proxy_logging_obj,
message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.",
)
return bool(result.applied or result.failed)
await run_batch(retry_only=False)
retry_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True):
break
await repository.cleanup_applied(
retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")),
batch_size=batch_size * 2,
)
metrics: Final = await repository.metrics()
record_deliveries(metrics, time.monotonic() - started)
if metrics.batch_count >= int(
os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000")
) or metrics.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")):
await _alert(
proxy_logging_obj,
message=f"Accounting delivery storage requires review. Retained batches={metrics.batch_count}, table bytes={metrics.table_bytes}.",
)
if metrics.oldest_pending_seconds > 300:
await _alert(
proxy_logging_obj,
message=f"Accounting deliveries are delayed. Pending={metrics.pending_count}, oldest={metrics.oldest_pending_seconds:.0f}s.",
)
finally:
if acquired:
await writer.pod_lock_manager.release_lock(cronjob_id=_JOB)
async def _alert(alerts: _AccountingAlerts, message: str) -> None:
await alerts.alerting_handler(message=message, level="High", alert_type=AlertType.failed_tracking_spend)

View file

@ -620,7 +620,7 @@ class DualCache(BaseCache):
increment_list: list["RedisPipelineIncrementOperation"],
local_only: bool = False,
parent_otel_span: Span | None = None,
**kwargs,
**kwargs: object, # kwargs-ok: accepted and unused, as the cache interface's other calls take them
) -> list[float] | None:
result: list[float] | None = None
try:

View file

@ -1673,6 +1673,9 @@ OPENAI_FINISH_REASONS: Final = [
]
HUMANLOOP_PROMPT_CACHE_TTL_SECONDS: Final = get_env_int("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60) # 1 minute
RESPONSE_FORMAT_TOOL_NAME = "json_tool_call" # default tool name used when converting response format to tool call
RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION: Final = (
"When you are ready to answer, call this tool with your final answer as its input. Do not reply in plain text."
)
########################### Logging Callback Constants ###########################
AZURE_STORAGE_MSFT_VERSION: Final = "2019-07-07"

View file

@ -616,11 +616,17 @@ def cost_per_token(
)
if per_second is not None:
return per_second
if custom_llm_provider == "anthropic":
return anthropic_cost_per_token(
model=model, usage=usage_block, service_tier=service_tier, model_info=requested_model_info_override
)
return generic_cost_per_token(
model=model,
usage=usage_block,
custom_llm_provider=custom_llm_provider or "",
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
model_info=requested_model_info_override,
)
@ -3323,15 +3329,35 @@ def batch_cost_calculator(
return total_prompt_cost * vertex_uplift, total_completion_cost * vertex_uplift
def _standard_rate(model_info: ModelInfo, key: str, prompt_tokens: int) -> float | None:
return get_cost_per_unit(model_info, get_modality_token_cost_key(model_info, key, prompt_tokens), None)
def _cached_audio_batch_fallback(
model_info: ModelInfo, batch_cache_read_rate: float | None, cache_read_rate: float | None, prompt_tokens: int
) -> float:
if batch_cache_read_rate is not None:
return batch_cache_read_rate
standard_audio_rate: Final = _standard_rate(model_info, "cache_read_input_audio_token_cost", prompt_tokens)
if standard_audio_rate is not None:
return standard_audio_rate
return 0.0 if cache_read_rate is None else cache_read_rate
def _batch_prompt_cost(
usage: Usage,
model_info: ModelInfo,
input_rate: float,
cache_read_rate: float | None,
batch_cache_read_rate: float | None,
cache_creation_rate: float | None,
cache_creation_rates_by_ttl: Mapping[str, float],
) -> float:
details: Final = parse_prompt_tokens_details(usage)
cache_read_rate: Final = (
_standard_rate(model_info, "cache_read_input_token_cost", usage.prompt_tokens)
if batch_cache_read_rate is None
else batch_cache_read_rate
)
cached_tokens: Final = details["cache_hit_tokens"] if cache_read_rate is not None else 0
cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cache_read_rate is not None else 0
bills_cache_writes: Final = cache_creation_rate is not None or bool(cache_creation_rates_by_ttl)
@ -3352,7 +3378,11 @@ def _batch_prompt_cost(
+ video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate)
+ (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0)
+ cached_audio_tokens
* _batch_rate(model_info, "cache_read_input_audio_token_cost_batches", cache_read_rate or 0.0)
* _batch_rate(
model_info,
"cache_read_input_audio_token_cost_batches",
_cached_audio_batch_fallback(model_info, batch_cache_read_rate, cache_read_rate, usage.prompt_tokens),
)
+ (
calculate_cache_writing_cost(
written_tokens,

View file

@ -0,0 +1,181 @@
"""Model deprecation emails with provider identity, a migration comparison and next steps."""
from collections.abc import Mapping
from html import escape
from types import MappingProxyType
from typing import Final
from urllib.parse import urlencode, urlsplit
from litellm.types.integrations.model_deprecation_email import (
DeprecationEmailMilestone,
ModelDeprecationEmail,
RenderedDeprecationEmail,
)
from litellm.types.proxy.model_deprecation import SuccessorModelSource
_LOGO_BASE_URL: Final = (
"https://raw.githubusercontent.com/BerriAI/litellm/main/ui/litellm-dashboard/public/assets/logos"
)
_EMAIL_LOGO_URL: Final = f"{_LOGO_BASE_URL}/litellm_logo.png"
# Reuse dashboard brand assets; PNGs also render in email clients that block SVG.
_PROVIDER_BRANDS: Final[Mapping[str, tuple[str, str]]] = MappingProxyType(
{
"gemini": ("Google", "google_pse.png"),
"vertex_ai": ("Google", "google_pse.png"),
"openai": ("OpenAI", "openai.png"),
"anthropic": ("Anthropic", "anthropic.png"),
"azure": ("Microsoft Azure", "microsoft_azure.png"),
"bedrock": ("AWS Bedrock", "aws.png"),
"bedrock_converse": ("AWS Bedrock", "aws.png"),
}
)
def _recommendation_source(source: SuccessorModelSource | None) -> str:
if source is SuccessorModelSource.DEPLOYMENT_CONFIG:
return "Gateway configuration"
if source is SuccessorModelSource.MODEL_COST_MAP:
return "LiteLLM model cost map"
return "Not provided"
def _http_url(value: str | None) -> str | None:
if value is None:
return None
try:
parsed: Final = urlsplit(value)
except ValueError:
return None
return value if parsed.scheme in ("http", "https") and parsed.netloc and not parsed.username else None
def _provider(model: str, fallback: str | None = None) -> str:
prefix, separator, _ = model.partition("/")
return prefix if separator else fallback or ""
def _provider_name(provider: str) -> str:
brand: Final = _PROVIDER_BRANDS.get(provider)
if brand:
return brand[0]
return provider.replace("_", " ") or "The provider"
def _model_html(model: str, provider: str) -> str:
brand: Final = _PROVIDER_BRANDS.get(provider)
logo: Final = (
f'<td width="38" valign="middle"><img src="{_LOGO_BASE_URL}/{brand[1]}" width="24" height="24" alt="{escape(brand[0], quote=True)}" style="display:block;object-fit:contain;border:0;"></td>'
if brand
else ""
)
return (
f'<table role="presentation" width="100%" cellspacing="0" cellpadding="0"><tr>{logo}'
f'<td><div style="font-size:11px;line-height:18px;color:#526079;">{escape(_provider_name(provider))}</div>'
f"<div style=\"font:14px 'Courier New',monospace;line-height:22px;word-break:break-word;overflow-wrap:anywhere;color:#17213b;\">{escape(model)}</div></td></tr></table>"
)
def render_model_deprecation_email(email: ModelDeprecationEmail, *, base_url: str) -> RenderedDeprecationEmail:
gateway: Final = _http_url(base_url)
if gateway is None:
raise ValueError("Model deprecation emails require a public HTTP(S) PROXY_BASE_URL")
urgent: Final = email.milestone is DeprecationEmailMilestone.ONE_DAY
days: Final = int(email.milestone)
timing: Final = "tomorrow" if urgent else "in 30 days"
provider: Final = _provider(email.usage_model, email.model.litellm_provider)
provider_name: Final = _provider_name(provider)
subject: Final = f"[LiteLLM] Model being deprecated by {provider_name} {timing}: {email.usage_model}"
deadline: Final = email.model.deprecation_date.strftime("%d %b %Y")
announcement: Final = f"{provider_name} is deprecating this model on {deadline} (UTC)."
preheader: Final = f"{announcement} Review the model change and migrate your affected keys."
query: Final = urlencode(
{
"tab": "models",
"from": email.start_date.isoformat(),
"to": email.end_date.isoformat(),
}
)
review_url: Final = f"{gateway.rstrip('/')}/ui/usage?{query}"
replacement: Final = email.model.successor_model
after: Final = (
_model_html(replacement, _provider(replacement, provider))
if replacement
else '<p style="margin:0;font-size:14px;line-height:22px;color:#17213b;">No recommended model listed. Choose a supported replacement.</p>'
)
source: Final = _recommendation_source(email.model.successor_model_source)
notice: Final = _http_url(email.model.deprecation_source)
source_link: Final = (
f'<a href="{escape(notice, quote=True)}" style="color:#263cdb;text-decoration:underline;">Read the provider notice &#8599;</a>'
if notice
else ""
)
team_names: Final = ", ".join(email.team_names[:3]) or "Personal keys"
more_teams: Final = f" + {len(email.team_names) - 3} more" if len(email.team_names) > 3 else ""
day_label: Final = "DAY" if urgent else "DAYS"
label: Final = "FINAL REMINDER" if urgent else "DEPRECATION NOTICE"
test_step: Final = (
"Test the recommended model with your workload."
if replacement
else "Choose a supported replacement and test it with your workload."
)
switch_step: Final = (
f"Switch your app or gateway routing before {deadline}. Then check that requests to the old model reach zero."
)
html: Final = f"""<!doctype html>
<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>{escape(subject)}</title>
<style>@import url('https://fonts.googleapis.com/css2?family=DM+Sans:wght@400;500;600;700&display=swap');
@media only screen and (max-width:620px){{.pad{{padding:24px!important}}.headline{{font-size:30px!important;line-height:36px!important}}.brand{{padding:20px 24px!important}}.brand img{{width:96px!important}}.comparison-half{{display:block!important;width:auto!important;border-right:0!important}}.comparison-after{{border-top:1px solid #e2e6ee!important}}}}
</style></head>
<body style="margin:0;padding:0;background:#f7f8fa;color:#17213b;font-family:'DM Sans',Arial,sans-serif;">
<div style="display:none;max-height:0;overflow:hidden;mso-hide:all;">{escape(preheader)}</div>
<table role="presentation" width="100%" cellspacing="0" cellpadding="0" border="0"><tr><td style="padding:0;">
<table role="presentation" width="100%" cellspacing="0" cellpadding="0" border="0" style="width:100%;background:#ffffff;">
<tr><td class="brand" style="padding:24px 36px;border-bottom:1px solid #e2e6ee;">
<table role="presentation" width="100%" cellspacing="0" cellpadding="0"><tr><td><img src="{_EMAIL_LOGO_URL}" width="112" alt="LiteLLM" style="display:block;border:0;height:auto;"></td>
<td align="right" style="font:11px 'Courier New',monospace;letter-spacing:1.5px;color:#526079;">/ MODEL DEPRECATIONS</td></tr></table></td></tr>
<tr><td class="pad" style="padding:32px 36px;background:#263cdb;color:#ffffff;">
<table role="presentation" width="100%" cellspacing="0" cellpadding="0"><tr><td style="font:11px 'Courier New',monospace;letter-spacing:1.8px;color:#ffffff;">{label}</td>
<td align="right" style="font:11px 'Courier New',monospace;color:#ffffff;">{days:02d} {day_label} UNTIL DEPRECATION</td></tr></table>
<h1 class="headline" style="margin:22px 0 16px;font-size:36px;line-height:42px;font-weight:600;letter-spacing:-1px;color:#ffffff;">Model being deprecated</h1>
<p style="margin:0;font-size:14px;line-height:22px;color:#ffffff;">{escape(announcement)}</p>
</td></tr>
<tr><td class="pad" style="padding:28px 36px;">
<table role="presentation" width="100%" cellspacing="0" cellpadding="0" style="border:1px solid #e2e6ee;border-collapse:collapse;">
<tr><td class="comparison-half" width="50%" valign="top" style="padding:16px 18px;border-right:1px solid #e2e6ee;">
<div style="margin-bottom:8px;font:10px 'Courier New',monospace;letter-spacing:1px;color:#526079;">BEFORE &middot; MODEL TO REPLACE</div>
{_model_html(email.usage_model, provider)}
</td>
<td class="comparison-half comparison-after" width="50%" valign="top" style="padding:16px 18px;background:#f5f7ff;">
<div style="margin-bottom:8px;font:10px 'Courier New',monospace;letter-spacing:1px;color:#263cdb;">AFTER &middot; RECOMMENDED MODEL</div>
{after}
</td></tr></table>
<p style="margin:10px 0 24px;font-size:12px;line-height:20px;color:#526079;">Recommendation source: {escape(source)}<br>{source_link}</p>
<h2 style="margin:0 0 12px;font-size:17px;font-weight:600;color:#17213b;">Your affected usage</h2>
<table role="presentation" width="100%" cellspacing="0" cellpadding="0" style="border-top:1px solid #e2e6ee;border-bottom:1px solid #e2e6ee;">
<tr><td width="50%" style="padding:14px 0;"><div style="font:24px 'Courier New',monospace;color:#17213b;">{email.key_count:,}</div><div style="margin-top:5px;font-size:12px;color:#526079;">Affected keys</div></td>
<td width="50%" style="padding:14px 0 14px 20px;border-left:1px solid #e2e6ee;"><div style="font:24px 'Courier New',monospace;color:#17213b;">{email.requests:,}</div><div style="margin-top:5px;font-size:12px;color:#526079;">Requests &middot; last 30 days</div></td></tr></table>
<p style="margin:10px 0 24px;font-size:12px;line-height:20px;color:#526079;">{escape(team_names + more_teams)}<br>Your keys and teams you administer.</p>
<h2 style="margin:0 0 10px;font-size:17px;font-weight:600;color:#17213b;">What to do</h2>
<ol style="margin:0 0 24px;padding-left:20px;font-size:13px;line-height:22px;color:#526079;">
<li style="padding:0 0 8px 4px;">Review the affected keys in your usage dashboard.</li>
<li style="padding:0 0 8px 4px;">{escape(test_step)}</li>
<li style="padding-left:4px;">{escape(switch_step)}</li></ol>
<table role="presentation" cellspacing="0" cellpadding="0"><tr><td bgcolor="#263cdb" style="border-radius:4px;mso-padding-alt:14px 20px;">
<a href="{escape(review_url, quote=True)}" style="display:inline-block;padding:14px 20px;color:#ffffff;font-size:13px;font-weight:600;line-height:18px;text-decoration:none;">Review affected usage &#8599;</a></td></tr></table>
</td></tr>
<tr><td class="pad" style="padding:20px 36px;background:#fafbfc;border-top:1px solid #e2e6ee;font-size:11px;line-height:18px;color:#526079;">
Your gateway admin enabled these emails. Reminder schedule: 30 days and 1 day before deprecation.<br><br>LiteLLM &middot; AI Gateway</td></tr>
</table></td></tr></table></body></html>"""
text: Final = (
f"{subject}\n\n{announcement}\n\n"
f"BEFORE - MODEL TO REPLACE\n{email.usage_model}\n\n"
f"AFTER - RECOMMENDED MODEL\n{replacement or 'No recommended model listed. Choose a supported replacement.'}\n"
f"Recommendation source: {source}\nProvider notice: {notice or 'Not provided'}\n\n"
f"YOUR AFFECTED USAGE\n{email.key_count} keys | {email.requests} requests in the last 30 days\n"
f"{team_names}{more_teams}\nYour keys and teams you administer.\n\n"
f"WHAT TO DO\n1. Review the affected keys in your usage dashboard.\n2. {test_step}\n3. {switch_step}\n\n"
f"Review affected usage: {review_url}\n\n"
"Your gateway admin enabled these emails. Reminder schedule: 30 days and 1 day before deprecation."
)
return RenderedDeprecationEmail(subject=subject, html=html, text=text)

View file

@ -0,0 +1,267 @@
"""Two owner-scoped retirement reminders, with durable per-recipient delivery receipts."""
from __future__ import annotations
import asyncio
import hashlib
import json
import os
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import date, datetime, timedelta, timezone
from enum import Enum
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from pydantic import AliasPath, BaseModel, ConfigDict, Field, TypeAdapter
from litellm.integrations.email_templates.model_deprecation_email import (
render_model_deprecation_email,
)
from litellm.models.team import LiteLLM_TeamTable
from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.table_repositories import DailyUserSpendRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
)
from litellm.types.integrations.model_deprecation_email import (
MODEL_DEPRECATION_EMAIL_EVENT,
DeprecationEmailMilestone,
ModelDeprecationEmail,
)
from litellm.types.proxy.model_deprecation import (
DEPRECATION_IDLE_POLL_SECONDS,
ModelDeprecationInfo,
)
if TYPE_CHECKING:
from litellm import Router
from litellm.proxy.utils import PrismaClient
_RECEIPT_PREFIX: Final = "model_deprecation_email:"
_PAGE_SIZE: Final = 100
class _OptIn(BaseModel):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
enabled: bool = Field(
False,
strict=True,
validation_alias=AliasPath("email_settings", MODEL_DEPRECATION_EMAIL_EVENT),
)
class _RequestCount(BaseModel):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
api_requests: int
class _UsedKey(BaseModel):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
api_key: str
metrics: _RequestCount = Field(alias="_sum")
_USED_KEYS: Final = TypeAdapter(tuple[_UsedKey, ...])
class _DeliveryStatus(str, Enum):
SENDING = "sending"
SENT = "sent"
class _Receipt(BaseModel):
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
revision: int = 0
status: _DeliveryStatus = _DeliveryStatus.SENDING
lease_until: datetime
@dataclass(frozen=True, slots=True)
class _Impact:
key_count: int = 0
requests: int = 0
teams: frozenset[str] = frozenset()
def _admin_ids(team: LiteLLM_TeamTable | None) -> frozenset[str]:
if team is None:
return frozenset()
return frozenset(
member.user_id for member in (team.members_with_roles or ()) if member.role == "admin" and member.user_id
)
def _recipient_emails(
user_id: str | None, team: LiteLLM_TeamTable | None, addresses: Mapping[str, str]
) -> frozenset[str]:
owner_ids: Final[frozenset[str]] = frozenset((user_id,)) if user_id else frozenset()
resolved: Final = frozenset(addresses[owner] for owner in owner_ids | _admin_ids(team) if owner in addresses)
if team is None:
return resolved
return resolved | frozenset(
member.user_email.strip()
for member in (team.members_with_roles or ())
if member.role == "admin"
and member.user_email
and member.user_email.strip()
and member.user_id not in addresses
)
def _usage_model(model: ModelDeprecationInfo) -> str:
name: Final = model.litellm_model or model.model_name
provider: Final = model.litellm_provider
if provider and not name.startswith(f"{provider}/"):
return f"{provider}/{name}"
return name
def _receipt_key(email: ModelDeprecationEmail) -> str:
identity: Final = json.dumps(
[
email.usage_model,
email.model.deprecation_date.isoformat(),
int(email.milestone),
email.recipient_email.casefold(),
]
)
return _RECEIPT_PREFIX + hashlib.sha256(identity.encode()).hexdigest()
class ModelDeprecationEmailSender:
def __init__(self, prisma_client: PrismaClient) -> None:
self.prisma: Final = prisma_client
self.config: Final = ConfigRepository(prisma_client, use_writer=True)
async def _impacts(self, model: str, start: date, end: date) -> Mapping[str, _Impact]:
impacts: Final[dict[str, _Impact]] = {} # mutable-ok: combines per-recipient counts across bounded usage pages
offset: int = 0 # rebind-ok: each page advances the stable key grouping cursor
while True:
raw = await DailyUserSpendRepository(self.prisma).table.group_by(
by=["api_key"],
where={
"model": model,
"date": {"gte": start.isoformat(), "lte": end.isoformat()},
"api_requests": {"gt": 0},
},
sum={"api_requests": True},
order={"api_key": "asc"},
skip=offset,
take=_PAGE_SIZE,
)
usage = _USED_KEYS.validate_python(raw)
if not usage:
return MappingProxyType(impacts)
keys = await VerificationTokenRepository(self.prisma).find_many_in("token", (row.api_key for row in usage))
teams = await TeamRepository(self.prisma).find_many_in(
"team_id", frozenset(key.team_id for key in keys if key.team_id)
)
team_by_id = {team.team_id: team for team in teams}
ids = frozenset(key.user_id for key in keys if key.user_id) | frozenset(
chain.from_iterable(_admin_ids(team) for team in teams)
)
users = await UserRepository(self.prisma).find_many_in("user_id", ids)
addresses = {
user.user_id: user.user_email.strip() for user in users if user.user_email and user.user_email.strip()
}
counts = {row.api_key: row.metrics.api_requests for row in usage}
for key in keys:
team = team_by_id.get(key.team_id) if key.team_id else None
if key.token is None:
continue
recipients = _recipient_emails(key.user_id, team, addresses)
for address in recipients:
previous = impacts.get(address, _Impact())
names: frozenset[str] = frozenset((team.team_alias or team.team_id,)) if team else frozenset()
impacts[address] = _Impact(
key_count=previous.key_count + 1,
requests=previous.requests + counts[key.token],
teams=previous.teams | names,
)
if len(usage) < _PAGE_SIZE:
return MappingProxyType(impacts)
offset += _PAGE_SIZE
async def _deliver(self, email: ModelDeprecationEmail, now: datetime, base_url: str) -> bool:
from litellm.proxy.utils import send_email
key: Final = _receipt_key(email)
record: Final = await self.config.get_param(key)
previous: Final = _Receipt.model_validate(record.param_value) if record else None
if previous and (previous.status is _DeliveryStatus.SENT or previous.lease_until > now):
return False
revision: Final = previous.revision if previous else 0
claim: Final = _Receipt(revision=revision + 1, lease_until=now + timedelta(minutes=5))
if not await self.config.set_param_if_revision(key, claim.model_dump(mode="json"), revision):
return False
content: Final = render_model_deprecation_email(email, base_url=base_url)
if not await send_email(
receiver_email=email.recipient_email,
subject=content.subject,
html=content.html,
text=content.text,
message_id=f"<litellm-deprecation-{key.removeprefix(_RECEIPT_PREFIX)}@alerts.litellm.ai>",
):
return False
sent: Final = claim.model_copy(update={"revision": claim.revision + 1, "status": _DeliveryStatus.SENT})
if not await self.config.set_param_if_revision(key, sent.model_dump(mode="json"), claim.revision):
raise RuntimeError("Could not record the delivered model deprecation reminder")
return True
async def run(self, llm_router: Router, *, now: datetime | None = None) -> int:
current: Final = now or datetime.now(timezone.utc)
settings: Final = await self.config.get_param("general_settings")
if settings is None or not _OptIn.model_validate(settings.param_value).enabled:
return 0
base_url: Final = os.getenv("PROXY_BASE_URL", "")
if not base_url:
return 0
await self.config.delete_params_older_than(_RECEIPT_PREFIX, current - timedelta(days=60))
snapshot: Final = collect_model_deprecations(llm_router=llm_router, today=current.date())
start: Final = current.date() - timedelta(days=29)
seen: Final[set[tuple[str, date]]] = set() # mutable-ok: deduplicates aliases in this pass
sent: int = 0 # rebind-ok: counts successfully delivered reminders in this pass
for model in snapshot.imminent:
if model.days_until_deprecation not in tuple(DeprecationEmailMilestone):
continue
name = _usage_model(model)
identity = (name, model.deprecation_date)
if identity in seen:
continue
seen.add(identity)
impacts = await self._impacts(name, start, current.date())
for address, impact in sorted(impacts.items()):
email = ModelDeprecationEmail(
model=model,
usage_model=name,
recipient_email=address,
milestone=DeprecationEmailMilestone(model.days_until_deprecation),
key_count=impact.key_count,
requests=impact.requests,
team_names=tuple(sorted(impact.teams)),
start_date=start,
end_date=current.date(),
)
if await self._deliver(email, now or datetime.now(timezone.utc), base_url):
sent += 1
return sent
async def run_scheduled_model_deprecation_emails() -> None:
"""Poll the email opt-in independently; a large owner mailing never delays Slack checks."""
from litellm._logging import verbose_proxy_logger
while True:
try:
from litellm.proxy.proxy_server import llm_router, prisma_client
if llm_router is not None and prisma_client is not None:
await ModelDeprecationEmailSender(prisma_client).run(llm_router)
except Exception: # noqa: BLE001 # transient SMTP/DB failures must not terminate the reminder loop
verbose_proxy_logger.exception("Error sending model deprecation email reminders")
await asyncio.sleep(DEPRECATION_IDLE_POLL_SECONDS)

View file

@ -338,10 +338,12 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_ToolTable",
"LiteLLM_TelemetryReport",
"LiteLLM_AccessGroupTable",
"LiteLLM_AccountingDelivery",
"LiteLLM_AccountingEvent",
"LiteLLM_SpendLogDelivery",
"LiteLLM_ClaudeCodePluginTable",
"LiteLLM_ClaudeCodePluginUpload",
"LiteLLM_ClaudeCodePluginVersion",
"LiteLLM_SkillSubmission",
"LiteLLM_ClaudeCodePluginDownload",
"LiteLLM_MemoryTable",
"LiteLLM_AdaptiveRouterState",
"LiteLLM_AdaptiveRouterSession",

View file

@ -202,7 +202,7 @@ def response_total_token_count(response: object) -> int:
def apply_remaining_usage_headers(
additional_headers: dict[str, object],
remaining_usage: dict[str, int],
remaining_usage: Mapping[str, object],
) -> None:
for header, value in remaining_usage.items():
if value is not None and header not in additional_headers:

View file

@ -108,6 +108,17 @@ def as_object_list(value: object) -> Sequence[object] | None:
return value if isinstance(value, list) else None
_Key = TypeVar("_Key") # rebind-ok: a TypeVar cannot be Final
_Value = TypeVar("_Value") # rebind-ok: a TypeVar cannot be Final
def without(
mapping: Mapping[_Key, _Value], *keys: object
) -> dict[_Key, _Value]: # mutable-ok: a fresh copy the caller owns
"""``mapping`` without ``keys``, as a new dict in the same order; ``mapping`` is unchanged."""
return {key: value for key, value in mapping.items() if key not in keys}
def decoded_json_object(value: object) -> Mapping[str, object]:
mapping: Final = as_str_mapping(value)
if mapping is not None:

View file

@ -7143,6 +7143,7 @@ class StandardLoggingPayloadSetup:
init_response_obj: object,
api_base: str | None = None,
router_model_id: str | None = None,
custom_model_info: ModelInfo | None = None,
) -> StandardLoggingModelInformation:
model_cost_name: Final = select_model_name_for_cost_calc(
model=base_model if custom_pricing else None,
@ -7161,6 +7162,8 @@ class StandardLoggingPayloadSetup:
if router_model_id is not None
else model_cost_name
)
if model_cost_name is not None and custom_model_info is not None:
return StandardLoggingModelInformation(model_map_key=model_cost_name, model_map_value=custom_model_info)
if model_cost_name is None:
model_cost_information = StandardLoggingModelInformation(model_map_key="", model_map_value=None)
else:
@ -7911,6 +7914,7 @@ def get_standard_logging_object_payload(
init_response_obj=init_response_obj,
api_base=litellm_params.get("api_base"),
router_model_id=logging_obj.get_router_model_id(),
custom_model_info=logging_obj.custom_model_info if custom_pricing else None,
)
error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(

View file

@ -1081,4 +1081,8 @@ def convert_to_model_response_object(
end_time=end_time,
convert_tool_call_to_json_mode=convert_tool_call_to_json_mode,
)
raise Exception(f"Invalid response object {traceback.format_exc()}\n\nreceived_args={received_args}")
# The conversion's own traceback, without the errors being handled when it ran: a router
# retry runs inside the handler of the attempt before it, and that attempt's error text
# (a rate limit) would otherwise decide how this one is classified.
own_traceback: Final = "".join(traceback.format_exception(type(e), e, e.__traceback__, chain=False))
raise Exception(f"Invalid response object {own_traceback}\n\nreceived_args={received_args}")

View file

@ -378,7 +378,8 @@ def filter_team_based_models(
ids_to_remove: Final = set()
if isinstance(healthy_deployments, dict):
return healthy_deployments
# A deployment asked for by id is filtered like any other: another team's is not served.
return [] if _owner_team(healthy_deployments) not in (None, request_team_id) else healthy_deployments # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # the legacy untyped deployment dict
for deployment in healthy_deployments:
_model_info = deployment.get("model_info") or {}
model_team_id = _model_info.get("team_id")
@ -431,6 +432,16 @@ def _deployment_supports_web_search(deployment: dict) -> bool:
return True
def is_web_search_request(request_kwargs: Mapping[str, object]) -> bool:
tools: Final[object] = request_kwargs.get("tools")
return isinstance(tools, (list, tuple)) and any(
isinstance(tool, dict)
and isinstance(tool_type := tool.get("type"), str)
and tool_type in _WEB_SEARCH_TOOL_TYPES
for tool in tools
)
def filter_web_search_deployments(
healthy_deployments: list[dict] | dict,
request_kwargs: dict | None = None,
@ -445,14 +456,7 @@ def filter_web_search_deployments(
if isinstance(healthy_deployments, dict):
return healthy_deployments
tools: Final[object] = request_kwargs.get("tools")
is_web_search_request: Final = isinstance(tools, (list, tuple)) and any(
isinstance(tool, dict)
and isinstance(tool_type := tool.get("type"), str)
and tool_type in _WEB_SEARCH_TOOL_TYPES
for tool in tools
)
if not is_web_search_request:
if not is_web_search_request(request_kwargs):
return healthy_deployments
# Filter out deployments that don't support web search
@ -607,3 +611,9 @@ def warn_on_provider_credential_mismatch(model_name: str, litellm_params: Mappin
)
verbose_router_logger.warning(warning)
return warning
def _owner_team(deployment: Mapping[str, object]) -> object:
"""The team a deployment belongs to (``model_info.team_id``), or None."""
model_info: Final = deployment.get("model_info")
return model_info.get("team_id") if isinstance(model_info, Mapping) else None # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a deployment dict's model_info

View file

@ -734,6 +734,7 @@ class ModelResponseIterator:
self._split_answer_blocks: tuple[str, ...] = ()
self._withheld_answer_close: str = ""
self._streamed_answer_keys: frozenset[str] | None = None
self._held_prose: str = ""
# For handling partial JSON chunks from fragmentation
# See: https://github.com/BerriAI/litellm/issues/17473
@ -826,8 +827,13 @@ class ModelResponseIterator:
thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock] = []
if "text" in content_block["delta"]:
text = content_block["delta"]["text"]
return text, tool_use, thinking_blocks, provider_specific_fields, reasoning_content
return (
self._hold_prose(content_block["delta"]["text"]),
tool_use,
thinking_blocks,
provider_specific_fields,
reasoning_content,
)
self.content_blocks.append(content_block)
if "partial_json" in content_block["delta"]:
# Only emit tool calls if we're in a tool_use or server_tool_use block
@ -1025,7 +1031,7 @@ class ModelResponseIterator:
if self.current_content_block_type in _THINKING_BLOCK_TYPES:
self.thinking_block_index += 1
if content_block_start["content_block"]["type"] == "text":
text = self._release_answer_close() + content_block_start["content_block"]["text"]
text = self._release_answer_close() + self._hold_prose(content_block_start["content_block"]["text"])
elif (
content_block_start["content_block"]["type"] == "tool_use"
or content_block_start["content_block"]["type"] == "server_tool_use"
@ -1161,7 +1167,7 @@ class ModelResponseIterator:
elif type_chunk == "message_delta":
finish_reason, usage, container, stop_reason_fields = self._handle_message_delta(chunk)
provider_specific_fields.update(stop_reason_fields)
text = self._release_answer_close() + self._drain_split_answer_blocks()
text = self._release_answer_close() + self._drain_split_answer_blocks() + self._release_held_prose()
if container:
provider_specific_fields["container"] = container
# Store container_id and re-emit code_interpreter_results
@ -1211,7 +1217,9 @@ class ModelResponseIterator:
status_code=500, # it looks like Anthropic API does not return a status code in the chunk error - default to 500
)
text, tool_use = self._handle_json_mode_chunk(text=text, tool_use=tool_use)
text, tool_use = self._handle_json_mode_chunk(
text=text + self._settle_held_prose(tool_use), tool_use=tool_use
)
refusal_fields: Final = refusal_params(provider_specific_fields)
returned_chunk: Final = ModelResponseStream(
@ -1309,6 +1317,24 @@ class ModelResponseIterator:
self._response_format_verdict = _envelope_verdict(buffered, schema, RESPONSE_FORMAT_TOOL_NAME)
return self._withhold_answer_close(buffered) if self._response_format_verdict == "pass_through" else ""
def _hold_prose(self, text: str) -> str:
if not self.json_mode or self.response_format_schema is None:
return text
self._held_prose += text
return ""
def _settle_held_prose(self, tool_use: ChatCompletionToolCallChunk | None) -> str:
if tool_use is None or tool_use.get("id") is None or self.current_content_block_type != "tool_use":
return ""
held: Final = self._held_prose
self._held_prose = ""
return "" if tool_use.get("function", {}).get("name") == RESPONSE_FORMAT_TOOL_NAME else held
def _release_held_prose(self) -> str:
held: Final = self._held_prose
self._held_prose = ""
return "" if self.converted_response_format_tool else held
def _withhold_answer_close(self, text: str) -> str:
pending: Final = self._withheld_answer_close + text
body: Final = pending.rstrip()
@ -1362,7 +1388,9 @@ class ModelResponseIterator:
def cut_stream_chunk(self) -> ModelResponseStream | None:
unsent: Final = self._flush_response_format_buffer() if self._response_format_buffer.getvalue() else ""
text: Final = unsent + self._release_answer_close() + self._drain_split_answer_blocks()
text: Final = (
unsent + self._release_answer_close() + self._drain_split_answer_blocks() + self._release_held_prose()
)
if not text:
return None
return ModelResponseStream(

View file

@ -26,6 +26,7 @@ from litellm.constants import (
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET,
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
RESPONSE_FORMAT_TOOL_NAME,
RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION,
)
from litellm.litellm_core_utils.asyncify import is_event_loop_running
from litellm.litellm_core_utils.core_helpers import (
@ -1786,7 +1787,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
json_schema: Mapping[str, object] | None,
request_params: Mapping[str, object],
) -> bool:
open_object_takes_forced_json_tool: Final = (
open_object_takes_json_tool: Final = (
json_schema is not None
and schema_has_open_object(json_schema)
and not self.is_thinking_requested(request_params)
@ -1794,7 +1795,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
and not request_params.get("tools")
and not AnthropicModelInfo.forced_tool_use_unsupported(model)
)
return self.uses_native_structured_output(model) and not open_object_takes_forced_json_tool
return self.uses_native_structured_output(model) and not open_object_takes_json_tool
def native_output_format_schema(self, json_schema: Mapping[str, object]) -> Mapping[str, object] | None:
return self.filter_anthropic_output_schema(dict(self._defs_unpacked(json_schema)))
@ -1809,7 +1810,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
return AnthropicOutputSchema(type="json_schema", schema=dict(native_schema))
def map_response_format_to_anthropic_tool(
self, value: dict | None, optional_params: dict, is_thinking_enabled: bool
self, value: dict | None, optional_params: dict, tool_is_forced: bool
) -> AnthropicMessagesTool | None:
ignore_response_format_types: Final = ["text"]
if value is None or value["type"] in ignore_response_format_types: # value is a no-op
@ -1827,6 +1828,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
_tool: Final = self._create_json_tool_call_for_response_format(
json_schema=json_schema,
description=None if tool_is_forced else RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION,
)
return _tool
@ -1929,7 +1931,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
drop_params=drop_params,
runs_hosted_mcp=self.runs_hosted_mcp,
)
is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params=request_params)
response_format_tool_is_forced: Final = (
"response_format" in request_params
and AnthropicModelInfo.response_format_tool_can_be_forced(
model, self._resolved_provider, REQUEST_PARAMS.validate_python(request_params)
)
)
# NB: ``map_openai_params`` deliberately does NOT sanitize tool names
# here. Names are the *original* OpenAI names at this stage, and must
@ -1989,10 +1996,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
if _output_format is not None:
optional_params["output_format"] = _output_format
else:
_tool = self.map_response_format_to_anthropic_tool(value, optional_params, is_thinking_enabled)
_tool = self.map_response_format_to_anthropic_tool(
value, optional_params, response_format_tool_is_forced
)
if _tool is None:
continue
if not is_thinking_enabled and not AnthropicModelInfo.forced_tool_use_unsupported(model):
if response_format_tool_is_forced:
_tool_choice = {
"name": RESPONSE_FORMAT_TOOL_NAME,
"type": "tool",
@ -2121,12 +2130,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
def _create_json_tool_call_for_response_format(
self,
json_schema: dict | None = None,
description: str | None = None,
) -> AnthropicMessagesTool:
"""
Handles creating a tool call for getting responses in JSON format.
Args:
json_schema (Optional[dict]): The JSON schema the response should be in
description (Optional[str]): The tool description, left off the tool when None
Returns:
AnthropicMessagesTool: The tool call to send to Anthropic API to get responses in JSON format
@ -2143,11 +2154,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
else:
_input_schema.update(cast(AnthropicInputSchema, json_schema))
_tool: Final = AnthropicMessagesTool(
name=RESPONSE_FORMAT_TOOL_NAME,
input_schema=_input_schema,
if description is None:
return AnthropicMessagesTool(name=RESPONSE_FORMAT_TOOL_NAME, input_schema=_input_schema)
return AnthropicMessagesTool(
name=RESPONSE_FORMAT_TOOL_NAME, input_schema=_input_schema, description=description
)
return _tool
def should_strip_billing_metadata(self) -> bool:
"""

View file

@ -1141,6 +1141,34 @@ class AnthropicModelInfo(BaseLLMModelInfo):
"""
return AnthropicModelInfo._supports_model_capability(model, "thinking_always_on", custom_llm_provider)
@staticmethod
def thinks_by_default(model: str, custom_llm_provider: str) -> bool:
return AnthropicModelInfo._is_always_on_thinking_model(
model, custom_llm_provider
) or AnthropicModelInfo._supports_model_capability(model, "thinking_on_by_default", custom_llm_provider)
@staticmethod
def will_think(model: str, custom_llm_provider: str, request_params: Mapping[str, object]) -> bool:
if AnthropicModelInfo._is_always_on_thinking_model(model, custom_llm_provider):
return True
match request_params:
case {"thinking": True | {"type": "enabled" | "adaptive"}}:
return True
case {"reasoning_effort": effort} if effort is not None and effort != "none":
return True
case {"thinking": {"type": "disabled"}} | {"reasoning_effort": "none"}:
return False
case _:
return AnthropicModelInfo.thinks_by_default(model, custom_llm_provider)
@staticmethod
def response_format_tool_can_be_forced(
model: str, custom_llm_provider: str, request_params: Mapping[str, object]
) -> bool:
return not AnthropicModelInfo.forced_tool_use_unsupported(model) and not AnthropicModelInfo.will_think(
model, custom_llm_provider, request_params
)
@staticmethod
def _supports_legacy_thinking(model: str, custom_llm_provider: str) -> bool:
"""Whether ``model`` is an adaptive-thinking model that still accepts legacy

View file

@ -670,6 +670,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
# Synthesized compaction block from compact_20260112 polyfill (streaming).
self.compaction_block = compaction_block
self.iterations_usage = iterations_usage
# The provider's error behind this stream's error frame, for a router that accounts for it.
self.stream_failure: Exception | None = None
self._refusal_text: str = ""
self._paused_turn: bool = False
self._claude_refusal_stop_details: AnthropicStopDetails | None = None
@ -1396,6 +1398,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
verbose_logger.exception("Anthropic Adapter - mid-stream error: %s", e)
logging_obj: Final = self.litellm_logging_obj
provider_error: Final = _provider_error(e)
self.stream_failure = provider_error
if logging_obj is not None and logging_obj.on_detached_stream_failure is not None:
if provider_error is e:
raise
@ -1619,6 +1622,10 @@ class AnthropicSSEStream(AsyncIterator[bytes]):
self._byte_stream: Final[AsyncIterator[bytes]] = anthropic_wrapper.async_anthropic_sse_wrapper()
self._hidden_params: dict[str, object] = {}
@property
def stream_failure(self) -> Exception | None:
return self._anthropic_wrapper.stream_failure
@property
def chunks(self) -> "list[ModelResponseStream] | None":
return self._anthropic_wrapper.chunks

View file

@ -141,6 +141,8 @@ class AnthropicResponsesStreamWrapper:
if litellm_logging_obj is not None:
litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id)
self._unreported_failure: Exception | None = None
# The failure behind this stream's error frame, for a router that accounts for it.
self.stream_failure: Exception | None = None
self._current_block_index: int = -1
# Map item_id -> content_block_index so we can stop the right block later
self._item_id_to_block_index: dict[str, int] = {}
@ -467,6 +469,7 @@ class AnthropicResponsesStreamWrapper:
llm_provider=getattr(self.responses_stream, "custom_llm_provider", None),
model=self.model,
)
self.stream_failure = self._unreported_failure
self._chunk_queue.append(_anthropic_error_chunk(status_code, message))
async def _report_unreported_failure(self) -> None:
@ -578,5 +581,9 @@ class AnthropicResponsesSSEByteStream:
async def aclose(self) -> None:
await self._frames.aclose()
@property
def stream_failure(self) -> Exception | None:
return self._wrapper.stream_failure
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
return self._wrapper.partial_response_on_disconnect()

View file

@ -3,11 +3,10 @@ from collections.abc import Callable, Mapping
from copy import deepcopy
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from urllib.parse import urlsplit
import httpx
from openai.types.responses import ResponseReasoningItem
from pydantic import TypeAdapter
from litellm._logging import verbose_logger
@ -16,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import prompt_cach
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.llms.base_llm.responses.transformation import SUMMARY_ONLY_REASONING_FIELD_VALUES, ReasoningFieldValues
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.types.llms.openai import *
@ -72,6 +72,21 @@ def _with_responses_api_version(api_base: str | None, litellm_params: Mapping[st
return MappingProxyType({**litellm_params, "api_version": _resolve_responses_api_version(api_base, litellm_params)})
_AZURE_REJECTED_INPUT_FIELDS: Final = MappingProxyType(
{"message": frozenset({"status"}), "reasoning": frozenset({"status"})}
)
_AZURE_REJECTED_INPUT_FIELDS_UNDER_DROP_PARAMS: Final = MappingProxyType(
{**_AZURE_REJECTED_INPUT_FIELDS, "reasoning": frozenset({"status", "content"})}
)
def _without_fields_azure_rejects(item: object, rejected_fields: Mapping[str, frozenset[str]]) -> object:
if not is_str_keyed_mapping(item):
return item
rejected: Final = rejected_fields.get(str(item.get("type")), frozenset())
return {key: value for key, value in item.items() if key not in rejected} if rejected else item
class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
# Parameters not supported by Azure Responses API
AZURE_UNSUPPORTED_PARAMS = ["context_management"]
@ -122,68 +137,17 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
def get_stripped_model_name(self, model: str) -> str:
return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "")
def _handle_reasoning_item(self, item: dict[str, Any]) -> Mapping[str, object]:
"""
Handle reasoning items to filter out the status field.
Issue: https://github.com/BerriAI/litellm/issues/13484
Azure OpenAI API does not accept 'status' field in reasoning input items.
"""
if item.get("type") == "reasoning":
try:
# Ensure required fields are present for ResponseReasoningItem
item_data: Final = dict(item)
if "summary" not in item_data:
item_data["summary"] = (
item_data.get("reasoning_content", "")[:100] + "..."
if len(item_data.get("reasoning_content", "")) > 100
else item_data.get("reasoning_content", "")
)
# Create ResponseReasoningItem object from the item data
reasoning_item: Final = ResponseReasoningItem.model_validate(item_data)
# Convert back to dict with exclude_none=True to exclude None fields
dict_reasoning_item: Final = reasoning_item.model_dump(exclude_none=True)
dict_reasoning_item.pop("status", None)
return dict_reasoning_item
except Exception as e:
verbose_logger.debug("Failed to create ResponseReasoningItem, falling back to manual filtering: %s", e)
# Fallback: manually filter out known None fields
filtered_item: Final = {
k: v
for k, v in item.items()
if v is not None or k not in {"status", "content", "encrypted_content"}
}
return filtered_item
return item
def _validate_input_param(
self, input: str | ResponseInputParam, *, drops_set_created_by: bool = False
self, input: str | ResponseInputParam, *, drop_params: bool = False
) -> str | ResponseInputParam:
"""
Override parent method to also filter out 'status' field from message items.
Azure OpenAI API does not accept 'status' field in input messages.
"""
from typing import cast
# First call parent's validation
validated_input: Final = super()._validate_input_param(input, drops_set_created_by=drops_set_created_by)
# Then filter out status from message items
if isinstance(validated_input, list):
filtered_input: Final[list[object]] = []
for item in validated_input:
if isinstance(item, dict) and item.get("type") == "message":
# Filter out status field from message items
filtered_item = {k: v for k, v in item.items() if k != "status"}
filtered_input.append(filtered_item)
else:
filtered_input.append(item)
return cast(ResponseInputParam, filtered_input)
return validated_input
validated_input: Final = super()._validate_input_param(input, drop_params=drop_params)
if not isinstance(validated_input, list):
return validated_input
rejected_fields: Final = (
_AZURE_REJECTED_INPUT_FIELDS_UNDER_DROP_PARAMS if drop_params else _AZURE_REJECTED_INPUT_FIELDS
)
sanitized_items: Final = [_without_fields_azure_rejects(item, rejected_fields) for item in validated_input]
return cast("ResponseInputParam", sanitized_items) # cast-ok: items keep their shape minus rejected fields
def transform_responses_api_request(
self,

View file

@ -22,6 +22,7 @@ from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_hea
from litellm.constants import (
BEDROCK_MIN_THINKING_BUDGET_TOKENS,
RESPONSE_FORMAT_TOOL_NAME,
RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION,
)
from litellm.integrations.anthropic_cache_control_hook import (
MAX_CACHE_CONTROL_BLOCKS,
@ -950,19 +951,24 @@ class AmazonConverseConfig(BaseConfig):
def _supports_native_structured_outputs(
model: str,
custom_llm_provider: str | None,
will_think: bool,
thinking_requested: bool,
model_info: Mapping[str, object] | None,
) -> bool:
"""Check if the Bedrock model takes native structured outputs (outputConfig.textFormat) for this request.
``supports_native_structured_output`` in ``litellm.model_cost`` turns it on for every request, and
``supports_bedrock_native_structured_output`` (the deployment's ``model_info`` first) also when the
request asks for thinking, which the json tool path would drop. Without thinking, a model that cannot
force the json tool gets it unless that key is false.
``supports_bedrock_native_structured_output`` (the deployment's ``model_info`` first) also when Claude
will think, which the json tool path would drop. When the request does not ask for thinking, a model
that cannot force the json tool gets it unless that key is false.
Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html
"""
return bedrock_chat_uses_native_structured_output(
model, thinking_requested, custom_llm_provider or "bedrock_converse", model_info
model,
will_think=will_think,
thinking_requested=thinking_requested,
custom_llm_provider=custom_llm_provider or "bedrock_converse",
model_info=model_info,
)
@staticmethod
@ -1005,11 +1011,13 @@ class AmazonConverseConfig(BaseConfig):
self,
model: str,
json_schema: Mapping[str, object],
*,
will_think: bool,
thinking_requested: bool,
model_info: Mapping[str, object] | None,
) -> Mapping[str, object] | None:
if not self._supports_native_structured_outputs(
model, self.custom_llm_provider, thinking_requested, model_info
model, self.custom_llm_provider, will_think, thinking_requested, model_info
):
return None
if BedrockModelInfo.get_base_model(model).startswith("anthropic"):
@ -1095,7 +1103,6 @@ class AmazonConverseConfig(BaseConfig):
request_params: Final = tool_params_without_hosted_mcp(
non_default_params, model=model, route="bedrock", drop_params=drop_params
)
is_thinking_enabled: Final = self.is_thinking_enabled(request_params)
base_model: Final = BedrockModelInfo.get_base_model(model)
drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model)
drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param(
@ -1109,7 +1116,6 @@ class AmazonConverseConfig(BaseConfig):
model=model,
optional_params=optional_params,
non_default_params=request_params,
is_thinking_enabled=is_thinking_enabled,
model_info=model_info,
)
if param == "max_tokens" or param == "max_completion_tokens":
@ -1314,7 +1320,6 @@ class AmazonConverseConfig(BaseConfig):
model: str,
optional_params: dict,
non_default_params: dict,
is_thinking_enabled: bool,
model_info: Mapping[str, object] | None = None,
) -> dict:
"""
@ -1339,9 +1344,16 @@ class AmazonConverseConfig(BaseConfig):
if "type" in value and value["type"] == "text":
return optional_params
thinking_requested: Final = self.is_thinking_requested(REQUEST_PARAMS.validate_python(non_default_params))
request_params: Final = REQUEST_PARAMS.validate_python(non_default_params)
thinking_requested: Final = self.is_thinking_requested(request_params)
native_schema: Final = (
self._native_output_schema(model, json_schema, thinking_requested, model_info)
self._native_output_schema(
model,
json_schema,
will_think=AnthropicModelInfo.will_think(model, "bedrock", request_params),
thinking_requested=thinking_requested,
model_info=model_info,
)
if json_schema is not None
else None
)
@ -1359,17 +1371,24 @@ class AmazonConverseConfig(BaseConfig):
elif json_schema is not None:
# Fallback: translate to a synthetic tool call
# https://docs.anthropic.com/en/docs/build-with-claude/tool-use#json-mode
tool_is_forced: Final = litellm.utils.supports_tool_choice(
model=model, custom_llm_provider=self.custom_llm_provider
) and AnthropicModelInfo.response_format_tool_can_be_forced(model, "bedrock", request_params)
schema_description: Final = description if isinstance(description, str) else None
instructs_unforced_tool: Final = not tool_is_forced and BedrockModelInfo.get_base_model(model).startswith(
"anthropic"
)
_tool: Final = self._create_json_tool_call_for_response_format(
json_schema=json_schema,
description=description,
description=(
" ".join(filter(None, (RESPONSE_FORMAT_UNFORCED_TOOL_DESCRIPTION, schema_description)))
if instructs_unforced_tool
else schema_description
),
)
optional_params = self._add_tools_to_optional_params(optional_params=optional_params, tools=[_tool])
if (
litellm.utils.supports_tool_choice(model=model, custom_llm_provider=self.custom_llm_provider)
and not is_thinking_enabled
and not AnthropicModelInfo.forced_tool_use_unsupported(model)
):
if tool_is_forced:
optional_params["tool_choice"] = ToolChoiceValuesBlock(
tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME)
)
@ -1380,8 +1399,8 @@ class AmazonConverseConfig(BaseConfig):
# returns {} instead of the requested JSON. The model already
# returns JSON when the prompt asks for it.
# A request that goes native only because the model can't force the json tool keeps the json tool's faked
# stream, so it still needs only bedrock:InvokeModel.
# A request that goes native without asking for thinking keeps the json tool's faked stream, so it still
# needs only bedrock:InvokeModel, unless the model takes native output on every request.
streams_natively: Final = native_schema is not None and bedrock_chat_opts_into_native_structured_output(
model, thinking_requested, self.custom_llm_provider or "bedrock_converse", model_info
)
@ -2219,8 +2238,8 @@ class AmazonConverseConfig(BaseConfig):
}
if additional_request_params:
data["additionalModelRequestFields"] = additional_request_params
if "thinking" in additional_request_params:
data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",)
if "thinking" in additional_request_params or AnthropicModelInfo.thinks_by_default(model, "bedrock"):
data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",)
if system_content_blocks:
data["system"] = system_content_blocks

View file

@ -142,14 +142,19 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
json_schema: Mapping[str, object] | None,
request_params: Mapping[str, object],
) -> bool:
return self.chat_uses_native_structured_output(
model, self.is_thinking_requested(request_params), self._deployment_model_info
) and (json_schema is None or bedrock_cleaned_output_schema(model, json_schema) is not None)
return self.chat_uses_native_structured_output(model, request_params, self._deployment_model_info) and (
json_schema is None or bedrock_cleaned_output_schema(model, json_schema) is not None
)
def chat_uses_native_structured_output(
self, model: str, thinking_requested: bool, model_info: Mapping[str, object] | None
self, model: str, request_params: Mapping[str, object], model_info: Mapping[str, object] | None
) -> bool:
return bedrock_chat_uses_native_structured_output(model, thinking_requested, model_info=model_info)
return bedrock_chat_uses_native_structured_output(
model,
will_think=self.will_think(model, "bedrock", request_params),
thinking_requested=self.is_thinking_requested(request_params),
model_info=model_info,
)
def uses_opt_in_native_output(
self, model: str, optional_params: Mapping[str, object], litellm_params: Mapping[str, object]
@ -359,7 +364,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
request_body=anthropic_request,
model_info=model_info,
chat_uses_native_output=lambda: self.chat_uses_native_structured_output(
model, self.is_thinking_requested(_NON_DEFAULT_PARAMS.validate_python(filtered_params)), model_info
model, _NON_DEFAULT_PARAMS.validate_python(filtered_params), model_info
),
)
strip_unsupported_bedrock_invoke_output_config_keys(

View file

@ -55,9 +55,11 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
)
def chat_uses_native_structured_output(
self, model: str, thinking_requested: bool, model_info: Mapping[str, object] | None
self, model: str, request_params: Mapping[str, object], model_info: Mapping[str, object] | None
) -> bool:
return bedrock_chat_opts_into_native_structured_output(model, thinking_requested, model_info=model_info)
return bedrock_chat_opts_into_native_structured_output(
model, self.will_think(model, "bedrock", request_params), model_info=model_info
)
def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]:
return (f"bedrock_mantle/{model}", *super()._default_max_tokens_models(model, litellm_params))

View file

@ -356,13 +356,15 @@ def bedrock_chat_opts_into_native_structured_output(
def bedrock_chat_uses_native_structured_output(
model: str,
*,
will_think: bool,
thinking_requested: bool,
custom_llm_provider: str = "bedrock",
model_info: Mapping[str, object] | None = None,
) -> bool:
return bedrock_chat_opts_into_native_structured_output(
model, thinking_requested, custom_llm_provider, model_info
) or (not thinking_requested and _bedrock_native_output_replaces_unforceable_tool(model, model_info))
return bedrock_chat_opts_into_native_structured_output(model, will_think, custom_llm_provider, model_info) or (
not thinking_requested and _bedrock_native_output_replaces_unforceable_tool(model, model_info)
)
def bedrock_native_output_is_deployment_opt_in(

View file

@ -31,6 +31,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.base_utils import (
type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # base_utils helper returns an untyped dict
)
from litellm.llms.base_llm.chat.transformation import REQUEST_PARAMS
from litellm.types.llms.anthropic import AllAnthropicToolsValues
from litellm.types.llms.databricks import (
AllDatabricksContentValues,
@ -541,12 +542,12 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
model: str,
value: dict | None,
optional_params: dict,
is_thinking_enabled: bool,
tool_is_forced: bool,
) -> DatabricksTool | None:
if value is None:
return None
tool: Final = self.map_response_format_to_anthropic_tool(value, optional_params, is_thinking_enabled)
tool: Final = self.map_response_format_to_anthropic_tool(value, optional_params, tool_is_forced)
databricks_tool: Final = self.convert_anthropic_tool_to_databricks_tool(tool)
return databricks_tool
@ -579,7 +580,6 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
if "claude" in model
else non_default_params
)
is_thinking_enabled: Final = self.is_thinking_enabled(request_params)
mapped_params: Final = super().map_openai_params(request_params, optional_params, model, drop_params)
if "claude" in model:
AnthropicConfig.translate_legacy_thinking_for_adaptive_model(
@ -596,17 +596,20 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
mapped_params.pop("max_completion_tokens", None)
if "response_format" in request_params and "claude" in model:
tool_is_forced: Final = AnthropicConfig.response_format_tool_can_be_forced(
model, "databricks", REQUEST_PARAMS.validate_python(request_params)
)
_tool: Final = self.map_response_format_to_databricks_tool(
model,
request_params["response_format"],
mapped_params,
is_thinking_enabled,
tool_is_forced,
)
if _tool is not None:
self._add_tools_to_optional_params(optional_params=optional_params, tools=[_tool])
optional_params["json_mode"] = True
if not is_thinking_enabled:
if tool_is_forced:
_tool_choice: Final = ChatCompletionToolChoiceObjectParam(
type="function",
function=ChatCompletionToolChoiceFunctionParam(name=RESPONSE_FORMAT_TOOL_NAME),

View file

@ -6,6 +6,7 @@ Why separate file? Make it easy to see how transformation works
Docs - https://docs.mistral.ai/api/
"""
import reprlib
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from itertools import accumulate, count
from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload
@ -45,7 +46,13 @@ from litellm.types.llms.mistral import (
MistralThinkingBlock,
MistralToolCallMessage,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionTextObject
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
ChatCompletionTextObject,
ChatCompletionToolChoiceFunctionParam,
ChatCompletionToolChoiceObjectParam,
)
from litellm.types.utils import ModelResponse, ModelResponseStream
from litellm.utils import convert_to_model_response_object, supports_reasoning
@ -79,6 +86,10 @@ _REASONING_PROMPT_REASON: Final = "litellm sends Magistral a reasoning system pr
_NON_LEVEL_EFFORT_REASON: Final = (
"the Mistral API takes it only as a level such as 'high', so Mistral applies its default effort."
)
_UNUSABLE_TOOL_CHOICE_REASON: Final = (
"the Mistral API takes tool_choice only as 'auto', 'none', 'any', 'required', "
"or {'type': 'function', 'function': {'name': <tool name>}}."
)
_DOCUMENTED_MAGISTRAL_EFFORTS: Final = frozenset({"low", "medium", "high"})
_REASONING_ON_THINKING_TYPES: Final = frozenset({"enabled", "adaptive"})
@ -118,7 +129,7 @@ class MistralConfig(OpenAIGPTConfig):
- `tools` (list or null): A list of available tools for the model. Use this to specify functions for which the model can generate JSON inputs.
- `tool_choice` (string - 'auto'/'any'/'none' or null): Specifies if/how functions are called. If set to none the model won't call a function and will generate a message instead. If set to auto the model can choose to either generate a message or call a function. If set to any the model is forced to call a function. Default - 'auto'.
- `tool_choice` (string - 'auto'/'any'/'none', an object naming one function such as {"type": "function", "function": {"name": "my_function"}}, or null): Specifies if/how functions are called. If set to none the model won't call a function and will generate a message instead. If set to auto the model can choose to either generate a message or call a function. If set to any the model is forced to call a function. Default - 'auto'.
- `stop` (string or array of strings): Stop generation if this token is detected. Or if one of these tokens is detected when providing an array
@ -135,7 +146,7 @@ class MistralConfig(OpenAIGPTConfig):
top_p: int | None = None
max_tokens: int | None = None
tools: list | None = None
tool_choice: Literal["auto", "any", "none"] | None = None
tool_choice: Literal["auto", "any", "none"] | ChatCompletionToolChoiceObjectParam | None = None
random_seed: int | None = None
safe_prompt: bool | None = None
response_format: dict | None = None
@ -147,7 +158,7 @@ class MistralConfig(OpenAIGPTConfig):
top_p: int | None = None,
max_tokens: int | None = None,
tools: list | None = None,
tool_choice: Literal["auto", "any", "none"] | None = None,
tool_choice: Literal["auto", "any", "none"] | ChatCompletionToolChoiceObjectParam | None = None,
random_seed: int | None = None,
safe_prompt: bool | None = None,
response_format: dict | None = None,
@ -192,13 +203,40 @@ class MistralConfig(OpenAIGPTConfig):
*(("prompt_cache_key",) if accepts_prompt_cache_key else ()),
]
def _map_tool_choice(self, tool_choice: str) -> str:
if tool_choice == "auto" or tool_choice == "none":
return tool_choice
elif tool_choice == "required":
return "any"
else: # openai 'tool_choice' object param not supported by Mistral API
return "any"
@staticmethod
def _map_tool_choice(tool_choice: object) -> str | ChatCompletionToolChoiceObjectParam | None:
match tool_choice:
case str(choice) if choice in ("auto", "none"):
return choice
case "required" | "any":
return "any"
case {"type": "function", "function": {"name": str(name)}} if name:
return ChatCompletionToolChoiceObjectParam(
type="function", function=ChatCompletionToolChoiceFunctionParam(name=name)
)
case _:
return None
def _tool_choice_to_send(
self, model: str, tool_choice: object, drop_params: bool
) -> str | ChatCompletionToolChoiceObjectParam | UnsupportedParamsError | None:
mapped: Final = self._map_tool_choice(tool_choice)
if mapped is not None:
return mapped
if drop_params or litellm.drop_params is True:
warn_dropped_reasoning(
"tool_choice", tool_choice, model, _UNUSABLE_TOOL_CHOICE_REASON, drop_params_gated=True
)
return None
return UnsupportedParamsError(
message=(
f"Unsupported tool_choice={reprlib.repr(tool_choice)} for model={model}, since "
f"{_UNUSABLE_TOOL_CHOICE_REASON} "
"To drop it, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n"
),
llm_provider=self.custom_llm_provider,
model=model,
)
@staticmethod
def _get_mistral_reasoning_system_prompt() -> str:
@ -244,18 +282,20 @@ class MistralConfig(OpenAIGPTConfig):
optional_params["tools"] = self._clean_tool_schema_for_mistral(value)
if param == "stream" and value is True:
optional_params["stream"] = value
if param == "temperature":
optional_params["temperature"] = value
if param == "top_p":
optional_params["top_p"] = value
if param == "stop":
optional_params["stop"] = value
if param == "tool_choice" and isinstance(value, str):
optional_params["tool_choice"] = self._map_tool_choice(tool_choice=value)
if param == "tool_choice" and "tools" in non_default_params and non_default_params["tools"]:
match self._tool_choice_to_send(
model,
cast(object, value), # cast-ok: untyped request params, the match narrows it
drop_params,
):
case UnsupportedParamsError() as unusable_tool_choice:
raise unusable_tool_choice
case None:
pass
case tool_choice:
optional_params["tool_choice"] = tool_choice
if param == "seed":
optional_params["extra_body"] = {"random_seed": value}
if param == "response_format":
optional_params["response_format"] = value
if (
param == "reasoning_effort"
and "magistral" not in model.lower()
@ -271,7 +311,7 @@ class MistralConfig(OpenAIGPTConfig):
# Flag that we need to add reasoning system prompt
optional_params["_add_reasoning_prompt"] = True
_warn_unless_reasoning_prompt_honors(param, OPAQUE_VALUE.validate_python(value), model)
if param in ("parallel_tool_calls", "prompt_cache_key"):
if param in ("temperature", "top_p", "stop", "response_format", "parallel_tool_calls", "prompt_cache_key"):
optional_params[param] = value
return optional_params

View file

@ -13,24 +13,26 @@ The handler is intentionally a thin orchestration layer:
2. Build ``RecognitionConfig`` / ``StreamingRecognitionConfig`` protobufs
from the structured dict produced by
:class:`NvidiaRivaAudioTranscriptionConfig`.
3. Construct ``riva.client.Auth`` honoring NVCF (function-id metadata + TLS)
vs self-hosted (any host:port, optional TLS) modes.
4. Stream the audio through Riva's ``streaming_response_generator`` and
aggregate ``is_final`` results into a single transcript, or, when
``riva_offline`` is set, send the whole buffer through the unary
``offline_recognize`` RPC for NIMs that only ship offline engines.
3. Open a gRPC channel honoring NVCF (TLS by default) vs self-hosted (any
host:port, optional TLS) modes, and attach the ``authorization`` and NVCF
``function-id`` metadata to each call exactly once.
4. Stream the audio through Riva's ``StreamingRecognize`` RPC and aggregate
``is_final`` results into a single transcript, or, when ``riva_offline``
is set, send the whole buffer through the unary ``Recognize`` RPC for NIMs
that only ship offline engines.
5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed
on ``_hidden_params`` so cost calculation works.
``riva-client`` is imported lazily so ``litellm`` core remains usable
without the optional STT extras installed.
``grpcio`` and ``protobuf`` are imported lazily so ``litellm`` core remains
usable without the optional STT extras installed.
"""
from __future__ import annotations
import asyncio
import inspect
from collections.abc import Callable, Iterable, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol
from collections.abc import Iterable, Sequence
from contextlib import closing
from typing import TYPE_CHECKING, Final
import httpx
@ -64,89 +66,19 @@ if TYPE_CHECKING:
Logging as LiteLLMLoggingObj,
)
# Stream audio to Riva in ~50 ms slices (1600 samples at 16 kHz). Matches
# NVIDIA's recommended chunk size for streaming ASR — small enough for
# responsive endpointing, large enough to keep per-RPC overhead low.
_DEFAULT_CHUNK_SAMPLES: Final = 1600
_DEFAULT_CHUNK_BYTES: Final = _DEFAULT_CHUNK_SAMPLES * 2 # int16 = 2 bytes/sample
_RIVA_INSTALL_HINT = "NVIDIA Riva client is not installed. Install with `pip install 'litellm[stt-nvidia-riva]'`."
class _RivaAuth(Protocol):
"""Opaque ``riva.client.Auth`` handle."""
class _RecognizeFuture(Protocol):
def result(self, timeout: float | None = None) -> object: ...
def done(self) -> bool: ...
def cancel(self) -> bool: ...
class _AsrService(Protocol):
@property
def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ...
def offline_recognize(self, audio_bytes: bytes, config: object, future: bool) -> _RecognizeFuture: ...
class _EndpointingConfig(Protocol):
"""Opaque ``EndpointingConfig`` protobuf message."""
class _EndpointingConfigField(Protocol):
CopyFrom: Callable[[_EndpointingConfig], None]
class _RecognitionConfig(Protocol):
@property
def endpointing_config(self) -> _EndpointingConfigField: ...
class _StreamingRecognitionConfig(Protocol):
"""Opaque ``StreamingRecognitionConfig`` protobuf message."""
class _AudioEncoding(Protocol):
@property
def LINEAR_PCM(self) -> object: ...
class _RivaClientModule(Protocol):
"""The ``riva.client`` entry points this handler calls."""
@property
def Auth(self) -> Callable[..., _RivaAuth]: ...
@property
def ASRService(self) -> Callable[[_RivaAuth], _AsrService]: ...
class _RivaAsrModule(Protocol):
"""The protobuf constructors this handler calls, from whichever module exposes them."""
@property
def AudioEncoding(self) -> _AudioEncoding: ...
@property
def RecognitionConfig(self) -> Callable[..., _RecognitionConfig]: ...
@property
def StreamingRecognitionConfig(self) -> Callable[..., _StreamingRecognitionConfig]: ...
@property
def EndpointingConfig(self) -> Callable[..., _EndpointingConfig]: ...
def _auth_factory(riva_module: _RivaClientModule) -> Callable[..., _RivaAuth]:
return riva_module.Auth
def _audio_encoding(riva_asr_module: _RivaAsrModule) -> _AudioEncoding:
return riva_asr_module.AudioEncoding
def _use_ssl(optional_params: dict) -> bool:
"""
NVCF endpoints are TLS-only, so ``nvcf_function_id`` defaults ``use_ssl``
to True. Self-hosted Riva defaults to plaintext but honors an explicit
``use_ssl``, since TLS-terminating ingresses in front of it are common.
"""
use_ssl_override: Final = optional_params.get("use_ssl")
if use_ssl_override is not None:
return bool(use_ssl_override)
return bool(optional_params.get("nvcf_function_id"))
def _parse_riva_offline(value: object, drop_params: bool) -> bool:
@ -192,7 +124,7 @@ class NvidiaRivaAudioTranscription:
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float | str | httpx.Timeout | None,
logging_obj: "LiteLLMLoggingObj",
logging_obj: LiteLLMLoggingObj,
api_key: str | None,
api_base: str | None,
atranscription: bool = False,
@ -237,13 +169,11 @@ class NvidiaRivaAudioTranscription:
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float | str | httpx.Timeout | None,
logging_obj: "LiteLLMLoggingObj",
logging_obj: LiteLLMLoggingObj,
api_key: str | None,
api_base: str | None,
provider_config: NvidiaRivaAudioTranscriptionConfig | None = None,
) -> TranscriptionResponse:
# ``riva-client`` exposes a sync streaming generator, so we offload
# the blocking call to a worker thread to keep the event loop free.
return await asyncio.to_thread(
self._run_sync,
model=model,
@ -267,7 +197,7 @@ class NvidiaRivaAudioTranscription:
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float | str | httpx.Timeout | None,
logging_obj: "LiteLLMLoggingObj",
logging_obj: LiteLLMLoggingObj,
api_key: str | None,
api_base: str | None,
provider_config: NvidiaRivaAudioTranscriptionConfig,
@ -317,18 +247,12 @@ class NvidiaRivaAudioTranscription:
) or litellm.drop_params is True
riva_offline: Final = _parse_riva_offline(optional_params.get("riva_offline"), drop_params=drop_params)
riva_module, riva_asr_module = _import_riva()
auth_obj: Final = self._construct_auth(
riva_module=riva_module,
api_base=api_base,
api_key=api_key,
optional_params=optional_params,
)
try:
from litellm.llms.nvidia_riva.audio_transcription import riva_grpc
except ImportError as e:
raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e
recognition_config: Final = self._build_recognition_config_proto(
riva_asr_module=riva_asr_module,
recognition_config_dict=recognition_config_dict,
)
recognition_config: Final = riva_grpc.recognition_config(recognition_config_dict)
logging_obj.pre_call(
input=None,
api_key=api_key,
@ -344,26 +268,18 @@ class NvidiaRivaAudioTranscription:
},
)
metadata: Final = riva_grpc.call_metadata(api_key, optional_params.get("nvcf_function_id"))
try:
asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj)
responses: Final = (
self._offline_responses(
asr_service=asr_service,
pcm_bytes=resampled.pcm_bytes,
recognition_config=recognition_config,
timeout=None if timeout is None else _deadline_seconds(timeout, drop_params),
deadline: Final = None if timeout is None else _deadline_seconds(timeout, drop_params)
with closing(riva_grpc.open_channel(api_base, _use_ssl(optional_params))) as channel:
responses: Final = (
riva_grpc.recognize(channel, recognition_config, resampled.pcm_bytes, metadata, deadline)
if riva_offline
else riva_grpc.streaming_recognize(
channel, recognition_config, resampled.pcm_bytes, metadata, deadline
)
)
if riva_offline
else self._streaming_responses(
asr_service=asr_service,
riva_asr_module=riva_asr_module,
pcm_bytes=resampled.pcm_bytes,
recognition_config=recognition_config,
timeout=timeout,
drop_params=drop_params,
)
)
final_results: Final = self._collect_final_results(responses, require_final=not riva_offline)
final_results: Final = self._collect_final_results(responses, require_final=not riva_offline)
except NvidiaRivaException:
raise
except Exception as e:
@ -400,142 +316,6 @@ class NvidiaRivaAudioTranscription:
return final_response
def _construct_auth(
self,
riva_module: _RivaClientModule,
api_base: str,
api_key: str | None,
optional_params: dict,
) -> _RivaAuth:
"""
Build a ``riva.client.Auth`` object.
- When ``nvcf_function_id`` is provided we attach the NVCF
``function-id`` and bearer ``authorization`` metadata, and default
``use_ssl`` to True (NVCF endpoints are TLS-only).
- Otherwise (self-hosted) we default ``use_ssl`` to False but still
honor an explicit override — self-hosted Riva behind an ingress
with TLS termination is a real deployment topology.
"""
nvcf_function_id: Final = optional_params.get("nvcf_function_id")
use_ssl_override: Final = optional_params.get("use_ssl")
use_ssl: Final = bool(use_ssl_override) if use_ssl_override is not None else bool(nvcf_function_id)
metadata: Final[list[tuple[str, str]]] = []
if nvcf_function_id:
metadata.append(("function-id", str(nvcf_function_id)))
if api_key:
metadata.append(("authorization", f"Bearer {api_key}"))
try:
return _auth_factory(riva_module)(uri=api_base, use_ssl=use_ssl, metadata_args=metadata)
except TypeError:
# Older riva-client signatures used positional-only args.
return _auth_factory(riva_module)(None, use_ssl, api_base, metadata)
def _build_recognition_config_proto(
self, riva_asr_module: _RivaAsrModule, recognition_config_dict: dict[str, Any]
) -> _RecognitionConfig:
encoding_name: Final = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper()
encoding_enum: Final[object] = getattr(
_audio_encoding(riva_asr_module),
encoding_name,
_audio_encoding(riva_asr_module).LINEAR_PCM,
)
config: Final[_RecognitionConfig] = riva_asr_module.RecognitionConfig(
encoding=encoding_enum,
sample_rate_hertz=int(recognition_config_dict["sample_rate_hertz"]),
language_code=recognition_config_dict["language_code"],
audio_channel_count=int(recognition_config_dict["audio_channel_count"]),
enable_automatic_punctuation=bool(recognition_config_dict.get("enable_automatic_punctuation", True)),
enable_word_time_offsets=bool(recognition_config_dict.get("enable_word_time_offsets", False)),
max_alternatives=int(recognition_config_dict.get("max_alternatives", 1)),
model=recognition_config_dict.get("model", "") or "",
verbatim_transcripts=bool(recognition_config_dict.get("verbatim_transcripts", False)),
profanity_filter=bool(recognition_config_dict.get("profanity_filter", False)),
)
endpointing: Final = recognition_config_dict.get("endpointing_config")
if isinstance(endpointing, dict) and endpointing:
try:
ep: Final[_EndpointingConfig] = riva_asr_module.EndpointingConfig(**endpointing)
config.endpointing_config.CopyFrom(ep)
except Exception:
# If the user supplied an unknown EndpointingConfig field
# (older Riva server), fall back to Riva's defaults rather
# than failing the whole request.
pass
return config
def _streaming_responses(
self,
asr_service: _AsrService,
riva_asr_module: _RivaAsrModule,
pcm_bytes: bytes,
recognition_config: _RecognitionConfig,
timeout: float | str | httpx.Timeout | None,
drop_params: bool,
) -> Iterable[object]:
streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig(
config=recognition_config, interim_results=False
)
return asr_service.streaming_response_generator(
audio_chunks=self._iter_audio_chunks(pcm_bytes),
streaming_config=streaming_config,
**self._timeout_kwargs(asr_service.streaming_response_generator, timeout, drop_params),
)
def _offline_responses(
self,
asr_service: _AsrService,
pcm_bytes: bytes,
recognition_config: _RecognitionConfig,
timeout: float | None,
) -> Iterable[object]:
pending: Final = asr_service.offline_recognize(audio_bytes=pcm_bytes, config=recognition_config, future=True)
try:
return (pending.result(timeout=timeout),)
except Exception as e:
if pending.done():
raise
pending.cancel()
raise NvidiaRivaException(
status_code=408,
message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds",
) from e
@classmethod
def _timeout_kwargs(
cls, rpc: Callable[..., object], timeout: float | str | httpx.Timeout | None, drop_params: bool
) -> Mapping[str, float]:
# Forward the deadline so the call cannot block forever if the
# server stalls. Older riva-client versions do not accept a
# ``timeout`` kwarg, so pass it only when supported.
if timeout is None or not cls._supports_timeout_kwarg(rpc):
return MappingProxyType({})
return MappingProxyType({"timeout": _deadline_seconds(timeout, drop_params)})
@staticmethod
def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool:
try:
sig: Final = inspect.signature(callable_obj)
except (TypeError, ValueError):
return False
params: Final = sig.parameters
if "timeout" in params:
return True
return any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
@staticmethod
def _iter_audio_chunks(pcm_bytes: bytes):
for offset in range(0, len(pcm_bytes), _DEFAULT_CHUNK_BYTES):
chunk = pcm_bytes[offset : offset + _DEFAULT_CHUNK_BYTES]
if not chunk:
continue
yield chunk
@staticmethod
def _collect_final_results(stream: Iterable[object], require_final: bool = True) -> Sequence[RivaFinalResult]:
"""
@ -570,27 +350,3 @@ class NvidiaRivaAudioTranscription:
}
)
return final_results
def _import_riva() -> tuple[_RivaClientModule, _RivaAsrModule]:
"""
Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``.
We try the SDK first (preferred) and fall back to importing the proto
module separately when the SDK packaging changes between versions.
"""
try:
import riva.client as riva_client
except ImportError as e:
raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e
riva_asr_module = riva_client
if not hasattr(riva_asr_module, "RecognitionConfig"):
try:
from riva.client.proto import riva_asr_pb2
riva_asr_module = riva_asr_pb2
except ImportError as e:
raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e
return riva_client, riva_asr_module

View file

@ -0,0 +1,199 @@
"""
The part of NVIDIA Riva's speech recognition gRPC schema that LiteLLM sends and reads.
Names and numbers match ``riva/proto/riva_asr.proto`` and ``riva/proto/riva_audio.proto`` in
https://github.com/nvidia-riva/common. Fields LiteLLM never sets or reads are left out, and the
wire format skips them. The messages live in their own descriptor pool so they never collide
with nvidia-riva-client's copies when both are imported in one process.
"""
from dataclasses import dataclass
from enum import Enum, IntEnum
from typing import Final
from google.protobuf import descriptor_pb2, descriptor_pool, message_factory
from google.protobuf.message import Message
_PACKAGE: Final = "nvidia.riva.asr"
_SERVICE: Final = f"/{_PACKAGE}.RivaSpeechRecognition"
RECOGNIZE_METHOD: Final = f"{_SERVICE}/Recognize"
STREAMING_RECOGNIZE_METHOD: Final = f"{_SERVICE}/StreamingRecognize"
class _Kind(Enum):
FLOAT = "TYPE_FLOAT"
INT32 = "TYPE_INT32"
BOOL = "TYPE_BOOL"
STRING = "TYPE_STRING"
MESSAGE = "TYPE_MESSAGE"
BYTES = "TYPE_BYTES"
ENUM = "TYPE_ENUM"
class AudioEncoding(IntEnum):
ENCODING_UNSPECIFIED = 0
LINEAR_PCM = 1
LINEAR_PCM: Final = AudioEncoding.LINEAR_PCM.value
@dataclass(frozen=True, slots=True)
class _Field:
name: str
number: int
kind: _Kind
type_name: str = ""
repeated: bool = False
optional: bool = False
oneof: str = ""
_MESSAGES: Final[tuple[tuple[str, tuple[_Field, ...]], ...]] = (
(
"RecognitionConfig",
(
_Field("encoding", 1, _Kind.ENUM, "AudioEncoding"),
_Field("sample_rate_hertz", 2, _Kind.INT32),
_Field("language_code", 3, _Kind.STRING),
_Field("max_alternatives", 4, _Kind.INT32),
_Field("profanity_filter", 5, _Kind.BOOL),
_Field("audio_channel_count", 7, _Kind.INT32),
_Field("enable_word_time_offsets", 8, _Kind.BOOL),
_Field("enable_automatic_punctuation", 11, _Kind.BOOL),
_Field("model", 13, _Kind.STRING),
_Field("verbatim_transcripts", 14, _Kind.BOOL),
_Field("endpointing_config", 25, _Kind.MESSAGE, "EndpointingConfig", optional=True),
),
),
(
"EndpointingConfig",
(
_Field("start_history", 1, _Kind.INT32, optional=True),
_Field("start_threshold", 2, _Kind.FLOAT, optional=True),
_Field("stop_history", 3, _Kind.INT32, optional=True),
_Field("stop_threshold", 4, _Kind.FLOAT, optional=True),
_Field("stop_history_eou", 5, _Kind.INT32, optional=True),
_Field("stop_threshold_eou", 6, _Kind.FLOAT, optional=True),
),
),
(
"StreamingRecognitionConfig",
(
_Field("config", 1, _Kind.MESSAGE, "RecognitionConfig"),
_Field("interim_results", 2, _Kind.BOOL),
),
),
(
"StreamingRecognizeRequest",
(
_Field("streaming_config", 1, _Kind.MESSAGE, "StreamingRecognitionConfig", oneof="streaming_request"),
_Field("audio_content", 2, _Kind.BYTES, oneof="streaming_request"),
),
),
(
"WordInfo",
(
_Field("start_time", 1, _Kind.INT32),
_Field("end_time", 2, _Kind.INT32),
_Field("word", 3, _Kind.STRING),
),
),
(
"SpeechRecognitionAlternative",
(
_Field("transcript", 1, _Kind.STRING),
_Field("words", 3, _Kind.MESSAGE, "WordInfo", repeated=True),
),
),
(
"StreamingRecognitionResult",
(
_Field("alternatives", 1, _Kind.MESSAGE, "SpeechRecognitionAlternative", repeated=True),
_Field("is_final", 2, _Kind.BOOL),
),
),
(
"StreamingRecognizeResponse",
(_Field("results", 1, _Kind.MESSAGE, "StreamingRecognitionResult", repeated=True),),
),
(
"RecognizeRequest",
(
_Field("config", 1, _Kind.MESSAGE, "RecognitionConfig"),
_Field("audio", 2, _Kind.BYTES),
),
),
(
"SpeechRecognitionResult",
(_Field("alternatives", 1, _Kind.MESSAGE, "SpeechRecognitionAlternative", repeated=True),),
),
(
"RecognizeResponse",
(_Field("results", 1, _Kind.MESSAGE, "SpeechRecognitionResult", repeated=True),),
),
)
def _oneof_names(fields: tuple[_Field, ...]) -> tuple[str, ...]:
declared: Final = tuple(dict.fromkeys(field.oneof for field in fields if field.oneof))
synthetic: Final = tuple(f"_{field.name}" for field in fields if field.optional)
return declared + synthetic
def _field_proto(field: _Field, oneofs: tuple[str, ...]) -> descriptor_pb2.FieldDescriptorProto:
oneof: Final = f"_{field.name}" if field.optional else field.oneof
return descriptor_pb2.FieldDescriptorProto(
name=field.name,
number=field.number,
type=descriptor_pb2.FieldDescriptorProto.Type.Value(field.kind.value),
label=(
descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED
if field.repeated
else descriptor_pb2.FieldDescriptorProto.LABEL_OPTIONAL
),
type_name=f".{_PACKAGE}.{field.type_name}" if field.type_name else None,
oneof_index=oneofs.index(oneof) if oneof else None,
proto3_optional=field.optional or None,
)
def _message_proto(name: str, fields: tuple[_Field, ...]) -> descriptor_pb2.DescriptorProto:
oneofs: Final = _oneof_names(fields)
return descriptor_pb2.DescriptorProto(
name=name,
field=[_field_proto(field, oneofs) for field in fields],
oneof_decl=[descriptor_pb2.OneofDescriptorProto(name=oneof) for oneof in oneofs],
)
_FILE: Final = descriptor_pb2.FileDescriptorProto(
name="litellm/nvidia_riva/riva_asr.proto",
package=_PACKAGE,
syntax="proto3",
enum_type=[
descriptor_pb2.EnumDescriptorProto(
name="AudioEncoding",
value=[
descriptor_pb2.EnumValueDescriptorProto(name=encoding.name, number=encoding.value)
for encoding in AudioEncoding
],
)
],
message_type=[_message_proto(name, fields) for name, fields in _MESSAGES],
)
_POOL: Final = descriptor_pool.DescriptorPool()
_POOL.AddSerializedFile(_FILE.SerializeToString())
def _message_class(name: str) -> type[Message]:
return message_factory.GetMessageClass(_POOL.FindMessageTypeByName(f"{_PACKAGE}.{name}"))
RecognitionConfig: Final = _message_class("RecognitionConfig")
EndpointingConfig: Final = _message_class("EndpointingConfig")
StreamingRecognitionConfig: Final = _message_class("StreamingRecognitionConfig")
StreamingRecognizeRequest: Final = _message_class("StreamingRecognizeRequest")
StreamingRecognizeResponse: Final = _message_class("StreamingRecognizeResponse")
RecognizeRequest: Final = _message_class("RecognizeRequest")
RecognizeResponse: Final = _message_class("RecognizeResponse")

View file

@ -0,0 +1,113 @@
"""
gRPC calls to NVIDIA Riva's speech recognition service.
The ``authorization`` and NVCF ``function-id`` metadata ride on each call, never on the channel, so every
request carries them exactly once over both plaintext and TLS.
"""
from collections.abc import Iterator, Mapping, Sequence
from typing import Any, Final, TypeAlias
import grpc
from google.protobuf.message import Message
from pydantic import TypeAdapter
from litellm.llms.nvidia_riva.audio_transcription import riva_asr_schema as schema
from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException
# Stream audio to Riva in ~50 ms slices (1600 samples at 16 kHz). Matches
# NVIDIA's recommended chunk size for streaming ASR, small enough for
# responsive endpointing, large enough to keep per-RPC overhead low.
_CHUNK_SAMPLES: Final = 1600
_CHUNK_BYTES: Final = _CHUNK_SAMPLES * 2
CallMetadata: TypeAlias = tuple[tuple[str, str], ...]
_ENDPOINTING_FIELDS: Final = TypeAdapter(Mapping[str, object])
def call_metadata(api_key: str | None, nvcf_function_id: object) -> CallMetadata:
function_id: Final = (("function-id", str(nvcf_function_id)),) if nvcf_function_id else ()
authorization: Final = (("authorization", f"Bearer {api_key}"),) if api_key else ()
return function_id + authorization
def open_channel(api_base: str, use_ssl: bool) -> grpc.Channel:
if use_ssl:
return grpc.secure_channel(api_base, grpc.ssl_channel_credentials())
return grpc.insecure_channel(api_base)
def recognition_config(config: Mapping[str, Any]) -> Message:
endpointing: Final = _endpointing_config(config.get("endpointing_config"))
return schema.RecognitionConfig(
encoding=schema.LINEAR_PCM,
sample_rate_hertz=int(config["sample_rate_hertz"]),
language_code=config["language_code"],
audio_channel_count=int(config["audio_channel_count"]),
enable_automatic_punctuation=bool(config.get("enable_automatic_punctuation", True)),
enable_word_time_offsets=bool(config.get("enable_word_time_offsets", False)),
max_alternatives=int(config.get("max_alternatives", 1)),
model=config.get("model", "") or "",
verbatim_transcripts=bool(config.get("verbatim_transcripts", False)),
profanity_filter=bool(config.get("profanity_filter", False)),
**({} if endpointing is None else {"endpointing_config": endpointing}),
)
def _endpointing_config(endpointing: object) -> Message | None:
try:
fields: Final = _ENDPOINTING_FIELDS.validate_python(endpointing)
return schema.EndpointingConfig(**fields) if fields else None
except (ValueError, TypeError):
# A missing, unknown or mistyped field falls back to Riva's default
# endpointing rather than failing the whole request.
return None
def _streaming_requests(config: Message, pcm_bytes: bytes) -> Iterator[Message]:
yield schema.StreamingRecognizeRequest(
streaming_config=schema.StreamingRecognitionConfig(config=config, interim_results=False)
)
for offset in range(0, len(pcm_bytes), _CHUNK_BYTES):
yield schema.StreamingRecognizeRequest(audio_content=pcm_bytes[offset : offset + _CHUNK_BYTES])
def streaming_recognize(
channel: grpc.Channel,
config: Message,
pcm_bytes: bytes,
metadata: CallMetadata,
timeout: float | None,
) -> Iterator[Message]:
rpc: Final = channel.stream_stream(
schema.STREAMING_RECOGNIZE_METHOD,
request_serializer=schema.StreamingRecognizeRequest.SerializeToString,
response_deserializer=schema.StreamingRecognizeResponse.FromString,
)
return rpc(_streaming_requests(config, pcm_bytes), metadata=metadata, timeout=timeout)
def recognize(
channel: grpc.Channel,
config: Message,
pcm_bytes: bytes,
metadata: CallMetadata,
timeout: float | None,
) -> Sequence[Message]:
rpc: Final = channel.unary_unary(
schema.RECOGNIZE_METHOD,
request_serializer=schema.RecognizeRequest.SerializeToString,
response_deserializer=schema.RecognizeResponse.FromString,
)
pending: Final = rpc.future(schema.RecognizeRequest(config=config, audio=pcm_bytes), metadata=metadata)
try:
return (pending.result(timeout=timeout),)
except Exception as e:
if pending.done():
raise
pending.cancel()
raise NvidiaRivaException(
status_code=408,
message=f"NVIDIA Riva offline recognition did not finish within {timeout} seconds",
) from e

View file

@ -669,7 +669,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
prompt_cache_options: object = None,
) -> tuple[str | ResponseInputParam, Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None]:
validated_input: Final = self._validate_input_param(
self._drop_bridge_minted_reasoning_items(input), drops_set_created_by=drops_params(litellm_params)
self._drop_bridge_minted_reasoning_items(input), drop_params=drops_params(litellm_params)
)
reads_breakpoints: Final = self._reads_prompt_cache_breakpoints(
model=model, litellm_params=litellm_params, prompt_cache_options=prompt_cache_options
@ -951,7 +951,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
return tools if all(new is old for new, old in zip(sanitized, tools, strict=True)) else sanitized
def _validate_input_param(
self, input: str | ResponseInputParam, *, drops_set_created_by: bool = False
self, input: str | ResponseInputParam, *, drop_params: bool = False
) -> str | ResponseInputParam:
"""
Ensure all input fields if pydantic are converted to dict
@ -965,9 +965,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
# if it's pydantic, convert to dict
if isinstance(item, BaseModel):
validated_input.append(
replayable_responses_item(
item.model_dump(exclude_none=True), drops_set_created_by=drops_set_created_by
)
replayable_responses_item(item.model_dump(exclude_none=True), drops_set_created_by=drop_params)
)
elif isinstance(item, dict):
validated_input.append(
@ -975,7 +973,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
self._handle_reasoning_item(cast(dict[str, object], item))
if item.get("type") == "reasoning"
else item,
drops_set_created_by=drops_set_created_by,
drops_set_created_by=drop_params,
)
)
else:

View file

@ -89,7 +89,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig):
"""Handle preset/ model prefix: send as {"preset": name} instead of {"model": name}."""
input = self._ensure_message_type(input)
if model.startswith("preset/"):
input = self._validate_input_param(input, drops_set_created_by=drops_params(litellm_params))
input = self._validate_input_param(input, drop_params=drops_params(litellm_params))
data: Final[dict] = {
"preset": model[len("preset/") :],
"input": input,

View file

@ -2125,6 +2125,7 @@
"anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
@ -2213,6 +2214,7 @@
"global.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
@ -2301,6 +2303,7 @@
"us.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
@ -2389,6 +2392,7 @@
"eu.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
@ -2472,6 +2476,7 @@
"au.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
@ -2560,6 +2565,7 @@
"jp.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_1hr": 1.1e-05,
@ -2958,6 +2964,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -3002,6 +3009,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -3046,6 +3054,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -3090,6 +3099,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -3134,6 +3144,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -3178,6 +3189,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -4063,6 +4075,7 @@
"deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement",
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"input_cost_per_token": 5e-06,
"output_cost_per_token": 2.5e-05,
"litellm_provider": "azure_ai",
@ -4204,6 +4217,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -16507,6 +16521,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -17042,6 +17057,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -21910,6 +21926,7 @@
"prompt_cache_min_tokens": 512,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_mid_conversation_system": true,
@ -22068,6 +22085,7 @@
"prompt_cache_min_tokens": 1024,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_mid_conversation_system": true,
@ -36308,6 +36326,7 @@
"input_cost_per_image_token_batches": 4e-06,
"input_cost_per_token_batches": 2.5e-06,
"output_cost_per_image_token": 3.2e-05,
"output_cost_per_image_token_batches": 1.6e-05,
"output_cost_per_token_batches": 5e-06,
"source": "https://developers.openai.com/api/docs/pricing",
"supported_endpoints": [
@ -36331,6 +36350,7 @@
"input_cost_per_image_token_batches": 4e-06,
"input_cost_per_token_batches": 2.5e-06,
"output_cost_per_image_token": 3.2e-05,
"output_cost_per_image_token_batches": 1.6e-05,
"output_cost_per_token_batches": 5e-06,
"source": "https://developers.openai.com/api/docs/pricing",
"supported_endpoints": [
@ -36350,6 +36370,7 @@
"input_cost_per_image_token_batches": 4e-06,
"input_cost_per_token_batches": 2.5e-06,
"output_cost_per_image_token": 3e-05,
"output_cost_per_image_token_batches": 1.5e-05,
"source": "https://developers.openai.com/api/docs/pricing",
"supported_endpoints": [
"/v1/images/generations",
@ -36385,6 +36406,7 @@
"input_cost_per_image_token_batches": 4e-06,
"input_cost_per_token_batches": 2.5e-06,
"output_cost_per_image_token": 3e-05,
"output_cost_per_image_token_batches": 1.5e-05,
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
@ -36420,6 +36442,7 @@
"input_cost_per_image_token_batches": 4e-06,
"input_cost_per_token_batches": 2.5e-06,
"output_cost_per_image_token": 3e-05,
"output_cost_per_image_token_batches": 1.5e-05,
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
@ -39774,6 +39797,7 @@
"litellm_provider": "openai",
"mode": "image_generation",
"output_cost_per_image_token": 8e-06,
"output_cost_per_image_token_batches": 4e-06,
"source": "https://developers.openai.com/api/docs/pricing",
"supported_endpoints": [
"/v1/images/generations",
@ -50486,6 +50510,7 @@
"prompt_cache_min_tokens": 1024,
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -50563,6 +50588,7 @@
"prompt_cache_min_tokens": 512,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -52968,6 +52994,7 @@
"regional_endpoint_uplift_multiplier": 1.1,
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_creation_input_token_cost_batches": 3.125e-06,
@ -53006,6 +53033,7 @@
"regional_endpoint_uplift_multiplier": 1.1,
"supports_mid_conversation_system": true,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"cache_creation_input_token_cost": 6.25e-06,
"cache_creation_input_token_cost_above_1hr": 1e-05,
"cache_creation_input_token_cost_batches": 3.125e-06,
@ -53255,6 +53283,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -59906,6 +59935,7 @@
"litellm_provider": "openai",
"mode": "image_generation",
"output_cost_per_image_token": 3.2e-05,
"output_cost_per_image_token_batches": 1.6e-05,
"output_cost_per_token": 1e-05,
"output_cost_per_token_batches": 5e-06,
"source": "https://developers.openai.com/api/docs/pricing",
@ -60629,6 +60659,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -64719,7 +64750,8 @@
{
"name": "claude-family-baseline",
"pattern": "claude-[a-z]+-\\d+(?:[-.]\\d+)?(?:-\\d{8})?",
"description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number.",
"fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"],
"description": "Any Claude family-major id with an optional minor and an optional 8-digit date suffix, under any provider namespace (bare, bedrock-dotted, vertex, databricks, ...), so bare majors like claude-newfamily-5 get the same baseline as claude-newfamily-5-1. Carries the model-family facts every Claude shares; it carries no pricing, so cost stays on the standard unpriced behavior rather than a guessed number. On mapped Claude entries from Anthropic, Azure Foundry, Bedrock and Vertex it only fills the keys an entry leaves out, so the entry's own values always win.",
"model_info": {
"mode": "chat",
"max_input_tokens": 200000,
@ -64762,6 +64794,14 @@
"thinking_always_on": true
}
},
{
"name": "claude-thinking-on-by-default",
"pattern": "claude-[a-z]+-[5-9](?!\\d)",
"description": "Claude at major 5 or higher, in any id shape that contains claude-<family>-<major>. Anthropic's thinking docs list every Claude 5+ model as thinking by default (on or always on) when the request sends no thinking param, and every 4.x as off. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. The response_format json tool is then left unforced, since a forced tool call skips the model's thinking.",
"model_info": {
"thinking_on_by_default": true
}
},
{
"name": "claude-mid-conversation-system",
"pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)",
@ -69285,6 +69325,7 @@
"output_cost_per_token": 1.2e-05,
"prompt_cache_min_tokens": 1024,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -69360,6 +69401,7 @@
"output_cost_per_token": 3e-05,
"prompt_cache_min_tokens": 512,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -69616,6 +69658,7 @@
"output_cost_per_token": 1.2e-05,
"prompt_cache_min_tokens": 1024,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -69681,6 +69724,7 @@
"output_cost_per_token": 3e-05,
"prompt_cache_min_tokens": 512,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -75432,6 +75476,7 @@
"apac.anthropic.claude-opus-5": {
"bedrock_converse_supports_strict_tools": false,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"litellm_provider": "bedrock_converse",
"supports_tool_search": true,
@ -75548,6 +75593,7 @@
"max_output_tokens": 128000,
"mode": "chat",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77052,6 +77098,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77104,6 +77151,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77148,6 +77196,7 @@
"thinking_always_on": false,
"prompt_cache_min_tokens": 512,
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -77195,6 +77244,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77231,6 +77281,7 @@
"max_output_tokens": 128000,
"mode": "chat",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77287,6 +77338,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77333,6 +77385,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -77373,6 +77426,7 @@
"prompt_cache_min_tokens": 512,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -77414,6 +77468,7 @@
"prompt_cache_min_tokens": 512,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -77460,6 +77515,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77511,6 +77567,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77557,6 +77614,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77617,6 +77675,7 @@
"prompt_cache_min_tokens": 512,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -77668,6 +77727,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -77718,6 +77778,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -77762,6 +77823,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
@ -79521,6 +79583,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -79735,6 +79798,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -79901,6 +79965,7 @@
"search_context_size_medium": 0.01
},
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_mid_conversation_system": true,
"supports_assistant_prefill": false,
"supports_computer_use": true,
@ -81886,6 +81951,7 @@
"prompt_cache_min_tokens": 512.0,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_adaptive_thinking": true,
"thinking_on_by_default": true,
"supports_anthropic_thinking_payload": true,
"supports_assistant_prefill": false,
"supports_function_calling": true,
@ -85880,5 +85946,235 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"gpt-image-2.6-sunburst": {
"cache_read_input_image_token_cost": 2e-06,
"cache_read_input_token_cost": 1.25e-06,
"input_cost_per_image_token": 8e-06,
"input_cost_per_token": 5e-06,
"litellm_provider": "openai",
"mode": "image_generation",
"output_cost_per_image_token": 3e-05,
"source": "https://developers.openai.com/api/docs/pricing",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
],
"supports_pdf_input": true,
"supports_vision": true
},
"gpt-image-2.6-flare": {
"cache_read_input_image_token_cost": 2e-06,
"cache_read_input_token_cost": 1.25e-06,
"input_cost_per_image_token": 8e-06,
"input_cost_per_token": 5e-06,
"litellm_provider": "openai",
"mode": "image_generation",
"output_cost_per_image_token": 3e-05,
"source": "https://developers.openai.com/api/docs/pricing",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
],
"supports_pdf_input": true,
"supports_vision": true
},
"bedrock/us-gov-east-1/moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 4.95e-06,
"cache_creation_input_token_cost_flex": 2.475e-06,
"cache_creation_input_token_cost_priority": 8.6625e-06,
"cache_read_input_token_cost": 3.96e-07,
"cache_read_input_token_cost_flex": 1.98e-07,
"cache_read_input_token_cost_priority": 6.93e-07,
"input_cost_per_token": 3.96e-06,
"input_cost_per_token_batches": 1.98e-06,
"input_cost_per_token_flex": 1.98e-06,
"input_cost_per_token_priority": 6.93e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 868928,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.98e-05,
"output_cost_per_token_batches": 9.9e-06,
"output_cost_per_token_flex": 9.9e-06,
"output_cost_per_token_priority": 3.465e-05,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_cache_breakpoint": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_regex_lookaround": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"token_count_framing": 64.0,
"token_count_tokenizer": "kimi_k3"
},
"bedrock/us-gov-west-1/moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 4.95e-06,
"cache_creation_input_token_cost_flex": 2.475e-06,
"cache_creation_input_token_cost_priority": 8.6625e-06,
"cache_read_input_token_cost": 3.96e-07,
"cache_read_input_token_cost_flex": 1.98e-07,
"cache_read_input_token_cost_priority": 6.93e-07,
"input_cost_per_token": 3.96e-06,
"input_cost_per_token_batches": 1.98e-06,
"input_cost_per_token_flex": 1.98e-06,
"input_cost_per_token_priority": 6.93e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 868928,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.98e-05,
"output_cost_per_token_batches": 9.9e-06,
"output_cost_per_token_flex": 9.9e-06,
"output_cost_per_token_priority": 3.465e-05,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_cache_breakpoint": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_regex_lookaround": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"token_count_framing": 64.0,
"token_count_tokenizer": "kimi_k3"
},
"bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0": {
"deprecation_date": "2026-09-10",
"deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-3-haiku.md",
"input_cost_per_token": 3e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 195904,
"max_output_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-east-1/index.json",
"successor_model": "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0",
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/us-gov-east-1/xai.grok-4.7": {
"cache_read_input_token_cost": 6.6e-07,
"cache_read_input_token_cost_flex": 3.3e-07,
"cache_read_input_token_cost_priority": 1.155e-06,
"input_cost_per_token": 2.64e-06,
"input_cost_per_token_batches": 1.32e-06,
"input_cost_per_token_flex": 1.32e-06,
"input_cost_per_token_priority": 4.62e-06,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 524288,
"max_output_tokens": 524288,
"mode": "chat",
"output_cost_per_token": 7.92e-06,
"output_cost_per_token_batches": 3.96e-06,
"output_cost_per_token_flex": 3.96e-06,
"output_cost_per_token_priority": 1.386e-05,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-east-1/index.json",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"use_openai_responses_path": true
},
"bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0": {
"deprecation_date": "2026-09-10",
"deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-3-haiku.md",
"input_cost_per_token": 3e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 195904,
"max_output_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/us-gov-west-1/index.json",
"successor_model": "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0",
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"bedrock_mantle/us-gov-west-1/xai.grok-4.7": {
"cache_read_input_token_cost": 6.6e-07,
"cache_read_input_token_cost_flex": 3.3e-07,
"cache_read_input_token_cost_priority": 1.155e-06,
"input_cost_per_token": 2.64e-06,
"input_cost_per_token_batches": 1.32e-06,
"input_cost_per_token_flex": 1.32e-06,
"input_cost_per_token_priority": 4.62e-06,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 524288,
"max_output_tokens": 524288,
"mode": "chat",
"output_cost_per_token": 7.92e-06,
"output_cost_per_token_batches": 3.96e-06,
"output_cost_per_token_flex": 3.96e-06,
"output_cost_per_token_priority": 1.386e-05,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-gov-west-1/index.json",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"use_openai_responses_path": true
},
"bedrock_mantle/xai.grok-4.7": {
"cache_read_input_token_cost": 5.5e-07,
"cache_read_input_token_cost_flex": 2.75e-07,
"cache_read_input_token_cost_priority": 9.625e-07,
"input_cost_per_token": 2.2e-06,
"input_cost_per_token_batches": 1.1e-06,
"input_cost_per_token_flex": 1.1e-06,
"input_cost_per_token_priority": 3.85e-06,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 524288,
"max_output_tokens": 524288,
"mode": "chat",
"output_cost_per_token": 6.6e-06,
"output_cost_per_token_batches": 3.3e-06,
"output_cost_per_token_flex": 3.3e-06,
"output_cost_per_token_priority": 1.155e-05,
"source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"use_openai_responses_path": true
},
"together_ai/ricursive-ai/Ripple-65B-Preview": {
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 131072,
"max_output_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://api.together.ai/v1/models",
"supports_prompt_caching": true
}
}

View file

@ -1,6 +1,23 @@
from collections.abc import Mapping
from collections.abc import AsyncIterator, Mapping
from dataclasses import dataclass
from typing import Protocol
from datetime import datetime
from enum import Enum
from types import MappingProxyType
from typing import Literal, Protocol
def has_agent_skill_metadata(manifest: Mapping[str, object]) -> bool:
return isinstance(manifest.get("agent_skill"), dict)
def uploaded_plugin_source(base_url: str, name: str, digest: str) -> Mapping[str, str]:
return MappingProxyType(
{
"source": "archive",
"url": f"{base_url.rstrip('/')}/claude-code/plugins/{name}/archive/{digest}.zip",
"sha256": digest,
}
)
@dataclass(frozen=True, slots=True)
@ -19,11 +36,130 @@ class PluginArchiveMetadata:
enabled: bool
source: Mapping[str, str]
manifest_json: str | None
id: str | None = None
version: str | None = None
description: str | None = None
has_agent_skill_metadata: bool = False
updated_at: datetime | None = None
@dataclass(frozen=True, slots=True)
class PluginPublication:
name: str
version: str | None
description: str | None
manifest: Mapping[str, object]
created_by: str | None
content: bytes | None = None
@dataclass(frozen=True, slots=True)
class PluginPublicationBase:
id: str
updated_at: datetime | None
@dataclass(frozen=True, slots=True)
class PublishedPlugin:
id: str
name: str
version: str
enabled: bool
updated_at: datetime
@dataclass(frozen=True, slots=True)
class PluginPublicationError:
kind: Literal["invalid_version", "conflict", "not_found"]
message: str
@dataclass(frozen=True, slots=True)
class PluginVersion:
id: str
version: str | None
description: str | None
manifest: Mapping[str, object]
published_at: datetime | None
published_by: str | None
archive_available: bool
@dataclass(frozen=True, slots=True)
class PluginSnapshot:
id: str
name: str
version: str | None
description: str | None
manifest: Mapping[str, object]
enabled: bool
created_at: datetime | None
updated_at: datetime | None
created_by: str | None
@dataclass(frozen=True, slots=True)
class PluginDownloader:
authenticated: bool
user_id: str | None
user_email: str | None
key_alias: str | None
team_id: str | None
team_alias: str | None
class PluginDownloadResult(Enum):
RECORDED = "recorded"
SKILL_NOT_FOUND = "skill_not_found"
@dataclass(frozen=True, slots=True)
class PluginDownloadEvent:
id: str
downloaded_at: datetime
version: str | None
digest: str
authenticated: bool
user_id: str | None
user_email: str | None
key_alias: str | None
team_id: str | None
team_alias: str | None
@dataclass(frozen=True, slots=True)
class PluginDownloadsResponse:
total_downloads: int
anonymous_downloads: int
downloads: tuple[PluginDownloadEvent, ...]
next_cursor: str | None
class PluginUploadStore(Protocol):
async def create(self, upload: PluginUpload) -> str | None: ...
async def get_plugin(self, name: str) -> PluginSnapshot | None: ...
async def publish(
self, publication: PluginPublication, expected: PluginPublicationBase | None = None
) -> PublishedPlugin | PluginPublicationError: ...
async def versions(self, metadata: PluginArchiveMetadata, limit: int, offset: int) -> tuple[PluginVersion, ...]: ...
async def version(self, metadata: PluginArchiveMetadata, version_id: str) -> PluginVersion | None: ...
async def version_content(self, metadata: PluginArchiveMetadata, version_id: str) -> bytes | None: ...
async def get_metadata(self, name: str) -> PluginArchiveMetadata | None: ...
async def get_content(self, metadata: PluginArchiveMetadata) -> bytes | None: ...
def public_metadata(self) -> AsyncIterator[PluginArchiveMetadata]: ...
async def record_download(
self, metadata: PluginArchiveMetadata, digest: str, downloader: PluginDownloader
) -> PluginDownloadResult: ...
async def downloads(self, name: str, limit: int, cursor: int | None) -> PluginDownloadsResponse | None: ...
async def download_counts(self, plugin_ids: tuple[str, ...]) -> Mapping[str, int]: ...

View file

@ -0,0 +1,110 @@
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Literal, Protocol, TypeAlias
from litellm.models.plugin_upload import (
PluginPublication,
PluginPublicationBase,
PluginPublicationError,
PublishedPlugin,
)
from litellm.types.proxy.skill_scanning import SkillScan
SubmissionStatus: TypeAlias = Literal["pending_review", "approved", "rejected"]
SubmissionOperation: TypeAlias = Literal["create", "update"]
@dataclass(frozen=True, slots=True)
class SkillProposal:
name: str
version: str | None
description: str | None
manifest: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class SkillSubmission:
id: str
plugin_name: str
operation: SubmissionOperation
status: SubmissionStatus
proposal: SkillProposal
has_archive: bool
submitted_by: str
created_at: datetime
base_plugin_id: str | None
base_updated_at: datetime | None
reviewed_by: str | None
reviewed_at: datetime | None
rejection_reason: str | None
@dataclass(frozen=True, slots=True)
class SubmissionPage:
submissions: tuple[SkillSubmission, ...]
has_more: bool
@dataclass(frozen=True, slots=True)
class SubmissionLimitExceeded:
message: str
@dataclass(frozen=True, slots=True)
class SubmissionMissing:
pass
@dataclass(frozen=True, slots=True)
class SubmissionConflict:
message: str
@dataclass(frozen=True, slots=True)
class SubmissionReviewed:
submission: SkillSubmission
plugin: PublishedPlugin | None
published_source: Mapping[str, str] | None = None
SubmissionReviewResult: TypeAlias = SubmissionReviewed | SubmissionMissing | SubmissionConflict | PluginPublicationError
@dataclass(frozen=True, slots=True)
class SkillScanInput:
scan: SkillScan
source: Mapping[str, str]
content: bytes | None
class SkillSubmissionStore(Protocol):
async def get_scan(self, submission_id: str) -> SkillScan | SubmissionMissing | None: ...
async def start_scan(
self, submission_id: str, scan: SkillScan
) -> SkillScanInput | SubmissionMissing | SubmissionConflict: ...
async def finish_scan(self, submission_id: str, scan: SkillScan) -> None: ...
async def submit(
self, upload: PluginPublication, submitted_by: str, base: PluginPublicationBase | None
) -> SkillSubmission | SubmissionLimitExceeded: ...
async def list(
self, submitted_by: str | None, status: SubmissionStatus | None, limit: int, offset: int
) -> SubmissionPage: ...
async def get(self, submission_id: str, submitted_by: str | None) -> SkillSubmission | None: ...
async def get_archive(self, submission_id: str, submitted_by: str | None) -> bytes | None: ...
async def review(
self,
submission_id: str,
reviewer: str,
approve: bool,
reason: str | None,
*,
publication_base_url: str | None = None,
) -> SubmissionReviewResult: ...

View file

@ -48,6 +48,15 @@ def _mount_app(prefix: str, attr_name: str = "app") -> Callable[["FastAPI", obje
return _register
def _include_skill_hub(app: "FastAPI", module: object) -> None:
from litellm.proxy.discovery_endpoints.skill_hub_submissions import router
from litellm.proxy.discovery_endpoints.skill_scanning import router as scan_router
_include_router()(app, module)
app.include_router(router)
app.include_router(scan_router)
@dataclass(frozen=True)
class LazyFeature:
name: str
@ -283,6 +292,12 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
module_path="litellm.proxy.decisions_endpoints.endpoints",
path_prefixes=("/v1/decisions", "/decisions", "/v1/systemone", "/systemone", "/v1/extraction", "/extraction"),
),
LazyFeature(
name="skill_hub",
module_path="litellm.proxy.discovery_endpoints.skill_hub_endpoints",
path_prefixes=("/skill-hub",),
register_fn=_include_skill_hub,
),
LazyFeature(
name="claude_code_marketplace",
module_path="litellm.proxy.anthropic_endpoints.claude_code_endpoints",

File diff suppressed because it is too large Load diff

View file

@ -89,6 +89,7 @@ from litellm.types.proxy.carried_budget_state import (
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
from litellm.types.proxy.ldap_sync import LDAPSyncSettings
from litellm.types.proxy.request_tag_mappings import RequestTagMapping
from litellm.types.proxy.skill_scanning import SkillScannerConfig
from litellm.types.proxy.spend_capture_rate import SpendCaptureRateCheckSettings
from litellm.types.router import AllowedModelRegion, RouterErrors, UpdateRouterConfig
from litellm.types.router_weights import validate_router_settings_dict
@ -450,6 +451,28 @@ class LiteLLMRoutes(enum.Enum):
"/v1/traces/{trace_id}/spans/{span_id}/error",
)
skill_hub_routes = (
"/skill-hub/skills",
"/skill-hub/skills/upload",
"/skill-hub/skills/import",
"/skill-hub/skills/{name}",
"/skill-hub/skills/{name}/archive",
"/skill-hub/skills/{name}/content",
"/skill-hub/skills/{plugin_name}/downloads",
"/skill-hub/skills/{name}/versions",
"/skill-hub/skills/{name}/versions/{version_id}",
"/skill-hub/skills/{name}/versions/{version_id}/archive",
"/skill-hub/skills/{name}/publish",
"/skill-hub/skills/{name}/unpublish",
"/skill-hub/skills/{name}/upload",
"/skill-hub/submissions",
"/skill-hub/submissions/{submission_id}",
"/skill-hub/submissions/{submission_id}/archive",
"/skill-hub/submissions/{submission_id}/approve",
"/skill-hub/submissions/{submission_id}/reject",
"/skill-hub/submissions/{submission_id}/scan",
)
litellm_native_routes = [
"/rag/ingest",
"/v1/rag/ingest",
@ -469,6 +492,7 @@ class LiteLLMRoutes(enum.Enum):
"/claude-code/plugins",
"/claude-code/plugins/upload",
"/claude-code/plugins/{plugin_name}/archive/{digest}.zip",
"/claude-code/plugins/{plugin_name}/downloads",
"/claude-code/plugins/{plugin_name}",
]
@ -699,6 +723,7 @@ class LiteLLMRoutes(enum.Enum):
+ key_management_routes
+ mcp_management_routes
+ list(agent_management_routes)
+ list(skill_hub_routes)
)
spend_tracking_routes = [
@ -941,6 +966,7 @@ class LiteLLMRoutes(enum.Enum):
# Agent registry - reads are role-scoped and writes are proxy-admin-gated
# inside agent_endpoints/endpoints.py
*agent_management_routes,
*skill_hub_routes,
] # routes that manage their own allowed/disallowed logic
## Org Admin Routes ##
@ -2486,6 +2512,11 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
description="Periodically synchronize LDAP groups to teams and existing users' memberships. Requires Redis.",
)
use_spend_logs_ledger: bool = Field(
False,
description="Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.",
)
request_tag_mappings: Sequence[RequestTagMapping] | None = Field(
None,
description="Map W3C baggage values to request tags. Exact value/tag rules take precedence over tag_prefix fallbacks.",
@ -3020,6 +3051,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="If True, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription.",
)
skill_scanner: SkillScannerConfig | None = Field(
None, description="Opt-in external scanner for manual admin review of skill submissions"
)
mcp_required_fields: list[str] | None = Field(
None,
description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).",
@ -3806,6 +3840,9 @@ class SpendLogsMetadata(TypedDict):
class SpendLogsPayload(TypedDict):
accounting_payload: NotRequired[Mapping[str, object]] # writable-ok: capture before buffering
accounting_log_disabled: NotRequired[bool] # writable-ok: capture logging mode before buffering
accounting_delivery_id: NotRequired[str] # writable-ok: assigned once before buffering
request_id: str
call_type: str
api_key: str

View file

@ -0,0 +1 @@
"""HTTP endpoints for gateway alerts."""

View file

@ -0,0 +1,109 @@
"""Shared acknowledgement state for model migration alerts."""
import hashlib
import json
from datetime import date, datetime, timezone
from typing import Annotated, Final
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import TypeAdapter
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps.database import DatabaseDep
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.daily_activity_buckets import api_key_bucket, api_key_from_bucket
from litellm.repositories.table_repositories import DailyTagSpendRepository
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
from litellm.types.proxy.model_deprecation import DeprecationAlertState, DeprecationAlertUpdate, DeprecationTagUsage
router: Final = APIRouter()
def alert_state_repository(database: DatabaseDep) -> ConfigRepository:
if database is None:
raise HTTPException(status_code=503, detail="Alert state requires a database connection")
return ConfigRepository(database.prisma, use_writer=True)
def alert_tag_repository(database: DatabaseDep) -> DailyTagSpendRepository:
if database is None:
raise HTTPException(status_code=503, detail="Model usage tags require a database connection")
return DailyTagSpendRepository(database.prisma)
def _state_key(model: str, deprecation_date: date) -> str:
identity: Final = json.dumps([model, deprecation_date.isoformat()])
digest: Final = hashlib.sha256(identity.encode()).hexdigest()
return f"model_deprecation_alert:{digest}"
def _require_admin(user: UserAPIKeyAuth, *, write: bool) -> None:
if user.user_role == LitellmUserRoles.PROXY_ADMIN:
return
if not write and user.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY:
return
raise HTTPException(status_code=403, detail="Only gateway admins can manage model deprecation alerts")
@router.get("/alerting/model-deprecations/state", tags=["alerting"], response_model=DeprecationAlertState)
async def get_deprecation_alert_state(
model: Annotated[str, Query(min_length=1, max_length=1024)],
deprecation_date: date,
user: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
repository: Annotated[ConfigRepository, Depends(alert_state_repository)],
) -> DeprecationAlertState:
_require_admin(user, write=False)
record: Final = await repository.get_param(_state_key(model, deprecation_date))
if record is None:
return DeprecationAlertState(model=model, deprecation_date=deprecation_date)
return DeprecationAlertState.model_validate(record.param_value)
@router.put("/alerting/model-deprecations/state", tags=["alerting"], response_model=DeprecationAlertState)
async def update_deprecation_alert_state(
data: DeprecationAlertUpdate,
user: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
repository: Annotated[ConfigRepository, Depends(alert_state_repository)],
) -> DeprecationAlertState:
_require_admin(user, write=True)
state: Final = DeprecationAlertState(
**data.model_dump(), updated_at=datetime.now(timezone.utc), updated_by=user.user_id
)
await repository.set_param(_state_key(data.model, data.deprecation_date), state.model_dump(mode="json"))
return state
@router.get("/alerting/model-deprecations/tags", tags=["alerting"], response_model=tuple[DeprecationTagUsage, ...])
async def get_deprecation_usage_tags(
model: Annotated[str, Query(min_length=1, max_length=1024)],
api_keys: Annotated[tuple[str, ...], Query(min_length=1, max_length=100)],
start_date: date,
end_date: date,
user: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
repository: Annotated[DailyTagSpendRepository, Depends(alert_tag_repository)],
) -> tuple[DeprecationTagUsage, ...]:
"""Distinct tags observed with this model for a page of affected reporting key IDs."""
_require_admin(user, write=False)
if not 0 <= (end_date - start_date).days < 30:
raise HTTPException(status_code=422, detail="Select up to 30 consecutive calendar dates")
rows: Final = await repository.table.group_by(
by=["api_key", "tag"],
where={
"api_key": {
# bounded-ok: query validation limits api_keys to 100 IDs.
"in": [api_key_from_bucket(key) for key in api_keys]
},
"model": model,
"date": {"gte": start_date.isoformat(), "lte": end_date.isoformat()},
"api_requests": {"gt": 0},
"tag": {"not": None},
},
)
tags: Final = TypeAdapter(tuple[DeprecationTagUsage, ...]).validate_python(rows)
return tuple(
sorted(
(DeprecationTagUsage(api_key=api_key_bucket(usage.api_key), tag=usage.tag) for usage in tags),
key=lambda usage: (usage.api_key, usage.tag),
)
)

Some files were not shown because too many files have changed in this diff Show more