diff --git a/.github/PULL_REQUEST_TEMPLATE/general.md b/.github/PULL_REQUEST_TEMPLATE/general.md index 4f1a18e925e..f815601a29c 100644 --- a/.github/PULL_REQUEST_TEMPLATE/general.md +++ b/.github/PULL_REQUEST_TEMPLATE/general.md @@ -6,19 +6,12 @@ ## Handwritten 1-3 sentence description - -`` - -`` +[insert handwritten description here]
Details @@ -100,6 +93,7 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac Otherwise proof must be fully e2e, no mocks, with real provider calls where applicable. Pytest alone is insufficient. The following format applies to live proof Show only the latest proof. Before uses the merge base and After the tested tip. Rerun affected QA after behavior changes. Verified neutral diffs may carry proof with the tested SHA, current tip, equivalence reason and focused checks, following AGENTS.md and /qa Use matching Before/After case headings with numbered commands and observed output. Put shared setup above them. For one case, omit case headings + A reviewer must be able to rerun every case from the body alone, so the shared setup gives the proxy start command, the config.yaml, env var names (never values) and the client code or curl used, all exactly as run. Wrap any block longer than about 20 lines in a collapsed
with a one-line summary ### Before () diff --git a/.github/workflows/interactions-spec-canary.yml b/.github/workflows/interactions-spec-canary.yml index a2e4646dce5..89ebdd87a7c 100644 --- a/.github/workflows/interactions-spec-canary.yml +++ b/.github/workflows/interactions-spec-canary.yml @@ -49,9 +49,9 @@ jobs: id: compliance run: | uv run --no-sync pytest tests/unit/interactions/test_openapi_compliance.py \ - --interactions-spec-url=https://ai.google.dev/static/api/interactions.openapi.json \ + --interactions-spec-url=https://aistudio.google.com/api-reference/interactions.openapi.json \ --tb=short -vv -p no:rerunfailures -p no:pytest-retry \ - --allow-hosts=127.0.0.1,::1,localhost,ai.google.dev + --allow-hosts=127.0.0.1,::1,localhost,ai.google.dev,aistudio.google.com tracking-issue: name: tracking-issue diff --git a/AGENTS.md b/AGENTS.md index 9aa9ab32528..01127303da3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -45,7 +45,7 @@ Run QA proxies with `python litellm/proxy/proxy_cli.py --config litellm/proxy/de Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at `, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030) -For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce +For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. The body alone must let a reviewer rerun it, with the proxy command, config, env var names and client code exactly as run, and long blocks collapsed in
. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y: - don't use emojis diff --git a/cookbook/skill_scanner/README.md b/cookbook/skill_scanner/README.md new file mode 100644 index 00000000000..2a68e116002 --- /dev/null +++ b/cookbook/skill_scanner/README.md @@ -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 diff --git a/cookbook/skill_scanner/server.py b/cookbook/skill_scanner/server.py new file mode 100644 index 00000000000..44a1cecb00f --- /dev/null +++ b/cookbook/skill_scanner/server.py @@ -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"))) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql new file mode 100644 index 00000000000..ecfa7321d87 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql @@ -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"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql new file mode 100644 index 00000000000..e7975258ca9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql @@ -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); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql new file mode 100644 index 00000000000..6b70cca4a9a --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql @@ -0,0 +1,3 @@ +SET LOCAL lock_timeout = '10s'; + +ALTER TABLE "LiteLLM_SkillSubmission" ADD COLUMN IF NOT EXISTS "scan_json" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql new file mode 100644 index 00000000000..43e26b324a0 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql @@ -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 $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 78f16d9f1fd..818f2dc4858 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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 diff --git a/litellm-rust/.agents/skills/rust-gateway/SKILL.md b/litellm-rust/.agents/skills/rust-gateway/SKILL.md index 4fee0b54a68..923c9bf5ab6 100644 --- a/litellm-rust/.agents/skills/rust-gateway/SKILL.md +++ b/litellm-rust/.agents/skills/rust-gateway/SKILL.md @@ -7,7 +7,7 @@ metadata: # Rust gateway -`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching +`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. `gateway-forward` is the reverse proxy that relays every request to the Python proxy and has no routes of its own. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching - [surfaces](references/surfaces.md) - Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 23ebfc7d3cf..107b55e89cb 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -547,7 +547,7 @@ dependencies = [ "http 1.4.2", "http-body 0.4.6", "hyper 0.14.32", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-rustls 0.24.2", "hyper-rustls 0.27.9", "hyper-util", @@ -736,7 +736,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "itoa", "matchit", @@ -1000,7 +1000,7 @@ dependencies = [ "home", "http 1.4.2", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-named-pipe", "hyper-rustls 0.27.9", "hyper-util", @@ -2255,7 +2255,7 @@ dependencies = [ "chrono", "http 1.4.2", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-rustls 0.27.9", "hyper-util", "ring", @@ -2391,7 +2391,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "lazy_static", "opentelemetry", "opentelemetry-semantic-conventions", @@ -2833,9 +2833,9 @@ dependencies = [ [[package]] name = "hyper" -version = "1.10.1" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +checksum = "2c3e324da4c95177d6291d4c8730197c0d1822f8a9766814a4a44fa5ab797c9c" dependencies = [ "atomic-waker", "bytes", @@ -2860,7 +2860,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fab3637d6b04a8037af8a266fdf6cf92ea957e8c53981a2bf6136572531025bf" dependencies = [ "hex", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "pin-project-lite", "tokio", @@ -2889,7 +2889,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ "http 1.4.2", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "rustls 0.23.42", "rustls-native-certs", @@ -2904,7 +2904,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" dependencies = [ - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "pin-project-lite", "tokio", @@ -2923,7 +2923,7 @@ dependencies = [ "futures-util", "http 1.4.2", "http-body 1.1.0", - "hyper 1.10.1", + "hyper 1.12.0", "ipnet", "libc", "percent-encoding", @@ -2942,7 +2942,7 @@ checksum = "986c5ce3b994526b3cd75578e62554abd09f0899d6206de48b3e96ab34ccc8c7" dependencies = [ "hex", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "pin-project-lite", "tokio", @@ -3769,6 +3769,7 @@ dependencies = [ "litellm-config", "litellm-error", "litellm-gateway-auth", + "litellm-gateway-forward", "litellm-gateway-inference", "litellm-gateway-mcp", "litellm-gateway-ui", @@ -3815,6 +3816,29 @@ dependencies = [ "veil", ] +[[package]] +name = "litellm-gateway-forward" +version = "0.1.0" +dependencies = [ + "axum", + "bytes", + "envy", + "futures-util", + "http-body-util", + "httparse", + "hyper 1.12.0", + "hyper-util", + "proptest", + "rstest", + "serde", + "thiserror 2.0.19", + "tokio", + "tokio-tungstenite", + "tokio-util", + "tower-service", + "tracing", +] + [[package]] name = "litellm-gateway-inference" version = "0.1.0" @@ -6031,7 +6055,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-rustls 0.27.9", "hyper-util", "js-sys", @@ -7512,6 +7536,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -7591,7 +7616,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-timeout", "hyper-util", "percent-encoding", @@ -8167,9 +8192,9 @@ dependencies = [ [[package]] name = "want" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +checksum = "ec4cdd0dd910afe868b7ef477227d8d538b46b3075031afee8a9f2acb0a2ed0b" dependencies = [ "try-lock", ] @@ -8551,7 +8576,7 @@ dependencies = [ "futures", "http 1.4.2", "http-body-util", - "hyper 1.10.1", + "hyper 1.12.0", "hyper-util", "log", "once_cell", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 71914a0bffe..7121df06bc0 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -26,6 +26,7 @@ litellm-inference-ocr = { path = "crates/inference-ocr" } litellm-inference-testing = { path = "crates/inference-testing" } litellm-gateway-mcp = { path = "crates/gateway-mcp" } litellm-gateway = { path = "crates/gateway" } +litellm-gateway-forward = { path = "crates/gateway-forward" } litellm-gateway-inference = { path = "crates/gateway-inference" } litellm-gateway-auth = { path = "crates/gateway-auth" } litellm-gateway-management = { path = "crates/gateway-management" } diff --git a/litellm-rust/crates/gateway-forward/AGENTS.md b/litellm-rust/crates/gateway-forward/AGENTS.md new file mode 100644 index 00000000000..2135980266f --- /dev/null +++ b/litellm-rust/crates/gateway-forward/AGENTS.md @@ -0,0 +1,37 @@ +# Forwarding gateway + +`Forward` is a reverse proxy in front of the Python proxy on loopback. Python must see each request as if the client had connected to it directly, because every Python route (inference, MCP, pass-through, UI, WebSocket) is served behind it without Rust knowing which one it is + +- Serve HTTP/1.1 only, as uvicorn does. The supervisor never puts Rust in front of a Python server that speaks HTTP/2 +- Never branch on a route or path. A behavior one endpoint needs is a protocol behavior every endpoint gets +- Forward the method, request target, headers, and body byte for byte within the [URL contract](FORWARDING.md#url-fragments-and-literal-hashes). An unescaped `#` starts a fragment and is omitted with everything after it, while `%23` is preserved. Only drop hop-by-hop headers (RFC 9110 section 7.6.1) and set `x-litellm-rust-peer`, which replaces any client copy +- Stream both bodies. Never buffer, cap, or time out a body, since uploads and provider streams have no upper bound +- A client disconnect closes the upstream connection at every stage, because that is how Python learns to stop a provider stream +- An upstream that dies mid-response fails the client body. Never end it cleanly, or SSE clients treat a truncated stream as complete +- Never invent a Host. Python builds absolute URLs from it, so a missing one must never become the private upstream address +- A request Python could not have received unchanged (HTTP/1.1 without Host, a transfer coding other than chunked, a target without a path) gets the status uvicorn would have given, from Rust +- Pooled upstream connections are retired well before Python's keep-alive timeout, so a request never lands on a connection Python is closing +- Idle client connections close after Python's keep-alive timeout, as uvicorn closed them, so a client that vanishes cannot hold a socket forever. Only the wait for the next request is timed, never an exchange in progress +- WebSocket is the one protocol switch Python serves. Upgrades are relayed as raw bytes after the 101 and drained on shutdown + +## Layout + +Every forwarding decision has one home. `rewrite.rs` owns what Python and the client receive, as pure functions over requests and responses. `Rejection` in `error.rs` is the only place that picks a status for a request Rust answers itself. `settings.rs` parses the environment into typed values, so an `Upstream` is a loopback origin by construction. `server.rs` accepts client connections and owns their timeouts, since `axum::serve` exposes none, and `connections.rs` tracks every connection and relay for shutdown. `forward.rs` and `websocket.rs` only move requests, bodies and upgraded connections and never touch a header themselves. `connector.rs` keeps hyper reading Python's answer after Python stops reading the request + +## Environment + +`Settings` reads only these variables, and only when `LITELLM_FORWARD_UPSTREAM` is set. The Python supervisor in `litellm/proxy/rust_forward/` sets every one of them from the values Python actually runs with, so all of them are required and a missing or invalid one fails startup naming it. The crate never reads Python's own settings (`HOST`, `PORT`, `KEEPALIVE_TIMEOUT`, `GRACEFUL_SHUTDOWN_TIMEOUT`) and keeps no copy of Python's defaults, since Python can get those values from CLI flags the environment never shows + +| Variable | Meaning | +|---|---| +| `LITELLM_FORWARD_UPSTREAM` | Python's loopback origin, e.g. `http://127.0.0.1:53211` | +| `LITELLM_FORWARD_HOST` | public bind host | +| `LITELLM_FORWARD_PORT` | public bind port | +| `LITELLM_FORWARD_UPSTREAM_KEEPALIVE` | seconds uvicorn keeps an idle connection open. Pooled connections are retired at half of it, and idle client connections are closed after it | +| `LITELLM_FORWARD_SHUTDOWN_DRAIN` | seconds to drain on shutdown, already including Python's graceful timeout and teardown headroom | + +## Tests + +Test through real sockets with `tests/support`, never `oneshot`. `Gateway` runs the forwarder on a listener like production, `recording` is a well-behaved upstream that captures what Python would receive, and `raw_upstream` and `Raw` drive bytes by hand for framing, timing and disconnects a real server would not produce. Waits measured in minutes use paused tokio time instead of real sleeps + +The forwarder sees HTTP/1.1 messages, not routes, so the tests are split into one file per traffic shape uvicorn can emit: `exchange` (plain request and response), `framing` (requests whose framing a parser could read two ways), `slow_head` (minutes before response headers), `stream` (incremental response bodies), `idle_stream` (long-lived streams that go quiet), `upload`, `download`, `websocket`, `client_identity`, `lifecycle` (disconnects, pooling, idle timeouts, shutdown) and `roundtrip` (the forwarding invariant for generated messages). A new route never gets its own test. Add its behavior to the shape it already has, and when it needs a wire behavior none of the files covers, add a new shape file for it diff --git a/litellm-rust/crates/gateway-forward/Cargo.toml b/litellm-rust/crates/gateway-forward/Cargo.toml new file mode 100644 index 00000000000..2fef19dca34 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "litellm-gateway-forward" +version = "0.1.0" +edition.workspace = true +rust-version.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +axum.workspace = true +envy = "0.4.2" +hyper = { version = "1.12.0", features = ["http1", "server"] } +hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] } +serde.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["io-util"] } +tokio-util = { version = "0.7", features = ["rt"] } +tower-service = "0.3.3" +tracing.workspace = true + +[dev-dependencies] +bytes.workspace = true +futures-util.workspace = true +http-body-util = "0.1" +httparse = "1" +hyper = { version = "1.12.0", features = ["client"] } +proptest.workspace = true +rstest.workspace = true +tokio = { workspace = true, features = ["io-util", "sync", "test-util"] } +tokio-tungstenite.workspace = true diff --git a/litellm-rust/crates/gateway-forward/FORWARDING.md b/litellm-rust/crates/gateway-forward/FORWARDING.md new file mode 100644 index 00000000000..532f341aba5 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/FORWARDING.md @@ -0,0 +1,31 @@ +# Python forwarding compatibility + +The opt-in forwarding gateway serves HTTP/1.1 and WebSockets in front of Python. Python still owns every route and policy. Its protocol contract and environment settings are in [AGENTS.md](AGENTS.md) + +## URL fragments and literal hashes + +The gateway follows Hyper's `http::Uri` parsing. If a raw client sends an unescaped `#` in the request target, only the path and query before the first `#` reach Python. For example, `/v1/files?purpose=batch#retry&limit=1` becomes `/v1/files?purpose=batch`, and `/gemini/v1beta/files/a#b?key=k` becomes `/gemini/v1beta/files/a` + +This is accepted behavior because `#` starts a URL fragment, which is separated before dereferencing ([RFC 3986 section 3.5](https://www.rfc-editor.org/rfc/rfc3986.html#section-3.5)). HTTP origin-form contains only the path and optional query, not a fragment ([RFC 9112 section 3.2.1](https://www.rfc-editor.org/rfc/rfc9112.html#section-3.2.1)). Browsers and URL-based clients omit the fragment before sending the request, so UI hash routing is unaffected + +A literal hash in path or query data must be encoded as `%23`, which the gateway forwards unchanged. Direct Python with h11 also accepts unescaped hashes as data, but that permissive raw-target behavior is outside this URL contract. Active socket tests cover both [fragment omission](tests/exchange.rs) and [encoded-hash preservation](tests/exchange.rs) + +## Known differences from uvicorn with h11 + +These cases remain executable as ignored socket tests. They are not passing parity checks + +| Input | Direct Python | Through Rust | Regression test | +|---|---|---|---| +| Header after a removed hop-by-hop field | Keeps the remaining header order | `HeaderMap::remove` can move a later header | [dropping_a_hop_by_hop_header_keeps_the_others_in_order](tests/exchange.rs) | +| Repeated header names interleaved with other names | Keeps the interleaving | Groups values by name, keeping their order within that name | [interleaved_repeated_headers_reach_python_interleaved](tests/exchange.rs) | +| Obsolete folded header, such as `x-a: b\r\n c` | Unfolds to `b c` | Rejects with 400 | [a_folded_header_reaches_python_unfolded](tests/exchange.rs) | +| Identical Content-Length list, such as `7, 7` | Accepts length 7 | Rejects with 400 | [a_content_length_list_of_one_value_reaches_python](tests/framing.rs) | +| HTTP/1.0 request with chunked transfer coding | Decodes the body | Rejects with 400 | [an_http10_chunked_upload_reaches_python](tests/framing.rs) | + +Clients can avoid the parser rejections by sending unfolded headers, one Content-Length value, and HTTP/1.1 for chunked uploads. Integrations must not depend on header ordering across different field names + +Run `cargo test -p litellm-gateway-forward --test exchange --test framing -- --ignored` from this directory to reproduce the differences + +## Connection header handling + +The gateway strips standard hop-by-hop fields and fields named by Connection, except routing, representation, credential, forwarding, and WebSocket handshake headers. A client cannot nominate those end-to-end fields for removal. The private `x-litellm-rust-peer` always replaces any client-supplied copy diff --git a/litellm-rust/crates/gateway-forward/src/connections.rs b/litellm-rust/crates/gateway-forward/src/connections.rs new file mode 100644 index 00000000000..10a4867fa51 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/connections.rs @@ -0,0 +1,58 @@ +use tokio_util::{ + sync::{CancellationToken, DropGuard}, + task::TaskTracker, +}; + +/// Every client connection and WebSocket relay the gateway runs, so shutdown can wait for them +/// and then close whatever outlives the drain. +#[derive(Clone, Default)] +pub(crate) struct Connections { + tasks: TaskTracker, + draining: CancellationToken, + closed: CancellationToken, +} + +impl Connections { + /// Runs `work` until it finishes or the gateway closes every connection. + pub fn spawn(&self, work: impl Future + Send + 'static) { + let closed = self.closed.clone(); + self.tasks.spawn(async move { + tokio::select! { + () = closed.cancelled() => {} + () = work => {} + } + }); + } + + pub fn begin_shutdown(&self) { + self.draining.cancel(); + } + + pub async fn draining(&self) { + self.draining.cancelled().await; + } + + pub fn is_draining(&self) -> bool { + self.draining.is_cancelled() + } + + pub async fn until_closed(&self, future: impl Future) -> Option { + tokio::select! { + () = self.closed.cancelled() => None, + output = future => Some(output), + } + } + + pub fn close(&self) { + self.closed.cancel(); + } + + pub fn close_on_drop(&self) -> DropGuard { + self.closed.clone().drop_guard() + } + + pub async fn wait(&self) { + self.tasks.close(); + self.tasks.wait().await; + } +} diff --git a/litellm-rust/crates/gateway-forward/src/connector.rs b/litellm-rust/crates/gateway-forward/src/connector.rs new file mode 100644 index 00000000000..5eb998db0c0 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/connector.rs @@ -0,0 +1,122 @@ +use std::{ + io::{self, IoSlice}, + pin::Pin, + task::{Context, Poll, ready}, +}; + +use axum::http::Uri; +use hyper::rt::{Read, ReadBufCursor, Write}; +use hyper_util::{ + client::legacy::connect::{Connected, Connection, HttpConnector}, + rt::TokioIo, +}; +use tokio::net::TcpStream; + +/// Connects to Python like `HttpConnector`, but keeps hyper reading Python's answer after Python +/// stops reading the request. uvicorn rejects an upload (413, 401) as soon as it reads the head +/// and closes with the body unread, so writing the rest fails, and hyper would fail the request +/// and drop the answer it already received. Like curl and httpcore, the forwarder discards the +/// rest of the request instead and reads on. A connection that never answered still fails, on the +/// read, and no request is ever sent twice. +#[derive(Clone)] +pub(crate) struct Connector(pub HttpConnector); + +impl tower_service::Service for Connector { + type Response = AnswerAfterRefusal; + type Error = >::Error; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll> { + self.0.poll_ready(context) + } + + fn call(&mut self, uri: Uri) -> Self::Future { + let connecting = self.0.call(uri); + Box::pin(async move { + Ok(AnswerAfterRefusal { + stream: connecting.await?, + refused: false, + }) + }) + } +} + +pub(crate) struct AnswerAfterRefusal { + stream: TokioIo, + refused: bool, +} + +impl AnswerAfterRefusal { + fn written(&mut self, written: io::Result, length: usize) -> io::Result { + match written { + Err(error) + if matches!( + error.kind(), + io::ErrorKind::BrokenPipe | io::ErrorKind::ConnectionReset + ) => + { + self.refused = true; + Ok(length) + } + written => written, + } + } +} + +impl Read for AnswerAfterRefusal { + fn poll_read( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: ReadBufCursor<'_>, + ) -> Poll> { + Pin::new(&mut self.stream).poll_read(context, buffer) + } +} + +impl Write for AnswerAfterRefusal { + fn poll_write( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + if self.refused { + return Poll::Ready(Ok(buffer.len())); + } + let written = ready!(Pin::new(&mut self.stream).poll_write(context, buffer)); + Poll::Ready(self.written(written, buffer.len())) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffers: &[IoSlice<'_>], + ) -> Poll> { + let length = buffers.iter().map(|buffer| buffer.len()).sum(); + if self.refused { + return Poll::Ready(Ok(length)); + } + let written = ready!(Pin::new(&mut self.stream).poll_write_vectored(context, buffers)); + Poll::Ready(self.written(written, length)) + } + + fn is_write_vectored(&self) -> bool { + self.stream.is_write_vectored() + } + + fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + if self.refused { + return Poll::Ready(Ok(())); + } + Pin::new(&mut self.stream).poll_flush(context) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_shutdown(context) + } +} + +impl Connection for AnswerAfterRefusal { + fn connected(&self) -> Connected { + self.stream.connected() + } +} diff --git a/litellm-rust/crates/gateway-forward/src/error.rs b/litellm-rust/crates/gateway-forward/src/error.rs new file mode 100644 index 00000000000..7dbdde68a85 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/error.rs @@ -0,0 +1,91 @@ +use axum::{ + http::StatusCode, + response::{IntoResponse, Response}, +}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid forwarding upstream URI")] + UpstreamUri(#[source] axum::http::uri::InvalidUri), + #[error( + "LITELLM_FORWARD_UPSTREAM must be a loopback IPv4 http origin without a path, query or credentials" + )] + Upstream, + #[error("{0} must be set in forward mode")] + MissingSetting(String), + #[error("invalid forwarding setting: {0}")] + InvalidSetting(String), + #[error("forwarding gateway I/O failed: {0}")] + Io(#[from] std::io::Error), +} + +/// Why a request got an answer from Rust instead of from Python. +#[derive(Debug, thiserror::Error)] +pub(crate) enum Rejection { + #[error("HTTP/1.1 request without a Host header")] + MissingHost, + #[error("request target has no path to forward")] + NoPath, + #[error("request uses a transfer coding other than chunked")] + UnsupportedTransferCoding, + #[error("client request body failed: {0}")] + ClientBody(#[source] hyper_util::client::legacy::Error), + #[error("python proxy unreachable: {0}")] + Unreachable(#[source] hyper_util::client::legacy::Error), + #[error("python proxy switched protocols without being asked to")] + UnrequestedUpgrade, + #[error("python proxy websocket upgrade failed: {0}")] + UpgradeFailed(#[source] hyper::Error), + #[error("forwarding gateway is shutting down")] + ShuttingDown, + #[error("could not build the request for the python proxy")] + Internal, +} + +impl Rejection { + /// A client body that ends early or breaks its framing fails the upstream request too, but + /// the fault is the client's, not Python's. + pub fn from_upstream(error: hyper_util::client::legacy::Error) -> Self { + let client_body = std::error::Error::source(&error) + .and_then(|source| source.downcast_ref::()) + .is_some_and(hyper::Error::is_user); + match client_body { + true => Self::ClientBody(error), + false => Self::Unreachable(error), + } + } + + fn status(&self) -> StatusCode { + match self { + Self::MissingHost + | Self::NoPath + | Self::UnsupportedTransferCoding + | Self::ClientBody(_) => StatusCode::BAD_REQUEST, + Self::Unreachable(_) | Self::UnrequestedUpgrade | Self::UpgradeFailed(_) => { + StatusCode::BAD_GATEWAY + } + Self::ShuttingDown => StatusCode::SERVICE_UNAVAILABLE, + Self::Internal => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn is_gateway_fault(&self) -> bool { + matches!( + self, + Self::Unreachable(_) + | Self::UnrequestedUpgrade + | Self::UpgradeFailed(_) + | Self::Internal + ) + } +} + +impl IntoResponse for Rejection { + fn into_response(self) -> Response { + match self.is_gateway_fault() { + true => tracing::warn!(error = %self, "request not forwarded"), + false => tracing::debug!(error = %self, "request not forwarded"), + } + self.status().into_response() + } +} diff --git a/litellm-rust/crates/gateway-forward/src/forward.rs b/litellm-rust/crates/gateway-forward/src/forward.rs new file mode 100644 index 00000000000..69fb63ada7d --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/forward.rs @@ -0,0 +1,169 @@ +use std::{ + pin::Pin, + task::{Context, Poll}, + time::Duration, +}; + +use axum::{ + body::Body, + extract::Request, + http::{StatusCode, Uri}, + response::{IntoResponse, Response}, +}; +use hyper::body::{Frame, Incoming, SizeHint}; +use hyper_util::{ + client::legacy::{Client, connect::HttpConnector}, + rt::{TokioExecutor, TokioTimer}, +}; +use tokio::net::TcpListener; + +use crate::{ + Rejection, Upstream, connections::Connections, connector::Connector, rewrite, server, + server::ClientSocket, websocket, +}; + +// h11 counts no headers and only caps a head at 16 KiB, so it accepts as many headers as fit in +// that many shortest lines (`a:\r\n`). Python copies every provider header into its response +// too, so hyper's default of 100 would turn ordinary traffic into 431s and 502s. +pub(crate) const MAX_HEADERS: usize = 16 * 1024 / b"a:\r\n".len(); + +#[derive(Clone)] +pub struct Forward { + upstream: Uri, + client: Client, + connections: Connections, + keepalive: Duration, +} + +impl Forward { + /// `upstream_keepalive` is how long the upstream keeps an idle connection open. Pooled + /// connections are dropped well before that, so a request never lands on one it is closing, + /// and idle client connections are closed after it, as uvicorn did. + pub fn new(upstream: &Upstream, upstream_keepalive: Duration) -> Self { + let mut connector = HttpConnector::new(); + connector.set_nodelay(true); + let client = Client::builder(TokioExecutor::new()) + .http1_max_headers(MAX_HEADERS) + .pool_timer(TokioTimer::new()) + .pool_idle_timeout(upstream_keepalive / 2) + .set_host(false) + .build(Connector(connector)); + Self { + upstream: upstream.uri().clone(), + client, + connections: Connections::default(), + keepalive: upstream_keepalive, + } + } + + /// Serves until `shutdown` resolves, then lets open requests and WebSockets finish for up to + /// `drain` before closing whatever is left. Python shuts down alongside and ends long streams + /// itself, so `drain` only backstops a Python that never does. + pub async fn serve( + self, + listener: TcpListener, + shutdown: impl Future, + drain: Duration, + ) { + server::serve(self, listener, shutdown, drain).await; + } + + pub(crate) fn connections(&self) -> &Connections { + &self.connections + } + + pub(crate) fn keepalive(&self) -> Duration { + self.keepalive + } + + pub(crate) async fn handle(&self, client: ClientSocket, request: Request) -> Response { + self.forward(client, request) + .await + .unwrap_or_else(IntoResponse::into_response) + } + + async fn forward( + &self, + client: ClientSocket, + mut request: Request, + ) -> Result { + if let Some(downstream) = websocket::take_handshake(&self.connections, &mut request)? { + return self.forward_websocket(client, request, downstream).await; + } + let request = rewrite::upstream_request(request, &self.upstream, client, false)?; + let response = self + .client + .request(request) + .await + .map_err(Rejection::from_upstream)?; + http_response(response) + } + + async fn forward_websocket( + &self, + client: ClientSocket, + request: Request, + downstream: hyper::upgrade::OnUpgrade, + ) -> Result { + let request = rewrite::upstream_request(request, &self.upstream, client, true)?; + let response = self + .connections + .until_closed(self.client.request(request)) + .await + .ok_or(Rejection::ShuttingDown)? + .map_err(Rejection::from_upstream)?; + if response.status() == StatusCode::SWITCHING_PROTOCOLS { + return websocket::switch(&self.connections, downstream, response).await; + } + http_response(response) + } +} + +fn http_response(response: Response) -> Result { + if response.status() == StatusCode::SWITCHING_PROTOCOLS { + return Err(Rejection::UnrequestedUpgrade); + } + Ok( + rewrite::downstream_response(response) + .map(|body| Body::new(Relayed { body, failed: None })), + ) +} + +/// Python's response body on its way to the client. +struct Relayed { + body: Incoming, + failed: Option, +} + +impl hyper::body::Body for Relayed { + type Data = ::Data; + type Error = hyper::Error; + + // hyper aborts a client connection whose response body fails without flushing the head it + // buffered, so a head Python sent right before dying would never arrive and SDKs would retry + // a request Python already ran. Holding the error for one poll lets hyper flush the head. + fn poll_frame( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + if let Some(error) = self.failed.take() { + return Poll::Ready(Some(Err(error))); + } + match Pin::new(&mut self.body).poll_frame(context) { + Poll::Ready(Some(Err(error))) => { + self.failed = Some(error); + context.waker().wake_by_ref(); + Poll::Pending + } + frame => frame, + } + } + + fn is_end_stream(&self) -> bool { + self.failed.is_none() && self.body.is_end_stream() + } + + fn size_hint(&self) -> SizeHint { + self.body.size_hint() + } +} diff --git a/litellm-rust/crates/gateway-forward/src/lib.rs b/litellm-rust/crates/gateway-forward/src/lib.rs new file mode 100644 index 00000000000..4ac17b50544 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/lib.rs @@ -0,0 +1,29 @@ +mod connections; +mod connector; +mod error; +mod forward; +mod rewrite; +mod server; +mod settings; +mod websocket; + +pub use error::Error; +pub(crate) use error::Rejection; +pub use forward::Forward; +pub use rewrite::PEER_HEADER; +pub use settings::{Settings, Upstream}; + +/// Binds the public listener and forwards to Python until `shutdown` resolves, then drains for +/// `settings.shutdown_drain`. +pub async fn run( + settings: Settings, + shutdown: impl Future + Send + 'static, +) -> Result<(), Error> { + let forward = Forward::new(&settings.upstream, settings.upstream_keepalive); + let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?; + tracing::info!(address = %listener.local_addr()?, "forwarding gateway listening"); + forward + .serve(listener, shutdown, settings.shutdown_drain) + .await; + Ok(()) +} diff --git a/litellm-rust/crates/gateway-forward/src/rewrite.rs b/litellm-rust/crates/gateway-forward/src/rewrite.rs new file mode 100644 index 00000000000..fa22bf19037 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/rewrite.rs @@ -0,0 +1,161 @@ +use axum::{ + body::HttpBody, + extract::Request, + http::{HeaderMap, HeaderValue, Response, Uri, Version, header}, +}; + +use crate::{Rejection, server::ClientSocket}; + +pub const PEER_HEADER: &str = "x-litellm-rust-peer"; + +/// The request Python receives: the client's method, target, headers and body unchanged, sent to +/// `upstream` over HTTP/1.1 without the client connection's own headers, plus the client's TCP +/// peer in `PEER_HEADER`, which replaces any copy the client sent. +pub(crate) fn upstream_request( + mut request: Request, + upstream: &Uri, + client: ClientSocket, + websocket: bool, +) -> Result { + if !only_chunked(request.headers()) { + return Err(Rejection::UnsupportedTransferCoding); + } + let path_and_query = request + .uri() + .path_and_query() + .cloned() + .ok_or(Rejection::NoPath)?; + if let Some(host) = public_host(&request, client)? { + request.headers_mut().insert(header::HOST, host); + } + let mut uri = upstream.clone().into_parts(); + uri.path_and_query = Some(path_and_query); + *request.uri_mut() = Uri::from_parts(uri).map_err(|_| Rejection::Internal)?; + *request.version_mut() = Version::HTTP_11; + let chunked = request.body().size_hint().exact().is_none(); + let headers = request.headers_mut(); + strip_hop_by_hop(headers); + if chunked { + // hyper sends no body for a GET with an unknown length unless told it is chunked. + headers.insert( + header::TRANSFER_ENCODING, + HeaderValue::from_static("chunked"), + ); + } + if websocket { + restore_websocket_handshake(headers); + } + headers.insert( + PEER_HEADER, + HeaderValue::from_str(&client.peer.to_string()).map_err(|_| Rejection::Internal)?, + ); + Ok(request) +} + +// Python builds absolute URLs from Host and, without a usable one, from the address it listens +// on, which behind Rust is the private loopback port. A client that sent no usable Host gets the +// public address it connected to, which is what Python used before. h11 rejects HTTP/1.1 without +// Host. +fn public_host(request: &Request, client: ClientSocket) -> Result, Rejection> { + let host = request.headers().get(header::HOST); + if host.is_some_and(|host| !host.is_empty()) { + return Ok(None); + } + if host.is_none() && request.version() == Version::HTTP_11 { + return Err(Rejection::MissingHost); + } + HeaderValue::from_str(&client.local.to_string()) + .map(Some) + .map_err(|_| Rejection::Internal) +} + +// hyper only decodes chunked framing, so anything else would reach Python still coded but framed +// as a plain body. uvicorn rejects it with the 400 it sends for every h11 parse error. +fn only_chunked(headers: &HeaderMap) -> bool { + headers + .get_all(header::TRANSFER_ENCODING) + .iter() + .flat_map(|value| value.as_bytes().split(|byte| *byte == b',')) + .all(|coding| coding.trim_ascii().eq_ignore_ascii_case(b"chunked")) +} + +/// The response the client receives: Python's, without the upstream connection's own headers. +pub(crate) fn downstream_response(mut response: Response) -> Response { + strip_hop_by_hop(response.headers_mut()); + response +} + +/// Python's 101 as the client receives it, still naming the protocol both sides switch to. +pub(crate) fn switched_response(mut response: Response) -> Response { + strip_hop_by_hop(response.headers_mut()); + restore_websocket_handshake(response.headers_mut()); + response +} + +pub(crate) fn is_websocket_handshake(headers: &HeaderMap) -> bool { + headers + .get(header::UPGRADE) + .is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket")) + && headers + .get_all(header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .any(|token| token.trim().eq_ignore_ascii_case("upgrade")) +} + +const HOP_BY_HOP: [&str; 9] = [ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + "proxy-connection", +]; + +// Per-connection headers (RFC 9110 section 7.6.1) describe the client's connection to Rust, not +// Rust's connection to Python, so a proxy must drop them instead of forwarding them. +fn strip_hop_by_hop(headers: &mut HeaderMap) { + let named: Vec<_> = headers + .get_all(header::CONNECTION) + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .filter_map(|name| header::HeaderName::from_bytes(name.trim().as_bytes()).ok()) + .filter(|name| !must_forward(name)) + .collect(); + for name in named { + headers.remove(name); + } + for name in HOP_BY_HOP { + headers.remove(name); + } +} + +fn must_forward(name: &header::HeaderName) -> bool { + matches!( + name.as_str(), + "host" + | "authorization" + | "cookie" + | "set-cookie" + | "www-authenticate" + | "authentication-info" + | "api-key" + | "x-api-key" + | "litellm-api-key" + | "forwarded" + | PEER_HEADER + ) || name.as_str().starts_with("content-") + || name.as_str().starts_with("x-forwarded-") + || name.as_str().starts_with("sec-websocket-") +} + +// Stripping drops the handshake's Connection and Upgrade, so each hop gets them back. +fn restore_websocket_handshake(headers: &mut HeaderMap) { + headers.insert(header::CONNECTION, HeaderValue::from_static("upgrade")); + headers.insert(header::UPGRADE, HeaderValue::from_static("websocket")); +} diff --git a/litellm-rust/crates/gateway-forward/src/server.rs b/litellm-rust/crates/gateway-forward/src/server.rs new file mode 100644 index 00000000000..0388642446a --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/server.rs @@ -0,0 +1,116 @@ +use std::{convert::Infallible, io, net::SocketAddr, pin::pin, time::Duration}; + +use axum::{body::Body, extract::Request}; +use hyper::{body::Incoming, server::conn::http1, service::service_fn}; +use hyper_util::rt::{TokioIo, TokioTimer}; +use tokio::net::TcpListener; + +use crate::{Forward, forward::MAX_HEADERS}; + +/// The client's TCP connection to the public listener. +#[derive(Clone, Copy, Debug)] +pub(crate) struct ClientSocket { + pub peer: SocketAddr, + pub local: SocketAddr, +} + +/// Serves each public connection with hyper directly, because `axum::serve` cannot configure +/// the connection timeouts uvicorn applied to the same clients. +pub(crate) async fn serve( + forward: Forward, + listener: TcpListener, + shutdown: impl Future, + drain: Duration, +) { + let connections = forward.connections().clone(); + let _close_on_drop = connections.close_on_drop(); + let keepalive = forward.keepalive(); + let builder = builder(keepalive); + let mut shutdown = pin!(shutdown); + loop { + let (stream, peer) = tokio::select! { + () = &mut shutdown => break, + accepted = listener.accept() => match accepted { + Ok(accepted) => accepted, + Err(error) => { + back_off(error).await; + continue; + } + }, + }; + let local = match stream.local_addr() { + Ok(local) => local, + Err(error) => { + tracing::debug!(%error, "client connection closed before it was served"); + continue; + } + }; + if let Err(error) = stream.set_nodelay(true) { + tracing::warn!(%error, "could not disable Nagle on a client connection"); + } + let socket = ClientSocket { peer, local }; + let forward = forward.clone(); + let service = service_fn(move |request: Request| { + let forward = forward.clone(); + async move { Ok::<_, Infallible>(forward.handle(socket, request.map(Body::new)).await) } + }); + let connection = builder + .serve_connection(TokioIo::new(stream), service) + .with_upgrades(); + let draining = connections.clone(); + connections.spawn(async move { + let mut connection = pin!(connection); + let ended = tokio::select! { + ended = connection.as_mut() => ended, + () = draining.draining() => { + connection.as_mut().graceful_shutdown(); + connection.await + } + }; + if let Err(error) = ended { + tracing::debug!(%error, "client connection ended with an error"); + } + }); + } + drop(listener); + connections.begin_shutdown(); + if tokio::time::timeout(drain, connections.wait()) + .await + .is_err() + { + tracing::warn!( + ?drain, + "closing connections still open after the shutdown drain" + ); + } + connections.close(); + connections.wait().await; +} + +// uvicorn closes a client connection that sits idle for `--keepalive_timeout` between requests. +// Without the same limit, a client that vanishes without closing (NAT expiry, a crashed host) +// would hold its socket for the life of the process until the gateway runs out of descriptors. +// Operators already tune that timeout against their load balancer, so the public side keeps +// Python's value instead of introducing its own. hyper only runs the header timer while a +// connection waits for its next request, so it never cuts a long stream, upload or WebSocket. +fn builder(keepalive: Duration) -> http1::Builder { + let mut builder = http1::Builder::new(); + builder + .max_headers(MAX_HEADERS) + .timer(TokioTimer::new()) + .header_read_timeout(keepalive); + builder +} + +async fn back_off(error: io::Error) { + if matches!( + error.kind(), + io::ErrorKind::ConnectionRefused + | io::ErrorKind::ConnectionAborted + | io::ErrorKind::ConnectionReset + ) { + return; + } + tracing::error!(%error, "could not accept a client connection"); + tokio::time::sleep(Duration::from_secs(1)).await; +} diff --git a/litellm-rust/crates/gateway-forward/src/settings.rs b/litellm-rust/crates/gateway-forward/src/settings.rs new file mode 100644 index 00000000000..de01a96c519 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/settings.rs @@ -0,0 +1,113 @@ +use std::{ffi::OsString, net::SocketAddr, str::FromStr, time::Duration}; + +use axum::http::Uri; +use serde::Deserialize; + +use crate::Error; + +const UPSTREAM: &str = "LITELLM_FORWARD_UPSTREAM"; + +/// The Python proxy's private origin. Only a loopback IPv4 address parses, so nothing can point +/// the forwarder at another host. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Upstream(Uri); + +impl Upstream { + pub(crate) fn uri(&self) -> &Uri { + &self.0 + } +} + +impl FromStr for Upstream { + type Err = Error; + + fn from_str(value: &str) -> Result { + let uri: Uri = value.parse().map_err(Error::UpstreamUri)?; + let address = uri + .authority() + .filter(|_| uri.scheme_str() == Some("http")) + .filter(|authority| !authority.as_str().contains('@')) + .filter(|_| uri.path() == "/" && uri.query().is_none()) + .and_then(|authority| authority.as_str().parse::().ok()); + match address { + Some(SocketAddr::V4(address)) if address.ip().is_loopback() => Ok(Self(uri)), + _ => Err(Error::Upstream), + } + } +} + +#[derive(Deserialize)] +struct Environment { + litellm_forward_upstream: String, + litellm_forward_host: String, + litellm_forward_port: u16, + litellm_forward_upstream_keepalive: f64, + litellm_forward_shutdown_drain: f64, +} + +/// Everything forward mode runs with. The Python supervisor resolves each value from the settings +/// Python actually runs with, so none of them is guessed or defaulted here. +/// Setting `LITELLM_FORWARD_UPSTREAM` requires every other `LITELLM_FORWARD_*` setting below. +#[derive(Debug, Clone, PartialEq)] +pub struct Settings { + pub upstream: Upstream, + pub host: String, + pub port: u16, + pub upstream_keepalive: Duration, + pub shutdown_drain: Duration, +} + +impl Settings { + /// `None` when `LITELLM_FORWARD_UPSTREAM` is unset, which means the gateway is not forwarding. + pub fn from_env() -> Result, Error> { + Self::from_vars(std::env::vars_os()) + } + + pub fn from_vars( + variables: impl IntoIterator, + ) -> Result, Error> { + let variables: Vec<(String, String)> = variables + .into_iter() + .filter_map(|(key, value)| Some((key.into_string().ok()?, value.into_string().ok()?))) + .collect(); + if !variables.iter().any(|(key, _)| key == UPSTREAM) { + return Ok(None); + } + let environment: Environment = envy::from_iter(variables).map_err(|error| match error { + envy::Error::MissingValue(field) => Error::MissingSetting(field.to_uppercase()), + envy::Error::Custom(message) => Error::InvalidSetting(message), + })?; + Ok(Some(Self { + upstream: environment.litellm_forward_upstream.parse()?, + host: environment.litellm_forward_host, + port: environment.litellm_forward_port, + upstream_keepalive: positive_seconds( + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE", + environment.litellm_forward_upstream_keepalive, + )?, + shutdown_drain: seconds( + "LITELLM_FORWARD_SHUTDOWN_DRAIN", + environment.litellm_forward_shutdown_drain, + )?, + })) + } +} + +fn seconds(variable: &str, value: f64) -> Result { + Duration::try_from_secs_f64(value).map_err(|_| { + Error::InvalidSetting(format!( + "{variable} must be a non-negative number of seconds, got {value}" + )) + }) +} + +/// The keep-alive also bounds the wait for a client's first request, so zero would close every +/// client connection before it could send one. +fn positive_seconds(variable: &str, value: f64) -> Result { + match seconds(variable, value)? { + Duration::ZERO => Err(Error::InvalidSetting(format!( + "{variable} must be a positive number of seconds, got {value}" + ))), + duration => Ok(duration), + } +} diff --git a/litellm-rust/crates/gateway-forward/src/websocket.rs b/litellm-rust/crates/gateway-forward/src/websocket.rs new file mode 100644 index 00000000000..84d89c75607 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/websocket.rs @@ -0,0 +1,86 @@ +use axum::{ + body::Body, + extract::Request, + http::{Method, Version}, + response::Response, +}; +use hyper::{ + body::Incoming, + upgrade::{OnUpgrade, Upgraded}, +}; +use hyper_util::rt::TokioIo; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +use crate::{Rejection, connections::Connections, rewrite}; + +pub(crate) fn take_handshake( + connections: &Connections, + request: &mut Request, +) -> Result, Rejection> { + if request.method() != Method::GET + || request.version() != Version::HTTP_11 + || !rewrite::is_websocket_handshake(request.headers()) + { + return Ok(None); + } + if connections.is_draining() { + return Err(Rejection::ShuttingDown); + } + Ok(Some(hyper::upgrade::on(request))) +} + +pub(crate) async fn switch( + connections: &Connections, + downstream: OnUpgrade, + mut response: Response, +) -> Result { + if !rewrite::is_websocket_handshake(response.headers()) { + return Err(Rejection::UnrequestedUpgrade); + } + let upstream = connections + .until_closed(hyper::upgrade::on(&mut response)) + .await + .ok_or(Rejection::ShuttingDown)? + .map_err(Rejection::UpgradeFailed)?; + connections.spawn(async move { + match downstream.await { + Ok(downstream) => relay(downstream, upstream).await, + Err(_) => tracing::debug!("websocket upgrade did not complete"), + } + }); + Ok(rewrite::switched_response(response).map(|_| Body::empty())) +} + +// uvicorn never half-closes a WebSocket, so Python's EOF means the socket is gone and the relay +// ends. Only a client that is gone ends it from the other side: the client's EOF just half-closes +// Python's side, and a write Python refuses means it closed with its last frames still in flight +// to the client. +async fn relay(downstream: Upgraded, upstream: Upgraded) { + let (mut client_read, mut client_write) = tokio::io::split(TokioIo::new(downstream)); + let (mut python_read, mut python_write) = tokio::io::split(TokioIo::new(upstream)); + let to_python = async { + let mut buffer = vec![0; 8192]; + loop { + let read = match client_read.read(&mut buffer).await { + Ok(0) => break, + Ok(read) => read, + Err(_) => return, + }; + if python_write.write_all(&buffer[..read]).await.is_err() { + return std::future::pending().await; + } + } + python_write.shutdown().await.ok(); + std::future::pending().await + }; + let to_client = async { + tokio::io::copy(&mut python_read, &mut client_write) + .await + .ok(); + client_write.shutdown().await.ok(); + }; + tokio::select! { + () = to_python => tracing::debug!("websocket client disconnected"), + () = to_client => {} + } +} diff --git a/litellm-rust/crates/gateway-forward/tests/client_identity.rs b/litellm-rust/crates/gateway-forward/tests/client_identity.rs new file mode 100644 index 00000000000..c484c89dc1d --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/client_identity.rs @@ -0,0 +1,223 @@ +mod support; + +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use axum::{body::Body, http::StatusCode, response::Response}; +use litellm_gateway_forward::{Forward, PEER_HEADER}; +use rstest::rstest; +use support::{DRAIN, Gateway, Raw, UVICORN_KEEPALIVE, raw_upstream, recording, within}; +use tokio::net::{TcpListener, TcpStream}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +#[rstest] +#[case::ordinary("")] +#[case::connection_nominated( + "connection: X-Forwarded-For, X-Forwarded-Proto, Forwarded, X-Litellm-Rust-Peer\r\n" +)] +#[tokio::test] +async fn python_sees_forwarding_headers_untouched_and_the_tcp_peer_rust_saw( + #[case] connection: &str, +) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "GET /sso/callback HTTP/1.1\r\n\ + host: public.example\r\n\ + x-forwarded-for: 192.0.2.1\r\n\ + x-forwarded-for: 192.0.2.2\r\n\ + x-forwarded-proto: https\r\n\ + forwarded: for=192.0.2.1;proto=https\r\n\ + x-litellm-rust-peer: 10.1.2.3:1\r\n{connection}\r\n", + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + + let request = received.recv().await.unwrap(); + let forwarded_for: Vec<_> = request.headers.get_all("x-forwarded-for").iter().collect(); + assert_eq!(forwarded_for, ["192.0.2.1", "192.0.2.2"]); + assert_eq!(request.headers["x-forwarded-proto"], "https"); + assert_eq!(request.headers["forwarded"], "for=192.0.2.1;proto=https"); + let peers: Vec<_> = request.headers.get_all(PEER_HEADER).iter().collect(); + assert_eq!(peers, [raw.local_address().to_string().as_str()]); +} + +/// Python builds redirects and OAuth metadata from Host, and without one from the address it +/// listens on, which behind Rust is its private loopback port. So a request without a usable Host +/// carries the public address the client connected to, except HTTP/1.1, which h11 rejects. +#[rstest] +#[case::http11_without_host("HTTP/1.1", "", false)] +#[case::http10_without_host("HTTP/1.0", "", true)] +#[case::http11_with_an_empty_host("HTTP/1.1", "host:\r\n", true)] +#[tokio::test] +async fn a_request_without_a_usable_host_never_names_the_private_upstream( + #[case] version: &str, + #[case] host: &str, + #[case] forwarded: bool, +) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!("GET /sso/key/generate {version}\r\n{host}\r\n")) + .await; + + let status = within(raw.response_head()).await.unwrap().status; + let hosts = received.try_recv().ok().map(|request| { + request + .headers + .get_all("host") + .iter() + .map(|value| value.to_str().unwrap().to_owned()) + .collect::>() + }); + let expected = if forwarded { + (StatusCode::OK, Some(vec![gateway.address.to_string()])) + } else { + (StatusCode::BAD_REQUEST, None) + }; + assert_eq!((status, hosts), expected); +} + +#[rstest] +#[case::uppercase_with_port("PUBLIC.Example:4000")] +#[case::ipv6_literal("[::1]:4000")] +#[case::ipv4_with_port("192.0.2.10:8443")] +#[tokio::test] +async fn python_sees_the_host_header_byte_for_byte(#[case] host: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!("GET /ui HTTP/1.1\r\nHost: {host}\r\n\r\n")) + .await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + let hosts: Vec<_> = request.headers.get_all("host").iter().collect(); + assert_eq!(hosts, [host]); +} + +/// Python's IP allowlists and rate limits read the peer exactly as uvicorn would have seen it on +/// the public socket, so a dual-stack listener reports IPv4 clients in their mapped form. +#[rstest] +#[case::ipv6("[::1]:0", "[::1]", "[::1]")] +#[case::ipv4_on_a_dual_stack_listener("[::]:0", "127.0.0.1", "[::ffff:127.0.0.1]")] +#[tokio::test] +async fn ipv6_listeners_report_the_peer_bracketed_with_its_port( + #[case] bind: &str, + #[case] connect: &str, + #[case] peer_ip: &str, +) { + let (upstream, mut received) = recording(ok).await; + let forward = Forward::new( + &format!("http://{upstream}").parse().unwrap(), + UVICORN_KEEPALIVE, + ); + let public = TcpListener::bind(bind).await.unwrap(); + let port = public.local_addr().unwrap().port(); + tokio::spawn(forward.serve(public, std::future::pending(), DRAIN)); + let mut raw = Raw::new( + TcpStream::connect(format!("{connect}:{port}")) + .await + .unwrap(), + ); + + raw.send("GET / HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!( + received.recv().await.unwrap().headers[PEER_HEADER], + format!("{peer_ip}:{}", raw.local_address().port()).as_str() + ); +} + +/// Without the peer header Python would take Rust's loopback hop as the client, which allowlists +/// often trust, so no spelling, repetition or Connection token may remove or replace it. +#[rstest] +#[case::named_as_hop_by_hop( + "connection: x-litellm-rust-peer\r\nx-litellm-rust-peer: 127.0.0.1:1\r\n" +)] +#[case::mixed_case_copies( + "X-LiteLLM-Rust-Peer: 127.0.0.1:1\r\nx-litellm-rust-peer: 192.0.2.7:2\r\n" +)] +#[case::comma_joined("x-litellm-rust-peer: 127.0.0.1:1, 192.0.2.7:2\r\n")] +#[tokio::test] +async fn a_client_can_neither_strip_nor_spoof_the_peer(#[case] spoof: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!( + "GET /key/info HTTP/1.1\r\nhost: public.example\r\n{spoof}\r\n" + )) + .await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + let peers: Vec<_> = request.headers.get_all(PEER_HEADER).iter().collect(); + assert_eq!(peers, [raw.local_address().to_string().as_str()]); +} + +/// Pooled upstream connections carry requests from many clients in turn, so the peer belongs to +/// each request and never to the connection it travels on. +#[rstest] +#[tokio::test] +async fn each_request_carries_its_own_peer_over_a_shared_upstream_connection() { + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = connections.clone(); + let upstream = raw_upstream(move |mut python| { + accepted.fetch_add(1, Ordering::SeqCst); + async move { + while let Some(head) = python.request_head().await { + let peer = head.header(PEER_HEADER).unwrap().to_vec(); + python + .send( + [ + format!("HTTP/1.1 200 OK\r\ncontent-length: {}\r\n\r\n", peer.len()) + .as_bytes(), + &peer, + ] + .concat(), + ) + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut clients = [ + gateway.connect().await, + gateway.connect().await, + gateway.connect().await, + ]; + + for index in [0, 1, 2, 2, 1, 0] { + let raw = &mut clients[index]; + raw.send("GET /v1/models HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + let peer = raw.body(head.content_length().unwrap()).await; + assert_eq!(peer, raw.local_address().to_string().into_bytes()); + } + assert_eq!(connections.load(Ordering::SeqCst), 1); +} diff --git a/litellm-rust/crates/gateway-forward/tests/download.rs b/litellm-rust/crates/gateway-forward/tests/download.rs new file mode 100644 index 00000000000..9707d2f53b7 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/download.rs @@ -0,0 +1,192 @@ +mod support; + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{body::Body, extract::Request, http::StatusCode, response::Response}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, client, raw_upstream, recording, within}; +use tokio::sync::{Notify, mpsc}; + +#[rstest] +#[tokio::test] +async fn streams_a_body_without_a_content_length_as_python_produces_it() { + let first: &[u8] = &[0xff, 0xfb, 0x90, 0x64, 0x00, 0x0d, 0x0a]; + let rest = vec![0xa5; 64 * 1024]; + let finish = Arc::new(Notify::new()); + let finished = finish.clone(); + let tail = rest.clone(); + let upstream = raw_upstream(move |mut python| { + let finished = finished.clone(); + let tail = tail.clone(); + async move { + python.request_head().await.unwrap(); + python + .send([&b"HTTP/1.1 200 OK\r\ncontent-type: audio/mpeg\r\ntransfer-encoding: chunked\r\n\r\n7\r\n"[..], first, b"\r\n"].concat()) + .await; + finished.notified().await; + python + .send([format!("{:x}\r\n", tail.len()).as_bytes(), &tail, b"\r\n0\r\n\r\n"].concat()) + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut response = within( + client().request( + Request::post(gateway.url("/v1/audio/speech")) + .body(Body::empty()) + .unwrap(), + ), + ) + .await + .unwrap(); + + assert!(!response.headers().contains_key("content-length")); + let frame = within(response.body_mut().frame()).await.unwrap().unwrap(); + assert_eq!(frame.into_data().unwrap(), first); + finish.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, rest); +} + +#[rstest] +#[tokio::test] +async fn range_requests_and_partial_content_pass_unchanged() { + let (upstream, mut received) = recording(|| -> Response { + Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header("content-range", "bytes 2-5/10") + .header("accept-ranges", "bytes") + .body(Body::from("2345")) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/videos/video_1/content HTTP/1.1\r\nhost: x\r\nrange: bytes=2-5\r\nif-range: \"v1\"\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::PARTIAL_CONTENT); + assert_eq!(head.header("content-range"), Some(&b"bytes 2-5/10"[..])); + assert_eq!(raw.body(head.content_length().unwrap()).await, b"2345"); + let request = received.recv().await.unwrap(); + assert_eq!(request.headers["range"], "bytes=2-5"); + assert_eq!(request.headers["if-range"], "\"v1\""); +} + +#[rstest] +#[tokio::test] +async fn a_large_download_is_paced_by_the_client_and_keeps_its_content_length() { + const LENGTH: usize = 128 * 1024 * 1024; + let written = Arc::new(AtomicUsize::new(0)); + let writing = written.clone(); + let upstream = raw_upstream(move |mut python| { + let writing = writing.clone(); + async move { + python.request_head().await.unwrap(); + python + .send(format!("HTTP/1.1 200 OK\r\ncontent-type: application/octet-stream\r\ncontent-length: {LENGTH}\r\n\r\n")) + .await; + let chunk = [0xa5; 64 * 1024]; + while writing.load(Ordering::SeqCst) < LENGTH { + python.send(chunk).await; + writing.fetch_add(chunk.len(), Ordering::SeqCst); + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/files/file-1/content HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.content_length(), Some(LENGTH)); + assert_eq!(head.header("transfer-encoding"), None); + tokio::time::sleep(Duration::from_secs(1)).await; + assert!( + written.load(Ordering::SeqCst) < LENGTH / 4, + "the forwarder buffered a download the client was not reading" + ); + let mut remaining = LENGTH; + while remaining > 0 { + remaining -= within(raw.body(remaining.min(1 << 20))).await.len(); + } +} + +/// The UI revalidates its static assets, and a 304 must carry exactly the validators Python sent, +/// with no length invented for a body it never has. +#[rstest] +#[tokio::test] +async fn a_conditional_request_gets_python_304_with_its_validators_only() { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let event = event.clone(); + async move { + while let Some(head) = python.request_head().await { + event + .send(head.header("if-none-match").map(<[u8]>::to_vec)) + .unwrap(); + python + .send("HTTP/1.1 304 Not Modified\r\netag: \"5f1-abc\"\r\ncache-control: no-cache\r\nvary: Accept-Encoding\r\ndate: Tue, 03 Nov 2026 10:00:00 GMT\r\n\r\n") + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + for _ in 0..2 { + raw.send("GET /ui/_next/static/chunks/main.js HTTP/1.1\r\nhost: x\r\nif-none-match: \"5f1-abc\"\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::NOT_MODIFIED); + let mut names: Vec<_> = head.headers.iter().map(|(name, _)| name.as_str()).collect(); + names.sort_unstable(); + assert_eq!(names, ["cache-control", "date", "etag", "vary"]); + assert_eq!(head.header("etag"), Some(&b"\"5f1-abc\""[..])); + assert_eq!( + head.header("date"), + Some(&b"Tue, 03 Nov 2026 10:00:00 GMT"[..]) + ); + assert_eq!( + within(events.recv()).await.unwrap().as_deref(), + Some(&b"\"5f1-abc\""[..]) + ); + } +} + +/// A client knows a download was cut short only because fewer bytes than the Content-Length +/// arrived before the connection closed, so the length must survive and the connection must end. +#[rstest] +#[tokio::test] +async fn a_download_python_cuts_short_closes_the_client_connection_before_its_length() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send([&b"HTTP/1.1 200 OK\r\ncontent-type: application/octet-stream\r\ncontent-length: 1000\r\n\r\n"[..], &[0xa5; 500]].concat()) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/files/file-1/content HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.content_length(), Some(1000)); + assert_eq!(within(raw.until_closed()).await, [0xa5; 500]); +} diff --git a/litellm-rust/crates/gateway-forward/tests/exchange.rs b/litellm-rust/crates/gateway-forward/tests/exchange.rs new file mode 100644 index 00000000000..8fed7b18453 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/exchange.rs @@ -0,0 +1,855 @@ +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + http::{HeaderValue, StatusCode}, + response::Response, +}; +use litellm_gateway_forward::PEER_HEADER; +use rstest::rstest; +use support::{Gateway, Head, client, raw_upstream, recording, serve, within}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + sync::mpsc, +}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +/// An upstream that reports each bodiless request head exactly as it arrived on the wire. +async fn heads() -> (std::net::SocketAddr, mpsc::UnboundedReceiver) { + let (sender, receiver) = mpsc::unbounded_channel(); + let address = raw_upstream(move |mut python| { + let sender = sender.clone(); + async move { + while let Some(head) = python.request_head().await { + sender.send(head).unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + (address, receiver) +} + +#[rstest] +#[tokio::test] +async fn forwards_the_request_as_the_client_sent_it_minus_hop_headers() { + let (upstream, mut received) = recording(|| { + Response::builder() + .status(StatusCode::UNAUTHORIZED) + .header("connection", "x-upstream-hop") + .header("x-upstream-hop", "remove") + .header("x-response", "keep") + .body(Body::from("python error")) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send( + "PATCH /v1/messages?x=a%2Fb&x=c HTTP/1.1\r\n\ + host: public.example\r\n\ + authorization: Bearer test\r\n\ + connection: x-remove, keep-alive\r\n\ + x-remove: remove\r\n\ + proxy-authorization: remove\r\n\ + te: trailers\r\n\ + content-length: 7\r\n\r\npayload", + ) + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::UNAUTHORIZED); + assert_eq!(head.header("x-response"), Some(&b"keep"[..])); + assert_eq!(head.header("x-upstream-hop"), None); + let body = raw.body(head.content_length().unwrap()).await; + assert_eq!(body, b"python error"); + + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "PATCH"); + assert_eq!(request.uri, "/v1/messages?x=a%2Fb&x=c"); + assert_eq!(request.headers["host"], "public.example"); + assert_eq!(request.headers["authorization"], "Bearer test"); + assert!(!request.headers.contains_key("x-remove")); + assert!(!request.headers.contains_key("proxy-authorization")); + assert!(!request.headers.contains_key("te")); + assert!(!request.headers.contains_key("connection")); + assert_eq!(request.body, "payload"); +} + +#[rstest] +#[case::host("HoSt", "public.example")] +#[case::length("Content-Length", "7")] +#[case::authorization("AUTHORIZATION", "Bearer test")] +#[case::cookie("Cookie", "session=test")] +#[case::api_key("X-Api-Key", "test")] +#[case::azure_key("Api-Key", "test")] +#[case::litellm_key("Litellm-Api-Key", "test")] +#[case::forwarded("Forwarded", "for=192.0.2.1;proto=https")] +#[case::forwarded_for("X-Forwarded-For", "192.0.2.1")] +#[case::forwarded_proto("X-Forwarded-Proto", "https")] +#[case::forwarded_host("X-Forwarded-Host", "public.example")] +#[case::content_type("Content-Type", "application/octet-stream")] +#[case::content_encoding("Content-Encoding", "identity")] +#[tokio::test] +async fn connection_cannot_remove_end_to_end_headers(#[case] name: &str, #[case] value: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let host = if name.eq_ignore_ascii_case("host") { + "" + } else { + "host: public.example\r\n" + }; + let length = if name.eq_ignore_ascii_case("content-length") { + "" + } else { + "content-length: 7\r\n" + }; + raw.send(format!( + "POST / HTTP/1.1\r\n{host}{length}{name}: {value}\r\nconnection: x-remove\r\nconnection: {name}, keep-alive\r\nx-remove: remove\r\n\r\npayload" + )).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = within(received.recv()).await.unwrap(); + assert_eq!(request.headers[name], value); + assert_eq!(request.body, "payload"); + assert!(!request.headers.contains_key("connection")); + assert!(!request.headers.contains_key("x-remove")); +} + +#[rstest] +#[tokio::test] +async fn connection_cannot_remove_response_credentials_or_representation_metadata() { + let (upstream, _received) = recording(|| { + Response::builder() + .status(StatusCode::UNAUTHORIZED) + .header("connection", "Content-Length, Set-Cookie, WWW-Authenticate, Authentication-Info, Content-Type, Content-Encoding, x-remove") + .header("content-length", "2") + .header("set-cookie", "session=; Max-Age=0") + .header("www-authenticate", "Bearer") + .header("authentication-info", "nextnonce=test") + .header("content-type", "text/plain") + .header("content-encoding", "identity") + .header("x-remove", "remove") + .body(Body::from("ok")) + .unwrap() + }).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::UNAUTHORIZED); + assert_eq!(head.content_length(), Some(2)); + assert_eq!(head.header("set-cookie"), Some(&b"session=; Max-Age=0"[..])); + assert_eq!(head.header("www-authenticate"), Some(&b"Bearer"[..])); + assert_eq!( + head.header("authentication-info"), + Some(&b"nextnonce=test"[..]) + ); + assert_eq!(head.header("content-type"), Some(&b"text/plain"[..])); + assert_eq!(head.header("content-encoding"), Some(&b"identity"[..])); + assert_eq!(head.header("x-remove"), None); + assert_eq!(raw.body(2).await, b"ok"); +} + +#[rstest] +#[case::get("GET")] +#[case::post("POST")] +#[case::put("PUT")] +#[case::delete("DELETE")] +#[case::options("OPTIONS")] +#[case::extension("PURGE")] +#[case::webdav("PROPFIND")] +#[case::lowercase("get")] +#[case::punctuated("M-SEARCH")] +#[tokio::test] +async fn forwards_every_method(#[case] method: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let response = client() + .request( + Request::builder() + .method(method) + .uri(gateway.url("/mcp")) + .body(Body::from("body")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, method); + assert_eq!(request.body, "body"); +} + +#[rstest] +#[case::double_slash("//v1//models")] +#[case::encoded_slash("/v1/files/a%2Fb/content")] +#[case::encoded_hash_in_query("/v1/files?purpose=batch%23retry&limit=1")] +#[case::encoded_hash_in_path("/v1/files/report%231.pdf/content")] +#[case::encoded_hash_before_query("/gemini/v1beta/files/a%23b?key=k")] +#[case::encoded_query("/v1/models?q=%E2%9C%93&empty=&flag")] +#[case::utf8_bytes_escaped("/mcp/%F0%9F%A6%80")] +#[case::lowercase_encoded_slash("/v1/files/a%2fb/content")] +#[case::dot_segments("/openai/./v1/../v1/models")] +#[case::encoded_dot_segments("/openai/%2e%2e/v1/%2E/models")] +#[case::matrix_params("/langfuse/api;v=1/public;x=y/traces")] +#[case::plus_and_encoded_space("/gemini/v1beta/models?q=a+b&r=a%20b&s=%2B")] +#[case::empty_query("/v1/models?")] +#[case::question_marks_in_query("/v1/models??a=b?c")] +#[case::colons_in_segments( + "/v1/projects/p/locations/l/publishers/google/models/m:streamGenerateContent?alt=sse" +)] +#[case::sub_delims("/~u/@:!$&'()*+,=/x")] +#[case::raw_bytes_browsers_leave_unescaped("/a{b}|c^d\\e\"f?g={}|^`")] +#[case::invalid_percent_escapes("/a%zz%00b?%")] +#[tokio::test] +async fn forwards_the_request_target_byte_for_byte(#[case] target: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n")) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().uri, target); +} + +#[rstest] +#[tokio::test] +async fn forwards_a_query_string_as_long_as_uvicorn_accepts() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let target = format!("/vertex_ai/v1/models?key=%2F{}&end", "a%20b".repeat(3000)); + raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n")) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().uri, target.as_str()); +} + +#[rstest] +#[case::in_query("/v1/files?purpose=batch#retry&limit=1", "/v1/files?purpose=batch")] +#[case::in_path("/v1/files/report#1.pdf/content", "/v1/files/report")] +#[case::before_query("/gemini/v1beta/files/a#b?key=k", "/gemini/v1beta/files/a")] +#[tokio::test] +async fn forwards_only_the_path_and_query_before_a_fragment( + #[case] target: &str, + #[case] expected: &str, +) { + let (upstream, mut received) = heads().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!("GET {target} HTTP/1.1\r\nhost: x\r\n\r\n")) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(within(received.recv()).await.unwrap().path, expected); +} + +#[rstest] +#[tokio::test] +async fn absolute_form_targets_reach_python_as_origin_form() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send( + "GET http://elsewhere.example/v1/models?x=1 HTTP/1.1\r\nhost: elsewhere.example\r\n\r\n", + ) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert_eq!(request.uri, "/v1/models?x=1"); +} + +#[rstest] +#[tokio::test] +async fn relays_a_response_with_more_headers_than_hyper_allows_by_default() { + let upstream = serve(axum::Router::new().fallback(|| async { + let headers = (0..150).map(|index| (format!("llm_provider-x-{index}"), "v")); + let mut response = ok(); + for (name, value) in headers { + response.headers_mut().insert( + axum::http::HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_static(value), + ); + } + response + })) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send("GET /v1/chat/completions HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.header("llm_provider-x-149"), Some(&b"v"[..])); +} + +/// uvicorn's h11 parser has no header count limit, only a 16 KiB cap on the whole head. +#[rstest] +#[tokio::test] +async fn forwards_a_request_with_as_many_headers_as_uvicorn_accepts() { + let (sender, mut received) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |python| { + let sender = sender.clone(); + async move { + let mut stream = python.into_stream(); + let mut head = Vec::new(); + while !head.ends_with(b"\r\n\r\n") { + head.push(stream.read_u8().await.unwrap()); + } + sender.send(head).unwrap(); + stream + .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") + .await + .unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let headers: String = (0..1000).map(|index| format!("x-{index}: v\r\n")).collect(); + raw.send(format!( + "GET /v1/chat/completions HTTP/1.1\r\nhost: x\r\n{headers}\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let head = String::from_utf8(within(received.recv()).await.unwrap()).unwrap(); + assert!(head.contains(&headers)); +} + +async fn header_fields_python_receives(fields: &str) -> Vec<(String, Vec)> { + let (upstream, mut received) = heads().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "GET /anthropic/v1/messages HTTP/1.1\r\nhost: x\r\n{fields}\r\n" + )) + .await; + within(raw.response_head()).await.unwrap(); + within(received.recv()) + .await + .unwrap() + .headers + .into_iter() + .filter(|(name, _)| name.starts_with("x-") && name != PEER_HEADER) + .collect() +} + +fn fields(expected: &[(&str, &str)]) -> Vec<(String, Vec)> { + expected + .iter() + .map(|(name, value)| ((*name).to_owned(), value.as_bytes().to_vec())) + .collect() +} + +/// Python reads `scope["headers"]` as the list the client sent, in that order. +#[rstest] +#[ignore = "divergence: HeaderMap::remove can move the headers after a dropped hop-by-hop one"] +#[tokio::test] +async fn dropping_a_hop_by_hop_header_keeps_the_others_in_order() { + let received = header_fields_python_receives( + "x-a: 1\r\nconnection: keep-alive\r\nx-b: 2\r\nx-c: 3\r\nx-d: 4\r\n", + ) + .await; + assert_eq!( + received, + fields(&[("x-a", "1"), ("x-b", "2"), ("x-c", "3"), ("x-d", "4")]) + ); +} + +/// uvicorn speaks only HTTP/1.1, so an HTTP/2 prior-knowledge client never reaches Python through +/// Rust either. hyper closes the connection where h11 answers 400, which no HTTP/2 client reads. +#[rstest] +#[tokio::test] +async fn an_http2_client_never_reaches_python() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00") + .await; + + assert!(within(raw.until_closed()).await.is_empty()); + assert!(received.try_recv().is_err()); +} + +#[rstest] +#[ignore = "divergence: HeaderMap groups every value of a name at its first position"] +#[tokio::test] +async fn interleaved_repeated_headers_reach_python_interleaved() { + let received = header_fields_python_receives("x-a: 1\r\nx-b: 2\r\nx-a: 3\r\n").await; + assert_eq!( + received, + fields(&[("x-a", "1"), ("x-b", "2"), ("x-a", "3")]) + ); +} + +/// h11 trims optional whitespace around a value and keeps everything inside it, including empty +/// values and names using every token character. +#[rstest] +#[tokio::test] +async fn python_sees_header_values_as_h11_would_parse_them() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nx-empty:\r\nx-blank: \r\nx-inner: a \t b \t\r\nX-Token!#$%&'*+.^_`|~: v\r\nx-empty: second\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + let values = |name: &str| -> Vec<&[u8]> { + request + .headers + .get_all(name) + .iter() + .map(HeaderValue::as_bytes) + .collect() + }; + assert_eq!(values("x-empty"), [&b""[..], b"second"]); + assert_eq!(values("x-blank"), [&b""[..]]); + assert_eq!(values("x-inner"), [&b"a \t b"[..]]); + assert_eq!(values("x-token!#$%&'*+.^_`|~"), [&b"v"[..]]); +} + +/// RFC 9112 section 5.2 lets a server either reject obs-fold or unfold it, and h11 unfolds. +#[rstest] +#[ignore = "divergence: hyper rejects obs-fold with a 400 where uvicorn unfolds it to one space"] +#[tokio::test] +async fn a_folded_header_reaches_python_unfolded() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nx-a: b\r\n c\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().headers["x-a"], "b c"); +} + +#[rstest] +#[tokio::test] +async fn keeps_repeated_response_headers_and_opaque_header_bytes() { + let (upstream, mut received) = recording(|| { + Response::builder() + .header("set-cookie", "a=1; Path=/") + .header("set-cookie", "b=2; HttpOnly") + .header("x-opaque", HeaderValue::from_bytes(b"caf\xe9").unwrap()) + .body(Body::empty()) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(b"GET / HTTP/1.1\r\nhost: x\r\nx-opaque: caf\xe9\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + let cookies: Vec<_> = head + .headers + .iter() + .filter(|(name, _)| name == "set-cookie") + .map(|(_, value)| value.as_slice()) + .collect(); + assert_eq!(cookies, [&b"a=1; Path=/"[..], b"b=2; HttpOnly"]); + assert_eq!(head.header("x-opaque"), Some(&b"caf\xe9"[..])); + assert_eq!( + received.recv().await.unwrap().headers["x-opaque"].as_bytes(), + b"caf\xe9" + ); +} + +/// An HTTP/1.0 client cannot decode chunked framing, so a streamed answer reaches it delimited by +/// the connection closing, as uvicorn sends it. +#[rstest] +#[tokio::test] +async fn an_http10_client_gets_a_streamed_answer_unframed_and_python_gets_its_body() { + let (upstream, mut received) = recording(|| { + Response::new(Body::from_stream(futures_util::stream::iter( + ["str", "eamed"].map(|piece| Ok::<_, Infallible>(Bytes::from_static(piece.as_bytes()))), + ))) + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/completions HTTP/1.0\r\nhost: x\r\ncontent-length: 7\r\n\r\npayload") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.header("transfer-encoding"), None); + assert_eq!(within(raw.until_closed()).await, b"streamed"); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "POST"); + assert_eq!(request.body, "payload"); +} + +#[rstest] +#[tokio::test] +async fn answers_pipelined_requests_in_order() { + let upstream = serve( + Router::new().fallback(|request: Request| async move { request.uri().path().to_owned() }), + ) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /first HTTP/1.1\r\nhost: x\r\n\r\nHEAD /second HTTP/1.1\r\nhost: x\r\n\r\nGET /third HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let first = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(first.content_length().unwrap()).await, b"/first"); + let second = within(raw.response_head()).await.unwrap(); + assert_eq!(second.content_length(), Some(7)); + let third = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(third.content_length().unwrap()).await, b"/third"); +} + +#[rstest] +#[case::no_content(StatusCode::NO_CONTENT)] +#[case::not_modified(StatusCode::NOT_MODIFIED)] +#[tokio::test] +async fn bodiless_statuses_keep_the_connection_reusable(#[case] status: StatusCode) { + let upstream = serve(Router::new().fallback(move |request: Request| async move { + match request.uri().path() { + "/next" => (StatusCode::OK, "next"), + _ => (status, ""), + } + })) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /empty HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + assert_eq!(within(raw.response_head()).await.unwrap().status, status); + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); +} + +#[rstest] +#[tokio::test] +async fn informational_responses_keep_the_connection_reusable() { + let upstream = raw_upstream(|mut python| async move { + while let Some(head) = python.request_head().await { + match head.path.as_str() { + "/hints" => { + python + .send("HTTP/1.1 103 Early Hints\r\nlink: ; rel=preload\r\n\r\nHTTP/1.1 200 OK\r\ncontent-length: 5\r\n\r\nfinal") + .await + } + _ => { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext") + .await + } + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /hints HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let mut hinted = within(raw.response_head()).await.unwrap(); + while hinted.status.is_informational() { + hinted = within(raw.response_head()).await.unwrap(); + } + assert_eq!(hinted.status, StatusCode::OK); + assert_eq!(raw.body(hinted.content_length().unwrap()).await, b"final"); + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); +} + +#[rstest] +#[case::found( + StatusCode::FOUND, + "https://idp.example/authorize?state=a%2Fb&redirect_uri=https%3A%2F%2Fpublic.example%2Fsso%2Fcallback" +)] +#[case::see_other(StatusCode::SEE_OTHER, "/ui/?login=success")] +#[case::temporary(StatusCode::TEMPORARY_REDIRECT, "./mcp/oauth/authorize?client_id=x")] +#[tokio::test] +async fn redirects_reach_the_client_with_location_and_cookies_unchanged( + #[case] status: StatusCode, + #[case] location: &'static str, +) { + let (upstream, _received) = recording(move || { + Response::builder() + .status(status) + .header("location", location) + .header( + "set-cookie", + "token=a%2Fb; Path=/; HttpOnly; Secure; SameSite=Lax", + ) + .header("set-cookie", "state=; Max-Age=0") + .header( + "set-cookie", + "session=x; Expires=Wed, 04 Nov 2026 10:00:00 GMT; Path=/", + ) + .body(Body::empty()) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /sso/key/generate HTTP/1.1\r\nhost: public.example\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, status); + assert_eq!(head.header("location"), Some(location.as_bytes())); + let cookies: Vec<_> = head + .headers + .iter() + .filter(|(name, _)| name == "set-cookie") + .map(|(_, value)| value.as_slice()) + .collect(); + assert_eq!( + cookies, + [ + &b"token=a%2Fb; Path=/; HttpOnly; Secure; SameSite=Lax"[..], + b"state=; Max-Age=0", + b"session=x; Expires=Wed, 04 Nov 2026 10:00:00 GMT; Path=/", + ] + ); +} + +#[rstest] +#[tokio::test] +async fn non_websocket_upgrades_reach_python_as_plain_requests() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nconnection: upgrade, http2-settings\r\nupgrade: h2c\r\nhttp2-settings: AAMAAABkAAQAoAAAAAIAAAAA\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert!(!request.headers.contains_key("upgrade")); + assert!(!request.headers.contains_key("connection")); + assert!(!request.headers.contains_key("http2-settings")); +} + +#[rstest] +#[tokio::test] +async fn forwards_an_asterisk_form_options_request() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("OPTIONS * HTTP/1.1\r\nhost: x\r\n\r\n").await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "OPTIONS"); + assert_eq!(request.uri, "*"); +} + +#[rstest] +#[tokio::test] +async fn an_authority_form_target_is_a_400_python_never_sees() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("CONNECT api.example:443 HTTP/1.1\r\nhost: api.example:443\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); + assert!(received.try_recv().is_err()); +} + +#[rstest] +#[tokio::test] +async fn a_client_connection_close_ends_only_the_client_connection() { + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = connections.clone(); + let upstream = raw_upstream(move |mut python| { + accepted.fetch_add(1, Ordering::SeqCst); + async move { + while python.request_head().await.is_some() { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + for _ in 0..2 { + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\nconnection: close\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.header("connection"), Some(&b"close"[..])); + assert_eq!(raw.body(2).await, b"ok"); + within(raw.closed()).await; + } + assert_eq!(connections.load(Ordering::SeqCst), 1); +} + +#[rstest] +#[tokio::test] +async fn a_python_connection_close_ends_only_the_upstream_connection() { + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = connections.clone(); + let upstream = raw_upstream(move |mut python| { + accepted.fetch_add(1, Ordering::SeqCst); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 500 Internal Server Error\r\ncontent-length: 5\r\nconnection: close\r\n\r\nerror") + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + for _ in 0..2 { + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(head.header("connection"), None); + assert_eq!(raw.body(5).await, b"error"); + } + assert_eq!(connections.load(Ordering::SeqCst), 2); +} + +/// uvicorn answers HEAD on a streaming route with chunked framing and no body, so the client +/// must not be told the body is empty, and must not wait for one before its next request. +#[rstest] +#[tokio::test] +async fn head_answered_with_chunked_framing_claims_no_length_and_keeps_the_connection() { + let upstream = raw_upstream(|mut python| async move { + while let Some(head) = python.request_head().await { + match head.path.as_str() { + "/next" => { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext") + .await + } + _ => { + python + .send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n") + .await + } + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("HEAD /mcp HTTP/1.1\r\nhost: x\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + assert_eq!(head.content_length(), None); + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); +} + +const GZIP: &[u8] = &[ + 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x4b, 0x49, 0x2c, 0x49, 0xb4, 0x52, + 0x30, 0xe4, 0xe2, 0x02, 0x00, 0xce, 0xcd, 0x78, 0x54, 0x09, 0x00, 0x00, 0x00, +]; + +#[rstest] +#[case::fixed_length(false)] +#[case::streamed(true)] +#[tokio::test] +async fn content_codings_cross_unchanged_both_ways(#[case] streamed: bool) { + let (upstream, mut received) = recording(move || { + let body = match streamed { + true => Body::from_stream(futures_util::stream::iter( + GZIP.chunks(5) + .map(|piece| Ok::<_, Infallible>(Bytes::from_static(piece))), + )), + false => Body::from(GZIP), + }; + Response::builder() + .header("content-encoding", "gzip") + .header("vary", "Accept-Encoding") + .body(body) + .unwrap() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send( + [ + format!( + "POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-encoding: gzip\r\naccept-encoding: gzip, br;q=0.5\r\ncontent-length: {}\r\n\r\n", + GZIP.len() + ) + .as_bytes(), + GZIP, + ] + .concat(), + ) + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.header("content-encoding"), Some(&b"gzip"[..])); + assert_eq!(head.header("vary"), Some(&b"Accept-Encoding"[..])); + let body = match head.content_length() { + Some(length) => raw.body(length).await, + None => within(raw.chunked_body()).await, + }; + assert_eq!(body, GZIP); + assert_eq!(head.content_length().is_none(), streamed); + let request = received.recv().await.unwrap(); + assert_eq!(request.headers["content-encoding"], "gzip"); + assert_eq!(request.headers["accept-encoding"], "gzip, br;q=0.5"); + assert_eq!(request.body, GZIP); +} + +#[rstest] +#[tokio::test] +async fn python_sees_no_accept_encoding_the_client_did_not_send() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /metrics/ HTTP/1.1\r\nhost: x\r\n\r\n").await; + within(raw.response_head()).await.unwrap(); + assert!( + !received + .recv() + .await + .unwrap() + .headers + .contains_key("accept-encoding") + ); +} diff --git a/litellm-rust/crates/gateway-forward/tests/framing.rs b/litellm-rust/crates/gateway-forward/tests/framing.rs new file mode 100644 index 00000000000..32a2c3f2b8d --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/framing.rs @@ -0,0 +1,259 @@ +mod support; + +use std::{collections::BTreeSet, net::SocketAddr}; + +use axum::http::StatusCode; +use rstest::rstest; +use support::{Gateway, Head, Raw, raw_upstream, within}; +use tokio::sync::mpsc; + +struct Captured { + head: Head, + body: Vec, +} + +/// A keep-alive upstream that decodes framing like h11 and reports every request it reads. +async fn capturing() -> (SocketAddr, mpsc::UnboundedReceiver) { + let (sender, receiver) = mpsc::unbounded_channel(); + let address = raw_upstream(move |mut python: Raw| { + let sender = sender.clone(); + async move { + while let Some(head) = python.request_head().await { + let body = match head.header("transfer-encoding") { + Some(_) => python.chunked_body().await, + None => python.body(content_length(&head)).await, + }; + sender.send(Captured { head, body }).unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + (address, receiver) +} + +/// The length h11 reads from every Content-Length field, which may repeat one value. +fn content_length(head: &Head) -> usize { + let lengths: BTreeSet = head + .headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("content-length")) + .flat_map(|(_, value)| value.split(|byte| *byte == b',')) + .map(|length| { + std::str::from_utf8(length.trim_ascii()) + .unwrap() + .parse() + .unwrap() + }) + .collect(); + assert!( + lengths.len() < 2, + "python got conflicting lengths {lengths:?}" + ); + lengths.first().copied().unwrap_or(0) +} + +/// RFC 9112 section 6.3: Transfer-Encoding overrides Content-Length. Forwarding both would let +/// Python split the stream at a different place than Rust did. Section 6.1 then closes the +/// connection after the response, so a request smuggled behind it never runs. +#[rstest] +#[tokio::test] +async fn a_request_framed_two_ways_reaches_python_framed_by_transfer_encoding_and_closes() { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 3\r\ntransfer-encoding: chunked\r\n\r\n5\r\nhello\r\n0\r\n\r\nGET /next HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + raw.body(2).await; + within(raw.until_closed()).await; + + let upload = received.recv().await.unwrap(); + assert_eq!(upload.head.header("content-length"), None); + assert_eq!( + upload.head.header("transfer-encoding"), + Some(&b"chunked"[..]) + ); + assert_eq!(upload.body, b"hello"); + assert!( + received.try_recv().is_err(), + "python ran the smuggled request" + ); +} + +#[rstest] +#[case::conflicting_content_lengths("content-length: 4\r\ncontent-length: 40\r\n")] +#[case::content_length_list("content-length: 4, 40\r\n")] +#[case::signed_content_length("content-length: +4\r\n")] +#[case::chunked_not_last("transfer-encoding: chunked, identity\r\ncontent-length: 4\r\n")] +#[case::unknown_coding("transfer-encoding: xchunked\r\ncontent-length: 4\r\n")] +#[case::space_before_colon("transfer-encoding : chunked\r\ncontent-length: 4\r\n")] +#[case::folded_transfer_encoding("transfer-encoding:\r\n chunked\r\ncontent-length: 4\r\n")] +#[case::bare_cr("x-a: b\rx-b: c\r\ncontent-length: 4\r\n")] +#[case::nul_byte("x-a: b\x00c\r\ncontent-length: 4\r\n")] +#[tokio::test] +async fn ambiguous_framing_never_reaches_python(#[case] headers: &str) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\nbodyGET /smuggled HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); + within(raw.closed()).await; + assert!(received.try_recv().is_err(), "python received a request"); +} + +#[rstest] +#[case::bare_lf("4\nbody\n0\n\n")] +#[case::invalid_size("zz\r\nbody\r\n0\r\n\r\n")] +#[case::size_overflow("10000000000000004\r\nbody\r\n0\r\n\r\n")] +#[tokio::test] +async fn a_malformed_request_body_is_a_400_not_a_502(#[case] chunks: &str) { + let upstream = raw_upstream(|mut python| async move { python.closed().await }).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{chunks}" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); +} + +/// hyper decodes only the chunked layer, so any other coding would reach Python still coded but +/// framed as a plain body. uvicorn turns every h11 parse error into a 400, ignoring the 501 h11 +/// suggests for codings. +#[rstest] +#[case::listed_before_chunked("transfer-encoding: gzip, chunked\r\n")] +#[case::split_across_fields("transfer-encoding: gzip\r\ntransfer-encoding: chunked\r\n")] +#[tokio::test] +async fn a_transfer_coding_other_than_chunked_is_the_400_uvicorn_sends_python_never_sees( + #[case] headers: &str, +) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\n4\r\nbody\r\n0\r\n\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_REQUEST + ); + assert!( + received.try_recv().is_err(), + "python received a still-coded body" + ); +} + +async fn assert_python_reads_the_body_and_the_next_request(headers: &str) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{headers}\r\npayloadGET /next HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + for _ in 0..2 { + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + raw.body(2).await; + } + assert_eq!(received.recv().await.unwrap().body, b"payload"); + assert_eq!(received.recv().await.unwrap().head.path, "/next"); +} + +#[rstest] +#[case::repeated_fields("content-length: 7\r\ncontent-length: 7\r\n")] +#[case::leading_zeros("content-length: 007\r\n")] +#[tokio::test] +async fn a_content_length_h11_reads_one_way_reaches_python_framed_the_same(#[case] headers: &str) { + assert_python_reads_the_body_and_the_next_request(headers).await; +} + +/// h11 accepts a Content-Length list that repeats one value, as RFC 9110 section 8.6 allows. +#[rstest] +#[ignore = "divergence: hyper's server parses each Content-Length field as one number and 400s a list"] +#[tokio::test] +async fn a_content_length_list_of_one_value_reaches_python() { + assert_python_reads_the_body_and_the_next_request("content-length: 7, 7\r\n").await; +} + +#[rstest] +#[case::leading_zeros("0004\r\nbody\r\n0\r\n\r\n", "body")] +#[case::uppercase_hex("A\r\nbody\r\nbody\r\n0\r\n\r\n", "body\r\nbody")] +#[case::whitespace_after_size("4 \r\nbody\r\n0\r\n\r\n", "body")] +#[case::extension_and_trailer("4;a=\"b;c\"\r\nbody\r\n0\r\nx-sum: 1\r\n\r\n", "body")] +#[tokio::test] +async fn chunk_framing_h11_accepts_reaches_python_decoded( + #[case] chunks: &str, + #[case] body: &str, +) { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{chunks}GET /next HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + for _ in 0..2 { + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + raw.body(2).await; + } + assert_eq!(received.recv().await.unwrap().body, body.as_bytes()); + assert_eq!(received.recv().await.unwrap().head.path, "/next"); +} + +/// h11 reads a head whose lines end in a bare LF, so Python gets the same request through Rust. +#[rstest] +#[tokio::test] +async fn a_head_with_bare_lf_line_endings_reaches_python() { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/files HTTP/1.1\nhost: x\nx-a: b\ncontent-length: 4\n\nbody") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let upload = received.recv().await.unwrap(); + assert_eq!(upload.head.header("x-a"), Some(&b"b"[..])); + assert_eq!(upload.body, b"body"); +} + +/// h11 decodes a chunked body from an HTTP/1.0 client, and RFC 9112 section 6.1 only asks the +/// server to close the connection afterwards. +#[rstest] +#[ignore = "divergence: hyper rejects Transfer-Encoding on HTTP/1.0 with a 400 that uvicorn never sends"] +#[tokio::test] +async fn an_http10_chunked_upload_reaches_python() { + let (upstream, mut received) = capturing().await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/files HTTP/1.0\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n4\r\nbody\r\n0\r\n\r\n") + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, b"body"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/idle_stream.rs b/litellm-rust/crates/gateway-forward/tests/idle_stream.rs new file mode 100644 index 00000000000..a07c0466481 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/idle_stream.rs @@ -0,0 +1,102 @@ +mod support; + +use std::time::Duration; + +use axum::{ + body::Body, + extract::Request, + http::StatusCode, + response::{IntoResponse, Response}, +}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, client, raw_upstream, recording, within}; +use tokio::{sync::mpsc, time::Instant}; + +const SSE_HEAD: &str = "HTTP/1.1 200 OK\r\n\ + date: Tue, 03 Nov 2026 00:00:00 GMT\r\n\ + server: uvicorn\r\n\ + cache-control: no-cache, no-transform\r\n\ + connection: keep-alive\r\n\ + content-type: text/event-stream\r\n\ + mcp-session-id: session-1\r\n\ + x-accel-buffering: no\r\n\ + transfer-encoding: chunked\r\n\r\n"; + +const PING: &str = ": ping - 2026-11-03 00:20:00.000000+00:00\r\n\r\n"; + +const MESSAGE: &str = "event: message\r\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/tools/list_changed\"}\r\n\r\n"; + +fn chunk(data: &str) -> String { + format!("{:x}\r\n{data}\r\n", data.len()) +} + +/// An MCP client holds a GET stream open for server-initiated messages, which may not come for +/// many minutes, so neither hop may close a stream for being idle. +#[rstest] +#[tokio::test(start_paused = true)] +async fn an_idle_mcp_stream_stays_open_in_both_directions() { + let idle = Duration::from_secs(20 * 60); + let (session, mut sessions) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let session = session.clone(); + async move { + let head = python.request_head().await.unwrap(); + session + .send(head.header("mcp-session-id").map(<[u8]>::to_vec)) + .unwrap(); + python.send(SSE_HEAD).await; + tokio::time::sleep(idle).await; + python.send(chunk(PING)).await; + python.send(chunk(MESSAGE)).await; + python.send("0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let started = Instant::now(); + + let response = client() + .request( + Request::get(gateway.url("/github/mcp")) + .header("accept", "text/event-stream") + .header("mcp-session-id", "session-1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.headers()["mcp-session-id"], "session-1"); + assert_eq!(response.headers()["x-accel-buffering"], "no"); + assert!(!response.headers().contains_key("connection")); + assert_eq!( + sessions.recv().await.unwrap().as_deref(), + Some(&b"session-1"[..]) + ); + let body = response.into_body().collect().await.unwrap().to_bytes(); + assert_eq!(body, [PING, MESSAGE].concat()); + assert!(started.elapsed() >= idle); +} + +#[rstest] +#[tokio::test] +async fn ends_a_session_with_a_bodiless_delete() { + let (upstream, mut received) = recording(|| -> Response { + ([("mcp-session-id", "session-1")], StatusCode::NO_CONTENT).into_response() + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("DELETE /mcp HTTP/1.1\r\nhost: x\r\nmcp-session-id: session-1\r\n\r\n") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::NO_CONTENT); + assert_eq!(head.header("mcp-session-id"), Some(&b"session-1"[..])); + let request = received.recv().await.unwrap(); + assert_eq!(request.method, "DELETE"); + assert_eq!(request.headers["mcp-session-id"], "session-1"); + assert!(!request.headers.contains_key("transfer-encoding")); + assert!(request.body.is_empty()); +} diff --git a/litellm-rust/crates/gateway-forward/tests/lifecycle.rs b/litellm-rust/crates/gateway-forward/tests/lifecycle.rs new file mode 100644 index 00000000000..85d10a37be4 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/lifecycle.rs @@ -0,0 +1,473 @@ +mod support; + +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use axum::{Router, body::Body, extract::Request, http::StatusCode, response::Response}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{ + DRAIN, Gateway, Raw, client, raw_upstream, recording, serve, unreachable_upstream, within, +}; +use tokio::{ + net::{TcpListener, TcpStream}, + sync::{Notify, mpsc}, +}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +/// Python stops a provider stream when it sees the client go away, so a client disconnect must +/// close Rust's connection to Python at every stage of the exchange. +#[rstest] +#[case::waiting_for_headers( + "POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n", + false +)] +#[case::mid_stream( + "POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n", + true +)] +#[case::mid_upload( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 1000000\r\n\r\npartial", + false +)] +#[case::waiting_for_websocket_handshake( + "GET /v1/realtime HTTP/1.1\r\nhost: x\r\nconnection: Upgrade\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: dGhlIHNhbXBsZSBub25jZQ==\r\n\r\n", + false +)] +#[tokio::test] +async fn client_disconnect_closes_the_upstream_connection( + #[case] request: &'static str, + #[case] headers_sent: bool, +) { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let event = event.clone(); + async move { + python.request_head().await.unwrap(); + if headers_sent { + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n6\r\ndata: \r\n") + .await; + } + event.send("received").unwrap(); + python.closed().await; + event.send("closed").unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(request).await; + assert_eq!(within(events.recv()).await, Some("received")); + if headers_sent { + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + } + drop(raw); + assert_eq!(within(events.recv()).await, Some("closed")); +} + +/// Python treats EOF on its socket as the client leaving, so Rust must not half-close the +/// upstream connection once it has sent the request body. +#[rstest] +#[tokio::test] +async fn python_sees_no_eof_while_its_response_is_pending() { + let (verdict, mut verdicts) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let verdict = verdict.clone(); + async move { + let head = python.request_head().await.unwrap(); + python.body(head.content_length().unwrap()).await; + let eof = tokio::time::timeout(Duration::from_millis(200), python.closed()) + .await + .is_ok(); + verdict.send(eof).unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + + let response = within( + client().request( + Request::post(gateway.url("/mcp")) + .body(Body::from("{}")) + .unwrap(), + ), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(verdicts.recv().await, Some(false)); +} + +/// uvicorn closes a connection that sat idle past its keep-alive timeout. A request Rust sends on +/// that connection as uvicorn closes it is lost, so Rust must retire idle connections first. +#[rstest] +#[case::reused_while_fresh(Duration::from_millis(0))] +#[case::past_the_pool_timeout(Duration::from_millis(250))] +#[case::past_python_keepalive(Duration::from_millis(450))] +#[tokio::test] +async fn never_reuses_a_connection_python_is_closing_for_idleness(#[case] idle: Duration) { + let keepalive = Duration::from_millis(400); + let upstream = raw_upstream(move |mut python: Raw| async move { + let mut idle_since = None::; + while python.request_head().await.is_some() { + if idle_since.is_some_and(|since| since.elapsed() >= keepalive) { + return; + } + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + idle_since = Some(Instant::now()); + } + }) + .await; + let gateway = Gateway::with_keepalive(upstream, keepalive).await; + let client = client(); + let first = within(client.request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + first.into_body().collect().await.unwrap(); + tokio::time::sleep(idle).await; + let second = + within(client.request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + assert_eq!(second.status(), StatusCode::OK); +} + +#[rstest] +#[case::health("/health/liveliness")] +#[case::messages("/v1/messages")] +#[tokio::test] +async fn unreachable_python_is_a_502(#[case] path: &str) { + let gateway = Gateway::start(unreachable_upstream()).await; + let response = + within(client().request(Request::get(gateway.url(path)).body(Body::empty()).unwrap())) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); +} + +#[rstest] +#[case::bad_status_line(b"HTTP/1.1 OK\r\ncontent-length: 0\r\n\r\n".to_vec())] +#[case::header_without_colon(b"HTTP/1.1 200 OK\r\nbroken header\r\ncontent-length: 0\r\n\r\n".to_vec())] +#[case::not_http(b"SSH-2.0-OpenSSH_9.6\r\n".to_vec())] +#[case::conflicting_lengths( + b"HTTP/1.1 200 OK\r\ncontent-length: 2\r\ncontent-length: 3\r\n\r\nok".to_vec() +)] +#[case::oversized_head( + [&b"HTTP/1.1 200 OK\r\nx-big: "[..], &vec![b'a'; 1024 * 1024], b"\r\n\r\n"].concat() +)] +#[case::closed_without_reply(Vec::new())] +#[tokio::test] +async fn an_upstream_that_breaks_http_is_a_502_and_the_client_connection_survives( + #[case] reply: Vec, +) { + let upstream = raw_upstream(move |mut python: Raw| { + let reply = reply.clone(); + async move { + python.request_head().await.unwrap(); + if reply.is_empty() { + return; + } + python.send(reply).await; + std::future::pending::<()>().await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + for _ in 0..2 { + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::BAD_GATEWAY); + assert_eq!(head.content_length(), Some(0)); + } +} + +#[rstest] +#[tokio::test] +async fn python_restarting_is_a_502_until_it_listens_again() { + let address = unreachable_upstream(); + let gateway = Gateway::start(address).await; + let get = || client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()); + assert_eq!( + within(get()).await.unwrap().status(), + StatusCode::BAD_GATEWAY + ); + + let listener = TcpListener::bind(address).await.unwrap(); + let router = Router::new().fallback(|| async { "up" }); + tokio::spawn(async move { axum::serve(listener, router).await }); + + assert_eq!(within(get()).await.unwrap().status(), StatusCode::OK); +} + +/// A request Python already read may have run, so it is answered with a 502 and never replayed. +#[rstest] +#[tokio::test] +async fn a_request_python_read_before_closing_is_a_502_and_never_replayed() { + let (event, mut seen) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let event = event.clone(); + async move { + while let Some(head) = python.request_head().await { + python.body(head.content_length().unwrap()).await; + event.send(head.path.clone()).unwrap(); + if head.path == "/v1/chat/completions" { + return; + } + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 2\r\n\r\nok") + .await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /warm HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n") + .await; + let warm = within(raw.response_head()).await.unwrap(); + raw.body(warm.content_length().unwrap()).await; + + raw.send("POST /v1/chat/completions HTTP/1.1\r\nhost: x\r\ncontent-length: 2\r\n\r\n{}") + .await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::BAD_GATEWAY); + let paths: Vec<_> = std::iter::from_fn(|| seen.try_recv().ok()).collect(); + assert_eq!(paths, ["/warm", "/v1/chat/completions"]); +} + +#[rstest] +#[tokio::test] +async fn shutdown_lets_an_in_flight_stream_finish() { + let release = Arc::new(Notify::new()); + let released = release.clone(); + let upstream = raw_upstream(move |mut python: Raw| { + let released = released.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n") + .await; + released.notified().await; + python.send("4\r\nlast\r\n0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = + within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + let server = gateway.shutdown(); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(!server.is_finished()); + release.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, "firstlast"); + within(server).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn shutdown_stops_waiting_for_a_stream_still_open_after_the_drain() { + let upstream = raw_upstream(|mut python: Raw| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n") + .await; + std::future::pending::<()>().await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = + within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + let started = std::time::Instant::now(); + + within(gateway.shutdown()).await.unwrap(); + + assert!(started.elapsed() >= DRAIN); + drop(response); +} + +#[rstest] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_keep_alive_connections_each_get_their_own_responses() { + let upstream = serve(Router::new().fallback(|request: Request| async move { + let index: usize = request.uri().query().unwrap().parse().unwrap(); + Response::builder() + .header("etag", format!("\"{index}\"")) + .body(Body::from(index.to_string().repeat(1 + index * 1000))) + .unwrap() + })) + .await; + let gateway = Gateway::start(upstream).await; + let connections = 6; + + let tasks: Vec<_> = (0..connections) + .map(|first| { + let address = gateway.address; + tokio::spawn(async move { + let mut raw = Raw::new(TcpStream::connect(address).await.unwrap()); + for index in (first..50).step_by(connections) { + raw.send(format!( + "GET /ui/chunk.js?{index} HTTP/1.1\r\nhost: x\r\n\r\n" + )) + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.header("etag"), Some(format!("\"{index}\"").as_bytes())); + let body = raw.body(head.content_length().unwrap()).await; + assert_eq!( + body, + index.to_string().repeat(1 + index * 1000).into_bytes() + ); + } + }) + }) + .collect(); + + for task in tasks { + within(task).await.unwrap(); + } +} + +// Python's `--keepalive_timeout` governed idle client connections before Rust took the public +// port, and operators tune it against their load balancer's idle timeout, so it still does. +// Without it, a client that vanishes without closing its connection holds a socket forever and +// the gateway eventually runs out of file descriptors. +const IDLE: Duration = Duration::from_millis(400); + +#[rstest] +#[tokio::test] +async fn closes_a_client_connection_idle_past_python_keepalive() { + let (upstream, _received) = recording(ok).await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + for _ in 0..2 { + tokio::time::sleep(IDLE / 2).await; + raw.send("GET /health/liveliness HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK); + raw.body(head.content_length().unwrap()).await; + } + let idle_since = Instant::now(); + within(raw.closed()).await; + + assert!(idle_since.elapsed() >= IDLE * 3 / 4); +} + +/// uvicorn would hold a connection that never sends a request, which lets a client park sockets +/// for free, so the same limit applies before the first request too. +#[rstest] +#[tokio::test] +async fn closes_a_client_connection_that_never_sends_a_request() { + let (upstream, _received) = recording(ok).await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + within(raw.closed()).await; +} + +/// The limit only covers the wait for the next request. A response that goes quiet for longer +/// belongs to Python, which decides how long it may take. +#[rstest] +#[tokio::test] +async fn keeps_a_connection_open_while_a_response_goes_quiet() { + let upstream = raw_upstream(|mut python: Raw| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + tokio::time::sleep(IDLE * 3).await; + python.send("4\r\ndone\r\n0\r\n\r\n").await; + }) + .await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + raw.send("GET /v1/chat/completions HTTP/1.1\r\nhost: x\r\n\r\n") + .await; + + within(raw.response_head()).await.unwrap(); + assert_eq!(within(raw.chunked_body()).await, b"done"); +} + +/// Likewise for an upload that pauses longer than the limit before it finishes. +#[rstest] +#[tokio::test] +async fn keeps_a_connection_open_while_an_upload_goes_quiet() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::with_keepalive(upstream, IDLE).await; + let mut raw = gateway.connect().await; + + raw.send("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: 8\r\n\r\nhalf") + .await; + tokio::time::sleep(IDLE * 3).await; + raw.send("done").await; + + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, "halfdone"); +} + +/// uvicorn answers a request it rejects (413, 401) as soon as it reads the head and then closes, +/// leaving the upload unread. The client must get that answer, not a 502 or a reset. +#[rstest] +#[case::sized("content-length: 16777216\r\n", vec![b'x'; 256 * 1024])] +#[case::chunked( + "transfer-encoding: chunked\r\n", + [&b"40000\r\n"[..], &[b'x'; 256 * 1024], b"\r\n"].concat() +)] +#[tokio::test] +async fn python_rejecting_an_upload_unread_and_closing_still_answers_the_client( + #[case] framing: &'static str, + #[case] partial_upload: Vec, +) { + let upstream = raw_upstream(|mut python: Raw| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 413 Payload Too Large\r\ncontent-type: application/json\r\ncontent-length: 26\r\nconnection: close\r\n\r\n{\"error\":\"file too large\"}") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\n{framing}\r\n" + )) + .await; + raw.send(partial_upload).await; + + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::PAYLOAD_TOO_LARGE); + assert_eq!( + raw.body(head.content_length().unwrap()).await, + b"{\"error\":\"file too large\"}" + ); +} diff --git a/litellm-rust/crates/gateway-forward/tests/roundtrip.rs b/litellm-rust/crates/gateway-forward/tests/roundtrip.rs new file mode 100644 index 00000000000..c220c7fe8e9 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/roundtrip.rs @@ -0,0 +1,360 @@ +mod support; + +use std::{ + convert::Infallible, + sync::{Arc, Mutex}, +}; + +use axum::{ + body::{Body, Bytes}, + http::{HeaderMap, HeaderName, HeaderValue, Method, Request, StatusCode}, + response::Response, +}; +use http_body_util::BodyExt; +use litellm_gateway_forward::PEER_HEADER; +use proptest::{prelude::*, test_runner::TestRunner}; +use support::{Gateway, client, recording, within}; + +#[derive(Clone, Debug)] +struct Message { + headers: Vec<(String, Vec>)>, + body: Vec, + chunked: bool, +} + +#[derive(Clone, Debug)] +struct Exchange { + method: Method, + target: String, + request: Message, + status: StatusCode, + response: Message, +} + +/// Any value h11 reads back unchanged: visible ASCII, obs-text and interior whitespace. +fn header_value() -> impl Strategy> { + prop::collection::vec( + prop_oneof![8 => 0x21u8..=0x7e, 2 => 0x80u8..=0xff, 1 => Just(b' '), 1 => Just(b'\t')], + 0..24, + ) + .prop_map(|value| value.trim_ascii().to_vec()) +} + +fn method() -> impl Strategy { + prop_oneof![ + Just(Method::GET), + Just(Method::POST), + Just(Method::PUT), + Just(Method::PATCH), + Just(Method::DELETE), + Just(Method::OPTIONS), + "[A-Za-z0-9!#$%&'*+.^_`|~-]{1,16}" + .prop_filter("HEAD and CONNECT frame differently", |method| { + !["HEAD", "CONNECT"].contains(&method.as_str()) + }) + .prop_map(|method| Method::from_bytes(method.as_bytes()).unwrap()), + ] +} + +fn message() -> impl Strategy { + ( + prop::collection::btree_map( + "x-[a-z0-9!#$%&'*+.^_`|~-]{1,12}", + prop::collection::vec(header_value(), 1..4), + 0..12, + ), + prop::collection::vec(any::(), 0..32 * 1024), + any::(), + ) + .prop_map(|(headers, body, chunked)| Message { + headers: headers.into_iter().collect(), + body, + chunked, + }) +} + +fn target() -> impl Strategy { + let segment = prop_oneof![ + "[A-Za-z0-9._~!$&'()*+,;=:@-]{0,8}", + "%[0-9A-Fa-f]{2}", + Just(".".to_owned()), + Just("..".to_owned()), + Just(String::new()), + ]; + ( + prop::collection::vec(segment, 0..6), + prop::option::of("[A-Za-z0-9=&._~%/?:@!$'()*+,;-]{0,40}"), + ) + .prop_map(|(segments, query)| { + let path = format!("/{}", segments.join("/")); + match query { + Some(query) => format!("{path}?{query}"), + None => path, + } + }) +} + +fn exchange() -> impl Strategy { + let status = (200u16..600) + .prop_filter("bodiless status", |code| ![204, 205, 304].contains(code)) + .prop_map(|code| StatusCode::from_u16(code).unwrap()); + (method(), target(), message(), status, message()).prop_map( + |(method, target, request, status, response)| Exchange { + method, + target, + request, + status, + response, + }, + ) +} + +fn body(message: &Message) -> Body { + if !message.chunked { + return Body::from(message.body.clone()); + } + let chunks: Vec> = message + .body + .chunks(1000) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect(); + Body::from_stream(futures_util::stream::iter(chunks)) +} + +fn apply(headers: &mut HeaderMap, message: &Message) { + if message.chunked { + headers.insert("transfer-encoding", HeaderValue::from_static("chunked")); + } + for (name, values) in &message.headers { + for value in values { + headers.append( + HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_bytes(value).unwrap(), + ); + } + } +} + +fn assert_carried(headers: &HeaderMap, message: &Message, framing_and_hop: &[&str]) { + for (name, values) in &message.headers { + let carried: Vec<&[u8]> = headers + .get_all(name) + .iter() + .map(HeaderValue::as_bytes) + .collect(); + let sent: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect(); + assert_eq!(carried, sent, "header {name}"); + } + for name in headers.keys() { + assert!( + framing_and_hop.contains(&name.as_str()) + || message + .headers + .iter() + .any(|(sent, _)| sent == name.as_str()), + "unexpected header {name}" + ); + } +} + +/// Any request a client can send and any response Python can give cross the forwarder unchanged, +/// apart from the framing and peer headers the forwarder owns. +#[test] +fn every_exchange_crosses_the_forwarder_unchanged() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let current = Arc::new(Mutex::new(None::)); + let responding = current.clone(); + let (gateway, received) = runtime.block_on(async { + let (upstream, received) = recording(move || { + let exchange = responding.lock().unwrap().clone().unwrap(); + let mut response = Response::new(body(&exchange.response)); + *response.status_mut() = exchange.status; + apply(response.headers_mut(), &exchange.response); + response + }) + .await; + ( + Gateway::start(upstream).await, + tokio::sync::Mutex::new(received), + ) + }); + let client = client(); + TestRunner::new(ProptestConfig::with_cases(256)) + .run(&exchange(), |exchange| { + *current.lock().unwrap() = Some(exchange.clone()); + runtime.block_on(async { + while received.lock().await.try_recv().is_ok() {} + let mut request = Request::builder() + .method(exchange.method.clone()) + .uri(gateway.url(&exchange.target)) + .body(body(&exchange.request)) + .unwrap(); + apply(request.headers_mut(), &exchange.request); + let response = support::within(client.request(request)).await.unwrap(); + + let python = received.lock().await.recv().await.unwrap(); + assert_eq!(python.method, exchange.method); + assert_eq!(python.uri, exchange.target.as_str()); + assert_eq!(python.body, exchange.request.body); + assert_carried( + &python.headers, + &exchange.request, + &["host", "content-length", "transfer-encoding", PEER_HEADER], + ); + assert_eq!(python.headers.get_all(PEER_HEADER).iter().count(), 1); + + assert_eq!(response.status(), exchange.status); + let headers = response.headers().clone(); + assert_carried( + &headers, + &exchange.response, + &["date", "content-length", "transfer-encoding"], + ); + let body = support::within(response.into_body().collect()) + .await + .unwrap(); + assert_eq!(body.to_bytes(), exchange.response.body); + }); + Ok(()) + }) + .unwrap(); +} + +#[derive(Clone, Debug)] +enum Framing { + Unframed, + Length, + Chunked(usize), +} + +#[derive(Clone, Debug)] +struct WireRequest { + method: Method, + target: String, + headers: Vec<(String, Vec)>, + body: Vec, + framing: Framing, +} + +fn wire_request() -> impl Strategy { + let framing = prop_oneof![ + Just(Framing::Unframed), + Just(Framing::Length), + (1usize..2048).prop_map(Framing::Chunked), + ]; + ( + method(), + target(), + prop::collection::vec( + ("[xX]-[A-Za-z0-9!#$%&'*+.^_`|~-]{1,12}", header_value()), + 0..24, + ), + prop::collection::vec(any::(), 0..8 * 1024), + framing, + ) + .prop_map(|(method, target, headers, body, framing)| WireRequest { + method, + target, + headers, + body: match framing { + Framing::Unframed => Vec::new(), + _ => body, + }, + framing, + }) +} + +fn encode(request: &WireRequest) -> Vec { + let framing = match request.framing { + Framing::Unframed => String::new(), + Framing::Length => format!("Content-Length: {}\r\n", request.body.len()), + Framing::Chunked(_) => "Transfer-Encoding: chunked\r\n".to_owned(), + }; + let fields = request + .headers + .iter() + .flat_map(|(name, value)| [name.as_bytes(), b": ", value, b"\r\n"].concat()); + let body = match request.framing { + Framing::Chunked(size) => request + .body + .chunks(size) + .flat_map(|chunk| { + [format!("{:x}\r\n", chunk.len()).as_bytes(), chunk, b"\r\n"].concat() + }) + .chain(*b"0\r\n\r\n") + .collect(), + _ => request.body.clone(), + }; + [ + format!( + "{} {} HTTP/1.1\r\nHost: x\r\n", + request.method, request.target + ) + .into_bytes(), + fields.collect(), + framing.into_bytes(), + b"\r\n".to_vec(), + body, + ] + .concat() +} + +fn by_name<'a>(fields: impl Iterator) -> Vec<(String, &'a [u8])> { + let mut fields: Vec<_> = fields + .map(|(name, value)| (name.to_ascii_lowercase(), value)) + .collect(); + fields.sort_by(|(left, _), (right, _)| left.cmp(right)); + fields +} + +/// Pipelined requests in any framing, written in arbitrary pieces, reach Python one by one with +/// the bytes h11 would have read from the client directly. +#[test] +fn pipelined_wire_requests_reach_python_as_sent() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let (gateway, received) = runtime.block_on(async { + let (upstream, received) = recording(|| Response::new(Body::from("ok"))).await; + ( + Gateway::start(upstream).await, + tokio::sync::Mutex::new(received), + ) + }); + let requests = (prop::collection::vec(wire_request(), 1..5), 1usize..4096); + TestRunner::new(ProptestConfig::with_cases(256)) + .run(&requests, |(requests, piece)| { + runtime.block_on(async { + let mut received = received.lock().await; + while received.try_recv().is_ok() {} + let mut raw = gateway.connect().await; + let wire: Vec = requests.iter().flat_map(encode).collect(); + for bytes in wire.chunks(piece) { + raw.send(bytes).await; + } + for request in &requests { + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::OK, "{request:?}"); + assert_eq!(raw.body(2).await, b"ok"); + let python = within(received.recv()).await.unwrap(); + assert_eq!(python.method, request.method); + assert_eq!(python.uri, request.target.as_str()); + assert_eq!(python.body, request.body); + assert_eq!(python.headers["host"], "x"); + let forwarded = python.headers.iter().filter(|(name, _)| { + !["host", "content-length", "transfer-encoding", PEER_HEADER] + .contains(&name.as_str()) + }); + assert_eq!( + by_name(forwarded.map(|(name, value)| (name.as_str(), value.as_bytes()))), + by_name( + request + .headers + .iter() + .map(|(name, value)| (name.as_str(), value.as_slice())) + ) + ); + } + }); + Ok(()) + }) + .unwrap(); +} diff --git a/litellm-rust/crates/gateway-forward/tests/settings.rs b/litellm-rust/crates/gateway-forward/tests/settings.rs new file mode 100644 index 00000000000..ce0d95e723f --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/settings.rs @@ -0,0 +1,110 @@ +use std::{ffi::OsString, time::Duration}; + +use litellm_gateway_forward::{Error, Settings}; +use rstest::rstest; + +const COMPLETE: [(&str, &str); 5] = [ + ("LITELLM_FORWARD_UPSTREAM", "http://127.0.0.1:53211"), + ("LITELLM_FORWARD_HOST", "0.0.0.0"), + ("LITELLM_FORWARD_PORT", "4000"), + ("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "2"), + ("LITELLM_FORWARD_SHUTDOWN_DRAIN", "40"), +]; + +fn parse(variables: &[(&str, &str)]) -> Result, Error> { + Settings::from_vars( + variables + .iter() + .map(|(key, value)| (OsString::from(key), OsString::from(value))), + ) +} + +fn without(variable: &str) -> Vec<(&'static str, &'static str)> { + COMPLETE + .into_iter() + .filter(|(key, _)| *key != variable) + .collect() +} + +fn with<'a>(variable: &'a str, value: &'a str) -> Vec<(&'a str, &'a str)> { + without(variable) + .into_iter() + .chain([(variable, value)]) + .collect() +} + +#[rstest] +fn reads_every_variable() { + assert_eq!( + parse(&COMPLETE).unwrap(), + Some(Settings { + upstream: "http://127.0.0.1:53211".parse().unwrap(), + host: "0.0.0.0".into(), + port: 4000, + upstream_keepalive: Duration::from_secs(2), + shutdown_drain: Duration::from_secs(40), + }) + ); +} + +#[rstest] +fn without_an_upstream_the_gateway_is_not_forwarding() { + assert_eq!(parse(&without("LITELLM_FORWARD_UPSTREAM")).unwrap(), None); +} + +#[rstest] +#[case::host("LITELLM_FORWARD_HOST")] +#[case::port("LITELLM_FORWARD_PORT")] +#[case::upstream_keepalive("LITELLM_FORWARD_UPSTREAM_KEEPALIVE")] +#[case::shutdown_drain("LITELLM_FORWARD_SHUTDOWN_DRAIN")] +fn a_missing_variable_fails_naming_it(#[case] variable: &str) { + let error = parse(&without(variable)).unwrap_err(); + + assert!(error.to_string().contains(variable), "{error}"); +} + +#[rstest] +#[case::port_overflow("LITELLM_FORWARD_PORT", "65536")] +#[case::port_text("LITELLM_FORWARD_PORT", "http")] +#[case::keepalive_text("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "soon")] +#[case::keepalive_negative("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "-1")] +#[case::keepalive_zero("LITELLM_FORWARD_UPSTREAM_KEEPALIVE", "0")] +#[case::drain_text("LITELLM_FORWARD_SHUTDOWN_DRAIN", "soon")] +#[case::drain_negative("LITELLM_FORWARD_SHUTDOWN_DRAIN", "-1")] +#[case::drain_infinite("LITELLM_FORWARD_SHUTDOWN_DRAIN", "inf")] +fn an_invalid_variable_fails_naming_it(#[case] variable: &str, #[case] value: &str) { + let error = parse(&with(variable, value)).unwrap_err(); + + assert!(error.to_string().contains(variable), "{error}"); +} + +#[rstest] +fn fractional_drain_seconds_are_used_as_given() { + let settings = parse(&with("LITELLM_FORWARD_SHUTDOWN_DRAIN", "12.5")) + .unwrap() + .unwrap(); + + assert_eq!(settings.shutdown_drain, Duration::from_millis(12_500)); +} + +#[rstest] +#[case::keepalive("KEEPALIVE_TIMEOUT")] +#[case::graceful_shutdown("GRACEFUL_SHUTDOWN_TIMEOUT")] +#[case::host("HOST")] +#[case::port("PORT")] +fn python_settings_are_never_read(#[case] variable: &str) { + let mut variables = COMPLETE.to_vec(); + variables.push((variable, "9")); + + assert_eq!(parse(&variables).unwrap(), parse(&COMPLETE).unwrap()); +} + +#[rstest] +fn a_non_loopback_upstream_fails_startup() { + let error = parse(&with("LITELLM_FORWARD_UPSTREAM", "http://192.0.2.1:4000")).unwrap_err(); + + assert!( + error.to_string().contains("LITELLM_FORWARD_UPSTREAM"), + "{error}" + ); +} diff --git a/litellm-rust/crates/gateway-forward/tests/slow_head.rs b/litellm-rust/crates/gateway-forward/tests/slow_head.rs new file mode 100644 index 00000000000..764733aa524 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/slow_head.rs @@ -0,0 +1,41 @@ +mod support; + +use std::time::Duration; + +use axum::{body::Body, extract::Request, http::StatusCode}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, client, raw_upstream}; +use tokio::time::Instant; + +/// Reasoning models with `stream=false`, image and video generation, and model health checks send +/// no response headers until the provider finishes, which can take many minutes. +#[rstest] +#[tokio::test(start_paused = true)] +async fn waits_for_response_headers_as_long_as_python_takes() { + let thinking = Duration::from_secs(30 * 60); + let upstream = raw_upstream(move |mut python| async move { + python.request_head().await.unwrap(); + tokio::time::sleep(thinking).await; + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 5\r\n\r\nimage") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let started = Instant::now(); + + let response = client() + .request( + Request::post(gateway.url("/v1/images/generations")) + .body(Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert!(started.elapsed() >= thinking); + let body = response.into_body().collect().await.unwrap().to_bytes(); + assert_eq!(body, "image"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/stream.rs b/litellm-rust/crates/gateway-forward/tests/stream.rs new file mode 100644 index 00000000000..eaa12f20b77 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/stream.rs @@ -0,0 +1,313 @@ +mod support; + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{body::Body, extract::Request, http::StatusCode}; +use http_body_util::BodyExt; +use hyper::body::Incoming; +use rstest::rstest; +use support::{Gateway, client, raw_upstream, within}; +use tokio::{ + io::AsyncWriteExt, + sync::{Notify, mpsc}, +}; + +fn chunk(bytes: &[u8]) -> Vec { + [format!("{:x}\r\n", bytes.len()).as_bytes(), bytes, b"\r\n"].concat() +} + +const EVENTS: [&str; 3] = ["data: 1\n\n", "data: 2\n\n", "data: [DONE]\n\n"]; + +#[rstest] +#[tokio::test] +async fn each_event_reaches_the_client_before_python_produces_the_next( + #[values(true, false)] chunked: bool, +) { + let produce = Arc::new(Notify::new()); + let next = produce.clone(); + let upstream = raw_upstream(move |mut python| { + let next = next.clone(); + async move { + python.request_head().await.unwrap(); + let framing = match chunked { + true => "transfer-encoding: chunked".to_owned(), + false => format!("content-length: {}", EVENTS.concat().len()), + }; + python + .send(format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n{framing}\r\n\r\n" + )) + .await; + for event in EVENTS { + match chunked { + true => python.send(chunk(event.as_bytes())).await, + false => python.send(event).await, + } + next.notified().await; + } + if chunked { + python.send("0\r\n\r\n").await; + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut response = within( + client().request( + Request::post(gateway.url("/v1/chat/completions")) + .body(Body::from(r#"{"stream":true}"#)) + .unwrap(), + ), + ) + .await + .unwrap(); + + for event in EVENTS { + let frame = within(response.body_mut().frame()).await.unwrap().unwrap(); + assert_eq!(frame.into_data().unwrap(), event); + produce.notify_one(); + } + let rest = within(response.into_body().collect()).await.unwrap(); + assert!(rest.to_bytes().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn relays_a_binary_framed_stream_unchanged() { + let message: Vec = [ + &[ + 0x00, 0x00, 0x00, 0x3d, 0x00, 0x00, 0x00, 0x20, 0x8a, 0x0d, 0x0a, 0xff, + ][..], + b"\x0b:event-type\x07\x00\x05chunk\x0d:content-type\x07\x00\x10application/json", + b"0\r\n\r\n{\"bytes\":\"eyJ0eXBlIjoiY29udGVudCJ9\"}", + &[0xc3, 0x28, 0x00, 0x0d], + ] + .concat(); + let sent = message.clone(); + let upstream = raw_upstream(move |mut python| { + let sent = sent.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-type: application/vnd.amazon.eventstream\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + for piece in sent.chunks(7) { + python.send(chunk(piece)).await; + } + python.send("0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = within( + client().request( + Request::post(gateway.url("/bedrock/model/m/converse-stream")) + .body(Body::empty()) + .unwrap(), + ), + ) + .await + .unwrap(); + + assert_eq!( + response.headers()["content-type"], + "application/vnd.amazon.eventstream" + ); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, message); +} + +#[rstest] +#[tokio::test] +async fn relays_a_body_delimited_by_the_upstream_closing() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\nconnection: close\r\n\r\nuntil eof") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap()); + let body = within(response).await.unwrap().into_body().collect(); + assert_eq!(within(body).await.unwrap().to_bytes(), "until eof"); +} + +async fn drain(mut body: Incoming) -> (Vec, Result<(), hyper::Error>) { + let mut delivered = Vec::new(); + loop { + match within(body.frame()).await { + None => return (delivered, Ok(())), + Some(Err(error)) => return (delivered, Err(error)), + Some(Ok(frame)) => delivered.extend_from_slice(&frame.into_data().unwrap()), + } + } +} + +#[rstest] +#[case::after_the_head("transfer-encoding: chunked", "", "")] +#[case::between_events("transfer-encoding: chunked", "6\r\ndata: \r\n", "data: ")] +#[case::mid_event("transfer-encoding: chunked", "a\r\ndata:", "data:")] +#[case::short_of_its_content_length("content-length: 100", "data: 1\n\n", "data: 1\n\n")] +#[tokio::test] +async fn an_upstream_dying_mid_stream_fails_the_client_body_instead_of_ending_it( + #[case] framing: &'static str, + #[case] sent: &'static str, + #[case] delivered: &'static str, +) { + let upstream = raw_upstream(move |mut python| async move { + python.request_head().await.unwrap(); + python + .send(format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n{framing}\r\n\r\n{sent}" + )) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let response = + within(client().request(Request::get(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let (received, ended) = drain(response.into_body()).await; + assert_eq!(received, delivered.as_bytes()); + assert!(ended.is_err()); +} + +#[rstest] +#[tokio::test] +async fn response_headers_reach_the_client_before_python_sends_any_body() { + let release = Arc::new(Notify::new()); + let released = release.clone(); + let upstream = raw_upstream(move |mut python| { + let released = released.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream; charset=utf-8\r\nx-accel-buffering: no\r\nx-litellm-call-id: call-1\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + released.notified().await; + python.send(chunk(b"data: 1\n\n")).await; + python.send("0\r\n\r\n").await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + + let response = + within(client().request(Request::post(gateway.url("/")).body(Body::empty()).unwrap())) + .await + .unwrap(); + + assert_eq!(response.headers()["x-litellm-call-id"], "call-1"); + assert_eq!(response.headers()["x-accel-buffering"], "no"); + release.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, "data: 1\n\n"); +} + +/// A connection closed without any response looks retryable to client SDKs, which would re-run a +/// request Python already ran. Python dying right after its head races the relay, hence the loop. +#[rstest] +#[tokio::test] +async fn a_head_python_sent_right_before_dying_still_reaches_the_client() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n") + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + for _ in 0..20 { + let mut raw = gateway.connect().await; + raw.send("POST / HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n") + .await; + let head = within(raw.response_head()).await; + assert_eq!(head.map(|head| head.status), Some(StatusCode::OK)); + } +} + +#[rstest] +#[tokio::test] +async fn a_client_that_stops_reading_stalls_python_instead_of_being_buffered_for() { + let written = Arc::new(AtomicUsize::new(0)); + let counted = written.clone(); + let (failed, mut write_failed) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let counted = counted.clone(); + let failed = failed.clone(); + async move { + python.request_head().await.unwrap(); + let mut python = python.into_stream(); + python + .write_all(b"HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n") + .await + .unwrap(); + let piece = chunk(&[0x5a; 64 * 1024]); + while python.write_all(&piece).await.is_ok() { + counted.fetch_add(piece.len(), Ordering::SeqCst); + } + failed.send(()).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST / HTTP/1.1\r\nhost: x\r\ncontent-length: 0\r\n\r\n") + .await; + within(raw.response_head()).await.unwrap(); + + tokio::time::sleep(Duration::from_secs(1)).await; + let stalled = written.load(Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(250)).await; + + assert_eq!(written.load(Ordering::SeqCst), stalled); + assert!(stalled < 64 * 1024 * 1024, "{stalled} bytes in flight"); + drop(raw); + within(write_failed.recv()).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn an_http10_client_receives_a_streamed_body_delimited_by_close() { + let upstream = raw_upstream(|mut python| async move { + python.request_head().await.unwrap(); + python + .send([ + &b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n"[..], + &chunk(b"data: 1\n"), + &chunk(b"data: 2\n"), + b"0\r\n\r\n", + ] + .concat()) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET /v1/stream HTTP/1.0\r\nhost: x\r\n\r\n").await; + + let response = String::from_utf8(within(raw.until_closed()).await).unwrap(); + let (head, body) = response.split_once("\r\n\r\n").unwrap(); + assert!(head.starts_with("HTTP/1.0 200 OK"), "{head}"); + assert!( + !head.to_ascii_lowercase().contains("transfer-encoding"), + "{head}" + ); + assert_eq!(body, "data: 1\ndata: 2\n"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/support/mod.rs b/litellm-rust/crates/gateway-forward/tests/support/mod.rs new file mode 100644 index 00000000000..d39532a456e --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/support/mod.rs @@ -0,0 +1,311 @@ +#![allow( + dead_code, + reason = "each test binary uses a different subset of the harness" +)] + +use std::{future::Future, net::SocketAddr, time::Duration}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + http::{HeaderMap, Method, StatusCode, Uri}, +}; +use http_body_util::BodyExt; +use hyper_util::{ + client::legacy::{Client, connect::HttpConnector}, + rt::TokioExecutor, +}; +use litellm_gateway_forward::Forward; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::{mpsc, oneshot}, + task::JoinHandle, +}; + +pub const WAIT: Duration = Duration::from_secs(5); +pub const UVICORN_KEEPALIVE: Duration = Duration::from_secs(5); +pub const DRAIN: Duration = Duration::from_secs(1); + +pub async fn within(future: impl Future) -> T { + tokio::time::timeout(WAIT, future).await.expect("timed out") +} + +/// A running forwarder in front of `upstream`, listening on a real socket like production. +pub struct Gateway { + pub address: SocketAddr, + stop: oneshot::Sender<()>, + server: JoinHandle<()>, +} + +impl Gateway { + pub async fn start(upstream: SocketAddr) -> Self { + Self::with_keepalive(upstream, UVICORN_KEEPALIVE).await + } + + pub async fn with_keepalive(upstream: SocketAddr, keepalive: Duration) -> Self { + let forward = Forward::new(&format!("http://{upstream}").parse().unwrap(), keepalive); + let public = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = public.local_addr().unwrap(); + let (stop, stopped) = oneshot::channel::<()>(); + let server = tokio::spawn(forward.serve( + public, + async { + stopped.await.ok(); + }, + DRAIN, + )); + Self { + address, + stop, + server, + } + } + + pub fn url(&self, path_and_query: &str) -> Uri { + format!("http://{}{path_and_query}", self.address) + .parse() + .unwrap() + } + + pub async fn connect(&self) -> Raw { + Raw::new(TcpStream::connect(self.address).await.unwrap()) + } + + /// Starts graceful shutdown and returns the server task, which ends once in-flight work drains + /// or `DRAIN` passes. + pub fn shutdown(self) -> JoinHandle<()> { + self.stop.send(()).unwrap(); + self.server + } +} + +pub fn client() -> Client { + Client::builder(TokioExecutor::new()).build_http() +} + +/// What the upstream received, after the forwarder rewrote it. +pub struct Received { + pub method: Method, + pub uri: Uri, + pub headers: HeaderMap, + pub body: Bytes, +} + +/// An axum upstream that records every request it fully reads and answers with `respond`. +pub async fn recording( + respond: impl Fn() -> axum::response::Response + Clone + Send + Sync + 'static, +) -> (SocketAddr, mpsc::UnboundedReceiver) { + let (sender, receiver) = mpsc::unbounded_channel(); + let address = serve(Router::new().fallback(move |request: Request| { + let sender = sender.clone(); + let respond = respond.clone(); + async move { + let (parts, body) = request.into_parts(); + let body = body.collect().await.unwrap().to_bytes(); + sender + .send(Received { + method: parts.method, + uri: parts.uri, + headers: parts.headers, + body, + }) + .unwrap(); + respond() + } + })) + .await; + (address, receiver) +} + +pub async fn serve(router: Router) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + address +} + +/// An upstream that hands each accepted connection to `handle` as raw bytes, for behavior a +/// well-behaved HTTP server would never produce. +pub async fn raw_upstream(handle: F) -> SocketAddr +where + F: Fn(Raw) -> Fut + Send + Sync + 'static, + Fut: Future + Send + 'static, +{ + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.unwrap(); + tokio::spawn(handle(Raw::new(stream))); + } + }); + address +} + +pub fn unreachable_upstream() -> SocketAddr { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap() +} + +pub struct Head { + pub status: StatusCode, + pub headers: Vec<(String, Vec)>, + pub path: String, +} + +impl Head { + pub fn header(&self, name: &str) -> Option<&[u8]> { + self.headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_slice()) + } + + pub fn content_length(&self) -> Option { + self.header("content-length") + .map(|value| std::str::from_utf8(value).unwrap().parse().unwrap()) + } +} + +/// One TCP connection driven byte by byte, so tests control framing, timing and disconnects. +/// Bytes read past a message head stay buffered for the next read. +pub struct Raw { + stream: TcpStream, + buffer: Vec, +} + +impl Raw { + pub fn new(stream: TcpStream) -> Self { + Self { + stream, + buffer: Vec::new(), + } + } + + pub fn into_stream(self) -> TcpStream { + assert!(self.buffer.is_empty(), "buffered bytes would be lost"); + self.stream + } + + pub fn local_address(&self) -> SocketAddr { + self.stream.local_addr().unwrap() + } + + pub async fn send(&mut self, bytes: impl AsRef<[u8]>) { + self.stream.write_all(bytes.as_ref()).await.unwrap(); + } + + /// Reads one request head, or `None` when the peer closed the connection first. + pub async fn request_head(&mut self) -> Option { + self.head(|bytes| { + let mut headers = [httparse::EMPTY_HEADER; 256]; + let mut request = httparse::Request::new(&mut headers); + let httparse::Status::Complete(length) = request.parse(bytes).unwrap() else { + return None; + }; + Some(( + length, + Head { + status: StatusCode::OK, + headers: owned(request.headers), + path: request.path.unwrap().to_owned(), + }, + )) + }) + .await + } + + pub async fn response_head(&mut self) -> Option { + self.head(|bytes| { + let mut headers = [httparse::EMPTY_HEADER; 256]; + let mut response = httparse::Response::new(&mut headers); + let httparse::Status::Complete(length) = response.parse(bytes).unwrap() else { + return None; + }; + Some(( + length, + Head { + status: StatusCode::from_u16(response.code.unwrap()).unwrap(), + headers: owned(response.headers), + path: String::new(), + }, + )) + }) + .await + } + + async fn head(&mut self, parse: impl Fn(&[u8]) -> Option<(usize, Head)>) -> Option { + loop { + if let Some((length, head)) = parse(&self.buffer) { + self.buffer.drain(..length); + return Some(head); + } + if self.fill().await == 0 { + return None; + } + } + } + + /// Reads exactly `length` body bytes, using what is already buffered first. + pub async fn body(&mut self, length: usize) -> Vec { + while self.buffer.len() < length { + assert_ne!(self.fill().await, 0, "closed mid-body"); + } + self.buffer.drain(..length).collect() + } + + /// Reads a whole chunked body and returns its decoded bytes. + pub async fn chunked_body(&mut self) -> Vec { + let mut decoded = Vec::new(); + loop { + let size_line = self.line().await; + let size = + usize::from_str_radix(size_line.split(';').next().unwrap().trim(), 16).unwrap(); + if size == 0 { + while !self.line().await.is_empty() {} + return decoded; + } + decoded.extend(self.body(size).await); + assert_eq!(self.body(2).await, b"\r\n"); + } + } + + async fn line(&mut self) -> String { + loop { + if let Some(end) = self.buffer.windows(2).position(|window| window == b"\r\n") { + let line = String::from_utf8(self.buffer[..end].to_vec()).unwrap(); + self.buffer.drain(..end + 2); + return line; + } + assert_ne!(self.fill().await, 0, "closed mid-line"); + } + } + + async fn fill(&mut self) -> usize { + let mut chunk = [0; 8192]; + let read = self.stream.read(&mut chunk).await.unwrap_or(0); + self.buffer.extend_from_slice(&chunk[..read]); + read + } + + /// Everything the peer sends until it closes, including bytes already buffered. + pub async fn until_closed(&mut self) -> Vec { + while self.fill().await != 0 {} + std::mem::take(&mut self.buffer) + } + + /// Resolves once the peer closes the connection, discarding anything it still sends. + pub async fn closed(&mut self) { + let mut chunk = [0; 8192]; + while self.stream.read(&mut chunk).await.unwrap_or(0) != 0 {} + } +} + +fn owned(headers: &[httparse::Header<'_>]) -> Vec<(String, Vec)> { + headers + .iter() + .map(|header| (header.name.to_owned(), header.value.to_vec())) + .collect() +} diff --git a/litellm-rust/crates/gateway-forward/tests/upload.rs b/litellm-rust/crates/gateway-forward/tests/upload.rs new file mode 100644 index 00000000000..3dc865c9c86 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/upload.rs @@ -0,0 +1,339 @@ +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, Bytes}, + extract::Request, + http::StatusCode, + response::Response, +}; +use http_body_util::BodyExt; +use rstest::rstest; +use support::{Gateway, Raw, client, raw_upstream, recording, serve, within}; +use tokio::{ + io::AsyncWriteExt, + sync::{Notify, mpsc}, +}; + +fn ok() -> Response { + Response::new(Body::from("ok")) +} + +#[rstest] +#[tokio::test] +async fn streams_upload_and_response_incrementally_without_a_body_limit() { + let upstream = serve( + Router::new().fallback(|request: Request| async { Response::new(request.into_body()) }), + ) + .await; + let gateway = Gateway::start(upstream).await; + let (sender, receiver) = mpsc::channel::>(1); + let stream = futures_util::stream::unfold(receiver, |mut receiver| async { + receiver.recv().await.map(|chunk| (chunk, receiver)) + }); + sender.send(Ok(Bytes::from_static(b"first"))).await.unwrap(); + let mut response = within( + client().request( + Request::post(gateway.url("/v1/files")) + .body(Body::from_stream(stream)) + .unwrap(), + ), + ) + .await + .unwrap(); + let frame = within(response.body_mut().frame()).await.unwrap().unwrap(); + assert_eq!(frame.into_data().unwrap(), "first"); + let large = Bytes::from(vec![b'x'; 8 * 1024 * 1024]); + sender.send(Ok(large.clone())).await.unwrap(); + drop(sender); + let rest = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(rest, large); +} + +/// hyper sends the 100 as soon as the request is forwarded, where uvicorn waits for the app to +/// read the body, so an upload Python rejects on headers alone is still sent in full. +#[rstest] +#[tokio::test] +async fn answers_expect_continue_so_large_uploads_start() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("POST /v1/audio/transcriptions HTTP/1.1\r\nhost: x\r\nexpect: 100-continue\r\ncontent-length: 4\r\n\r\n") + .await; + let interim = within(raw.response_head()).await.unwrap(); + assert_eq!(interim.status, StatusCode::CONTINUE); + raw.send("data").await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, "data"); +} + +/// curl sends `expect: 100-continue` for any body over 1 MiB and waits for the 100 before sending +/// the multipart body, which must reach Python intact even where it looks like framing. +#[rstest] +#[tokio::test] +async fn expect_continue_delivers_a_large_multipart_upload_byte_for_byte() { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let boundary = "------------------------d74496d66958873e"; + let audio: Vec = (0..3 * 1024 * 1024) + .map(|index: u32| (index.wrapping_mul(2_654_435_761) >> 24) as u8) + .chain(*b"\r\n0\r\n\r\nPOST /smuggled HTTP/1.1\r\n\r\n--") + .collect(); + let body = [ + format!("--{boundary}\r\ncontent-disposition: form-data; name=\"model\"\r\n\r\nwhisper-1\r\n--{boundary}\r\ncontent-disposition: form-data; name=\"file\"; filename=\"a b.wav\"\r\ncontent-type: audio/wav\r\n\r\n").into_bytes(), + audio, + format!("\r\n--{boundary}--\r\n").into_bytes(), + ] + .concat(); + let content_type = format!("multipart/form-data; boundary={boundary}"); + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/audio/transcriptions HTTP/1.1\r\nhost: x\r\ncontent-type: {content_type}\r\nexpect: 100-continue\r\ncontent-length: {}\r\n\r\n", + body.len() + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::CONTINUE + ); + raw.send(&body).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = received.recv().await.unwrap(); + assert_eq!(request.headers["content-type"], content_type.as_str()); + assert_eq!(request.body, body); +} + +#[rstest] +#[case::get("GET")] +#[case::delete("DELETE")] +#[case::options("OPTIONS")] +#[case::post("POST")] +#[tokio::test] +async fn forwards_a_chunked_body_on_any_method(#[case] method: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "{method} /v1/search HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n4\r\nbody\r\n0\r\n\r\n" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + assert_eq!(received.recv().await.unwrap().body, "body"); +} + +#[rstest] +#[case::no_framing("", "", &[], "")] +#[case::empty_content_length("content-length: 0\r\n", "", &[("content-length", "0")], "")] +#[case::empty_chunked( + "transfer-encoding: chunked\r\n", + "0\r\n\r\n", + &[("transfer-encoding", "chunked")], + "" +)] +#[case::sized( + "content-length: 7\r\n", + "payload", + &[("content-length", "7")], + "payload" +)] +#[case::chunk_extensions_and_trailers( + "transfer-encoding: chunked\r\n", + "4;name=\"v\"\r\nbody\r\n0\r\nx-checksum: 1\r\n\r\n", + &[("transfer-encoding", "chunked")], + "body" +)] +#[tokio::test] +async fn a_body_reaches_python_with_the_framing_the_client_chose( + #[values("POST", "GET")] method: &str, + #[case] framing: &str, + #[case] wire_body: &str, + #[case] python_framing: &[(&str, &str)], + #[case] python_body: &str, +) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "{method} /v1/files HTTP/1.1\r\nhost: x\r\n{framing}\r\n{wire_body}" + )) + .await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + + let request = received.recv().await.unwrap(); + let forwarded: Vec<(&str, &str)> = ["content-length", "transfer-encoding"] + .into_iter() + .flat_map(|name| { + request + .headers + .get_all(name) + .iter() + .map(move |value| (name, value.to_str().unwrap())) + }) + .collect(); + assert_eq!(request.method, method); + assert_eq!(forwarded, python_framing); + assert_eq!(request.body, python_body); +} + +#[rstest] +#[tokio::test] +async fn a_large_upload_is_paced_by_python_reading_it() { + const LENGTH: usize = 128 * 1024 * 1024; + let start_reading = Arc::new(Notify::new()); + let reading = start_reading.clone(); + let upstream = raw_upstream(move |mut python| { + let reading = reading.clone(); + async move { + let head = python.request_head().await.unwrap(); + reading.notified().await; + let mut remaining = head.content_length().unwrap(); + while remaining > 0 { + remaining -= python.body(remaining.min(1 << 20)).await.len(); + } + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") + .await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let (read, mut write) = gateway.connect().await.into_stream().into_split(); + let sent = Arc::new(AtomicUsize::new(0)); + let sending = sent.clone(); + let upload = tokio::spawn(async move { + write + .write_all( + format!("POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: {LENGTH}\r\n\r\n") + .as_bytes(), + ) + .await + .unwrap(); + let chunk = [0x5a; 64 * 1024]; + while sending.load(Ordering::SeqCst) < LENGTH { + write.write_all(&chunk).await.unwrap(); + sending.fetch_add(chunk.len(), Ordering::SeqCst); + } + write + }); + + tokio::time::sleep(Duration::from_secs(1)).await; + assert!( + sent.load(Ordering::SeqCst) < LENGTH / 4, + "the forwarder buffered an upload Python was not reading" + ); + start_reading.notify_one(); + let write = tokio::time::timeout(Duration::from_secs(60), upload) + .await + .unwrap() + .unwrap(); + let mut raw = Raw::new(read.reunite(write).unwrap()); + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); +} + +#[rstest] +#[case::between_chunks("5\r\nhello\r\n")] +#[case::mid_chunk("100\r\nhello")] +#[tokio::test] +async fn a_client_abandoning_a_chunked_upload_never_ends_python_body_cleanly( + #[case] sent: &'static str, +) { + let (forwarded, mut received) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let forwarded = forwarded.clone(); + async move { + python.request_head().await.unwrap(); + forwarded.send(python.until_closed().await).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ntransfer-encoding: chunked\r\n\r\n{sent}" + )) + .await; + tokio::time::sleep(Duration::from_millis(100)).await; + drop(raw); + + let body = within(received.recv()).await.unwrap(); + assert!(body.starts_with(b"5\r\nhello\r\n")); + assert!( + !body.ends_with(b"0\r\n\r\n"), + "python saw a truncated upload end cleanly" + ); +} + +#[rstest] +#[tokio::test] +async fn an_early_answer_during_an_upload_keeps_both_connections_in_sync() { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python| { + let event = event.clone(); + async move { + while let Some(head) = python.request_head().await { + let length = head.content_length(); + event.send(head.path.into_bytes()).unwrap(); + let Some(length) = length else { + python + .send("HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nnext") + .await; + continue; + }; + python + .send("HTTP/1.1 413 Payload Too Large\r\ncontent-length: 2\r\n\r\n{}") + .await; + event.send(python.body(length).await).unwrap(); + } + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let upload = b"GET /smuggled HTTP/1.1\r\nhost: x\r\n\r\n".repeat(1024); + let (first_half, second_half) = upload.split_at(upload.len() / 2); + let mut raw = gateway.connect().await; + raw.send(format!( + "POST /v1/files HTTP/1.1\r\nhost: x\r\ncontent-length: {}\r\n\r\n", + upload.len() + )) + .await; + raw.send(first_half).await; + + let rejected = within(raw.response_head()).await.unwrap(); + assert_eq!(rejected.status, StatusCode::PAYLOAD_TOO_LARGE); + assert_eq!(raw.body(2).await, b"{}"); + raw.send(second_half).await; + raw.send("GET /next HTTP/1.1\r\nhost: x\r\n\r\n").await; + let next = within(raw.response_head()).await.unwrap(); + assert_eq!(raw.body(next.content_length().unwrap()).await, b"next"); + + assert_eq!(within(events.recv()).await.unwrap(), b"/v1/files"); + assert_eq!(within(events.recv()).await.unwrap(), upload); + assert_eq!(within(events.recv()).await.unwrap(), b"/next"); +} diff --git a/litellm-rust/crates/gateway-forward/tests/upstream.rs b/litellm-rust/crates/gateway-forward/tests/upstream.rs new file mode 100644 index 00000000000..25e4768a54e --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/upstream.rs @@ -0,0 +1,21 @@ +use litellm_gateway_forward::Upstream; +use rstest::rstest; + +#[rstest] +#[case::https("https://127.0.0.1:4000")] +#[case::public("http://192.0.2.1:4000")] +#[case::hostname("http://localhost:4000")] +#[case::ipv6("http://[::1]:4000")] +#[case::path("http://127.0.0.1:4000/api")] +#[case::query("http://127.0.0.1:4000/?x=y")] +#[case::credentials("http://user@127.0.0.1:4000")] +fn rejects_anything_but_a_loopback_http_origin(#[case] upstream: &str) { + assert!(upstream.parse::().is_err()); +} + +#[rstest] +#[case::standard("http://127.0.0.1:4000")] +#[case::loopback_range("http://127.0.0.2:4000/")] +fn accepts_a_loopback_http_origin(#[case] upstream: &str) { + assert!(upstream.parse::().is_ok()); +} diff --git a/litellm-rust/crates/gateway-forward/tests/websocket.rs b/litellm-rust/crates/gateway-forward/tests/websocket.rs new file mode 100644 index 00000000000..ff64ddf38d3 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/tests/websocket.rs @@ -0,0 +1,719 @@ +mod support; + +use std::{ + net::SocketAddr, + sync::Arc, + time::{Duration, Instant}, +}; + +use axum::{ + Router, + http::{HeaderValue, StatusCode}, + response::IntoResponse, +}; +use futures_util::{SinkExt, StreamExt, future::join_all}; +use litellm_gateway_forward::PEER_HEADER; +use rstest::rstest; +use support::{DRAIN, Gateway, Head, Raw, raw_upstream, serve, within}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, + sync::{Barrier, mpsc}, +}; +use tokio_tungstenite::{ + WebSocketStream, accept_hdr_async, client_async, + tungstenite::{ + Error, Message, + client::IntoClientRequest, + handshake::{ + derive_accept_key, + server::{ErrorResponse, Request as Handshake, Response as Accepted}, + }, + protocol::{CloseFrame, frame::coding::CloseCode}, + }, +}; + +const KEY: &str = "dGhlIHNhbXBsZSBub25jZQ=="; + +async fn connect( + address: SocketAddr, + connection: &'static str, +) -> Result, Box> { + let mut request = format!("ws://{address}/v1/realtime?model=m") + .into_client_request() + .unwrap(); + let headers = request.headers_mut(); + headers.insert("authorization", HeaderValue::from_static("Bearer test")); + headers.insert( + "sec-websocket-protocol", + HeaderValue::from_static("realtime"), + ); + headers.insert("connection", HeaderValue::from_static(connection)); + let stream = TcpStream::connect(address).await.unwrap(); + let (socket, response) = client_async(request, stream).await.map_err(Box::new)?; + assert_eq!(response.headers()["sec-websocket-protocol"], "realtime"); + Ok(socket) +} + +/// Greets, echoes every message, and closes with 4000 when told to "bye". +#[expect( + clippy::result_large_err, + reason = "tungstenite's handshake callback returns ErrorResponse" +)] +async fn echo() -> SocketAddr { + raw_upstream(|python| async move { + let callback = |request: &Handshake, mut response: Accepted| { + assert_eq!(request.uri(), "/v1/realtime?model=m"); + assert_eq!(request.headers()["authorization"], "Bearer test"); + assert!(request.headers().contains_key(PEER_HEADER)); + response.headers_mut().insert( + "sec-websocket-protocol", + HeaderValue::from_static("realtime"), + ); + Ok::<_, ErrorResponse>(response) + }; + let mut socket = accept_hdr_async(python.into_stream(), callback) + .await + .unwrap(); + socket.send(Message::text("greeting")).await.unwrap(); + while let Some(Ok(message)) = socket.next().await { + if message == Message::text("bye") { + socket + .close(Some(CloseFrame { + code: CloseCode::Library(4000), + reason: "done".into(), + })) + .await + .unwrap(); + continue; + } + if message.is_close() { + break; + } + socket.send(message).await.unwrap(); + } + }) + .await +} + +fn handshake(target: &str, headers: &str) -> String { + format!( + "GET {target} HTTP/1.1\r\nhost: public.example\r\nconnection: Upgrade\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nsec-websocket-key: {KEY}\r\n{headers}\r\n" + ) +} + +/// The 101 uvicorn writes for `handshake`, with `headers` added. +fn switching(handshake: &Head, headers: &str) -> String { + let accept = derive_accept_key(handshake.header("sec-websocket-key").unwrap()); + format!( + "HTTP/1.1 101 Switching Protocols\r\nupgrade: websocket\r\nconnection: Upgrade\r\nsec-websocket-accept: {accept}\r\n{headers}\r\n" + ) +} + +fn values(head: &Head, name: &str) -> Vec { + head.headers + .iter() + .filter(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| String::from_utf8(value.clone()).unwrap()) + .collect() +} + +/// Accepts with `headers` in the 101, hands the handshake head to the test, then echoes bytes. +async fn byte_echo(headers: &'static str) -> (SocketAddr, mpsc::UnboundedReceiver) { + let (sender, handshakes) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let sender = sender.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, headers)).await; + sender.send(handshake).unwrap(); + echo_bytes(python).await; + } + }) + .await; + (upstream, handshakes) +} + +async fn echo_bytes(python: Raw) { + let mut stream = python.into_stream(); + let (mut read, mut write) = stream.split(); + tokio::io::copy(&mut read, &mut write).await.ok(); +} + +async fn assert_echoes(socket: &mut TcpStream, message: &[u8]) { + socket.write_all(message).await.unwrap(); + let mut echoed = vec![0; message.len()]; + within(socket.read_exact(&mut echoed)).await.unwrap(); + assert_eq!(echoed, message); +} + +async fn open(gateway: &Gateway, handshake: String) -> (TcpStream, Head) { + let mut raw = gateway.connect().await; + raw.send(handshake).await; + let switched = within(raw.response_head()).await.unwrap(); + assert_eq!(switched.status, StatusCode::SWITCHING_PROTOCOLS); + (raw.into_stream(), switched) +} + +#[rstest] +#[tokio::test] +async fn relays_messages_both_ways( + #[values( + "Upgrade", + "keep-alive, Upgrade", + "upgrade", + "Upgrade, Sec-WebSocket-Key, Sec-WebSocket-Version" + )] + connection: &'static str, + #[values( + Message::text("hello"), + Message::binary(vec![0, 1, 2]), + Message::binary(vec![7; 4 * 1024 * 1024]) + )] + message: Message, +) { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, connection)).await.unwrap(); + assert_eq!( + socket.next().await.unwrap().unwrap(), + Message::text("greeting") + ); + socket.send(message.clone()).await.unwrap(); + assert_eq!(within(socket.next()).await.unwrap().unwrap(), message); + socket.close(None).await.unwrap(); + while socket.next().await.is_some() {} + drop(socket); + within(gateway.shutdown()).await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn relays_a_close_python_initiates() { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + socket.send(Message::text("bye")).await.unwrap(); + let Message::Close(Some(frame)) = within(socket.next()).await.unwrap().unwrap() else { + panic!("expected python's close frame"); + }; + assert_eq!(frame.code, CloseCode::Library(4000)); + assert_eq!(frame.reason, "done"); +} + +#[rstest] +#[tokio::test] +async fn passes_an_upstream_rejection_through() { + let upstream = serve( + Router::new() + .fallback(|| async { (StatusCode::FORBIDDEN, "origin not allowed").into_response() }), + ) + .await; + let gateway = Gateway::start(upstream).await; + let error = within(connect(gateway.address, "Upgrade")) + .await + .unwrap_err(); + let Error::Http(response) = *error else { + panic!("expected the upstream rejection, got {error}"); + }; + assert_eq!(response.status(), StatusCode::FORBIDDEN); + assert_eq!(response.body().as_deref(), Some(&b"origin not allowed"[..])); +} + +/// The idle timeout only covers a connection waiting for its next HTTP request, and a realtime +/// session may sit silent between turns. +#[rstest] +#[tokio::test] +async fn keeps_a_quiet_socket_open_past_the_keepalive() { + let gateway = Gateway::with_keepalive(echo().await, Duration::from_millis(400)).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + tokio::time::sleep(Duration::from_millis(1200)).await; + socket.send(Message::text("still here")).await.unwrap(); + assert_eq!( + within(socket.next()).await.unwrap().unwrap(), + Message::text("still here") + ); +} + +/// Browsers carry the API key in a subprotocol, Python accepts with the first one it was offered, +/// and uvicorn negotiates permessage-deflate, so every handshake header must cross unchanged. +#[rstest] +#[tokio::test] +async fn handshake_headers_cross_unchanged_both_ways() { + let (upstream, mut handshakes) = byte_echo( + "sec-websocket-protocol: realtime\r\nsec-websocket-extensions: permessage-deflate; server_max_window_bits=12; client_max_window_bits=12\r\nset-cookie: a=1\r\nset-cookie: b=2\r\nx-litellm-call-id: abc\r\n", + ) + .await; + let gateway = Gateway::start(upstream).await; + let target = "/v1/realtime?model=gpt-realtime&api-key=sk%2Fa%2Bb%3D&dup=1&dup=2"; + let (_socket, switched) = open( + &gateway, + handshake( + target, + "sec-websocket-protocol: realtime, openai-insecure-api-key.sk-websocket-test, openai-beta.realtime-v1\r\nsec-websocket-protocol: second-line\r\nsec-websocket-extensions: permessage-deflate; client_max_window_bits\r\ncookie: token=abc\r\n", + ), + ) + .await; + let handshake = within(handshakes.recv()).await.unwrap(); + + assert_eq!(handshake.path, target); + assert_eq!( + values(&handshake, "sec-websocket-protocol"), + [ + "realtime, openai-insecure-api-key.sk-websocket-test, openai-beta.realtime-v1", + "second-line" + ] + ); + assert_eq!( + values(&handshake, "sec-websocket-extensions"), + ["permessage-deflate; client_max_window_bits"] + ); + assert_eq!(values(&handshake, "cookie"), ["token=abc"]); + assert_eq!(values(&switched, "sec-websocket-protocol"), ["realtime"]); + assert_eq!( + values(&switched, "sec-websocket-extensions"), + ["permessage-deflate; server_max_window_bits=12; client_max_window_bits=12"] + ); + assert_eq!(values(&switched, "set-cookie"), ["a=1", "b=2"]); + assert_eq!(values(&switched, "x-litellm-call-id"), ["abc"]); +} + +#[rstest] +#[tokio::test] +async fn invents_no_subprotocol_or_extension() { + let (upstream, mut handshakes) = byte_echo("").await; + let gateway = Gateway::start(upstream).await; + let (_socket, switched) = open(&gateway, handshake("/v1/realtime", "")).await; + let handshake = within(handshakes.recv()).await.unwrap(); + for name in ["sec-websocket-protocol", "sec-websocket-extensions"] { + assert!(handshake.header(name).is_none()); + assert!(switched.header(name).is_none()); + } +} + +#[rstest] +#[case::capitalized_upgrade("Upgrade: WebSocket\r\nConnection: keep-alive, Upgrade\r\n")] +#[case::uppercase("UPGRADE: WEBSOCKET\r\nCONNECTION: UPGRADE\r\n")] +#[case::split_connection("Upgrade: websocket\r\nConnection: keep-alive\r\nConnection: Upgrade\r\n")] +#[tokio::test] +async fn accepts_any_spelling_of_the_upgrade_headers(#[case] upgrade: &str) { + let (upstream, _handshakes) = byte_echo("").await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open( + &gateway, + format!( + "GET /v1/realtime HTTP/1.1\r\nhost: x\r\n{upgrade}sec-websocket-version: 13\r\nsec-websocket-key: {KEY}\r\n\r\n" + ), + ) + .await; + socket.write_all(b"ping").await.unwrap(); + let mut echoed = [0; 4]; + within(socket.read_exact(&mut echoed)).await.unwrap(); + assert_eq!(&echoed, b"ping"); +} + +/// Compressed frames set RSV1 and every client frame is masked, so the socket is relayed as bytes, +/// never re-framed. +#[rstest] +#[tokio::test] +async fn relays_socket_bytes_exactly() { + let (upstream, _handshakes) = byte_echo("").await; + let gateway = Gateway::start(upstream).await; + let (socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let sent: Vec = (0..3_000_017u32) + .map(|index| index.wrapping_mul(2_654_435_761).to_be_bytes()[1]) + .collect(); + let (mut read, mut write) = socket.into_split(); + let writing = sent.clone(); + let writer = tokio::spawn(async move { + for chunk in writing.chunks(7919) { + write.write_all(chunk).await.unwrap(); + } + write + }); + let mut echoed = vec![0; sent.len()]; + within(read.read_exact(&mut echoed)).await.unwrap(); + assert!(echoed == sent); + drop(writer.await.unwrap()); +} + +/// uvicorn writes the 101, an error event and a close frame back to back and then closes, so they +/// share a segment with the FIN, and a client may send its first frame with the handshake. +#[rstest] +#[tokio::test] +async fn keeps_bytes_that_arrive_with_the_handshake() { + let (early, mut received) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let early = early.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + let mut segment = switching(&handshake, "").into_bytes(); + segment.extend(b"\x81\x05error\x88\x08\x03\xf0denied"); + python.send(segment).await; + early.send(within(python.body(6)).await).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(handshake("/v1/realtime", "") + "early!").await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::SWITCHING_PROTOCOLS + ); + assert_eq!(within(received.recv()).await.unwrap(), b"early!"); + assert_eq!( + within(raw.until_closed()).await, + b"\x81\x05error\x88\x08\x03\xf0denied" + ); +} + +/// Python frees a realtime session's budget and parallel slot when its socket closes, so a client +/// that disappears without a close frame must still close Python's socket. +#[rstest] +#[case::fin(false)] +#[case::reset(true)] +#[tokio::test] +async fn a_vanished_client_closes_the_python_socket(#[case] reset: bool) { + let (event, mut events) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let event = event.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + event.send("open").unwrap(); + python.closed().await; + event.send("closed").unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let (socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + assert_eq!(within(events.recv()).await, Some("open")); + if reset { + #[expect(deprecated, reason = "a zero linger is how a test sends a TCP reset")] + socket.set_linger(Some(Duration::ZERO)).unwrap(); + } + drop(socket); + assert_eq!(within(events.recv()).await, Some("closed")); +} + +#[rstest] +#[case::fin(false)] +#[case::reset(true)] +#[tokio::test] +async fn a_vanished_python_ends_the_client_socket(#[case] reset: bool) { + let upstream = raw_upstream(move |mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + let socket = python.into_stream(); + if reset { + #[expect(deprecated, reason = "a zero linger is how a test sends a TCP reset")] + socket.set_linger(Some(Duration::ZERO)).unwrap(); + } + drop(socket); + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let mut rest = Vec::new(); + assert!(matches!( + within(socket.read_to_end(&mut rest)).await, + Ok(0) | Err(_) + )); +} + +#[rstest] +#[tokio::test] +async fn shutdown_closes_open_sockets_after_drain() { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + let started = Instant::now(); + within(gateway.shutdown()).await.unwrap(); + assert!(started.elapsed() >= DRAIN); + assert!(!matches!(socket.next().await, Some(Ok(Message::Text(_))))); +} + +#[rstest] +#[tokio::test] +async fn shutdown_relays_the_close_python_sends_while_draining() { + let gateway = Gateway::start(echo().await).await; + let mut socket = within(connect(gateway.address, "Upgrade")).await.unwrap(); + socket.next().await.unwrap().unwrap(); + let server = gateway.shutdown(); + tokio::time::sleep(Duration::from_millis(100)).await; + socket.send(Message::text("bye")).await.unwrap(); + let Message::Close(Some(frame)) = within(socket.next()).await.unwrap().unwrap() else { + panic!("expected python's close frame"); + }; + assert_eq!(frame.code, CloseCode::Library(4000)); + while socket.next().await.is_some() {} + drop(socket); + let started = Instant::now(); + within(server).await.unwrap(); + assert!(started.elapsed() < DRAIN / 2); +} + +#[rstest] +#[tokio::test] +async fn refuses_new_sockets_while_draining() { + let gateway = Gateway::start(echo().await).await; + let mut open = within(connect(gateway.address, "Upgrade")).await.unwrap(); + open.next().await.unwrap().unwrap(); + let address = gateway.address; + let server = gateway.shutdown(); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!( + TcpStream::connect(address).await.is_err() || connect(address, "Upgrade").await.is_err() + ); + open.send(Message::text("still open")).await.unwrap(); + assert_eq!( + within(open.next()).await.unwrap().unwrap(), + Message::text("still open") + ); + drop(open); + within(server).await.unwrap(); +} + +/// uvicorn never half-closes: once Python closes the socket nothing more can cross it, so the +/// relay must end instead of waiting for the client to close its side too. +#[rstest] +#[tokio::test] +async fn relay_ends_when_python_closes_while_the_client_holds_its_socket() { + let upstream = raw_upstream(move |mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + assert_eq!(within(socket.read(&mut [0; 1])).await.unwrap(), 0); + let started = Instant::now(); + within(gateway.shutdown()).await.unwrap(); + assert!(started.elapsed() < DRAIN / 2); + drop(socket); +} + +/// uvicorn answers a handshake Python closes before accepting with a 403, and one Python denies +/// through `websocket.http.response` with any status, headers and body, then drops the socket. The +/// client gets that response whole and keeps a connection that still reaches Python. +#[rstest] +#[case::closed_before_accept("403 Forbidden", None, "")] +#[case::denial_response( + "401 Unauthorized", + Some("Bearer realm=\"litellm\""), + "{\"error\":\"invalid api key\"}" +)] +#[tokio::test] +async fn a_denied_handshake_reaches_the_client_whole( + #[case] status: &'static str, + #[case] authenticate: Option<&'static str>, + #[case] body: &'static str, +) { + let upstream = raw_upstream(move |mut python: Raw| async move { + python.request_head().await.unwrap(); + let authenticate = authenticate + .map(|value| format!("www-authenticate: {value}\r\n")) + .unwrap_or_default(); + python + .send(format!( + "HTTP/1.1 {status}\r\ndate: Tue, 03 Nov 2026 08:00:00 GMT\r\nconnection: close\r\ncontent-length: {}\r\ncontent-type: text/plain; charset=utf-8\r\n{authenticate}\r\n{body}", + body.len() + )) + .await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + for _ in 0..2 { + raw.send(handshake("/v1/realtime", "")).await; + let denied = within(raw.response_head()).await.unwrap(); + assert_eq!(denied.status.as_str(), &status[..3]); + assert_eq!( + values(&denied, "www-authenticate"), + Vec::from_iter(authenticate.map(String::from)) + ); + assert!(denied.header("sec-websocket-accept").is_none()); + let length = denied.content_length().unwrap(); + assert_eq!(within(raw.body(length)).await, body.as_bytes()); + } +} + +/// Only Python's EOF ends a relay, so a client that half-closes after its close frame still gets +/// everything Python sends before closing its own side. +#[rstest] +#[tokio::test] +async fn a_client_half_close_still_receives_what_python_sends_after_it() { + let upstream = raw_upstream(|mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + let received = python.until_closed().await; + python.send(received).await; + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let close = b"\x88\x82\x01\x02\x03\x04\x02\xea"; + socket.write_all(close).await.unwrap(); + socket.shutdown().await.unwrap(); + let mut replied = Vec::new(); + within(socket.read_to_end(&mut replied)).await.unwrap(); + assert_eq!(replied, close); +} + +/// Python authenticates every handshake before accepting it, so sessions opened together must all +/// reach Python before any of them is accepted, and each must stay on its own socket. +#[rstest] +#[tokio::test] +async fn opens_many_sockets_at_once_without_crossing_them() { + const SOCKETS: usize = 64; + let everyone = Arc::new(Barrier::new(SOCKETS)); + let upstream = raw_upstream(move |mut python: Raw| { + let everyone = everyone.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + everyone.wait().await; + python.send(switching(&handshake, "")).await; + echo_bytes(python).await; + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let gateway = &gateway; + within(join_all((0..SOCKETS).map(|index| async move { + let (mut socket, _) = open(gateway, handshake("/v1/realtime", "")).await; + assert_echoes(&mut socket, format!("socket {index}").as_bytes()).await; + }))) + .await; +} + +/// Python resolves the key, budget and model before it accepts, and nothing bounds how long that +/// takes, so the idle timeout must not cut a handshake waiting on its 101. +#[rstest] +#[tokio::test] +async fn waits_for_a_101_slower_than_the_keepalive() { + let upstream = raw_upstream(|mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + tokio::time::sleep(Duration::from_millis(1200)).await; + python.send(switching(&handshake, "")).await; + echo_bytes(python).await; + }) + .await; + let gateway = Gateway::with_keepalive(upstream, Duration::from_millis(400)).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + assert_echoes(&mut socket, b"late").await; +} + +/// A 101 the client did not ask for, or to a protocol other than the one it asked for, would hand +/// the client a socket it cannot speak, so it is a 502 and Python's socket is closed. +#[rstest] +#[case::plain_request("GET /v1/models HTTP/1.1\r\nhost: x\r\n\r\n".to_owned(), "websocket")] +#[case::other_protocol(handshake("/v1/realtime", ""), "h2c")] +#[tokio::test] +async fn an_unrequested_switch_is_a_502_that_closes_python( + #[case] request: String, + #[case] protocol: &'static str, +) { + let (closed, mut closes) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let closed = closed.clone(); + async move { + python.request_head().await.unwrap(); + python + .send(format!( + "HTTP/1.1 101 Switching Protocols\r\nconnection: upgrade\r\nupgrade: {protocol}\r\n\r\n" + )) + .await; + python.closed().await; + closed.send(()).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send(request).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::BAD_GATEWAY + ); + within(closes.recv()).await.unwrap(); +} + +/// A realtime session streams audio for as long as the client reads it, so a client that stops +/// reading must stall Python's writes instead of piling them up in the gateway. +#[rstest] +#[tokio::test] +async fn a_client_that_stops_reading_stalls_python() { + let (sent, mut totals) = mpsc::unbounded_channel(); + let upstream = raw_upstream(move |mut python: Raw| { + let sent = sent.clone(); + async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + python.body(2).await; + let mut stream = python.into_stream(); + let chunk = vec![0x82; 64 * 1024]; + let mut total = 0; + while let Ok(Ok(())) = + tokio::time::timeout(Duration::from_millis(300), stream.write_all(&chunk)).await + { + total += chunk.len(); + } + sent.send(total).unwrap(); + } + }) + .await; + let gateway = Gateway::start(upstream).await; + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + socket.write_all(b"go").await.unwrap(); + let total = within(totals.recv()).await.unwrap(); + assert!( + total < 64 << 20, + "python wrote {total} bytes into a stalled socket" + ); +} + +/// Python closes a session that is still streaming audio by sending an error event and a close +/// frame and dropping its socket with the client's next frames unread, which resets it. Those last +/// bytes reached the client when it talked to uvicorn directly, so they must reach it through the +/// relay too, even though writing the client's frames to Python now fails. +#[rstest] +#[tokio::test] +async fn python_closing_mid_stream_still_delivers_its_last_frames() { + const LAST: &[u8] = b"\x81\x05error\x88\x02\x03\xf0"; + let upstream = raw_upstream(|mut python: Raw| async move { + let handshake = python.request_head().await.unwrap(); + python.send(switching(&handshake, "")).await; + python.body(1).await; + tokio::time::sleep(Duration::from_millis(20)).await; + python.send(LAST).await; + }) + .await; + let gateway = Gateway::start(upstream).await; + for _ in 0..20 { + let (mut socket, _) = open(&gateway, handshake("/v1/realtime", "")).await; + let (mut read, mut write) = socket.split(); + let streaming = async { + while write.write_all(b"\x82\x80\0\0\0\0").await.is_ok() { + tokio::time::sleep(Duration::from_micros(200)).await; + } + std::future::pending().await + }; + let mut received = Vec::new(); + let receiving = async { + let mut chunk = [0; 4096]; + while let Ok(read @ 1..) = read.read(&mut chunk).await { + received.extend_from_slice(&chunk[..read]); + } + }; + within(async { + tokio::select! { + () = streaming => {} + () = receiving => {} + } + }) + .await; + assert_eq!(received, LAST); + } +} diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 6eeaa43e013..d0cedbb3b36 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -12,6 +12,7 @@ envy = "0.4.2" http-body-util = "0.1" litellm-inference.workspace = true litellm-gateway-inference.workspace = true +litellm-gateway-forward.workspace = true litellm-gateway-mcp.workspace = true tokio-util = "0.7" thiserror.workspace = true diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index 47c1fe83f50..ff528723095 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -10,7 +10,7 @@ use serde_json::json; mod settings; -use settings::{Settings, UiSettings}; +use settings::{LogSettings, Settings, UiSettings}; struct StderrSink { level: Level, @@ -41,6 +41,15 @@ impl Sink for StderrSink { #[tokio::main] async fn main() -> Result<(), Box> { + if let Some(forward) = litellm_gateway_forward::Settings::from_env()? { + let log: LogSettings = settings::from_iter(std::env::vars_os())?; + Logger::new(StderrSink { + level: log.log_level(), + }) + .install_global()?; + litellm_gateway_forward::run(forward, shutdown_signal()).await?; + return Ok(()); + } let settings: Settings = settings::from_iter(std::env::vars_os())?; let level = settings.log_level(); Logger::new(StderrSink { level }).install_global()?; @@ -126,7 +135,20 @@ async fn shutdown_signal() { } } } - #[cfg(not(unix))] + // Windows support exists for local CLI runs only, the proxy is never deployed on Windows. The supervisor + // stops the gateway with CTRL_BREAK_EVENT, which has to drain the same way SIGTERM does. + #[cfg(windows)] + { + match tokio::signal::windows::ctrl_break() { + Ok(mut ctrl_break) => { + tokio::select! { _ = tokio::signal::ctrl_c() => (), _ = ctrl_break.recv() => () } + } + Err(_) => { + let _ = tokio::signal::ctrl_c().await; + } + } + } + #[cfg(not(any(unix, windows)))] { let _ = tokio::signal::ctrl_c().await; } diff --git a/litellm-rust/crates/gateway/src/settings.rs b/litellm-rust/crates/gateway/src/settings.rs index 445071504ed..c2c8bc30934 100644 --- a/litellm-rust/crates/gateway/src/settings.rs +++ b/litellm-rust/crates/gateway/src/settings.rs @@ -4,6 +4,23 @@ use litellm_secret_value::SecretValue; use litellm_tracing::Level; use serde::{Deserialize, de::DeserializeOwned}; +#[derive(Deserialize)] +pub(super) struct LogSettings { + rust_log: Option, +} + +impl LogSettings { + pub fn log_level(&self) -> Level { + log_level(self.rust_log.as_deref()) + } +} + +fn log_level(rust_log: Option<&str>) -> Level { + rust_log + .and_then(|value| value.parse().ok()) + .unwrap_or(Level::INFO) +} + #[derive(Deserialize)] #[serde(default)] pub(super) struct Settings { @@ -26,10 +43,7 @@ impl Default for Settings { impl Settings { pub fn log_level(&self) -> Level { - self.rust_log - .as_deref() - .and_then(|value| value.parse().ok()) - .unwrap_or(Level::INFO) + log_level(self.rust_log.as_deref()) } } diff --git a/litellm-rust/crates/python-bridge/src/router/config.rs b/litellm-rust/crates/python-bridge/src/router/config.rs index 0b46f2208c5..de383a0ffa1 100644 --- a/litellm-rust/crates/python-bridge/src/router/config.rs +++ b/litellm-rust/crates/python-bridge/src/router/config.rs @@ -188,6 +188,8 @@ struct RequestSpec { web_search: Option, route_handler: Option, decisions_media: Option, + batch_create: Option, + embedding_batch: Option, wildcard_request: Option, unrouted: Option, passthrough_on_no_deployment: Option, @@ -255,6 +257,8 @@ pub(crate) fn parse_request(value: Value, model: Option<&str>) -> Result, + }, /// Tag routing: return the deployments that remain. Tags, /// `enable_pre_call_checks`: return the facts the filter decides by. @@ -79,6 +84,10 @@ pub(crate) enum Stage { RouteIncapable { handler: String }, /// `_decisions_deployment_refuses_media`: return the candidates that refuse the call's media. MediaRefusing, + /// `deployments_admitting_embedding_batch`: return the candidates the embedding pick keeps. + EmbeddingBatch, + /// The override strategy's pre-call check on a deployment named by id. + OverrideCheck(String), /// The routing plugins: run them when `run`, then return the candidate models they left. RoutingPlugins { run: bool }, /// The host's `CustomLogger` deployment filters: return the ids that remain, or a failure. @@ -163,6 +172,8 @@ pub(crate) enum HostCall { /// The name the attempt resolved: the requested name or the fallback target. model: String, model_group: String, + /// The group a pre-routing hook selected for `model`, when it differs. + routed: Option, /// The `order` level of a same-group order fallback. target_order: Option, /// Every deployment `model` resolved to for this caller, the picked one included. @@ -279,6 +290,8 @@ pub(crate) enum HostCallResponse { RefuseFallbacks(bool), /// The candidate ids that remain, pick first. Select(Vec), + /// A strategy's pick among candidates sharing an id: positions in the candidates, pick first. + Positions(Vec), /// The model group to resolve instead, a custom strategy's deployment, or `None`. PreRoute(Option), /// The facts a `pre_call` select step asked for. @@ -309,9 +322,11 @@ impl HostCallResponse { Self::Slept => StepKind::Sleep, Self::AllowFallback(_) => StepKind::AllowFallback, Self::RefuseFallbacks(_) => StepKind::RefuseFallbacks, - Self::Select(_) | Self::PreCall(_) | Self::Models(_) | Self::Filtered(_) => { - StepKind::Select - } + Self::Select(_) + | Self::Positions(_) + | Self::PreCall(_) + | Self::Models(_) + | Self::Filtered(_) => StepKind::Select, Self::PreRoute(_) => StepKind::PreRoute, Self::Cache(_) => StepKind::Cache, Self::Secrets(_) => StepKind::Secrets, @@ -343,6 +358,13 @@ impl HostCallResponse { { return Err(DriveError::UnknownCandidate { id: id.clone() }); } + if let (Self::Positions(positions), HostCall::Select { candidates, .. }) = (self, call) + && let Some(position) = positions.iter().find(|p| **p >= candidates.len()) + { + return Err(DriveError::UnknownCandidate { + id: format!("#{position}"), + }); + } if let (Self::Secrets(values), HostCall::Secrets { refs, .. }) = (self, call) && values.len() != refs.len() { @@ -654,6 +676,28 @@ impl RouterCache for StepCache { } /// A deployment as Python spells a `model_list` entry, with its id in `model_info`. +/// Whether two candidates share an id, which a pick answered by id can't tell apart. +fn shares_ids(candidates: &[Arc]) -> bool { + let mut seen = std::collections::HashSet::new(); + !candidates.iter().all(|d| seen.insert(d.id.as_str())) +} + +/// The candidates a host answer names, in its order. An id names every candidate that carries +/// it, so deployments sharing an id stay together through a filter, as Python's list keeps them. +fn answered_by_id(candidates: &[Arc], ids: &[String]) -> Vec> { + let mut taken = vec![false; candidates.len()]; + let mut kept = Vec::new(); + for id in ids { + for (position, deployment) in candidates.iter().enumerate() { + if !taken[position] && deployment.id == *id { + taken[position] = true; + kept.push(deployment.clone()); + } + } + } + kept +} + pub(crate) fn deployment_record(deployment: &Deployment) -> Value { deployment.record() } @@ -693,6 +737,7 @@ impl Host for StepHost { deployment: deployment_record(&attempt.deployment), model: attempt.name.clone(), model_group: attempt.model_group.clone(), + routed: attempt.routed.clone(), target_order: attempt.target_order, candidate_ids: attempt.candidate_ids.clone(), retry: attempt.retry, @@ -736,13 +781,22 @@ impl Host for StepHost { let call = HostCall::Select { stage: match stage { SelectStage::Plugins => Stage::Plugins, - SelectStage::Strategy(strategy) => Stage::Strategy(strategy.to_string()), + SelectStage::Strategy(strategy) => Stage::Strategy { + strategy: strategy.to_string(), + records: if shares_ids(&candidates) { + candidates.iter().map(|d| deployment_record(d)).collect() + } else { + Vec::new() + }, + }, SelectStage::Tags => Stage::Tags, SelectStage::TextOnly => Stage::TextOnly, SelectStage::RouteIncapable { handler } => Stage::RouteIncapable { handler: handler.clone(), }, SelectStage::MediaRefusing => Stage::MediaRefusing, + SelectStage::EmbeddingBatch => Stage::EmbeddingBatch, + SelectStage::OverrideCheck(strategy) => Stage::OverrideCheck(strategy.to_string()), }, pool: ctx.pool.iter().map(|d| d.id.clone()).collect(), model_group: ctx.model_group.to_owned(), @@ -751,13 +805,14 @@ impl Host for StepHost { group: group.map(str::to_owned), candidates: candidates.iter().map(|d| d.id.clone()).collect(), }; - let HostCallResponse::Select(ids) = answered(call_host(call).await)? else { - unreachable!("{CHECKED}"); - }; - Ok(ids - .iter() - .filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned()) - .collect()) + match answered(call_host(call).await)? { + HostCallResponse::Select(ids) => Ok(answered_by_id(&candidates, &ids)), + HostCallResponse::Positions(positions) => Ok(positions + .iter() + .map(|position| candidates[*position].clone()) + .collect()), + _ => unreachable!("{CHECKED}"), + } } async fn pre_route(&self, pre: &PreRoute<'_>) -> Result, Error> { @@ -855,10 +910,7 @@ impl Host for StepHost { unreachable!("{CHECKED}"); }; let ids = result.map_err(Error::Attempt)?; - Ok(ids - .iter() - .filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned()) - .collect()) + Ok(answered_by_id(&candidates, &ids)) } async fn allow_fallback( @@ -922,6 +974,7 @@ impl Host for StepHost { | RouterEvent::Blocked { .. } | RouterEvent::CooldownSet { .. } | RouterEvent::Selected { .. } + | RouterEvent::PickStarted { .. } ) { return; } @@ -1461,7 +1514,7 @@ mod tests { let mut driver = start(&router, "primary"); let step = store.resume(&mut driver, None); let Step::HostCall(HostCall::Select { - stage: Stage::Strategy(strategy), + stage: Stage::Strategy { strategy, records }, candidates, .. }) = step @@ -1470,6 +1523,7 @@ mod tests { }; assert_eq!(strategy, "least-busy"); assert_eq!(candidates, ["a1", "a2"]); + assert!(records.is_empty(), "distinct ids send no records"); let step = store.resume( &mut driver, Some(HostCallResponse::Select(vec!["a2".to_owned()])), @@ -1480,6 +1534,40 @@ mod tests { ); } + #[rstest] + fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() { + let mut second = deployment("primary", "a1"); + second["litellm_params"]["api_base"] = json!("https://second.example"); + let router = router( + json!({ + "model_list": [deployment("primary", "a1"), second], + "routing_strategy": "least-busy", + }), + json!({}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + let step = store.resume(&mut driver, None); + let Step::HostCall(HostCall::Select { + stage: Stage::Strategy { records, .. }, + .. + }) = step + else { + panic!("a host strategy should ask for a selection, got {step:?}"); + }; + assert_eq!( + records.len(), + 2, + "each candidate sharing an id sends its record" + ); + let step = store.resume(&mut driver, Some(HostCallResponse::Positions(vec![1]))); + assert!( + matches!(&step, Step::HostCall(HostCall::Attempt { deployment, .. }) + if deployment["litellm_params"]["api_base"] == "https://second.example"), + "expected the second entry, got {step:?}" + ); + } + #[rstest] fn a_selection_is_driven_through_cache_steps(pair: Router) { let request = RouterRequest { diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs index c69259fcf1c..8157d0f5a32 100644 --- a/litellm-rust/crates/python-bridge/src/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -28,7 +28,7 @@ use litellm_router_core::{ use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, - types::{PyBool, PyDict, PyList, PyString}, + types::{PyBool, PyDict, PyList, PyString, PyTuple}, }; use serde_json::Value; @@ -457,7 +457,7 @@ impl RouterCoreHandle { } /// Record a success a logging callback saw, driven like `report_failure`; it ends with - /// `done` (`key`: the success counter's key, or None for an unknown deployment). + /// `done` (`key`: the tokens counter's key when usage was added, else None). /// `usage_model` is the model its usage is keyed by, and `tracks_usage` the host's answer to /// whether it counts any (see `Output::tracks_usage`). `callbacks`: see `report_failure`. #[pyo3(signature = (id, total_tokens, requests, usage_model=None, tracks_usage=None, callbacks=None))] @@ -1193,6 +1193,12 @@ fn host_call_response( } Ok(HostCallResponse::Filtered(Ok(outcome.extract()?))) } + HostCall::Select { + stage: Stage::Strategy { records, .. }, + .. + } if !records.is_empty() => Ok(HostCallResponse::Positions( + outcome.ok_or_else(missing)?.extract()?, + )), HostCall::Select { .. } => Ok(HostCallResponse::Select( outcome.ok_or_else(missing)?.extract()?, )), @@ -1285,6 +1291,15 @@ fn pre_route_outcome(outcome: Option>) -> PyResult() { return Ok(Some(PreRouted::Model(model.to_str()?.to_owned()))); } + // `(model, group)`: the model to resolve, served as the group a session binding stamped. + if let Ok(pair) = outcome.cast::() + && pair.len() == 2 + { + return Ok(Some(PreRouted::Bound { + model: pair.get_item(0)?.extract()?, + group: pair.get_item(1)?.extract()?, + })); + } let spec: DeploymentSpec = serde_json::from_value(document(&outcome)?) .map_err(|error| PyValueError::new_err(error.to_string()))?; let deployment = Deployment::from_spec(spec) @@ -1401,7 +1416,7 @@ fn cause_items(dict: &Bound<'_, PyDict>, cause: RetryCause) -> PyResult<()> { fn pick_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> { error_items(dict, error)?; match error { - Error::NoHealthyDeployments { model } | Error::UnknownModel { model } => { + Error::NoHealthyDeployments { model, .. } | Error::UnknownModel { model } => { dict.set_item("failed_model", model) } _ => Ok(()), @@ -1443,6 +1458,10 @@ fn error_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> { dict.set_item("error_model", model)?; dict.set_item("error_detail", format!("{until} {timezone}")) } + Error::NoHealthyDeployments { + request_ids: Some(ids), + .. + } => dict.set_item("error_request_ids", PyList::new(dict.py(), ids)?), Error::NonNumericWeight { metric } => dict.set_item("error_detail", metric), Error::MissingRoutingGroup { group } => dict.set_item("error_model", group), Error::ModelInfo { exception, detail } => { @@ -1637,6 +1656,10 @@ fn events_list(py: Python<'_>, events: Vec) -> PyResult { + dict.set_item("kind", "pick_started")?; + dict.set_item("model", model)?; + } RouterEvent::Selected { deployment_id, pick, @@ -1685,6 +1708,7 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) -> deployment, model, model_group, + routed, target_order, candidate_ids, retry, @@ -1696,6 +1720,7 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) -> max_retries, model_group_size, } => { + dict.set_item("routed", routed)?; dict.set_item("scope", scope)?; dict.set_item("max_retries", max_retries)?; dict.set_item("model_group_size", model_group_size)?; @@ -1729,7 +1754,12 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) -> } => { let (stage, strategy) = match stage { Stage::Plugins => ("plugins", None), - Stage::Strategy(strategy) => ("strategy", Some(strategy)), + Stage::Strategy { strategy, records } => { + if !records.is_empty() { + dict.set_item("records", to_py(py, &Value::Array(records))?)?; + } + ("strategy", Some(strategy)) + } Stage::Tags => ("tags", None), Stage::PreCall => ("pre_call", None), Stage::TextOnly => ("text_only", None), @@ -1738,6 +1768,8 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) -> ("route_incapable", None) } Stage::MediaRefusing => ("media_refusing", None), + Stage::EmbeddingBatch => ("embedding_batch", None), + Stage::OverrideCheck(strategy) => ("override_check", Some(strategy)), Stage::ProviderRequest => ("provider_request", None), Stage::CallbackFilters { excluded_ids, diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs index 287bb98d8b7..a118cac28c1 100644 --- a/litellm-rust/crates/router-core/src/call.rs +++ b/litellm-rust/crates/router-core/src/call.rs @@ -4,13 +4,7 @@ //! `should_retry_this_error` and `_time_to_sleep_before_retry` in `litellm/router.py`, and //! `run_async_fallback` in `router_utils/fallback_event_handlers.py`. -use std::{ - borrow::Cow, - future::Future, - pin::Pin, - sync::Arc, - time::{Duration, Instant}, -}; +use std::{borrow::Cow, future::Future, pin::Pin, sync::Arc, time::Duration}; use serde_json::{Map, Value}; @@ -29,7 +23,7 @@ use crate::{ host_index::python_str, lookup::Lookup, mock::MockFailure, - outcome::OutcomeWrite, + outcome::{OutcomeCallbacks, OutcomeWrite}, prerouting::{PreRoute, PreRouted}, registry::{ResolveKind, ResolveScope}, retry::{RetryContext, resolve_retry_policy, should_retry}, @@ -63,6 +57,12 @@ pub struct RouterResponse { struct Failed { error: Box, deployment_id: Option, + /// The group a pre-routing hook rewrote the level's name to for the failed attempt (an auto + /// router's tier), which its order fallbacks and fallback lookup key off + /// (`get_pre_routing_selection`). + routed: Option, + /// The group the host stamped that attempt as served for ([`Candidates::served_as`]). + served_as: Option, } impl From for Failed { @@ -70,6 +70,8 @@ impl From for Failed { Self { error: Box::new(error), deployment_id: None, + routed: None, + served_as: None, } } } @@ -103,6 +105,8 @@ struct Target { /// `_failover_excluded_ids` in the level's metadata: the deployments weighted failover /// already excluded on the way here. failover_excluded: Arc<[String]>, + /// The group a pre-routing hook selected for the level's failed attempt, if it rewrote it. + routed: Option, } /// A step of a level's fallback chain: a configured target, or an order level of the group. @@ -166,10 +170,18 @@ fn unattempted_chain( target: &Target, at: FallbackHop, attempted: &[Tried], + served: (Option<&str>, Option<&str>), ) -> bool { - let mut lookup = vec![target.group.as_str()]; - if request.model != target.group { - lookup.push(request.model.as_str()); + let (routed, served_as) = served; + let mut lookup: Vec<&str> = Vec::new(); + for group in routed + .into_iter() + .chain(served_as) + .chain([target.group.as_str(), request.model.as_str()]) + { + if !lookup.contains(&group) { + lookup.push(group); + } } list.lookup_groups(&lookup, lists.keys) .is_some_and(|targets| { @@ -285,6 +297,7 @@ impl Core { params: None, excluded: Arc::from([]), failover_excluded: Arc::from([]), + routed: None, }; let at = FallbackHop::default(); // A fallbacks mock stands in for the requested group's retry loop: no pick, no attempt. @@ -343,6 +356,8 @@ impl Core { let Failed { error, deployment_id: failed_id, + routed, + served_as, } = failed; let error = *error; self.level_failed(at, &target, &error, false); @@ -360,7 +375,14 @@ impl Core { self.level_failed(at, &target, &raised, true); return Err(raised); } - let (targets, order_path) = self.fallback_targets(snapshot, request, &target, &error); + let (targets, order_path) = self.fallback_targets( + snapshot, + request, + &target, + routed.as_deref(), + served_as.as_deref(), + &error, + ); if !order_path { self.host.observe(&RouterEvent::FallingBack { level: at.level, @@ -368,6 +390,7 @@ impl Core { }); } let mut target = target; + target.routed = routed; // `_maybe_run_weighted_failover`: its own `run_async_fallback`, before the regular one. let failover = if order_path { None @@ -463,7 +486,12 @@ impl Core { if at.depth >= max_fallbacks { return Err(None); } - let failed = Tried::Target(FallbackTarget::Name(target.group.clone())); + let failed = Tried::Target(FallbackTarget::Name( + target + .routed + .clone() + .unwrap_or_else(|| target.group.clone()), + )); if !chain.attempted.contains(&failed) { chain.attempted.push(failed); } @@ -488,17 +516,34 @@ impl Core { let mut unroutable: Option = None; for next in targets { let (tried, group, order) = match &next { - Next::Order(group, order) => ( - Tried::Order(group.clone(), *order), - group.clone(), - Some(*order), - ), + Next::Order(group, order) => { + // An order level of the group a pre-routing hook selected passes the + // fallback gates as any group the level did not serve does. + if *group != target.group { + match self.host.allow_fallback(request, group).await { + Ok(true) => {} + Ok(false) => continue, + Err(raised) => return Err(Some(Error::Attempt(raised))), + } + } + ( + Tried::Order(group.clone(), *order), + group.clone(), + Some(*order), + ) + } Next::Failover(group, excluded) => { let mut key = excluded.to_vec(); key.sort(); (Tried::Failover(group.clone(), key), group.clone(), None) } - Next::Fallback(FallbackTarget::Name(name)) if *name == target.group => continue, + // The group that just failed: the pre-routing selection when there was one, so + // a subagent bound to another tier can fall back to the model it asked for. + Next::Fallback(FallbackTarget::Name(name)) + if *name == target.routed.as_deref().unwrap_or(&target.group) => + { + continue; + } // A provider-scoped request never leaves its group; a dict target with no // `model` names no group, so it is skipped too. Next::Fallback(fallback) @@ -573,6 +618,7 @@ impl Core { params: params.clone(), excluded, failover_excluded: target.failover_excluded.clone(), + routed: None, }; let parent = at; let at = FallbackHop { @@ -704,6 +750,8 @@ impl Core { snapshot: &RouterStateSnapshot, request: &RouterRequest, target: &Target, + routed: Option<&str>, + served_as: Option<&str>, error: &Error, ) -> (Option>, bool) { // A pick's error decides by its exception's classes as an attempt's does: a pre-call @@ -711,11 +759,19 @@ impl Core { let failure = retried_failure(error); let is_a = |name: &str| failure.as_ref().is_some_and(|f| f.is_a(name)); let lists = Lists::of(snapshot, request); - // `fallback_lookup_groups`: the group being served, then the group the request was for. - let mut lookup = vec![target.group.as_str()]; - if request.model != target.group { - lookup.push(request.model.as_str()); + // `fallback_lookup_groups`: the group a pre-routing hook selected, the group being served, + // then the group the request was for. + let mut lookup: Vec<&str> = routed.into_iter().collect(); + for group in served_as + .into_iter() + .chain([target.group.as_str(), request.model.as_str()]) + { + if !lookup.contains(&group) { + lookup.push(group); + } } + // `order_model_group`: the pre-routing selection's order levels, else the level's own. + let order_group = routed.unwrap_or(target.group.as_str()); let chain = |list: &Fallbacks| { if list.is_client_side(lists.keys) { Some(list.as_chain()) @@ -734,7 +790,7 @@ impl Core { if !skip_order { // A routing group's members, with a priority group's priorities as their order. let listing_team = request.listing_team_id.as_deref(); - let mut orders: Vec = group_members(snapshot, &target.group, listing_team) + let mut orders: Vec = group_members(snapshot, order_group, listing_team) .iter() .filter(|d| d.listed_for_team(listing_team)) .filter_map(|d| d.order()) @@ -742,11 +798,14 @@ impl Core { orders.sort_unstable(); orders.dedup(); if orders.len() > 1 { - let current = target.order.unwrap_or(orders[0]); + let current = target + .order + .filter(|_| routed.is_none()) + .unwrap_or(orders[0]); let mut targets: Vec = orders .into_iter() .filter(|order| *order > current) - .map(|order| Next::Order(target.group.clone(), order)) + .map(|order| Next::Order(order_group.to_owned(), order)) .collect(); if let Some(list) = lists.fallbacks { targets.extend(fallbacks_of(chain(list).unwrap_or_default())); @@ -879,6 +938,8 @@ impl Core { return Err(Failed { error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))), deployment_id: None, + routed: None, + served_as: None, }); } let Some(first_failure) = retried_failure(&failed.error) else { @@ -965,6 +1026,8 @@ impl Core { return Err(Failed { error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))), deployment_id: None, + routed: None, + served_as: None, }); } if num_retries == 0 { @@ -1162,7 +1225,7 @@ impl Core { .unrouted_attempt(snapshot, op, request, target, (retry, max_retries), at) .await; } - let pick_started = Instant::now(); + let pick_started = self.health.now_time(); let (deployment, candidates) = match self .picked(snapshot, op, request, target, inputs, retry, at) .await @@ -1179,7 +1242,11 @@ impl Core { self.host.observe(&RouterEvent::Selected { model_group: candidates.model_group.clone(), deployment_id: deployment.id.clone(), - pick: pick_started.elapsed(), + pick: self + .health + .now_time() + .duration_since(pick_started) + .unwrap_or_default(), }); if !self.host.supports(op, deployment.provider()) { return Err(Error::Config(ConfigError::Unsupported(Unsupported::Op { @@ -1205,7 +1272,7 @@ impl Core { }, None, snapshot.settings.dynamic.default_ref(), - jiff::Timestamp::now(), + self.health.timestamp(), ), Ok(None) ); @@ -1219,6 +1286,7 @@ impl Core { } else { candidates.model_group.clone() }, + routed: routed_name(target, &candidates), target_order: target.order, candidate_ids: candidates.all.iter().map(|d| d.id.clone()).collect(), retry, @@ -1236,23 +1304,43 @@ impl Core { // the stream's end. match self.host.invoke(&attempt).await { Ok(output) => { + // `_should_raise_content_policy_error`: a blocked response is raised only when + // a content-policy fallback would serve it, and counts as no failure. + let routed = routed_name(target, &candidates); + let served = (routed.as_deref(), candidates.served_as.as_deref()); + let raised = match output.blocked { + Some(Blocked::ContentFilter) => self.has_content_policy_fallback( + snapshot, request, target, at, attempted, served, + ), + Some(Blocked::Refusal) => self.refusal_fallback_available( + snapshot, request, target, at, attempted, served, + ), + None => false, + }; let tracks = output .tracks_usage .unwrap_or_else(|| deployment.tracks_usage()); + // A shared logging object logs the success of the response that serves the + // request alone, so a raised blocked response counts its request and no success. + let callbacks = match raised && request.outcome_callbacks.shared_logging { + true => OutcomeCallbacks { + success: false, + ..request.outcome_callbacks + }, + false => request.outcome_callbacks, + }; let recorded = match (output.deferred, tracks) { (true, true) => { let tokens = output.total_tokens as f64; - self.health.add_usage(&deployment, None, tokens, 1.0).await + self.health + .add_usage(&deployment, None, tokens, 1.0) + .await + .map(|_| ()) } (true, false) => Ok(()), (false, _) => { self.health - .record_success( - request.outcome_callbacks, - &deployment, - output.total_tokens, - tracks, - ) + .record_success(callbacks, &deployment, output.total_tokens, tracks) .await } }; @@ -1262,16 +1350,6 @@ impl Core { self.host.observe(&RouterEvent::Succeeded { deployment_id: deployment.id.clone(), }); - // `_should_raise_content_policy_error`: a blocked response is raised only when - // a content-policy fallback would serve it, and counts as no failure. - let raised = - match output.blocked { - Some(Blocked::ContentFilter) => self - .has_content_policy_fallback(snapshot, request, target, at, attempted), - Some(Blocked::Refusal) => self - .refusal_fallback_available(snapshot, request, target, at, attempted), - None => false, - }; if let Some(blocked) = output.blocked.filter(|_| raised) { self.host.observe(&RouterEvent::Blocked { level: at.level, @@ -1284,6 +1362,8 @@ impl Core { return Err(Failed { error: Box::new(Error::Attempt(failure)), deployment_id: Some(deployment.id.clone()), + routed: routed_name(target, &candidates), + served_as: candidates.served_as.clone(), }); } Ok(RouterResponse { @@ -1303,6 +1383,8 @@ impl Core { Err(Failed { error: Box::new(Error::Attempt(failure)), deployment_id: Some(deployment.id.clone()), + routed: routed_name(target, &candidates), + served_as: candidates.served_as.clone(), }) } } @@ -1325,6 +1407,9 @@ impl Core { &snapshot.settings.routing_strategy, request, )?; + self.host.observe(&RouterEvent::PickStarted { + model: target.group.clone(), + }); let routed = self .pre_routed(PreRoute { op, @@ -1334,22 +1419,25 @@ impl Core { retry, }) .await?; + let (routed, served_as) = match routed { + PreRouted::Bound { model, group } => (PreRouted::Model(model), Some(group)), + other => (other, None), + }; match routed { PreRouted::Model(name) => { snapshot.missing_routing_group(request, &name)?; - let candidates = self + let mut candidates = self .candidates(snapshot, op, request, &name, inputs) .await?; + candidates.served_as = served_as; if candidates.healthy.is_empty() { - return Err(Error::NoHealthyDeployments { - model: candidates.failed_model(&snapshot.registry), - } - .into()); + return Err(candidates.none_healthy(&snapshot.registry).into()); } let deployment = self.pick(snapshot, op, request, &candidates).await?; Ok((deployment, candidates)) } // A custom strategy's pick is served as is, for the requested name. + PreRouted::Bound { .. } => unreachable!("mapped to a model above"), PreRouted::Deployment(deployment) => { let candidates = Candidates { name: target.group.clone(), @@ -1358,6 +1446,7 @@ impl Core { all: vec![deployment.clone()], healthy: vec![deployment.clone()], wildcard_request: None, + served_as: None, }; Ok((deployment, candidates)) } @@ -1383,6 +1472,7 @@ impl Core { deployment, name: target.group.clone(), model_group: target.group.clone(), + routed: None, target_order: target.order, candidate_ids: Vec::new(), retry, @@ -1417,10 +1507,11 @@ impl Core { target: &Target, at: FallbackHop, attempted: &[Tried], + served: (Option<&str>, Option<&str>), ) -> bool { let lists = Lists::of(snapshot, request); match lists.content_policy { - Some(list) => unattempted_chain(list, &lists, request, target, at, attempted), + Some(list) => unattempted_chain(list, &lists, request, target, at, attempted, served), None => snapshot.settings.dynamic.fallbacks.has_star(), } } @@ -1434,20 +1525,22 @@ impl Core { target: &Target, at: FallbackHop, attempted: &[Tried], + served: (Option<&str>, Option<&str>), ) -> bool { if request.disable_fallbacks { return false; } let lists = Lists::of(snapshot, request); if lists.content_policy.is_some() { - return self.has_content_policy_fallback(snapshot, request, target, at, attempted); + return self + .has_content_policy_fallback(snapshot, request, target, at, attempted, served); } if snapshot.settings.dynamic.fallbacks.has_star() { return true; } - lists - .fallbacks - .is_some_and(|list| unattempted_chain(list, &lists, request, target, at, attempted)) + lists.fallbacks.is_some_and(|list| { + unattempted_chain(list, &lists, request, target, at, attempted, served) + }) } /// `async_deployment_callback_on_failure`, when it runs for this failure. @@ -1553,3 +1646,8 @@ impl Core { Ok(()) } } + +/// The name a pre-routing hook rewrote the level's group to for this pick, if it did. +fn routed_name(target: &Target, candidates: &Candidates) -> Option { + (candidates.name != target.group).then(|| candidates.name.clone()) +} diff --git a/litellm-rust/crates/router-core/src/call/stream.rs b/litellm-rust/crates/router-core/src/call/stream.rs index a0fce4372dd..f1dc71a9b18 100644 --- a/litellm-rust/crates/router-core/src/call/stream.rs +++ b/litellm-rust/crates/router-core/src/call/stream.rs @@ -37,6 +37,9 @@ pub struct StreamOrigin { pub attempted_retries: u32, /// The budget the first stream retry committed to (`committed_retry_budget_for_request`). pub committed_budget: Option, + /// The group a pre-routing hook selected for `group`, which the fallback lookup keys on first. + #[serde(default)] + pub routed: Option, } /// What may still take over a stream that fails before content. @@ -61,14 +64,22 @@ impl StreamOrigin { params: None, excluded: Arc::from([]), failover_excluded: Arc::from([]), + routed: self.routed.clone(), } } /// `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()); + let mut lookup: Vec<&str> = Vec::new(); + for group in self + .routed + .as_deref() + .into_iter() + .chain([self.group.as_str(), request.model.as_str()]) + { + if !lookup.contains(&group) { + lookup.push(group); + } } lookup } @@ -186,7 +197,10 @@ impl Router { scope: 0, }, &mut chain, - Failed::from(Error::Attempt(failure)), + Failed { + routed: origin.routed.clone(), + ..Failed::from(Error::Attempt(failure)) + }, ) .await } diff --git a/litellm-rust/crates/router-core/src/error.rs b/litellm-rust/crates/router-core/src/error.rs index b4112501c36..926772f435c 100644 --- a/litellm-rust/crates/router-core/src/error.rs +++ b/litellm-rust/crates/router-core/src/error.rs @@ -40,7 +40,13 @@ pub enum Error { #[error("no deployment serves model `{model}`")] UnknownModel { model: String }, #[error("no deployment of `{model}` is available")] - NoHealthyDeployments { model: String }, + NoHealthyDeployments { + model: String, + /// The ids of the deployments the request could route to before the cooldown filter + /// (`request_deployments`), when the candidate list came up empty; `None` when the + /// strategy refused every candidate. + request_ids: Option>, + }, #[error("the attempt failed")] Attempt(#[from] AttemptError), #[error("the routing cache failed")] @@ -201,8 +207,9 @@ impl Error { Self::UnknownModel { model } => Self::UnknownModel { model: model.clone(), }, - Self::NoHealthyDeployments { model } => Self::NoHealthyDeployments { + Self::NoHealthyDeployments { model, request_ids } => Self::NoHealthyDeployments { model: model.clone(), + request_ids: request_ids.clone(), }, Self::Paused { model } => Self::Paused { model: model.clone(), diff --git a/litellm-rust/crates/router-core/src/health.rs b/litellm-rust/crates/router-core/src/health.rs index 020a1a0f47d..eaa61dac7ba 100644 --- a/litellm-rust/crates/router-core/src/health.rs +++ b/litellm-rust/crates/router-core/src/health.rs @@ -247,6 +247,10 @@ impl Health { self.clock.now() } + pub(crate) fn timestamp(&self) -> jiff::Timestamp { + jiff::Timestamp::try_from(self.clock.now()).unwrap_or(jiff::Timestamp::UNIX_EPOCH) + } + fn now(&self) -> f64 { unix_seconds(self.clock.now()) } @@ -554,35 +558,35 @@ impl Health { return Ok(()); } callbacks - .record( - OutcomeWrite::AttemptUsage, - self.add_usage(deployment, None, total_tokens as f64, 1.0), - ) + .record(OutcomeWrite::AttemptUsage, async { + self.add_usage(deployment, None, total_tokens as f64, 1.0) + .await + .map(|_| ()) + }) .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. + /// deployment's own), unless there is nothing to add. Returns the tokens counter's key when + /// it added to it. pub async fn add_usage( &self, deployment: &Deployment, model: Option<&str>, total_tokens: f64, requests: f64, - ) -> Result<(), Error> { + ) -> Result, Error> { if total_tokens <= 0.0 && requests <= 0.0 { - return Ok(()); + return Ok(None); } let minute = current_minute(self.clock.now()); let model = model.unwrap_or(&deployment.params.model); + let tpm_key = usage_key(&deployment.id, model, UsageMetric::Tpm, &minute); self.cache .async_increment_usage( &[ - ( - usage_key(&deployment.id, model, UsageMetric::Tpm, &minute), - total_tokens, - ), + (tpm_key.clone(), total_tokens), ( usage_key(&deployment.id, model, UsageMetric::Rpm, &minute), requests, @@ -591,7 +595,7 @@ impl Health { Some(USAGE_TTL), ) .await?; - Ok(()) + Ok(Some(tpm_key)) } /// `async_deployment_callback_on_failure`: a failed request counts against the deployment's diff --git a/litellm-rust/crates/router-core/src/host.rs b/litellm-rust/crates/router-core/src/host.rs index 0a72be8f980..cfc97480bf1 100644 --- a/litellm-rust/crates/router-core/src/host.rs +++ b/litellm-rust/crates/router-core/src/host.rs @@ -83,6 +83,13 @@ pub struct RouterRequest { /// 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, + /// A batch create (`batch_create_routing`): when the caller may use a deployment declared + /// `model_info.mode: batch`, the Azure deployments that are not are skipped + /// (`without_deployments_that_cannot_batch`). + pub batch_create: bool, + /// An embedding pick: the host keeps the candidates whose `max_embedding_batch_size` admits + /// the input (`deployments_admitting_embedding_batch`). + pub embedding_batch: 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, @@ -291,6 +298,9 @@ pub struct Attempt { pub name: String, /// The model group being served, after fallbacks. pub model_group: String, + /// The group a pre-routing hook selected for `name`, when it differs (the tier stamp the + /// fallback lookup keys on first). + pub routed: Option, /// The `order` level an order fallback targets in the same group. pub target_order: Option, /// Every deployment `name` resolved to for this caller, the picked one included. @@ -383,11 +393,22 @@ pub enum SelectStage { /// `_decisions_deployment_refuses_media`: return the candidates that refuse the audio or /// videos a Decisions call carries. MediaRefusing, + /// `deployments_admitting_embedding_batch`: return the candidates that admit the embedding + /// input, or what the pick keeps when none does (`deployments_when_none_admits_embedding_batch`). + EmbeddingBatch, + /// `_async_override_selector_pre_call_check`: a deployment named by id under a per-request + /// strategy override still passes that strategy's pre-call check (rpm accounting); return the + /// candidates, or raise its refusal. + OverrideCheck(RoutingStrategy), } /// Routing events, delivered through [`Host::observe`]. #[derive(Clone, Debug, PartialEq)] pub enum RouterEvent { + /// A pick for `model` (the level's group) begins: its route phase opens. + PickStarted { + model: String, + }, Selected { model_group: String, deployment_id: String, diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 5f15b6b8df0..80c8f597157 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -366,7 +366,7 @@ impl Router { }) .await?; let name = match routed { - PreRouted::Model(name) => name, + PreRouted::Model(name) | PreRouted::Bound { model: name, .. } => name, PreRouted::Deployment(deployment) => return Ok(deployment), }; snapshot.missing_routing_group(request, &name)?; @@ -374,6 +374,10 @@ impl Router { .core .candidates(&snapshot, op, request, &name, &SelectInputs::default()) .await?; + // An empty candidate list names what the request could route to before the cooldown filter. + if candidates.healthy.is_empty() && !candidates.kind.is_single() { + return Err(candidates.none_healthy(&snapshot.registry)); + } self.core.pick(&snapshot, op, request, &candidates).await } @@ -398,9 +402,7 @@ impl Router { .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), - }); + return Err(candidates.none_healthy(&snapshot.registry)); } Ok(candidates) } @@ -434,7 +436,7 @@ impl Router { /// 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. + /// tokens counter's key when it added usage (`_increment_deployment_usage`), else `None`. pub async fn report_success( &self, id: &str, @@ -462,22 +464,22 @@ impl Router { let Some(deployment) = snapshot.registry.deployment(id) else { return Ok(None); }; - let key = self.core.health.count_success(id).await?; - if report + 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?; + return Ok(None); } - Ok(Some(key)) + self.core + .health + .add_usage( + &deployment, + report.usage_model.as_deref(), + total_tokens, + requests, + ) + .await } /// `sync_deployment_callback_on_success`: count a success against `id`, configured or not, diff --git a/litellm-rust/crates/router-core/src/lookup.rs b/litellm-rust/crates/router-core/src/lookup.rs index 39c235d6f7b..8fef0ba222f 100644 --- a/litellm-rust/crates/router-core/src/lookup.rs +++ b/litellm-rust/crates/router-core/src/lookup.rs @@ -433,6 +433,10 @@ impl<'a> Lookup<'a> { .collect() } + pub fn has_team_wildcards(&self, team_id: &str) -> bool { + self.host.team_patterns.iter().any(|(t, _)| t == team_id) + } + /// 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> { diff --git a/litellm-rust/crates/router-core/src/outcome.rs b/litellm-rust/crates/router-core/src/outcome.rs index 59e0d14b2c2..7403b819d6d 100644 --- a/litellm-rust/crates/router-core/src/outcome.rs +++ b/litellm-rust/crates/router-core/src/outcome.rs @@ -28,6 +28,10 @@ pub struct OutcomeCallbacks { pub failure: bool, /// The attempt runs through `make_call`, which counts its usage whatever the callbacks do. pub make_call: bool, + /// The request's attempts log on one logging object the caller passed (`litellm_logging_obj`), + /// which logs one success a request, the response that serves it: an attempt whose blocked + /// response is raised for a fallback counts no success. + pub shared_logging: bool, } impl Default for OutcomeCallbacks { @@ -36,6 +40,7 @@ impl Default for OutcomeCallbacks { success: true, failure: true, make_call: true, + shared_logging: false, } } } @@ -78,21 +83,25 @@ mod tests { success: true, failure: true, make_call: true, + shared_logging: false, }; const UNBOUND: OutcomeCallbacks = OutcomeCallbacks { success: false, failure: false, make_call: true, + shared_logging: false, }; const UNBOUND_WITHOUT_MAKE_CALL: OutcomeCallbacks = OutcomeCallbacks { success: false, failure: false, make_call: false, + shared_logging: false, }; const BOUND_SYNC: OutcomeCallbacks = OutcomeCallbacks { success: false, failure: true, make_call: true, + shared_logging: false, }; #[rstest] diff --git a/litellm-rust/crates/router-core/src/prerouting.rs b/litellm-rust/crates/router-core/src/prerouting.rs index d6e3df05e7a..8ea1377a5aa 100644 --- a/litellm-rust/crates/router-core/src/prerouting.rs +++ b/litellm-rust/crates/router-core/src/prerouting.rs @@ -33,6 +33,9 @@ pub struct PreRoute<'a> { pub enum PreRouted { /// Resolve this model group instead of the requested name. Model(String), + /// Resolve `model`, served as `group`: the group a session binding stamped on the request, + /// which the level's fallbacks also look up. + Bound { model: String, group: String }, /// A custom strategy's pick: serve this deployment, without resolving or filtering. Deployment(Arc), } diff --git a/litellm-rust/crates/router-core/src/registry.rs b/litellm-rust/crates/router-core/src/registry.rs index 78c7d3b67e3..594aff5f6dd 100644 --- a/litellm-rust/crates/router-core/src/registry.rs +++ b/litellm-rust/crates/router-core/src/registry.rs @@ -234,6 +234,23 @@ impl Deployment { /// `_deployment_supports_web_search`: `model_info.supports_web_search` when it is set, read /// for truthiness, else true. + /// `is_batch_capable`: declared `model_info.mode: batch`. + pub(crate) fn is_batch_capable(&self) -> bool { + self.model_info.raw.get("mode").and_then(Value::as_str) == Some("batch") + } + + /// `_cannot_run_batches`: an Azure (or Azure AI) deployment not declared batch-capable, by + /// its `custom_llm_provider`, else its model's provider prefix. + pub(crate) fn cannot_run_batches(&self) -> bool { + let provider = self.params.custom_llm_provider.as_deref().or_else(|| { + self.params + .model + .split_once('/') + .map(|(provider, _)| provider) + }); + matches!(provider, Some("azure" | "azure_ai")) && !self.is_batch_capable() + } + pub(crate) fn supports_web_search(&self) -> bool { self.model_info .raw diff --git a/litellm-rust/crates/router-core/src/resolve.rs b/litellm-rust/crates/router-core/src/resolve.rs index 9075fc36c53..d897716b1e5 100644 --- a/litellm-rust/crates/router-core/src/resolve.rs +++ b/litellm-rust/crates/router-core/src/resolve.rs @@ -223,14 +223,20 @@ pub(crate) fn wildcard_name<'a>( registry: &'a Registry, name: &'a str, specific_deployment: bool, + team_id: Option<&str>, ) -> Option<&'a str> { let model = registry.alias_of(name).unwrap_or(name); - (registry.has_wildcards() + let lookup = crate::lookup::Lookup::new(registry); + let team_patterns = team_id.is_some_and(|team| lookup.has_team_wildcards(team)); + let team_misses = || team_id.is_none_or(|team| lookup.team_wildcard(model, team).is_empty()); + let global_misses = || registry.has_wildcards() && registry.wildcard_for(model).is_none(); + ((team_patterns || 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()) + && team_misses() + && (team_patterns || global_misses())) .then_some(model) } @@ -339,15 +345,21 @@ pub fn resolve_request( 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 { + // `_core_wildcard_route_deployments`: the team's own pattern, else the proxy-wide one, + // each tried as spelled and then under the name's provider (`get_pattern`). + let team_owned = team_id.and_then(|team| { + let lookup = crate::lookup::Lookup::new(registry); + [Some(model), caller.wildcard_request] + .into_iter() + .flatten() + .map(|request| (request, lookup.team_wildcard(request, team))) + .find(|(_, deployments)| !deployments.is_empty()) + }); + if let Some((request, 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 pattern = public.unwrap_or(request); + routed_wildcard(pattern, request, std::slice::from_ref(deployment)) }); let kept = reserved(model, routed.collect(), caller, now)?; return resolved(ResolveKind::Wildcard, model, kept); @@ -428,6 +440,7 @@ pub fn resolve_request( mod tests { use std::collections::BTreeMap; + use rstest::rstest; use serde_json::{Value, json}; use super::*; @@ -544,8 +557,11 @@ mod tests { #[test] fn a_name_no_pattern_matches_is_tried_under_its_provider() { let registry = registry(wildcard_group()); - assert_eq!(wildcard_name(®istry, "gpt-4o", false), Some("gpt-4o")); - assert_eq!(wildcard_name(®istry, "openai/gpt-4o", false), None); + assert_eq!( + wildcard_name(®istry, "gpt-4o", false, None), + Some("gpt-4o") + ); + assert_eq!(wildcard_name(®istry, "openai/gpt-4o", false, None), None); let caller = Caller { wildcard_request: Some("openai/gpt-4o"), ..Caller::default() @@ -564,6 +580,48 @@ mod tests { assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o"); } + #[rstest] + fn a_bare_name_is_tried_under_its_provider_against_the_teams_own_wildcard() { + let registry = registry(json!([ + {"model_name": "model_name_tw_1", "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "tw", "team_id": "TW", "team_public_model_name": "openai/*"}}, + {"model_name": "g", "litellm_params": {"model": "openai/g"}, "model_info": {"id": "g"}}, + ])); + // No proxy-wide wildcard: only the team's patterns call for the provider spelling. + assert_eq!( + wildcard_name(®istry, "gpt-4o-mini", false, Some("TW")), + Some("gpt-4o-mini") + ); + assert_eq!(wildcard_name(®istry, "gpt-4o-mini", false, None), None); + assert_eq!( + wildcard_name(®istry, "openai/gpt-4o-mini", false, Some("TW")), + None + ); + let resolve = |team_id| { + let caller = Caller { + team_id, + wildcard_request: Some("openai/gpt-4o-mini"), + ..Caller::default() + }; + resolve_request( + ®istry, + "gpt-4o-mini", + false, + caller, + None, + None, + Timestamp::now(), + ) + .unwrap() + }; + let resolution = resolve(Some("TW")).unwrap(); + assert_eq!(resolution.kind, ResolveKind::Wildcard); + assert_eq!(resolution.group, "gpt-4o-mini"); + assert_eq!(ids(Some(resolution.clone())), ["tw"]); + assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o-mini"); + assert!(resolve(Some("OTHER")).is_none()); + } + #[test] fn a_deployment_named_by_id_refuses_the_endpoint_like_any_other() { let registry = registry(json!([ diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs index fefa7f25ad1..c7f3b1f9c31 100644 --- a/litellm-rust/crates/router-core/src/select.rs +++ b/litellm-rust/crates/router-core/src/select.rs @@ -16,7 +16,7 @@ use crate::{ host::{RouterRequest, RoutingContext, SelectStage}, registry::{Deployment, Registry, ResolveKind}, resolve::{Caller, resolve_request}, - strategy::{StrategyScope, routing_context}, + strategy::{StrategyScope, request_override, routing_context}, }; /// Randomness for the shuffle and retry jitter, injectable so tests are deterministic. @@ -77,11 +77,23 @@ pub struct Candidates { /// The provider-prefixed spelling the name resolved under (`get_pattern`), when the host /// was asked for one. pub wildcard_request: Option, + /// The group the host stamped the pick as served for when a pre-routing hook rewrote the name + /// (a Claude Code subagent's bound router): a fallback lookup key. + pub served_as: Option, } 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. + /// The error an empty candidate list raises: it names the deployments the request could + /// route to before the cooldown filter. + pub(crate) fn none_healthy(&self, registry: &Registry) -> Error { + Error::NoHealthyDeployments { + model: self.failed_model(registry), + request_ids: Some(self.all.iter().map(|d| d.id.clone()).collect()), + } + } + 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), @@ -269,8 +281,12 @@ impl Core { // 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_name = crate::resolve::wildcard_name( + registry, + name, + request.specific_deployment, + request.team_id.as_deref(), + ); let wildcard_request = match wildcard_name { Some(model) => { let scope = StrategyScope::Default; @@ -296,7 +312,7 @@ impl Core { caller, default_target, snapshot.settings.dynamic.default_ref(), - jiff::Timestamp::now(), + self.health.timestamp(), )? .ok_or_else(|| Error::UnknownModel { // The name `_common_checks_available_deployment` looked for, past the alias. @@ -466,6 +482,7 @@ impl Core { all, healthy, wildcard_request, + served_as: None, }); } @@ -503,6 +520,7 @@ impl Core { all, healthy, wildcard_request, + served_as: None, }); } @@ -602,6 +620,34 @@ impl Core { } } + // `deployments_admitting_embedding_batch`: an embedding pick keeps the deployments whose + // input limit admits the batch; when none does, the host keeps what the pick waits for or + // the 400 names. + if request.embedding_batch && !resolution.kind.is_single() && !healthy.is_empty() { + let mut ctx = self.context( + snapshot, + op, + request, + &resolution.group, + name, + StrategyScope::Default, + ); + ctx.pool = &all; + healthy = self + .host + .select(&SelectStage::EmbeddingBatch, &ctx, healthy) + .await?; + } + + // `without_deployments_that_cannot_batch`: a batch create on a group the caller may use a + // batch deployment of skips the Azure deployments that cannot run batches. + if request.batch_create + && !resolution.kind.is_single() + && healthy.iter().any(|d| d.is_batch_capable()) + { + healthy.retain(|d| !d.cannot_run_batches()); + } + if let Some(pinned) = request.pinned_deployment_id.as_deref() && healthy.iter().any(|d| d.id == pinned) { @@ -653,6 +699,7 @@ impl Core { all, healthy, wildcard_request, + served_as: None, }) } @@ -691,14 +738,43 @@ impl Core { ) -> Result, Error> { let no_deployment = || Error::NoHealthyDeployments { model: candidates.failed_model(&snapshot.registry), + request_ids: None, }; - // A deployment named by its id is returned as is, before any strategy runs. + // A deployment named by its id is returned as is, before any strategy runs, past a + // per-request strategy override's pre-call check. if candidates.kind.is_single() { - return candidates + let named = candidates .healthy .first() .cloned() - .ok_or_else(no_deployment); + .ok_or_else(no_deployment)?; + if let Some(strategy) = request_override(request.routing_strategy.as_deref()) + && !matches!( + strategy, + RoutingStrategy::SimpleShuffle | RoutingStrategy::Unset + ) + { + let ctx = self.context( + snapshot, + op, + request, + &candidates.model_group, + &candidates.name, + StrategyScope::Override, + ); + let kept = self + .host + .select( + &SelectStage::OverrideCheck(strategy), + &ctx, + vec![named.clone()], + ) + .await?; + if kept.is_empty() { + return Err(no_deployment()); + } + } + return Ok(named); } let (strategy, scope) = routing_context( &snapshot.registry, diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs index 9f41371f255..a2a741d6135 100644 --- a/litellm-rust/crates/router-core/tests/call_loop.rs +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -1755,3 +1755,57 @@ async fn a_listed_pass_through_pick_names_the_group_an_alias_resolves_to(config: Error::NoPassThroughDeployments { model } if model == "primary" )); } + +struct StoppedClock(u64); + +impl litellm_router_core::health::Clock for StoppedClock { + fn now(&self) -> std::time::SystemTime { + std::time::UNIX_EPOCH + Duration::from_secs(self.0) + } +} + +fn utc_hour(hour: u64) -> u64 { + 1_767_225_600 + hour * 3_600 +} + +#[rstest] +#[case::inside_the_window(utc_hour(10), vec!["open"])] +#[case::outside_the_window(utc_hour(20), vec!["open", "reserved"])] +#[tokio::test] +async fn access_windows_follow_the_routers_clock(#[case] now: u64, #[case] expected: Vec<&str>) { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "chat", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "open"}}, + {"model_name": "chat", "litellm_params": {"model": "openai/b"}, "model_info": { + "id": "reserved", + "access_windows": [{"start": "09:00", "end": "17:00", "timezone": "UTC", "team_ids": ["batch"]}], + }}, + ] + })) + .unwrap(); + let router = Router::builder( + config, + RouterDynamicConfig::default(), + Arc::new(ScriptedHost::default()), + ) + .randomness(Arc::new(SeededRandomness::new(7))) + .clock(Arc::new(StoppedClock(now))) + .build() + .unwrap(); + let mut picked = std::collections::BTreeSet::new(); + for _ in 0..40 { + let deployment = router + .select(&Op::from("acompletion"), &request_for("chat")) + .await + .unwrap(); + picked.insert(deployment.id.clone()); + } + assert_eq!(picked.into_iter().collect::>(), expected); +} + +fn request_for(model: &str) -> RouterRequest { + RouterRequest { + model: model.to_owned(), + ..RouterRequest::default() + } +} diff --git a/litellm/accounting/application.py b/litellm/accounting/application.py index 1e5f23e2997..c911f4e35ce 100644 --- a/litellm/accounting/application.py +++ b/litellm/accounting/application.py @@ -7,6 +7,7 @@ from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_b 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.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction if TYPE_CHECKING: @@ -25,6 +26,7 @@ async def apply_deliveries( from litellm.proxy.guardrails.usage_tracking import UsageUnitIncrement, UsageUnitKey, upsert_usage_unit_row contributions: Final = project_deliveries(deliveries) + await apply_rollup_lock_timeout(transaction) await writer.apply_accounting_entities( prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions) ) diff --git a/litellm/accounting/metrics.py b/litellm/accounting/metrics.py index ade715d5d14..d4650813549 100644 --- a/litellm/accounting/metrics.py +++ b/litellm/accounting/metrics.py @@ -1,30 +1,55 @@ -from typing import Final - -from prometheus_client import Gauge, Histogram +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final 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" -) +if TYPE_CHECKING: + from prometheus_client import Gauge, Histogram + + +@dataclass(frozen=True, slots=True) +class _Instruments: + pending: "Gauge" + parked: "Gauge" + oldest: "Gauge" + batch_time: "Histogram" + batch_count: "Gauge" + table_bytes: "Gauge" + + +def _instruments() -> _Instruments | None: + try: + # prometheus-client ships only in the proxy-runtime extra, so a plain litellm[proxy] install lands here + from prometheus_client import Gauge, Histogram + except ImportError: + return None + return _Instruments( + pending=Gauge("litellm_accounting_pending", "Pending SpendLogs deliveries", multiprocess_mode="livemax"), + parked=Gauge("litellm_accounting_parked", "Parked SpendLogs deliveries", multiprocess_mode="livemax"), + oldest=Gauge( + "litellm_accounting_oldest_pending_seconds", + "Age of the oldest pending SpendLogs delivery", + multiprocess_mode="livemax", + ), + batch_time=Histogram("litellm_accounting_batch_seconds", "SpendLogs delivery worker batch time"), + batch_count=Gauge( + "litellm_accounting_batch_deliveries", "Retained batch deliveries", multiprocess_mode="livemax" + ), + table_bytes=Gauge( + "litellm_accounting_delivery_table_bytes", "Delivery table size in bytes", multiprocess_mode="livemax" + ), + ) + + +_INSTRUMENTS: Final = _instruments() 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) + if _INSTRUMENTS is None: + return + _INSTRUMENTS.pending.set(metrics.pending_count) + _INSTRUMENTS.parked.set(metrics.parked_count) + _INSTRUMENTS.oldest.set(max(0.0, metrics.oldest_pending_seconds)) + _INSTRUMENTS.batch_time.observe(seconds) + _INSTRUMENTS.batch_count.set(metrics.batch_count) + _INSTRUMENTS.table_bytes.set(metrics.table_bytes) diff --git a/litellm/accounting/snapshot.py b/litellm/accounting/snapshot.py index d5652f79ed0..ab85df08320 100644 --- a/litellm/accounting/snapshot.py +++ b/litellm/accounting/snapshot.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from litellm.router import Router _FIELDS: Final = TypeAdapter(dict[str, object]) +_JSON_SAFE_INTEGER: Final = 2**53 - 1 def _mapping(value: object) -> Mapping[str, object] | None: @@ -86,6 +87,7 @@ def build_accounting_payload( "metadata": accounting_metadata, }, non_finite_as_null=True, + integer_limit=_JSON_SAFE_INTEGER, ) ) diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py index df71f2da31c..be18d8ac6b1 100644 --- a/litellm/accounting/worker.py +++ b/litellm/accounting/worker.py @@ -14,6 +14,7 @@ from litellm.accounting.metrics import record_deliveries from litellm.accounting.projection import Contribution from litellm.repositories.accounting_delivery_repository import ( AccountingDeliveryRepository, + DeliveryBatchResult, DeliveryProjection, LedgerDatabase, LedgerTransaction, @@ -34,6 +35,14 @@ class _AccountingAlerts(Protocol): _JOB: Final = "accounting_delivery" +def _filled(result: DeliveryBatchResult, batch_size: int) -> bool: + return result.applied + len(result.parked) >= batch_size + + +def _progressed(result: DeliveryBatchResult) -> bool: + return bool(result.applied or result.failed) + + async def flush_accounting_deliveries( writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging" ) -> None: @@ -54,9 +63,9 @@ async def flush_accounting_deliveries( ) repository: Final = AccountingDeliveryRepository(database) started: Final = time.monotonic() - batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "10000")) + batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500")) - async def run_batch(*, retry_only: bool) -> bool: + async def run_batch(*, retry_only: bool) -> DeliveryBatchResult: 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) @@ -68,17 +77,25 @@ async def flush_accounting_deliveries( proxy_logging_obj, message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.", ) - return bool(result.applied or result.failed) + return result - await run_batch(retry_only=False) + fresh_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline: + break retry_deadline: Final = time.monotonic() + 5 for _ in range(batch_size): - if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True): + if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)): + break + retention_days: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")) + cleanup_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if ( + await repository.cleanup_applied(retention_days=retention_days, batch_size=batch_size * 2) + < batch_size * 2 + or time.monotonic() >= cleanup_deadline + ): 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( diff --git a/litellm/constants.py b/litellm/constants.py index 50e090eabaf..8d05115e76e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -28,7 +28,7 @@ GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth" GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type" GITHUB_COPILOT_HOST_KEY: Final = "github_copilot_host" GITHUB_COPILOT_DEFAULT_HOST: Final = "github.com" -GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS: Final = 60 +GITHUB_COPILOT_USER_SESSION_TTL_SECONDS: Final = 1800 GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60 GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow" USER_PROVIDER_CREDENTIAL_CACHE_PREFIX: Final = "user_provider_credential" diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 0579eb15baf..385522c2435 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -898,6 +898,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def redacts_messages_itself(self) -> bool: return False + def exempt_from_message_redaction(self) -> bool: + """Router pre-call checks that key routing off the request and never log it return True, so + `turn_off_message_logging` leaves the success event they receive unredacted.""" + return False + def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict: """ Redacts or excludes fields from StandardLoggingPayload before callbacks receive it. diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 5b81fd9a331..0f98a8fce0e 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -341,6 +341,9 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_SpendLogDelivery", "LiteLLM_ClaudeCodePluginTable", "LiteLLM_ClaudeCodePluginUpload", + "LiteLLM_ClaudeCodePluginVersion", + "LiteLLM_SkillSubmission", + "LiteLLM_ClaudeCodePluginDownload", "LiteLLM_MemoryTable", "LiteLLM_AdaptiveRouterState", "LiteLLM_AdaptiveRouterSession", diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 0b8c128286a..d3cc04408b6 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -71,12 +71,12 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj _TERMINAL_STATUSES: Final = frozenset( - {"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action"} + {"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action", "requires_continuation"} ) _POLLABLE_STATUSES: Final = frozenset({"in_progress", "queued"}) -_STATUSES_THAT_PRODUCED_OUTPUT: Final = frozenset({"completed", "requires_action"}) +_STATUSES_THAT_PRODUCED_OUTPUT: Final = frozenset({"completed", "requires_action", "requires_continuation"}) SettlementOutcome: TypeAlias = Literal["billed", "released", "unsettled"] diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index acd43c9b830..cd144646f79 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -95,6 +95,7 @@ async def acreate( response_mime_type: str | None = None, # Continuation previous_interaction_id: str | None = None, + continuation_token: str | None = None, # Extra params extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, @@ -129,6 +130,7 @@ async def acreate( response_format: JSON schema for response format response_mime_type: MIME type of the response previous_interaction_id: ID of previous interaction for continuation + continuation_token: Token from a `requires_continuation` interaction to resume its long decode extra_headers: Additional headers extra_body: Additional body parameters timeout: Request timeout @@ -164,6 +166,7 @@ async def acreate( response_format=response_format, response_mime_type=response_mime_type, previous_interaction_id=previous_interaction_id, + continuation_token=continuation_token, extra_headers=extra_headers, extra_body=extra_body, timeout=timeout, @@ -224,6 +227,7 @@ def create( response_mime_type: str | None = None, # Continuation previous_interaction_id: str | None = None, + continuation_token: str | None = None, # Extra params extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, @@ -262,6 +266,7 @@ def create( response_format: JSON schema for response format response_mime_type: MIME type of the response previous_interaction_id: ID of previous interaction for continuation + continuation_token: Token from a `requires_continuation` interaction to resume its long decode extra_headers: Additional headers extra_body: Additional body parameters timeout: Request timeout diff --git a/litellm/interactions/utils.py b/litellm/interactions/utils.py index 3895a85061d..1d7465a99c1 100644 --- a/litellm/interactions/utils.py +++ b/litellm/interactions/utils.py @@ -20,6 +20,7 @@ INTERACTIONS_API_OPTIONAL_PARAMS: Final = { "response_format", "response_mime_type", "previous_interaction_id", + "continuation_token", "agent_config", } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ce57a9bd4a5..9b6ec8d28d6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4074,12 +4074,11 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, ) - if ( - callable(callback) is True and is_sync_request and customLogger is not None - ): # custom logger functions + if callable(callback) is True and is_sync_request: # custom logger functions + function_logger = customLogger if customLogger is not None else CustomLogger() print_verbose(f"success callbacks: Running Custom Callback Function - {callback}") - customLogger.log_event( + function_logger.log_event( kwargs=call_details, response_obj=result, start_time=start_time, @@ -4406,7 +4405,11 @@ class Logging(LiteLLMLoggingBaseClass): if isinstance(callback, CustomLogger): # custom logger class from litellm.integrations.shadow_eval_logger import ShadowEvalLogger - model_call_details: dict = self.model_call_details + model_call_details: dict = ( + self._model_call_details_with_request_messages() + if callback.exempt_from_message_redaction() + else self.model_call_details + ) ################################## # call redaction hook for custom logger model_call_details = callback.redact_standard_logging_payload_from_model_call_details( @@ -4970,6 +4973,21 @@ class Logging(LiteLLMLoggingBaseClass): ) -> None: executor.submit(self.success_handler, result, start_time, end_time, cache_hit) + def _model_call_details_with_request_messages(self) -> dict[str, object]: # mutable-ok: loggers take a dict + """The success-event payload with the request messages put back where `turn_off_message_logging` + redacted them in place, for a callback that routes on the request and never logs it.""" + request_messages: Final = cast(object, self.messages) # cast-ok: the request messages are only put back as-is + standard_logging_object: Final = self.model_call_details.get("standard_logging_object") + return { + **self.model_call_details, + "messages": request_messages, + **( + {"standard_logging_object": {**standard_logging_object, "messages": request_messages}} + if isinstance(standard_logging_object, dict) + else {} + ), + } + def _get_completed_success_callbacks( self, *, asynchronous: bool ) -> tuple[str | Callable[..., object] | CustomLogger, ...]: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 4c3907dc52f..eda98ae2c67 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5879,7 +5879,7 @@ def _bedrock_tools_pt( description = name defs = parameters.pop("$defs", {}) - defs_copy = copy.deepcopy(defs) if isinstance(defs, dict) else {} + defs_copy = defs if isinstance(defs, dict) else {} # Expand $ref references in parameters using the definitions # Note: We don't pre-flatten defs as that causes exponential memory growth # with circular references (see issue #19098). unpack_defs handles nested diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 909e1222ae6..ea26d549326 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -19,6 +19,12 @@ def _finite_or_none(value: float | None) -> float | None: return None if isinstance(value, float) and not math.isfinite(value) else value +def _bounded(value: float | None, integer_limit: int | None) -> float | None: + if integer_limit is None or isinstance(value, bool) or not isinstance(value, int): + return value + return max(-integer_limit, min(integer_limit, value)) + + def safe_json_structure( data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, @@ -26,12 +32,14 @@ def safe_json_structure( key: str | None = None, *, non_finite_as_null: bool = False, + integer_limit: int | None = None, ) -> object: """ Rebuild data out of JSON-native pieces while detecting circular references. If a circular reference is detected then a marker string is returned. NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. With non_finite_as_null, NaN and infinite floats become null, since jsonb rejects them. + With integer_limit, integers are clamped to that magnitude. value_transform, when given, is applied to every string leaf (and to the str() fallback for non-serializable objects) with the mapping key the leaf @@ -51,7 +59,8 @@ def safe_json_structure( cleaned = obj.replace("\x00", "") if "\x00" in obj else obj return _transform(key, cleaned) if isinstance(obj, (int, float, bool, type(None))): - return _finite_or_none(obj) if non_finite_as_null else obj + bounded: Final = _bounded(obj, integer_limit) + return _finite_or_none(bounded) if non_finite_as_null else bounded # Check for circular reference. if id(obj) in seen: return "CircularReference Detected" diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 4f890ed922f..ac4addb0847 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -827,11 +827,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return cls._raise_invalid_reasoning_effort(model, value, llm_provider) def get_supported_openai_params(self, model: str): + sampling_params: Final = ("temperature", "top_p") if self._supports_sampling_params(model) else () params: Final = [ "stream", "stop", - "temperature", - "top_p", + *sampling_params, "max_tokens", "max_completion_tokens", "tools", diff --git a/litellm/llms/anthropic/messages/adapters/handler.py b/litellm/llms/anthropic/messages/adapters/handler.py index 31a4e9d9d73..da120cad0e1 100644 --- a/litellm/llms/anthropic/messages/adapters/handler.py +++ b/litellm/llms/anthropic/messages/adapters/handler.py @@ -123,10 +123,9 @@ def _extract_proxy_litellm_metadata( def _requested_router(kwargs: Mapping[str, object]) -> "Router | None": from litellm.router import Router - from litellm.router_python_core import PythonCore requested: Final = kwargs.get("litellm_router") - return requested if isinstance(requested, (Router, PythonCore)) else None # pyright: ignore[reportReturnType] # the router passes its core, which carries Router's members + return requested if isinstance(requested, Router) else None def request_safeguards_evaluator( diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index 1895a97af7d..50da73cc8b9 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -126,6 +126,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): "response_format", "response_mime_type", "previous_interaction_id", + "continuation_token", ] def validate_environment( @@ -206,6 +207,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): "environment", "response_modalities", "previous_interaction_id", + "continuation_token", ] dropped_keys: Final = ( frozenset({"agent_config"}) @@ -301,7 +303,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): ) -> InteractionsAPIStreamingResponse: """Parse streaming chunk.""" verbose_logger.debug("Google AI Interactions streaming chunk: %s", parsed_chunk) - return InteractionsAPIStreamingResponse(**parsed_chunk) + return InteractionsAPIStreamingResponse.model_validate(parsed_chunk) # GET / DELETE / CANCEL - just build URLs, responses match spec directly diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index ee0e7af7e9b..3af3b50ed8f 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -1,62 +1,30 @@ +import hashlib import json import os import time -from dataclasses import dataclass -from datetime import datetime -from typing import Any, Final, TypeVar +from collections.abc import Mapping +from typing import Final, Protocol, cast # noqa: TID251 # pins untyped clients and caches to local Protocols import httpx -from pydantic import BaseModel, ConfigDict, TypeAdapter, with_config +from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import GITHUB_COPILOT_USER_SESSION_TTL_SECONDS from litellm.litellm_core_utils.asyncify import can_block_current_thread from litellm.llms.custom_httpx.http_handler import get_httpx_client from .common_utils import ( GetAccessTokenError, - GetAPIKeyError, GetDeviceCodeError, - RefreshAPIKeyError, ) # Constants (default values — overridable via environment variables at call time) DEFAULT_GITHUB_CLIENT_ID: Final = "Iv1.b507a08c87ecfe98" DEFAULT_GITHUB_DEVICE_CODE_URL: Final = "https://github.com/login/device/code" DEFAULT_GITHUB_ACCESS_TOKEN_URL: Final = "https://github.com/login/oauth/access_token" -DEFAULT_GITHUB_API_KEY_URL: Final = "https://api.github.com/copilot_internal/v2/token" - - -class CopilotEndpoints(BaseModel): - model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True) - - api: str | None = None - - -class CopilotCachedKey(BaseModel): - model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True) - - token: str - expires_at: float - endpoints: CopilotEndpoints | None = None - - -class CopilotCachedEndpoints(BaseModel): - model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True) - - endpoints: CopilotEndpoints | None = None - - -ApiKeyFileView = TypeVar("ApiKeyFileView", CopilotCachedKey, CopilotCachedEndpoints) - - -@dataclass(frozen=True, slots=True) -class _UnsavedApiKey: - access_token: str - key: CopilotCachedKey - - -_UNSAVED_API_KEYS: Final[dict[str, _UnsavedApiKey]] = {} # mutable-ok: refreshed keys api-key.json could not hold +DEFAULT_GITHUB_COPILOT_USER_URL: Final = "https://api.github.com/copilot_internal/user" def _read_token_file(path: str) -> str: @@ -78,11 +46,66 @@ class _AccessTokenPoll(TypedDict): error: ReadOnly[NotRequired[object]] -_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_JSON_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter( + Mapping[str, object], config=ConfigDict(strict=True) +) _DEVICE_CODE: Final = TypeAdapter(_DeviceCode) _ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPoll) +class _CopilotAPIBaseCache(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: str, ttl: float | None = None) -> object: ... + + +class _SyncGetClient(Protocol): + def get(self, url: str, *, headers: Mapping[str, str]) -> httpx.Response: ... + + +_COPILOT_API_BASE_CACHE: Final[_CopilotAPIBaseCache] = cast( # cast-ok: [LIT006] InMemoryCache methods lack type stubs + _CopilotAPIBaseCache, + InMemoryCache( + max_size_in_memory=1000, + default_ttl=GITHUB_COPILOT_USER_SESSION_TTL_SECONDS, + ), +) + + +def _copilot_api_base_for_token(access_token: str) -> str | None: + cache_key: Final = hashlib.sha256(access_token.encode()).hexdigest() + cached_api_base: Final[object] = _COPILOT_API_BASE_CACHE.get_cache(cache_key) + if isinstance(cached_api_base, str): + return cached_api_base + + try: + response: Final = cast( # cast-ok: [LIT006] HTTPHandler.get has no return annotation + _SyncGetClient, + get_httpx_client(), + ).get( + DEFAULT_GITHUB_COPILOT_USER_URL, + headers=github_api_headers(access_token), + ) + if not response.is_success: + return None + payload: Final[Mapping[str, object]] = _JSON_OBJECT.validate_python(response.json()) + endpoints: Final[Mapping[str, object]] = _JSON_OBJECT.validate_python(payload.get("endpoints")) + api_base: Final[object] = endpoints.get("api") + if not isinstance(api_base, str): + return None + from .github_host import github_host_endpoints + from .per_user_auth import validated_copilot_api_base + + validated: Final = validated_copilot_api_base(api_base, github_host_endpoints("github.com")) + if validated != api_base.rstrip("/"): + return None + _COPILOT_API_BASE_CACHE.set_cache(cache_key, validated) + except Exception: + return None + else: + return validated + + def github_api_headers( access_token: str | None = None, ) -> dict[str, str]: # mutable-ok: returned straight into httpx handlers whose headers params require dict @@ -111,7 +134,6 @@ class Authenticator: self.token_dir, os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_FILE", "access-token"), ) - self.api_key_file = os.path.join(self.token_dir, os.getenv("GITHUB_COPILOT_API_KEY_FILE", "api-key.json")) self._ensure_token_dir() def get_access_token(self) -> str: @@ -155,7 +177,7 @@ class Authenticator: except OSError: verbose_logger.error("Error saving access token to file") return access_token - except (GetDeviceCodeError, GetAccessTokenError, RefreshAPIKeyError) as e: + except (GetDeviceCodeError, GetAccessTokenError) as e: verbose_logger.warning("Failed attempt %s: %s", attempt + 1, e) continue @@ -165,69 +187,15 @@ class Authenticator: ) def get_api_key(self) -> str: - """ - Get the API key, refreshing if necessary. - - Returns: - str: The GitHub Copilot API key. - - Raises: - GetAPIKeyError: If unable to obtain an API key. - """ - unsaved: Final = self._unsaved_api_key() - if unsaved is not None: - return unsaved.token - cached: Final = self._read_api_key_file(CopilotCachedKey) - if cached is not None and cached.expires_at > datetime.now().timestamp(): - return cached.token - if cached is not None: - verbose_logger.warning("API key expired, refreshing") - - try: - api_key_info: Final = self._refresh_api_key() - except RefreshAPIKeyError as e: - raise GetAPIKeyError( - message=f"Failed to refresh API key: {e}", - status_code=401, - ) from e - serialized: Final = json.dumps(api_key_info) - try: - with open(self.api_key_file, "w") as f: - f.write(serialized) - except OSError as e: - verbose_logger.warning( - "Keeping the refreshed API key in memory, it could not be saved to %s: %s", self.api_key_file, e - ) - self._keep_unsaved_api_key(serialized) - token: Final = api_key_info.get("token") - if not token: - raise GetAPIKeyError( - message="API key response missing token", - status_code=401, - ) - return token + """Return the stored GitHub access token, loading it with device flow when needed.""" + return self.get_access_token() def get_api_base(self) -> str | None: - """ - Get the API endpoint of the current API key. - - Returns: - Optional[str]: The GitHub Copilot API endpoint, or None if not found. - """ - cached: Final = self._unsaved_api_key() or self._read_api_key_file(CopilotCachedEndpoints) - if cached is None or cached.endpoints is None: - return None - return cached.endpoints.api - - def _read_api_key_file(self, view: type[ApiKeyFileView]) -> ApiKeyFileView | None: - try: - return view.model_validate_json(_read_token_file(self.api_key_file)) - except OSError: - verbose_logger.warning("No API key file found or error opening file") - return None - except ValueError as e: - verbose_logger.warning("Ignoring %s, not a Copilot API key file: %s", self.api_key_file, e) + """Return the Copilot API endpoint for a stored token without starting device flow.""" + access_token: Final = self._stored_access_token() + if access_token is None: return None + return _copilot_api_base_for_token(access_token) def _stored_access_token(self) -> str | None: try: @@ -236,59 +204,6 @@ class Authenticator: return None return access_token or None - def _unsaved_api_key(self) -> CopilotCachedKey | None: - unsaved: Final = _UNSAVED_API_KEYS.get(self.api_key_file) - if unsaved is None or unsaved.key.expires_at <= time.time(): - return None - return unsaved.key if unsaved.access_token == self._stored_access_token() else None - - def _keep_unsaved_api_key(self, serialized: str) -> None: - access_token: Final = self._stored_access_token() - if access_token is None: - return - try: - key: Final = CopilotCachedKey.model_validate_json(serialized) - except ValueError: - return - _UNSAVED_API_KEYS[self.api_key_file] = _UnsavedApiKey(access_token=access_token, key=key) - - def _refresh_api_key(self) -> dict[str, Any]: - """ - Refresh the API key using the access token. - - Returns: - Dict[str, Any]: The API key information including token and expiration. - - Raises: - RefreshAPIKeyError: If unable to refresh the API key. - """ - access_token: Final = self.get_access_token() - headers: Final = self._get_github_headers(access_token) - api_key_url: Final = os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL) - - max_retries: Final = 3 - for attempt in range(max_retries): - try: - sync_client = get_httpx_client() - response = sync_client.get(api_key_url, headers=headers) - response.raise_for_status() - - response_json = response.json() - - if "token" in response_json: - return response_json - else: - verbose_logger.warning("API key response missing token: %s", response_json) - except httpx.HTTPStatusError as e: - verbose_logger.error("HTTP error refreshing API key (attempt %s/%s): %s", attempt + 1, max_retries, e) - except Exception as e: - verbose_logger.error("Unexpected error refreshing API key: %s", e) - - raise RefreshAPIKeyError( - message="Failed to refresh API key after maximum retries", - status_code=401, - ) - def _ensure_token_dir(self) -> None: """Ensure the token directory exists.""" if not os.path.exists(self.token_dir): diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 83bc8df23a6..0285bde2a87 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -16,7 +16,7 @@ from litellm.types.utils import ModelResponse from ..authenticator import Authenticator from ..common_utils import ( DEFAULT_GITHUB_COPILOT_API_BASE, - GetAPIKeyError, + GetAccessTokenError, get_copilot_default_headers, ) from ..per_user_auth import require_github_copilot_user_session @@ -56,7 +56,7 @@ class GithubCopilotConfig(OpenAIConfig): dynamic_api_base: Final = self.api_base_without_login(api_base) try: dynamic_api_key: Final = self.authenticator.get_api_key() - except GetAPIKeyError as e: + except GetAccessTokenError as e: raise AuthenticationError( model=model, llm_provider=custom_llm_provider, @@ -104,7 +104,7 @@ class GithubCopilotConfig(OpenAIConfig): return get_copilot_default_headers(session_token) try: return get_copilot_default_headers(self.authenticator.get_api_key()) - except GetAPIKeyError: + except GetAccessTokenError: return {} def validate_environment( diff --git a/litellm/llms/github_copilot/common_utils.py b/litellm/llms/github_copilot/common_utils.py index 515b7bf4fb5..fcadb0cb2ef 100644 --- a/litellm/llms/github_copilot/common_utils.py +++ b/litellm/llms/github_copilot/common_utils.py @@ -46,18 +46,6 @@ class GetAccessTokenError(GithubCopilotError): pass -class APIKeyExpiredError(GithubCopilotError): - pass - - -class RefreshAPIKeyError(GithubCopilotError): - pass - - -class GetAPIKeyError(GithubCopilotError): - pass - - def get_copilot_default_headers( api_key: str, ) -> dict[str, str]: # mutable-ok: callers merge caller headers into the returned dict diff --git a/litellm/llms/github_copilot/embedding/transformation.py b/litellm/llms/github_copilot/embedding/transformation.py index 4f308b01159..740397f0419 100644 --- a/litellm/llms/github_copilot/embedding/transformation.py +++ b/litellm/llms/github_copilot/embedding/transformation.py @@ -30,7 +30,7 @@ from litellm.utils import convert_to_model_response_object from ..authenticator import Authenticator from ..common_utils import ( DEFAULT_GITHUB_COPILOT_API_BASE, - GetAPIKeyError, + GetAccessTokenError, get_copilot_default_headers, pin_session_authorization, ) @@ -100,7 +100,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): return merged_headers - except GetAPIKeyError as e: + except GetAccessTokenError as e: raise AuthenticationError( model=model, llm_provider="github_copilot", diff --git a/litellm/llms/github_copilot/github_host.py b/litellm/llms/github_copilot/github_host.py index cfe122ebac4..09a5a79575a 100644 --- a/litellm/llms/github_copilot/github_host.py +++ b/litellm/llms/github_copilot/github_host.py @@ -7,8 +7,8 @@ from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST from .authenticator import ( DEFAULT_GITHUB_ACCESS_TOKEN_URL, - DEFAULT_GITHUB_API_KEY_URL, DEFAULT_GITHUB_CLIENT_ID, + DEFAULT_GITHUB_COPILOT_USER_URL, DEFAULT_GITHUB_DEVICE_CODE_URL, ) @@ -20,11 +20,9 @@ class GithubHostEndpoints: host: str device_code_url: str access_token_url: str - api_key_url: str copilot_user_url: str user_url: str client_id: str - exchanges_copilot_token: bool copilot_api_hosts: frozenset[str] copilot_api_host_suffix: str | None default_copilot_api_base: str @@ -51,17 +49,14 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints: if normalized is None: raise ValueError("host must be github.com or a single-subdomain GHE.com hostname") client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) - exchanges_copilot_token: Final = client_id == DEFAULT_GITHUB_CLIENT_ID if normalized == GITHUB_COPILOT_DEFAULT_HOST: return GithubHostEndpoints( host=normalized, device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL), access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL), - api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL), - copilot_user_url="https://api.github.com/copilot_internal/user", + copilot_user_url=DEFAULT_GITHUB_COPILOT_USER_URL, user_url="https://api.github.com/user", client_id=client_id, - exchanges_copilot_token=exchanges_copilot_token, copilot_api_hosts=frozenset({"githubcopilot.com"}), copilot_api_host_suffix=".githubcopilot.com", default_copilot_api_base="https://api.githubcopilot.com", @@ -74,11 +69,9 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints: host=normalized, device_code_url=f"https://{normalized}/login/device/code", access_token_url=f"https://{normalized}/login/oauth/access_token", - api_key_url=f"https://{api_host}/copilot_internal/v2/token", copilot_user_url=f"https://{api_host}/copilot_internal/user", user_url=f"https://{api_host}/user", client_id=client_id, - exchanges_copilot_token=exchanges_copilot_token, copilot_api_hosts=frozenset({copilot_api_host}), copilot_api_host_suffix=None, default_copilot_api_base=f"https://{copilot_api_host}", diff --git a/litellm/llms/github_copilot/messages/transformation.py b/litellm/llms/github_copilot/messages/transformation.py index bc2b68cd02b..bed855ca178 100644 --- a/litellm/llms/github_copilot/messages/transformation.py +++ b/litellm/llms/github_copilot/messages/transformation.py @@ -10,7 +10,7 @@ from litellm.llms.anthropic.messages.native.transformation import ( from ..authenticator import Authenticator from ..common_utils import ( DEFAULT_GITHUB_COPILOT_API_BASE, - GetAPIKeyError, + GetAccessTokenError, get_copilot_default_headers, pin_session_authorization, ) @@ -84,7 +84,7 @@ class GithubCopilotAnthropicMessagesConfig(AnthropicMessagesConfig): dynamic_api_key: Final = ( user_session.token if user_session is not None else self.authenticator.get_api_key() ) - except GetAPIKeyError as e: + except GetAccessTokenError as e: raise AuthenticationError( model=model, llm_provider="github_copilot", diff --git a/litellm/llms/github_copilot/per_user_auth.py b/litellm/llms/github_copilot/per_user_auth.py index e210414ce8f..54c64c48084 100644 --- a/litellm/llms/github_copilot/per_user_auth.py +++ b/litellm/llms/github_copilot/per_user_auth.py @@ -1,7 +1,6 @@ -"""Per-user GitHub Copilot OAuth: exchange a caller's stored GitHub token for a -short-lived Copilot token, plus the device-flow helpers the proxy endpoints use -to establish that connection. The shared device login (Authenticator) is never -touched on the per-user path. +"""Per-user GitHub Copilot OAuth sessions and the device-flow helpers the proxy +endpoints use to establish those connections. The shared device login +(Authenticator) is never touched on the per-user path. """ import asyncio @@ -28,7 +27,7 @@ from litellm.constants import ( GITHUB_COPILOT_AUTH_TYPE_KEY, GITHUB_COPILOT_HOST_KEY, GITHUB_COPILOT_PER_USER_AUTH_TYPE, - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS, + GITHUB_COPILOT_USER_SESSION_TTL_SECONDS, ) from litellm.exceptions import ( APIConnectionError, @@ -39,6 +38,7 @@ from litellm.exceptions import ( ) from .authenticator import github_api_headers +from .common_utils import get_copilot_default_headers from .github_host import GithubHostEndpoints, github_host_endpoints, normalized_github_host GITHUB_COPILOT_USER_SESSION_KWARG_KEY: Final = "github_copilot_user_session" @@ -46,7 +46,6 @@ _LLM_PROVIDER: Final = "github_copilot" _DEVICE_FLOW_SCOPE: Final = "read:user" _SESSION_CACHE_MAX_SIZE: Final = 1000 _SYNC_LOCK_COUNT: Final = 64 -_ACCOUNT_LOOKUP_TTL_SECONDS: Final = 1800 class _SessionCache(Protocol): @@ -88,19 +87,19 @@ _SESSION_CACHE: Final = cast( # cast-ok: pins the untyped module cache to _Sess _SessionCache, InMemoryCache( max_size_in_memory=_SESSION_CACHE_MAX_SIZE, - default_ttl=GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS, + default_ttl=GITHUB_COPILOT_USER_SESSION_TTL_SECONDS, ), ) -_EXCHANGE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT)) +_RESOLVE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT)) _IN_FLIGHT: Final[ # mutable-ok: single-flight dedup registry mutated under its own locking dict[str, tuple[str, "asyncio.Future[GithubCopilotUserSession]"]] ] = {} @with_config(ConfigDict(extra="allow", strict=True)) -class _CopilotTokenPayload(TypedDict): - token: ReadOnly[str] - expires_at: ReadOnly[NotRequired[object]] +class _CopilotUserPayload(TypedDict): + access_type_sku: ReadOnly[NotRequired[object]] + chat_enabled: ReadOnly[NotRequired[object]] endpoints: ReadOnly[NotRequired[object]] @@ -126,20 +125,12 @@ class _CopilotEndpointsPayload(TypedDict): api: ReadOnly[NotRequired[object]] -@with_config(ConfigDict(extra="allow", strict=True)) -class _CopilotAccountPayload(TypedDict): - access_type_sku: ReadOnly[NotRequired[object]] - chat_enabled: ReadOnly[NotRequired[object]] - endpoints: ReadOnly[NotRequired[object]] - - @with_config(ConfigDict(extra="allow", strict=True)) class _GithubUserPayload(TypedDict): login: ReadOnly[str] -_COPILOT_TOKEN: Final = TypeAdapter(_CopilotTokenPayload) -_COPILOT_ACCOUNT: Final = TypeAdapter(_CopilotAccountPayload) +_COPILOT_USER: Final = TypeAdapter(_CopilotUserPayload) _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload) _ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload) _GITHUB_USER: Final = TypeAdapter(_GithubUserPayload) @@ -237,13 +228,13 @@ def _cached_session(cache_key: str, github_token: str) -> GithubCopilotUserSessi return None if not hmac.compare_digest(cached_github_token, github_token): return None - if expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS <= time.time(): + if expires_at <= time.time(): return None return GithubCopilotUserSession(token=token, api_base=api_base) def _cache_session(cache_key: str, github_token: str, session: GithubCopilotUserSession, expires_at: float) -> None: - ttl: Final = expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS - time.time() + ttl: Final = expires_at - time.time() if ttl <= 0: return _SESSION_CACHE.set_cache(cache_key, (github_token, session.token, session.api_base, expires_at), ttl=ttl) @@ -272,57 +263,32 @@ def _session_from_response( raise _reject_connection(credential_name) if response.status_code == 429: raise CallerCredentialRateLimitError( - message="GitHub rate limited the GitHub Copilot token exchange", + message="GitHub rate limited the GitHub Copilot user lookup", llm_provider=_LLM_PROVIDER, model="", ) if not response.is_success: raise ServiceUnavailableError( - message=f"GitHub Copilot token exchange failed with status {response.status_code}", + message=f"GitHub Copilot user lookup failed with status {response.status_code}", llm_provider=_LLM_PROVIDER, model="", ) - if not endpoints.exchanges_copilot_token: - return _session_from_account(response, github_token, credential_name, cache_key, endpoints) try: - payload: Final = _COPILOT_TOKEN.validate_python(response.json()) + payload: Final = _COPILOT_USER.validate_python(response.json()) except Exception as e: raise ServiceUnavailableError( - message="GitHub Copilot token exchange returned an invalid response", + message="GitHub Copilot user lookup returned an invalid response", llm_provider=_LLM_PROVIDER, model="", ) from e - api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints) - session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base) - expires_at: Final = _expires_at_seconds(payload.get("expires_at")) - if expires_at > 0: - _cache_session(cache_key, github_token, session, expires_at) - return session - - -def _session_from_account( - response: httpx.Response, - github_token: str, - credential_name: str, - cache_key: str, - endpoints: GithubHostEndpoints, -) -> GithubCopilotUserSession: - try: - payload: Final = _COPILOT_ACCOUNT.validate_python(response.json()) - except Exception as e: - raise ServiceUnavailableError( - message="GitHub Copilot account lookup returned an invalid response", - llm_provider=_LLM_PROVIDER, - model="", - ) from e if payload.get("chat_enabled") is False or payload.get("access_type_sku") == "no_access": _SESSION_CACHE.delete_cache(cache_key) raise _reject_connection(credential_name) api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints) session: Final = GithubCopilotUserSession(token=github_token, api_base=api_base) - _cache_session(cache_key, github_token, session, time.time() + _ACCOUNT_LOOKUP_TTL_SECONDS) + _cache_session(cache_key, github_token, session, time.time() + GITHUB_COPILOT_USER_SESSION_TTL_SECONDS) return session @@ -332,10 +298,6 @@ def _endpoints_api(payload_endpoints: object) -> object: return _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api") -def _session_lookup_url(endpoints: GithubHostEndpoints) -> str: - return endpoints.api_key_url if endpoints.exchanges_copilot_token else endpoints.copilot_user_url - - def _to_int(raw: object) -> int: if isinstance(raw, bool): return 0 @@ -351,20 +313,7 @@ def _to_int(raw: object) -> int: return 0 -def _expires_at_seconds(raw: object) -> float: - if isinstance(raw, bool): - return 0.0 - if isinstance(raw, (int, float)): - return float(raw) - if isinstance(raw, str): - try: - return float(raw) - except ValueError: - return 0.0 - return 0.0 - - -def _do_exchange( +def _do_resolve( github_token: str, credential_name: str, cache_key: str, @@ -377,12 +326,12 @@ def _do_exchange( _SyncGetClient, litellm.module_level_client ) response: Final = client.get( - _session_lookup_url(endpoints), + endpoints.copilot_user_url, headers=github_api_headers(github_token), ) except Exception as e: raise APIConnectionError( - message="GitHub Copilot token exchange request failed", + message="GitHub Copilot user lookup request failed", llm_provider=_LLM_PROVIDER, model="", ) from e @@ -398,7 +347,7 @@ def _new_async_github_client() -> _AsyncGitHubClient: )(_LLM_PROVIDER) -async def _do_aexchange( +async def _do_aresolve( github_token: str, credential_name: str, cache_key: str, @@ -407,19 +356,19 @@ async def _do_aexchange( client: Final = _new_async_github_client() try: response: Final = await client.get( - _session_lookup_url(endpoints), + endpoints.copilot_user_url, headers=github_api_headers(github_token), ) except Exception as e: raise APIConnectionError( - message="GitHub Copilot token exchange request failed", + message="GitHub Copilot user lookup request failed", llm_provider=_LLM_PROVIDER, model="", ) from e return _session_from_response(response, github_token, credential_name, cache_key, endpoints) -def exchange_github_token( +def resolve_copilot_user_session( user_id: str, github_token: str, credential_name: str, @@ -429,15 +378,15 @@ def exchange_github_token( cached: Final = _cached_session(cache_key, github_token) if cached is not None: return cached - lock: Final = _EXCHANGE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT] + lock: Final = _RESOLVE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT] with lock: rechecked: Final = _cached_session(cache_key, github_token) if rechecked is not None: return rechecked - return _do_exchange(github_token, credential_name, cache_key, endpoints) + return _do_resolve(github_token, credential_name, cache_key, endpoints) -async def aexchange_github_token( +async def aresolve_copilot_user_session( user_id: str, github_token: str, credential_name: str, @@ -452,14 +401,14 @@ async def aexchange_github_token( if in_flight_entry is not None: in_flight_token, in_flight = in_flight_entry if not hmac.compare_digest(in_flight_token, github_token): - return await _do_aexchange(github_token, credential_name, cache_key, endpoints) + return await _do_aresolve(github_token, credential_name, cache_key, endpoints) if not in_flight.done() and in_flight.get_loop() is loop: return await asyncio.shield(in_flight) future: Final = loop.create_future() entry: Final = (github_token, future) _IN_FLIGHT[cache_key] = entry try: - session: Final = await _do_aexchange(github_token, credential_name, cache_key, endpoints) + session: Final = await _do_aresolve(github_token, credential_name, cache_key, endpoints) except BaseException as e: if not future.done(): future.set_exception(e) @@ -610,12 +559,41 @@ async def acheck_copilot_seat( credential_name: str, endpoints: GithubHostEndpoints, ) -> GithubCopilotUserSession: - return await aexchange_github_token( + session: Final = await aresolve_copilot_user_session( user_id=user_id, github_token=github_token, credential_name=credential_name, endpoints=endpoints, ) + client: Final = _new_async_github_client() + models_url: Final = f"{session.api_base}/models" + try: + response: Final = await client.get( + models_url, + headers=get_copilot_default_headers(session.token), + ) + except Exception as e: + raise APIConnectionError( + message="GitHub Copilot models check request failed", + llm_provider=_LLM_PROVIDER, + model="", + ) from e + if response.status_code in (401, 403, 404): + _SESSION_CACHE.delete_cache(_session_cache_key(user_id, credential_name, endpoints.host)) + raise _reject_connection(credential_name) + if response.status_code == 429: + raise CallerCredentialRateLimitError( + message="GitHub rate limited the GitHub Copilot models check", + llm_provider=_LLM_PROVIDER, + model="", + ) + if not response.is_success: + raise ServiceUnavailableError( + message=f"GitHub Copilot models check failed with status {response.status_code}", + llm_provider=_LLM_PROVIDER, + model="", + ) + return session def github_copilot_auth_mode(litellm_credential_name: object, auth_type_value: object) -> bool: @@ -804,7 +782,7 @@ def attach_github_copilot_user_session( credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else "" endpoints: Final = github_copilot_host_endpoints(credential_name) user_id, github_token = _caller_connection(kwargs, credential_name, endpoints) - session: Final = exchange_github_token( + session: Final = resolve_copilot_user_session( user_id=user_id, github_token=github_token, credential_name=credential_name, @@ -826,7 +804,7 @@ async def aattach_github_copilot_user_session( credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else "" endpoints: Final = github_copilot_host_endpoints(credential_name) user_id, github_token = _caller_connection(kwargs, credential_name, endpoints) - session: Final = await aexchange_github_token( + session: Final = await aresolve_copilot_user_session( user_id=user_id, github_token=github_token, credential_name=credential_name, diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index b87a3d543e7..f337d4ae0d3 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -38,7 +38,7 @@ from litellm.utils import cached_get_model_info_helper from ..authenticator import Authenticator from ..common_utils import ( DEFAULT_GITHUB_COPILOT_API_BASE, - GetAPIKeyError, + GetAccessTokenError, get_copilot_default_headers, pin_session_authorization, ) @@ -256,7 +256,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): return merged_headers - except GetAPIKeyError as e: + except GetAccessTokenError as e: raise AuthenticationError( model=model, llm_provider="github_copilot", diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 87c944a642a..f330d7fc6d4 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -1,3 +1,4 @@ +import inspect import time import types from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence @@ -379,6 +380,11 @@ _EXTRA_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str] | None) _EXTRA_QUERY_ADAPTER: Final = TypeAdapter(dict[str, object] | None) _NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({}) _SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body")) +_CLIENT_API_KEY_ONLY_AUTH: Final[Mapping[str, Mapping[str, bool]]] = ( + types.MappingProxyType({"security": types.MappingProxyType({"bearer_auth": True})}) + if "security" in inspect.signature(make_request_options).parameters + else types.MappingProxyType({}) +) def _raw_chat_response( @@ -398,6 +404,7 @@ def _request_without_sdk_transform( extra_query=_EXTRA_QUERY_ADAPTER.validate_python(data.get("extra_query")), extra_body=data.get("extra_body"), timeout=timeout, + **_CLIENT_API_KEY_ONLY_AUTH, # pyright: ignore[reportArgumentType] # openai<3 make_request_options has no security kwarg ) return body, options diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 76fd85d001b..a29f97ee9ea 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4046,6 +4046,8 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 2.5e-07, + "deprecation_date": "2027-09-01", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "litellm_provider": "azure_ai", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -4068,7 +4070,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, "prompt_cache_min_tokens": 512, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models" + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" }, "azure_ai/claude-opus-5": { "deprecation_date": "2027-07-08", @@ -13458,13 +13460,16 @@ "supports_reasoning": true }, "azure_ai/mistral-document-ai-2512": { + "deprecation_date": "2027-02-15", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.003, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" ], - "source": "https://ai.azure.com/catalog/models/mistral-document-ai-2512" + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", + "successor_model": "azure_ai/mistral-ocr-4-0" }, "azure_ai/mistral-ocr-4-0": { "litellm_provider": "azure_ai", @@ -32320,7 +32325,9 @@ "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", "successor_model": "gemini/gemini-3.8-flash", "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, - "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 + "input_cost_per_token_above_200k_tokens_flex": 1.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, + "output_cost_per_token_above_200k_tokens_flex": 7.5e-06 }, "gemini/gemini-2.5-computer-use-preview-10-2025": { "deprecation_date": "2026-07-28", @@ -32811,7 +32818,7 @@ "output_cost_per_token": 9e-06, "output_cost_per_video_token": 1.75e-05, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_endpoints": [ "/v1beta/interactions" ], @@ -32833,7 +32840,7 @@ "tpm": 800000, "deprecation_date": "2026-10-22", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", - "successor_model": "gemini/gemini-omni-1.1-flash" + "successor_model": "gemini/gemini-omni-1.2-flash" }, "gemini/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, @@ -32899,7 +32906,9 @@ "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, - "output_cost_per_token_above_200k_tokens_batches": 9e-06 + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -32965,7 +32974,9 @@ "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, - "output_cost_per_token_above_200k_tokens_batches": 9e-06 + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -33613,14 +33624,14 @@ "gemini/veo-3.1-fast-generate-preview": { "deprecation_date": "2026-10-22", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", - "successor_model": "gemini/gemini-omni-1.1-flash", + "successor_model": "gemini/gemini-omni-1.2-flash", "litellm_provider": "gemini", "max_input_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, "output_cost_per_second_1080p": 0.12, "output_cost_per_second_4k": 0.3, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_modalities": [ "text" ], @@ -33631,13 +33642,13 @@ "gemini/veo-3.1-generate-preview": { "deprecation_date": "2026-10-22", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", - "successor_model": "gemini/gemini-omni-1.1-flash", + "successor_model": "gemini/gemini-omni-1.2-flash", "litellm_provider": "gemini", "max_input_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, "output_cost_per_second_4k": 0.6, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_modalities": [ "text" ], @@ -33648,13 +33659,13 @@ "gemini/veo-3.1-lite-generate-preview": { "deprecation_date": "2026-10-22", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", - "successor_model": "gemini/gemini-omni-1.1-flash", + "successor_model": "gemini/gemini-omni-1.2-flash", "litellm_provider": "gemini", "max_input_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.05, "output_cost_per_second_1080p": 0.08, - "source": "https://ai.google.dev/gemini-api/docs/video", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_modalities": [ "text" ], @@ -40887,6 +40898,22 @@ "output_cost_per_token": 4e-07, "supports_tool_choice": true }, + "jina-reranker-v3.5": { + "input_cost_per_token": 5e-08, + "litellm_provider": "jina_ai", + "max_input_tokens": 131072, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" + }, + "jina-reranker-m0": { + "input_cost_per_token": 5e-08, + "litellm_provider": "jina_ai", + "max_input_tokens": 10240, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" + }, "jina-reranker-v2-base-multilingual": { "input_cost_per_token": 5e-08, "litellm_provider": "jina_ai", @@ -43765,7 +43792,8 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/zai-glm-5-3": { "cache_read_input_token_cost": 1.4e-07, @@ -45264,6 +45292,18 @@ "source": "https://tokenfactory.nebius.com/api/public/models_info", "cache_read_input_token_cost": 1.5e-07 }, + "nebius/moonshotai/Kimi-K3.1-Pro": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "nebius", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "supports_prompt_caching": true, + "source": "https://tokenfactory.nebius.com/api/public/models_info" + }, "nebius/moonshotai/Kimi-K3": { "max_input_tokens": 1048576, "max_output_tokens": 1024000, @@ -60260,6 +60300,7 @@ "rpm": 10, "gemini_audio_only_live": true, "input_cost_per_second": 8.33333333333e-05, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_response_schema": false, "supports_reasoning": true }, @@ -67653,17 +67694,20 @@ "source": "https://deepinfra.com/pricing" }, "gemini/gemini-omni-1.1-flash": { + "deprecation_date": "2026-12-03", + "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, "output_cost_per_video_token": 1.75e-05, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/models/gemini-omni-flash", + "successor_model": "gemini/gemini-omni-1.2-flash", "supported_endpoints": [ "/v1beta/interactions" ], @@ -71179,6 +71223,8 @@ }, "gpt-5.5-cyber-preview": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-02-01", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "input_cost_per_token": 1.25e-05, "litellm_provider": "openai", "mode": "chat", @@ -72138,8 +72184,10 @@ "output_cost_per_token": 4.5e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_audio_output": true, "tpm": 250000, "rpm": 10, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -72159,8 +72207,10 @@ "output_cost_per_token": 4.5e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_audio_output": true, "tpm": 250000, "rpm": 10, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -76457,6 +76507,33 @@ "supports_tool_choice": true, "supports_vision": true }, + "us-gov.xai.grok-4.7": { + "supports_regex_lookaround": false, + "supported_endpoints": ["/v1/chat/completions"], + "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_response_format": true, + "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_converse", + "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", + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "baseten/deepseek-ai/DeepSeek-V4.1-Flash-Fast": { "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 6e-07, @@ -78425,7 +78502,7 @@ "cache_creation_input_token_cost": 1e-05, "cache_creation_input_token_cost_above_1hr": 1.6e-05, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2027-10-01", + "deprecation_date": "2027-01-15", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "input_cost_per_token": 8e-06, "litellm_provider": "azure_ai", @@ -78435,7 +78512,7 @@ "output_cost_per_token": 4e-05, "prompt_cache_min_tokens": 512.0, "search_context_cost_per_query": {"search_context_size_high":0.02,"search_context_size_low":0.02,"search_context_size_medium":0.02}, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models", + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -82770,6 +82847,39 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true }, + "gpt-6-cyber": { + "cache_creation_input_token_cost": 2.5e-05, + "cache_creation_input_token_cost_above_12hr": 5e-05, + "cache_read_input_token_cost": 2e-06, + "input_cost_per_token": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.0001, + "source": "https://developers.openai.com/api/docs/models/gpt-6-cyber", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true + }, "gpt-6-astra-law": { "cache_read_input_token_cost": 1e-06, "input_cost_per_token": 1e-05, @@ -83130,6 +83240,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_prompt_caching": true, "supports_function_calling": true, "supports_response_schema": true, @@ -83146,6 +83261,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_prompt_caching": true, "supports_function_calling": true, "supports_response_schema": true, @@ -83155,6 +83275,7 @@ "litellm_provider": "empiriolabs", "mode": "chat", "max_input_tokens": 202745, + "max_output_tokens": 128000, "input_cost_per_token": 8.25e-07, "output_cost_per_token": 3.301e-06, "input_cost_per_token_above_32k_tokens": 1.1e-06, @@ -83178,6 +83299,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-2" @@ -83230,6 +83356,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-3" @@ -83245,6 +83376,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-3-flash" @@ -83342,6 +83478,28 @@ "supports_response_schema": true, "source": "https://empiriolabs.ai/models/kimi-k3" }, + "empiriolabs/kimi-k3-1-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1.8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/kimi-k3-1-pro" + }, "empiriolabs/mimo-v2-6-flash": { "litellm_provider": "empiriolabs", "mode": "chat", @@ -83721,6 +83879,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" }, @@ -83737,6 +83900,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-27b" }, @@ -83753,6 +83921,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" }, @@ -83769,6 +83942,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" }, @@ -83830,6 +84008,10 @@ "max_output_tokens": 65536, "input_cost_per_token": 2.9e-08, "output_cost_per_token": 2.87e-07, + "input_cost_per_token_above_128k_tokens": 1.15e-07, + "output_cost_per_token_above_128k_tokens": 1.147e-06, + "input_cost_per_token_above_256k_tokens": 1.72e-07, + "output_cost_per_token_above_256k_tokens": 1.72e-06, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -83852,6 +84034,11 @@ "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, "supports_audio_input": true, "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" }, @@ -83865,6 +84052,11 @@ "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, "supports_response_schema": true, "supports_audio_input": true, "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" @@ -83924,6 +84116,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-6-27b" }, @@ -84421,7 +84618,7 @@ "max_output_tokens": 131072, "input_cost_per_token": 1e-06, "output_cost_per_token": 2.7e-06, - "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost": 1e-07, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -85310,6 +85507,24 @@ "output_cost_per_token": 6.09e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/XiaomiMiMo/MiMo-V2.6-Flash": { + "cache_read_input_token_cost": 2.24e-09, + "input_cost_per_token": 1.12e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 2.24e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.6-Pro": { + "cache_read_input_token_cost": 2.88e-09, + "input_cost_per_token": 3.48e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6.96e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/deepseek-ai/DeepSeek-Prover-V2-671B": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -85434,6 +85649,15 @@ "output_cost_per_token": 3.168e-06, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/deepseek-ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 4.4e-07, @@ -85443,6 +85667,24 @@ "output_cost_per_token": 1.32e-06, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/deepseek-ai/deepseek-v4.1-flash-1017": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/gmicloud/Ripple-65B-Preview": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/google/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, @@ -85759,6 +86001,14 @@ "output_cost_per_token": 2.5e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/openai/gpt-oss-20b": { + "input_cost_per_token": 4e-08, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/stepfun-ai/Step-3.7-Flash": { "cache_read_input_token_cost": 4e-08, "input_cost_per_token": 2e-07, @@ -85803,6 +86053,15 @@ "output_cost_per_token": 6e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/tencent/hy4": { + "cache_read_input_token_cost": 2.1e-08, + "input_cost_per_token": 4.17e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.2505e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/tencent/hy4-preview": { "cache_read_input_token_cost": 4.2e-08, "input_cost_per_token": 8.34e-07, @@ -86176,5 +86435,100 @@ "output_cost_per_token": 6e-07, "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true + }, + "vertex_ai/gemini-omni-1.2-flash": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 260640, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": false, + "supports_audio_output": false, + "supports_computer_use": false, + "supports_function_calling": false, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_tool_choice": false, + "supports_url_context": false, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": false + }, + "gemini/gemini-omni-1.2-flash": { + "input_cost_per_image_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "input_cost_per_video_token": 1.5e-06, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.2e-05, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "video", + "audio" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_reasoning": true, + "supports_system_messages": false, + "supports_video_input": true, + "supports_vision": true + }, + "databricks/databricks-kimi-k3-1-pro": { + "cache_creation_input_token_cost": 1.99997e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "metadata": {"notes":"Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference."}, + "mode": "chat", + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.000142857, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 100.0 } } diff --git a/litellm/models/plugin_upload.py b/litellm/models/plugin_upload.py index 21a55d44960..bd5e85c79cc 100644 --- a/litellm/models/plugin_upload.py +++ b/litellm/models/plugin_upload.py @@ -1,12 +1,25 @@ 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) class PluginUpload: name: str @@ -27,13 +40,126 @@ class PluginArchiveMetadata: 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]: ... diff --git a/litellm/models/skill_submission.py b/litellm/models/skill_submission.py new file mode 100644 index 00000000000..09d722fa69f --- /dev/null +++ b/litellm/models/skill_submission.py @@ -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: ... diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index c8fd9d860a1..01c22dc6725 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -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 @@ -287,6 +296,7 @@ LAZY_FEATURES: Final[tuple[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", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 9a61a2d3f7a..8e7a172cb10 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -7689,6 +7689,145 @@ "title": "PluginAuthor", "type": "object" }, + "PluginDownloadEvent": { + "properties": { + "authenticated": { + "title": "Authenticated", + "type": "boolean" + }, + "digest": { + "title": "Digest", + "type": "string" + }, + "downloaded_at": { + "format": "date-time", + "title": "Downloaded At", + "type": "string" + }, + "id": { + "title": "Id", + "type": "string" + }, + "key_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Alias" + }, + "team_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Alias" + }, + "team_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Id" + }, + "user_email": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Email" + }, + "user_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Id" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "id", + "downloaded_at", + "version", + "digest", + "authenticated", + "user_id", + "user_email", + "key_alias", + "team_id", + "team_alias" + ], + "title": "PluginDownloadEvent", + "type": "object" + }, + "PluginDownloadsResponse": { + "properties": { + "anonymous_downloads": { + "title": "Anonymous Downloads", + "type": "integer" + }, + "downloads": { + "items": { + "$ref": "#/components/schemas/PluginDownloadEvent" + }, + "title": "Downloads", + "type": "array" + }, + "next_cursor": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Next Cursor" + }, + "total_downloads": { + "title": "Total Downloads", + "type": "integer" + } + }, + "required": [ + "total_downloads", + "anonymous_downloads", + "downloads", + "next_cursor" + ], + "title": "PluginDownloadsResponse", + "type": "object" + }, "PluginListItem": { "description": "Plugin item in list responses.", "properties": { @@ -8060,7 +8199,7 @@ } ], "default": "1.0.0", - "description": "Semantic version", + "description": "Semantic version, defaults to 1.0.0 on creation", "title": "Version" } }, @@ -8098,7 +8237,7 @@ "type": "object" }, "UpdatePluginRequest": { - "description": "Request body for replacing an existing plugin.\n\nThe plugin name is the resource identity and is supplied as the path\nparameter, so it cannot be changed here. This is a full replace: omitted\nfields reset to their defaults, so version is cleared rather than\ndefaulting to the create-time \"1.0.0\".", + "description": "Request body for replacing an existing plugin.\n\nThe plugin name is the resource identity and is supplied as the path\nparameter, so it cannot be changed here. This is a full replace: omitted\nfields reset to their defaults, except version which automatically increments\nthe published patch version when omitted or null.", "properties": { "author": { "anyOf": [ @@ -8147,6 +8286,19 @@ "description": "Skill domain (e.g., 'Productivity')", "title": "Domain" }, + "expected_updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Published updated_at read by the editor. A stale value returns 409", + "title": "Expected Updated At" + }, "homepage": { "anyOf": [ { @@ -8220,7 +8372,7 @@ "type": "null" } ], - "description": "Semantic version; cleared if omitted", + "description": "Explicit higher SemVer, or omit/null for an automatic patch bump", "title": "Version" } }, @@ -8351,7 +8503,7 @@ } ], "default": "1.0.0", - "description": "Semantic version", + "description": "Semantic version, defaults to 1.0.0 on creation", "title": "Version" } }, @@ -8683,7 +8835,7 @@ ] }, "put": { - "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default (e.g. an omitted version is cleared, not kept).\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: Name of the plugin to update (path parameter)\n - source: Plugin source reference (github, url, git-subdir, or archive format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n - installation_preference: 'available', 'auto_install', or 'required' (optional)\n\nReturns:\n Update status (action is always \"updated\") and plugin information.\n\nExample:\n ```bash\n curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"2.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", + "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default, except version which automatically increments\nthe published patch version when omitted or null.\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: Name of the plugin to update (path parameter)\n - source: Plugin source reference (github, url, git-subdir, or archive format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n - installation_preference: 'available', 'auto_install', or 'required' (optional)\n\nReturns:\n Update status (action is always \"updated\") and plugin information.\n\nExample:\n ```bash\n curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"2.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", "operationId": "update_plugin_claude_code_plugins__plugin_name__put", "parameters": [ { @@ -8838,6 +8990,87 @@ ] } }, + "/claude-code/plugins/{plugin_name}/downloads": { + "get": { + "description": "Admin-only ZIP download totals and authenticated downloader history.\n\nAnonymous downloads contribute only to totals. Counts include repeat requests and client prefetches,\nstarting when tracking is deployed. External sources and installation completion are not observed.", + "operationId": "get_plugin_downloads_claude_code_plugins__plugin_name__downloads_get", + "parameters": [ + { + "in": "path", + "name": "plugin_name", + "required": true, + "schema": { + "title": "Plugin Name", + "type": "string" + } + }, + { + "description": "Maximum number of download events", + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 25, + "description": "Maximum number of download events", + "maximum": 100, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "description": "Previous next_cursor", + "in": "query", + "name": "cursor", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[1-9][0-9]{0,18}$", + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Previous next_cursor", + "title": "Cursor" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/PluginDownloadsResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Plugin Downloads", + "tags": [ + "claude_code_marketplace" + ] + } + }, "/claude-code/plugins/{plugin_name}/enable": { "post": { "description": "Enable a disabled plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: The name of the plugin to enable", @@ -59000,6 +59233,44 @@ "skill_hub": { "components": { "schemas": { + "Body_replace_skill_archive_skill_hub_skills__name__upload_put": { + "properties": { + "expected_updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Expected Updated At" + }, + "file": { + "contentMediaType": "application/octet-stream", + "description": "Replacement Agent Skills ZIP with the same declared skill name", + "title": "File", + "type": "string" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "file" + ], + "title": "Body_replace_skill_archive_skill_hub_skills__name__upload_put", + "type": "object" + }, "Body_upload_skill_skill_hub_skills_upload_post": { "properties": { "file": { @@ -59085,7 +59356,7 @@ "type": "null" } ], - "description": "Version label for the imported snapshot", + "description": "Semantic version, or omit/null to start at 1.0.0", "title": "Version" } }, @@ -59095,6 +59366,180 @@ "title": "ImportSkillRequest", "type": "object" }, + "PluginDownloadEvent": { + "properties": { + "authenticated": { + "title": "Authenticated", + "type": "boolean" + }, + "digest": { + "title": "Digest", + "type": "string" + }, + "downloaded_at": { + "format": "date-time", + "title": "Downloaded At", + "type": "string" + }, + "id": { + "title": "Id", + "type": "string" + }, + "key_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Alias" + }, + "team_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Alias" + }, + "team_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Id" + }, + "user_email": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Email" + }, + "user_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Id" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "id", + "downloaded_at", + "version", + "digest", + "authenticated", + "user_id", + "user_email", + "key_alias", + "team_id", + "team_alias" + ], + "title": "PluginDownloadEvent", + "type": "object" + }, + "PluginDownloadsResponse": { + "properties": { + "anonymous_downloads": { + "title": "Anonymous Downloads", + "type": "integer" + }, + "downloads": { + "items": { + "$ref": "#/components/schemas/PluginDownloadEvent" + }, + "title": "Downloads", + "type": "array" + }, + "next_cursor": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Next Cursor" + }, + "total_downloads": { + "title": "Total Downloads", + "type": "integer" + } + }, + "required": [ + "total_downloads", + "anonymous_downloads", + "downloads", + "next_cursor" + ], + "title": "PluginDownloadsResponse", + "type": "object" + }, + "SkillHubAgentSkill": { + "properties": { + "description": { + "title": "Description", + "type": "string" + }, + "digest": { + "title": "Digest", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "origin": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Origin" + } + }, + "required": [ + "name", + "description", + "digest", + "origin" + ], + "title": "SkillHubAgentSkill", + "type": "object" + }, "SkillHubList": { "properties": { "count": { @@ -59154,6 +59599,19 @@ ], "title": "Domain" }, + "expected_updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Published updated_at read by the editor. A stale value returns 409", + "title": "Expected Updated At" + }, "keywords": { "anyOf": [ { @@ -59188,18 +59646,196 @@ "type": "null" } ], + "description": "Explicit higher SemVer, or omit/null for an automatic patch bump", "title": "Version" } }, "title": "SkillHubMetadata", "type": "object" }, + "SkillHubPendingSkill": { + "properties": { + "action": { + "const": "submitted", + "default": "submitted", + "title": "Action", + "type": "string" + }, + "proposal": { + "$ref": "#/components/schemas/SkillHubProposal" + }, + "status": { + "const": "pending_review", + "default": "pending_review", + "title": "Status", + "type": "string" + }, + "submission_id": { + "title": "Submission Id", + "type": "string" + } + }, + "required": [ + "submission_id", + "proposal" + ], + "title": "SkillHubPendingSkill", + "type": "object" + }, + "SkillHubProposal": { + "properties": { + "agent_skill": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubAgentSkill" + }, + { + "type": "null" + } + ] + }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "domain": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Domain" + }, + "keywords": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Keywords" + }, + "name": { + "title": "Name", + "type": "string" + }, + "namespace": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Namespace" + }, + "source": { + "additionalProperties": { + "type": "string" + }, + "title": "Source", + "type": "object" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "name", + "description", + "version", + "source" + ], + "title": "SkillHubProposal", + "type": "object" + }, "SkillHubPublication": { "additionalProperties": false, "properties": {}, "title": "SkillHubPublication", "type": "object" }, + "SkillHubRejection": { + "additionalProperties": false, + "properties": { + "reason": { + "anyOf": [ + { + "maxLength": 2000, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Reason" + } + }, + "title": "SkillHubRejection", + "type": "object" + }, + "SkillHubReviewResponse": { + "properties": { + "action": { + "enum": [ + "approved", + "rejected" + ], + "title": "Action", + "type": "string" + }, + "skill": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "type": "null" + } + ] + }, + "status": { + "const": "success", + "default": "success", + "title": "Status", + "type": "string" + }, + "submission": { + "$ref": "#/components/schemas/SkillHubSubmission" + } + }, + "required": [ + "action", + "submission", + "skill" + ], + "title": "SkillHubReviewResponse", + "type": "object" + }, "SkillHubSkill": { "properties": { "content_url": { @@ -59224,6 +59860,29 @@ ], "title": "Description" }, + "domain": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Domain" + }, + "download_count": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "description": "Gateway ZIP downloads, included in proxy-admin catalog listings", + "title": "Download Count" + }, "download_url": { "anyOf": [ { @@ -59239,10 +59898,35 @@ "title": "Id", "type": "string" }, + "keywords": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Keywords" + }, "name": { "title": "Name", "type": "string" }, + "namespace": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Namespace" + }, "source": { "additionalProperties": { "type": "string" @@ -59250,6 +59934,18 @@ "title": "Source", "type": "object" }, + "updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Updated At" + }, "version": { "anyOf": [ { @@ -59279,6 +59975,533 @@ "title": "SkillHubSkill", "type": "object" }, + "SkillHubSubmission": { + "properties": { + "base_skill_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Base Skill Id" + }, + "base_updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Base Updated At" + }, + "created_at": { + "format": "date-time", + "title": "Created At", + "type": "string" + }, + "has_archive": { + "title": "Has Archive", + "type": "boolean" + }, + "id": { + "title": "Id", + "type": "string" + }, + "operation": { + "enum": [ + "create", + "update" + ], + "title": "Operation", + "type": "string" + }, + "proposal": { + "$ref": "#/components/schemas/SkillHubProposal" + }, + "rejection_reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Rejection Reason" + }, + "reviewed_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Reviewed At" + }, + "reviewed_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Reviewed By" + }, + "skill_name": { + "title": "Skill Name", + "type": "string" + }, + "status": { + "enum": [ + "pending_review", + "approved", + "rejected" + ], + "title": "Status", + "type": "string" + }, + "submitted_by": { + "title": "Submitted By", + "type": "string" + } + }, + "required": [ + "id", + "skill_name", + "operation", + "status", + "proposal", + "has_archive", + "submitted_by", + "created_at", + "base_skill_id", + "base_updated_at", + "reviewed_by", + "reviewed_at", + "rejection_reason" + ], + "title": "SkillHubSubmission", + "type": "object" + }, + "SkillHubSubmissionDetail": { + "properties": { + "current_skill": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "type": "null" + } + ] + }, + "submission": { + "$ref": "#/components/schemas/SkillHubSubmission" + } + }, + "required": [ + "submission", + "current_skill" + ], + "title": "SkillHubSubmissionDetail", + "type": "object" + }, + "SkillHubSubmissionList": { + "properties": { + "count": { + "title": "Count", + "type": "integer" + }, + "has_more": { + "title": "Has More", + "type": "boolean" + }, + "submissions": { + "items": { + "$ref": "#/components/schemas/SkillHubSubmission" + }, + "title": "Submissions", + "type": "array" + } + }, + "required": [ + "submissions", + "count", + "has_more" + ], + "title": "SkillHubSubmissionList", + "type": "object" + }, + "SkillHubVersion": { + "properties": { + "content_kind": { + "enum": [ + "stored_archive", + "source_reference" + ], + "title": "Content Kind", + "type": "string" + }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "download_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Download Url" + }, + "id": { + "title": "Id", + "type": "string" + }, + "manifest": { + "additionalProperties": true, + "title": "Manifest", + "type": "object" + }, + "published_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Published At" + }, + "published_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Published By" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "id", + "version", + "description", + "manifest", + "published_at", + "published_by", + "content_kind", + "download_url" + ], + "title": "SkillHubVersion", + "type": "object" + }, + "SkillHubVersions": { + "properties": { + "next_offset": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Next Offset" + }, + "versions": { + "items": { + "$ref": "#/components/schemas/SkillHubVersion" + }, + "title": "Versions", + "type": "array" + } + }, + "required": [ + "versions", + "next_offset" + ], + "title": "SkillHubVersions", + "type": "object" + }, + "SkillScan": { + "properties": { + "archive_sha256": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Archive Sha256" + }, + "completed_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Completed At" + }, + "error": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Error" + }, + "expires_at": { + "format": "date-time", + "title": "Expires At", + "type": "string" + }, + "id": { + "title": "Id", + "type": "string" + }, + "report": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillScanReport" + }, + { + "type": "null" + } + ] + }, + "requested_by": { + "title": "Requested By", + "type": "string" + }, + "started_at": { + "format": "date-time", + "title": "Started At", + "type": "string" + }, + "status": { + "enum": [ + "running", + "completed", + "failed" + ], + "title": "Status", + "type": "string" + } + }, + "required": [ + "id", + "status", + "requested_by", + "started_at", + "expires_at" + ], + "title": "SkillScan", + "type": "object" + }, + "SkillScanFinding": { + "properties": { + "explanation": { + "anyOf": [ + { + "maxLength": 8000, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Explanation" + }, + "file": { + "maxLength": 2048, + "title": "File", + "type": "string" + }, + "message": { + "maxLength": 8000, + "title": "Message", + "type": "string" + }, + "remediation": { + "anyOf": [ + { + "maxLength": 8000, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Remediation" + }, + "rule_id": { + "maxLength": 200, + "title": "Rule Id", + "type": "string" + }, + "severity": { + "enum": [ + "LOW", + "MEDIUM", + "HIGH", + "CRITICAL" + ], + "title": "Severity", + "type": "string" + }, + "start_line": { + "anyOf": [ + { + "minimum": 1.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Start Line" + } + }, + "required": [ + "rule_id", + "severity", + "message", + "file" + ], + "title": "SkillScanFinding", + "type": "object" + }, + "SkillScanReport": { + "properties": { + "complete": { + "title": "Complete", + "type": "boolean" + }, + "findings": { + "items": { + "$ref": "#/components/schemas/SkillScanFinding" + }, + "maxItems": 500, + "title": "Findings", + "type": "array" + }, + "limitations": { + "default": [], + "items": { + "maxLength": 2000, + "type": "string" + }, + "maxItems": 100, + "title": "Limitations", + "type": "array" + }, + "risk_score": { + "maximum": 100.0, + "minimum": 0.0, + "title": "Risk Score", + "type": "integer" + }, + "scan_mode": { + "enum": [ + "static-only", + "static+llm" + ], + "title": "Scan Mode", + "type": "string" + }, + "scanner_version": { + "maxLength": 100, + "minLength": 1, + "title": "Scanner Version", + "type": "string" + } + }, + "required": [ + "scanner_version", + "complete", + "scan_mode", + "risk_score", + "findings" + ], + "title": "SkillScanReport", + "type": "object" + }, + "SkillScanRequest": { + "additionalProperties": false, + "properties": {}, + "title": "SkillScanRequest", + "type": "object" + }, + "SkillScanResponse": { + "properties": { + "configured": { + "title": "Configured", + "type": "boolean" + }, + "scan": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillScan" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "configured", + "scan" + ], + "title": "SkillScanResponse", + "type": "object" + }, "ValidationError": { "properties": { "ctx": { @@ -59413,7 +60636,15 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/SkillHubSkill" + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "$ref": "#/components/schemas/SkillHubPendingSkill" + } + ], + "title": "Response Import Skill Skill Hub Skills Import Post" } } }, @@ -59460,7 +60691,15 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/SkillHubSkill" + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "$ref": "#/components/schemas/SkillHubPendingSkill" + } + ], + "title": "Response Upload Skill Skill Hub Skills Upload Post" } } }, @@ -59575,7 +60814,7 @@ ] }, "put": { - "description": "Replace registry metadata, clearing omitted fields and preserving skill files.", + "description": "Replace registry metadata, preserving files. Members submit an immutable proposal for admin review.", "operationId": "replace_skill_metadata_skill_hub_skills__name__put", "parameters": [ { @@ -59603,7 +60842,15 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/SkillHubSkill" + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "$ref": "#/components/schemas/SkillHubPendingSkill" + } + ], + "title": "Response Replace Skill Metadata Skill Hub Skills Name Put" } } }, @@ -59859,6 +61106,743 @@ "skill_hub" ] } + }, + "/skill-hub/skills/{name}/upload": { + "put": { + "description": "Replace skill files, retaining registry fields. Members propose an update against the original editor timestamp.", + "operationId": "replace_skill_archive_skill_hub_skills__name__upload_put", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "multipart/form-data": { + "schema": { + "$ref": "#/components/schemas/Body_replace_skill_archive_skill_hub_skills__name__upload_put" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "$ref": "#/components/schemas/SkillHubPendingSkill" + } + ], + "title": "Response Replace Skill Archive Skill Hub Skills Name Upload Put" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Replace Skill Archive", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/versions": { + "get": { + "description": "Read published metadata history. External references are not retained artifacts.", + "operationId": "list_skill_versions_skill_hub_skills__name__versions_get", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + }, + { + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 20, + "maximum": 100, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "in": "query", + "name": "offset", + "required": false, + "schema": { + "default": 0, + "minimum": 0, + "title": "Offset", + "type": "integer" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubVersions" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Skill Versions", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/versions/{version_id}": { + "get": { + "description": "Read a retained publication under the current skill's visibility policy.", + "operationId": "get_skill_version_skill_hub_skills__name__versions__version_id__get", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + }, + { + "in": "path", + "name": "version_id", + "required": true, + "schema": { + "title": "Version Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubVersion" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Skill Version", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/versions/{version_id}/archive": { + "get": { + "description": "Download a retained publication as a portable ZIP with SKILL.md at its root.", + "operationId": "download_skill_version_skill_hub_skills__name__versions__version_id__archive_get", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + }, + { + "in": "path", + "name": "version_id", + "required": true, + "schema": { + "title": "Version Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/zip": { + "schema": { + "format": "binary", + "type": "string" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Download Skill Version", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{plugin_name}/downloads": { + "get": { + "description": "Admin-only ZIP download totals and authenticated downloader history.\n\nAnonymous downloads contribute only to totals. Counts include repeat requests and client prefetches,\nstarting when tracking is deployed. External sources and installation completion are not observed.", + "operationId": "get_plugin_downloads_skill_hub_skills__plugin_name__downloads_get", + "parameters": [ + { + "in": "path", + "name": "plugin_name", + "required": true, + "schema": { + "title": "Plugin Name", + "type": "string" + } + }, + { + "description": "Maximum number of download events", + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 25, + "description": "Maximum number of download events", + "maximum": 100, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "description": "Previous next_cursor", + "in": "query", + "name": "cursor", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[1-9][0-9]{0,18}$", + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Previous next_cursor", + "title": "Cursor" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/PluginDownloadsResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Plugin Downloads", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions": { + "get": { + "description": "Read a page of owned proposals. Proxy admins and admin viewers can read every owner's proposals.", + "operationId": "list_submissions_skill_hub_submissions_get", + "parameters": [ + { + "in": "query", + "name": "status", + "required": false, + "schema": { + "anyOf": [ + { + "enum": [ + "pending_review", + "approved", + "rejected" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Status" + } + }, + { + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 20, + "maximum": 100, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "in": "query", + "name": "offset", + "required": false, + "schema": { + "default": 0, + "minimum": 0, + "title": "Offset", + "type": "integer" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubSubmissionList" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Submissions", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}": { + "get": { + "description": "Read an immutable proposal and its currently visible published skill.", + "operationId": "get_submission_skill_hub_submissions__submission_id__get", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubSubmissionDetail" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Submission", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}/approve": { + "put": { + "description": "Publish a pending proposal and record its approval atomically. Only proxy admins may review.", + "operationId": "approve_submission_skill_hub_submissions__submission_id__approve_put", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubReviewResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Approve Submission", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}/archive": { + "get": { + "description": "Download an owned proposal as a private portable ZIP. Legacy nonportable archives return 404.", + "operationId": "download_submission_skill_hub_submissions__submission_id__archive_get", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/zip": { + "schema": { + "format": "binary", + "type": "string" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Download Submission", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}/reject": { + "put": { + "description": "Reject a pending proposal without changing the published skill. Only proxy admins may review.", + "operationId": "reject_submission_skill_hub_submissions__submission_id__reject_put", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubRejection" + }, + { + "type": "null" + } + ], + "title": "Body" + } + } + } + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubReviewResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Reject Submission", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}/scan": { + "get": { + "description": "Read the latest manual security scan. Only proxy admins may read reports.", + "operationId": "get_skill_scan_skill_hub_submissions__submission_id__scan_get", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillScanResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Skill Scan", + "tags": [ + "skill_hub" + ] + }, + "post": { + "description": "Start an optional scan of a pending submission. Scans never change approval decisions.", + "operationId": "start_skill_scan_skill_hub_submissions__submission_id__scan_post", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillScanRequest" + }, + { + "type": "null" + } + ], + "title": " Request" + } + } + } + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillScanResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Start Skill Scan", + "tags": [ + "skill_hub" + ] + } } } }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 25849634764..d3ac264518e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -88,6 +88,7 @@ from litellm.types.proxy.carried_budget_state import ( ) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry 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 @@ -449,15 +450,29 @@ class LiteLLMRoutes(enum.Enum): "/v1/traces/{trace_id}/spans/{span_id}/error", ) - litellm_native_routes = [ + 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", "/rag/query", @@ -476,6 +491,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}", ] @@ -706,6 +722,7 @@ class LiteLLMRoutes(enum.Enum): + key_management_routes + mcp_management_routes + list(agent_management_routes) + + list(skill_hub_routes) ) spend_tracking_routes = [ @@ -948,6 +965,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 ## @@ -2489,8 +2507,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): """ 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.", + True, + description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.", ) request_tag_mappings: Sequence[RequestTagMapping] | None = Field( @@ -3027,6 +3045,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']).", diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 0faefaeaa41..26cb0a2f3e6 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -27,14 +27,22 @@ from datetime import datetime, timezone from types import MappingProxyType from typing import Annotated, Final, Protocol, TypedDict, get_args -from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile from fastapi.responses import JSONResponse, Response from pydantic import Json, TypeAdapter -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, assert_never from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.url_utils import is_https_url_with_host -from litellm.models.plugin_upload import PluginUpload, has_agent_skill_metadata +from litellm.models.plugin_upload import ( + PluginDownloader, + PluginDownloadResult, + PluginDownloadsResponse, + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + has_agent_skill_metadata, +) from litellm.proxy._types import CommonProxyErrors, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( SkillVisibility, @@ -208,8 +216,8 @@ async def upload_plugin( } spec: Final = RegisterPluginRequest.model_validate({**metadata.model_dump(), "source": source}) _validate_installation_preference(spec) - plugin_id: Final = await uploads.create( - PluginUpload( + published: Final = await uploads.publish( + PluginPublication( name=spec.name, version=spec.version, description=spec.description, @@ -218,15 +226,15 @@ async def upload_plugin( created_by=user_api_key_dict.user_id, ) ) - if plugin_id is None: - raise _name_conflict_error(spec.name) + if isinstance(published, PluginPublicationError): + raise _publication_error(published) return RegisterPluginResponse( status="success", action="created", plugin=PluginResponse( - id=plugin_id, + id=published.id, name=spec.name, - version=spec.version, + version=published.version, description=spec.description, source=source, installation_preference=spec.installation_preference, @@ -235,6 +243,27 @@ async def upload_plugin( ) +def require_download_recorded(result: PluginDownloadResult) -> None: + match result: + case PluginDownloadResult.RECORDED: + return + case PluginDownloadResult.SKILL_NOT_FOUND: + raise _error_response(404, "Skill archive not found") + case _: + assert_never(result) + + +def plugin_downloader(caller: UserAPIKeyAuth | None) -> PluginDownloader: + return PluginDownloader( + authenticated=caller is not None, + user_id=caller.user_id if caller is not None else None, + user_email=caller.user_email if caller is not None else None, + key_alias=caller.key_alias if caller is not None else None, + team_id=caller.team_id if caller is not None else None, + team_alias=caller.team_alias if caller is not None else None, + ) + + @router.get( "/claude-code/plugins/{plugin_name}/archive/{digest}.zip", tags=["Claude Code Marketplace"], @@ -268,6 +297,7 @@ async def download_plugin_archive( content: Final = await uploads.get_content(archive) if content is None: raise _error_response(404, "Skill archive not found") + require_download_recorded(await uploads.record_download(archive, digest, plugin_downloader(caller))) return Response( content=content, media_type="application/zip", @@ -279,6 +309,35 @@ async def download_plugin_archive( ) +@router.get( + "/claude-code/plugins/{plugin_name}/downloads", + tags=["Claude Code Marketplace"], + response_model=PluginDownloadsResponse, +) +async def get_plugin_downloads( + plugin_name: str, + uploads: PluginUploadRepositoryDep, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + limit: Annotated[int, Query(ge=1, le=100, description="Maximum number of download events")] = 25, + cursor: Annotated[str | None, Query(pattern=r"^[1-9][0-9]{0,18}$", description="Previous next_cursor")] = None, +) -> PluginDownloadsResponse: + """Admin-only ZIP download totals and authenticated downloader history. + + Anonymous downloads contribute only to totals. Counts include repeat requests and client prefetches, + starting when tracking is deployed. External sources and installation completion are not observed. + """ + _require_proxy_admin(user_api_key_dict) + if uploads is None: + raise _error_response(500, CommonProxyErrors.db_not_connected_error.value) + before: Final = int(cursor) if cursor is not None else None + if before is not None and before > 9223372036854775807: + raise _error_response(422, "Invalid download cursor") + result: Final = await uploads.downloads(plugin_name, limit, before) + if result is None: + raise _error_response(404, "Plugin not found") + return result + + @router.get( "/claude-code/marketplace.json", tags=["Claude Code Marketplace"], @@ -472,10 +531,14 @@ def _error_response(status_code: int, message: str) -> HTTPException: return HTTPException(status_code=status_code, detail={"error": message}) -def _name_conflict_error(name: str) -> HTTPException: - return _error_response( - 409, f"A skill named '{name}' already exists. Update the existing skill instead of adding it again." - ) +def _publication_error(error: PluginPublicationError) -> HTTPException: + match error.kind: + case "invalid_version": + return _error_response(400, error.message) + case "not_found": + return _error_response(404, error.message) + case "conflict": + return _error_response(409, error.message) def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: @@ -498,6 +561,7 @@ def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: ) async def register_plugin( request: RegisterPluginRequest, + uploads: PluginUploadRepositoryDep, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -540,12 +604,11 @@ async def register_plugin( }' ``` """ - from prisma.errors import UniqueViolationError - try: _require_proxy_admin(user_api_key_dict) - prisma_client: Final = await get_prisma_client() + if uploads is None: + raise _error_response(500, CommonProxyErrors.db_not_connected_error.value) if not re.match(r"^[a-z0-9-]+$", request.name): raise HTTPException( @@ -556,30 +619,19 @@ async def register_plugin( _validate_plugin_source(request.source) _validate_installation_preference(request) - existing: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( - where={"name": request.name} - ) - if existing: - raise _name_conflict_error(request.name) - manifest: Final[Mapping[str, object]] = _build_plugin_manifest(request.name, request) - try: - plugin: Final[_PluginRecord] = await ClaudeCodePluginRepository(prisma_client).table.create( - data={ - "name": request.name, - "version": request.version, - "description": request.description, - "manifest_json": json.dumps(manifest), - "files_json": "{}", - "enabled": True, - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - "created_by": user_api_key_dict.user_id, - } + plugin: Final = await uploads.publish( + PluginPublication( + name=request.name, + version=request.version, + description=request.description, + manifest=manifest, + created_by=user_api_key_dict.user_id, ) - except UniqueViolationError: - raise _name_conflict_error(request.name) + ) + if isinstance(plugin, PluginPublicationError): + raise _publication_error(plugin) verbose_proxy_logger.info("Plugin %s created successfully", request.name) @@ -590,7 +642,7 @@ async def register_plugin( id=plugin.id, name=plugin.name, version=plugin.version, - description=plugin.description, + description=request.description, source=request.source, installation_preference=request.installation_preference, enabled=plugin.enabled, @@ -755,6 +807,7 @@ async def get_plugin( async def update_plugin( plugin_name: str, request: UpdatePluginRequest, + uploads: PluginUploadRepositoryDep, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ @@ -763,7 +816,8 @@ async def update_plugin( The plugin is identified by its name in the path, which is the resource identity and cannot be changed here. This is a full replace, not a merge: the manifest is rebuilt from the request body, so any optional field left - out is reset to its default (e.g. an omitted version is cleared, not kept). + out is reset to its default, except version which automatically increments + the published patch version when omitted or null. Send the full desired state. Returns 404 if no plugin with the given name exists; use @@ -801,6 +855,8 @@ async def update_plugin( try: _require_proxy_admin(user_api_key_dict) + if uploads is None: + raise _error_response(500, CommonProxyErrors.db_not_connected_error.value) prisma_client: Final = await get_prisma_client() @@ -808,6 +864,8 @@ async def update_plugin( existing: Final[_PluginRecord | None] = await repository.table.find_unique(where={"name": plugin_name}) if not existing: raise _error_response(404, f"Plugin '{plugin_name}' not found") + if request.expected_updated_at is not None and request.expected_updated_at != existing.updated_at: + raise _error_response(409, "The published skill changed. Reload it before publishing an update") current_manifest: Final = MANIFEST_FIELDS.validate_json(existing.manifest_json or "{}") if request.source != current_manifest.get("source"): @@ -823,18 +881,18 @@ async def update_plugin( ), } - plugin: Final[_PluginRecord | None] = await repository.table.update( - where={"name": plugin_name}, - data={ - "version": request.version, - "description": request.description, - "manifest_json": json.dumps(manifest), - "files_json": "{}", - "updated_at": datetime.now(timezone.utc), - }, + plugin: Final = await uploads.publish( + PluginPublication( + name=plugin_name, + version=request.version, + description=request.description, + manifest=manifest, + created_by=user_api_key_dict.user_id, + ), + PluginPublicationBase(existing.id, existing.updated_at), ) - if plugin is None: - raise _error_response(404, f"Plugin '{plugin_name}' not found") + if isinstance(plugin, PluginPublicationError): + raise _publication_error(plugin) verbose_proxy_logger.info("Plugin %s updated successfully", plugin_name) @@ -845,7 +903,7 @@ async def update_plugin( id=plugin.id, name=plugin.name, version=plugin.version, - description=plugin.description, + description=request.description, source=request.source, installation_preference=request.installation_preference, enabled=plugin.enabled, diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 59fd00cec43..9c433af487e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -909,13 +909,14 @@ class ResetBudgetJob: db_span("reset_budget_cascade", "LiteLLM_BudgetTable"), budget_cascade_unit_of_work(self._new_batch) as uow, ): + uow.row_locks.queue_row_locks(cascade.budget_ids) + _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) uow.team_memberships.queue_total_spend_repair(cascade.budget_ids) _queue_budget_linked_resets(uow.team_memberships, cascade) _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) - _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) + _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) - _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 80f09df6c86..ff57596232f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -328,13 +328,13 @@ _SpendTableName = Literal[ "agent_list_transactions", ] _SPEND_TABLE_COMMIT_ORDER: Final[tuple[_SpendTableName, ...]] = ( - "user_list_transactions", - "end_user_list_transactions", - "key_list_transactions", - "team_list_transactions", - "team_member_list_transactions", "org_list_transactions", "org_member_list_transactions", + "team_list_transactions", + "user_list_transactions", + "end_user_list_transactions", + "team_member_list_transactions", + "key_list_transactions", "project_list_transactions", "tag_list_transactions", "model_access_group_list_transactions", @@ -563,6 +563,24 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, """ +def _org_members_in_key_order(spend_by_member_key: Mapping[str, float]) -> tuple[tuple[tuple[str, str], float], ...]: + """Sorted by the decoded (organization_id, user_id), the order every other writer locks these rows in, since + the percent-encoded key does not sort the same way.""" + decoded: Final = ( + ((unquote(key.split("::")[1]), unquote(key.split("::")[3])), cost) for key, cost in spend_by_member_key.items() + ) + return tuple(sorted(decoded)) + + +def _charged_team_ids(updates: DBSpendUpdateTransactions) -> tuple[str, ...]: + members: Final = updates.get("team_member_list_transactions") or {} + return tuple( + sorted( + {*(updates.get("team_list_transactions") or {}), *(key.split("::")[1] for key in members)}, + ) + ) + + async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None: # key is "team_id::::user_id::"; locks are taken in sorted team_id order like the team endpoints rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items()) @@ -712,7 +730,7 @@ class DBSpendUpdateWriter: proxy_settings: Final = cast( # cast-ok: proxy settings use string keys Mapping[str, object], general_settings ) - use_ledger: Final = proxy_settings.get("use_spend_logs_ledger") is True + use_ledger: Final = proxy_settings.get("use_spend_logs_ledger", True) is not False request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {}) if kwargs is not None: kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker @@ -863,8 +881,12 @@ class DBSpendUpdateWriter: payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs ) if "accounting_payload" in payload: + from litellm.proxy.utils import accounting_only_spend_log + await self._insert_spend_log_to_db( - payload={**payload, "accounting_log_disabled": disable_spend_logs}, + payload=accounting_only_spend_log(payload) + if disable_spend_logs + else {**payload, "accounting_log_disabled": False}, prisma_client=prisma_client, llm_router=llm_router, ) @@ -2294,6 +2316,10 @@ class DBSpendUpdateWriter: transaction: _SpendTransaction, updates: DBSpendUpdateTransactions, ) -> None: + """Every team advisory lock is taken before the first row lock, the order the team roster writers use, so + applying a batch of charges cannot deadlock with a concurrent roster change.""" + for team_id in _charged_team_ids(updates): + _ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id) await self._commit_spend_updates_to_db( prisma_client, 0, proxy_logging_obj, updates, outer_transaction=transaction ) @@ -2313,6 +2339,104 @@ class DBSpendUpdateWriter: """ from litellm.proxy.utils import ProxyUpdateSpend + ### UPDATE ORG TABLE ### + org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"] + verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions) + if org_list_transactions is not None and len(org_list_transactions.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with _spend_update_tx( + prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction + ) as transaction: + async with transaction.batch_() as batcher: + # Sort by org_id for consistent lock ordering across pods to prevent deadlocks. + for org_id, response_cost in sorted(org_list_transactions.items()): + batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists + where={"organization_id": org_id}, + data={"spend": {"increment": response_cost}}, + ) + break + except Exception as e: + if outer_transaction is not None: + raise + await self.handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + if on_table_committed is not None: + on_table_committed("org_list_transactions") + + org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions") + verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions) + if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with ( + _spend_update_tx( + prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction + ) as transaction, + transaction.batch_() as batcher, + ): + for (org_id, member_id), response_cost in _org_members_in_key_order( + org_member_list_transactions + ): + batcher.litellm_organizationmembership.update_many( + where={"organization_id": org_id, "user_id": member_id}, + data={"spend": {"increment": response_cost}}, + ) + break + except Exception as e: + if outer_transaction is not None: + raise + await self.handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + if on_table_committed is not None: + on_table_committed("org_member_list_transactions") + + ### UPDATE TEAM TABLE ### + team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"] + verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions) + if team_list_transactions is not None and len(team_list_transactions.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with _spend_update_tx( + prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction + ) as transaction: + async with transaction.batch_() as batcher: + # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. + for team_id, response_cost in sorted(team_list_transactions.items()): + verbose_proxy_logger.debug( + "Updating spend for team id=%s by %s", team_id, response_cost + ) + batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists + where={"team_id": team_id}, + data={"spend": {"increment": response_cost}}, + ) + break + except Exception as e: + if outer_transaction is not None: + raise + await self.handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + if on_table_committed is not None: + on_table_committed("team_list_transactions") + ### UPDATE USER TABLE ### user_list_transactions: Final = db_spend_update_transactions["user_list_transactions"] verbose_proxy_logger.debug("User Spend transactions: %s", user_list_transactions) @@ -2359,75 +2483,6 @@ class DBSpendUpdateWriter: ) if on_table_committed is not None: on_table_committed("end_user_list_transactions") - ### UPDATE KEY TABLE ### - key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"] - verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions) - if key_list_transactions is not None and len(key_list_transactions.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with _spend_update_tx( - prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction - ) as transaction: - async with transaction.batch_() as batcher: - # Sort by token for consistent lock ordering across pods to prevent deadlocks. - for token, response_cost in sorted(key_list_transactions.items()): - spend_increment: _SpendIncrement = {"increment": response_cost} - batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists - where={"token": token}, - data={ - "spend": spend_increment, - "total_spend": spend_increment, - "last_active": datetime.now(timezone.utc), - }, - ) - break - except Exception as e: - if outer_transaction is not None: - raise - await self.handle_spend_update_failure( - e=e, - attempt=i, - n_retry_times=n_retry_times, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - if on_table_committed is not None: - on_table_committed("key_list_transactions") - - ### UPDATE TEAM TABLE ### - team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"] - verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions) - if team_list_transactions is not None and len(team_list_transactions.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with _spend_update_tx( - prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction - ) as transaction: - async with transaction.batch_() as batcher: - # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. - for team_id, response_cost in sorted(team_list_transactions.items()): - verbose_proxy_logger.debug( - "Updating spend for team id=%s by %s", team_id, response_cost - ) - batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists - where={"team_id": team_id}, - data={"spend": {"increment": response_cost}}, - ) - break - except Exception as e: - if outer_transaction is not None: - raise - await self.handle_spend_update_failure( - e=e, - attempt=i, - n_retry_times=n_retry_times, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - if on_table_committed is not None: - on_table_committed("team_list_transactions") ### UPDATE TEAM Membership TABLE with spend ### team_member_list_transactions: Final = db_spend_update_transactions["team_member_list_transactions"] @@ -2478,22 +2533,27 @@ class DBSpendUpdateWriter: elif on_table_committed is not None: on_table_committed("team_member_list_transactions") - ### UPDATE ORG TABLE ### - org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"] - verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions) - if org_list_transactions is not None and len(org_list_transactions.keys()) > 0: + ### UPDATE KEY TABLE ### + key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"] + verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions) + if key_list_transactions is not None and len(key_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: async with _spend_update_tx( - prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction + prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction ) as transaction: async with transaction.batch_() as batcher: - # Sort by org_id for consistent lock ordering across pods to prevent deadlocks. - for org_id, response_cost in sorted(org_list_transactions.items()): - batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists - where={"organization_id": org_id}, - data={"spend": {"increment": response_cost}}, + # Sort by token for consistent lock ordering across pods to prevent deadlocks. + for token, response_cost in sorted(key_list_transactions.items()): + spend_increment: _SpendIncrement = {"increment": response_cost} + batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists + where={"token": token}, + data={ + "spend": spend_increment, + "total_spend": spend_increment, + "last_active": datetime.now(timezone.utc), + }, ) break except Exception as e: @@ -2507,39 +2567,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) if on_table_committed is not None: - on_table_committed("org_list_transactions") - - org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions") - verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions) - if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with ( - _spend_update_tx( - prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction - ) as transaction, - transaction.batch_() as batcher, - ): - for key, response_cost in sorted(org_member_list_transactions.items()): - _, quoted_org_id, _, quoted_user_id = key.split("::") - batcher.litellm_organizationmembership.update_many( - where={"organization_id": unquote(quoted_org_id), "user_id": unquote(quoted_user_id)}, - data={"spend": {"increment": response_cost}}, - ) - break - except Exception as e: - if outer_transaction is not None: - raise - await self.handle_spend_update_failure( - e=e, - attempt=i, - n_retry_times=n_retry_times, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - if on_table_committed is not None: - on_table_committed("org_member_list_transactions") + on_table_committed("key_list_transactions") ### UPDATE PROJECT TABLE ### project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions") diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py index 8f929119dee..323cf57dfec 100644 --- a/litellm/proxy/deps/repositories.py +++ b/litellm/proxy/deps/repositories.py @@ -5,11 +5,13 @@ from fastapi import Depends from litellm.models.deployment_warning import DeploymentWarningIgnoreStore from litellm.models.gateway_activity import GatewayActivityStore from litellm.models.plugin_upload import PluginUploadStore +from litellm.models.skill_submission import SkillSubmissionStore from litellm.models.user_banner import UserBannerStore from litellm.proxy.deps.database import DatabaseDep from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository from litellm.repositories.gateway_activity_repository import GatewayActivityRepository from litellm.repositories.plugin_upload_repository import PluginUploadRepository +from litellm.repositories.skill_submission_repository import SkillSubmissionRepository from litellm.repositories.user_banner_repository import UserBannerRepository @@ -43,3 +45,12 @@ async def get_plugin_upload_repository(database: DatabaseDep) -> PluginUploadSto PluginUploadRepositoryDep: TypeAlias = Annotated[PluginUploadStore | None, Depends(get_plugin_upload_repository)] + + +async def get_skill_submission_repository(database: DatabaseDep) -> SkillSubmissionStore | None: + return None if database is None else SkillSubmissionRepository(database.prisma) + + +SkillSubmissionRepositoryDep: TypeAlias = Annotated[ + SkillSubmissionStore | None, Depends(get_skill_submission_repository) +] diff --git a/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py b/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py index faef3b9e4d4..93d1ace435d 100644 --- a/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py +++ b/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py @@ -1,5 +1,8 @@ import asyncio import hashlib +from collections.abc import Mapping +from dataclasses import dataclass, replace +from datetime import datetime from typing import Annotated, Final from urllib.parse import quote @@ -10,18 +13,31 @@ import litellm from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory has untyped optional params ) -from litellm.models.plugin_upload import PluginArchiveMetadata, PluginUpload, PluginUploadStore +from litellm.models.plugin_upload import ( + PluginArchiveMetadata, + PluginDownloader, + PluginDownloadsResponse, + PluginPublication, + PluginSnapshot, + PluginUploadStore, + PluginVersion, + PublishedPlugin, + has_agent_skill_metadata, +) +from litellm.models.skill_submission import SkillSubmission, SkillSubmissionStore from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( MANIFEST_FIELDS, delete_plugin, disable_plugin, enable_plugin, + get_plugin_downloads, header_key_caller, list_plugins, plugin_archive_source_matches, plugin_archive_url, - update_plugin, + plugin_downloader, + require_download_recorded, ) from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import resolve_skill_visibility from litellm.proxy.anthropic_endpoints.claude_code_endpoints.skill_upload_archive import build_uploaded_skill_archive @@ -31,11 +47,17 @@ from litellm.proxy.common_utils.claude_code_marketplace_auth import ( marketplace_key_header_present, ) from litellm.proxy.common_utils.resource_ownership import is_proxy_admin -from litellm.proxy.deps.repositories import PluginUploadRepositoryDep +from litellm.proxy.deps.repositories import PluginUploadRepositoryDep, SkillSubmissionRepositoryDep from litellm.proxy.discovery_endpoints.agent_skills_archive import ( SkillArchive, build_agent_skill_archive, ) +from litellm.proxy.discovery_endpoints.skill_hub_mutations import ( + SkillMutationError, + archive_writer, + write_skill, + write_target, +) from litellm.proxy.discovery_endpoints.skill_import import ( MAX_IMPORT_BYTES, SkillImportError, @@ -43,18 +65,33 @@ from litellm.proxy.discovery_endpoints.skill_import import ( parse_skill_source, select_imported_skill, ) -from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 -from litellm.types.proxy.claude_code_endpoints import PluginListItem, UpdatePluginRequest +from litellm.types.proxy.claude_code_endpoints import PluginListItem from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import AgentSkillsIndexEntry from litellm.types.proxy.skill_hub import ( ImportSkillRequest, SkillHubList, SkillHubMetadata, + SkillHubMutationResponse, + SkillHubPendingSkill, + SkillHubProposal, SkillHubPublication, SkillHubSkill, + SkillHubVersion, + SkillHubVersions, ) router: Final = APIRouter(prefix="/skill-hub/skills", tags=["Skills Hub"]) +_EDITABLE_METADATA: Final = frozenset(SkillHubMetadata.model_fields) +router.add_api_route( + "/{plugin_name}/downloads", get_plugin_downloads, methods=["GET"], response_model=PluginDownloadsResponse +) + + +@dataclass(frozen=True, slots=True) +class VisibleSkillArchive: + archive: SkillArchive + metadata: PluginArchiveMetadata + downloader: PluginDownloader def _url(request: Request, route: str) -> str: @@ -67,7 +104,7 @@ def _download_url(request: Request, name: str) -> str: return _url(request, f"skill-hub/skills/{quote(name, safe='')}/archive") -def _skill_response(request: Request, plugin: PluginListItem) -> SkillHubSkill: +def _skill_response(request: Request, plugin: PluginListItem, download_count: int | None = None) -> SkillHubSkill: local_archive: Final = plugin_archive_source_matches( request, plugin.name, @@ -80,16 +117,22 @@ def _skill_response(request: Request, plugin: PluginListItem) -> SkillHubSkill: name=plugin.name, description=plugin.description, version=plugin.version, + domain=plugin.domain, + namespace=plugin.namespace, + keywords=None if plugin.keywords is None else tuple(plugin.keywords), + updated_at=None if plugin.updated_at is None else datetime.fromisoformat(plugin.updated_at), visibility="public" if plugin.enabled else "private", source={"source": "archive", "url": _download_url(request, plugin.name)} if local_archive else plugin.source, download_url=_download_url(request, plugin.name) if local_archive else None, content_url=_url(request, f"skill-hub/skills/{quote(plugin.name, safe='')}/content") if local_archive else None, + download_count=download_count, ) @router.get("", response_model=SkillHubList) async def list_skills( request: Request, + uploads: PluginUploadRepositoryDep, caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], query: str = "", offset: Annotated[int, Query(ge=0)] = 0, @@ -103,8 +146,14 @@ async def list_skills( for plugin in sorted(catalog.plugins, key=lambda item: item.name) if needle in " ".join((plugin.name, plugin.description or "", *(plugin.keywords or ()))).casefold() ) + page: Final = matched[offset : offset + limit] + counts: Final[Mapping[str, int]] = ( + await _require_store(uploads).download_counts(tuple(plugin.id for plugin in page)) + if is_proxy_admin(caller) + else {} + ) return SkillHubList( - skills=tuple(_skill_response(request, plugin) for plugin in matched[offset : offset + limit]), + skills=tuple(_skill_response(request, plugin, counts.get(plugin.id)) for plugin in page), count=len(matched), next_offset=offset + limit if offset + limit < len(matched) else None, ) @@ -122,22 +171,28 @@ async def get_skill( metadata: Final = await store.get_metadata(name) if metadata is None or not (await resolve_skill_visibility(caller)).allows(metadata): raise HTTPException(404, "Skill not found") - return _metadata_response(request, metadata) + return skill_metadata_response(request, metadata) -def _metadata_response(request: Request, metadata: PluginArchiveMetadata) -> SkillHubSkill: +def skill_metadata_response(request: Request, metadata: PluginArchiveMetadata) -> SkillHubSkill: + manifest: Final = MANIFEST_FIELDS.validate_json(metadata.manifest_json or "{}") return _skill_response( request, - PluginListItem( - id=metadata.id or metadata.name, - name=metadata.name, - description=metadata.description, - version=metadata.version, - source=dict(metadata.source), - has_agent_skill_metadata=metadata.has_agent_skill_metadata, - enabled=metadata.enabled, - created_at=None, - updated_at=None, + PluginListItem.model_validate( + { + "id": metadata.id or metadata.name, + "name": metadata.name, + "description": metadata.description, + "version": metadata.version, + "source": dict(metadata.source), + "has_agent_skill_metadata": metadata.has_agent_skill_metadata, + "enabled": metadata.enabled, + "created_at": None, + "updated_at": metadata.updated_at.isoformat() if metadata.updated_at is not None else None, + "domain": manifest.get("domain"), + "namespace": manifest.get("namespace"), + "keywords": manifest.get("keywords"), + } ), ) @@ -153,57 +208,33 @@ def _require_admin(caller: UserAPIKeyAuth) -> None: raise HTTPException(403, "Only proxy admins may publish skills") -async def skill_publisher( - caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], -) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import proxy_logging_obj - - _require_admin(caller) - limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter") - if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): - raise HTTPException(503, "Skills Hub publishing requires the default v3 rate limiter") - key_identity: Final = caller.api_key or caller.token - identity: Final = ( - f"user:{caller.user_id}" if caller.user_id else f"key:{key_identity}" if key_identity else "local-admin" - ) - result: Final = await limiter.atomic_check_and_increment_by_n( - descriptors=[ - { - "key": "skill_hub_publish", - "value": hashlib.sha256(identity.encode()).hexdigest(), - "rate_limit": {"requests_per_unit": 10, "window_size": 60}, - } - ], - increments=[{"requests": 1}], - parent_otel_span=caller.parent_otel_span, - ) - if result["overall_code"] == "OVER_LIMIT": - raise HTTPException( - 429, "Skills Hub allows 10 uploads or imports per minute per user", headers={"Retry-After": "60"} - ) - return caller - - -@router.put("/{name}", response_model=SkillHubSkill) +@router.put("/{name}", response_model=SkillHubMutationResponse) async def replace_skill_metadata( request: Request, name: str, body: SkillHubMetadata, uploads: PluginUploadRepositoryDep, + submissions: SkillSubmissionRepositoryDep, caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], -) -> SkillHubSkill: - """Replace registry metadata, clearing omitted fields and preserving skill files.""" - _require_admin(caller) - store: Final = _require_store(uploads) - metadata: Final = await store.get_metadata(name) - if metadata is None: +) -> SkillHubMutationResponse: + """Replace registry metadata, preserving files. Members submit an immutable proposal for admin review.""" + existing: Final = await write_target(name, True, uploads, caller, body.expected_updated_at) + if isinstance(existing, SkillMutationError): + raise HTTPException(existing.status, existing.message) + if existing is None: raise HTTPException(404, "Skill not found") - manifest: Final = MANIFEST_FIELDS.validate_json(metadata.manifest_json or "{}") - updated: Final = UpdatePluginRequest.model_validate( - {**manifest, **body.model_dump(), "source": dict(metadata.source)} + publication: Final = PluginPublication( + name=name, + version=body.version, + description=body.description, + manifest={ + **{key: value for key, value in existing.manifest.items() if key not in _EDITABLE_METADATA}, + **body.model_dump(exclude_none=True), + "name": name, + }, + created_by=caller.user_id, ) - await update_plugin(plugin_name=name, request=updated, user_api_key_dict=caller) - return await get_skill(request, name, uploads, caller) + return await _mutate(request, publication, existing, _require_store(uploads), submissions, caller) @router.post("/{name}/publish", response_model=SkillHubSkill) @@ -242,14 +273,15 @@ async def remove_skill(name: str, caller: Annotated[UserAPIKeyAuth, Depends(user return Response(status_code=204) -@router.post("/upload", response_model=SkillHubSkill) +@router.post("/upload", response_model=SkillHubMutationResponse) async def upload_skill( request: Request, uploads: PluginUploadRepositoryDep, - caller: Annotated[UserAPIKeyAuth, Depends(skill_publisher)], + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(archive_writer)], file: Annotated[UploadFile, File(description="Agent Skills ZIP with SKILL.md and optional supporting files")], version: Annotated[str | None, Form()] = None, -) -> SkillHubSkill: +) -> SkillHubMutationResponse: """Publish a portable skill while retaining Claude marketplace compatibility.""" store: Final = _require_store(uploads) content: Final = await file.read(MAX_IMPORT_BYTES + 1) @@ -258,21 +290,22 @@ async def upload_skill( archive: Final = await asyncio.to_thread(build_agent_skill_archive, content) if archive is None: raise HTTPException(422, "Upload a valid Agent Skill with name and description in SKILL.md") - return await _publish(request, store, caller, archive, version, None) + return await _publish(request, store, submissions, caller, archive, version, None) def import_http_client() -> httpx.AsyncClient: return get_async_httpx_client(llm_provider="skill_hub_import").client -@router.post("/import", response_model=SkillHubSkill) +@router.post("/import", response_model=SkillHubMutationResponse) async def import_skill( request: Request, body: ImportSkillRequest, uploads: PluginUploadRepositoryDep, - caller: Annotated[UserAPIKeyAuth, Depends(skill_publisher)], + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(archive_writer)], client: Annotated[httpx.AsyncClient, Depends(import_http_client)], -) -> SkillHubSkill: +) -> SkillHubMutationResponse: """Import a public skills.sh or GitHub skill as a portable snapshot.""" store: Final = _require_store(uploads) source: Final = parse_skill_source(body.url, skill=body.skill, ref=body.ref, path=body.path) @@ -284,17 +317,122 @@ async def import_skill( archive: Final = await asyncio.to_thread(select_imported_skill, content, source) if isinstance(archive, SkillImportError): raise HTTPException(archive.status, archive.message) - return await _publish(request, store, caller, archive, body.version, body.url) + return await _publish(request, store, submissions, caller, archive, body.version, body.url) async def _publish( request: Request, store: PluginUploadStore, + submissions: SkillSubmissionStore | None, caller: UserAPIKeyAuth, archive: SkillArchive, version: str | None, origin: str | None, +) -> SkillHubMutationResponse: + existing: Final = await write_target(archive.declared_name or "", False, store, caller) + if isinstance(existing, SkillMutationError): + raise HTTPException(existing.status, existing.message) + publication: Final = await prepare_skill_publication(request, archive, version, origin, caller.user_id) + return await _mutate(request, publication, existing, store, submissions, caller) + + +async def _mutate( + request: Request, + publication: PluginPublication, + existing: PluginSnapshot | None, + store: PluginUploadStore, + submissions: SkillSubmissionStore | None, + caller: UserAPIKeyAuth, +) -> SkillHubMutationResponse: + result: Final = await write_skill(publication, existing, store, submissions, caller) + if isinstance(result, SkillMutationError): + raise HTTPException(result.status, result.message) + if isinstance(result, SkillSubmission): + return SkillHubPendingSkill(submission_id=result.id, proposal=skill_proposal_response(request, result)) + return skill_publication_response(request, publication, result) + + +def skill_proposal_response(request: Request, submission: SkillSubmission) -> SkillHubProposal: + proposal: Final = submission.proposal + return SkillHubProposal.model_validate( + { + **proposal.manifest, + "name": proposal.name, + "description": proposal.description, + "version": proposal.version, + **( + { + "source": { + "source": "archive", + "url": _url(request, f"skill-hub/submissions/{quote(submission.id, safe='')}/archive"), + } + } + if submission.has_archive + else {} + ), + } + ) + + +@router.put("/{name}/upload", response_model=SkillHubMutationResponse) +async def replace_skill_archive( + request: Request, + name: str, + uploads: PluginUploadRepositoryDep, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(archive_writer)], + file: Annotated[UploadFile, File(description="Replacement Agent Skills ZIP with the same declared skill name")], + version: Annotated[str | None, Form()] = None, + expected_updated_at: Annotated[datetime | None, Form()] = None, +) -> SkillHubMutationResponse: + """Replace skill files, retaining registry fields. Members propose an update against the original editor timestamp.""" + existing: Final = await write_target(name, True, uploads, caller, expected_updated_at) + if isinstance(existing, SkillMutationError): + raise HTTPException(existing.status, existing.message) + if existing is None: + raise HTTPException(404, "Skill not found") + content: Final = await file.read(MAX_IMPORT_BYTES + 1) + if len(content) > MAX_IMPORT_BYTES: + raise HTTPException(413, "Skill ZIP must be 10 MiB or smaller") + archive: Final = await asyncio.to_thread(build_agent_skill_archive, content) + if archive is None or archive.declared_name != name: + raise HTTPException(422, "Replacement SKILL.md must declare the existing skill name and a description") + prepared: Final = await prepare_skill_publication(request, archive, version, None, caller.user_id) + publication: Final = replace(prepared, manifest={**existing.manifest, **prepared.manifest}) + return await _mutate(request, publication, existing, _require_store(uploads), submissions, caller) + + +def skill_publication_response( + request: Request, publication: PluginPublication, published: PublishedPlugin ) -> SkillHubSkill: + return _skill_response( + request, + PluginListItem.model_validate( + { + "source": publication.manifest.get("source", {}), + "id": published.id, + "name": published.name, + "description": publication.description, + "version": published.version, + "enabled": published.enabled, + "has_agent_skill_metadata": has_agent_skill_metadata(publication.manifest), + "created_at": None, + "updated_at": published.updated_at.isoformat(), + "domain": publication.manifest.get("domain"), + "namespace": publication.manifest.get("namespace"), + "keywords": publication.manifest.get("keywords"), + } + ), + ) + + +async def prepare_skill_publication( + request: Request, + archive: SkillArchive, + version: str | None, + origin: str | None, + created_by: str | None, +) -> PluginPublication: name: Final = archive.declared_name if name is None: raise HTTPException(422, "SKILL.md must declare a skill name") @@ -322,20 +460,119 @@ async def _publish( "origin": origin, }, } - identifier: Final = await store.create( - PluginUpload(name, version, archive.declared_description, manifest, plugin, caller.user_id) + return PluginPublication( + name=name, + version=version, + description=archive.declared_description, + manifest=manifest, + content=plugin, + created_by=created_by, ) - if identifier is None: - raise HTTPException(409, "A skill with this name already exists") - return _metadata_response( - request, - PluginArchiveMetadata( - name, True, source, None, identifier, version, archive.declared_description, has_agent_skill_metadata=True + + +async def _version_metadata(name: str, store: PluginUploadStore, caller: UserAPIKeyAuth) -> PluginArchiveMetadata: + metadata: Final = await store.get_metadata(name) + if metadata is None or metadata.id is None or not (await resolve_skill_visibility(caller)).allows(metadata): + raise HTTPException(404, "Skill not found") + return metadata + + +def _version_response(request: Request, name: str, version: PluginVersion) -> SkillHubVersion: + return SkillHubVersion( + id=version.id, + version=version.version, + description=version.description, + manifest=version.manifest, + published_at=version.published_at, + published_by=version.published_by, + content_kind="stored_archive" if version.archive_available else "source_reference", + download_url=( + _url(request, f"skill-hub/skills/{quote(name, safe='')}/versions/{quote(version.id, safe='')}/archive") + if version.archive_available + else None ), ) -async def _visible_archive(request: Request, name: str, uploads: PluginUploadStore | None) -> SkillArchive: +@router.get("/{name}/versions", response_model=SkillHubVersions) +async def list_skill_versions( + request: Request, + name: str, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + limit: Annotated[int, Query(ge=1, le=100)] = 20, + offset: Annotated[int, Query(ge=0)] = 0, +) -> SkillHubVersions: + """Read published metadata history. External references are not retained artifacts.""" + store: Final = _require_store(uploads) + metadata: Final = await _version_metadata(name, store, caller) + versions: Final = await store.versions(metadata, limit + 1, offset) + return SkillHubVersions( + versions=tuple(_version_response(request, name, version) for version in versions[:limit]), + next_offset=offset + limit if len(versions) > limit else None, + ) + + +@router.get("/{name}/versions/{version_id}", response_model=SkillHubVersion) +async def get_skill_version( + request: Request, + name: str, + version_id: str, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillHubVersion: + """Read a retained publication under the current skill's visibility policy.""" + store: Final = _require_store(uploads) + metadata: Final = await _version_metadata(name, store, caller) + version: Final = await store.version(metadata, version_id) + if version is None: + raise HTTPException(404, "Skill version not found") + return _version_response(request, name, version) + + +@router.get( + "/{name}/versions/{version_id}/archive", + response_class=Response, + responses={200: {"content": {"application/zip": {"schema": {"type": "string", "format": "binary"}}}}}, +) +async def download_skill_version( + name: str, + version_id: str, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> Response: + """Download a retained publication as a portable ZIP with SKILL.md at its root.""" + store: Final = _require_store(uploads) + metadata: Final = await _version_metadata(name, store, caller) + version: Final = await store.version(metadata, version_id) + if version is None: + raise HTTPException(404, "Skill version not found") + content: Final = await store.version_content(metadata, version_id) + archive: Final = ( + await asyncio.to_thread(build_agent_skill_archive, content, plugin_name=name) if content is not None else None + ) + if archive is None: + raise HTTPException(404, "This version has no retained Agent Skills archive") + require_download_recorded( + await store.record_download( + replace(metadata, version=version.version), + archive.digest.removeprefix("sha256:"), + plugin_downloader(caller), + ) + ) + return Response( + archive.content, + media_type="application/zip", + headers={ + "Content-Disposition": f'attachment; filename="{name}.zip"', + "Cache-Control": "private, no-store", + "X-Content-Type-Options": "nosniff", + "ETag": f'"{archive.digest}"', + }, + ) + + +async def _visible_archive(request: Request, name: str, uploads: PluginUploadStore | None) -> VisibleSkillArchive: anonymous_allowed: Final = litellm.public_skills_index is True and not marketplace_auth_required() caller: Final = ( await header_key_caller(request) if marketplace_key_header_present(request) or not anonymous_allowed else None @@ -358,7 +595,7 @@ async def _visible_archive(request: Request, name: str, uploads: PluginUploadSto ) if archive is None: raise HTTPException(404, "This skill has no stored Agent Skills archive") - return archive + return VisibleSkillArchive(archive, metadata, plugin_downloader(caller)) @router.get( @@ -370,9 +607,15 @@ async def download_skill( request: Request, name: str, uploads: PluginUploadRepositoryDep, digest: str | None = None ) -> Response: """Download the complete portable skill, with SKILL.md at the ZIP root.""" - archive: Final = await _visible_archive(request, name, uploads) + visible: Final = await _visible_archive(request, name, uploads) + archive: Final = visible.archive if digest is not None and digest != archive.digest: raise HTTPException(404, "Skill archive version not found") + require_download_recorded( + await _require_store(uploads).record_download( + visible.metadata, archive.digest.removeprefix("sha256:"), visible.downloader + ) + ) return Response( archive.content, media_type="application/zip", @@ -388,9 +631,9 @@ async def download_skill( @router.get("/{name}/content", response_class=Response, responses={200: {"content": {"text/markdown": {}}}}) async def skill_content(request: Request, name: str, uploads: PluginUploadRepositoryDep) -> Response: """Read the original SKILL.md without downloading supporting files.""" - archive: Final = await _visible_archive(request, name, uploads) + visible: Final = await _visible_archive(request, name, uploads) return Response( - archive.manifest_content, + visible.archive.manifest_content, media_type="text/markdown", headers={"Cache-Control": "private, no-store", "X-Content-Type-Options": "nosniff"}, ) diff --git a/litellm/proxy/discovery_endpoints/skill_hub_mutations.py b/litellm/proxy/discovery_endpoints/skill_hub_mutations.py new file mode 100644 index 00000000000..63c519923c6 --- /dev/null +++ b/litellm/proxy/discovery_endpoints/skill_hub_mutations.py @@ -0,0 +1,172 @@ +import hashlib +from dataclasses import dataclass +from datetime import datetime +from typing import Annotated, Final + +from fastapi import Depends, HTTPException +from typing_extensions import assert_never + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PluginSnapshot, + PluginUploadStore, + PublishedPlugin, +) +from litellm.models.skill_submission import ( + SkillSubmission, + SkillSubmissionStore, + SubmissionConflict, + SubmissionLimitExceeded, + SubmissionMissing, + SubmissionReviewed, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import resolve_skill_visibility +from litellm.proxy.auth.auth_exception_handler import DB_UNAVAILABLE_FALLBACK_USER_ID +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin +from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 +from litellm.repositories.plugin_publication import resolve_plugin_version +from litellm.types.proxy.auth.user_roles import LitellmUserRoles + + +@dataclass(frozen=True, slots=True) +class SkillMutationError: + status: int + message: str + + +def writer_error(caller: UserAPIKeyAuth) -> SkillMutationError | None: + if is_proxy_admin(caller): + return None + if ( + not caller.user_id + or caller.user_id == DB_UNAVAILABLE_FALLBACK_USER_ID + or caller.user_role in (None, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) + ): + return SkillMutationError(403, "Submitting a skill requires a writable user identity.") + return None + + +def submission_owner(caller: UserAPIKeyAuth) -> str | None | SkillMutationError: + if is_proxy_admin(caller) or caller.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY: + return None + if not caller.user_id or caller.user_id == DB_UNAVAILABLE_FALLBACK_USER_ID: + return SkillMutationError(403, "Reading submissions requires a user identity.") + return caller.user_id + + +async def archive_writer(caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import proxy_logging_obj + + error: Final = writer_error(caller) + if error is not None: + raise HTTPException(error.status, error.message) + limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter") + if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): + raise HTTPException(503, "Skills Hub publishing requires the default v3 rate limiter") + key_identity: Final = caller.api_key or caller.token + identity: Final = ( + f"user:{caller.user_id}" if caller.user_id else f"key:{key_identity}" if key_identity else "local-admin" + ) + result: Final = await limiter.atomic_check_and_increment_by_n( + descriptors=[ + { + "key": "skill_hub_publish", + "value": hashlib.sha256(identity.encode()).hexdigest(), + "rate_limit": {"requests_per_unit": 10, "window_size": 60}, + } + ], + increments=[{"requests": 1}], + parent_otel_span=caller.parent_otel_span, + ) + if result["overall_code"] == "OVER_LIMIT": + raise HTTPException( + 429, "Skills Hub allows 10 uploads or imports per minute per user", headers={"Retry-After": "60"} + ) + return caller + + +async def write_target( + name: str, + updating: bool, + uploads: PluginUploadStore | None, + caller: UserAPIKeyAuth, + expected_updated_at: datetime | None = None, +) -> PluginSnapshot | SkillMutationError | None: + error: Final = writer_error(caller) + if error is not None: + return error + if uploads is None: + return SkillMutationError(503, "Skills Hub requires a connected database") + existing: Final = await uploads.get_plugin(name) + if updating and (existing is None or not (await resolve_skill_visibility(caller)).allows(existing)): + return SkillMutationError(404, "Skill not found") + if existing is not None and expected_updated_at is not None and existing.updated_at != expected_updated_at: + return SkillMutationError(409, "The published skill changed. Reload it before submitting.") + if not updating and existing is not None: + return SkillMutationError(409, f"A skill named '{name}' already exists. Update the existing skill instead.") + return existing + + +def publication_error(error: PluginPublicationError) -> SkillMutationError: + return SkillMutationError({"invalid_version": 400, "conflict": 409, "not_found": 404}[error.kind], error.message) + + +async def write_skill( + publication: PluginPublication, + existing: PluginSnapshot | None, + uploads: PluginUploadStore, + submissions: SkillSubmissionStore | None, + caller: UserAPIKeyAuth, +) -> PublishedPlugin | SkillSubmission | SkillMutationError: + error: Final = writer_error(caller) + if error is not None: + return error + validation: Final = resolve_plugin_version( + publication.version, None if existing is None else existing.version, updating=existing is not None + ) + if isinstance(validation, PluginPublicationError): + return publication_error(validation) + expected: Final = None if existing is None else PluginPublicationBase(existing.id, existing.updated_at) + if is_proxy_admin(caller): + published: Final = await uploads.publish(publication, expected) + return publication_error(published) if isinstance(published, PluginPublicationError) else published + if submissions is None or not caller.user_id: + return SkillMutationError(503, "Skills Hub requires a connected database") + submission: Final = await submissions.submit(publication, caller.user_id, expected) + return ( + SkillMutationError(429, submission.message) if isinstance(submission, SubmissionLimitExceeded) else submission + ) + + +async def review_skill( + submission_id: str, + submissions: SkillSubmissionStore | None, + caller: UserAPIKeyAuth, + approve: bool, + reason: str | None, + publication_base_url: str | None = None, +) -> SubmissionReviewed | SkillMutationError: + if not is_proxy_admin(caller): + return SkillMutationError(403, "Only proxy admins may review skill submissions.") + if submissions is None: + return SkillMutationError(503, "Skills Hub requires a connected database") + result: Final = await submissions.review( + submission_id, caller.user_id or "proxy-admin", approve, reason, publication_base_url=publication_base_url + ) + match result: + case SubmissionMissing(): + return SkillMutationError(404, "Skill submission not found") + case SubmissionConflict(message=message): + return SkillMutationError(409, message) + case PluginPublicationError(kind="invalid_version") as error: + return publication_error(error) + case PluginPublicationError(message=message): + return SkillMutationError(409, message) + case SubmissionReviewed(): + return result + case _: + assert_never(result) diff --git a/litellm/proxy/discovery_endpoints/skill_hub_submissions.py b/litellm/proxy/discovery_endpoints/skill_hub_submissions.py new file mode 100644 index 00000000000..c1bb422b2a0 --- /dev/null +++ b/litellm/proxy/discovery_endpoints/skill_hub_submissions.py @@ -0,0 +1,192 @@ +import asyncio +from typing import Annotated, Final + +from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response + +from litellm.models.plugin_upload import PluginPublication +from litellm.models.skill_submission import SkillSubmission, SkillSubmissionStore, SubmissionStatus +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import resolve_skill_visibility +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import PluginUploadRepositoryDep, SkillSubmissionRepositoryDep +from litellm.proxy.discovery_endpoints.agent_skills_archive import build_agent_skill_archive +from litellm.proxy.discovery_endpoints.skill_hub_endpoints import ( + skill_metadata_response, + skill_proposal_response, + skill_publication_response, +) +from litellm.proxy.discovery_endpoints.skill_hub_mutations import SkillMutationError, review_skill, submission_owner +from litellm.types.proxy.skill_hub import ( + SkillHubRejection, + SkillHubReviewResponse, + SkillHubSubmission, + SkillHubSubmissionDetail, + SkillHubSubmissionList, +) + +router: Final = APIRouter(prefix="/skill-hub/submissions", tags=["Skills Hub"]) + + +def _store(submissions: SkillSubmissionStore | None) -> SkillSubmissionStore: + if submissions is None: + raise HTTPException(503, "Skills Hub requires a connected database") + return submissions + + +def _owner(caller: UserAPIKeyAuth) -> str | None: + owner: Final = submission_owner(caller) + if isinstance(owner, SkillMutationError): + raise HTTPException(owner.status, owner.message) + return owner + + +def _response(request: Request, submission: SkillSubmission) -> SkillHubSubmission: + return SkillHubSubmission( + id=submission.id, + skill_name=submission.plugin_name, + operation=submission.operation, + status=submission.status, + proposal=skill_proposal_response(request, submission), + has_archive=submission.has_archive, + submitted_by=submission.submitted_by, + created_at=submission.created_at, + base_skill_id=submission.base_plugin_id, + base_updated_at=submission.base_updated_at, + reviewed_by=submission.reviewed_by, + reviewed_at=submission.reviewed_at, + rejection_reason=submission.rejection_reason, + ) + + +@router.get("", response_model=SkillHubSubmissionList) +async def list_submissions( + request: Request, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + status: SubmissionStatus | None = None, + limit: Annotated[int, Query(ge=1, le=100)] = 20, + offset: Annotated[int, Query(ge=0)] = 0, +) -> SkillHubSubmissionList: + """Read a page of owned proposals. Proxy admins and admin viewers can read every owner's proposals.""" + owner: Final = _owner(caller) + page: Final = await _store(submissions).list(owner, status, limit, offset) + return SkillHubSubmissionList( + submissions=tuple(_response(request, row) for row in page.submissions), + count=len(page.submissions), + has_more=page.has_more, + ) + + +@router.get("/{submission_id}", response_model=SkillHubSubmissionDetail) +async def get_submission( + request: Request, + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillHubSubmissionDetail: + """Read an immutable proposal and its currently visible published skill.""" + owner: Final = _owner(caller) + submission: Final = await _store(submissions).get(submission_id, owner) + if submission is None: + raise HTTPException(404, "Skill submission not found") + if uploads is None: + raise HTTPException(503, "Skills Hub requires a connected database") + current: Final = await uploads.get_metadata(submission.plugin_name) + visible: Final = current is not None and (owner is None or (await resolve_skill_visibility(caller)).allows(current)) + return SkillHubSubmissionDetail( + submission=_response(request, submission), + current_skill=skill_metadata_response(request, current) if current is not None and visible else None, + ) + + +@router.get( + "/{submission_id}/archive", + response_class=Response, + responses={200: {"content": {"application/zip": {"schema": {"type": "string", "format": "binary"}}}}}, +) +async def download_submission( + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> Response: + """Download an owned proposal as a private portable ZIP. Legacy nonportable archives return 404.""" + owner: Final = _owner(caller) + store: Final = _store(submissions) + submission: Final = await store.get(submission_id, owner) + if submission is None: + raise HTTPException(404, "Skill submission not found") + content: Final = await store.get_archive(submission_id, owner) + archive: Final = ( + await asyncio.to_thread(build_agent_skill_archive, content, plugin_name=submission.plugin_name) + if content is not None + else None + ) + if archive is None: + raise HTTPException(404, "This submission has no portable Agent Skills archive") + return Response( + archive.content, + media_type="application/zip", + headers={ + "Content-Disposition": f'attachment; filename="{submission.plugin_name}.zip"', + "Cache-Control": "private, no-store", + "X-Content-Type-Options": "nosniff", + "ETag": f'"{archive.digest}"', + }, + ) + + +async def _review( + request: Request, + submission_id: str, + submissions: SkillSubmissionStore | None, + caller: UserAPIKeyAuth, + approve: bool, + reason: str | None, +) -> SkillHubReviewResponse: + from litellm.proxy.utils import get_custom_url + + result: Final = await review_skill( + submission_id, submissions, caller, approve, reason, get_custom_url(str(request.base_url)) + ) + if isinstance(result, SkillMutationError): + raise HTTPException(result.status, result.message) + proposal: Final = result.submission.proposal + publication: Final = PluginPublication( + name=proposal.name, + version=proposal.version, + description=proposal.description, + manifest={ + **proposal.manifest, + **({"source": dict(result.published_source)} if result.published_source else {}), + }, + created_by=result.submission.submitted_by, + ) + return SkillHubReviewResponse( + action="approved" if approve else "rejected", + submission=_response(request, result.submission), + skill=None if result.plugin is None else skill_publication_response(request, publication, result.plugin), + ) + + +@router.put("/{submission_id}/approve", response_model=SkillHubReviewResponse) +async def approve_submission( + request: Request, + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillHubReviewResponse: + """Publish a pending proposal and record its approval atomically. Only proxy admins may review.""" + return await _review(request, submission_id, submissions, caller, True, None) + + +@router.put("/{submission_id}/reject", response_model=SkillHubReviewResponse) +async def reject_submission( + request: Request, + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + body: SkillHubRejection | None = None, +) -> SkillHubReviewResponse: + """Reject a pending proposal without changing the published skill. Only proxy admins may review.""" + return await _review(request, submission_id, submissions, caller, False, None if body is None else body.reason) diff --git a/litellm/proxy/discovery_endpoints/skill_scanning.py b/litellm/proxy/discovery_endpoints/skill_scanning.py new file mode 100644 index 00000000000..ee236482fbf --- /dev/null +++ b/litellm/proxy/discovery_endpoints/skill_scanning.py @@ -0,0 +1,186 @@ +import asyncio +import io +import json +from datetime import datetime, timedelta, timezone +from typing import Annotated, Final +from uuid import uuid4 + +import httpx +from fastapi import APIRouter, BackgroundTasks, Body, Depends, HTTPException +from pydantic import ValidationError + +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory has untyped optional params +) +from litellm.models.skill_submission import SkillScanInput, SkillSubmissionStore, SubmissionConflict, SubmissionMissing +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin +from litellm.proxy.deps.repositories import SkillSubmissionRepositoryDep +from litellm.secret_managers.main import get_secret_str +from litellm.types.proxy.skill_scanning import ( + SkillScan, + SkillScannerConfig, + SkillScanReport, + SkillScanRequest, + SkillScanResponse, +) + +router: Final = APIRouter() +_MAX_REPORT_BYTES: Final = 1024 * 1024 + + +def scanner_config() -> SkillScannerConfig | str | None: + from litellm.proxy.proxy_server import general_settings_view + + value: Final = general_settings_view().get("skill_scanner") + if value is None: + return None + try: + return SkillScannerConfig.model_validate(value) + except ValidationError: + return "The skill scanner configuration is invalid" + + +async def scan_report( + client: httpx.AsyncClient, config: SkillScannerConfig, api_key: str, target: SkillScanInput +) -> SkillScanReport | str: + try: + async with client.stream( + "POST", + config.endpoint, + headers={"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"}, + data={ + "source": json.dumps(dict(target.source)), + "use_llm": str(config.use_llm).lower(), + "expires_at": target.scan.expires_at.isoformat(), + }, + files=None if target.content is None else {"file": ("skill.zip", target.content, "application/zip")}, + timeout=httpx.Timeout(config.timeout_seconds, connect=10), + follow_redirects=False, + ) as response: + if response.status_code != 200: + return f"Scanner returned HTTP {response.status_code}. Check the scanner service and authentication." + output: Final = io.BytesIO() + async for chunk in response.aiter_bytes(chunk_size=65536): + if output.tell() + len(chunk) > _MAX_REPORT_BYTES: + return "Scanner report exceeded the 1 MiB limit." + output.write(chunk) + report: Final = SkillScanReport.model_validate_json(output.getvalue()) + if config.use_llm and report.scan_mode != "static+llm": + return report.model_copy( + update={ + "complete": False, + "limitations": (*report.limitations[:99], "Requested semantic analysis did not run."), + } + ) + return report + except (httpx.HTTPError, ValidationError): + return "The scanner was unreachable or returned an invalid report." + + +async def complete_scan( + submissions: SkillSubmissionStore, + submission_id: str, + target: SkillScanInput, + config: SkillScannerConfig, + api_key: str, +) -> None: + try: + client: Final = get_async_httpx_client( + "skill_scanner", params={"ssl_verify": True, "follow_redirects": False} + ).client + report: Final = await asyncio.wait_for( + scan_report(client, config, api_key, target), + timeout=(target.scan.expires_at - datetime.now(timezone.utc)).total_seconds(), + ) + completed: Final = target.scan.model_copy( + update={ + "status": "completed" if isinstance(report, SkillScanReport) else "failed", + "report": report if isinstance(report, SkillScanReport) else None, + "error": report if isinstance(report, str) else None, + "completed_at": datetime.now(timezone.utc), + } + ) + await submissions.finish_scan(submission_id, completed) + except (TimeoutError, asyncio.CancelledError) as error: + await submissions.finish_scan( + submission_id, + target.scan.model_copy( + update={ + "status": "failed", + "error": "Scan timed out or was interrupted. Try scanning again.", + "completed_at": datetime.now(timezone.utc), + } + ), + ) + if isinstance(error, asyncio.CancelledError): + raise + + +@router.get( + "/skill-hub/submissions/{submission_id}/scan", + tags=["Skills Hub"], + response_model=SkillScanResponse, +) +async def get_skill_scan( + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillScanResponse: + """Read the latest manual security scan. Only proxy admins may read reports.""" + if not is_proxy_admin(caller): + raise HTTPException(403, "Proxy admin access is required to review security scans") + if submissions is None: + raise HTTPException(503, "Database is not connected") + scan: Final = await submissions.get_scan(submission_id) + if isinstance(scan, SubmissionMissing): + raise HTTPException(404, "Skill submission not found") + return SkillScanResponse( + configured=isinstance(scanner_config(), SkillScannerConfig), + scan=None if scan is None else scan.at(datetime.now(timezone.utc)), + ) + + +@router.post( + "/skill-hub/submissions/{submission_id}/scan", + tags=["Skills Hub"], + response_model=SkillScanResponse, +) +async def start_skill_scan( + submission_id: str, + background_tasks: BackgroundTasks, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + _request: Annotated[SkillScanRequest | None, Body()] = None, +) -> SkillScanResponse: + """Start an optional scan of a pending submission. Scans never change approval decisions.""" + if not is_proxy_admin(caller): + raise HTTPException(403, "Proxy admin access is required to scan skills") + if submissions is None: + raise HTTPException(503, "Database is not connected") + config: Final = scanner_config() + if config is None: + raise HTTPException(503, "Security scanning is not configured") + if isinstance(config, str): + raise HTTPException(503, config) + api_key: Final = get_secret_str(config.api_key_env) + if not api_key: + raise HTTPException(503, "The scanner authentication secret is not configured") + now: Final = datetime.now(timezone.utc) + target: Final = await submissions.start_scan( + submission_id, + SkillScan( + id=str(uuid4()), + status="running", + requested_by=caller.user_id or "proxy-admin", + started_at=now, + expires_at=now + timedelta(seconds=config.timeout_seconds), + ), + ) + if isinstance(target, SubmissionMissing): + raise HTTPException(404, "Skill submission not found") + if isinstance(target, SubmissionConflict): + raise HTTPException(409, target.message) + background_tasks.add_task(complete_scan, submissions, submission_id, target, config, api_key) + return SkillScanResponse(configured=True, scan=target.scan) diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index e8f0dc68d83..326f6400ba3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -4,7 +4,8 @@ # # +-------------------------------------------------------------+ import os -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from itertools import chain, dropwhile, takewhile from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException @@ -14,6 +15,10 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -218,15 +223,14 @@ class ZscalerAIGuard(CustomGuardrail): zscaler_ai_guard_result = None direction: Final = "OUT" if input_type == "response" else "IN" verbose_proxy_logger.debug("direction: %s", direction) - # Concatenate all texts and send to Zscaler AI Guard - if texts: - concatenated_text: Final = " ".join(texts) + content: Final = self._content_to_scan(inputs, input_type) + if content: zscaler_ai_guard_result = await self.make_zscaler_ai_guard_api_call( zscaler_ai_guard_url=self.zscaler_ai_guard_url, api_key=self.api_key, policy_id=policy_id, direction=direction, - content=concatenated_text, + content=content, user=user_for_payload, **kwargs, ) @@ -244,6 +248,37 @@ class ZscalerAIGuard(CustomGuardrail): verbose_proxy_logger.debug("ZscalerAIGuard: Successfully applied guardrail.") return inputs + def _content_to_scan( + self, + inputs: GenericGuardrailAPIInputs, + input_type: Literal["request", "response"], + ) -> str | None: + structured_messages: Final = inputs.get("structured_messages") or [] + scans_last_user_turn: Final = ( + input_type == "request" + and not effective_scan_only_tool_results_for_guardrail(self) + and any("content" in message for message in structured_messages) + ) + texts: Final = ( + self._last_user_turn_texts(structured_messages) if scans_last_user_turn else inputs.get("texts", []) + ) + return " ".join(texts) if texts else None + + @staticmethod + def _last_user_turn_texts(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + from_last_user_message: Final = dropwhile(lambda message: message.get("role") != "user", reversed(messages)) + last_user_turn: Final = tuple(takewhile(lambda message: message.get("role") == "user", from_last_user_message)) + texts: Final = tuple(chain.from_iterable(map(ZscalerAIGuard._content_texts, reversed(last_user_turn)))) + return texts if any(text.strip() for text in texts) else () + + @staticmethod + def _content_texts(message: Mapping[str, object]) -> tuple[str, ...]: + content: Final = message.get("content") + if isinstance(content, str): + return (content,) + parts: Final = (as_str_mapping(part) for part in as_object_list(content) or ()) + return tuple(text for part in parts if part is not None and isinstance(text := part.get("text"), str)) + def extract_blocking_info(self, response): """ Extracts transaction ID and blocking detector details from a response. diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 6558c82f80c..4cbc2155af0 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -347,47 +347,45 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None: # --------------------------------------------------------------------------- -async def _sync_add_access_group_to_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None: - """Add access_group_id to each team's access_group_ids (idempotent).""" - for team_id in team_ids: +async def _sync_access_group_on_teams( + tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str +) -> None: + """Rows are locked in sorted team_id order, the order every spend writer locks team rows in, so this + cannot deadlock with them.""" + to_add: Final = frozenset(add) + for team_id in sorted(to_add | frozenset(remove)): team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) - if team is not None and access_group_id not in (team.access_group_ids or []): + if team is None: + continue + current = list(team.access_group_ids or []) + if team_id in to_add and access_group_id not in current: + await tx.litellm_teamtable.update( + where={"team_id": team_id}, data={"access_group_ids": [*current, access_group_id]} + ) + elif team_id not in to_add and access_group_id in current: await tx.litellm_teamtable.update( where={"team_id": team_id}, - data={"access_group_ids": list(team.access_group_ids or []) + [access_group_id]}, + data={"access_group_ids": list(_without_access_group(current, access_group_id))}, ) -async def _sync_remove_access_group_from_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None: - """Remove access_group_id from each team's access_group_ids (idempotent).""" - for team_id in team_ids: - team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) - if team is not None and access_group_id in (team.access_group_ids or []): - await tx.litellm_teamtable.update( - where={"team_id": team_id}, - data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]}, - ) - - -async def _sync_add_access_group_to_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None: - """Add access_group_id to each key's access_group_ids (idempotent).""" - for token in key_tokens: +async def _sync_access_group_on_keys( + tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str +) -> None: + """Rows are locked in sorted token order, the order every spend writer locks key rows in.""" + to_add: Final = frozenset(add) + for token in sorted(to_add | frozenset(remove)): key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) - if key is not None and access_group_id not in (key.access_group_ids or []): + if key is None: + continue + current = list(key.access_group_ids or []) + if token in to_add and access_group_id not in current: await tx.litellm_verificationtoken.update( - where={"token": token}, - data={"access_group_ids": list(key.access_group_ids or []) + [access_group_id]}, + where={"token": token}, data={"access_group_ids": [*current, access_group_id]} ) - - -async def _sync_remove_access_group_from_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None: - """Remove access_group_id from each key's access_group_ids (idempotent).""" - for token in key_tokens: - key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) - if key is not None and access_group_id in (key.access_group_ids or []): + elif token not in to_add and access_group_id in current: await tx.litellm_verificationtoken.update( - where={"token": token}, - data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]}, + where={"token": token}, data={"access_group_ids": list(_without_access_group(current, access_group_id))} ) @@ -582,8 +580,8 @@ async def create_access_group( ) # Sync team and key tables to reference the new access group - await _sync_add_access_group_to_teams(tx, data.assigned_team_ids or [], record.access_group_id) - await _sync_add_access_group_to_keys(tx, data.assigned_key_ids or [], record.access_group_id) + await _sync_access_group_on_teams(tx, data.assigned_team_ids or [], (), record.access_group_id) + await _sync_access_group_on_keys(tx, data.assigned_key_ids or [], (), record.access_group_id) except HTTPException: raise except Exception as e: @@ -718,10 +716,8 @@ async def update_access_group( data=update_data, ) - await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id) - await _sync_remove_access_group_from_teams(tx, teams_to_remove, access_group_id) - await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id) - await _sync_remove_access_group_from_keys(tx, keys_to_remove, access_group_id) + await _sync_access_group_on_teams(tx, teams_to_add, teams_to_remove, access_group_id) + await _sync_access_group_on_keys(tx, keys_to_add, keys_to_remove, access_group_id) except HTTPException: raise except Exception as e: @@ -790,25 +786,8 @@ async def delete_access_group( ) affected_key_tokens = list(all_affected_key_tokens) - # Update teams returned by find_many directly — we already have their data. - for team in teams_with_group: - await tx.litellm_teamtable.update( - where={"team_id": team.team_id}, - data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]}, - ) - # Use _sync_remove only for out-of-sync teams not found by the hasSome query. - out_of_sync_team_ids: Final = set(existing.assigned_team_ids or []) - {t.team_id for t in teams_with_group} - await _sync_remove_access_group_from_teams(tx, list(out_of_sync_team_ids), access_group_id) - - # Update keys returned by find_many directly — we already have their data. - for key in keys_with_group: - await tx.litellm_verificationtoken.update( - where={"token": key.token}, - data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]}, - ) - # Use _sync_remove only for out-of-sync keys not found by the hasSome query. - out_of_sync_key_tokens: Final = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group} - await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id) + await _sync_access_group_on_teams(tx, (), affected_team_ids, access_group_id) + await _sync_access_group_on_keys(tx, (), affected_key_tokens, access_group_id) detached_agent_ids: Final = await _detach_access_group_from_agents(tx, access_group_id) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index c3839d60496..66493743685 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -287,12 +287,12 @@ def compute_tag_metadata_totals(records: Sequence[DailySpendRecord]) -> SpendMet def _entity_metadata( - entity_metadata_field: Mapping[str, dict[str, object]] | None, + entity_metadata_field: Mapping[str, Mapping[str, object]] | None, entity_id: str, ) -> dict[str, object]: """The metadata payload for one entity breakdown bucket, empty when the caller passed none.""" stored: Final = entity_metadata_field.get(entity_id) if entity_metadata_field else None - return stored if stored is not None else {} + return dict(stored) if stored is not None else {} def provider_bucket_metadata(provider: str, base: Mapping[str, object]) -> Mapping[str, object]: @@ -1043,7 +1043,7 @@ def _fold_entity_rollups_sync( results: Sequence[DailySpendData], entity_rows: Sequence[EntityRollupRow], api_key_metadata: Mapping[str, KeyMetadataRow], - entity_metadata_field: Mapping[str, dict[str, object]] | None, # mutable-ok: shared field shape + entity_metadata_field: Mapping[str, Mapping[str, object]] | None, ) -> None: """Write breakdown.entities onto the already-built per-day results.""" by_date: Final = {day.date.strftime("%Y-%m-%d"): day for day in results} @@ -1117,6 +1117,19 @@ async def _guardrail_spend_or_empty( return () +async def _user_entity_metadata( + repository: DailyActivityRepository, entity_rows: Sequence[EntityRollupRow] +) -> Mapping[str, Mapping[str, object]]: + user_ids: Final = frozenset(row.entity_id for row in entity_rows if row.entity_id) + users: Final = await repository.user_metadata(user_ids) + return MappingProxyType( + { + user_id: MappingProxyType({"user_email": user.user_email, "user_alias": user.user_alias}) + for user_id, user in users.items() + } + ) + + async def get_daily_activity_aggregated( repository: DailyActivityRepository, scope: DailyActivityScope, @@ -1161,12 +1174,17 @@ async def get_daily_activity_aggregated( if entity_api_keys else MappingProxyType({}) ) + user_entity_metadata: Final[Mapping[str, Mapping[str, object]]] = ( + await _user_entity_metadata(repository, entity_records) + if scope.table is DailyActivityTable.USER + else MappingProxyType({}) + ) await asyncio.to_thread( _fold_entity_rollups_sync, results=aggregated["results"], entity_rows=entity_records, api_key_metadata=entity_key_metadata, - entity_metadata_field=entity_metadata_field, + entity_metadata_field=MappingProxyType({**user_entity_metadata, **(entity_metadata_field or {})}), ) results: Final = ( list(await asyncio.to_thread(_with_guardrail_breakdown, aggregated["results"], guardrail_rows)) diff --git a/litellm/proxy/management_endpoints/daily_activity_scopes.py b/litellm/proxy/management_endpoints/daily_activity_scopes.py index 7ff73ac6e2e..1eb6ab3210b 100644 --- a/litellm/proxy/management_endpoints/daily_activity_scopes.py +++ b/litellm/proxy/management_endpoints/daily_activity_scopes.py @@ -363,7 +363,7 @@ USER_RESOLVER = EntityScopeResolver( tags=("Budget & Spend Tracking", "Internal User management"), query=_user_query, resolve=_resolve_user, - include_entity_breakdown=False, + include_entity_breakdown=True, include_guardrail_breakdown=True, operation_names=MappingProxyType( { diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d2aac9cae38..93a429a2fcc 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -83,6 +83,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: handle_update_object_permission_common, set_object_permission, ) +from litellm.proxy.management_helpers.team_roster_sync import user_tx_db from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import handle_exception_on_proxy, hash_password from litellm.repositories.organization_repository import OrganizationRepository @@ -125,6 +126,11 @@ _USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) _USER_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 +_LOCK_ALL_USERS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_UserTable" WHERE TRUE ORDER BY user_id COLLATE "C" FOR UPDATE' +) + + @with_config(ConfigDict(extra="allow", strict=True)) class _AdminTeamAndKeyRows(TypedDict): keys: ReadOnly[Sequence[object] | None] @@ -2011,14 +2017,16 @@ async def bulk_user_update( try: # Perform bulk database update - await UserRepository(prisma_client).table.update_many( - where={}, - data=( - {**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])} - if "model_max_budget" in non_default_values - else non_default_values - ), - ) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ALL_USERS_SQL) + _ = await user_tx_db(transaction).update_many( + where={}, + data=( + {**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])} + if "model_max_budget" in non_default_values + else non_default_values + ), + ) for start in range(0, len(all_users_in_db), _USER_CACHE_INVALIDATION_BATCH_SIZE): await asyncio.gather( diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 8a1b9c112dc..4e511aee24a 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -1047,12 +1047,12 @@ async def delete_organization( deleted_orgs: Final = [] for organization_id in data.organization_ids: - # delete all teams in the organization - await _table(TeamRepository(prisma_client)).delete_many(where={"organization_id": organization_id}) - # delete all members in the organization - await _table(OrganizationMembershipRepository(prisma_client)).delete_many( - where={"organization_id": organization_id} - ) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ORG_TEAMS_SQL, organization_id) + _ = await transaction.litellm_teamtable.delete_many(where={"organization_id": organization_id}) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ORG_MEMBERS_SQL, organization_id) + _ = await transaction.litellm_organizationmembership.delete_many(where={"organization_id": organization_id}) await _delete_organization_keys( organization_id=organization_id, prisma_client=prisma_client, @@ -1074,6 +1074,19 @@ async def delete_organization( return deleted_orgs +_LOCK_ORG_TEAMS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE organization_id = $1 ORDER BY team_id COLLATE "C" FOR UPDATE' +) +_LOCK_ORG_MEMBERS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_OrganizationMembership" WHERE organization_id = $1 ' + 'ORDER BY organization_id COLLATE "C", user_id COLLATE "C" FOR UPDATE' +) +_LOCK_ORG_KEYS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE organization_id = $1 ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) + + async def _delete_organization_keys( organization_id: str, prisma_client: PrismaClient, @@ -1087,7 +1100,9 @@ async def _delete_organization_keys( hashed_tokens=hashed_tokens_to_delete, prisma_client=prisma_client, ) - await _table(VerificationTokenRepository(prisma_client)).delete_many(where=key_filter) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ORG_KEYS_SQL, organization_id) + _ = await transaction.litellm_verificationtoken.delete_many(where={"organization_id": organization_id}) await delete_cache_key_objects( hashed_tokens=hashed_tokens_to_delete, user_api_key_cache=user_api_key_cache, diff --git a/litellm/proxy/management_helpers/access_group_model_sync.py b/litellm/proxy/management_helpers/access_group_model_sync.py index 890715bf9c2..20ea2ed0338 100644 --- a/litellm/proxy/management_helpers/access_group_model_sync.py +++ b/litellm/proxy/management_helpers/access_group_model_sync.py @@ -6,6 +6,7 @@ the arrays alone strands every group on a name nothing serves any more. """ from collections.abc import Sequence +from contextlib import AbstractAsyncContextManager from typing import Final, Protocol from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper @@ -53,6 +54,15 @@ _REMOVE_MODEL_NAME_SQL: Final = ( ) +class RawTransactions(Protocol): + def tx(self) -> AbstractAsyncContextManager[RawExecutor]: ... + + +def raw_transactions(prisma_client: object) -> RawTransactions: + db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin + + def raw_executor(prisma_client: object) -> RawExecutor: db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin diff --git a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py index f1a9faf6b2e..0e65e98c557 100644 --- a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py +++ b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py @@ -6,18 +6,27 @@ Those allowlists store public model names, not ids, so a deployment rename that alone denies the new name while the old entry grants a name nothing serves any more. """ -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final +from pydantic import ConfigDict + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, still_backed +from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, raw_transactions, still_backed +from litellm.proxy.management_helpers.access_group_team_sync import TEAM_ADVISORY_LOCK_SQL from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel +class _ListingTeam(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + team_id: str + + class _TouchedRow(LiteLLMBaseModel): kind: str object_id: str @@ -80,11 +89,56 @@ def _rewrite_sql(set_clause: str, where_clause: str) -> str: return f"WITH {ctes} {rows}" +_TEAMS_LISTING_SQL: Final = ( + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE $1 = ANY("models") ORDER BY team_id COLLATE "C"' +) + +_LOCK_ORGS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_OrganizationTable" WHERE $1 = ANY("models") ' + 'ORDER BY organization_id COLLATE "C" FOR UPDATE' +) +_LOCK_TEAMS_LISTING_SQL: Final = ( + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE $1 = ANY("models") ORDER BY team_id COLLATE "C" FOR UPDATE' +) +_LOCK_USERS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_UserTable" WHERE $1 = ANY("models") ORDER BY user_id COLLATE "C" FOR UPDATE' +) +_LOCK_KEYS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE $1 = ANY("models") ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_PROJECTS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_ProjectTable" WHERE $1 = ANY("models") ORDER BY project_id COLLATE "C" FOR UPDATE' +) + _REPLACE_SQL: Final = _rewrite_sql('array_replace(array_remove("models", $2), $1, $2)', '$1 = ANY("models")') _APPEND_SQL: Final = _rewrite_sql('array_append("models", $2)', '$1 = ANY("models") AND NOT ($2 = ANY("models"))') +def _team_ids(rows: Sequence[object]) -> frozenset[str]: + return frozenset(_ListingTeam.model_validate(row).team_id for row in rows) + + +async def _rewrite_allowlists( + prisma_client: object, old_name: str, new_name: str, *, append: bool, attempts: int +) -> Sequence[object]: + """Team advisory locks come from a snapshot, so a team that starts listing the name before its row lock would + be locked without one. That attempt commits having written nothing, and the next one locks the fresh set.""" + async with raw_transactions(prisma_client).tx() as transaction: + listed: Final = _team_ids(await transaction.query_raw(_TEAMS_LISTING_SQL, old_name)) + for team_id in sorted(listed): + _ = await transaction.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) + _ = await transaction.query_raw(_LOCK_ORGS_LISTING_SQL, old_name) + locked: Final = _team_ids(await transaction.query_raw(_LOCK_TEAMS_LISTING_SQL, old_name)) + retry: Final = attempts > 1 and not locked <= listed + if not retry: + _ = await transaction.query_raw(_LOCK_USERS_LISTING_SQL, old_name) + _ = await transaction.query_raw(_LOCK_KEYS_LISTING_SQL, old_name) + _ = await transaction.query_raw(_LOCK_PROJECTS_LISTING_SQL, old_name) + return await transaction.query_raw(_APPEND_SQL if append else _REPLACE_SQL, old_name, new_name) + return await _rewrite_allowlists(prisma_client, old_name, new_name, append=append, attempts=attempts - 1) + + async def sync_model_allowlists_for_renamed_model( prisma_client: object, *, @@ -98,8 +152,8 @@ async def sync_model_allowlists_for_renamed_model( return executor: Final = raw_executor(prisma_client) old_name_still_backed: Final = await still_backed(executor, llm_router, old_name, model_id) - touched_rows: Final = await executor.query_raw( - _APPEND_SQL if old_name_still_backed else _REPLACE_SQL, old_name, new_name + touched_rows: Final = await _rewrite_allowlists( + prisma_client, old_name, new_name, append=old_name_still_backed, attempts=3 ) touched: Final = tuple(_TouchedRow.model_validate(row) for row in touched_rows) await evict_and_broadcast( diff --git a/litellm/proxy/middleware/request_size_limit_middleware.py b/litellm/proxy/middleware/request_size_limit_middleware.py index f2883168d65..329a3c38d91 100644 --- a/litellm/proxy/middleware/request_size_limit_middleware.py +++ b/litellm/proxy/middleware/request_size_limit_middleware.py @@ -5,6 +5,7 @@ from typing import Final from fastapi import HTTPException from starlette._utils import get_route_path +from starlette.routing import compile_path from starlette.types import ASGIApp, Message, Receive, Scope, Send from litellm.proxy.common_utils.anthropic_error_payload import route_error_content @@ -25,7 +26,20 @@ def effective_max_request_size_mb( ) -> float | None: global_limit: Final = configured if configured is not None and configured > 0 else None method: Final = scope.get("method") - route_limit: Final = route_limits.get((method, get_route_path(scope))) if isinstance(method, str) else None + path: Final = get_route_path(scope) + exact_limit: Final = route_limits.get((method, path)) if isinstance(method, str) else None + route_limit: Final = ( + exact_limit + if exact_limit is not None + else next( + ( + limit + for (route_method, route), limit in route_limits.items() + if route_method == method and "{" in route and compile_path(route)[0].fullmatch(path) is not None + ), + None, + ) + ) if route_limit is None: return global_limit return min(route_limit, global_limit) if global_limit is not None else route_limit diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index debefec4a30..c1af5d6f980 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2934,7 +2934,11 @@ def configured_max_request_size_mb() -> int | float | None: _ROUTE_MAX_REQUEST_SIZE_MB: Final[Mapping[tuple[str, str], float]] = MappingProxyType( - {("POST", "/claude-code/plugins/upload"): 11, ("POST", "/skill-hub/skills/upload"): 11} + { + ("POST", "/claude-code/plugins/upload"): 11, + ("POST", "/skill-hub/skills/upload"): 11, + ("PUT", "/skill-hub/skills/{name}/upload"): 11, + } ) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 148d465099f..08942851af9 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -271,6 +271,10 @@ REQUIRED_ONE_OF_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, str]]] = Map {"acreate_interaction": ("model", "agent"), "acreate_agent": ("id", "name")} ) +BODY_PARAM_ALTERNATIVES_BY_ROUTE: Final[Mapping[str, Mapping[str, tuple[str, ...]]]] = MappingProxyType( + {"acreate_interaction": MappingProxyType({"input": ("continuation_token",)})} +) + JSON_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) _OPTIONAL_JSON_OBJECT_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter( Mapping[str, object] | None @@ -320,8 +324,11 @@ def _find_missing_required_body_param( ) if missing_merge_base_param is not None: return MissingBodyParam(name=missing_merge_base_param, model_deployments_loaded=True) + alternatives: Final = BODY_PARAM_ALTERNATIVES_BY_ROUTE.get(route_type, {}) missing_present_params: Final = tuple( - param for param in REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE.get(route_type, ()) if param not in data + param + for param in REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE.get(route_type, ()) + if param not in data and not any(alternative in data for alternative in alternatives.get(param, ())) ) if not missing_present_params: return None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 78f16d9f1fd..818f2dc4858 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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 diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index fcb9eb4a280..0348d8e1f82 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -77,9 +77,6 @@ from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_stat from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SpendLogsRepository from litellm.repositories.team_repository import TeamRepository -from litellm.repositories.verification_token_repository import ( - VerificationTokenRepository, -) from litellm.types.proxy.spend_capture_rate import CaptureRateReport, SpendCaptureProvider if TYPE_CHECKING: @@ -292,24 +289,22 @@ class _TeamTable(Protocol): async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ... -class _VerificationTokenTable(Protocol): - """The subset of the Prisma verification token table API this module uses.""" - - async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ... - - def _spend_logs_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_SpendLogs"]: return SpendLogsRepository(prisma_client).table +_LOCK_ALL_TEAMS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE TRUE ORDER BY team_id COLLATE "C" FOR UPDATE' +) +_LOCK_ALL_KEYS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE TRUE ORDER BY token COLLATE "C" FOR UPDATE' +) + + def _team_table(prisma_client: PrismaClient) -> _TeamTable: return TeamRepository(prisma_client).table -def _verification_token_table(prisma_client: PrismaClient) -> _VerificationTokenTable: - return VerificationTokenRepository(prisma_client).table - - def _spend_logs_daily_summary_sql( *, start_date_iso: str, @@ -3626,8 +3621,11 @@ async def global_spend_reset(): code=status.HTTP_401_UNAUTHORIZED, ) - await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={}) - await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={}) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ALL_TEAMS_SQL) + _ = await transaction.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) + _ = await transaction.query_raw(_LOCK_ALL_KEYS_SQL) + _ = await transaction.litellm_verificationtoken.update_many(data={"spend": 0.0}, where={}) return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a2928bba287..00187790901 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5545,6 +5545,20 @@ class _StaleReadEngine: return self.wrapper is current and self.generation == current.engine_generation +_LOCK_KEYS_BY_TOKEN_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE token = ANY($1::text[]) ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_KEYS_BY_TEAM_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE team_id = ANY($1::text[]) ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_TEAMS_BY_ID_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE team_id = ANY($1::text[]) ' + 'ORDER BY team_id COLLATE "C" FOR UPDATE' +) + + class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() @@ -6845,24 +6859,26 @@ class PrismaClient: else: hashed_token = token hashed_tokens.append(hashed_token) - filter_query: dict[str, object] = {} - if user_id is not None: - filter_query = {"AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}]} - else: - filter_query = {"token": {"in": hashed_tokens}} - - deleted_tokens: Final[int] = await VerificationTokenRepository(self).table.delete_many( - where=filter_query - ) + async with self.tx() as transaction: + _ = await transaction.query_raw(_LOCK_KEYS_BY_TOKEN_SQL, hashed_tokens) + deleted_tokens: Final[int] = await transaction.litellm_verificationtoken.delete_many( + where={"token": {"in": hashed_tokens}, "user_id": user_id} + if user_id is not None + else {"token": {"in": hashed_tokens}} + ) verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens) return {"deleted_keys": deleted_tokens} elif table_name == "team" and team_id_list is not None and isinstance(team_id_list, list): # admin only endpoint -> `/team/delete` - await TeamRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) + async with self.tx() as transaction: + _ = await transaction.query_raw(_LOCK_TEAMS_BY_ID_SQL, team_id_list) + _ = await transaction.litellm_teamtable.delete_many(where={"team_id": {"in": team_id_list}}) return {"deleted_teams": team_id_list} elif table_name == "key" and team_id_list is not None and isinstance(team_id_list, list): # admin only endpoint -> `/team/delete` - await VerificationTokenRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) + async with self.tx() as transaction: + _ = await transaction.query_raw(_LOCK_KEYS_BY_TEAM_SQL, team_id_list) + _ = await transaction.litellm_verificationtoken.delete_many(where={"team_id": {"in": team_id_list}}) except Exception as e: import traceback diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py index 7dd92aecb2d..d5bbe891089 100644 --- a/litellm/repositories/accounting_delivery_repository.py +++ b/litellm/repositories/accounting_delivery_repository.py @@ -5,7 +5,7 @@ from itertools import chain from types import TracebackType from typing import TYPE_CHECKING, Final, Literal, Protocol -from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.repositories.prisma_protocols import TableActions @@ -41,7 +41,10 @@ class _Delivery(BaseModel): delivery_id: str attempts: int - accounting_payload: JsonValue + + +class _ClaimedDelivery(_Delivery): + accounting_payload: str @dataclass(frozen=True, slots=True) @@ -68,9 +71,11 @@ class DeliveryMetrics(BaseModel): _DELIVERIES: Final = TypeAdapter(tuple[_Delivery, ...]) +_CLAIMED: Final = TypeAdapter(tuple[_ClaimedDelivery, ...]) _PAYLOAD: Final = TypeAdapter(dict[str, object]) _TIMEOUT: Final = timedelta(seconds=60) -_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload FROM "LiteLLM_SpendLogDelivery" +_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload::text AS accounting_payload +FROM "LiteLLM_SpendLogDelivery" WHERE state = 'pending' AND next_attempt_at <= CURRENT_TIMESTAMP AND (attempts > 0) = $2::boolean ORDER BY next_attempt_at, delivery_id LIMIT $1::integer FOR UPDATE SKIP LOCKED""" @@ -79,10 +84,20 @@ SET state = 'applied', applied_at = CURRENT_TIMESTAMP, last_error = NULL WHERE delivery_id = ANY($1::text[]) AND state = 'pending' """ _FAILED: Final = """UPDATE "LiteLLM_SpendLogDelivery" SET attempts = attempts + 1, last_error = $3, - next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)))::integer), + next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)) * (0.5 + random()))), state = $4::"AccountingDeliveryState" WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer -RETURNING delivery_id, attempts, accounting_payload""" +RETURNING delivery_id, attempts""" +_UNDECODABLE: Final = """UPDATE "LiteLLM_SpendLogDelivery" +SET state = 'parked', attempts = attempts + 1, last_error = 'UndecodablePayload' +WHERE delivery_id = ANY($1::text[]) AND state = 'pending'""" + + +def _decoded(delivery: _ClaimedDelivery) -> DeliveryProjection | None: + try: + return DeliveryProjection(_PAYLOAD.validate_json(delivery.accounting_payload)) + except ValueError: + return None def _is_permanent_failure(error: Exception) -> bool: @@ -138,26 +153,25 @@ class AccountingDeliveryRepository: try: async with self._database.tx(timeout=_TIMEOUT) as transaction: retries: Final = ( - _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, 1, True)) + _CLAIMED.validate_python(await transaction.query_raw(_CLAIM, 1, True)) if retry_only is not False else () ) - claimed = retries or ( # rebind-ok: retain claimed identities after rollback - _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, batch_size, False)) + rows: Final = retries or ( + _CLAIMED.validate_python(await transaction.query_raw(_CLAIM, batch_size, False)) if retry_only is not True else () ) - if not claimed: - return DeliveryBatchResult() - await apply( - transaction, - tuple( - DeliveryProjection(_PAYLOAD.validate_python(delivery.accounting_payload)) - for delivery in claimed - ), - ) - await transaction.execute_raw(_APPLIED, tuple(delivery.delivery_id for delivery in claimed)) - return DeliveryBatchResult(applied=len(claimed)) + decoded: Final = tuple((row, _decoded(row)) for row in rows) + undecodable: Final = tuple(row.delivery_id for row, projection in decoded if projection is None) + applicable: Final = tuple((row, projection) for row, projection in decoded if projection is not None) + if undecodable: + await transaction.execute_raw(_UNDECODABLE, undecodable) + claimed = tuple(row for row, _ in applicable) # rebind-ok: retain claimed identities after rollback + if claimed: + await apply(transaction, tuple(projection for _, projection in applicable)) + await transaction.execute_raw(_APPLIED, tuple(row.delivery_id for row in claimed)) + return DeliveryBatchResult(applied=len(claimed), parked=undecodable) except Exception as error: # noqa: BLE001 # roll back all destinations before recording a retry if not claimed: raise diff --git a/litellm/repositories/plugin_publication.py b/litellm/repositories/plugin_publication.py new file mode 100644 index 00000000000..f5fab458c6b --- /dev/null +++ b/litellm/repositories/plugin_publication.py @@ -0,0 +1,175 @@ +import hashlib +import json +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Final, Protocol + +from pydantic import TypeAdapter +from semver import Version + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PublishedPlugin, +) +from litellm.repositories.prisma_protocols import TableActions + +if TYPE_CHECKING: + from prisma.models import ( + LiteLLM_ClaudeCodePluginTable, + LiteLLM_ClaudeCodePluginUpload, + LiteLLM_ClaudeCodePluginVersion, + ) + +_MANIFEST: Final = TypeAdapter(dict[str, object]) +_SOURCE: Final = TypeAdapter(dict[str, str]) + + +class PluginPublicationTransaction(Protocol): + @property + def litellm_claudecodeplugintable(self) -> TableActions["LiteLLM_ClaudeCodePluginTable"]: ... + + @property + def litellm_claudecodepluginupload(self) -> TableActions["LiteLLM_ClaudeCodePluginUpload"]: ... + + @property + def litellm_claudecodepluginversion(self) -> TableActions["LiteLLM_ClaudeCodePluginVersion"]: ... + + async def query_raw(self, query: str, *args: object) -> object: ... + + +def resolve_plugin_version( + requested: str | None, current: str | None = None, *, updating: bool = False +) -> str | PluginPublicationError: + try: + previous: Final = Version.parse(current) if current is not None else None + except ValueError: + return _legacy_version(requested) + if updating and previous is None: + return _legacy_version(requested) + try: + candidate: Final = Version.parse(requested) if requested is not None else None + except ValueError: + return PluginPublicationError("invalid_version", "Version must be a valid semantic version, for example 1.0.1") + if candidate is None: + return str(previous.bump_patch()) if previous is not None else "1.0.0" + if previous is not None and candidate.compare(previous) <= 0: + return PluginPublicationError("invalid_version", f"Version must be strictly higher than {current}") + return str(candidate) + + +def _legacy_version(requested: str | None) -> str | PluginPublicationError: + if requested is None: + return PluginPublicationError( + "conflict", + "The published version is missing or is not SemVer. Supply an explicit semantic version to update it", + ) + return resolve_plugin_version(requested) + + +def verified_archive(content: bytes | None, manifest_json: str | None) -> bytes | None: + if content is None: + return None + manifest: Final = _MANIFEST.validate_json(manifest_json or "{}") + source: Final = _SOURCE.validate_python(manifest.get("source", {})) + if source.get("source") != "archive" or source.get("sha256", "").lower() != hashlib.sha256(content).hexdigest(): + return None + return content + + +async def _snapshot( + tx: PluginPublicationTransaction, + plugin: "LiteLLM_ClaudeCodePluginTable", + content: bytes | None, + published_by: str | None, +) -> None: + from prisma import Base64 + + archive: Final = verified_archive(content, plugin.manifest_json) + await tx.litellm_claudecodepluginversion.create( + data={ + "plugin_id": plugin.id, + "version": plugin.version, + "description": plugin.description, + "manifest_json": plugin.manifest_json, + "published_at": plugin.updated_at, + "published_by": published_by, + "content": Base64.encode(archive) if archive is not None else None, + "archive_available": archive is not None, + } + ) + + +async def publish_plugin( + tx: PluginPublicationTransaction, + publication: PluginPublication, + expected: PluginPublicationBase | None = None, +) -> PublishedPlugin | PluginPublicationError: + from prisma import Base64 + + await tx.query_raw('SELECT id FROM "LiteLLM_ClaudeCodePluginTable" WHERE name = $1 FOR UPDATE', publication.name) + current: Final = await tx.litellm_claudecodeplugintable.find_unique(where={"name": publication.name}) + if current is not None and expected is None: + return PluginPublicationError("conflict", f"A skill named '{publication.name}' already exists") + if expected is not None and current is None: + return PluginPublicationError("not_found", f"Plugin '{publication.name}' not found") + if current is not None and ( + expected is not None and (current.id != expected.id or current.updated_at != expected.updated_at) + ): + return PluginPublicationError("conflict", "The published skill changed. Reload it before publishing an update") + version: Final = resolve_plugin_version( + publication.version, current.version if current is not None else None, updating=expected is not None + ) + if isinstance(version, PluginPublicationError): + return version + manifest_json: Final = json.dumps({**publication.manifest, "version": version}) + if publication.content is not None and verified_archive(publication.content, manifest_json) is None: + return PluginPublicationError("conflict", "Uploaded archive does not match its published source digest") + timestamp: Final = datetime.now(timezone.utc) + data: Final = { + "version": version, + "description": publication.description, + "manifest_json": manifest_json, + "files_json": "{}", + "updated_at": timestamp, + } + if current is None: + created: Final = await tx.litellm_claudecodeplugintable.create( + data={**data, "name": publication.name, "created_by": publication.created_by} + ) + if publication.content is not None: + await tx.litellm_claudecodepluginupload.create( + data={"plugin_name": publication.name, "content": Base64.encode(publication.content)} + ) + await _snapshot(tx, created, publication.content, publication.created_by) + return PublishedPlugin(created.id, created.name, version, created.enabled, created.updated_at or timestamp) + previous_archive: Final = await tx.litellm_claudecodepluginupload.find_unique(where={"plugin_name": current.name}) + previous_content: Final = ( + verified_archive(previous_archive.content.decode(), current.manifest_json) + if previous_archive is not None + else None + ) + history_count: Final = await tx.litellm_claudecodepluginversion.count(where={"plugin_id": current.id}, take=1) + if history_count == 0: + await _snapshot(tx, current, previous_content, None) + unchanged_content: Final = ( + previous_content + if _MANIFEST.validate_json(current.manifest_json or "{}").get("source") == publication.manifest.get("source") + else None + ) + content: Final = publication.content if publication.content is not None else unchanged_content + updated: Final = await tx.litellm_claudecodeplugintable.update(where={"id": current.id}, data=data) + if updated is None: + raise RuntimeError("Locked plugin disappeared during publication") + if content is None: + await tx.litellm_claudecodepluginupload.delete_many(where={"plugin_name": current.name}) + else: + await tx.litellm_claudecodepluginupload.upsert( + where={"plugin_name": current.name}, + data={ + "create": {"plugin_name": current.name, "content": Base64.encode(content)}, + "update": {"content": Base64.encode(content)}, + }, + ) + await _snapshot(tx, updated, content, publication.created_by) + return PublishedPlugin(updated.id, updated.name, version, updated.enabled, updated.updated_at or timestamp) diff --git a/litellm/repositories/plugin_upload_repository.py b/litellm/repositories/plugin_upload_repository.py index f37b191e3d6..99e226e442b 100644 --- a/litellm/repositories/plugin_upload_repository.py +++ b/litellm/repositories/plugin_upload_repository.py @@ -1,50 +1,171 @@ -import json -from collections.abc import AsyncIterator, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager +from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # Prisma dynamically exposes its typed table actions +from typing import ( + TYPE_CHECKING, + Final, + Protocol, + cast, # noqa: TID251 # Prisma dynamically exposes its typed table actions +) from pydantic import TypeAdapter, ValidationError -from litellm.models.plugin_upload import PluginArchiveMetadata, PluginUpload, has_agent_skill_metadata +from litellm.models.plugin_upload import ( + PluginArchiveMetadata, + PluginDownloader, + PluginDownloadEvent, + PluginDownloadResult, + PluginDownloadsResponse, + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PluginSnapshot, + PluginUpload, + PluginVersion, + PublishedPlugin, + has_agent_skill_metadata, +) +from litellm.repositories.plugin_publication import PluginPublicationTransaction, publish_plugin, verified_archive from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ClaudeCodePluginRepository if TYPE_CHECKING: from prisma import models as prisma_models + from litellm.proxy.db.create_views import SupportsExecuteRaw from litellm.proxy.utils import PrismaClient _MANIFEST: Final = TypeAdapter(dict[str, object]) _SOURCE: Final = TypeAdapter(dict[str, str]) PAGE_SIZE: Final = 100 +_INCREMENT_ANONYMOUS_DOWNLOADS: Final = """ +UPDATE "LiteLLM_ClaudeCodePluginTable" +SET "anonymous_download_count" = "anonymous_download_count" + 1 +WHERE "id" = $1 +""" +_DOWNLOAD_COUNTS_QUERY: Final = """ +SELECT p.id, p.anonymous_download_count + COUNT(d.id) AS download_count +FROM "LiteLLM_ClaudeCodePluginTable" p +LEFT JOIN "LiteLLM_ClaudeCodePluginDownload" d ON d.plugin_id = p.id +WHERE p.id = ANY($1::text[]) +GROUP BY p.id +""" + + +@dataclass(frozen=True, slots=True) +class _DownloadCount: + id: str + download_count: int + + +_DOWNLOAD_COUNTS: Final = TypeAdapter(tuple[_DownloadCount, ...]) +_VERSIONS: Final = TypeAdapter(tuple[PluginVersion, ...]) +_VERSION_QUERY: Final = """ + SELECT v.id, v.version, v.description, COALESCE(v.manifest_json, '{}')::jsonb AS manifest, + v.published_at, v.published_by, v.archive_available + FROM "LiteLLM_ClaudeCodePluginVersion" v JOIN "LiteLLM_ClaudeCodePluginTable" p ON p.id = v.plugin_id + WHERE p.id = $1 AND p.enabled = $2 +""" + + +class _PublicationDatabase(PluginPublicationTransaction, Protocol): + def tx(self) -> AbstractAsyncContextManager[PluginPublicationTransaction]: ... class PluginUploadRepository: def __init__(self, prisma_client: "PrismaClient") -> None: + self._db = cast( # cast-ok: the generated Prisma transaction delegate is untyped at this boundary + _PublicationDatabase, prisma_client.writer_db + ) + self._writer = cast( # cast-ok: PrismaWrapper dynamically delegates raw statements to its generated client + "SupportsExecuteRaw", prisma_client.writer_db + ) self._plugins = ClaudeCodePluginRepository(prisma_client, use_writer=True).table self._archives = cast( # cast-ok: Prisma exposes generated table actions through an untyped delegate TableActions["prisma_models.LiteLLM_ClaudeCodePluginUpload"], prisma_client.writer_db.litellm_claudecodepluginupload, ) + self._downloads = cast( # cast-ok: Prisma exposes generated table actions through an untyped delegate + TableActions["prisma_models.LiteLLM_ClaudeCodePluginDownload"], + prisma_client.writer_db.litellm_claudecodeplugindownload, + ) async def create(self, upload: PluginUpload) -> str | None: - from prisma import Base64 + result: Final = await self.publish( + PluginPublication( + name=upload.name, + version=upload.version, + description=upload.description, + manifest=upload.manifest, + content=upload.content, + created_by=upload.created_by, + ) + ) + return result.id if isinstance(result, PublishedPlugin) else None + + async def publish( + self, publication: PluginPublication, expected: PluginPublicationBase | None = None + ) -> PublishedPlugin | PluginPublicationError: from prisma.errors import UniqueViolationError try: - plugin: Final = await self._plugins.create( - data={ - "name": upload.name, - "version": upload.version, - "description": upload.description, - "manifest_json": json.dumps(dict(upload.manifest)), - "created_by": upload.created_by, - "uploaded_archive": {"create": {"content": Base64.encode(upload.content)}}, - } - ) + async with self._db.tx() as tx: + return await publish_plugin(tx, publication, expected) except UniqueViolationError: + return PluginPublicationError("conflict", "A skill with this name or version already exists") + + async def versions(self, metadata: PluginArchiveMetadata, limit: int, offset: int) -> tuple[PluginVersion, ...]: + rows: Final = await self._db.query_raw( + _VERSION_QUERY + " ORDER BY v.published_at DESC NULLS LAST, v.id DESC LIMIT $3 OFFSET $4", + metadata.id, + metadata.enabled, + limit, + offset, + ) + return _VERSIONS.validate_python(rows) + + async def version(self, metadata: PluginArchiveMetadata, version_id: str) -> PluginVersion | None: + rows: Final = await self._db.query_raw( + _VERSION_QUERY + " AND v.id = $3", + metadata.id, + metadata.enabled, + version_id, + ) + versions: Final = _VERSIONS.validate_python(rows) + return versions[0] if versions else None + + async def version_content(self, metadata: PluginArchiveMetadata, version_id: str) -> bytes | None: + version: Final = await self._db.litellm_claudecodepluginversion.find_first( + where={ + "id": version_id, + "plugin_id": metadata.id, + "plugin": { + "is": { + "enabled": metadata.enabled, + } + }, + } + ) + if version is None or version.content is None: return None - return plugin.id + return verified_archive(version.content.decode(), version.manifest_json) + + async def get_plugin(self, name: str) -> PluginSnapshot | None: + row: Final = await self._plugins.find_unique(where={"name": name}) + if row is None: + return None + return PluginSnapshot( + id=row.id, + name=row.name, + version=row.version, + description=row.description, + manifest=MappingProxyType(_MANIFEST.validate_json(row.manifest_json or "{}")), + enabled=row.enabled, + created_at=row.created_at, + updated_at=row.updated_at, + created_by=row.created_by, + ) async def get_metadata(self, name: str) -> PluginArchiveMetadata | None: plugin: Final = await self._plugins.find_unique(where={"name": name}) @@ -78,6 +199,7 @@ class PluginUploadRepository: "plugin_name": metadata.name, "plugin": { "is": { + "id": metadata.id, "manifest_json": metadata.manifest_json, "enabled": metadata.enabled, } @@ -86,6 +208,72 @@ class PluginUploadRepository: ) return None if archive is None else archive.content.decode() + async def record_download( + self, metadata: PluginArchiveMetadata, digest: str, downloader: PluginDownloader + ) -> PluginDownloadResult: + from prisma.errors import ForeignKeyViolationError + + if metadata.id is None: + return PluginDownloadResult.SKILL_NOT_FOUND + if not downloader.authenticated: + # A counter bounds anonymous storage; raw SQL preserves Prisma's @updatedAt metadata. + affected: Final = await self._writer.execute_raw(_INCREMENT_ANONYMOUS_DOWNLOADS, metadata.id) + return PluginDownloadResult.RECORDED if affected == 1 else PluginDownloadResult.SKILL_NOT_FOUND + try: + await self._downloads.create( + data={ + "plugin_id": metadata.id, + "version": metadata.version, + "digest": digest, + "authenticated": downloader.authenticated, + "user_id": downloader.user_id, + "user_email": downloader.user_email, + "key_alias": downloader.key_alias, + "team_id": downloader.team_id, + "team_alias": downloader.team_alias, + } + ) + except ForeignKeyViolationError: + return PluginDownloadResult.SKILL_NOT_FOUND + return PluginDownloadResult.RECORDED + + async def downloads(self, name: str, limit: int, cursor: int | None) -> PluginDownloadsResponse | None: + plugin: Final = await self._plugins.find_unique(where={"name": name}) + if plugin is None: + return None + total: Final = await self._downloads.count(where={"plugin_id": plugin.id}) + rows: Final = await self._downloads.find_many( + where={"plugin_id": plugin.id, **({"id": {"lt": cursor}} if cursor is not None else {})}, + order={"id": "desc"}, + take=limit + 1, + ) + return PluginDownloadsResponse( + total_downloads=total + plugin.anonymous_download_count, + anonymous_downloads=plugin.anonymous_download_count, + downloads=tuple( + PluginDownloadEvent( + id=str(row.id), + downloaded_at=row.downloaded_at, + version=row.version, + digest=row.digest, + authenticated=row.authenticated, + user_id=row.user_id, + user_email=row.user_email, + key_alias=row.key_alias, + team_id=row.team_id, + team_alias=row.team_alias, + ) + for row in rows[:limit] + ), + next_cursor=str(rows[limit - 1].id) if len(rows) > limit else None, + ) + + async def download_counts(self, plugin_ids: tuple[str, ...]) -> Mapping[str, int]: + if not plugin_ids: + return MappingProxyType({}) + rows: Final = _DOWNLOAD_COUNTS.validate_python(await self._db.query_raw(_DOWNLOAD_COUNTS_QUERY, plugin_ids)) + return MappingProxyType({row.id: row.download_count for row in rows}) + def _archive_metadata(plugin: "prisma_models.LiteLLM_ClaudeCodePluginTable") -> PluginArchiveMetadata: manifest: Final = _MANIFEST.validate_json(plugin.manifest_json or "{}") @@ -98,4 +286,5 @@ def _archive_metadata(plugin: "prisma_models.LiteLLM_ClaudeCodePluginTable") -> version=plugin.version, description=plugin.description, has_agent_skill_metadata=has_agent_skill_metadata(manifest), + updated_at=plugin.updated_at, ) diff --git a/litellm/repositories/skill_submission_repository.py b/litellm/repositories/skill_submission_repository.py new file mode 100644 index 00000000000..829aa0aadbf --- /dev/null +++ b/litellm/repositories/skill_submission_repository.py @@ -0,0 +1,368 @@ +import hashlib +import json +from collections.abc import Mapping +from contextlib import AbstractAsyncContextManager +from dataclasses import replace +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # Prisma's generated delegates are dynamic +from urllib.parse import urlsplit + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PublishedPlugin, + uploaded_plugin_source, +) +from litellm.models.skill_submission import ( + SkillProposal, + SkillScanInput, + SkillSubmission, + SubmissionConflict, + SubmissionLimitExceeded, + SubmissionMissing, + SubmissionOperation, + SubmissionPage, + SubmissionReviewed, + SubmissionReviewResult, + SubmissionStatus, +) +from litellm.repositories.plugin_publication import PluginPublicationTransaction, publish_plugin, verified_archive +from litellm.repositories.prisma_protocols import TableActions +from litellm.types.proxy.skill_scanning import SkillScan + +if TYPE_CHECKING: + from prisma import models as prisma_models + + from litellm.proxy.utils import PrismaClient + + +class _StoredProposal(BaseModel): + model_config = ConfigDict(frozen=True) + + manifest: Mapping[str, object] + version: str | None + + def proposal(self, name: str) -> SkillProposal: + return SkillProposal( + name, self.version, _DESCRIPTION.validate_python(self.manifest.get("description")), self.manifest + ) + + +class _PendingUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + count: int + storage_bytes: int + + +class _SubmissionRow(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + plugin_name: str + operation: SubmissionOperation + status: SubmissionStatus + proposal_json: str + 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 + + def submission(self) -> SkillSubmission: + stored: Final = _StoredProposal.model_validate_json(self.proposal_json) + return SkillSubmission( + id=self.id, + plugin_name=self.plugin_name, + operation=self.operation, + status=self.status, + proposal=stored.proposal(self.plugin_name), + has_archive=self.has_archive, + submitted_by=self.submitted_by, + created_at=self.created_at, + base_plugin_id=self.base_plugin_id, + base_updated_at=self.base_updated_at, + reviewed_by=self.reviewed_by, + reviewed_at=self.reviewed_at, + rejection_reason=self.rejection_reason, + ) + + +_SCAN_ROWS: Final = TypeAdapter(tuple[Mapping[str, str | None], ...]) +_ROWS: Final = TypeAdapter(tuple[_SubmissionRow, ...]) +_DESCRIPTION: Final = TypeAdapter[str | None](str | None) +_SOURCE: Final = TypeAdapter(dict[str, str]) +_USAGE: Final = TypeAdapter(tuple[_PendingUsage, ...]) +_MAX_PENDING_SUBMISSIONS: Final = 20 +_MAX_PENDING_BYTES: Final = 50 * 1024 * 1024 +_SELECT: Final = """ +SELECT id, plugin_name, operation, status, proposal_json, + archive_content IS NOT NULL AS has_archive, submitted_by, created_at, + base_plugin_id, base_updated_at, reviewed_by, reviewed_at, rejection_reason +FROM "LiteLLM_SkillSubmission" +""" + + +class SubmissionTransaction(PluginPublicationTransaction, Protocol): + @property + def litellm_skillsubmission(self) -> TableActions["prisma_models.LiteLLM_SkillSubmission"]: ... + + +class SubmissionDatabase(SubmissionTransaction, Protocol): + def tx(self) -> AbstractAsyncContextManager[SubmissionTransaction]: ... + + +def _submission(row: "prisma_models.LiteLLM_SkillSubmission") -> SkillSubmission: + return _SubmissionRow.model_validate( + {**row.model_dump(), "has_archive": row.archive_content is not None} + ).submission() + + +class SkillSubmissionRepository: + def __init__(self, prisma_client: "PrismaClient") -> None: + self._database = cast( # cast-ok: Prisma dynamically exposes generated delegates and transactions + SubmissionDatabase, prisma_client.writer_db + ) + + async def get_scan(self, submission_id: str) -> SkillScan | SubmissionMissing | None: + rows: Final = _SCAN_ROWS.validate_python( + await self._database.query_raw( + 'SELECT scan_json FROM "LiteLLM_SkillSubmission" WHERE id = $1', submission_id + ) + ) + if not rows: + return SubmissionMissing() + value: Final = rows[0]["scan_json"] + return None if value is None else SkillScan.model_validate_json(value) + + async def start_scan( + self, submission_id: str, scan: SkillScan + ) -> SkillScanInput | SubmissionMissing | SubmissionConflict: + async with self._database.tx() as transaction: + await transaction.query_raw( + 'SELECT id FROM "LiteLLM_SkillSubmission" WHERE id = $1 FOR UPDATE', submission_id + ) + row: Final = await transaction.litellm_skillsubmission.find_unique(where={"id": submission_id}) + if row is None: + return SubmissionMissing() + if row.status != "pending_review": + return SubmissionConflict("Only pending submissions can be scanned.") + previous: Final = None if row.scan_json is None else SkillScan.model_validate_json(row.scan_json) + if previous is not None and previous.at(scan.started_at).status == "running": + return SubmissionConflict("A scan is already running for this submission.") + stored: Final = _StoredProposal.model_validate_json(row.proposal_json) + source: Final = _SOURCE.validate_python(stored.manifest.get("source", {})) + inherited: Final = ( + await transaction.litellm_claudecodepluginupload.find_unique(where={"plugin_name": row.plugin_name}) + if row.archive_content is None and source.get("source") == "archive" + else None + ) + content: Final = ( + row.archive_content.decode() + if row.archive_content is not None + else verified_archive( + None if inherited is None else inherited.content.decode(), json.dumps(stored.manifest) + ) + ) + managed_archive: Final = urlsplit(source.get("url", "")).path.endswith( + f"/claude-code/plugins/{row.plugin_name}/archive/{source.get('sha256', '')}.zip" + ) + if source.get("source") == "archive" and managed_archive and content is None: + return SubmissionConflict("The proposal's archive is no longer available. Submit a new ZIP to scan it.") + started: Final = scan.model_copy( + update={"archive_sha256": None if content is None else hashlib.sha256(content).hexdigest()} + ) + await transaction.litellm_skillsubmission.update( + where={"id": submission_id}, data={"scan_json": started.model_dump_json()} + ) + return SkillScanInput(started, source, content) + + async def finish_scan(self, submission_id: str, scan: SkillScan) -> None: + async with self._database.tx() as transaction: + rows: Final = _SCAN_ROWS.validate_python( + await transaction.query_raw( + 'SELECT scan_json FROM "LiteLLM_SkillSubmission" WHERE id = $1 FOR UPDATE', submission_id + ) + ) + value: Final = rows[0]["scan_json"] if rows else None + if value is None: + return + current: Final = SkillScan.model_validate_json(value) + if current.id != scan.id or current.status != "running": + return + result: Final = ( + current.at(scan.completed_at) + if scan.completed_at is not None and scan.completed_at >= current.expires_at + else scan + ) + await transaction.litellm_skillsubmission.update( + where={"id": submission_id}, data={"scan_json": result.model_dump_json()} + ) + + async def submit( + self, upload: PluginPublication, submitted_by: str, base: PluginPublicationBase | None + ) -> SkillSubmission | SubmissionLimitExceeded: + from prisma import Base64 + + proposal_json: Final = _StoredProposal(manifest=dict(upload.manifest), version=upload.version).model_dump_json() + incoming_bytes: Final = len(proposal_json.encode("utf-8")) + len(upload.content or b"") + async with self._database.tx() as transaction: + await transaction.query_raw( + "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtext('litellm.skill_submissions'), hashtext($1))", + submitted_by, + ) + usage: Final = _USAGE.validate_python( + await transaction.query_raw( + """SELECT COUNT(*) AS count, + COALESCE(SUM(octet_length(proposal_json) + + COALESCE(octet_length(archive_content), 0)), 0) AS storage_bytes + FROM "LiteLLM_SkillSubmission" WHERE submitted_by = $1 AND status = 'pending_review'""", + submitted_by, + ) + )[0] + if usage.count >= _MAX_PENDING_SUBMISSIONS: + return SubmissionLimitExceeded( + "You already have 20 pending skill submissions. " + "Wait for an admin to review one before submitting another." + ) + if usage.storage_bytes + incoming_bytes > _MAX_PENDING_BYTES: + return SubmissionLimitExceeded( + "Pending skill submissions cannot exceed 50 MiB of stored metadata and archives per user. " + "Wait for an admin review or submit less data." + ) + row: Final = await transaction.litellm_skillsubmission.create( + data={ + "plugin_name": upload.name, + "operation": "create" if base is None else "update", + "proposal_json": proposal_json, + "archive_content": None if upload.content is None else Base64.encode(upload.content), + "submitted_by": submitted_by, + "base_plugin_id": None if base is None else base.id, + "base_updated_at": None if base is None else base.updated_at, + } + ) + return _submission(row) + + async def list( + self, submitted_by: str | None, status: SubmissionStatus | None, limit: int, offset: int + ) -> SubmissionPage: + rows: Final = _ROWS.validate_python( + await self._database.query_raw( + _SELECT + + """WHERE ($1::text IS NULL OR submitted_by = $1) + AND ($2::text IS NULL OR status = $2) + ORDER BY created_at DESC, id DESC LIMIT $3 OFFSET $4""", + submitted_by, + status, + limit + 1, + offset, + ) + ) + return SubmissionPage(tuple(row.submission() for row in rows[:limit]), has_more=len(rows) > limit) + + async def get(self, submission_id: str, submitted_by: str | None) -> SkillSubmission | None: + rows: Final = _ROWS.validate_python( + await self._database.query_raw( + _SELECT + "WHERE id = $1 AND ($2::text IS NULL OR submitted_by = $2)", + submission_id, + submitted_by, + ) + ) + return None if not rows else rows[0].submission() + + async def get_archive(self, submission_id: str, submitted_by: str | None) -> bytes | None: + row: Final = await self._database.litellm_skillsubmission.find_first( + where={"id": submission_id, **({"submitted_by": submitted_by} if submitted_by is not None else {})} + ) + return None if row is None or row.archive_content is None else row.archive_content.decode() + + async def review( + self, + submission_id: str, + reviewer: str, + approve: bool, + reason: str | None, + *, + publication_base_url: str | None = None, + ) -> SubmissionReviewResult: + from prisma.errors import UniqueViolationError + + try: + async with self._database.tx() as transaction: + await transaction.query_raw( + 'SELECT id FROM "LiteLLM_SkillSubmission" WHERE id = $1 FOR UPDATE', submission_id + ) + row: Final = await transaction.litellm_skillsubmission.find_unique(where={"id": submission_id}) + if row is None: + return SubmissionMissing() + submission: Final = _submission(row) + if submission.status != "pending_review": + return SubmissionConflict("This submission has already been reviewed.") + content: Final = None if row.archive_content is None else row.archive_content.decode() + if approve and content is not None and not publication_base_url: + return SubmissionConflict("Uploaded skill approval requires a trusted publication URL.") + source: Final = ( + uploaded_plugin_source(publication_base_url, row.plugin_name, hashlib.sha256(content).hexdigest()) + if approve and content is not None and publication_base_url + else _SOURCE.validate_python(submission.proposal.manifest.get("source")) + ) + publication: Final = await _publish_submission(transaction, row, source, content) if approve else None + if isinstance(publication, PluginPublicationError): + return publication + reviewed_at: Final = datetime.now(timezone.utc) + status: Final[SubmissionStatus] = "approved" if approve else "rejected" + await transaction.litellm_skillsubmission.update( + where={"id": submission_id}, + data={ + "status": status, + "reviewed_by": reviewer, + "reviewed_at": reviewed_at, + "rejection_reason": None if approve else reason, + }, + ) + return SubmissionReviewed( + submission=replace( + submission, + status=status, + reviewed_by=reviewer, + reviewed_at=reviewed_at, + rejection_reason=None if approve else reason, + ), + plugin=publication, + published_source=source if approve else None, + ) + except UniqueViolationError: + return SubmissionConflict("A skill with this name is already published. Submit a new update proposal.") + + +async def _publish_submission( + transaction: PluginPublicationTransaction, + row: "prisma_models.LiteLLM_SkillSubmission", + source: Mapping[str, str], + content: bytes | None, +) -> PublishedPlugin | PluginPublicationError: + stored: Final = _StoredProposal.model_validate_json(row.proposal_json) + proposal: Final = stored.proposal(row.plugin_name) + base: Final = ( + None + if row.operation == "create" + else PluginPublicationBase(id=row.base_plugin_id or "", updated_at=row.base_updated_at) + ) + return await publish_plugin( + transaction, + PluginPublication( + name=row.plugin_name, + version=stored.version, + description=proposal.description, + manifest={**stored.manifest, "source": dict(source)}, + content=content, + created_by=row.submitted_by, + ), + base, + ) diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index f6aad883479..b78f7e5a9b7 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -109,6 +109,46 @@ class TeamMembershipSpendResetWrites: self.batch.execute_raw(_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL, list(budget_ids)) +_LOCK_BUDGET_ORGS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_OrganizationTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY organization_id COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_MEMBERSHIPS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamMembership" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY team_id COLLATE "C", user_id COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_KEYS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_PROJECTS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_ProjectTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY project_id COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_TAGS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TagTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY tag_name COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_ACCESS_GROUPS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_ModelAccessGroupBudgetTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY access_group_name COLLATE "C" FOR UPDATE' +) + + +@dataclass(frozen=True, slots=True) +class BudgetLinkedRowLocks: + batch: PrismaBatch + + def queue_row_locks(self, budget_ids: Sequence[str]) -> None: + """Locks every row a cascade may reset, in the shared row lock order, before the filtered resets run.""" + self.batch.execute_raw(_LOCK_BUDGET_ORGS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_MEMBERSHIPS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_KEYS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_PROJECTS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_TAGS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_ACCESS_GROUPS_SQL, list(budget_ids)) + + @dataclass(frozen=True, slots=True) class BudgetWindowWrites: table: BatchTable @@ -137,6 +177,7 @@ class BudgetCascadeUnitOfWork: ``enduser_spend_reset_unit_of_work`` because one tier can gate millions. """ + row_locks: BudgetLinkedRowLocks team_memberships: TeamMembershipSpendResetWrites keys: LinkedSpendResetWrites organizations: LinkedSpendResetWrites @@ -163,6 +204,7 @@ async def budget_cascade_unit_of_work( ) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]: batch = new_batch() yield BudgetCascadeUnitOfWork( + row_locks=BudgetLinkedRowLocks(batch=batch), team_memberships=TeamMembershipSpendResetWrites( resets=LinkedSpendResetWrites(table=batch.litellm_teammembership), batch=batch ), diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py index 5cb256a8ac0..be331a91475 100644 --- a/litellm/responses/mcp/stream_round_billing.py +++ b/litellm/responses/mcp/stream_round_billing.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Callable, Coroutine, Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass from datetime import datetime from typing import Final, TypeAlias @@ -297,12 +297,8 @@ class StreamedRoundBilling: self._close(parked) self._bank((self._dispatch(row, billed, self._logging_obj.attempt_scoped_copy()),), row) - def _dispatch( - self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj - ) -> BilledRoundsDispatch | Coroutine[object, object, None]: + def _dispatch(self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj) -> BilledRoundsDispatch: dispatch: Final = self._log_success(_SuccessLog(logger, row, billed.chat_usage)) - if billed.chat_usage is None: - return dispatch return BilledRoundsDispatch(dispatch, billed.spend) @staticmethod diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 249c6370b86..288e4a7119d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -53,7 +53,7 @@ from litellm.litellm_core_utils.internal_call_metadata import ( requests_background, without_budget_reservation, ) -from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers @@ -638,6 +638,24 @@ class BaseResponsesAPIStreamingIterator: self._handle_failure(e) raise + def _deferred_logging_dispatch( + self, logging_coroutine: Coroutine[object, object, None], logging_response: object + ) -> BilledRoundsDispatch | Coroutine[object, object, None]: + response: Final[object] = getattr(logging_response, "response", None) + if not isinstance(response, ResponsesAPIResponse): + return logging_coroutine + price: Final = ( + UncachedPrice(0.0, None) + if self._completed_response_cache_hit + else self.logging_obj.uncached_price(response) + ) + spend: Final = BilledRoundsSpend( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response.usage), + price.cost or 0.0, + price.cost_breakdown, + ) + return BilledRoundsDispatch(logging_coroutine, spend) + def _log_completed_response(self, *, is_async: bool) -> None: if self._completed_response_logged: return @@ -662,7 +680,9 @@ class BaseResponsesAPIStreamingIterator: if deferred_dispatch_armed: # End-of-stream guardrail scans write guardrail_information after # the terminal event; dispatching now would snapshot metadata early. - self.logging_obj._deferred_stream_complete_args = (logging_coroutine,) + self.logging_obj._deferred_stream_complete_args = ( + self._deferred_logging_dispatch(logging_coroutine, logging_response), + ) else: spawn_logging_task( asyncio.get_running_loop(), logging_coroutine, task_name="Responses stream success logging" diff --git a/litellm/router.py b/litellm/router.py index 6a60d5d958e..a5d72daea1e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1,41 +1,208 @@ """ litellm.Router: load balancing, retries, fallbacks and cooldowns across deployments. -``Router`` holds no routing logic. It builds its configuration (``litellm/router_config.py``), picks -one ``RouterCore`` at construction (``litellm/router_core.py``) and delegates every member to it: +``Router`` holds no routing logic. It builds its configuration (``litellm/router_config.py``), builds +its ``RouterCore`` from it (the Rust core, ``litellm/router_core.py``) and delegates every member to it: the endpoint methods (acompletion, aembedding, aresponses, aanthropic_messages, ...) and the entry points (registry, settings, health, select, apply, observe, register_plugin, retire). The members outside code used before the entry points live in ``DeprecatedRouterMembers`` until they are removed. """ +import functools import inspect import weakref -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence from types import FunctionType, MappingProxyType -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Any, Final, Literal -from litellm import router_python_core from litellm.router_config import RouterConfig, build_router_config from litellm.router_core import RouterCore, make_core from litellm.router_deprecated import DeprecatedRouterMembers -from litellm.router_python_core import ( - PythonCore, - as_output_cap, - model_info_is_active_for_environment, - stamped_client_ceiling, -) from litellm.router_python_host import PythonHost, python_host +from litellm.router_rust.steps.prerouting import as_output_cap, stamped_client_ceiling +from litellm.router_rust.surface.registry import model_info_is_active_for_environment +from litellm.types.litellm_params import RoutingStrategyName from litellm.types.router import ( + AlertingConfig, AllowedFailsPolicy, AssistantsTypedDict, + AutoRouterCapabilityLimit, CustomRoutingStrategyBase, Deployment, + DeploymentTypedDict, + FallbackAccessCheck, + FallbackBudgetCheck, + GuardrailTypedDict, LiteLLM_Params, ModelGroupInfo, + OptionalPreCallChecks, RetryPolicy, + RouterGeneralSettings, + RouterModelGroupAliasItem, + RoutingGroup, + RoutingPlugin, + SearchToolTypedDict, ) +from litellm.types.utils import GenericBudgetConfigType -_CORE_SIGNATURE: Final = inspect.signature(PythonCore.__init__) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # reflection reads only the parameters' names +if TYPE_CHECKING: + from litellm.router_rust.core import RustCore + + +# ``Router(...)``'s parameters, their defaults and its docstring: ``router_arguments`` binds a call +# against them and ``Router.__init__`` reports them. The body is never run. +def _router_init( + self, + model_list: list[DeploymentTypedDict] | list[dict[str, Any]] | None = None, + ## ASSISTANTS API ## + assistants_config: AssistantsTypedDict | None = None, + ## SEARCH API ## + search_tools: list[SearchToolTypedDict] | None = None, + ## GUARDRAIL API ## + guardrail_list: list[GuardrailTypedDict] | None = None, + ## CACHING ## + redis_url: str | None = None, + redis_host: str | None = None, + redis_port: int | None = None, + redis_password: str | None = None, + redis_db: int | None = None, + cache_responses: bool | None = False, + cache_kwargs: dict = {}, # additional kwargs to pass to RedisCache (see caching.py) + caching_groups: list[tuple] | None = None, # if you want to cache across model groups + client_ttl: int = 3600, # ttl for cached clients - will re-initialize after this time in seconds + ## SCHEDULER ## + polling_interval: float | None = None, + default_priority: int | None = None, + ## RELIABILITY ## + num_retries: int | None = None, + max_fallbacks: int | None = None, # max fallbacks to try before exiting the call. Defaults to 5. + timeout: float | None = None, + stream_timeout: float | None = None, + default_litellm_params: dict | None = None, # default params for Router.chat.completion.create + default_max_parallel_requests: int | None = None, + set_verbose: bool = False, + debug_level: Literal["DEBUG", "INFO"] = "INFO", + default_fallbacks: list[str] | None = None, # generic fallbacks, works across all deployments + fallbacks: list = [], + context_window_fallbacks: list = [], + content_policy_fallbacks: list = [], + model_group_alias: dict[str, str | RouterModelGroupAliasItem] | None = {}, + enable_pre_call_checks: bool = False, + enable_tag_filtering: bool = False, + tag_filtering_match_any: bool = True, + tag_routing_prefix: str = "", + plugins: list[RoutingPlugin] | None = None, + retry_after: int = 0, # min time to wait before retrying a failed request + retry_policy: RetryPolicy | dict | None = None, # set custom retries for different exceptions + model_group_retry_policy: dict[str, RetryPolicy] = {}, # set custom retry policies based on model group + allowed_fails: int | None = None, # Number of times a deployment can failbefore being added to cooldown + allowed_fails_policy: AllowedFailsPolicy | None = None, # set custom allowed fails policy + cooldown_time: float | None = None, # (seconds) time to cooldown a deployment after failure + disable_cooldowns: bool | None = None, + fallback_on_unsupported_params: bool = True, + routing_strategy: RoutingStrategyName = "simple-shuffle", + optional_pre_call_checks: OptionalPreCallChecks | None = None, + routing_strategy_args: dict = {}, # just for latency-based + routing_groups: list[RoutingGroup | dict] | None = None, + provider_budget_config: GenericBudgetConfigType | None = None, + alerting_config: AlertingConfig | None = None, + router_general_settings: RouterGeneralSettings | None = RouterGeneralSettings(), + deployment_affinity_ttl_seconds: int = 3600, + model_group_affinity_config: dict[str, list[str]] | None = None, + ignore_invalid_deployments: bool = False, + enable_health_check_routing: bool = False, + health_check_staleness_threshold: int | None = None, + health_check_ignore_transient_errors: bool = False, + background_health_check_model_groups: Sequence[str] | None = None, + enable_weighted_failover: bool = False, + enable_mid_stream_fallback_continuation: bool = False, + fallback_access_check: FallbackAccessCheck | None = None, + fallback_budget_check: FallbackBudgetCheck | None = None, + auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, + session_affinity_repin_on_failover: bool = False, + _host: PythonHost = python_host, +) -> None: + """ + Initialize the Router class with the given parameters for caching, reliability, and routing strategy. + + Args: + model_list (Optional[list]): List of models to be used. Defaults to None. + redis_url (Optional[str]): URL of the Redis server. Defaults to None. + redis_host (Optional[str]): Hostname of the Redis server. Defaults to None. + redis_port (Optional[int]): Port of the Redis server. Defaults to None. + redis_password (Optional[str]): Password of the Redis server. Defaults to None. + cache_responses (Optional[bool]): Flag to enable caching of responses. Defaults to False. + cache_kwargs (dict): Additional kwargs to pass to RedisCache. Defaults to {}. + caching_groups (Optional[List[tuple]]): List of model groups for caching across model groups. Defaults to None. + client_ttl (int): Time-to-live for cached clients in seconds. Defaults to 3600. + polling_interval: (Optional[float]): frequency of polling queue. Only for '.scheduler_acompletion()'. Default is 3ms. + default_priority: (Optional[int]): the priority given to every chat or text completion request that sets no 'priority' of its own; '.acompletion()' and '.atext_completion()' queue it through the scheduler. Default is None. + num_retries (Optional[int]): Number of retries for failed requests. Defaults to 2. + timeout (Optional[float]): Timeout for requests. Defaults to None. + default_litellm_params (dict): Default parameters for Router.chat.completion.create. Defaults to {}. + set_verbose (bool): Flag to set verbose mode. Defaults to False. + debug_level (Literal["DEBUG", "INFO"]): Debug level for logging. Defaults to "INFO". + fallbacks (List): List of fallback options. Defaults to []. + context_window_fallbacks (List): List of context window fallback options. Defaults to []. + enable_pre_call_checks (boolean): Filter out deployments which are outside context window limits for a given prompt + model_group_alias (Optional[dict]): Alias for model groups. Defaults to {}. + retry_after (int): Minimum time to wait before retrying a failed request. Defaults to 0. + allowed_fails (Optional[int]): Number of allowed fails before adding to cooldown. Defaults to None. + cooldown_time (float): Time to cooldown a deployment after failure in seconds. Defaults to 1. + routing_strategy (Literal["simple-shuffle", "least-busy", "usage-based-routing", "latency-based-routing", "cost-based-routing"]): Routing strategy used for the implicit "default" group (any model not claimed by an entry in `routing_groups`). Defaults to "simple-shuffle". + routing_strategy_args (dict): Additional args for the default group's routing strategy (e.g. latency window). Defaults to {}. + routing_groups (Optional[List[RoutingGroup]]): Named subsets of `model_name`s with a group routing strategy. + Priority groups apply only to group calls and may overlap. Other groups supply their members' default + strategy, with at most one such group per model. Unclaimed models use the top-level strategy. + alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None. + provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None. + deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. + session_affinity_repin_on_failover (bool): After a successful replacement response, move a session pin + whose deployment is in cooldown or removed. Async routes only. Defaults to False. Applied at startup. + ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. + enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. + enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False. + fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). + fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback). + Returns: + Router: An instance of the litellm.Router class. + + Example Usage: + ```python + from litellm import Router + model_list = [ + { + "model_name": "azure-gpt-3.5-turbo", # model alias + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/", + "api_key": , + "api_version": , + "api_base": + }, + }, + { + "model_name": "azure-gpt-3.5-turbo", # model alias + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/", + "api_key": , + "api_version": , + "api_base": + }, + }, + { + "model_name": "openai-gpt-3.5-turbo", # model alias + "litellm_params": { # params for litellm completion/embedding call + "model": "gpt-3.5-turbo", + "api_key": , + }, + ] + + router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}]) + ``` + """ + + +_CORE_SIGNATURE: Final = inspect.signature(_router_init) def router_arguments( @@ -55,55 +222,129 @@ def _class_descriptor(cls: type, name: str) -> bool: if TYPE_CHECKING: + from typing import overload + + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.types.llms.openai import AllMessageValues, OpenAIFileObject + from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, ModelResponse + + class Router(RustCore, DeprecatedRouterMembers): + """What type checkers see: the constructor's parameters and the endpoints' return types. The + endpoints are served by the core at run time (litellm/router_rust/invoke/endpoints.py).""" + + __init__ = _router_init # pyright: ignore[reportAssignmentType, reportIncompatibleMethodOverride] # Router(...)'s parameters - class Router(PythonCore, DeprecatedRouterMembers): @classmethod def from_config(cls, config: RouterConfig, host: PythonHost = python_host) -> "Router": ... + @overload + async def acompletion( + self, model: str, messages: list[AllMessageValues], stream: Literal[True], **kwargs: object + ) -> CustomStreamWrapper: ... + @overload + async def acompletion( + self, model: str, messages: list[AllMessageValues], stream: Literal[False] = False, **kwargs: object + ) -> ModelResponse: ... + @overload + async def acompletion( + self, model: str, messages: list[AllMessageValues], stream: bool = False, **kwargs: object + ) -> CustomStreamWrapper | ModelResponse: ... + async def acompletion( # pyright: ignore[reportIncompatibleMethodOverride] # the endpoint's typed signature + self, model: str, messages: list[AllMessageValues], stream: bool = False, **kwargs: object + ) -> CustomStreamWrapper | ModelResponse: ... + def completion( # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion + self, model: str, messages: list[dict[str, str]], **kwargs: object + ) -> ModelResponse | CustomStreamWrapper: ... + def embedding( # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion + self, model: str, input: str | list[object], is_async: bool | None = False, **kwargs: object + ) -> EmbeddingResponse: ... + async def aembedding( # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion + self, model: str, input: str | list[object], is_async: bool | None = True, **kwargs: object + ) -> EmbeddingResponse: ... + async def acreate_file(self, model: str, **kwargs: object) -> OpenAIFileObject: ... # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion + async def acreate_batch(self, model: str, **kwargs: object) -> LiteLLMBatch: ... # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion + async def aretrieve_batch(self, model: str | None = None, **kwargs: object) -> LiteLLMBatch: ... # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion + def __getattr__(self, name: str) -> Any: ... # the endpoints served by name + else: _MISSING: Final = object() _replaced_core_members: Final[ # mutable-ok: originals of patched core members, restored on delete dict[tuple[type, str], list[object]] ] = {} + def _core_class() -> type: + from litellm.router_rust.core import RustCore + + return RustCore + class _RouterType(type): """ - Class-level reads of core members (``Router._helper(...)``) go to ``PythonCore``, and so do - class-level patches, which put the original back when the patch is removed. - ``isinstance`` and ``dir`` see ``PythonCore`` too. + Class-level reads of core members (``Router._helper``) go to the core's class, and so do + class-level patches, which put the original back when the patch is removed. ``dir`` lists + the core's members and endpoints too, so ``Mock(spec=Router)`` accepts them. """ def __getattr__(cls, name: str) -> object: - return getattr(PythonCore, name) + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + core: Final = _core_class() + try: + return getattr(core, name) + except AttributeError: + if not _is_endpoint(name): + raise + return _endpoint_function(name) def __setattr__(cls, name: str, value: object) -> None: - if any(name in vars(base) for base in cls.__mro__) or not hasattr(PythonCore, name): + core: Final = _core_class() + if any(name in vars(base) for base in cls.__mro__) or not (hasattr(core, name) or _is_endpoint(name)): super().__setattr__(name, value) return - _replaced_core_members.setdefault((PythonCore, name), []).append(vars(PythonCore).get(name, _MISSING)) - setattr(PythonCore, name, value) + _replaced_core_members.setdefault((core, name), []).append(vars(core).get(name, _MISSING)) + setattr(core, name, value) def __delattr__(cls, name: str) -> None: - replaced: Final = _replaced_core_members.get((PythonCore, name)) + core: Final = _core_class() + replaced: Final = _replaced_core_members.get((core, name)) if not replaced or any(name in vars(base) for base in cls.__mro__): super().__delattr__(name) return original: Final = replaced.pop() if original is _MISSING: - delattr(PythonCore, name) + delattr(core, name) return - setattr(PythonCore, name, original) - - def __instancecheck__(cls, instance: object) -> bool: - return type.__instancecheck__(cls, instance) or isinstance(instance, PythonCore) + setattr(core, name, original) def __dir__(cls) -> list[str]: # mutable-ok: dir() returns a list - return sorted({*type.__dir__(cls), *dir(PythonCore)}) + from litellm.router_rust.invoke.endpoints import endpoint_names + + return sorted({*type.__dir__(cls), *dir(_core_class()), *endpoint_names()}) + + def _is_endpoint(name: str) -> bool: + from litellm.router_rust.invoke.endpoints import endpoint_names + + return name in endpoint_names() + + @functools.cache + def _endpoint_function(name: str) -> Callable[..., object]: + """An endpoint the core serves by name, read off the class (``Router.aretrieve_batch``): a + function that calls the instance's endpoint, so a class-level patch of it replaces it.""" + + def endpoint( + self: object, + *args: object, + **kwargs: object, # kwargs-ok: the endpoint's own arguments + ) -> object: + served: Final = _core_class().__getattr__(self, name) # pyright: ignore[reportAny] # the core's endpoint lookup + return served(*args, **kwargs) # pyright: ignore[reportAny] # see above + + endpoint.__name__ = endpoint.__qualname__ = name + return endpoint def _patchable_method(core: object, name: str) -> bool: """A method of the core: an instance-level replacement of it is also kept on the facade, so ``mock.patch.object(router, name)`` finds it there and nested patches unwind in order. A core - that serves its members from another core (``__getattr__``) is asked for the bound method.""" + that serves its members from ``__getattr__`` is asked for the bound method.""" static: Final = inspect.getattr_static(type(core), name, None) if static is not None: return isinstance(static, FunctionType | staticmethod | classmethod) @@ -183,17 +424,13 @@ else: delattr(self._core, name) def __dir__(self) -> list[str]: # mutable-ok: dir() returns a list - return sorted({*super().__dir__(), *dir(self.__dict__.get("_core", PythonCore))}) + core: Final = self.__dict__.get("_core") + return sorted({*super().__dir__(), *(dir(core) if core is not None else ())}) Router.__init__.__signature__ = _CORE_SIGNATURE.replace( # pyright: ignore[reportFunctionMemberAccess] # inspect reads it parameters=[parameter for parameter in _CORE_SIGNATURE.parameters.values() if parameter.name != "_host"] ) - Router.__init__.__doc__ = PythonCore.__init__.__doc__ - - -def __getattr__(name: str) -> object: - """Module-level helpers that lived in this module before the core moved out.""" - return getattr(router_python_core, name) + Router.__init__.__doc__ = _router_init.__doc__ __all__ = [ diff --git a/litellm/router_api.py b/litellm/router_api.py index 82660439ae9..dc95b742d56 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1491,7 +1491,9 @@ class RoutingContext: return self._router._params_sent_to_deployment(request, deployment) def default_litellm_params(self) -> Mapping[str, object]: - return self._router._active_default_litellm_params() + """The router's ``default_litellm_params`` that are set (a None value is no default).""" + configured: Final = self.settings().default_litellm_params + return {key: value for key, value in (configured or {}).items() if value is not None} def settings(self) -> RouterSettings: return self._router.settings() diff --git a/litellm/router_attempt.py b/litellm/router_attempt.py index 33e0a19c530..b7b8d2ac87d 100644 --- a/litellm/router_attempt.py +++ b/litellm/router_attempt.py @@ -3,8 +3,7 @@ One router attempt, as the host runs it: the request for the deployment the rout The Rust router core picks a deployment and hands the attempt to its host; ``PythonHost`` builds the keyword arguments for that one deployment here and calls ``litellm.`` with them. Nothing -here selects, retries, falls back or records cooldowns: the core does all of that. Nothing here -reads a Python router either, so this module outlives ``router_python_core.py``. +here selects, retries, falls back or records cooldowns: the core does all of that. """ import re @@ -61,6 +60,16 @@ POSITIONAL: Final = MappingProxyType( ) +# The arguments an endpoint takes without a default, which a call that leaves one out raises +# ``TypeError`` for before it routes. +REQUIRED: Final = MappingProxyType( + { + "atext_completion": ("model", "prompt"), + "text_completion": ("model", "prompt"), + } +) + + def request_from_call(op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]: """The call's arguments as one keyword request. Raises ``TypeError`` like the endpoint would.""" names: Final = POSITIONAL.get(op, ()) @@ -71,6 +80,12 @@ def request_from_call(op: str, args: tuple[object, ...], kwargs: Mapping[str, ob if name in request: raise TypeError(f"{op}() got multiple values for argument {name!r}") request[name] = value + missing: Final = [name for name in REQUIRED.get(op, ()) if name not in request] + if missing: + listed: Final = " and ".join(repr(name) for name in missing) + raise TypeError( + f"{op}() missing {len(missing)} required positional argument{'s' if len(missing) > 1 else ''}: {listed}" + ) return request diff --git a/litellm/router_core.py b/litellm/router_core.py index c077b3f0015..a906e4d86fb 100644 --- a/litellm/router_core.py +++ b/litellm/router_core.py @@ -1,16 +1,13 @@ """ -``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which picks one. +``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which builds it. -``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. ``RustCore`` -(``litellm/router_rust/core.py`` over litellm-rust/crates/router-core) is a separate implementation -that reaches outside itself only through its host; it does not use ``PythonCore``. The core is -picked once, when the router is built, from its ``RouterConfig``. +The core is ``RustCore`` (``litellm/router_rust/core.py`` over litellm-rust/crates/router-core), +which reaches outside itself only through its host. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Protocol, runtime_checkable -from litellm._logging import verbose_router_logger from litellm.router_config import RouterConfig if TYPE_CHECKING: @@ -52,25 +49,7 @@ class RouterCore(Protocol): def make_core(config: RouterConfig, host: "PythonHost") -> RouterCore: - """ - The Rust core when the route catalog's router rule allows it and the Rust core serves this - config, else the Python core. Decided once, for the router's life, from the config alone: the - Rust core never builds, wraps or falls back to a Python core. - """ - from litellm.router_python_core import PythonCore - from litellm.rust_bridge.catalog import Route, RouteContext, decision - from litellm.rust_bridge.configuration import Decision + """The router's core: the Rust core, built once for the router's life from its config.""" + from litellm.router_rust.core import RustCore - chosen: Final = decision(RouteContext(Route.ROUTER)) - if chosen is not Decision.PYTHON: - from litellm.router_rust.core import RustCore - from litellm.rust_bridge._native import RustBridgeDeclined - - try: - return RustCore.build(config, host) # pyright: ignore[reportReturnType] # RustCore serves the RouterCore entry points; other members raise RustCoreUnsupported - except RustBridgeDeclined as declined: - # A config the Rust core cannot take keeps the router on Python while both cores exist. - if chosen is Decision.RUST_REQUIRED: - raise - verbose_router_logger.info("Router served by the Python core: %s", declined) - return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature + return RustCore.build(config, host) # pyright: ignore[reportReturnType] # RustCore serves the RouterCore entry points; other members raise RustCoreUnsupported diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 9d915fe0982..c2e1c197259 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -92,10 +92,9 @@ def _note(name: str, replacement: str) -> None: def _as_router(obj: "DeprecatedRouterMembers") -> "Router": """The Router this mixin is part of (it is mixed into ``litellm.Router`` only).""" from litellm.router import Router - from litellm.router_python_core import PythonCore - if isinstance(obj, (Router, PythonCore)): - return obj # pyright: ignore[reportReturnType] # a PythonCore carries every member the mixin reads + if isinstance(obj, Router): + return obj raise TypeError(f"{type(obj).__name__} is not a litellm.Router; DeprecatedRouterMembers is mixed into Router only") @@ -112,9 +111,12 @@ def _public_attribute_error(obj: object, core_name: str, public_name: str, error class _DeprecatedAttribute(Generic[_T]): """A former Router attribute, read through an entry point. Assignment still replaces it.""" - def __init__(self, replacement: str, read: Callable[["Router"], _T]) -> None: + def __init__(self, replacement: str, read: Callable[["Router"], _T], unbound: object = None) -> None: self.replacement = replacement self.read = read + # What a class-level read returns when the core keeps no member of this name: a former + # method reads back as a function taking the router, as it did on the Router class. + self.unbound = unbound self.name = "" def __set_name__(self, owner: type, name: str) -> None: @@ -129,7 +131,7 @@ class _DeprecatedAttribute(Generic[_T]): def __get__(self, obj: "DeprecatedRouterMembers | None", owner: type) -> object: if obj is None: - return getattr(owner, "_core_" + self.name, self) + return getattr(owner, "_core_" + self.name, self if self.unbound is None else self.unbound) _note(self.name, self.replacement) router: Final = _as_router(obj) try: @@ -216,6 +218,23 @@ def _unreserved_picker(router: "Router") -> Callable[..., Awaitable[object]]: return picker +async def _async_get_available_deployment( + router: "Router", + model: str, + request_kwargs: dict, # mutable-ok: the old parameter type, unchanged + messages: list[dict[str, str]] | None = None, # mutable-ok: the old parameter type, unchanged + input: str | list | None = None, # mutable-ok: the old parameter type, unchanged + specific_deployment: bool | None = False, +) -> object: + """``Router.async_get_available_deployment`` read on the class: the router's pick, unbound.""" + return await _unreserved_picker(router)(model, request_kwargs, messages, input, specific_deployment) + + +_async_get_available_deployment.__name__ = _async_get_available_deployment.__qualname__ = ( + "async_get_available_deployment" +) + + def _slack_alerting_logger(router: "Router") -> "SlackAlerting": """``observe().slack_alerting_logger``; a router without one has no such attribute, as before.""" logger: Final = router.observe().slack_alerting_logger @@ -282,7 +301,9 @@ class DeprecatedRouterMembers: cooldown_cache = _attr("health().cooldowns", lambda r: r.health().cooldowns) health_state_cache = _attr("health().health_checks", lambda r: r.health().health_checks) deployment_latency_map = _attr("health().latency", lambda r: r.health().latency) - async_get_available_deployment = _attr("select(model, request_kwargs)", lambda r: _unreserved_picker(r)) + async_get_available_deployment = _DeprecatedAttribute( + "select(model, request_kwargs)", lambda r: _unreserved_picker(r), unbound=_async_get_available_deployment + ) # -- observers ------------------------------------------------------------------------------ slack_alerting_logger = _attr("observe().slack_alerting_logger", _slack_alerting_logger) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py deleted file mode 100644 index 8ca2afb9a99..00000000000 --- a/litellm/router_python_core.py +++ /dev/null @@ -1,18017 +0,0 @@ -# +-----------------------------------------------+ -# | | -# | Give Feedback / Get Help | -# | https://github.com/BerriAI/litellm/issues/new | -# | | -# +-----------------------------------------------+ -# -# Thank you ! We ❤️ you! - Krrish & Ishaan - -""" -``PythonCore``: the Python implementation of the router core behind ``litellm.Router``. - -It holds all routing logic (registry, selection, retries, fallbacks, cooldowns, usage) and reaches -outside itself only through its ``PythonHost``: inference calls, sleeps, callback registration and -the process settings (``litellm/router_config.py``). The Rust core in litellm-rust/crates/router-core -implements the same interface. ``litellm/router.py`` picks one at construction and delegates to it. -""" - -import asyncio -import contextlib -import copy -import enum -import inspect -import json -import logging -import re -import threading -import time -import traceback -import weakref -from collections import defaultdict -from collections.abc import ( - AsyncGenerator, - AsyncIterable, - AsyncIterator, - Awaitable, - Callable, - Generator, - Iterable, - Iterator, - Mapping, - MutableMapping, - Sequence, -) -from contextvars import ContextVar -from dataclasses import dataclass -from datetime import datetime, timezone -from functools import lru_cache, partial -from itertools import chain -from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast - -import anyio -import httpx -import openai -from openai import AsyncOpenAI -from pydantic import BaseModel, TypeAdapter, ValidationError -from typing_extensions import assert_never, overload - -from litellm import exceptions as litellm_exceptions -from litellm import get_secret_str -from litellm import utils as litellm_utils -from litellm._internal_context import in_mirrored_call, mirrored_call, service_target, with_service_target -from litellm._logging import verbose_router_logger -from litellm._uuid import uuid -from litellm.caching.caching import Cache as LitellmCache -from litellm.caching.caching import ( - DualCache, - InMemoryCache, - RedisCache, - RedisClusterCache, -) -from litellm.caching.redis_cache import log_redis_failure -from litellm.constants import ( - BAGGAGE_MAPPED_TAGS_METADATA_KEY, - CLIENT_OUTPUT_CEILING_METADATA_KEY, - CONSUMED_REQUEST_TAGS_METADATA_KEY, - DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, - DEFAULT_HEALTH_CHECK_INTERVAL, - DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, - DEFAULT_MAX_LRU_CACHE_SIZE, - INTERNAL_CALL_ORIGIN_METADATA_KEY, - KUBERNETES_POD_ROUTING_KEY, - MID_STREAM_FAILURE_LOG_WAIT_SECONDS, - OUTPUT_TOKEN_CEILING_PARAMS, - PROMPT_CACHE_PIN_CACHE_MAX_SIZE, - REALTIME_SESSION_ENDED_ABNORMALLY_KEY, - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, - ROUTING_REQUEST_TAGS_METADATA_KEY, - RUNTIME_UPDATABLE_ROUTER_SETTINGS, - SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, -) -from litellm.integrations.custom_guardrail import CustomGuardrail, is_guardrail_intervention -from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.otel.routing import routing_decision_attributes -from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span -from litellm.litellm_core_utils import exception_mapping_utils -from litellm.litellm_core_utils.access_windows import access_windows_config_error, filter_reserved_deployments -from litellm.litellm_core_utils.add_retry_fallback_headers import ( - add_fallback_headers_to_response, - add_retry_headers_to_response, - apply_quality_router_decision_headers, - apply_remaining_usage_headers, - apply_response_model_id, - complexity_router_decision_headers, - ensure_response_additional_headers, - get_hidden_params_dict, - prepare_response_for_header_attachment, - replace_complexity_router_headers, - response_total_token_count, -) -from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure -from litellm.litellm_core_utils.asyncify import run_async_function -from litellm.litellm_core_utils.batch_utils import ( - batch_create_routing, - deployment_supported_endpoints, - get_router_metadata_variable_name, - is_batch_retrieve_call_type, - replace_model_in_jsonl, - should_replace_model_in_jsonl, - without_deployments_that_cannot_batch, -) -from litellm.litellm_core_utils.clientside_credential_handler import ( - get_dynamic_litellm_params, - is_clientside_credential, -) -from litellm.litellm_core_utils.completion_timeout import CompletionTimeout -from litellm.litellm_core_utils.core_helpers import ( - as_str_mapping, - coerce_token_limit, - get_litellm_metadata_from_kwargs, - get_metadata_variable_name_from_kwargs, - get_or_create_metadata_bucket, - get_parent_otel_span_from_kwargs, -) -from litellm.litellm_core_utils.coroutine_checker import coroutine_checker -from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference -from litellm.litellm_core_utils.dd_tracing import tracer -from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models, deployment_pricing_models -from litellm.litellm_core_utils.get_litellm_params import SUPPLEMENTAL_KWARGS_KEYS -from litellm.litellm_core_utils.get_llm_provider_logic import ( - declared_authenticating_provider, - inferred_provider, - is_registered_custom_provider, - resolve_model_and_provider_for_metadata, -) -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap -from litellm.litellm_core_utils.get_retry_from_policy import ( - get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, -) -from litellm.litellm_core_utils.get_retry_from_policy import resolve_retry_policy -from litellm.litellm_core_utils.hidden_params import ( - HIDDEN_PARAMS_ATTR as _HIDDEN_PARAMS_ATTR, -) -from litellm.litellm_core_utils.hidden_params import ( - get_hidden_params, - set_hidden_params, -) -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES -from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter -from litellm.litellm_core_utils.ptu_pricing import ( - PTU_COST_ATTRIBUTION_ENV_VAR, - declares_ptu, - declares_ptu_shares, - is_ptu_cost_attribution_enabled, - ptu_config_error, - ptu_identity_error, - ptu_terms, - zeroed_ptu_pricing, -) -from litellm.litellm_core_utils.ptu_shares import ( - PTUTeamCeiling, - filter_ptu_shared_deployments, - ptu_capacity_warning, - ptu_share_teams, - share_holder_request_kwargs, - team_ptu_ceiling, -) -from litellm.litellm_core_utils.reasoning_effort_capability import ( - deployment_is_catalog_mapped, - intersect_supported_reasoning_efforts, - resolve_supported_reasoning_efforts, -) -from litellm.litellm_core_utils.request_timeout_resolver import ( - get_configured_request_timeout, - get_service_tier_timeout, -) -from litellm.litellm_core_utils.routing_common_utils import ( - MID_STREAM_CONTINUATION_KWARG, - as_retry_skipped_deployment_ids, - deployment_model_id_pin, - deployments_the_caller_may_use, - filter_team_based_models, - filter_web_search_deployments, - format_content_policy_fallback_hint, - format_content_policy_no_fallback_note, - format_fallback_outcome_message, - format_no_fallback_group_message, - get_request_team_id, - is_proxy_admin_request, - provider_for_generic_call, - request_user_api_key_auth, - resolve_model_group_alias, - team_may_use_deployment, - truncate_fallback_error_detail, - warn_on_ignored_model_capability_flags, - warn_on_provider_credential_mismatch, - with_admin_api_key_secret_reference, - without_model_capability_flags, - without_router_only_kwargs, -) -from litellm.litellm_core_utils.routing_groups import ( - apply_routing_group_priority, - parse_routing_groups, - validate_routing_strategy, -) -from litellm.litellm_core_utils.secret_redaction import redact_string -from litellm.litellm_core_utils.sensitive_data_masker import ( - SensitiveDataMasker, - mask_sensitive_structure, -) -from litellm.litellm_core_utils.token_counter import ( - offload_token_count, - served_openai_models, - single_encoding_openai_model, -) -from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy -from litellm.llms.anthropic.common_utils import AnthropicModelInfo, anthropic_error_frame_exception -from litellm.llms.base_llm.passthrough.transformation import replace_path_segment -from litellm.llms.base_llm.vector_store.transformation import ( - RouterVectorStoreEmbeddingExecutor, - vector_store_request_metadata, -) -from litellm.llms.bedrock.common_utils import get_bedrock_base_model -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client -from litellm.llms.openai_like.json_loader import JSONProviderRegistry -from litellm.llms.openai_like.model_info import ( - MODEL_INFO_DISCOVERY_PROVIDERS, - MODEL_INFO_REFRESH_CONCURRENCY, - MODEL_INFO_REFRESH_SECONDS, - get_openai_compatible_model_info, -) -from litellm.main import Chat as LitellmChat -from litellm.responses.bridge_check import keep_registered_mode -from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults -from litellm.router_api import ( - PLAIN_SETTING_ATTRIBUTES, - SETTING_FIELDS, - ApplyResult, - ChangeSet, - RouterHealth, - RouterRegistry, - RouterSettings, - RoutingObservers, - SelectKind, - aapply_changes, - apply_changes, -) -from litellm.router_api import RoutingContext as RouterContext -from litellm.router_api import deployment_id as _deployment_id -from litellm.router_deprecated import DeprecatedRouterMembers -from litellm.router_host import router_host -from litellm.router_python_host import PythonHost, python_host -from litellm.router_strategy.auto_router.auto_router_model_naming import ( - AUTO_ROUTER_MODEL_PREFIX, - STRATEGY_ROUTER_PARAM_FIELDS, - GatedAutoRouterCapability, - capability_limit_violation, - claimed_capabilities, - classify_strategy_router_model, - count_capability_routers, -) -from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy -from litellm.router_strategy.budget_limiter import RouterBudgetLimiting -from litellm.router_strategy.complexity_router.context_compaction import ( - arm_compaction, - compact_to_fit, - compaction_pending, - initialize_compaction_state, - is_native_compaction_call, - reject_recursive_compactor, - surface_for_call, -) -from litellm.router_strategy.least_busy import ( - CALL_IDENTITY, - SELECTION_RESERVATION_MAX_KEYS, - LeastBusyLoggingHandler, - call_id_of, - owning_picks, - pick_reserves, - picks_reserve, -) -from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler -from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler -from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler -from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage -from litellm.router_strategy.simple_shuffle import simple_shuffle -from litellm.router_strategy.tag_based_routing import ( - baggage_stamp_for_serving_deployment, - get_deployments_for_tag, - get_routing_tags_from_request_kwargs, - get_tags_from_request_kwargs, - is_valid_deployment_tag, - tags_with_serving_deployment_tags, -) -from litellm.router_utils.client_initalization_utils import InitalizeCachedClient, MaxParallelRequestsLimit -from litellm.router_utils.cooldown_cache import CooldownCache -from litellm.router_utils.cooldown_handlers import ( - DEFAULT_COOLDOWN_TIME_SECONDS, - _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper across router_utils submodules, matching the other cooldown_handlers imports on this line - async_get_cooldown_deployments, - async_get_cooldown_deployments_with_debug_info, - get_cooldown_deployment_info, - get_cooldown_deployments, - is_background_response_poll_not_found, - is_caller_entra_credential_failure, - is_caller_scoped_auth_failure, - is_caller_timeout_408, - is_failure_outside_selected_deployment, - is_video_lookup_not_found, - mark_deployment_selection_failure, - set_cooldown_deployments, -) -from litellm.router_utils.deployment_numbers import refuse_unservable_numbers -from litellm.router_utils.embedding_batch_size_check import ( - async_log_rejection_before_provider_call, - clear_rejected_before_provider_call, - deployments_admitting_embedding_batch, - deployments_when_none_admits_embedding_batch, - embedding_batch_too_large_error, - embedding_deployment_selection, - log_rejection_before_provider_call, -) -from litellm.router_utils.endpoint_eligibility import ( - eligible_deployments, - endpoint_selection, - warn_on_unenforced_endpoint_flag, -) -from litellm.router_utils.fallback_event_handlers import ( - MID_STREAM_FALLBACK_CONTROLS_KEY, - AttemptedFallbackTargets, - attempted_fallback_targets_from, - attempted_retries_for_request, - carry_over_pre_routing_selection, - carry_over_routed_deployment, - carry_superseded_attempt_usage, - check_non_standard_fallback_format, - clear_pre_routing_selection, - committed_retry_budget_for_request, - continuation_fallback_targets, - fallback_lookup_groups, - fallback_target_name, - fallback_target_overrides_request_params, - fallbacks_disabled_for_request, - get_fallback_model_group_for_lookup_groups, - get_pre_routing_selection, - has_unattempted_fallback_target, - is_bad_request_without_fallback, - mid_stream_fallback_hop_kwargs, - mid_stream_fallback_snapshot_kwargs, - mid_stream_retry_kwargs, - per_request_fallback_controls, - record_disable_fallbacks, - record_pre_routing_selection, - record_retry_attempt, - routed_deployment_id, - run_async_fallback, -) -from litellm.router_utils.global_callback_holders import ( - live_check_owners, - routers_holding_global_callbacks, - routers_holding_global_callbacks_lock, -) -from litellm.router_utils.handle_error import ( - async_raise_no_deployment_exception, - send_llm_exception_alert, -) -from litellm.router_utils.health_state_cache import DeploymentHealthCache -from litellm.router_utils.kubernetes_pod_discovery import ( - KubernetesPodDiscovery, - async_resolve_pods_after, - async_resolve_pods_after_bound, - resolve_pods_after, - resolve_pods_after_bound, -) -from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_MARKER, - RESPONSES_CONTINUATION_KWARG, - ContinuationOutput, - ContinuationPrefillDeploymentCheck, - ContinuedPartial, - ResponsesContinuation, - ResponsesContinuationSeam, - constrains_continuation, - constrains_responses_continuation, - continuation_deployments, - continuation_output_ceilings, - continued_text, - deployment_can_continue, - responses_continuation_input, - responses_continuation_input_for, -) -from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( - DeploymentAffinityCheck, - warn_on_unknown_model_group_affinity_flags, -) -from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, -) -from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - async_admit_io_token_turn, - async_io_token_refund_failure, - build_io_token_rate_limit_headers, - deployment_has_io_token_limits, - refund_stale_reservation_before_retry, - set_io_token_rate_limit_request_kwargs, - without_io_token_limits, -) -from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( - ModelRateLimitingCheck, -) -from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( - PromptCachingDeploymentCheck, -) -from litellm.router_utils.request_priority import ( - InvalidPriority, - parse_default_priority, - priority_error, - priority_header_fields, - request_drops_params, - resolve_request_priority, - unqueueable_priority_message, -) -from litellm.router_utils.router_callbacks.track_deployment_metrics import ( - find_deployment_metadata, - get_counted_usage_tokens, - increment_deployment_failures_for_current_minute, - increment_deployment_successes_for_current_minute, -) -from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch -from litellm.router_utils.stream_usage_injection import router_config_requests_stream_usage -from litellm.scheduler import FlowItem, Scheduler -from litellm.secret_managers.main import str_to_bool -from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD, RoutingStrategyName -from litellm.types.llms.openai import ( - AllMessageValues, - ChatCompletionToolParam, - FileTypes, - OpenAIFileObject, - OpenAIFilesPurpose, -) -from litellm.types.router import ( - CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, - VALID_LITELLM_ENVIRONMENTS, - AlertingConfig, - AllowedFailsPolicy, - AssistantsTypedDict, - AutoRouterCapabilityLimit, - ConsumedRequestTagsStamp, - CredentialLiteLLMParams, - CustomRoutingStrategyBase, - Deployment, - DeploymentModelListingInfo, - DeploymentTypedDict, - DiscoveredDeploymentModelInfo, - FallbackAccessCheck, - FallbackBudgetCheck, - GenericLiteLLMParams, - GuardrailTypedDict, - LiteLLM_Params, - LiteLLMParamsTypedDict, - MockRouterTestingParams, - ModelGroupInfo, - OptionalPreCallChecks, - PreRoutingStrategy, - RetryPolicy, - RouterCacheEnum, - RouterErrors, - RouterGeneralSettings, - RouterModelGroupAliasItem, - RouterRateLimitError, - RouterRateLimitErrorBasic, - RoutingContext, - RoutingGroup, - RoutingPlugin, - RoutingStrategy, - SearchToolTypedDict, - TaggedPreRoutingStrategy, - holds_secret_pointer, - validate_litellm_params_strictly, -) -from litellm.types.services import ServiceTypes -from litellm.types.utils import ( - AUTOROUTER_CLASSIFIER_CALL_ORIGIN, - COST_MAP_LOOKUP_KEY, - DEPLOYMENT_SCOPED_TOKEN_LIMIT_FIELDS, - PROMPT_QUOTING_ROUTING_DECISION_FIELDS, - CustomPricingLiteLLMParams, - GenericBudgetConfigType, - LiteLLMBatch, - LlmProviders, - LlmProvidersSet, - ModelInfo, - ModelResponseStream, - StandardLoggingPayload, - StandardLoggingRoutingDecision, - Usage, - all_litellm_params, - echoed_cost_map_fields, - is_server_derived_pricing_key, - shared_backend_model_info, -) -from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES, decode_video_id_with_provider -from litellm.utils import ( - CustomStreamWrapper, - EmbeddingResponse, - ModelInfoLookupMiss, - ModelResponse, - ProviderConfigManager, - Rules, - TextCompletionResponse, - _restore_correlation_context_if_supported, # pyright: ignore[reportPrivateUsage] # same task-local cleanup as the SDK wrapper - aresolve_api_key_secret_reference, - catalog_with_runtime_registrations, - function_setup, - get_llm_provider, - get_non_default_completion_params, - get_potential_model_names, - get_secret, - get_utc_datetime, - is_generalized_model_info, - is_region_allowed, - lookup_model_info_for_registration, - model_info_with_pricing_overrides, - provider_rejectable_params, - set_live_deployment_replay, - token_counter, -) - -if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span - - from litellm.exceptions import MidStreamFallbackError - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - from litellm.router_api import UntypedDict - from litellm.router_strategy.adaptive_router.adaptive_router import ( - AdaptiveRouter, - ) - from litellm.router_strategy.auto_router.auto_router import ( - AutoRouter, - PreRoutingHookResponse, - ) - from litellm.router_strategy.complexity_router.complexity_router import ( - ComplexityRouter, - ) - from litellm.router_strategy.quality_router.quality_router import ( - QualityRouter, - ) - from litellm.types.llms.anthropic_messages.anthropic_response import ( - AnthropicMessagesResponse, - ) - from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject - from litellm.types.llms.openai import ( - ResponseAPIUsage, - ResponseIncompleteEvent, - ResponsesAPIResponse, - ResponsesAPIStreamingResponse, - ) - - Span = _Span -else: - Span = Any - AutoRouter = Any - ComplexityRouter = Any - AdaptiveRouter = Any - QualityRouter = Any - PreRoutingHookResponse = Any - -RouterStrategySelector: TypeAlias = ( - LeastBusyLoggingHandler - | LowestCostLoggingHandler - | LowestLatencyLoggingHandler - | LowestTPMLoggingHandler - | LowestTPMLoggingHandler_v2 -) - - -def _cost_value_as_float(value: str | float | None) -> float | None: - if value is None: - return None - try: - return float(value) - except (TypeError, ValueError): - return None - - -def model_info_is_active_for_environment(model_info: Mapping[str, object] | None) -> bool: - """Single owner of the environment-gating rule: a deployment whose model_info names - `supported_environments` loads only on pods whose LITELLM_ENVIRONMENT is in that list. - `Router.deployment_is_active_for_environment` delegates here, and the model-write - endpoints consult the same rule to tell a deliberately inactive model from one that - was dropped by a failed reload.""" - if model_info is None: - return True - supported_environments: Final = model_info.get("supported_environments") - if supported_environments is None: - return True - if not isinstance(supported_environments, (list, tuple)): - raise ValueError( - f"supported_environments must be a list of {VALID_LITELLM_ENVIRONMENTS}. " - f"but set as: {supported_environments} for model_info: {model_info}" - ) - litellm_environment: Final = get_secret_str(secret_name="LITELLM_ENVIRONMENT") - if litellm_environment is None: - raise ValueError("Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set in .env") - - if litellm_environment not in VALID_LITELLM_ENVIRONMENTS: - raise ValueError( - f"LITELLM_ENVIRONMENT must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {litellm_environment}" - ) - - for _env in supported_environments: - if _env not in VALID_LITELLM_ENVIRONMENTS: - raise ValueError( - f"supported_environments must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {_env} " - f"for model_info: {model_info}" - ) - - if litellm_environment in supported_environments: - return True - return False - - -_PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") -_CallbackT = TypeVar("_CallbackT") -_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) - - -_ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( - {"model", "api_base", "api_key", "api_version", "responses_api_version"} -) -_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" -_FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS: Final = frozenset({"chatgpt", "anthropic"}) -_CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") -_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 -CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" - -_PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) -_STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS: Final = frozenset({"anthropic_messages", "aresponses", "responses"}) -_video_call_deployment_id: Final[ContextVar[str | None]] = ContextVar("video_call_deployment_id", default=None) -_routed_call_type: Final[ContextVar[str | None]] = ContextVar("routed_call_type", default=None) -_DECISIONS_MEDIA_CALL_TYPES: Final = frozenset({"adecisions", "decisions", "asystemone", "systemone"}) -_ROUTE_CAPABILITY_RESOLVERS: Final[Mapping[str, Callable[[str, LlmProviders], object | None]]] = MappingProxyType( - { - "adecisions": ProviderConfigManager.get_provider_decisions_config, - "decisions": ProviderConfigManager.get_provider_decisions_config, - } -) -_RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( - { - "prompt_caching": PromptCachingDeploymentCheck, - "enforce_model_rate_limits": ModelRateLimitingCheck, - "encrypted_content_affinity": EncryptedContentAffinityCheck, - } -) - - -def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) -> bool: - for chunk in chunks: - if not chunk.choices: - continue - delta = chunk.choices[0].delta - if ( - delta.get("content") - or delta.get("tool_calls") - or delta.get("function_call") - or delta.get("reasoning_content") - or delta.get("thinking_blocks") - or delta.get("reasoning_items") - or delta.get("audio") - or delta.get("images") - or delta.get("annotations") - ): - return True - return False - - -_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) -_NO_DEFAULT_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) -_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_CALL_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str]) -_RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None) -_ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) -_FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_DEPLOYMENT_PARAMS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | GenericLiteLLMParams]] = TypeAdapter( - Mapping[str, object] | GenericLiteLLMParams -) -_OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) -_DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected" -_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) -_DEPLOYMENTS_ADAPTER: Final = TypeAdapter(Sequence[Mapping[str, object]]) - - -def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]: - """Continue the last text block in place so the prefill reads as the text the caller already received.""" - last: Final = blocks[-1] if blocks else None - if last is not None and last.get("type") == "text" and isinstance(last_text := last.get("text"), str): - return [*blocks[:-1], {**last, "text": f"{last_text}{text}"}] - return [*blocks, {"type": "text", "text": text}] - - -def _is_fallback_hop(request_kwargs: Mapping[str, object]) -> bool: - fallback_depth: Final = request_kwargs.get("fallback_depth") - return isinstance(fallback_depth, int) and fallback_depth > 0 - - -def _deployment_that_just_failed(request_metadata: object) -> str | None: - try: - model_info: Final = _FALLBACK_HOP_ADAPTER.validate_python( - _FALLBACK_HOP_ADAPTER.validate_python(request_metadata).get("model_info") - ) - except ValidationError: - return None - model_id: Final = model_info.get("id") - return model_id if isinstance(model_id, str) else None - - -def _drops_malformed_request_tools( - deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object] -) -> bool: - if isinstance(request.get("tools") or [], (list, tuple)) or not deployment_params.get("tools"): - return False - set_defaults: Final = {name: value for name, value in defaults.items() if value is not None} - return drop_policy({**deployment_params, **set_defaults, **request}).drops("tools") - - -def _malformed_request_params_the_deployment_sets( - deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object] -) -> frozenset[str]: - overrides: Final = { - key: value - for key, value in request.items() - if key in SUPPLEMENTAL_KWARGS_KEYS and deployment_params.get(key) is not None - } - if not overrides: - return frozenset() - try: - validate_litellm_params_strictly(GenericLiteLLMParams, overrides) - except ValidationError as error: - set_defaults: Final = {name: value for name, value in defaults.items() if value is not None} - policy: Final = drop_policy({**deployment_params, **set_defaults, **request}) - malformed: Final = (str(detail["loc"][0]) for detail in error.errors() if detail["loc"]) - return frozenset(key for key in malformed if policy.drops(key)) - return frozenset() - - -def _drop_malformed_request_params_the_deployment_sets( - kwargs: dict[str, object], # mutable-ok: drops the keys from the router's request kwargs in place - deployment_params: Mapping[str, object], - defaults: Mapping[str, object], -) -> None: - for key in _malformed_request_params_the_deployment_sets(deployment_params, defaults, kwargs): - kwargs.pop(key) - - -def _merged_tools(deployment_tools: Sequence[object], request_tools: object) -> list[object] | None: - if not isinstance(request_tools, (list, tuple)): - return None - return [*deployment_tools, *request_tools] if deployment_tools or request_tools else None - - -def _deployment_pick_attributes(model: str, request_kwargs: Mapping[str, object] | None) -> Mapping[str, str | int]: - """Bounded attributes for one deployment pick; attempt is 1-based within the current model group.""" - kwargs: Final = request_kwargs or {} - metadata: Final = kwargs.get("litellm_metadata", kwargs.get("metadata")) - attempted_retries: Final = metadata.get("attempted_retries") if isinstance(metadata, Mapping) else None - retries: Final = attempted_retries if isinstance(attempted_retries, int) else 0 - reason: Final = "retry" if retries > 0 else "fallback" if _is_fallback_hop(kwargs) else "initial" - return MappingProxyType( - { - "litellm.deployment.attempt": retries + 1, - "litellm.deployment.reason": reason, - "litellm.deployment.model_group": model, - } - ) - - -def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, object]: - return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) - - -def _operator_max_input_tokens(deployment: Mapping[str, object]) -> int | None: - model_info: Final = _configured_model_info(deployment) - catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) - catalog_entry: Final = ( - GetModelCostMap.loaded_model_cost_map().get(catalog_key) if isinstance(catalog_key, str) else None - ) - if catalog_entry is not None and "max_input_tokens" in echoed_cost_map_fields(model_info, catalog_entry): - return None - return coerce_token_limit(model_info.get("max_input_tokens")) - - -def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None: - from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request - - configured: Final = _configured_model_info(deployment).get("supports_search_tool") - if isinstance(configured, bool): - return configured - params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) - model: Final = params.get("model") - provider: Final = params.get("custom_llm_provider") - if not isinstance(model, str): - return None - try: - resolved_model, resolved_provider = resolve_model_and_provider_for_metadata( - model=model, custom_llm_provider=provider if isinstance(provider, str) else None - ) - chat_config: Final = ProviderConfigManager.get_provider_chat_config( - model=resolved_model, provider=LlmProviders(resolved_provider) - ) - except (litellm_exceptions.BadRequestError, ValueError): - return None - return True if chat_config_builds_anthropic_request(chat_config) else None - - -def _every_deployment_refused( - deployments: Sequence[Mapping[str, object]], refused_deployment_ids: tuple[str, ...] -) -> bool: - return all(_configured_model_info(deployment).get("id") in refused_deployment_ids for deployment in deployments) - - -def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]: - if silent_model is None or in_mirrored_call(): - return () - try: - targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model) - except ValidationError: - verbose_router_logger.warning( - "silent_model must be a model name or a list of model names, got %r; skipping shadow traffic", - silent_model, - ) - return () - return (targets,) if isinstance(targets, str) else tuple(targets) - - -def _silent_experiment_kwargs_snapshot(kwargs: Mapping[str, object]) -> Mapping[str, object]: - metadata: Final = kwargs.get("metadata") - if not isinstance(metadata, Mapping): - return MappingProxyType({**kwargs}) - return MappingProxyType({**kwargs, "metadata": dict(metadata)}) - - -def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: - """ - Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still - holds the pre-routing model group name, so it has to follow the deployment the router just picked. - - Returns kwargs to merge into the downstream call, empty when there is no session model to resolve. - """ - try: - typed_session: Final = _SESSION_ADAPTER.validate_python(session) - except ValidationError: - return _NO_SESSION_KWARGS - if "model" not in typed_session and typed_session.get("type") != "transcription": - return _NO_SESSION_KWARGS - from litellm.realtime_api.main import with_resolved_session_model - - return MappingProxyType({"session": with_resolved_session_model(dict(typed_session), model_name)}) - - -# Router._aanthropic_messages_streaming_iterator buffers lifecycle chunks -# until real content commits the primary stream, and only while a fallback -# can still take over; a hostile or slow-starting upstream that never emits -# content or an error could otherwise grow that buffer without bound, so -# hitting this cap forces an early commit instead. -MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS: Final = 200 - - -def _anthropic_stream_forwards_ping_live(chunk: object, has_generated_content: bool) -> bool: - """A `ping` keepalive reaches the client live whenever the stream has not committed: it carries no - lifecycle, so it cannot create overlapping lifecycles on the wire, and it keeps the connection alive - while lifecycle frames sit buffered for a possible fallback during a long thinking pass.""" - from litellm.llms.anthropic.messages.native.streaming_iterator import is_anthropic_ping_chunk - - return not has_generated_content and is_anthropic_ping_chunk(chunk) - - -def _is_retriable_anthropic_status(status_code: int) -> bool: - return status_code == 429 or status_code >= 500 - - -def _without_line_breaks(value: object) -> str: - return str(value).replace("\r", "").replace("\n", "") - - -def _anthropic_stream_error_is_gateway_verdict(chunk: object) -> bool: - """AgenticAnthropicStreamingIterator's own retrieval-failure frame is the gateway's verdict, not a provider - failure: another deployment would rerun the same failed hook, so it reaches the client instead of falling back.""" - from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import ( - is_server_fulfilled_tool_leak_error, - ) - - return is_server_fulfilled_tool_leak_error(chunk) - - -def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error: "MidStreamFallbackError") -> bool: - """ - A MidStreamFallbackError raised directly by the source iterator (the - completion-bridge path's CustomStreamWrapper, e.g. on a transport drop) - carries its own pre_first_chunk bookkeeping - gated the same way a - detected SSE error event is, so a fallback is never appended after real - content already reached the client on either path. - """ - return has_generated_content or not error.is_pre_first_chunk - - -def _stream_raised_error_status(error: Exception) -> int | None: - raw_status: Final = getattr(error, "status_code", None) - if isinstance(raw_status, int): - return raw_status - if isinstance(raw_status, str) and raw_status.isdigit(): - return int(raw_status) - response_status: Final = getattr(getattr(error, "response", None), "status_code", None) - return response_status if isinstance(response_status, int) else None - - -def _stream_fallback_error_for_raised( - error: Exception, model: str, has_generated_content: bool, llm_provider: str = "anthropic" -) -> "MidStreamFallbackError | None": - """The pre-stream retry rule (408, 409, 429, 5xx); None means the raise propagates unchanged.""" - if has_generated_content: - return None - status_code: Final = _stream_raised_error_status(error) - if status_code is None: - return _stream_pre_content_error(error, model, llm_provider) - retriable: Final = python_host.invoke("should_retry")(status_code) - return _stream_pre_content_error(error, model, llm_provider) if retriable else None - - -def _stream_pre_content_error( - error: Exception, model: str, llm_provider: str = "anthropic" -) -> "MidStreamFallbackError": - """The envelope the fallback chain judges a failure by when the client has received no content yet.""" - from litellm.exceptions import MidStreamFallbackError - - return MidStreamFallbackError( - message=str(error), - model=model, - llm_provider=llm_provider, - original_exception=error, - is_pre_first_chunk=True, - ) - - -def _context_window_stream_fallback_error( - error: "MidStreamFallbackError | litellm_exceptions.ContextWindowExceededError", - model: object, - llm_provider: str, - has_output: bool, -) -> "MidStreamFallbackError": - if isinstance(error, litellm_exceptions.ContextWindowExceededError): - if has_output: - raise error - return _stream_pre_content_error(error, model if isinstance(model, str) else "", llm_provider) - if isinstance(error.original_exception, litellm_exceptions.ContextWindowExceededError) and ( - has_output or error.generated_content or error.emitted_disqualifying_content - ): - raise error.original_exception - return error - - -def _deployment_num_retries(deployment: "Deployment | None") -> int | None: - """The deployment's own num_retries litellm_param, an int or a digit string the way the config loader leaves it.""" - configured: Final = getattr(deployment.litellm_params, "num_retries", None) if deployment is not None else None - if isinstance(configured, bool) or not isinstance(configured, (int, str)): - return None - return int(configured) if str(configured).isdigit() else None - - -def _request_fallback_list( - kwargs: Mapping[str, object], key: str, router_default: "list[object] | None" -) -> "list[object] | None": # mutable-ok: should_retry_this_error's own parameter type - return cast("list[object] | None", kwargs.get(key, router_default)) # cast-ok: same type as the router attribute - - -def _request_model_group(kwargs: Mapping[str, object]) -> str | None: - model_group: Final = kwargs.get("model") - return model_group if isinstance(model_group, str) else None - - -def _mid_stream_error_before_output(error: "MidStreamFallbackError") -> bool: - return not error.generated_content or error.is_pre_first_chunk - - -def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception: - """The provider's own error, which is what the retry policy and should_retry_this_error classify.""" - return error.original_exception if error.original_exception is not None else error - - -_RestT: Final = TypeVar("_RestT", Iterator[object], AsyncIterator[object]) -_RouteDeploymentsT: Final = TypeVar("_RouteDeploymentsT", bound=Sequence[Mapping[str, object]] | Mapping[str, object]) - - -@dataclass(frozen=True, slots=True) -class _FirstChunk(Generic[_RestT]): - """ - The chunk Router._acompletion / _completion already read off a stream (None: the - provider closed it empty) and the iterator it was read from, which carries on after it. - """ - - chunk: ModelResponseStream | None - rest: _RestT - - -async def _aread_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[AsyncIterator[object]]: - """ - Read the first chunk while Router._acompletion still holds the request, so a - stream that fails before any token raises there, inside async_function_with_retries - and the fallback chain, instead of surfacing to the caller mid-iteration. - - Iterates the way `async for` does, one __aiter__ call and then __anext__ on what it - returned: a CustomStreamWrapper subclass may hand iteration to another object and - leave completion_stream None (the MCP gateway's MCPStreamWrapper). - """ - from litellm.exceptions import MidStreamFallbackError - - rest: Final = stream.__aiter__() - try: - return _FirstChunk(chunk=await rest.__anext__(), rest=rest) - except StopAsyncIteration: - return _FirstChunk(chunk=None, rest=rest) - except MidStreamFallbackError as e: - raise _mid_stream_retry_trigger(e) from e - - -def _read_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[Iterator[object]]: - """Sync twin of _aread_first_chunk for Router._completion.""" - from litellm.exceptions import MidStreamFallbackError - - rest: Final = iter(stream) - try: - return _FirstChunk(chunk=next(rest), rest=rest) - except StopIteration: - return _FirstChunk(chunk=None, rest=rest) - except MidStreamFallbackError as e: - raise _mid_stream_retry_trigger(e) from e - - -def _chunks_with_first( - model_response: Iterable[object], taken: _FirstChunk[Iterator[object]] | None -) -> Iterator[object]: - """The stream with the chunk _read_first_chunk already read put back in front (None: the stream as is).""" - if taken is None: - yield from model_response - return - if taken.chunk is None: - return - yield taken.chunk - yield from taken.rest - - -async def _achunks_with_first( - model_response: AsyncIterable[object], taken: _FirstChunk[AsyncIterator[object]] | None -) -> AsyncIterator[object]: - """Async twin of _chunks_with_first.""" - if taken is None: - async for item in model_response: - yield item - return - if taken.chunk is None: - return - yield taken.chunk - async for item in taken.rest: - yield item - - -def _stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool: - """ - Whether `chunk` should make Router._aanthropic_messages_streaming_iterator - commit to the primary Anthropic stream (real content arrived, or the - pre-content buffer cap was hit) rather than keep buffering lifecycle - frames toward a possible fallback. - """ - from litellm.llms.anthropic.messages.native.streaming_iterator import ( - is_anthropic_content_delta_chunk, - ) - - if has_generated_content: - return False - return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS - - -@dataclass(frozen=True, slots=True) -class _StreamRetryOpened: - response: object - attempted_retries: int - max_retries: int - - -@dataclass(frozen=True, slots=True) -class _StreamRetriesExhausted: - error: "MidStreamFallbackError" - - -_StreamRetryOutcome: TypeAlias = _StreamRetryOpened | _StreamRetriesExhausted - - -MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200 - - -def _retry_policy_ceiling(policy: RetryPolicy) -> int: - """The most retries any error class under this policy can be granted.""" - retry_counts: Final[tuple[object, ...]] = tuple(policy.model_dump().values()) - return max((retries for retries in retry_counts if isinstance(retries, int)), default=0) - - -def _responses_stream_holds_event(item: object, held_event_count: int) -> bool: - from litellm.responses.streaming_iterator import PRE_OUTPUT_LIFECYCLE_EVENT_TYPES - - if held_event_count >= MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: - return False - return getattr(item, "type", None) in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES - - -def _responses_event_announces_background_run(item: object) -> bool: - return getattr(getattr(item, "response", None), "background", None) is True - - -class FallbackAwareAnthropicMessagesStream: - """ - Bare async generators can't carry the `_hidden_params` attribute the - proxy reads response headers off of (see - router_utils.add_retry_fallback_headers.get_hidden_params_dict), so this - thin wrapper carries it through from the source iterator - mirrors - AnthropicMessagesStreamingResponse. Used by - Router._aanthropic_messages_streaming_iterator. - """ - - def __init__(self, async_generator: AsyncGenerator[bytes, None], source_iterator: object) -> None: - self._async_generator = async_generator - self._source_iterator = source_iterator - self.fallback_headers_adopted = False - self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) - self._followed_source_params: object = None - - @property - def has_buffered_provider_output(self) -> bool: - return getattr(self._source_iterator, "has_buffered_provider_output", False) is True - - @property - def chunks(self) -> list[ModelResponseStream] | None: - return cast( # cast-ok: chunks is a list of ModelResponseStream on the inner stream - "list[ModelResponseStream] | None", getattr(self._source_iterator, "chunks", None) - ) - - @property - def messages(self) -> list[AllMessageValues] | None: - return cast( # cast-ok: messages is a list of AllMessageValues on the inner stream - "list[AllMessageValues] | None", getattr(self._source_iterator, "messages", None) - ) - - @property - def model(self) -> str | None: - return cast( # cast-ok: model is a str on the inner stream - "str | None", getattr(self._source_iterator, "model", None) - ) - - def adopt_fallback_source(self, fallback_response: object) -> None: - self._source_iterator = fallback_response - self.fallback_headers_adopted = True - - def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": - from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect - - return partial_responses_incomplete_event_on_disconnect(self._source_iterator) - - def follow_source_attribution(self) -> None: - """ - A retry's or a fallback's stream carries a wrapper of its own, so a hop it makes before its - first byte lands on that wrapper while the proxy reads the headers off this one. Mirrors the - source's current attribution onto this wrapper whenever the source has adopted a new one. - """ - source: Final = self._source_iterator - if not getattr(source, "fallback_headers_adopted", False): - return - source_params: Final = getattr(source, "_hidden_params", None) - if source_params is None or source_params is self._followed_source_params: - return - self._followed_source_params = source_params - hidden_params, headers = PythonCore._prepare_fallback_hidden_params(source) # pyright: ignore[reportPrivateUsage] # this wrapper is the Router's own stream type - self.merge_fallback_hidden_params(hidden_params, headers) - - def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream": - return self - - async def __anext__(self) -> bytes: - return await self._async_generator.__anext__() - - async def aclose(self) -> None: - await self._async_generator.aclose() - - def merge_fallback_hidden_params( - self, - fallback_hidden_params: Mapping[str, object], - fallback_headers: Mapping[str, object], - ) -> None: - """ - Raw bytes can't carry their own _hidden_params the way a - ModelResponseStream/ResponsesAPI event can, so a mid-stream - fallback's provider headers (e.g. Bedrock's x-amzn-requestid) are - merged onto the wrapper itself instead - mirrors - Router._apply_fallback_hidden_params_to_item's merge shape. - """ - existing_headers: Final = cast( # cast-ok: additional_headers is always a dict[str, object] when present - "dict[str, object]", self._hidden_params.get("additional_headers") or {} - ) - self._hidden_params = { - **self._hidden_params, - **fallback_hidden_params, - "additional_headers": dict(replace_complexity_router_headers(existing_headers, fallback_headers)), - } - - -class RoutingArgs(enum.Enum): - ttl = 60 # 1min (RPM/TPM expire key) - - -# Routers that are still in use, so a price data reload can rebuild the cost-map -# entries their deployments own. Weak so a router nothing references any more, such -# as the per-request one built from a caller-supplied user_config, drops out on its -# own rather than leaving entries behind that nothing can withdraw. -_live_routers: Final["weakref.WeakSet[PythonCore]"] = weakref.WeakSet() -# Shared with the other router cores of the process (litellm/router_utils/global_callback_holders.py). -_routers_holding_global_callbacks: Final["weakref.WeakSet[PythonCore]"] = routers_holding_global_callbacks # pyright: ignore[reportAssignmentType] # every member serves the holder members -_routers_holding_global_callbacks_lock: Final = routers_holding_global_callbacks_lock - - -def _replay_live_router_model_cost() -> None: - """Re-assert every live router's deployments after the cost map is refreshed.""" - for router in tuple(_live_routers): - router.replay_model_cost_registrations() - - -set_live_deployment_replay(_replay_live_router_model_cost) - - -RETRY_BREADCRUMB_LIMIT: Final = 4 - - -class FallbackAwareStreamWrapper(CustomStreamWrapper): - """Base for the Router's chat-completion stream wrappers, which are built around the - attempt the Router picked first and have to repoint themselves when a fallback takes over.""" - - fallback_headers_adopted: bool = False - - def adopt_fallback_response_headers( - self, - fallback_response: object, - prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]], - ) -> None: - """Repoint this wrapper at the deployment that served the stream. - - Replaces rather than merges, so the failed attempt's `x-request-id`, rate limit - counters, `model_id` and `api_base` cannot reach the proxy's response headers or - its callbacks. - """ - self._response_headers = getattr(fallback_response, "_response_headers", None) - fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params - if fallback_hidden_params: - self._hidden_params = { - **fallback_hidden_params, - # dict() because add_retry_fallback_headers mutates additional_headers in place - "additional_headers": dict(fallback_headers), - } - self._base_hidden_params = { - **self._hidden_params, - "response_cost": None, - } - self.fallback_headers_adopted = True - - -def as_output_cap(value: object) -> int | None: - """A client-sent output cap coerced to an int: ints, floats and numeric strings, never bools - or negatives.""" - if isinstance(value, bool) or not isinstance(value, (int, float, str)): - return None - try: - cap: Final = int(float(value)) - except (ValueError, OverflowError): - return None - return cap if cap >= 0 else None - - -def stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None: - """None means unstamped; an empty mapping means the original caller supplied no limit.""" - stamped: Final = next( - ( - bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) - for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) - if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket - ), - None, - ) - if not isinstance(stamped, dict): - return None - return MappingProxyType( - { - carrier: cap - for carrier, value in stamped.items() - if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None - } - ) - - -def _exact_base_model_max_input_tokens(base_model: str) -> int | None: - try: - model_info: Final = python_host.invoke("get_model_info")(model=base_model) - except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model - return None - if is_generalized_model_info(model_info): - return None - return model_info.get("max_input_tokens") - - -def _provider_model_id(model_name: str, custom_llm_provider: str) -> str: - model_id: Final = model_name.rsplit("/", 1)[-1] - return get_bedrock_base_model(model_id) if custom_llm_provider.startswith("bedrock") else model_id - - -def _names_route_model(base_model: str, route_key: str, custom_llm_provider: str) -> bool: - route_model_id: Final = _provider_model_id(route_key, custom_llm_provider) - base_model_id: Final = _provider_model_id(base_model, custom_llm_provider) - return route_model_id in (base_model_id, base_model_id.split("@", 1)[0]) - - -def _declared_base_model(user_model_info: Mapping[str, object]) -> str | None: - match user_model_info: - case {"base_model": str() as base_model}: - return base_model - case _: - return None - - -def _base_model_widened_window( - route_model_info: ModelMapInfo, base_model: str | None, custom_llm_provider: str -) -> Mapping[str, int]: - route_max_input_tokens: Final = route_model_info.get("max_input_tokens") - if base_model is None or route_max_input_tokens is None: - return MappingProxyType({}) - if not _names_route_model(base_model, route_model_info["key"], custom_llm_provider): - return MappingProxyType({}) - base_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model) - if base_max_input_tokens is None or base_max_input_tokens <= route_max_input_tokens: - return MappingProxyType({}) - return MappingProxyType({"max_input_tokens": base_max_input_tokens}) - - -_SETTING_ATTRIBUTES: Final = frozenset("_core_" + name for name in SETTING_FIELDS) - - -class PythonCore(DeprecatedRouterMembers): - _core_model_names: set[str] = set() - _core_cache_responses: bool | None = False - _registry_version: int = 0 - # Bumped by every ``apply``: the static arguments plus ``_core_applied`` are the settings. - _settings_version: int = 0 - _core_applied: Mapping[str, object] = MappingProxyType({}) - - # -- entry points: everything outside code needs besides the endpoint methods --------------- - # See litellm/router_api.py. The members outside code used before these live in - # DeprecatedRouterMembers (litellm/router_deprecated.py). - - def registry(self) -> RouterRegistry: - """Deployment lookups: by id, group, alias, wildcard and team, plus configured plugins.""" - cached: Final = self.__dict__.get("_registry_view") - if isinstance(cached, RouterRegistry): - return cached - view: Final = RouterRegistry(self) - self.__dict__["_registry_view"] = view - return view - - def settings(self) -> RouterSettings: - """ - A frozen snapshot of the router's configuration. It is built once and reused until a - setting attribute is rebound or deleted (``__setattr__``/``__delattr__`` drop it), so the - request path can read it cheaply. - """ - cached: Final = self.__dict__.get("_settings_snapshot") - if isinstance(cached, RouterSettings): - return cached - snapshot: Final = RouterSettings.of(self) - self.__dict__["_settings_snapshot"] = snapshot - return snapshot - - def __setattr__(self, name: str, value: object) -> None: - if ( - (name in PLAIN_SETTING_ATTRIBUTES or name in _SETTING_ATTRIBUTES) - and self.__dict__.get("_core_ready") - and not self.__dict__.get("_core_applying") - ): - # Settings change only through a changeset, so every core serving this router sees it. - self.apply(ChangeSet(assign={name.removeprefix("_core_"): value})) - return - if name in _SETTING_ATTRIBUTES: - self.__dict__.pop("_settings_snapshot", None) - elif name == "_core_model_list" and "_core_model_list" in self.__dict__: - # A whole new deployment list (``router.model_list = ...``) is a deployment change too. - self.__dict__["_registry_version"] = self._registry_version + 1 - super().__setattr__(name, value) - - def __delattr__(self, name: str) -> None: - if name in _SETTING_ATTRIBUTES: - self.__dict__.pop("_settings_snapshot", None) - super().__delattr__(name) - - def health(self) -> RouterHealth: - """Cooldowns, blocked deployments, health checks and usage.""" - cached: Final = self.__dict__.get("_health_view") - if isinstance(cached, RouterHealth): - return cached - view: Final = RouterHealth(self) - self.__dict__["_health_view"] = view - return view - - @overload - async def select( - self, - model: str, - request_kwargs: dict[str, object] | None = None, # mutable-ok: the selection cores take a dict - *, - kind: Literal["pick", "pass_through"] = "pick", - messages: list[dict[str, str]] | None = None, # mutable-ok: the selection cores take a list - input: str | list[object] | None = None, # mutable-ok: the selection cores take a list - specific_deployment: bool | None = False, - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ) -> "UntypedDict": ... # a mutable dict, as before: the cores' untyped deployment dict, unchanged - - @overload - async def select( - self, - model: str, - request_kwargs: dict[str, object] | None = None, # mutable-ok: the selection cores take a dict - *, - kind: Literal["healthy"], - messages: list[dict[str, str]] | None = None, # mutable-ok: the selection cores take a list - input: str | list[object] | None = None, # mutable-ok: the selection cores take a list - specific_deployment: bool | None = False, - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ) -> "list[UntypedDict] | UntypedDict": ... # mutable-ok: the core's type (list[dict] | dict), unchanged - - async def select( # pyright: ignore[reportUnknownParameterType] # the selection cores return untyped deployment dicts - self, - model: str, - request_kwargs: dict[str, object] | None = None, # mutable-ok: the selection cores take a dict - *, - kind: SelectKind = "pick", - messages: list[dict[str, str]] | None = None, # mutable-ok: the selection cores take a list - input: str | list[object] | None = None, # mutable-ok: the selection cores take a list - specific_deployment: bool | None = False, - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ): - """ - Pick a deployment for ``model`` without calling it (``pick``), list the deployments that - are healthy for the request (``healthy``), or pick one for a pass-through route - (``pass_through``). - """ - if kind == "pick": - with picks_reserve(False): - return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - request_kwargs=request_kwargs if request_kwargs is not None else {}, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - if kind == "healthy": - return await self._core_async_get_healthy_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - request_kwargs=request_kwargs if request_kwargs is not None else {}, - messages=messages, - input=input, - specific_deployment=specific_deployment, - parent_otel_span=parent_otel_span, - health_check_probe=health_check_probe, - ) - if kind == "pass_through": - with picks_reserve(False): - return self._core_get_available_deployment_for_pass_through( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - raise ValueError(f"unknown select kind {kind!r}") - - def apply(self, changes: ChangeSet) -> ApplyResult: - """Apply runtime changes: settings, deployments, the routing state's Redis cache.""" - with self._core_applying_changes(changes): - return apply_changes(self, changes) - - async def aapply(self, changes: ChangeSet) -> ApplyResult: - """``apply``, plus the changes that need I/O (``refresh_model_info``).""" - with self._core_applying_changes(changes): - return await aapply_changes(self, changes) - - @contextlib.contextmanager - def _core_applying_changes(self, changes: ChangeSet) -> Generator[None, None, None]: - """Mark a changeset in progress, then record its settings and bump ``_settings_version``.""" - depth: Final = self.__dict__.get("_core_applying", 0) - self.__dict__["_core_applying"] = depth + 1 - try: - yield - finally: - self.__dict__["_core_applying"] = depth - if changes.assign or changes.settings: - self.__dict__["_core_applied"] = MappingProxyType( - {**self._core_applied, **changes.assign, **changes.settings} - ) - self.__dict__["_settings_version"] = self._settings_version + 1 - - def observe(self, logging_obj: "LiteLLMLogging | None" = None, *, request_scoped: bool = False) -> RoutingObservers: - """ - The router's routing observers. ``request_scoped=True`` stops registering them globally, - so each request binds them instead; passing a request's ``logging_obj`` binds them to it. - """ - if request_scoped: - self._core_serve_through_request_logging() - if logging_obj is not None: - self._core_bind_to_request_logging(logging_obj) - return RoutingObservers(slack_alerting_logger=getattr(self, "_core_slack_alerting_logger", None)) - - def register_plugin(self, plugin: CustomRoutingStrategyBase) -> None: - """Route with a custom strategy instead of the configured one.""" - self._core_set_custom_routing_strategy(plugin) - - def retire(self) -> None: - """Release the router's global callback registrations and routing state.""" - self._core_discard() - - def _core_public(self) -> "PythonCore": - """The ``litellm.Router`` this core serves, for objects the router hands out; the core itself if none.""" - owner: Final = self.__dict__.get("_core_owner") - public: Final = owner() if owner is not None else None - return self if public is None else public - - async def call(self, op: str, request: Mapping[str, object]) -> object: - """Run one endpoint (``acompletion``, ``aembedding``, ...) with the request's keyword payload.""" - result: Final = getattr(self, op)(**request) - return await result if inspect.isawaitable(result) else result - - default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour - tenacity = None - leastbusy_logger: LeastBusyLoggingHandler | None = None - lowesttpm_logger: LowestTPMLoggingHandler | None = None - optional_callbacks: list[CustomLogger | Callable | str] | None = None - - def __init__( - self, - model_list: list[DeploymentTypedDict] | list[dict[str, Any]] | None = None, - ## ASSISTANTS API ## - assistants_config: AssistantsTypedDict | None = None, - ## SEARCH API ## - search_tools: list[SearchToolTypedDict] | None = None, - ## GUARDRAIL API ## - guardrail_list: list[GuardrailTypedDict] | None = None, - ## CACHING ## - redis_url: str | None = None, - redis_host: str | None = None, - redis_port: int | None = None, - redis_password: str | None = None, - redis_db: int | None = None, - cache_responses: bool | None = False, - cache_kwargs: dict = {}, # additional kwargs to pass to RedisCache (see caching.py) - caching_groups: list[tuple] | None = None, # if you want to cache across model groups - client_ttl: int = 3600, # ttl for cached clients - will re-initialize after this time in seconds - ## SCHEDULER ## - polling_interval: float | None = None, - default_priority: int | None = None, - ## RELIABILITY ## - num_retries: int | None = None, - max_fallbacks: int | None = None, # max fallbacks to try before exiting the call. Defaults to 5. - timeout: float | None = None, - stream_timeout: float | None = None, - default_litellm_params: dict | None = None, # default params for Router.chat.completion.create - default_max_parallel_requests: int | None = None, - set_verbose: bool = False, - debug_level: Literal["DEBUG", "INFO"] = "INFO", - default_fallbacks: list[str] | None = None, # generic fallbacks, works across all deployments - fallbacks: list = [], - context_window_fallbacks: list = [], - content_policy_fallbacks: list = [], - model_group_alias: dict[str, str | RouterModelGroupAliasItem] | None = {}, - enable_pre_call_checks: bool = False, - enable_tag_filtering: bool = False, - tag_filtering_match_any: bool = True, - tag_routing_prefix: str = "", - plugins: list[RoutingPlugin] | None = None, - retry_after: int = 0, # min time to wait before retrying a failed request - retry_policy: RetryPolicy | dict | None = None, # set custom retries for different exceptions - model_group_retry_policy: dict[str, RetryPolicy] = {}, # set custom retry policies based on model group - allowed_fails: int | None = None, # Number of times a deployment can failbefore being added to cooldown - allowed_fails_policy: AllowedFailsPolicy | None = None, # set custom allowed fails policy - cooldown_time: float | None = None, # (seconds) time to cooldown a deployment after failure - disable_cooldowns: bool | None = None, - fallback_on_unsupported_params: bool = True, - routing_strategy: RoutingStrategyName = "simple-shuffle", - optional_pre_call_checks: OptionalPreCallChecks | None = None, - routing_strategy_args: dict = {}, # just for latency-based - routing_groups: list[RoutingGroup | dict] | None = None, - provider_budget_config: GenericBudgetConfigType | None = None, - alerting_config: AlertingConfig | None = None, - router_general_settings: RouterGeneralSettings | None = RouterGeneralSettings(), - deployment_affinity_ttl_seconds: int = 3600, - model_group_affinity_config: dict[str, list[str]] | None = None, - ignore_invalid_deployments: bool = False, - enable_health_check_routing: bool = False, - health_check_staleness_threshold: int | None = None, - health_check_ignore_transient_errors: bool = False, - background_health_check_model_groups: Sequence[str] | None = None, - enable_weighted_failover: bool = False, - enable_mid_stream_fallback_continuation: bool = False, - fallback_access_check: FallbackAccessCheck | None = None, - fallback_budget_check: FallbackBudgetCheck | None = None, - auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, - session_affinity_repin_on_failover: bool = False, - _host: PythonHost = python_host, - ) -> None: - """ - Initialize the Router class with the given parameters for caching, reliability, and routing strategy. - - Args: - model_list (Optional[list]): List of models to be used. Defaults to None. - redis_url (Optional[str]): URL of the Redis server. Defaults to None. - redis_host (Optional[str]): Hostname of the Redis server. Defaults to None. - redis_port (Optional[int]): Port of the Redis server. Defaults to None. - redis_password (Optional[str]): Password of the Redis server. Defaults to None. - cache_responses (Optional[bool]): Flag to enable caching of responses. Defaults to False. - cache_kwargs (dict): Additional kwargs to pass to RedisCache. Defaults to {}. - caching_groups (Optional[List[tuple]]): List of model groups for caching across model groups. Defaults to None. - client_ttl (int): Time-to-live for cached clients in seconds. Defaults to 3600. - polling_interval: (Optional[float]): frequency of polling queue. Only for '.scheduler_acompletion()'. Default is 3ms. - default_priority: (Optional[int]): the priority given to every chat or text completion request that sets no 'priority' of its own; '.acompletion()' and '.atext_completion()' queue it through the scheduler. Default is None. - num_retries (Optional[int]): Number of retries for failed requests. Defaults to 2. - timeout (Optional[float]): Timeout for requests. Defaults to None. - default_litellm_params (dict): Default parameters for Router.chat.completion.create. Defaults to {}. - set_verbose (bool): Flag to set verbose mode. Defaults to False. - debug_level (Literal["DEBUG", "INFO"]): Debug level for logging. Defaults to "INFO". - fallbacks (List): List of fallback options. Defaults to []. - context_window_fallbacks (List): List of context window fallback options. Defaults to []. - enable_pre_call_checks (boolean): Filter out deployments which are outside context window limits for a given prompt - model_group_alias (Optional[dict]): Alias for model groups. Defaults to {}. - retry_after (int): Minimum time to wait before retrying a failed request. Defaults to 0. - allowed_fails (Optional[int]): Number of allowed fails before adding to cooldown. Defaults to None. - cooldown_time (float): Time to cooldown a deployment after failure in seconds. Defaults to 1. - routing_strategy (Literal["simple-shuffle", "least-busy", "usage-based-routing", "latency-based-routing", "cost-based-routing"]): Routing strategy used for the implicit "default" group (any model not claimed by an entry in `routing_groups`). Defaults to "simple-shuffle". - routing_strategy_args (dict): Additional args for the default group's routing strategy (e.g. latency window). Defaults to {}. - routing_groups (Optional[List[RoutingGroup]]): Named subsets of `model_name`s with a group routing strategy. - Priority groups apply only to group calls and may overlap. Other groups supply their members' default - strategy, with at most one such group per model. Unclaimed models use the top-level strategy. - alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None. - provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None. - deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. - session_affinity_repin_on_failover (bool): After a successful replacement response, move a session pin - whose deployment is in cooldown or removed. Async routes only. Defaults to False. Applied at startup. - ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. - enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. - enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False. - fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). - fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback). - Returns: - Router: An instance of the litellm.Router class. - - Example Usage: - ```python - from litellm import Router - model_list = [ - { - "model_name": "azure-gpt-3.5-turbo", # model alias - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/", - "api_key": , - "api_version": , - "api_base": - }, - }, - { - "model_name": "azure-gpt-3.5-turbo", # model alias - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/", - "api_key": , - "api_version": , - "api_base": - }, - }, - { - "model_name": "openai-gpt-3.5-turbo", # model alias - "litellm_params": { # params for litellm completion/embedding call - "model": "gpt-3.5-turbo", - "api_key": , - }, - ] - - router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}]) - ``` - """ - self._host: Final = _host - self._core_ctx: Final = RouterContext(self) - - self.set_verbose = set_verbose - self.ignore_invalid_deployments = ignore_invalid_deployments - self.auto_router_capability_limit = auto_router_capability_limit - self._core_fallback_access_check: Final = fallback_access_check - self._core_fallback_budget_check: Final = fallback_budget_check - self.debug_level = debug_level - self._core_enable_pre_call_checks = enable_pre_call_checks - self._core_enable_tag_filtering = enable_tag_filtering - self._core_tag_filtering_match_any = tag_filtering_match_any - self._core_tag_routing_prefix = tag_routing_prefix - from litellm._service_logger import ServiceLogging - - self.service_logger_obj: ServiceLogging = ServiceLogging() - self._host.suppress_debug_info() # prevents 'Give Feedback/Get help' message from being emitted on Router - Relevant Issue: https://github.com/BerriAI/litellm/issues/5942 - if self.set_verbose is True: - if debug_level == "INFO": - verbose_router_logger.setLevel(logging.INFO) - elif debug_level == "DEBUG": - verbose_router_logger.setLevel(logging.DEBUG) - self._core_router_general_settings: RouterGeneralSettings = router_general_settings or RouterGeneralSettings() - - self.assistants_config = assistants_config - self._core_search_tools = search_tools or [] - self._core_guardrail_list = guardrail_list or [] - self._core_deployment_names: list[str] = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 - self._core_deployment_latency_map = {} - ### CACHING ### - cache_type: Literal["local", "redis", "redis-semantic", "s3"] = "local" # default to an in-memory cache - redis_cache = None - cache_config: Final[dict[str, Any]] = {} - - self.client_ttl = client_ttl - self._core_kubernetes_pod_discovery: Final = KubernetesPodDiscovery() - if redis_url is not None or (redis_host is not None and redis_port is not None): - cache_type = "redis" - - if redis_url is not None: - cache_config["url"] = redis_url - - if redis_host is not None: - cache_config["host"] = redis_host - - if redis_port is not None: - cache_config["port"] = str(redis_port) - - if redis_password is not None: - cache_config["password"] = redis_password - - if redis_db is not None: - verbose_router_logger.warning( - "Deprecated 'redis_db' argument used. Please remove 'redis_db' from your config/database and use 'cache_kwargs' instead." - ) - cache_config["db"] = str(redis_db) - - # Add additional key-value pairs from cache_kwargs - cache_config.update(cache_kwargs) - redis_cache = self._create_redis_cache(cache_config) - - if cache_responses: - if self._host.global_cache() is None: - # the cache can be initialized on the proxy server. We should not overwrite it - self._host.set_global_cache(LitellmCache(type=cache_type, **cache_config)) - self._core_cache_responses = cache_responses - self._core_cache = DualCache( - redis_cache=redis_cache, in_memory_cache=InMemoryCache() - ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. - self._core_least_busy_reservations: InMemoryCache = InMemoryCache( - max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS - ) - self._claude_code_session_router_cache: DualCache = DualCache( - redis_cache=redis_cache, - in_memory_cache=InMemoryCache(), - ) - self._core_prompt_caching_pin_cache: DualCache = DualCache( - redis_cache=redis_cache, - in_memory_cache=InMemoryCache(max_size_in_memory=PROMPT_CACHE_PIN_CACHE_MAX_SIZE), - ) - - ### SCHEDULER ### - self.scheduler = Scheduler(polling_interval=polling_interval, redis_cache=redis_cache) - parsed_default_priority: Final = parse_default_priority(default_priority) - if isinstance(parsed_default_priority, InvalidPriority): - verbose_router_logger.warning( - "router_settings.default_priority is ignored, no request is queued by default: %s", - parsed_default_priority.message, - ) - self.default_priority = ( - None if isinstance(parsed_default_priority, InvalidPriority) else parsed_default_priority - ) - self._core_default_deployment = ( - None # use this to track the users default deployment, when they want to use model = * - ) - self._core_default_max_parallel_requests = default_max_parallel_requests - self._core_provider_default_deployment_ids: list[str] = [] - self._core_pattern_router = PatternMatchRouter() - self._core_team_pattern_routers: dict[str, PatternMatchRouter] = {} # {"TEAM_ID": PatternMatchRouter} - self._core_auto_routers: dict[str, list[TaggedPreRoutingStrategy[AutoRouter]]] = {} - self._core_complexity_routers: dict[str, list[TaggedPreRoutingStrategy[ComplexityRouter]]] = {} - self._core_adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] = {} - self._core_quality_routers: dict[str, list[TaggedPreRoutingStrategy[QualityRouter]]] = {} - self._core_routing_plugins: list[RoutingPlugin] = list(plugins) if plugins else [] - - # Initialize model_group_alias early since it's used in set_model_list - self._core_model_group_alias: dict[str, str | RouterModelGroupAliasItem] = ( - model_group_alias or {} - ) # dict to store aliases for router, ex. {"gpt-4": "gpt-3.5-turbo"}, all requests with gpt-4 -> get routed to gpt-3.5-turbo group - - # Initialize model ID to deployment index mapping for O(1) lookups - self.model_id_to_deployment_index_map: dict[str, int] = {} - # Initialize model name to deployment indices mapping for O(1) lookups - # Maps model_name -> list of indices in model_list - self._core_model_name_to_deployment_indices: dict[str, list[int]] = {} - # Maps (team_id, team_public_model_name) -> list of indices in model_list - self._core_team_model_to_deployment_indices: dict[tuple[str, str], list[int]] = {} - self._core_team_public_model_names: frozenset[str] = frozenset() - # Guards model_list and the three index maps above: writers change them, and readers turn an - # index into a deployment, under it, so a lookup never lands between a removal and its re-index - self._deployment_index_lock = threading.RLock() - - # Initialize cache attributes that ``_invalidate_model_group_info_cache`` - # and ``_invalidate_access_groups_cache`` touch *before* the first - # ``set_model_list`` below (which calls those invalidations as part of - # building the model index) and before ``_init_routing_groups(None)`` - # (which calls them on every group rebuild). - self._access_groups_cache: dict[str, list[str]] | None = None - # Per-router cache for the proxy auth-layer "is this model explicitly - # zero-cost?" check. Lives on the router so it is invalidated alongside - # ``_cached_get_model_group_info`` and dies with the router (no - # ``id()``-reuse risk after GC). See - # ``litellm.proxy.auth.auth_checks._is_model_cost_zero``. - self._zero_cost_cache: dict[str, bool] = {} - self._core_cached_deployment_model_info = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)( - self._core_get_deployment_model_info - ) - self._ptu_share_teams = lru_cache(maxsize=1)(self._collect_ptu_share_teams) - self._cached_team_ptu_ceiling = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(self._resolve_team_ptu_ceiling) - self._cached_deployment_tracks_usage = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)( - self._deployment_tracks_usage - ) - self.served_openai_models = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(self._served_openai_models) - self._discovered_model_info_cache: InMemoryCache = InMemoryCache( - max_size_in_memory=max(len(model_list or ()), 1), - default_ttl=2 * MODEL_INFO_REFRESH_SECONDS, - ) - self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None - self._init_routing_groups(None) - self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = () - - self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds - self.session_affinity_repin_on_failover = session_affinity_repin_on_failover - self.model_group_affinity_config = model_group_affinity_config - warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) - - if model_list is not None: - # set_model_list will build indices automatically - self.set_model_list(model_list) - self.healthy_deployments: list = self._core_model_list - for m in model_list: - if "model" in m["litellm_params"]: - self._core_deployment_latency_map[m["litellm_params"]["model"]] = 0 - else: - self._core_model_list: list = [] # initialize an empty list - to allow _add_deployment and delete_deployment to work - # Track this router so a price data reload can rebuild its deployments' - # cost-map entries from the list it is serving at that moment, whether - # those deployments came from model_list or were added later from the DB. - _live_routers.add(self) - live_check_owners.add(self) - - if allowed_fails is not None: - self._core_allowed_fails = allowed_fails - else: - self._core_allowed_fails = self._host.dynamic().allowed_fails - self.explicit_allowed_fails = allowed_fails - self._core_cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS - self._core_cooldown_cache = CooldownCache( - cache=self._core_cache, default_cooldown_time=self._core_cooldown_time - ) - self._core_disable_cooldowns = disable_cooldowns - self.fallback_on_unsupported_params = fallback_on_unsupported_params - self._core_enable_health_check_routing = enable_health_check_routing - self.enable_weighted_failover = enable_weighted_failover - self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enable_mid_stream_fallback_continuation) - self._core_health_check_ignore_transient_errors = health_check_ignore_transient_errors - self._core_background_health_check_model_groups: frozenset[str] | None = ( - frozenset(background_health_check_model_groups) - if background_health_check_model_groups is not None - else None - ) - _staleness: Final = health_check_staleness_threshold or ( - DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER - ) - self._core_health_state_cache = DeploymentHealthCache( - cache=self._core_cache, staleness_threshold=float(_staleness) - ) - - if num_retries is not None: - self.num_retries = num_retries - elif self._host.dynamic().num_retries is not None: - self.num_retries = self._host.dynamic().num_retries - else: - self.num_retries = openai.DEFAULT_MAX_RETRIES - - if max_fallbacks is not None: - self.max_fallbacks = max_fallbacks - elif self._host.dynamic().max_fallbacks is not None: - self.max_fallbacks = self._host.dynamic().max_fallbacks - else: - self.max_fallbacks = self._host.dynamic().router_max_fallbacks - - self._explicit_timeout = timeout # None when user did not pass timeout - self.timeout = timeout or self._host.dynamic().request_timeout - # Per-attempt request_timeout, independent of router_settings.timeout. - # Only stored when a router timeout is also set, since otherwise - # request_timeout already flows through self.timeout above. - self.request_timeout = get_configured_request_timeout() if timeout is not None else None - self.stream_timeout = stream_timeout - - self.retry_after = retry_after - self._core_routing_strategy = self._normalize_strategy(routing_strategy) - self._routing_groups_input: list[RoutingGroup | dict] | None = routing_groups - - ## SETTING FALLBACKS ## - ### validate if it's set + in correct format - _fallbacks = fallbacks or self._host.dynamic().fallbacks - - self.validate_fallbacks(fallback_param=_fallbacks) - ### set fallbacks - self._core_fallbacks = _fallbacks - - if default_fallbacks is not None or self._host.dynamic().default_fallbacks is not None: - _fallbacks = default_fallbacks or self._host.dynamic().default_fallbacks - if isinstance(_fallbacks, str): - raise TypeError( - f"default_fallbacks must be a list of model names, got '{_fallbacks}'. Use ['{_fallbacks}']" - ) - if self._core_fallbacks is not None: - self._core_fallbacks.append({"*": _fallbacks}) - else: - self._core_fallbacks = [{"*": _fallbacks}] - - self._core_context_window_fallbacks = context_window_fallbacks or self._host.dynamic().context_window_fallbacks - - _content_policy_fallbacks: Final = content_policy_fallbacks or self._host.dynamic().content_policy_fallbacks - self.validate_fallbacks(fallback_param=_content_policy_fallbacks) - self._core_content_policy_fallbacks = _content_policy_fallbacks - self.total_calls: defaultdict = defaultdict(int) # dict to store total calls made to each model - self.fail_calls: defaultdict = defaultdict(int) # dict to store fail_calls made to each model - self.success_calls: defaultdict = defaultdict(int) # dict to store success_calls made to each model - - # make Router.chat.completions.create compatible for openai.chat.completions.create - default_litellm_params = default_litellm_params or {} - self.chat = LitellmChat(params=default_litellm_params, router_obj=self._core_public()) - - # default litellm args - self._core_default_litellm_params = default_litellm_params - self._core_default_litellm_params.setdefault("timeout", timeout) - self._core_default_litellm_params.setdefault("max_retries", 0) - self._core_default_litellm_params.setdefault("metadata", {}).update({"caching_groups": caching_groups}) - - self.deployment_stats: dict = {} # used for debugging load balancing - """ - deployment_stats = { - "122999-2828282-277: - { - "model": "gpt-3", - "api_base": "http://localhost:4000", - "num_requests": 20, - "avg_latency": 0.001, - "num_failures": 0, - "num_successes": 20 - } - } - """ - - ### ROUTING SETUP ### - if self._normalize_strategy(routing_strategy) == "lar1": - from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - - apply_lar1_routing_strategy(self._core_ctx, routing_strategy_args) - else: - self.routing_strategy_init( - routing_strategy=routing_strategy, - routing_strategy_args=routing_strategy_args, - ) - self._init_routing_groups(self._routing_groups_input) - self._override_selectors: dict[str, RouterStrategySelector | None] = {} - self._override_selectors_lock = threading.Lock() - self._request_scoped_selectors: tuple[CustomLogger, ...] = () - self._request_scoped_policies: tuple[CustomLogger, ...] = () - self.access_groups = None - with _routers_holding_global_callbacks_lock: - _routers_holding_global_callbacks.add(self) - ## USAGE TRACKING ## - self._host.callbacks.add_litellm_async_success_callback(self.deployment_callback_on_success) - if isinstance(self._host.callback_list("success_callback"), list): - self._host.callbacks.add_litellm_success_callback(self.sync_deployment_callback_on_success) - else: - self._host.set_callback_list("success_callback", [self.sync_deployment_callback_on_success]) - self._host.callbacks.add_litellm_async_failure_callback(self.async_deployment_callback_on_failure) - ## COOLDOWNS ## - if isinstance(self._host.callback_list("failure_callback"), list): - self._host.callbacks.add_litellm_failure_callback(self.deployment_callback_on_failure) - else: - self._host.set_callback_list("failure_callback", [self.deployment_callback_on_failure]) - self.routing_strategy_args = routing_strategy_args - self._core_provider_budget_config = provider_budget_config - self.router_budget_logger: RouterBudgetLimiting | None = None - if RouterBudgetLimiting.should_init_router_budget_limiter( - model_list=model_list, provider_budget_config=self._core_provider_budget_config - ): - if optional_pre_call_checks is not None: - optional_pre_call_checks.append("router_budget_limiting") - else: - optional_pre_call_checks = ["router_budget_limiting"] - self.retry_policy: RetryPolicy | None = None - if retry_policy is not None: - if isinstance(retry_policy, dict): - self.retry_policy = RetryPolicy(**retry_policy) - elif isinstance(retry_policy, RetryPolicy): - self.retry_policy = retry_policy - if self.retry_policy is not None: - verbose_router_logger.info( - "\x1b[32mRouter Custom Retry Policy Set:\n%s\x1b[0m", - self.retry_policy.model_dump(exclude_none=True), - ) - - self.model_group_retry_policy: dict[str, RetryPolicy] | None = model_group_retry_policy - - self._core_allowed_fails_policy: AllowedFailsPolicy | None = None - if allowed_fails_policy is not None: - if isinstance(allowed_fails_policy, dict): - self._core_allowed_fails_policy = AllowedFailsPolicy(**allowed_fails_policy) - elif isinstance(allowed_fails_policy, AllowedFailsPolicy): - self._core_allowed_fails_policy = allowed_fails_policy - - if self._core_allowed_fails_policy is not None: - verbose_router_logger.info( - "\x1b[32mRouter Custom Allowed Fails Policy Set:\n%s\x1b[0m", - self._core_allowed_fails_policy.model_dump(exclude_none=True), - ) - - self.alerting_config: AlertingConfig | None = alerting_config - - if optional_pre_call_checks is not None: - self.add_optional_pre_call_checks(optional_pre_call_checks) - - # If model_group_affinity_config is set but no global affinity checks were - # enabled, we still need the DeploymentAffinityCheck callback (with global - # flags all False) so per-group config can activate affinity per model group. - if self.model_group_affinity_config: - self._ensure_deployment_affinity_callback() - - if self.alerting_config is not None: - self._initialize_alerting() - - self.initialize_assistants_endpoint() - self.initialize_router_endpoints() - self.apply_default_settings() - # From here on a settings write is a changeset (``__setattr__``). - self.__dict__["_core_ready"] = True - - @staticmethod - def get_valid_args() -> list[str]: - """ - Returns a list of valid arguments for the Router.__init__ method. - """ - arg_spec: Final = inspect.getfullargspec(PythonCore.__init__) - return [arg for arg in (*arg_spec.args, *arg_spec.kwonlyargs) if arg not in ("self", "_host")] - - def apply_default_settings(self): - """ - Apply the default settings to the router. - """ - - default_pre_call_checks: Final[OptionalPreCallChecks] = [] - self.add_optional_pre_call_checks(default_pre_call_checks) - - if self.enable_mid_stream_fallback_continuation: - self._register_continuation_prefill_check() - - def _register_continuation_prefill_check(self) -> None: - # Process-global on purpose: Router._core_discard() must not drop a filter another router still needs. - self._host.callbacks.add_litellm_callback(ContinuationPrefillDeploymentCheck()) - - def _core_discard(self) -> None: - """ - Pseudo-destructor to be invoked to clean up global data structures when router is no longer used. - For now, unhook router's callbacks from all lists - """ - # Stop contributing to cost-map rebuilds straight away rather than waiting - # for this router to be collected. - _live_routers.discard(self) - live_check_owners.discard(self) - self.retire_routing_state() - self._release_global_callbacks() - - def _release_global_callbacks(self) -> None: - with _routers_holding_global_callbacks_lock: - _routers_holding_global_callbacks.discard(self) - slack_alerting_loggers: Final = self._slack_alerting_loggers() - adaptive_router_hooks: Final = self._adaptive_router_hooks() - vacated_keys: Final = self._host.callbacks.drop_callbacks( - ( - *self._callback_selectors(), - *self._override_selectors.values(), - *(self.optional_callbacks or ()), - *slack_alerting_loggers, - *adaptive_router_hooks, - ), - method_owners=(self, *slack_alerting_loggers), - ) - with _routers_holding_global_callbacks_lock: - survivors: Final = tuple(_routers_holding_global_callbacks) - candidates: Final = chain.from_iterable(router._global_callback_candidates() for router in survivors) - for candidate in self._host.callbacks.with_custom_logger_keys(candidates, vacated_keys): - self._register_router_selector(candidate) - adaptive_keys: Final = frozenset(map(self._host.callbacks.get_custom_logger_key, adaptive_router_hooks)) - adaptive_survivor: Final = next((router for router in survivors if router._core_adaptive_routers), None) - if adaptive_survivor is not None and vacated_keys & adaptive_keys: - adaptive_survivor._sync_adaptive_router_hooks() - - def _core_serve_through_request_logging(self) -> None: - self._request_scoped_selectors = (*self._callback_selectors(), *self._request_adaptive_router_hooks()) - self._request_scoped_policies = (*self._optional_loggers(), *self._slack_alerting_loggers()) - self._release_global_callbacks() - - def _core_bind_to_request_logging(self, logging_obj: LiteLLMLogging) -> None: - logging_obj.add_dynamic_function_callbacks( - async_success=self.deployment_callback_on_success, - failure=self.deployment_callback_on_failure, - async_failure=self.async_deployment_callback_on_failure, - ) - for selector in self._request_scoped_selectors: - logging_obj.replace_global_callback(selector) - for policy in self._request_policies_not_set_globally(): - logging_obj.add_dynamic_callback(policy) - - def _request_policies_not_set_globally(self) -> tuple[CustomLogger, ...]: - global_keys: Final = frozenset( - map( - self._host.callbacks.get_custom_logger_key, - (callback for callback in self._host.callback_list("callbacks") if isinstance(callback, CustomLogger)), - ) - ) - return tuple( - policy - for policy in self._request_scoped_policies - if self._host.callbacks.get_custom_logger_key(policy) not in global_keys - ) - - def _request_adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: - from litellm.router_strategy.adaptive_router.hooks import ( - AdaptiveRouterPostCallHook, - ) - - return tuple( - AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) - for tagged in chain.from_iterable(self._core_adaptive_routers.values()) - ) - - def _routing_callbacks(self) -> Iterable[object]: - if not (self._request_scoped_selectors or self._request_scoped_policies): - return self._host.callback_list("callbacks") - replaced_keys: Final = frozenset( - map(self._host.callbacks.get_custom_logger_key, self._request_scoped_selectors) - ) - return chain( - self._host.callbacks.without_custom_logger_keys(self._host.callback_list("callbacks"), replaced_keys), - self._request_scoped_selectors, - self._request_policies_not_set_globally(), - ) - - def retire_routing_state(self) -> None: - for strategy in (*self._global_callback_candidates(), *self._override_selectors.values()): - if isinstance(strategy, BaseRoutingStrategy | RouterBudgetLimiting): - strategy.retire() - - def _slack_alerting_loggers(self) -> tuple[CustomLogger, ...]: - slack_alerting_logger: Final = getattr(self, "_core_slack_alerting_logger", None) - return (slack_alerting_logger,) if isinstance(slack_alerting_logger, CustomLogger) else () - - def _adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: - from litellm.router_strategy.adaptive_router.hooks import ( - AdaptiveRouterPostCallHook, - ) - - own_strategy_ids: Final = frozenset( - id(tagged.strategy) for tagged in chain.from_iterable(self._core_adaptive_routers.values()) - ) - hooks: Final = self._host.callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) - return tuple( - hook - for hook in hooks - if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own_strategy_ids - ) - - def _optional_loggers(self) -> tuple[CustomLogger, ...]: - return tuple(callback for callback in self.optional_callbacks or () if isinstance(callback, CustomLogger)) - - def _global_callback_candidates(self) -> tuple[CustomLogger, ...]: - return (*self._callback_selectors(), *self._optional_loggers(), *self._slack_alerting_loggers()) - - @staticmethod - def _create_redis_cache( - cache_config: dict[str, Any], - ) -> RedisCache | RedisClusterCache: - """ - Initializes either a RedisCache or RedisClusterCache based on the cache_config. - """ - startup_nodes = cache_config.get("startup_nodes") - if not startup_nodes: - _env_cluster_nodes: Final = get_secret("REDIS_CLUSTER_NODES") - if _env_cluster_nodes is not None and isinstance(_env_cluster_nodes, str): - startup_nodes = _OPAQUE_VALUE_ADAPTER.validate_python(json.loads(_env_cluster_nodes)) - - if startup_nodes: - return RedisClusterCache(**{**cache_config, "startup_nodes": startup_nodes}) - else: - return RedisCache(**cache_config) - - def _core_update_redis_cache(self, cache: RedisCache) -> None: - """ - Update the redis cache for the router, if none set. - - Allows proxy user to just do - ```yaml - litellm_settings: - cache: true - ``` - and caching to just work. - """ - self._core_cache.attach_redis_cache(cache) - self._claude_code_session_router_cache.attach_redis_cache(cache) - self._core_prompt_caching_pin_cache.attach_redis_cache(cache) - - # Maps a routing strategy string to the attribute on `self` that holds - # the default group's strategy selector for that strategy. (The selectors - # double as `CustomLogger` callbacks, hence the legacy `*_logger` attrs.) - _DEFAULT_SELECTOR_ATTR_BY_STRATEGY: dict[str, str] = { - "least-busy": "leastbusy_logger", - "usage-based-routing": "lowesttpm_logger", - "usage-based-routing-v2": "lowesttpm_logger_v2", - "latency-based-routing": "lowestlatency_logger", - "cost-based-routing": "lowestcost_logger", - } - - @staticmethod - def _normalize_strategy( - strategy: RoutingStrategy | str | None, - ) -> str | None: - if strategy is None: - return None - if isinstance(strategy, RoutingStrategy): - return strategy.value - return strategy - - @staticmethod - def _validate_routing_strategy(routing_strategy: RoutingStrategy | str | None) -> None: - validate_routing_strategy(routing_strategy) - - def _build_strategy_selector( - self, - strategy: RoutingStrategy | str, - routing_strategy_args: dict, - register_callbacks: bool = True, - ) -> RouterStrategySelector | None: - """ - Constructs a strategy selector for a given strategy. - Returns None for `simple-shuffle` (no selector needed) and unknown - strategies. - """ - selector: RouterStrategySelector | None = None - match self._normalize_strategy(strategy): - case RoutingStrategy.LEAST_BUSY.value: - selector = LeastBusyLoggingHandler( - router_cache=self._core_cache, reservation_store=self._core_least_busy_reservations - ) - case RoutingStrategy.USAGE_BASED_ROUTING.value: - selector = LowestTPMLoggingHandler( - router_cache=self._core_cache, - routing_args=routing_strategy_args, - ) - case RoutingStrategy.USAGE_BASED_ROUTING_V2.value: - selector = LowestTPMLoggingHandler_v2( - router_cache=self._core_cache, - routing_args=routing_strategy_args, - ) - case RoutingStrategy.LATENCY_BASED.value: - selector = LowestLatencyLoggingHandler( - router_cache=self._core_cache, - routing_args=routing_strategy_args, - ) - case RoutingStrategy.COST_BASED.value: - selector = LowestCostLoggingHandler( - router_cache=self._core_cache, - routing_args={}, - ) - case _: - pass - - if selector is not None and register_callbacks: - self._register_router_selector(selector) - - return selector - - def _register_router_selector(self, selector: CustomLogger) -> None: - if isinstance(selector, LeastBusyLoggingHandler): - if isinstance(self._host.callback_list("input_callback"), list): - self._host.callbacks.add_litellm_input_callback(selector) - else: - self._host.set_callback_list("input_callback", [selector]) - if isinstance(self._host.callback_list("callbacks"), list): - self._host.callbacks.add_litellm_callback(selector) - - def _callback_selectors(self) -> tuple[RouterStrategySelector, ...]: - return tuple( - selector - for selector in ( - *(getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()), - *chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()), - ) - if selector is not None - ) - - def _unregister_router_selectors(self, selectors: Sequence[object]) -> None: - """ - Drop router-owned strategy selectors from litellm's global callback - lists. Used before re-init (`routing_strategy_init` / - `_init_routing_groups`) and by `discard` so dead selectors stop - receiving callback events. `function_setup` copies `litellm.callbacks` - into the per-event lists on every request, so those are purged too. - """ - for selector in selectors: - if isinstance(selector, BaseRoutingStrategy): - selector.retire() - self._host.callbacks.drop_callbacks(selectors) - - def _apply_updated_routing_strategy_args(self) -> None: - """ - Re-link the default group's selector to the current `routing_strategy_args`. - - Selectors freeze their `RoutingArgs` at construction, so a runtime args - update would otherwise keep serving the boot-time values until restart. - Latency/usage state survives the rebuild: it lives in the shared router - cache, not on the selector. - """ - strategy: Final = self._normalize_strategy(self._core_routing_strategy) - if strategy == "lar1": - from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - - apply_lar1_routing_strategy(self._core_ctx, self.routing_strategy_args) - return - - attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "") - current: Final = getattr(self, attr, None) if attr is not None else None - if attr is None or current is None: - return - - try: - rebuilt: Final = self._build_strategy_selector( - strategy=strategy or "", - routing_strategy_args=self.routing_strategy_args, - ) - except (TypeError, ValidationError): - verbose_router_logger.exception( - "Invalid routing_strategy_args %s for '%s'; keeping the previous ones", - self.routing_strategy_args, - strategy, - ) - return - - self._unregister_router_selectors((current,)) - setattr(self, attr, rebuilt) - - def routing_strategy_init(self, routing_strategy: RoutingStrategy | str, routing_strategy_args: dict): - verbose_router_logger.info("Routing strategy: %s", routing_strategy) - self._validate_routing_strategy(routing_strategy) - self._reset_custom_routing_strategy() - - self._unregister_router_selectors( - [getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()] - + list(getattr(self, "_override_selectors", {}).values()) - ) - self._override_selectors = {} - - self.leastbusy_logger: LeastBusyLoggingHandler | None = None - self.lowesttpm_logger: LowestTPMLoggingHandler | None = None - self.lowesttpm_logger_v2: LowestTPMLoggingHandler_v2 | None = None - self.lowestlatency_logger: LowestLatencyLoggingHandler | None = None - self.lowestcost_logger: LowestCostLoggingHandler | None = None - - selector: Final = self._build_strategy_selector( - strategy=routing_strategy, - routing_strategy_args=routing_strategy_args, - ) - attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(self._normalize_strategy(routing_strategy) or "") - # TODO: legacy `self._logger` attributes are read directly by - # `get_settings()` and external callers. Fold the default group into - # `self._group_selectors["default"]` and drop these attribute writes — - # the dual storage is an antipattern preserved only for back-compat. - if attr is not None: - setattr(self, attr, selector) - - def _init_routing_groups( - self, - groups_input: list[RoutingGroup | dict] | None, - ) -> None: - """ - Validates and indexes `routing_groups`. Each `model_name` may belong to - at most one non-priority group. Constructs per-group strategy selectors so - groups with different `routing_strategy_args` track independent state. - - Models not claimed by a non-priority group are served by the implicit - `"default"` group, whose selectors are the `self._logger` - attributes set up in `routing_strategy_init`. - """ - if not groups_input: - self._replace_routing_groups(()) - return - - known_model_names: Final = frozenset( - m["model_name"] for m in (self._core_model_list or ()) if m.get("model_name") - ) - groups: Final = parse_routing_groups(groups_input, known_model_names=known_model_names) - - alias_names: Final = frozenset(self._core_model_group_alias or ()) - for group in groups: - if group.group_name in known_model_names or group.group_name in alias_names: - if group.routing_strategy == "priority": - raise ValueError("Priority routing group names must not shadow a model or alias") - verbose_router_logger.warning( - "routing_groups: group_name '%s' is shadowed by an existing model_name or model_group_alias; " - "the group's strategy still applies to its members, but the name is not callable until renamed.", - group.group_name, - ) - - built: Final = tuple( - ( - group, - self._build_strategy_selector( - strategy="simple-shuffle" if group.routing_strategy == "priority" else group.routing_strategy, - routing_strategy_args=group.routing_strategy_args or {}, - register_callbacks=False, - ), - ) - for group in groups - ) - self._replace_routing_groups(built) - - def _replace_routing_groups( - self, - built: tuple[tuple[RoutingGroup, RouterStrategySelector | None], ...], - ) -> None: - previous_selectors: Final[Mapping[str, Mapping[str, RouterStrategySelector]]] = getattr( - self, "_group_selectors", {} - ) - self._unregister_router_selectors( - tuple(sel for selectors in previous_selectors.values() for sel in selectors.values()) - ) - for _, selector in built: - if selector is not None: - self._register_router_selector(selector) - - self._core_routing_groups: dict[str, RoutingGroup] = {group.group_name: group for group, _ in built} - self._model_to_group: dict[str, str] = { - model_name: group.group_name - for group, _ in built - if group.routing_strategy != "priority" - for model_name in group.models - } - self._group_selectors: dict[str, dict[str, RouterStrategySelector]] = { - group.group_name: ( - {} if selector is None else {self._normalize_strategy(group.routing_strategy) or "": selector} - ) - for group, selector in built - } - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - - def get_routing_group(self, model_name: str) -> RoutingGroup | None: - """ - The routing group callable as `model_name`, or None. A real deployment - `model_name` added after init shadows a same-named group (mirroring - `_try_early_resolve_deployments_for_model_not_in_names`, where concrete - models win over indirection); config-time collisions are rejected by - `_init_routing_groups`. - """ - if not self._core_routing_groups: - return None - group: Final = self._core_routing_groups.get(model_name) - if ( - group is None - or model_name in self._core_model_name_to_deployment_indices - or model_name in (self._core_model_group_alias or {}) - ): - return None - return group - - def _get_routing_group_deployments( - self, model: str, team_id: str | None = None - ) -> list[DeploymentTypedDict] | None: # mutable-ok: list matches _get_all_deployments' contract for callers - """ - The union of member deployments for a routing group called as `model`, - or None when `model` is not a callable group. The requested name stays - the group name so strategy selectors key their state by it. - - `_common_checks_available_deployment` consults this BEFORE its - early-resolve step so a wildcard `default_deployment` or pattern route - cannot hijack a group call. Overall resolution precedence there: - specific deployment > model id > model_group_alias > routing group > - model_name > team/pattern/default fallbacks. - """ - if not self._core_routing_groups: - return None - routing_group: Final = self.get_routing_group(model) - if routing_group is None: - return None - return [ - apply_routing_group_priority(routing_group, member, deployment) - for member in routing_group.models - for deployment in self._core_get_all_deployments(model_name=member, team_id=team_id) - ] - - def _is_priority_routing_group(self, model: str) -> bool: - resolved: Final = self._core_get_model_from_alias(model=model) or model - group: Final = self.get_routing_group(resolved) - return group is not None and group.routing_strategy == "priority" - - def _core_is_recognized_model(self, model: str) -> bool: - """ - Whether `model` names something this router serves directly: a - deployment model_name, a deployment id, a `model_group_alias`, or a - callable routing group. Proxy request gates share this predicate so a - new virtual-model kind cannot be forgotten at one of them; wildcard, - default-deployment, and deployment-name fallbacks stay caller policy. - """ - return ( - model in self._core_model_names - or self._core_has_model_id(model) - or (self._core_model_group_alias is not None and model in self._core_model_group_alias) - or self.get_routing_group(model) is not None - ) - - def _core_routing_group_has_alternatives(self, model_group: str | None) -> bool: - """ - True when `model_group` names a callable routing group whose member - union spans more than one deployment. Cooldown handling passes the - FAILING REQUEST's model group here: a 429 on a group call cools the - member down so selection moves to the group's alternatives, while a - direct call to a single-deployment member keeps the - single-deployment-model-group cooldown exemption. - """ - if model_group is None: - return False - resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group - group: Final = self.get_routing_group(resolved) - if group is None: - return False - return sum(len(self._core_model_name_to_deployment_indices.get(member) or ()) for member in group.models) > 1 - - def _core_team_model_has_alternatives(self, deployment_id: str) -> bool: - deployment: Final = self._core_get_deployment(model_id=deployment_id) - if deployment is None: - return False - team_id: Final = deployment.model_info.team_id - public_model_name: Final = deployment.model_info.team_public_model_name - if team_id is None or public_model_name is None: - return False - siblings: Final = self._team_deployments_named(team_id, public_model_name) or () - routable_siblings: Final = self._filter_blocked_deployments([dict(sibling) for sibling in siblings]) - return len(routable_siblings) > 1 - - _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) - - def _get_request_routing_strategy_override(self, request_kwargs: dict | None) -> str | None: - """ - Reads a per-request `routing_strategy` override (forwarded by the proxy - from key/team `router_settings`) out of the request kwargs. - - Only strategies with a per-request-capable selector are honored; - anything else (unknown strings, `lar1`, `provider-budget-routing`) is - ignored with a warning so a bad value stored on a key or team can - never take down that caller's traffic. - """ - if not request_kwargs: - return None - raw_strategy: Final = request_kwargs.get("routing_strategy") - if raw_strategy is None: - return None - strategy = self._normalize_strategy(raw_strategy) if isinstance(raw_strategy, (str, RoutingStrategy)) else None - if not isinstance(strategy, str) or strategy not in self._OVERRIDABLE_ROUTING_STRATEGIES: - verbose_router_logger.warning( - "Ignoring per-request routing_strategy override '%s'; supported overrides: %s.", - raw_strategy, - sorted(self._OVERRIDABLE_ROUTING_STRATEGIES), - ) - return None - return strategy - - def _get_override_strategy_selector(self, strategy: str) -> RouterStrategySelector | None: - """ - Returns the selector for a per-request strategy override. - - Reuses the default group's selector when the override matches the - router's configured strategy (so shared state keeps accumulating in - one place); otherwise lazily builds one selector per strategy and - caches it for the router's lifetime so its usage/latency state - persists across requests. - """ - if strategy == self._normalize_strategy(self._core_routing_strategy): - attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy) - return getattr(self, attr, None) if attr is not None else None - with self._override_selectors_lock: - if strategy not in self._override_selectors: - self._override_selectors[strategy] = self._build_strategy_selector( - strategy=strategy, - routing_strategy_args={}, - register_callbacks=False, - ) - return self._override_selectors[strategy] - - def _override_selector_pre_call_check( - self, strategy: str | None, selector: RouterStrategySelector | None, deployment: dict - ) -> None: - """ - Override selectors are not in `litellm.callbacks`, so the pre-call check that - `routing_strategy_pre_call_checks` runs for the router's own selectors (rpm - accounting for `usage-based-routing-v2`) runs here, for the overriding request only. - """ - if selector is None or strategy is None or selector is not self._override_selectors.get(strategy): - return - selector.pre_call_check(deployment) - - async def _async_override_selector_pre_call_check( - self, - strategy: str | None, - selector: RouterStrategySelector | None, - deployment: dict, - parent_otel_span: Span | None, - ) -> None: - if selector is None or strategy is None or selector is not self._override_selectors.get(strategy): - return - await selector.async_pre_call_check(deployment, parent_otel_span) - - def _bind_override_selector_to_request( - self, strategy: str, selector: RouterStrategySelector | None, request_kwargs: Mapping[str, object] | None - ) -> None: - if selector is None or request_kwargs is None or strategy in self._globally_registered_strategies(): - return - logging_obj: Final = request_kwargs.get("litellm_logging_obj") - if isinstance(logging_obj, LiteLLMLogging): - logging_obj.add_dynamic_callback(selector) - - def _globally_registered_strategies(self) -> frozenset[str]: - configured: Final = ( - self._core_routing_strategy, - *(group.routing_strategy for group in self._core_routing_groups.values()), - ) - return frozenset( - normalized for normalized in map(self._normalize_strategy, configured) if normalized is not None - ) - - def _core_arm_routing_read_prefetch(self, model: str, request_kwargs: dict[str, object] | None = None) -> None: - """Declare the cooldown read (and, for usage-based routing, the usage read) that - `async_get_available_deployment` will make for `model` on the request's Redis batch, so admission's - flush carries it. A miss (alias, no batch) costs nothing: routing then reads as it always has.""" - try: - strategy, selector = self._get_routing_context(model, request_kwargs) - usage_selector: Final = ( - selector - if strategy == "usage-based-routing-v2" and isinstance(selector, LowestTPMLoggingHandler_v2) - else None - ) - deployments: Final = self._core_get_model_list(model_name=model) - if deployments: - RoutingPrefetch.arm(self._core_ctx, usage_selector, deployments) - except Exception as e: # noqa: BLE001 # a prefetch is an optimisation, never a reason to fail the request - verbose_router_logger.debug( - "routing read prefetch not armed for %s: %s", _without_line_breaks(model), _without_line_breaks(e) - ) - - def _get_routing_context( - self, model: str, request_kwargs: dict | None = None - ) -> tuple[str | None, RouterStrategySelector | None]: - """ - Resolves the routing strategy and selector to use for the given model. - - A per-request `routing_strategy` in `request_kwargs` (forwarded by the - proxy from key/team `router_settings`) takes precedence over both the - model's routing group and the router's top-level strategy, since it is - the most specific expression of caller intent. - - Otherwise every model belongs to exactly one group: an explicit entry - from `routing_groups` (either because `model` IS a callable group name, - or because it is a member of one), or the implicit `"default"` group - driven by the router's top-level `routing_strategy` / - `routing_strategy_args`. - - `self._core_routing_strategy` may be either a string or a `RoutingStrategy` - enum member (the constructor accepts both), so it is normalized to a - string here. Downstream call sites and `_select_deployment_*` arms - compare against string literals. - """ - override: Final = self._get_request_routing_strategy_override(request_kwargs) - if override is not None: - verbose_router_logger.debug("routing_group=request-override model=%s strategy=%s", model, override) - override_selector: Final = self._get_override_strategy_selector(override) - self._bind_override_selector_to_request(override, override_selector, request_kwargs) - return override, override_selector - - resolved_model: Final = self._core_get_model_from_alias(model=model) or model - group_name: Final = ( - resolved_model - if self.get_routing_group(resolved_model) is not None - else self._model_to_group.get(resolved_model) - ) - if group_name is None: - strategy = self._normalize_strategy(self._core_routing_strategy) - attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "") - selector = getattr(self, attr, None) if attr is not None else None - verbose_router_logger.debug("routing_group=default model=%s strategy=%s", model, strategy) - return strategy, selector - - group: Final = self._core_routing_groups[group_name] - if group.routing_strategy == "priority": - return "simple-shuffle", None - strategy = self._normalize_strategy(group.routing_strategy) - selector = self._group_selectors.get(group_name, {}).get(strategy or "") - verbose_router_logger.debug("routing_group=%s model=%s strategy=%s", group_name, model, strategy) - return strategy, selector - - async def _select_deployment_async( - self, - *, - strategy: str | None, - selector: Any | None, - model: str, - healthy_deployments: list, - messages: list[dict[str, str]] | None, - input: str | list | None, - request_kwargs: dict | None, - reserve: bool, - ) -> Any | None: - """ - Asks the strategy selector for a deployment. Caller handles - `simple-shuffle` separately (it does not flow through a selector). - Returns None for unknown strategies or when the selector is missing. - """ - if selector is None: - return None - match strategy: - case "least-busy": - return await selector.async_get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - reserve=reserve, - call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), - ) - case "usage-based-routing": - # `LowestTPMLoggingHandler` (v1) only exposes the sync - # `get_available_deployments`. Mirror the pre-routing-groups - # top-level fallback by calling it inline so groups using v1 - # still work from async callers. - return selector.get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - ) - case "usage-based-routing-v2" | "cost-based-routing": - return await selector.async_get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - ) - case "latency-based-routing": - return await selector.async_get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - ) - case _: - return None - - def _select_deployment_sync( - self, - *, - strategy: str | None, - selector: Any | None, - model: str, - healthy_deployments: list, - messages: list[dict[str, str]] | None, - input: str | list | None, - request_kwargs: dict | None, - reserve: bool, - ) -> Any | None: - """ - Sync sibling of `_select_deployment_async`. Caller handles - `simple-shuffle` separately. - """ - if selector is None: - return None - - # `cost-based-routing` is intentionally omitted — - # `LowestCostLoggingHandler` only implements - # `async_get_available_deployments` - match strategy: - case "least-busy": - return selector.get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - reserve=reserve, - call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), - ) - case "usage-based-routing" | "usage-based-routing-v2": - return selector.get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - ) - case "latency-based-routing": - return selector.get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - ) - case _: - return None - - def initialize_assistants_endpoint(self): - ## INITIALIZE PASS THROUGH ASSISTANTS ENDPOINT ## - self.acreate_assistants = self.factory_function(self._host.invoke("acreate_assistants")) - self.adelete_assistant = self.factory_function(self._host.invoke("adelete_assistant")) - self.aget_assistants = self.factory_function(self._host.invoke("aget_assistants")) - self.acreate_thread = self.factory_function(self._host.invoke("acreate_thread")) - self.aget_thread = self.factory_function(self._host.invoke("aget_thread")) - self.a_add_message = self.factory_function(self._host.invoke("a_add_message")) - self.aget_messages = self.factory_function(self._host.invoke("aget_messages")) - self.arun_thread = self.factory_function(self._host.invoke("arun_thread")) - - def _initialize_core_endpoints(self): - """Helper to initialize core router endpoints.""" - self.amoderation = self.factory_function(self._host.invoke("amoderation"), call_type="moderation") - self.aanthropic_messages = self.factory_function( - self._host.invoke("anthropic_messages"), call_type="anthropic_messages" - ) - self.anthropic_messages = self.factory_function( - self._host.invoke("anthropic_messages"), call_type="anthropic_messages" - ) - self.agenerate_content = self.factory_function( - self._host.invoke("agenerate_content"), call_type="agenerate_content" - ) - self.aadapter_generate_content = self.factory_function( - self._host.invoke("aadapter_generate_content"), call_type="aadapter_generate_content" - ) - self.aresponses = self.factory_function(self._host.invoke("aresponses"), call_type="aresponses") - self.afile_delete = self.factory_function(self._host.invoke("afile_delete"), call_type="afile_delete") - self.afile_content = self.factory_function(self._host.invoke("afile_content"), call_type="afile_content") - self.responses = self.factory_function(self._host.invoke("responses"), call_type="responses") - self.aget_responses = self.factory_function(self._host.invoke("aget_responses"), call_type="aget_responses") - self.acancel_responses = self.factory_function( - self._host.invoke("acancel_responses"), call_type="acancel_responses" - ) - self.acompact_responses = self.factory_function( - self._host.invoke("acompact_responses"), call_type="acompact_responses" - ) - self.adelete_responses = self.factory_function( - self._host.invoke("adelete_responses"), call_type="adelete_responses" - ) - self.alist_input_items = self.factory_function( - self._host.invoke("alist_input_items"), call_type="alist_input_items" - ) - self._arealtime = self.factory_function(self._host.invoke("_arealtime"), call_type="_arealtime") - self.acreate_realtime_client_secret = self.factory_function( - self._host.invoke("acreate_realtime_client_secret"), call_type="acreate_realtime_client_secret" - ) - self.arealtime_calls = self.factory_function(self._host.invoke("arealtime_calls"), call_type="arealtime_calls") - self.acreate_realtime_transcription_session = self.factory_function( - self._host.invoke("acreate_realtime_transcription_session"), - call_type="acreate_realtime_transcription_session", - ) - self.acreate_realtime_translation_client_secret = self.factory_function( - self._host.invoke("acreate_realtime_translation_client_secret"), - call_type="acreate_realtime_translation_client_secret", - ) - self.arealtime_translation_calls = self.factory_function( - self._host.invoke("arealtime_translation_calls"), - call_type="arealtime_translation_calls", - ) - self._aresponses_websocket = self.factory_function( - self._host.invoke("_aresponses_websocket"), call_type="_aresponses_websocket" - ) - self.acreate_fine_tuning_job = self.factory_function( - self._host.invoke("acreate_fine_tuning_job"), call_type="acreate_fine_tuning_job" - ) - self.acancel_fine_tuning_job = self.factory_function( - self._host.invoke("acancel_fine_tuning_job"), call_type="acancel_fine_tuning_job" - ) - self.alist_fine_tuning_jobs = self.factory_function( - self._host.invoke("alist_fine_tuning_jobs"), call_type="alist_fine_tuning_jobs" - ) - self.aretrieve_fine_tuning_job = self.factory_function( - self._host.invoke("aretrieve_fine_tuning_job"), call_type="aretrieve_fine_tuning_job" - ) - self.afile_list = self.factory_function(self._host.invoke("afile_list"), call_type="alist_files") - self.aimage_edit = self.factory_function(self._host.invoke("aimage_edit"), call_type="aimage_edit") - self.allm_passthrough_route = self.factory_function( - self._host.invoke("allm_passthrough_route"), call_type="allm_passthrough_route" - ) - # Note: acancel_batch is defined as a method on the Router class (not using factory_function) - # to properly handle model-to-provider mapping like acreate_batch and aretrieve_batch - - def _initialize_vector_store_endpoints(self): - """Initialize vector store endpoints.""" - from litellm.vector_stores.main import ( - adelete, - alist, - aretrieve, - asearch, - aupdate, - create, - delete, - list, - retrieve, - search, - update, - ) - - self.avector_store_search = self.factory_function(asearch, call_type="avector_store_search") - self.vector_store_search = self.factory_function(search, call_type="vector_store_search") - self.vector_store_create = self.factory_function(create, call_type="vector_store_create") - self.avector_store_retrieve = self.factory_function(aretrieve, call_type="avector_store_retrieve") - self.vector_store_retrieve = self.factory_function(retrieve, call_type="vector_store_retrieve") - self.avector_store_list = self.factory_function(alist, call_type="avector_store_list") - self.vector_store_list = self.factory_function(list, call_type="vector_store_list") - self.avector_store_update = self.factory_function(aupdate, call_type="avector_store_update") - self.vector_store_update = self.factory_function(update, call_type="vector_store_update") - self.avector_store_delete = self.factory_function(adelete, call_type="avector_store_delete") - self.vector_store_delete = self.factory_function(delete, call_type="vector_store_delete") - - def _initialize_vector_store_file_endpoints(self): - """Initialize vector store file endpoints.""" - from litellm.vector_store_files.main import ( - acreate as avector_store_file_create_fn, - ) - from litellm.vector_store_files.main import ( - adelete as avector_store_file_delete_fn, - ) - from litellm.vector_store_files.main import alist as avector_store_file_list_fn - from litellm.vector_store_files.main import ( - aretrieve as avector_store_file_retrieve_fn, - ) - from litellm.vector_store_files.main import ( - aretrieve_content as avector_store_file_content_fn, - ) - from litellm.vector_store_files.main import ( - aupdate as avector_store_file_update_fn, - ) - from litellm.vector_store_files.main import ( - create as vector_store_file_create_fn, - ) - from litellm.vector_store_files.main import ( - delete as vector_store_file_delete_fn, - ) - from litellm.vector_store_files.main import list as vector_store_file_list_fn - from litellm.vector_store_files.main import ( - retrieve as vector_store_file_retrieve_fn, - ) - from litellm.vector_store_files.main import ( - retrieve_content as vector_store_file_content_fn, - ) - from litellm.vector_store_files.main import ( - update as vector_store_file_update_fn, - ) - - self.avector_store_file_create = self.factory_function( - avector_store_file_create_fn, call_type="avector_store_file_create" - ) - self.vector_store_file_create = self.factory_function( - vector_store_file_create_fn, call_type="vector_store_file_create" - ) - self.avector_store_file_list = self.factory_function( - avector_store_file_list_fn, call_type="avector_store_file_list" - ) - self.vector_store_file_list = self.factory_function( - vector_store_file_list_fn, call_type="vector_store_file_list" - ) - self.avector_store_file_retrieve = self.factory_function( - avector_store_file_retrieve_fn, call_type="avector_store_file_retrieve" - ) - self.vector_store_file_retrieve = self.factory_function( - vector_store_file_retrieve_fn, call_type="vector_store_file_retrieve" - ) - self.avector_store_file_content = self.factory_function( - avector_store_file_content_fn, call_type="avector_store_file_content" - ) - self.vector_store_file_content = self.factory_function( - vector_store_file_content_fn, call_type="vector_store_file_content" - ) - self.avector_store_file_update = self.factory_function( - avector_store_file_update_fn, call_type="avector_store_file_update" - ) - self.vector_store_file_update = self.factory_function( - vector_store_file_update_fn, call_type="vector_store_file_update" - ) - self.avector_store_file_delete = self.factory_function( - avector_store_file_delete_fn, call_type="avector_store_file_delete" - ) - self.vector_store_file_delete = self.factory_function( - vector_store_file_delete_fn, call_type="vector_store_file_delete" - ) - - def _initialize_google_genai_endpoints(self): - """Initialize Google GenAI endpoints.""" - from litellm.google_genai import ( - agenerate_content, - agenerate_content_stream, - generate_content, - generate_content_stream, - ) - - self.agenerate_content = self.factory_function(agenerate_content, call_type="agenerate_content") - self.generate_content = self.factory_function(generate_content, call_type="generate_content") - self.agenerate_content_stream = self.factory_function( - agenerate_content_stream, call_type="agenerate_content_stream" - ) - self.generate_content_stream = self.factory_function( - generate_content_stream, call_type="generate_content_stream" - ) - - def _initialize_ocr_search_endpoints(self): - """Initialize OCR and search endpoints.""" - from litellm.ocr import aocr, ocr - - self.aocr = self.factory_function(aocr, call_type="aocr") - self.ocr = self.factory_function(ocr, call_type="ocr") - - from litellm.search import asearch, search - - self.asearch = self.factory_function(asearch, call_type="asearch") - self.search = self.factory_function(search, call_type="search") - - from litellm.decisions import adecisions, decisions - - self.adecisions = self.factory_function(adecisions, call_type="adecisions") - self.decisions = self.factory_function(decisions, call_type="decisions") - - from litellm.systemone import asystemone, systemone - - self.asystemone = self.factory_function(asystemone, call_type="asystemone") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - self.systemone = self.factory_function(systemone, call_type="systemone") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - - from litellm.extraction import aextraction, extraction - - self.aextraction = self.factory_function(aextraction, call_type="aextraction") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - self.extraction = self.factory_function(extraction, call_type="extraction") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - - def _initialize_video_endpoints(self): - """Initialize video endpoints.""" - from litellm.videos import ( - avideo_content, - avideo_create_character, - avideo_edit, - avideo_extension, - avideo_generation, - avideo_get_character, - avideo_list, - avideo_remix, - avideo_status, - video_content, - video_create_character, - video_edit, - video_extension, - video_generation, - video_get_character, - video_list, - video_remix, - video_status, - ) - - self.avideo_generation = self.factory_function(avideo_generation, call_type="avideo_generation") - self.video_generation = self.factory_function(video_generation, call_type="video_generation") - self.avideo_list = self.factory_function(avideo_list, call_type="avideo_list") - self.video_list = self.factory_function(video_list, call_type="video_list") - self.avideo_status = self.factory_function(avideo_status, call_type="avideo_status") - self.video_status = self.factory_function(video_status, call_type="video_status") - self.avideo_content = self.factory_function(avideo_content, call_type="avideo_content") - self.video_content = self.factory_function(video_content, call_type="video_content") - self.avideo_remix = self.factory_function(avideo_remix, call_type="avideo_remix") - self.video_remix = self.factory_function(video_remix, call_type="video_remix") - self.avideo_create_character = self.factory_function( - avideo_create_character, call_type="avideo_create_character" - ) - self.video_create_character = self.factory_function(video_create_character, call_type="video_create_character") - self.avideo_get_character = self.factory_function(avideo_get_character, call_type="avideo_get_character") - self.video_get_character = self.factory_function(video_get_character, call_type="video_get_character") - self.avideo_edit = self.factory_function(avideo_edit, call_type="avideo_edit") - self.video_edit = self.factory_function(video_edit, call_type="video_edit") - self.avideo_extension = self.factory_function(avideo_extension, call_type="avideo_extension") - self.video_extension = self.factory_function(video_extension, call_type="video_extension") - - def _initialize_container_endpoints(self): - """Initialize container endpoints.""" - from litellm.containers import ( - acreate_container, - adelete_container, - alist_containers, - aretrieve_container, - create_container, - delete_container, - list_containers, - retrieve_container, - ) - from litellm.containers.endpoint_factory import ( - _generated_endpoints as container_file_endpoints, - ) - - self.acreate_container = self.factory_function(acreate_container, call_type="acreate_container") - self.create_container = self.factory_function(create_container, call_type="create_container") - self.alist_containers = self.factory_function(alist_containers, call_type="alist_containers") - self.list_containers = self.factory_function(list_containers, call_type="list_containers") - self.aretrieve_container = self.factory_function(aretrieve_container, call_type="aretrieve_container") - self.retrieve_container = self.factory_function(retrieve_container, call_type="retrieve_container") - self.adelete_container = self.factory_function(adelete_container, call_type="adelete_container") - self.delete_container = self.factory_function(delete_container, call_type="delete_container") - - # Auto-register JSON-generated container file endpoints - for name, func in container_file_endpoints.items(): - setattr(self, name, self.factory_function(func, call_type=name)) - - def _initialize_skills_endpoints(self): - """Initialize Anthropic Skills API endpoints.""" - self.acreate_skill = self.factory_function(self._host.invoke("acreate_skill"), call_type="acreate_skill") - self.alist_skills = self.factory_function(self._host.invoke("alist_skills"), call_type="alist_skills") - self.aget_skill = self.factory_function(self._host.invoke("aget_skill"), call_type="aget_skill") - self.adelete_skill = self.factory_function(self._host.invoke("adelete_skill"), call_type="adelete_skill") - - def _initialize_interactions_endpoints(self): - """Initialize Google Interactions API endpoints.""" - from litellm.interactions import acancel as acancel_interaction - from litellm.interactions import acreate as acreate_interaction - from litellm.interactions import adelete as adelete_interaction - from litellm.interactions import aget as aget_interaction - from litellm.interactions import cancel as cancel_interaction - from litellm.interactions import create as create_interaction - from litellm.interactions import delete as delete_interaction - from litellm.interactions import get as get_interaction - - self.acreate_interaction = self.factory_function(acreate_interaction, call_type="acreate_interaction") - self.create_interaction = self.factory_function(create_interaction, call_type="create_interaction") - self.aget_interaction = self.factory_function(aget_interaction, call_type="aget_interaction") - self.get_interaction = self.factory_function(get_interaction, call_type="get_interaction") - self.adelete_interaction = self.factory_function(adelete_interaction, call_type="adelete_interaction") - self.delete_interaction = self.factory_function(delete_interaction, call_type="delete_interaction") - self.acancel_interaction = self.factory_function(acancel_interaction, call_type="acancel_interaction") - self.cancel_interaction = self.factory_function(cancel_interaction, call_type="cancel_interaction") - - def _initialize_managed_agents_endpoints(self): - """Initialize Google Managed Agents API endpoints (v1beta/agents).""" - from litellm.interactions.agents import acreate as acreate_agent - from litellm.interactions.agents import adelete as adelete_agent - from litellm.interactions.agents import aget as aget_agent - from litellm.interactions.agents import alist as alist_agents - from litellm.interactions.agents import alist_versions as alist_agent_versions - from litellm.interactions.agents import create as create_agent - from litellm.interactions.agents import delete as delete_agent - from litellm.interactions.agents import get as get_agent - from litellm.interactions.agents import list as list_agents - from litellm.interactions.agents import list_versions as list_agent_versions - - self.acreate_agent = self.factory_function(acreate_agent, call_type="acreate_agent") - self.create_agent = self.factory_function(create_agent, call_type="create_agent") - self.alist_agents = self.factory_function(alist_agents, call_type="alist_agents") - self.list_agents = self.factory_function(list_agents, call_type="list_agents") - self.aget_agent = self.factory_function(aget_agent, call_type="aget_agent") - self.get_agent = self.factory_function(get_agent, call_type="get_agent") - self.adelete_agent = self.factory_function(adelete_agent, call_type="adelete_agent") - self.delete_agent = self.factory_function(delete_agent, call_type="delete_agent") - self.alist_agent_versions = self.factory_function(alist_agent_versions, call_type="alist_agent_versions") - self.list_agent_versions = self.factory_function(list_agent_versions, call_type="list_agent_versions") - - def _initialize_specialized_endpoints(self): - """Helper to initialize specialized router endpoints (vector store, OCR, search, video, container, skills, interactions).""" - self._initialize_vector_store_endpoints() - self._initialize_vector_store_file_endpoints() - self._initialize_google_genai_endpoints() - self._initialize_ocr_search_endpoints() - # Override vector store methods with router-aware implementations - self._override_vector_store_methods_for_router() - self._initialize_video_endpoints() - self._initialize_container_endpoints() - self._initialize_skills_endpoints() - self._initialize_interactions_endpoints() - self._initialize_managed_agents_endpoints() - - def initialize_router_endpoints(self): - self._initialize_core_endpoints() - self._initialize_specialized_endpoints() - - def validate_fallbacks(self, fallback_param: Sequence[Mapping[str, object]] | None): - """ - Validate the fallbacks parameter. - """ - if fallback_param is None: - return - for fallback_dict in fallback_param: - if not isinstance(fallback_dict, dict): - raise ValueError(f"Item '{fallback_dict}' is not a dictionary.") - if len(fallback_dict) != 1: - raise ValueError( - f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." - ) - ((model_group, targets),) = fallback_dict.items() - if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=[fallback_dict]): - continue - raise TypeError( - f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. " - f"Use {{'{model_group}': ['{targets}']}}" - ) - - def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None: - def _move_before_deployment_affinity( - callback_list: list[_CallbackT], - callback_to_move: _CallbackT, - ) -> None: - if callback_to_move not in callback_list: - return - callback_list.remove(callback_to_move) - insert_index: Final = next( - (idx for idx, callback in enumerate(callback_list) if isinstance(callback, DeploymentAffinityCheck)), - len(callback_list), - ) - callback_list.insert(insert_index, callback_to_move) - - if enable_global_affinity or EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - self.model_group_affinity_config - ): - if self.optional_callbacks is None: - self.optional_callbacks = [] - - existing_ec_callback: EncryptedContentAffinityCheck | None = None - for cb in self.optional_callbacks: - if isinstance(cb, EncryptedContentAffinityCheck): - existing_ec_callback = cb - break - - if existing_ec_callback is not None: - existing_ec_callback.ctx = self._core_ctx - existing_ec_callback.enable_global_affinity = ( - existing_ec_callback.enable_global_affinity or enable_global_affinity - ) - existing_ec_callback.model_group_affinity_config = self.model_group_affinity_config or {} - ec_callback = existing_ec_callback - else: - ec_callback = EncryptedContentAffinityCheck( - ctx=self._core_ctx, - enable_global_affinity=enable_global_affinity, - model_group_affinity_config=self.model_group_affinity_config, - ) - self.optional_callbacks.append(ec_callback) - self._host.callbacks.add_litellm_callback(ec_callback) - - _move_before_deployment_affinity(self.optional_callbacks, ec_callback) - _move_before_deployment_affinity(self._host.callback_list("callbacks"), ec_callback) - - def _ensure_deployment_affinity_callback(self) -> None: - """Register the DeploymentAffinityCheck callback (global flags all False) if absent. - - Needed when nothing enabled a global affinity flag but affinity can still - activate per request: per-group `model_group_affinity_config` entries, or the - session-affinity marker a complexity router stamps at pre-routing time. - """ - if any(isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])): - return - if self.optional_callbacks is None: - self.optional_callbacks = [] - affinity_callback: Final = DeploymentAffinityCheck( - cache=self._core_cache, - ttl_seconds=self.deployment_affinity_ttl_seconds, - enable_user_key_affinity=False, - enable_responses_api_affinity=False, - enable_session_id_affinity=False, - model_group_affinity_config=self.model_group_affinity_config, - is_priority_group=self._is_priority_routing_group, - session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, - ) - self.optional_callbacks.append(affinity_callback) - self._host.callbacks.add_litellm_callback(affinity_callback) - - def add_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None): - if optional_pre_call_checks is None: - return - - # --------------------------------------------------------------------- - # Unified deployment affinity (session stickiness) - # --------------------------------------------------------------------- - enable_user_key_affinity: Final = "deployment_affinity" in optional_pre_call_checks - enable_responses_api_affinity: Final = "responses_api_deployment_check" in optional_pre_call_checks - enable_session_id_affinity: Final = "session_affinity" in optional_pre_call_checks - if enable_user_key_affinity or enable_responses_api_affinity or enable_session_id_affinity: - if self.optional_callbacks is None: - self.optional_callbacks = [] - - existing_affinity_callback: DeploymentAffinityCheck | None = None - for cb in self.optional_callbacks: - if isinstance(cb, DeploymentAffinityCheck): - existing_affinity_callback = cb - break - - if existing_affinity_callback is not None: - existing_affinity_callback.enable_user_key_affinity = ( - existing_affinity_callback.enable_user_key_affinity or enable_user_key_affinity - ) - existing_affinity_callback.enable_responses_api_affinity = ( - existing_affinity_callback.enable_responses_api_affinity or enable_responses_api_affinity - ) - existing_affinity_callback.enable_session_id_affinity = ( - existing_affinity_callback.enable_session_id_affinity or enable_session_id_affinity - ) - existing_affinity_callback.ttl_seconds = self.deployment_affinity_ttl_seconds - if self.model_group_affinity_config: - existing_affinity_callback.model_group_affinity_config = self.model_group_affinity_config - else: - affinity_callback: Final = DeploymentAffinityCheck( - cache=self._core_cache, - ttl_seconds=self.deployment_affinity_ttl_seconds, - enable_user_key_affinity=enable_user_key_affinity, - enable_responses_api_affinity=enable_responses_api_affinity, - enable_session_id_affinity=enable_session_id_affinity, - model_group_affinity_config=self.model_group_affinity_config, - is_priority_group=self._is_priority_routing_group, - session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, - ) - self.optional_callbacks.append(affinity_callback) - self._host.callbacks.add_litellm_callback(affinity_callback) - - # --------------------------------------------------------------------- - # Encrypted content affinity - # --------------------------------------------------------------------- - self._add_encrypted_content_affinity_check( - enable_global_affinity=("encrypted_content_affinity" in optional_pre_call_checks) - ) - - # --------------------------------------------------------------------- - # Remaining optional pre-call checks - # --------------------------------------------------------------------- - for pre_call_check in optional_pre_call_checks: - _callback: CustomLogger | None = None - if pre_call_check in ( - "deployment_affinity", - "responses_api_deployment_check", - "session_affinity", - "encrypted_content_affinity", - ): - continue - if pre_call_check == "prompt_caching": - _callback = PromptCachingDeploymentCheck( - cache=self._core_prompt_caching_pin_cache, is_priority_group=self._is_priority_routing_group - ) - elif pre_call_check == "router_budget_limiting": - if self._core_get_router_deployment_budget_limiter() is not None: - continue - _callback = RouterBudgetLimiting( - dual_cache=self._core_cache, - provider_budget_config=self._core_provider_budget_config, - model_list=self._core_model_list, - ) - self.router_budget_logger = _callback - elif pre_call_check == "enforce_model_rate_limits": - _callback = ModelRateLimitingCheck(dual_cache=self._core_cache) - - if _callback is None: - continue - - if self.optional_callbacks is not None and any( - isinstance(callback, type(_callback)) for callback in self.optional_callbacks - ): - continue - if self.optional_callbacks is None: - self.optional_callbacks = [] - self.optional_callbacks.append(_callback) - self._host.callbacks.add_litellm_callback(_callback) - - def set_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None) -> None: - if optional_pre_call_checks is None: - return - requested: Final = frozenset(optional_pre_call_checks) - for name, callback_cls in _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS.items(): - if name not in requested: - self._remove_optional_callbacks_of_type(callback_cls) - self.add_optional_pre_call_checks(optional_pre_call_checks) - - def _remove_optional_callbacks_of_type(self, callback_cls: type[CustomLogger]) -> None: - if self.optional_callbacks is None or not any(type(cb) is callback_cls for cb in self.optional_callbacks): - return - self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls] - if any( - router is not self and any(type(cb) is callback_cls for cb in (router.optional_callbacks or [])) - for router in (*_live_routers, *(owner for owner in live_check_owners if not isinstance(owner, PythonCore))) - ): - return - for cb in tuple(self._host.callback_list("callbacks")): - if type(cb) is callback_cls: - self._host.callbacks.remove_callback_from_list_by_object( - self._host.callback_list("callbacks"), cb, require_self=False - ) - - def print_deployment(self, deployment: dict): - """ - returns a copy of the deployment with the api key masked - - Only returns 2 characters of the api key and masks the rest with * (10 *). - """ - try: - _deployment_copy: Final = copy.deepcopy(deployment) - litellm_params: Final[dict] = _deployment_copy["litellm_params"] - - if self._host.dynamic().redact_user_api_key_info: - masker: Final = SensitiveDataMasker(visible_prefix=2, visible_suffix=0) - _deployment_copy["litellm_params"] = masker.mask_dict(litellm_params) - elif "api_key" in litellm_params: - litellm_params["api_key"] = litellm_params["api_key"][:2] + "*" * 10 - - return _deployment_copy - except Exception as e: - verbose_router_logger.debug("Error occurred while printing deployment - %s", e) - raise e - - @staticmethod - def _deployment_params_with_request_reasoning_override( - deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object] - ) -> dict[str, object]: # mutable-ok: litellm's request pipeline consumes a mutable kwargs mapping - """Return deployment params whose equivalent effort controls cannot outrank a request override. - - Providers expose the same setting through several native carriers. A request-level - ``reasoning_effort`` is the portable override, so a deployment's ``thinking`` or nested - ``*.effort`` must not remain beside it and either win or trigger a conflicting-params 400. - Every changed mapping is copied so the Router's shared deployment config stays immutable. - """ - sanitized: Final = dict(deployment_params) - if request_kwargs.get("reasoning_effort") is None: - return sanitized - - sanitized.pop("thinking", None) - PythonCore._pop_effort_from_nested_carrier(sanitized, "output_config") - PythonCore._pop_effort_from_nested_carrier(sanitized, "reasoning") - - extra_body: Final = sanitized.get("extra_body") - if isinstance(extra_body, Mapping): - sanitized_extra_body: Final = dict(extra_body) - sanitized_extra_body.pop("reasoning_effort", None) - sanitized_extra_body.pop("thinking", None) - PythonCore._pop_effort_from_nested_carrier(sanitized_extra_body, "output_config") - PythonCore._pop_effort_from_nested_carrier(sanitized_extra_body, "reasoning") - if sanitized_extra_body: - sanitized["extra_body"] = sanitized_extra_body - else: - sanitized.pop("extra_body", None) - return sanitized - - @staticmethod - def _is_classifier_internal_call(kwargs: Mapping[str, object]) -> bool: - metadata: Final = kwargs.get("metadata") - litellm_metadata: Final = kwargs.get("litellm_metadata") - return any( - isinstance(candidate, Mapping) - and candidate.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == AUTOROUTER_CLASSIFIER_CALL_ORIGIN - for candidate in (metadata, litellm_metadata) - ) - - def _drop_unsupported_classifier_reasoning_effort( - self, - deployment: DeploymentTypedDict, - model: str, - kwargs: dict[str, object], # mutable-ok: fallback must update the active request and its log body together - ) -> None: - """Let a classifier fallback without reasoning support remain a usable fallback. - - The dashboard only offers explicitly advertised levels, but an existing config can outlive - a model change and fallbacks can target a different group. Unknown capability fails open; - only a provider that explicitly rejects the parameter has it removed. - """ - if kwargs.get("reasoning_effort") is None or not self._is_classifier_internal_call(kwargs): - return - if self._deployment_accepts_param(deployment, model, "reasoning_effort"): - return - verbose_router_logger.warning( - "litellm.router.py: dropping classifier reasoning_effort for model=%s because the selected deployment does not support it", - model, - ) - kwargs.pop("reasoning_effort", None) - proxy_server_request: Final = kwargs.get("proxy_server_request") - if not isinstance(proxy_server_request, dict): - return - body: Final = proxy_server_request.get("body") - if isinstance(body, dict): - body.pop("reasoning_effort", None) - - ### COMPLETION, EMBEDDING, IMG GENERATION FUNCTIONS - - def completion(self, model: str, messages: list[dict[str, str]], **kwargs) -> ModelResponse | CustomStreamWrapper: - """ - Example usage: - response = router.completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, how's it going?"}] - """ - try: - verbose_router_logger.debug("router.completion(model=%s,..)", model) - kwargs["model"] = model - kwargs["messages"] = messages - kwargs["original_function"] = self._completion - controls: Final = per_request_fallback_controls(kwargs) - kwargs[MID_STREAM_FALLBACK_CONTROLS_KEY] = controls # rebind-ok: forwarded to every hop - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - - response: Final = self.function_with_fallbacks(**kwargs) - return response - except Exception as e: - raise e - - def _completion(self, model: str, messages: list[dict[str, str]], **kwargs) -> ModelResponse | CustomStreamWrapper: - model_name = None - deployment = None - try: - controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None) - input_kwargs_for_streaming_fallback: Final = mid_stream_fallback_snapshot_kwargs( - model=model, controls=controls, kwargs=kwargs - ) - - # pick the one that is available (lowest TPM/RPM) - with picks_reserve(True): - deployment = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._drop_unsupported_classifier_reasoning_effort( - deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment - model=model, - kwargs=kwargs, - ) - # Check for silent model experiment - # Make a local copy of litellm_params to avoid mutating the Router's state - litellm_params: Final = self._deployment_params_with_request_reasoning_override( - deployment["litellm_params"], kwargs - ) - silent_model: Final = litellm_params.pop("silent_model", None) - - for silent_target in _silent_experiment_targets(silent_model): - # Mirroring traffic to a secondary model - # Use threading.Thread (not ThreadPoolExecutor) - executor.submit() - # requires pickling args, which fails when kwargs contain unpicklable - # objects (e.g. _thread.RLock from OTEL spans, loggers) in deployment. - threading.Thread( - target=self._silent_experiment_completion, - args=(silent_target, messages), - kwargs=_silent_experiment_kwargs_snapshot(kwargs), - daemon=True, - ).start() - - kwargs.setdefault("messages", messages) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) - kwargs.pop("silent_model", None) # Ensure it's not in kwargs either - model_name = litellm_params["model"] - potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs) - # check if provided keys == client keys # - dynamic_api_key: Final = kwargs.get("api_key", None) - if ( - dynamic_api_key is not None - and potential_model_client is not None - and dynamic_api_key != potential_model_client.api_key - ): - model_client = None - else: - model_client = potential_model_client - - ### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit) - ## only run if model group given, not model id - if model in self._core_model_names or not self._core_has_model_id(model): - self.routing_strategy_pre_call_checks(deployment=deployment) - - input_kwargs: Final = without_router_only_kwargs( - { - **litellm_params, - "messages": messages, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - response: Final = self._host.invoke("completion")(**input_kwargs) - verbose_router_logger.info("litellm.completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - ## CHECK CONTENT FILTER ERROR ## - if isinstance(response, ModelResponse): - _should_raise = self._should_raise_content_policy_error(model=model, response=response, kwargs=kwargs) - if _should_raise: - raise litellm_exceptions.ContentPolicyViolationError( - message="Response output was blocked.", - model=model, - llm_provider="", - ) - - if ( - isinstance(response, CustomStreamWrapper) - and response.completion_stream is None - and response.make_call is not None - ): - response.fetch_sync_stream() - - # Wrap streaming responses so MidStreamFallbackError (raised - # during iteration) triggers the Router's fallback chain. - if isinstance(response, CustomStreamWrapper): - return self._completion_streaming_iterator( - model_response=response, - messages=messages, - initial_kwargs=input_kwargs_for_streaming_fallback, - first_chunk=_read_first_chunk(response), - ) - - return response - except Exception as e: - verbose_router_logger.info("litellm.completion(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - # Set per-deployment num_retries on exception for retry logic - if deployment is not None: - self._set_deployment_num_retries_on_exception(e, deployment) - self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) - raise e - - def _get_silent_experiment_kwargs(self, **kwargs) -> dict: - """ - Prepare kwargs for a silent experiment by ensuring isolation from the primary call. - - Guarantee metadata isolation: safe_deep_copy falls back to the original - reference when deepcopy fails (e.g. metadata contains UserAPIKeyAuth with - parent_otel_span — an OTel Span that is not deepcopy-able). Force a shallow - copy of the metadata dict so mutations (model_group, is_silent_experiment) - never corrupt the main call's metadata. - """ - from litellm.litellm_core_utils.core_helpers import safe_deep_copy - - silent_kwargs: Final = safe_deep_copy(kwargs) - - # safe_deep_copy may fall back to the original metadata reference when - # deepcopy fails (UserAPIKeyAuth.parent_otel_span is not deepcopy-able). - # Detect this via identity check and force a shallow copy so that setting - # model_group / is_silent_experiment on the silent dict doesn't corrupt - # the primary call's metadata. - original_metadata: Final = kwargs.get("metadata") - if original_metadata is not None and silent_kwargs.get("metadata") is original_metadata: - silent_kwargs["metadata"] = dict(original_metadata) - - if "metadata" not in silent_kwargs: - silent_kwargs["metadata"] = {} - - # OTel spans are not safe to use across event loops. The silent - # experiment runs in a new event loop, so strip the span to prevent - # cross-loop tracing races or span corruption. - silent_kwargs["metadata"].pop("litellm_parent_otel_span", None) - - silent_kwargs["metadata"]["is_silent_experiment"] = True - - # Pop logging objects and call IDs to ensure a fresh logging context - # This prevents collisions in the Proxy's database (spend_logs) - silent_kwargs.pop("litellm_call_id", None) - silent_kwargs.pop("litellm_logging_obj", None) - silent_kwargs.pop("standard_logging_object", None) - # DON'T pop proxy_server_request — it's needed for spend log metadata - - return silent_kwargs - - async def _run_silent_experiment( - self, silent_model: str, messages: Sequence[Mapping[str, str]], silent_kwargs: Mapping[str, object] - ) -> None: - remaining_kwargs: Final = MappingProxyType( - {key: value for key, value in silent_kwargs.items() if key != "stream"} - ) - with mirrored_call(): - response: Final = await self.acompletion( - model=silent_model, - messages=cast(list[AllMessageValues], messages), - stream=bool(silent_kwargs.get("stream", False)), - **remaining_kwargs, - ) - if not isinstance(response, CustomStreamWrapper): - return - async for _ in response: - pass - - def _silent_experiment_completion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): - """ - Run a silent experiment in the background (thread). - """ - try: - # Prevent infinite recursion if silent model also has a silent model - if kwargs.get("metadata", {}).get("is_silent_experiment", False): - return - - messages = copy.deepcopy(messages) - - verbose_router_logger.info("Starting silent experiment for model %s", silent_model) - - silent_kwargs: Final = self._get_silent_experiment_kwargs(**kwargs) - - # Override model_group to correctly attribute metrics to the silent model - silent_kwargs["metadata"]["model_group"] = silent_model - - # Create a new event loop for this thread so that async success - # callbacks (e.g. _ProxyDBLogger) can schedule and run DB writes. - loop: Final = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - try: - - async def _run_silent_completion(): - await self._run_silent_experiment(silent_model, messages, silent_kwargs) - # Drain any fire-and-forget tasks (e.g. alerting hooks) - # scheduled via asyncio.create_task during acompletion. - pending: Final = asyncio.all_tasks() - current: Final = asyncio.current_task() - if current is not None: - pending.discard(current) - if pending: - await asyncio.gather(*pending, return_exceptions=True) - - loop.run_until_complete(_run_silent_completion()) - finally: - loop.close() - except Exception as e: - verbose_router_logger.error("Silent experiment failed for model %s: %s", silent_model, e) - - # fmt: off - - @overload - async def acompletion( - self, model: str, messages: list[AllMessageValues], stream: Literal[True], **kwargs - ) -> CustomStreamWrapper: - ... - - @overload - async def acompletion( - self, model: str, messages: list[AllMessageValues], stream: Literal[False] = False, **kwargs - ) -> ModelResponse: - ... - - @overload - async def acompletion( - self, model: str, messages: list[AllMessageValues], stream: Literal[True, False] = False, **kwargs - ) -> CustomStreamWrapper | ModelResponse: - ... - - # fmt: on - - # The actual implementation of the function - async def acompletion( - self, - model: str, - messages: list[AllMessageValues], - stream: bool = False, - **kwargs, - ): - try: - kwargs["model"] = model - kwargs["messages"] = messages - kwargs["stream"] = stream - kwargs["original_function"] = self._acompletion - - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - start_time: Final = time.time() - is_prompt_management_model: Final = self._is_prompt_management_model(model) - if not is_prompt_management_model: - controls: Final = per_request_fallback_controls(kwargs) - kwargs[MID_STREAM_FALLBACK_CONTROLS_KEY] = controls # rebind-ok: forwarded to every hop - request_priority, request_kwargs = self._core_split_request_priority(model=model, kwargs=kwargs) - original_function: Final[Callable[..., Awaitable[object]]] = ( - self._prompt_management_acompletion(model=model, messages=messages) - if is_prompt_management_model - else self.async_function_with_fallbacks - ) - if request_priority is None: - response = await original_function(**request_kwargs) - else: - response = await self._schedule_factory( - model=model, - priority=request_priority, - original_function=original_function, - args=(), - kwargs=request_kwargs, - ) - end_time: Final = time.time() - _duration: Final = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.ROUTER, - duration=_duration, - call_type="acompletion", - start_time=start_time, - end_time=end_time, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) - ) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - @staticmethod - def _combine_fallback_usage( - fallback_item: ModelResponseStream, - complete_response_object_usage: Usage | None, - ) -> None: - """Merge partial-stream usage with fallback-stream usage on the chunk.""" - from litellm.cost_calculator import BaseTokenUsageProcessor - - usage: Final = cast(Usage | None, getattr(fallback_item, "usage", None)) - usage_objects: Final = [usage] if usage is not None else [] - if complete_response_object_usage is not None: - usage_objects.append(complete_response_object_usage) - combined_usage: Final = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=usage_objects) - setattr(fallback_item, "usage", combined_usage) - - @staticmethod - def _price_cut_partial( - cut_stream: CustomStreamWrapper, - partial_response: ModelResponse | TextCompletionResponse | None, - ) -> "ContinuedPartial | None": - """Price the cut stream's partial, plus any partial it inherited, before a fallback can reuse its logging.""" - from litellm.cost_calculator import BaseTokenUsageProcessor - - cut_logging: Final = cut_stream.logging_obj - partial_usage: Final = getattr(partial_response, "usage", None) - if partial_response is None or not isinstance(partial_usage, Usage): - return None - inherited_usage: Final = cut_logging.model_call_details.get("continued_partial_usage") - inherited_cost: Final = cut_logging.model_call_details.get("continued_partial_cost") - inherited_id: Final = cut_logging.model_call_details.get("continued_response_id") - priced_partial: Final = partial_response.model_copy( - update={"model": cut_stream.model} if cut_stream.model else {} - ) - priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor - key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost" - } - return ContinuedPartial( - usage=BaseTokenUsageProcessor.combine_usage_objects( - usage_objects=[partial_usage, inherited_usage] - if isinstance(inherited_usage, Usage) - else [partial_usage] - ), - response_cost=(cut_logging.response_cost_calculator(result=priced_partial) or 0.0) - + (inherited_cost if isinstance(inherited_cost, float) else 0.0), - response_id=inherited_id if isinstance(inherited_id, str) else partial_response.id, - ) - - @staticmethod - def _carry_partial_spend_to_fallback(fallback_response: object, cut_partial: "ContinuedPartial | None") -> None: - """Bill what the cut legs consumed on the continuation's success log.""" - fallback_logging: Final = getattr(fallback_response, "logging_obj", None) - if cut_partial is None or not isinstance(fallback_logging, LiteLLMLogging): - return - fallback_logging.record_continued_partial( - usage=cut_partial.usage, response_cost=cut_partial.response_cost, response_id=cut_partial.response_id - ) - - @staticmethod - def _prepare_fallback_hidden_params( - fallback_response: object, - ) -> tuple[dict[str, object], dict[str, object]]: - fallback_hidden_params: Final = get_hidden_params_dict(fallback_response) - fallback_headers: Final = fallback_hidden_params.get("additional_headers") - if not isinstance(fallback_headers, dict): - return fallback_hidden_params, {} - return fallback_hidden_params, cast("dict[str, object]", fallback_headers) - - @staticmethod - def _adopt_fallback_response_headers( - wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]", - fallback_response: object, - ) -> tuple[dict[str, object], dict[str, object]]: - """Repoint the wrapper at `fallback_response`, returning its prepared hidden params.""" - prepared: Final = PythonCore._prepare_fallback_hidden_params(fallback_response) - adopting_wrapper: Final = wrapper_ref() - if adopting_wrapper is not None: - adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared) - return prepared - - @staticmethod - def _apply_fallback_hidden_params_to_item( - fallback_item: object, - prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]], - ) -> None: - if fallback_item is None or not hasattr(fallback_item, _HIDDEN_PARAMS_ATTR): - return - - fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params - item_hidden_params: Final = get_hidden_params_dict(fallback_item) - item_headers = item_hidden_params.get("additional_headers") - if not isinstance(item_headers, dict): - item_headers = {} - - set_hidden_params( - fallback_item, - { - **item_hidden_params, - **fallback_hidden_params, - "additional_headers": {**item_headers, **fallback_headers}, - }, - ) - - async def _acompletion_streaming_iterator( - self, - model_response: CustomStreamWrapper, - messages: list[dict[str, str]], - initial_kwargs: dict[str, Any], - deployment_slot: contextlib.AsyncExitStack | None = None, - first_chunk: _FirstChunk[AsyncIterator[object]] | None = None, - ) -> CustomStreamWrapper: - """ - Helper to iterate over a streaming response. - - Catches errors for fallbacks using the router's fallback system - - `deployment_slot` holds the deployment's max_parallel_requests semaphore; it is - released when the stream is exhausted, closed, or falls back to another deployment - - `first_chunk` carries the chunk _acompletion already read (None: iterate from the start) - """ - from litellm.exceptions import MidStreamFallbackError - - held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack() - - class FallbackStreamWrapper(FallbackAwareStreamWrapper): - def __init__(self, async_generator: AsyncGenerator): - # Copy attributes from the original model_response - super().__init__( - completion_stream=async_generator, - model=model_response.model, - custom_llm_provider=model_response.custom_llm_provider, - logging_obj=model_response.logging_obj, - _response_headers=getattr(model_response, "_response_headers", None), - ) - self._async_generator = async_generator - inner_chunks: Final[object] = getattr(model_response, "chunks", None) - if isinstance(inner_chunks, list): - self.chunks = inner_chunks - # Preserve hidden params (including litellm_overhead_time_ms) from original response - model_response_hidden_params: Final = get_hidden_params(model_response) - if model_response_hidden_params is not None: - self._hidden_params = dict(model_response_hidden_params) - - def __aiter__(self): - return self - - async def __anext__(self): - return await self._async_generator.__anext__() - - async def close_model_response() -> None: - if not hasattr(model_response, "aclose"): - return - try: - await model_response.aclose() - except BaseException as e: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing model_response: %s", - e, - ) - - async def stream_with_fallbacks(): - fallback_response = None # Track for cleanup in finally - try: - async for item in _achunks_with_first(model_response, first_chunk): - yield item - except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: - e: Final = ( - _context_window_stream_fallback_error( - stream_error, - getattr(model_response, "model", ""), - model_response.custom_llm_provider or "", - bool(getattr(model_response, "response_uptil_now", "")) - or _stream_chunks_have_generated_content(model_response.chunks), - ) - if isinstance(stream_error, litellm_exceptions.ContextWindowExceededError) - or isinstance(stream_error.original_exception, litellm_exceptions.ContextWindowExceededError) - else stream_error - ) - with anyio.CancelScope(shield=True): - await close_model_response() - await held_slot.aclose() - committed: Final = bool( - not e.is_pre_first_chunk - and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) - ) - cut_hidden_params: Final = get_hidden_params(model_response) or {} - continue_after_content: Final = committed and self._mid_stream_continuation_eligible( - e=e, request_kwargs=initial_kwargs, cut_deployment_id=cut_hidden_params.get("model_id") - ) - if committed and not continue_after_content: - self._raise_original_mid_stream_error(e) - - from litellm.main import stream_chunk_builder - - complete_response_object: Final = stream_chunk_builder( - chunks=model_response.chunks, messages=messages if continue_after_content else None - ) - complete_response_object_usage: Final = cast( - Usage | None, - getattr(complete_response_object, "usage", None), - ) - try: - # Use the router's fallback system - model_group: Final = cast(str, initial_kwargs.get("model")) - fallbacks: Final[list | None] = initial_kwargs.get("fallbacks", self._core_fallbacks) - context_window_fallbacks: Final[list | None] = initial_kwargs.get( - "context_window_fallbacks", self._core_context_window_fallbacks - ) - content_policy_fallbacks: Final[list | None] = initial_kwargs.get( - "content_policy_fallbacks", self._core_content_policy_fallbacks - ) - initial_kwargs["original_function"] = self._acompletion - continuation_messages: Final = ( - self._build_completion_continuation_input(messages, e.generated_content) - if continue_after_content - else None - ) - continuation_output: Final = ( - ContinuationOutput( - e.generated_content, - continuation_messages[-1], - model_group, - complete_response_object, - checks_post_call_rules=initial_kwargs.get(MID_STREAM_CONTINUATION_KWARG) - is not MID_STREAM_CONTINUATION_MARKER, - ) - if continuation_messages is not None - else None - ) - cut_partial: Final = ( - PythonCore._price_cut_partial(model_response, complete_response_object) - if continuation_output is not None - else None - ) - if cut_partial is not None: - model_response.logging_obj.mark_stash_billed_by_continuation() - if continue_after_content: - emitted_tokens: Final = int( - getattr(complete_response_object_usage, "completion_tokens", 0) or 0 - ) - reduced_ceilings: Final = continuation_output_ceilings(initial_kwargs, emitted_tokens) - if reduced_ceilings is None or continuation_messages is None: - self._raise_original_mid_stream_error(e) - await model_response.wait_for_failure_logging(timeout=MID_STREAM_FAILURE_LOG_WAIT_SECONDS) - initial_kwargs.update(reduced_ceilings) - initial_kwargs["messages"] = continuation_messages - initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = MID_STREAM_CONTINUATION_MARKER - else: - initial_kwargs["messages"] = messages - self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) - fallback_response = await self.async_function_with_fallbacks_common_utils( - e=e.original_exception - if isinstance(e.original_exception, litellm_exceptions.ContextWindowExceededError) - else e, - disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, - ) - - # If fallback returns a streaming response, iterate over it - if hasattr(fallback_response, "__aiter__"): - PythonCore._carry_partial_spend_to_fallback(fallback_response, cut_partial) - prepared_fallback_hidden_params = PythonCore._adopt_fallback_response_headers( - wrapper_ref, fallback_response - ) - fallback_headers_are_settled = False - fallback_items: Final = ( - continuation_output.splice(fallback_response) - if continuation_output is not None - else fallback_response - ) - async for fallback_item in fallback_items: - if not fallback_headers_are_settled: - fallback_headers_are_settled = True - # a fallback that failed over again only repoints itself once it yields - prepared_fallback_hidden_params = PythonCore._adopt_fallback_response_headers( - wrapper_ref, fallback_response - ) - PythonCore._apply_fallback_hidden_params_to_item( - fallback_item, prepared_fallback_hidden_params - ) - if ( - fallback_item - and isinstance(fallback_item, ModelResponseStream) - and hasattr(fallback_item, "usage") - ): - self._combine_fallback_usage( - fallback_item, - complete_response_object_usage if continuation_output is not None else None, - ) - yield fallback_item - else: - # If fallback returns a non-streaming response, yield None - yield None - - except Exception as fallback_error: - # If fallback also fails, log and re-raise original error - verbose_router_logger.error("Fallback also failed: %s", fallback_error) - # No fallback handled the mid-stream error, so surface the - # real provider exception (e.g. RateLimitError) instead of - # leaking the internal MidStreamFallbackError to the client - if ( - isinstance(fallback_error, MidStreamFallbackError) - and fallback_error.original_exception is not None - ): - raise fallback_error.original_exception from fallback_error - raise fallback_error - finally: - # Close the underlying streams to release HTTP connections - # back to the connection pool when the generator is closed - # (e.g. on client disconnect). - # Shield from anyio cancellation so the awaits can complete. - with anyio.CancelScope(shield=True): - await close_model_response() - await held_slot.aclose() - if fallback_response is not None and hasattr(fallback_response, "aclose"): - try: - await fallback_response.aclose() - except BaseException as close_error: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing fallback_response: %s", - close_error, - ) - - wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) - # weak, so the generator closing over it does not keep the wrapper out of - # refcount teardown and delay the `finally` that releases the deployment slot - wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) - return wrapped_response - - @staticmethod - def _extract_partial_responses_usage( - source_iterator: "BaseResponsesAPIStreamingIterator", - ) -> Optional["ResponseAPIUsage"]: - """ - Best-effort: pull partial token usage from a Responses-API streaming - iterator that errored mid-stream, normalized to ResponseAPIUsage so - the caller can combine without crossing token-naming conventions. - - Two sources, in priority order: - 1. The bridge path (LiteLLMCompletionStreamingIterator) accumulates - chat-completion chunks while streaming — feed them through - stream_chunk_builder to recover chat Usage, then translate - (prompt_tokens → input_tokens, completion_tokens → output_tokens). - 2. The native path (ResponsesAPIStreamingIterator) only has a - completed_response object if the stream reached - RESPONSE_COMPLETED before erroring — uncommon mid-stream but - worth checking. Already ResponseAPIUsage-shaped. - - Returns None when no partial usage is recoverable. - """ - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - from litellm.types.llms.openai import ( - ResponseAPIUsage, - ResponseCompletedEvent, - ResponseFailedEvent, - ResponseIncompleteEvent, - ) - - # Bridge subclass is the only iterator that accumulates chat-completion - # chunks. isinstance narrows the type so we can read the attribute - # directly instead of getattr-ing on the base class. - if isinstance(source_iterator, LiteLLMCompletionStreamingIterator): - chunks: Final = source_iterator.collected_chat_completion_chunks - if chunks: - try: - from litellm.main import stream_chunk_builder - - built: Final = stream_chunk_builder(chunks=chunks) - # stream_chunk_builder returns ModelResponse | - # TextCompletionResponse | None. ModelResponse sets .usage - # in __init__ rather than declaring it as a class field, so - # static narrowing doesn't expose it. Mirror the sync path - # (_completion_streaming_iterator) and pull via getattr. - chat: Final[object | None] = getattr(built, "usage", None) if built is not None else None - if chat is not None: - # getattr-with-default because the test path may - # substitute a SimpleNamespace lacking some fields; - # real Usage instances always have them. - prompt: Final = int(getattr(chat, "prompt_tokens", 0) or 0) - completion: Final = int(getattr(chat, "completion_tokens", 0) or 0) - total: Final = int(getattr(chat, "total_tokens", prompt + completion) or (prompt + completion)) - return ResponseAPIUsage( - input_tokens=prompt, - output_tokens=completion, - total_tokens=total, - ) - except Exception: - # Builder is best-effort — fall through to native path. - pass - - # Native path: completed_response is set only if RESPONSE_COMPLETED - # arrived before the error (uncommon mid-stream but worth checking). - # Already ResponseAPIUsage-shaped — return as-is. - completed: Final = source_iterator.completed_response - if isinstance( - completed, - (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent), - ): - return completed.response.usage - return None - - @staticmethod - def _combine_responses_fallback_usage( - fallback_item: "BaseLiteLLMOpenAIResponseObject", - partial_usage: "ResponseAPIUsage", - ) -> None: - """ - Merge partial-stream usage with fallback-stream usage on a - Responses-API streaming event. - - Only mutates events that carry a `response` with a `usage` field - (response.completed / response.failed / response.incomplete). Other - events pass through unchanged. - - Both inputs are ResponseAPIUsage-shaped (see - _extract_partial_responses_usage which normalizes the bridge path), - so we can sum input_tokens / output_tokens / total_tokens directly - and produce a clean ResponseAPIUsage — no token-naming split, no - setattr bypass. - """ - from litellm.types.llms.openai import ( - ResponseAPIUsage, - ResponseCompletedEvent, - ResponseFailedEvent, - ResponseIncompleteEvent, - ) - - if not isinstance( - fallback_item, - (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent), - ): - return - response: Final = fallback_item.response - if response.usage is None: - return - - fb: Final = response.usage - response.usage = ResponseAPIUsage( - input_tokens=(partial_usage.input_tokens or 0) + (fb.input_tokens or 0), - output_tokens=(partial_usage.output_tokens or 0) + (fb.output_tokens or 0), - total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0), - ) - - @staticmethod - def _raise_original_mid_stream_error(e: "MidStreamFallbackError") -> NoReturn: - """Surface the provider error the stream wrapper carried instead of the internal MidStreamFallbackError.""" - if e.original_exception is not None: - raise e.original_exception from e - raise e - - def _active_default_litellm_params(self) -> Mapping[str, object]: - default_params: Final = _REQUEST_KWARGS_ADAPTER.validate_python(self._core_default_litellm_params) - return {key: value for key, value in default_params.items() if value is not None} - - def _params_sent_to_deployment( - self, request_kwargs: Mapping[str, object], deployment: Deployment | None - ) -> Mapping[str, object]: - """The request as the deployment received it, merged in the order _update_kwargs_with_deployment applies.""" - deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} - defaults: Final = self._active_default_litellm_params() - deployment_tool_choice: Final = ( - {"tool_choice": deployment_params["tool_choice"]} - if deployment_params.get("tool_choice") is not None - else {} - ) - return MappingProxyType({**deployment_params, **defaults, **deployment_tool_choice, **request_kwargs}) - - def _names_deployment_that_cannot_continue(self, target: object) -> bool: - name: Final = fallback_target_name(target) - if name is None or name in self._core_model_names or not self._core_has_model_id(name): - return False - deployment: Final = self._core_get_deployment(model_id=name) - return deployment is not None and not deployment_can_continue(deployment.model_dump(exclude_none=True)) - - def _mid_stream_continuation_eligible( - self, - e: "MidStreamFallbackError", - request_kwargs: Mapping[str, object], - cut_deployment_id: object = None, - ) -> bool: - """Whether a stream that broke after plain assistant text may continue via prefill.""" - if not self.enable_mid_stream_fallback_continuation: - return False - if not e.generated_content or e.emitted_disqualifying_content: - return False - if isinstance(e.original_exception, litellm_exceptions.ContentPolicyViolationError): - return False - cut_deployment: Final = ( - self._core_get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None - ) - if isinstance(cut_deployment_id, str) and cut_deployment is None: - return False - if constrains_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)): - return False - requested_group: Final = request_kwargs.get("model") - lookup_groups: Final = fallback_lookup_groups( - request_kwargs, requested_group if isinstance(requested_group, str) else None - ) - targets: Final = continuation_fallback_targets( - request_kwargs.get("fallbacks", self._core_fallbacks), lookup_groups - ) - if any(fallback_target_overrides_request_params(target) for target in targets): - return False - if any(self._names_deployment_that_cannot_continue(target) for target in targets): - return False - return not any( - isinstance(callback, CustomGuardrail) - and callback._event_hook_is_event_type( # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks - GuardrailEventHooks.post_call - ) - for callback in self._host.callback_list("callbacks") - ) - - @staticmethod - def _build_completion_continuation_input( - messages: Sequence[Mapping[str, object]], - generated_content: str, - ) -> Sequence[Mapping[str, object]] | None: - """Append the partial output as an assistant prefill, or extend a trailing `prefix: True` assistant turn in - place so the request never ends in two assistant messages. None when the request ends in a plain assistant - turn, in a prefill turn whose content has a shape this cannot extend, or when the cut came mid-echo of the - prefill.""" - last: Final = messages[-1] if messages else None - if last is None or last.get("role") != "assistant": - return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}] - if last.get("prefix") is not True: - return None - new_text: Final = continued_text(last, generated_content) - if new_text is None: - return None - content: Final = last.get("content") - if content is None or isinstance(content, str): - return [*messages[:-1], {**last, "content": f"{content or ''}{new_text}", "prefix": True}] - try: - blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content) - except ValidationError: - return None - extended: Final = _extend_trailing_text_block(blocks, new_text) - return [*messages[:-1], {**last, "content": extended, "prefix": True}] - - async def _aresponses_streaming_iterator( - self, - response: "BaseResponsesAPIStreamingIterator", - initial_kwargs: dict[str, object], - ) -> "BaseResponsesAPIStreamingIterator": - """ - Wrap a Responses-API streaming iterator so MidStreamFallbackError - triggers the Router's fallback chain (parity with - _acompletion_streaming_iterator for the chat-completions path). - - The Responses-API streaming path goes through - _ageneric_api_call_with_fallbacks rather than _acompletion, so the - returned iterator is never wrapped by the chat completions - fallback handler. Without this wrapper, MidStreamFallbackError - raised mid-stream from the underlying CustomStreamWrapper (used by - LiteLLMCompletionStreamingIterator when the Responses API is - served via the completion bridge) propagates unhandled and the - configured cross-provider fallback never fires. - - Full parity with the chat-completions path: - - Pre-first-chunk: retry with the original input unchanged. - - Partial content: continue from the generated text, as an - assistant prefill on a fallback that takes one and as assistant - history plus a user request to continue on any other. A request - whose output is constrained (a JSON schema, a forced tool), or a - stream that began a tool call, is not continued, and the primary's - error surfaces instead. - - Usage combining: merge partial-stream usage onto the fallback's - response.completed event so accounting reflects both attempts. - - Stream cleanup: shielded aclose() on both source and fallback - iterators on terminate. - """ - from litellm.exceptions import MidStreamFallbackError - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - get_openai_response_types, - ) - - source_iterator: Final = response - - # Pre-resolve the set of terminal stream event types so the - # per-chunk type check inside FallbackResponsesStreamWrapper - # stays cheap; mirrors the source-iterator filter at - # responses/streaming_iterator.py:243-247. - _openai_types: Final = get_openai_response_types() - _RESPONSES_TERMINAL_EVENT_TYPES: Final = ( - _openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - _openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, - _openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, - ) - - class FallbackResponsesStreamWrapper(BaseResponsesAPIStreamingIterator): - """ - Subclasses BaseResponsesAPIStreamingIterator only for isinstance - compatibility (proxy + interactions code paths check the type). - Bypasses the parent constructor and delegates iteration to an - async generator. - """ - - fallback_headers_adopted: bool = False - - def __init__(self, async_generator: "AsyncGenerator[ResponsesAPIStreamingResponse, None]"): - import time - from datetime import datetime - - self._async_generator = async_generator - self._billing_source: object = source_iterator - self._last_sequence_number = -1 - # Mirror every attribute BaseResponsesAPIStreamingIterator.__init__ - # would have set. The wrapper bypasses super().__init__ (it has no - # httpx.Response of its own and no provider config to drive), so - # we copy from source_iterator where applicable and use safe - # defaults elsewhere. This keeps inherited methods (e.g. - # _check_max_streaming_duration, _handle_failure) safe to call. - # - # The bridge path (LiteLLMCompletionStreamingIterator used by - # Anthropic/Bedrock/Vertex) does not call super().__init__ and - # is missing many of these attributes — use getattr fallbacks - # so wrapper construction never raises AttributeError. The - # bridge stores the logging object as `litellm_logging_obj`. - # base class declares non-Optional types for these - # fields but the bridge path (LiteLLMCompletionStreamingIterator) - # can legitimately omit them at runtime — keep the None - # fallback. Same lines passed mypy on the pre-fix file - # because the surrounding function body wasn't fully - # type-narrowed; the new typed terminal-event tuple above - # is what made these surface. - self.response = getattr(source_iterator, "response", None) - self.model = getattr(source_iterator, "model", None) - self.logging_obj = getattr( - source_iterator, - "logging_obj", - getattr(source_iterator, "litellm_logging_obj", None), - ) - self.finished = False - self.responses_api_provider_config = getattr(source_iterator, "responses_api_provider_config", None) - self.completed_response = None - self.start_time = getattr(source_iterator, "start_time", datetime.now()) - self._failure_handled = False - self._yielded_first_chunk = False - self._generated_content = "" - self._completed_response_cached = False - self._completed_response_logged = False - self._completed_response_cache_hit = None - self._persist_completed_response_before_logging = True - self._stream_created_time = time.time() - self.litellm_metadata = getattr(source_iterator, "litellm_metadata", None) - self.custom_llm_provider = getattr(source_iterator, "custom_llm_provider", None) - self.request_data = getattr(source_iterator, "request_data", {}) or {} - self.call_type = getattr(source_iterator, "call_type", None) - # Preserve hidden params so response headers (model_id, - # api_base, additional_headers) keep flowing. - self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) - - def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]: - prepared: Final = PythonCore._prepare_fallback_hidden_params(fallback_response) - self._billing_source = fallback_response - self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} - self.fallback_headers_adopted = True - return prepared - - def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": - """ - The event to bill on a client disconnect. After a mid-stream hop the primary - stream's relayed tokens are folded into the fallback's event, the way - _combine_responses_fallback_usage does for a natural end of stream. - """ - from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect - - current: Final = partial_responses_incomplete_event_on_disconnect(self._billing_source) - if self._billing_source is source_iterator: - return current - primary: Final = partial_responses_incomplete_event_on_disconnect(source_iterator) - if current is None or primary is None: - return current or primary - if primary.response.usage is not None: - PythonCore._combine_responses_fallback_usage(current, primary.response.usage) - return current - - def __aiter__(self): - return self - - @property - def yielded_first_chunk(self) -> bool: - return self._yielded_first_chunk - - async def __anext__(self): - try: - chunk: Final = await self._async_generator.__anext__() - except StopAsyncIteration: - # The inner generator is exhausted. If we never sniffed a - # terminal event off a chunk (the bridge path emits the - # final response.completed via common_done_event_logic, - # which raises StopAsyncIteration after returning it), - # fall back to whatever the source iterator latched so - # the proxy's container-ownership hook still sees a - # completed_response instead of logging a spurious - # "no completed_response" warning. - if self.completed_response is None: - self.completed_response = getattr(source_iterator, "completed_response", None) - raise - # Sniff the terminal stream event off each forwarded chunk - # so ``self.completed_response`` is populated regardless of - # which inner iterator produced it (source_iterator, - # fallback_iterator, or any future wrapper). Without this - # the proxy's container-ownership hook (which reads - # ``getattr(stream_response, "completed_response", None)`` - # via extract_completed_responses_response) silently - # records nothing on streaming /v1/responses calls — every - # follow-up /v1/containers//files call then 403s for - # the very key that created the container (#30210). - if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: - self.completed_response = chunk - self._yielded_first_chunk = True - sequence_number: Final = getattr(chunk, "sequence_number", None) - if not isinstance(sequence_number, int): - return chunk - if not self.fallback_headers_adopted: - self._last_sequence_number = max(self._last_sequence_number, sequence_number) - return chunk - self._last_sequence_number = max(self._last_sequence_number + 1, sequence_number) - if self._last_sequence_number == sequence_number: - return chunk - return chunk.model_copy(update={"sequence_number": self._last_sequence_number}) - - async def aclose(self): - # async generators always expose aclose — no defensive check needed. - await self._async_generator.aclose() - - async def stream_with_fallbacks(): - held_lifecycle_events: tuple[object, ...] = () # rebind-ok: flushed at first output, dropped on fallback - response_id_exposed = False # rebind-ok: set once, before an early-released lifecycle event is yielded - try: - async for item in cast( # cast-ok: response iterators are async streams - AsyncIterator[object], source_iterator - ): - if not response_id_exposed and self._aresponses_background_stream_can_release( - item, source_iterator, initial_kwargs - ): - response_id_exposed = True - if not response_id_exposed and _responses_stream_holds_event(item, len(held_lifecycle_events)): - held_lifecycle_events = (*held_lifecycle_events, item) - continue - for held_event in held_lifecycle_events: - yield held_event - held_lifecycle_events = () - yield item - for held_event in held_lifecycle_events: - yield held_event - except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: - e: Final = ( - _context_window_stream_fallback_error( - stream_error, - source_iterator.model or "", - source_iterator.custom_llm_provider or "", - wrapper.yielded_first_chunk, - ) - if isinstance(stream_error, litellm_exceptions.ContextWindowExceededError) - or isinstance(stream_error.original_exception, litellm_exceptions.ContextWindowExceededError) - else stream_error - ) - if response_id_exposed: - self._raise_released_stream_error(e, initial_kwargs) - async with contextlib.aclosing( - self._aresponses_fallback_attempt( - e, source_iterator, initial_kwargs, wrapper.adopt_fallback_headers, held_lifecycle_events - ) - ) as fallback_stream: - async for fallback_item in fallback_stream: - yield fallback_item - except Exception: - for held_event in held_lifecycle_events: - yield held_event - raise - finally: - with anyio.CancelScope(shield=True): - if hasattr(source_iterator, "aclose"): - try: - await source_iterator.aclose() - except Exception as exc: - verbose_router_logger.debug( - "stream_with_fallbacks(aresponses): error closing source: %s", - exc, - ) - - wrapper: Final = FallbackResponsesStreamWrapper( - stream_with_fallbacks() # pyright: ignore[reportArgumentType] # relays response events typed as object - ) - return wrapper - - def _aresponses_background_stream_can_release( - self, - item: object, - source_iterator: "BaseResponsesAPIStreamingIterator", - kwargs: Mapping[str, object], - ) -> bool: - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - - if ( - type(source_iterator) is not ResponsesAPIStreamingIterator - or source_iterator.custom_llm_provider != "openai" - ): - return False - if not _responses_event_announces_background_run(item): - return False - model_group: Final = kwargs.get("model") - return isinstance(model_group, str) and not self._stream_can_fall_back(model_group, kwargs) - - def _raise_released_stream_error( - self, e: "MidStreamFallbackError", initial_kwargs: Mapping[str, object] - ) -> NoReturn: - original: Final = e.original_exception - model_group: Final = initial_kwargs.get("model") - if ( - self._host.dynamic().expose_router_debug_in_errors - and isinstance(original, litellm_exceptions.ContentPolicyViolationError) - and isinstance(model_group, str) - and not fallbacks_disabled_for_request(initial_kwargs) - and _mid_stream_error_before_output(e) - ): - hop_depth: Final = initial_kwargs.get("fallback_depth") - original.message += format_content_policy_no_fallback_note( - model_group, - fallback_lookup_groups(initial_kwargs, model_group), - _request_fallback_list( - initial_kwargs, - "content_policy_fallbacks", - cast( # cast-ok: untyped router attribute - "list[object] | None", self._core_content_policy_fallbacks - ), - ), - _request_fallback_list( - initial_kwargs, - "fallbacks", - cast("list[object] | None", self._core_fallbacks), # cast-ok: untyped router attribute - ), - isinstance(hop_depth, int) and hop_depth > 0, - ) - PythonCore._raise_original_mid_stream_error(e) - - def _responses_continuation( - self, - e: "MidStreamFallbackError", - source_iterator: "BaseResponsesAPIStreamingIterator", - request_kwargs: Mapping[str, object], - ) -> ResponsesContinuation | None: - """What a fallback continues from once a Responses stream broke after text, raising the primary's error - when the request constrains its output or the stream began a tool call, since a continued JSON document or - tool call would not parse, and when a content filter cut it, so no other model finishes filtered text.""" - if _mid_stream_error_before_output(e): - return None - if source_iterator.started_tool_call or isinstance( - e.original_exception, litellm_exceptions.ContentPolicyViolationError - ): - PythonCore._raise_original_mid_stream_error(e) - cut_deployment_id: Final = (get_hidden_params(source_iterator) or {}).get("model_id") - cut_deployment: Final = ( - self._core_get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None - ) - if constrains_responses_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)): - PythonCore._raise_original_mid_stream_error(e) - earlier: Final = request_kwargs.get(RESPONSES_CONTINUATION_KWARG) - if isinstance(earlier, ResponsesContinuation): - return earlier.extended(e.generated_content) - return ResponsesContinuation(request_kwargs.get("input"), e.generated_content) - - @staticmethod - def _responses_continuation_ceilings( - e: "MidStreamFallbackError", request_kwargs: Mapping[str, object] - ) -> Mapping[str, object]: - """The caller's output ceilings less what the cut stream already emitted, raising the primary's error once - nothing is left.""" - emitted_tokens: Final = token_counter(model=e.model or "", text=e.generated_content) - reduced: Final = continuation_output_ceilings(request_kwargs, emitted_tokens) - if reduced is None: - PythonCore._raise_original_mid_stream_error(e) - return reduced - - async def _aresponses_fallback_attempt( - self, - e: "MidStreamFallbackError", - source_iterator: "BaseResponsesAPIStreamingIterator", - initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain - adopt_headers: Callable[[object], tuple[dict[str, object], dict[str, object]]], # mutable-ok: hidden params - held_lifecycle_events: tuple[object, ...], - ) -> AsyncGenerator[object, None]: - """ - Re-enters the Router's fallback chain for a mid-stream Responses API error and yields - whatever the fallback attempt produces. The lifecycle events the primary stream held - back reach the client only when no fallback lands, so the client sees exactly one - response announced, the one whose id completes. Split out of - _aresponses_streaming_iterator to keep each function's cyclomatic complexity within - the repo's C901 budget. - """ - from litellm.exceptions import MidStreamFallbackError - - continuation: Final = self._responses_continuation(e, source_iterator, initial_kwargs) - reduced_ceilings: Final[Mapping[str, object]] = ( - PythonCore._responses_continuation_ceilings(e, initial_kwargs) if continuation is not None else {} - ) - seam: Final = ResponsesContinuationSeam(continuation.emitted_text if continuation is not None else "") - partial_usage: Final = PythonCore._extract_partial_responses_usage(source_iterator) - fallback_response = None # rebind-ok: pre-init so finally can close it if a fallback was actually attempted - fallback_yielded = False # rebind-ok: flipped on the first fallback item so a fallback that dies before its first event still replays the primary's held announcement - try: - model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: model group - fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the common_utils list|None param - "fallbacks", self._core_fallbacks - ) - context_window_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below - "context_window_fallbacks", self._core_context_window_fallbacks - ) - content_policy_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below - "content_policy_fallbacks", self._core_content_policy_fallbacks - ) - initial_kwargs["original_function"] = ( # rebind-ok: the fallback chain re-enters on the same kwargs - self._ageneric_api_call_with_fallbacks_responses_attempt - ) - if continuation is not None: - initial_kwargs.update(reduced_ceilings) - initial_kwargs[RESPONSES_CONTINUATION_KWARG] = continuation # rebind-ok: shaped per fallback deployment - initial_kwargs["input"] = responses_continuation_input(continuation) # rebind-ok: fallback hop input - # The Responses-API path stores observability metadata - # under "litellm_metadata" (not the default "metadata") — - # see _ageneric_api_call_with_fallbacks. Mirroring that - # here ensures model_group, model_group_alias, and trace - # ids land in the same key litellm.aresponses reads from. - self._update_kwargs_before_fallbacks( - model=model_group, - kwargs=initial_kwargs, - metadata_variable_name="litellm_metadata", - ) - # The content-policy dispatch branch matches on the trigger's own type, so a refusal's - # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. - fallback_trigger: Final[Exception] = ( - e.original_exception - if isinstance( - e.original_exception, - (litellm_exceptions.ContentPolicyViolationError, litellm_exceptions.ContextWindowExceededError), - ) - else e - ) - fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success - e=fallback_trigger, - disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, - ) - prepared_fallback_hidden_params: Final = adopt_headers(fallback_response) - if hasattr(fallback_response, "__aiter__"): - async for fallback_item in seam.splice(fallback_response): - PythonCore._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) - if partial_usage is not None: - PythonCore._combine_responses_fallback_usage(fallback_item, partial_usage) - fallback_yielded = True - yield fallback_item - else: - fallback_yielded = True # rebind-ok: see the pre-init above - yield fallback_response - except Exception as fallback_error: - verbose_router_logger.error("Responses streaming fallback also failed: %s", fallback_error) - if not fallback_yielded: - for held_event in held_lifecycle_events: - yield held_event - if isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None: - raise fallback_error.original_exception from fallback_error - raise - finally: - if fallback_response is not None and hasattr(fallback_response, "aclose"): - with anyio.CancelScope(shield=True): - try: - await fallback_response.aclose() - except Exception as exc: - verbose_router_logger.debug( - "stream_with_fallbacks(aresponses): error closing fallback: %s", - exc, - ) - - def _completion_streaming_iterator( - self, - model_response: CustomStreamWrapper, - messages: list[dict[str, str]], - initial_kwargs: dict, - first_chunk: _FirstChunk[Iterator[object]] | None = None, - ) -> CustomStreamWrapper: - """ - Sync equivalent of _acompletion_streaming_iterator. - - Wraps a sync streaming response so that MidStreamFallbackError - (raised by CustomStreamWrapper.__next__) triggers the Router's - fallback chain instead of surfacing directly to the caller. - """ - from litellm.exceptions import MidStreamFallbackError - - class SyncFallbackStreamWrapper(FallbackAwareStreamWrapper): - def __init__(self, sync_generator: Generator): - super().__init__( - completion_stream=sync_generator, - model=model_response.model, - custom_llm_provider=model_response.custom_llm_provider, - logging_obj=model_response.logging_obj, - _response_headers=getattr(model_response, "_response_headers", None), - ) - self._sync_generator = sync_generator - model_response_hidden_params: Final = get_hidden_params(model_response) - if model_response_hidden_params is not None: - self._hidden_params = dict(model_response_hidden_params) - - def __iter__(self): - return self - - def __next__(self): - return next(self._sync_generator) - - router_self: Final = self - - def stream_with_fallbacks(): - fallback_response = None - try: - for item in _chunks_with_first(model_response, first_chunk): - yield item - except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: - e: Final = ( - _context_window_stream_fallback_error( - stream_error, - getattr(model_response, "model", ""), - model_response.custom_llm_provider or "", - bool(getattr(model_response, "response_uptil_now", "")) - or _stream_chunks_have_generated_content(model_response.chunks), - ) - if isinstance(stream_error, litellm_exceptions.ContextWindowExceededError) - or isinstance(stream_error.original_exception, litellm_exceptions.ContextWindowExceededError) - else stream_error - ) - if fallbacks_disabled_for_request(initial_kwargs) or ( - not e.is_pre_first_chunk - and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) - ): - if e.original_exception is not None: - raise e.original_exception from e - raise - - try: - model_group: Final = cast(str, initial_kwargs.get("model")) - fallbacks: Final[list | None] = initial_kwargs.get("fallbacks", router_self._core_fallbacks) - context_window_fallbacks: Final[list | None] = initial_kwargs.get( - "context_window_fallbacks", - router_self._core_context_window_fallbacks, - ) - content_policy_fallbacks: Final[list | None] = initial_kwargs.get( - "content_policy_fallbacks", - router_self._core_content_policy_fallbacks, - ) - initial_kwargs["original_function"] = router_self._completion - initial_kwargs["messages"] = messages - router_self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) - fallback_response = run_async_function( - router_self.async_function_with_fallbacks_common_utils, - e=e.original_exception - if isinstance(e.original_exception, litellm_exceptions.ContextWindowExceededError) - else e, - disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, - ) - - if hasattr(fallback_response, "__iter__"): - prepared_fallback_hidden_params = PythonCore._adopt_fallback_response_headers( - wrapper_ref, fallback_response - ) - fallback_headers_are_settled = False - for fallback_item in cast(Iterable[object], fallback_response): # cast-ok: __iter__ was checked - if not fallback_headers_are_settled: - fallback_headers_are_settled = True - # a fallback that failed over again only repoints itself once it yields - prepared_fallback_hidden_params = PythonCore._adopt_fallback_response_headers( - wrapper_ref, fallback_response - ) - PythonCore._apply_fallback_hidden_params_to_item( - fallback_item, prepared_fallback_hidden_params - ) - if ( - fallback_item - and isinstance(fallback_item, ModelResponseStream) - and hasattr(fallback_item, "usage") - ): - router_self._combine_fallback_usage(fallback_item, None) - yield fallback_item - else: - yield None - - except Exception as fallback_error: - verbose_router_logger.error("Fallback also failed: %s", fallback_error) - if ( - isinstance(fallback_error, MidStreamFallbackError) - and fallback_error.original_exception is not None - ): - raise fallback_error.original_exception from fallback_error - raise fallback_error - finally: - if hasattr(model_response, "close"): - try: - model_response.close() - except BaseException as close_err: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing model_response: %s", - close_err, - ) - if fallback_response is not None and hasattr(fallback_response, "close"): - try: - fallback_response.close() - except BaseException as close_err: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing fallback_response: %s", - close_err, - ) - - wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks()) - # weak, for the same reason as the async twin - wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) - return wrapped_response - - async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): - """ - Run a silent experiment in the background. - """ - try: - # Prevent infinite recursion if silent model also has a silent model - if kwargs.get("metadata", {}).get("is_silent_experiment", False): - return - - messages = copy.deepcopy(messages) - - verbose_router_logger.info("Starting silent experiment for model %s", silent_model) - - silent_kwargs: Final = self._get_silent_experiment_kwargs(**kwargs) - # Override model_group to correctly attribute metrics to the silent model - silent_kwargs["metadata"]["model_group"] = silent_model - - # Trigger the silent request - await self._run_silent_experiment(silent_model, messages, silent_kwargs) - except Exception as e: - verbose_router_logger.error("Silent experiment failed for model %s: %s", silent_model, e) - - async def _acompletion( - self, model: str, messages: list[dict[str, str]], **kwargs - ) -> ModelResponse | CustomStreamWrapper: - """ - - Get an available deployment - - call it with a semaphore over the call - - semaphore specific to it's rpm - - in the semaphore, make a check against it's local rpm before running - """ - model_name = None - deployment = None - _timeout_debug_deployment_dict = {} # this is a temporary dict to debug timeout issues - try: - controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None) - input_kwargs_for_streaming_fallback: Final = mid_stream_fallback_snapshot_kwargs( - model=model, controls=controls, kwargs=kwargs - ) - - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - start_time: Final = time.time() - deployment = await self._core_async_get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._drop_unsupported_classifier_reasoning_effort( - deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment - model=model, - kwargs=kwargs, - ) - - _timeout_debug_deployment_dict = deployment - end_time: Final = time.time() - _duration: Final = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.ROUTER, - duration=_duration, - call_type="async_get_available_deployment", - start_time=start_time, - end_time=end_time, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) - ) - - # debug how often this deployment picked - - self._track_deployment_metrics(deployment=deployment, parent_otel_span=parent_otel_span) - - # Check for silent model experiment - # Make a local copy of litellm_params to avoid mutating the Router's state - litellm_params: Final = self._deployment_params_with_request_reasoning_override( - deployment["litellm_params"], kwargs - ) - silent_model: Final = litellm_params.pop("silent_model", None) - - for silent_target in _silent_experiment_targets(silent_model): - # Mirroring traffic to a secondary model - # This is a silent experiment, so we don't want to block the primary request - asyncio.create_task( - self._silent_experiment_acompletion( - silent_model=silent_target, - messages=messages, # Use messages instead of *args - **_silent_experiment_kwargs_snapshot(kwargs), - ) - ) - - kwargs.setdefault("messages", messages) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) - kwargs.pop("silent_model", None) # Ensure it's not in kwargs either - - model_name = litellm_params["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - input_kwargs: Final = without_router_only_kwargs( - { - **litellm_params, - "messages": messages, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - logging_obj: Final[LiteLLMLogging | None] = kwargs.get("litellm_logging_obj", None) - - max_parallel_requests_limit: Final = self._get_client( - deployment=deployment, - kwargs=kwargs, - client_type="max_parallel_requests", - ) - compacted_input: Final = await compact_to_fit(self._core_ctx, deployment, input_kwargs, "chat") - async with contextlib.AsyncExitStack() as deployment_slot: - if isinstance(max_parallel_requests_limit, MaxParallelRequestsLimit): - deployment_slot.enter_context(max_parallel_requests_limit) - await self.async_routing_strategy_pre_call_checks( - deployment=deployment, - logging_obj=logging_obj, - parent_otel_span=parent_otel_span, - ) - response = await self._host.invoke("acompletion")(**compacted_input) - - ## CHECK CONTENT FILTER ERROR ## - if isinstance(response, ModelResponse): - _should_raise = self._should_raise_content_policy_error( - model=model, response=response, kwargs=kwargs - ) - if _should_raise: - raise litellm_exceptions.ContentPolicyViolationError( - message="Response output was blocked.", - model=model, - llm_provider="", - ) - - if ( - isinstance(response, CustomStreamWrapper) - and response.completion_stream is None - and response.make_call is not None - ): - await response.fetch_stream() - first_chunk: Final = ( - await _aread_first_chunk(response) if isinstance(response, CustomStreamWrapper) else None - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - # debug how often this deployment picked - self._track_deployment_metrics( - deployment=deployment, - response=response, - parent_otel_span=parent_otel_span, - ) - - if isinstance(response, CustomStreamWrapper): - return await self._acompletion_streaming_iterator( - model_response=response, - messages=messages, - initial_kwargs=input_kwargs_for_streaming_fallback, - deployment_slot=deployment_slot.pop_all(), - first_chunk=first_chunk, - ) - - return response - except litellm_exceptions.Timeout as e: - deployment_litellm_params: Final[Mapping[str, object]] = _timeout_debug_deployment_dict.get( - "litellm_params", {} - ) - if self._host.dynamic().expose_router_debug_in_errors: - e.message += ( - f"\n\nDeployment Info: request_timeout: {deployment_litellm_params.get('request_timeout')}" - f"\ntimeout: {deployment_litellm_params.get('timeout')}" - f"\nstream_timeout: {deployment_litellm_params.get('stream_timeout')}" - ) - # Set per-deployment num_retries on exception for retry logic - if deployment is not None: - self._set_deployment_num_retries_on_exception(e, deployment) - self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) - raise e - except Exception as e: - verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - # Set per-deployment num_retries on exception for retry logic - if deployment is not None: - self._set_deployment_num_retries_on_exception(e, deployment) - self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) - raise e - - def _update_kwargs_before_fallbacks( - self, - model: str, - kwargs: dict, - metadata_variable_name: str | None = "metadata", - ) -> None: - """ - Adds/updates to kwargs: - - litellm_trace_id - - metadata - - num_retries is deliberately left as the caller passed it, with no default filled in. - async_function_with_retries resolves the router/global default itself, and it can only - honour the precedence request > deployment litellm_params > litellm_settings while an - absent num_retries still means "the caller did not ask for one". - """ - kwargs.setdefault("litellm_trace_id", str(uuid.uuid4())) - model_group_alias: str | None = None - if self._core_get_model_from_alias(model=model): - model_group_alias = model - kwargs.setdefault(metadata_variable_name, {}).update( - {"model_group": model, "model_group_alias": model_group_alias} - ) - - def _set_deployment_num_retries_on_exception(self, exception: Exception, deployment: dict) -> None: - """ - Set num_retries from deployment litellm_params on the exception. - - This allows the retry logic in async_function_with_retries to use - per-deployment retry settings instead of the global setting. - """ - # Only set if exception doesn't already have num_retries - if hasattr(exception, "num_retries") and exception.num_retries is not None: - return - - litellm_params: Final = deployment.get("litellm_params", {}) - dep_num_retries: Final = litellm_params.get("num_retries") - if dep_num_retries is not None: - try: - exception.num_retries = int(dep_num_retries) # Handle both int and str - except (ValueError, TypeError): - pass # Skip if value can't be converted to int - - def _set_failed_deployment_id_on_exception(self, exception: Exception, deployment: Mapping[str, Any]) -> None: - """ - Stamp the failed deployment's `model_info.id` on the exception so the - fallback layer can exclude it from subsequent re-picks within the same - request (used by weighted-routing failover). - - Idempotent: never overwrites an existing value, so the id of the - deployment that *first* failed in a chain is preserved if multiple - layers re-raise. - """ - if getattr(exception, "failed_deployment_id", None): - return - deployment_id: Final = (deployment.get("model_info") or {}).get("id") - if deployment_id: - try: - exception.failed_deployment_id = deployment_id - except Exception: - pass - - def _stamp_failed_deployment_id_with_effective_model_info( - self, exception: Exception, deployment: Mapping[str, object], kwargs: Mapping[str, object] - ) -> None: - # A client-side-credential call gets a dynamic deployment id generated inside - # _update_kwargs_with_deployment and stamped into kwargs["model_info"]; stamping - # the static shared deployment's id instead would let one tenant's bad credentials - # cool down the deployment every other tenant sharing this config relies on. - if getattr(exception, "failed_deployment_id", None): - return - effective_model_info: Final = kwargs.get("model_info") or deployment.get("model_info") or MappingProxyType({}) - self._set_failed_deployment_id_on_exception(exception, MappingProxyType({"model_info": effective_model_info})) - configured_deployment_id: Final = _MODEL_INFO_ADAPTER.validate_python(effective_model_info).get( - "original_model_id" - ) - if isinstance(configured_deployment_id, str) and configured_deployment_id: - exception.failed_configured_deployment_id = configured_deployment_id # pyright: ignore[reportAttributeAccessIssue] # rebind-ok: the exception carries the failure to _maybe_run_weighted_failover - - def _stamp_failed_attempt(self, exception: Exception, kwargs: Mapping[str, object], *, picked: bool) -> None: - if picked: - self._stamp_failed_deployment_id_with_effective_model_info(exception, MappingProxyType({}), kwargs) - - @staticmethod - def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None: - effective_model_info: Final = kwargs.get("model_info") - if not isinstance(effective_model_info, Mapping): - return - model_info: Final = _MODEL_INFO_ADAPTER.validate_python(effective_model_info) - configured_deployment_id: Final = ( - model_info.get("original_model_id") - if isinstance(exception, litellm_exceptions.UnsupportedEndpointError) - else None - ) - deployment_id: Final = configured_deployment_id or model_info.get("id") - if isinstance(deployment_id, str) and deployment_id: - exception.retry_skip_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # dynamic stamp, read by _deployment_ids_to_skip_on_retry - - def _update_kwargs_with_default_litellm_params( - self, kwargs: dict, metadata_variable_name: str | None = "metadata" - ) -> None: - """ - Adds default litellm params to kwargs, if set. - - Handles inserting this as either "metadata" or "litellm_metadata" depending on the metadata_variable_name - """ - # 1) copy your defaults and pull out metadata - defaults: Final = dict(without_retrieve_stream_defaults(self._core_default_litellm_params, kwargs)) - metadata_defaults: Final = defaults.pop("metadata", {}) or {} - - # 2) add any non-metadata defaults that aren't already in kwargs - for key, value in defaults.items(): - if value is None: - continue - kwargs.setdefault(key, value) - - # 3) merge in metadata, this handles inserting this as either "metadata" or "litellm_metadata" - kwargs.setdefault(metadata_variable_name, {}).update(metadata_defaults) - - def _handle_clientside_credential( - self, deployment: dict, kwargs: dict, function_name: str | None = None - ) -> Deployment: - """ - Build a per-request Deployment carrying the caller-supplied api_key/api_base, - with its own stable id for cooldown, logging, and cost-map identity. - - This deployment is deliberately never registered with the router (no - upsert_deployment/add_deployment call): doing so used to add it to - self._core_model_list under the shared model_name, which made a request-scoped, - caller-supplied provider credential a permanent, load-balanced deployment - that every other caller of that model group could be routed onto. Its - pricing is still registered directly, so a custom price configured on the - underlying deployment still applies to this call. - """ - model_info: Final = deployment.get("model_info", {}).copy() - litellm_params: Final = deployment["litellm_params"].copy() - dynamic_litellm_params: Final = get_dynamic_litellm_params(litellm_params=litellm_params, request_kwargs=kwargs) - # Use deployment model_name as model_group for generating model_id - metadata_variable_name: Final = get_router_metadata_variable_name( - function_name=function_name, - ) - model_group: Final = kwargs.get(metadata_variable_name, {}).get("model_group") - _model_id: Final = self._core_generate_model_id(model_group=model_group, litellm_params=dynamic_litellm_params) - original_model_id: Final = model_info.get("id") - model_info["id"] = _model_id - model_info["original_model_id"] = original_model_id - deployment_pydantic_obj: Final = Deployment( - model_name=model_group, - litellm_params=LiteLLM_Params.model_validate(dynamic_litellm_params), - model_info=model_info, - ) - PythonCore._register_deployment_pricing(deployment=deployment_pydantic_obj) - return deployment_pydantic_obj - - @staticmethod - def _merge_tools_from_deployment( - deployment: dict, kwargs: dict, defaults: Mapping[str, object] = _NO_DEFAULT_PARAMS - ) -> None: - """ - Merge tools from deployment litellm_params with request kwargs. - When both have tools, concatenate them (deployment tools first, then request tools). - tool_choice: use request value if provided, else deployment's. - """ - dep_params_raw: Final = deployment.get("litellm_params", {}) or {} - if isinstance(dep_params_raw, dict): - dep_params = dep_params_raw - else: - dep_params = dep_params_raw.model_dump(exclude_none=True) - if _drops_malformed_request_tools(dep_params, defaults, kwargs): - for name in TOOL_FIELDS: - kwargs.pop(name, None) - merged_tools: Final = _merged_tools(dep_params.get("tools") or [], kwargs.get("tools") or []) - if merged_tools is not None: - kwargs["tools"] = merged_tools - if "tool_choice" not in kwargs and dep_params.get("tool_choice") is not None: - kwargs["tool_choice"] = dep_params["tool_choice"] - - def _update_kwargs_with_deployment( - self, - deployment: dict, - kwargs: dict, - function_name: str | None = None, - splits_stream_timeout: bool = False, - opens_responses_socket: bool = False, - ) -> None: - """ - 5 jobs: - - Adds selected deployment, model_info and api_base to kwargs["metadata"] (used for logging) - - Adds default litellm params to kwargs, if set. - - Merges tools from deployment with request (proxy-configured tools + request tools). - - On a fallback hop, drops the encrypted reasoning this deployment cannot decrypt, keeping its summary. - - On a Responses continuation, sends the partial response as prefill where this deployment takes one. - """ - for key in self._forwarded_alias_marker_keys_the_deployment_sets( - deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) - ): - kwargs.pop(key, None) - kwargs.pop("daily_reset_time", None) # pyright: ignore[reportUnknownMemberType] # kwargs is an untyped dict - defaults: Final = _LITELLM_PARAMS_ADAPTER.validate_python(self._core_default_litellm_params) - self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=defaults) - - model_info = deployment.get("model_info", {}).copy() - deployment_litellm_model_name = deployment["litellm_params"]["model"] - deployment_api_base = deployment["litellm_params"].get("api_base") - deployment_model_name: Final = deployment["model_name"] - from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils - - request_context: Final = cast( # cast-ok: the router's legacy request dict has string keys - Mapping[str, object], kwargs - ) - raw_deployment_params: Final[object] = deployment["litellm_params"] - deployment_context: Final = _LITELLM_PARAMS_ADAPTER.validate_python( - raw_deployment_params.model_dump(exclude_none=True) - if isinstance(raw_deployment_params, LiteLLM_Params) - else raw_deployment_params - ) - _drop_malformed_request_params_the_deployment_sets( - kwargs, # pyright: ignore[reportUnknownArgumentType] # kwargs is the router's untyped request dict - deployment_context, - defaults, - ) - declared_provider: Final = deployment_context.get("custom_llm_provider") - declared_model: Final = deployment_context.get("model") - deployment_provider: Final = ( - declared_provider - if isinstance(declared_provider, str) - else inferred_provider(declared_model if isinstance(declared_model, str) else None) - ) - forwarded_authorization: Final = ( - ProviderSpecificHeaderUtils.forwarded_authorization( - request_context.get("provider_specific_header"), deployment_provider - ) - if deployment_provider in _FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS - else None - ) - if is_clientside_credential(request_kwargs=kwargs) or forwarded_authorization is not None: - deployment_pydantic_obj: Final = self._handle_clientside_credential( - deployment=deployment, - kwargs={ - **kwargs, - **({"api_key": forwarded_authorization} if forwarded_authorization is not None else {}), - }, - function_name=function_name, - ) - model_info = deployment_pydantic_obj.model_info.model_dump() - deployment_litellm_model_name = deployment_pydantic_obj.litellm_params.model - deployment_api_base = deployment_pydantic_obj.litellm_params.api_base - - metadata_variable_name: Final = get_router_metadata_variable_name( - function_name=function_name, - ) - deployment_that_just_failed: Final = _deployment_that_just_failed( - _FALLBACK_HOP_ADAPTER.validate_python(kwargs).get(metadata_variable_name) - ) - routing: Final = deployment.get(KUBERNETES_POD_ROUTING_KEY) - - kwargs.setdefault(metadata_variable_name, {}).update( - { - "deployment": deployment_litellm_model_name, - "model_info": model_info, - "api_base": deployment_api_base, - "deployment_model_name": deployment_model_name, - KUBERNETES_POD_ROUTING_KEY: routing, - } - ) - other_bucket: Final = "metadata" if metadata_variable_name == "litellm_metadata" else "litellm_metadata" - other_bucket_metadata: Final = kwargs.get(other_bucket) - if isinstance(other_bucket_metadata, Mapping) and KUBERNETES_POD_ROUTING_KEY in other_bucket_metadata: - kwargs[other_bucket] = { # rebind-ok: this method updates the router's request kwargs in place - **other_bucket_metadata, - KUBERNETES_POD_ROUTING_KEY: routing, - } - - # A retry/fallback reuses this same kwargs dict for the next deployment. - # Refund and clear any reservation the previous deployment attempt left - # here before it's wiped below, instead of relying on that attempt's - # (possibly still-pending) failure event to do it. - refund_stale_reservation_before_retry(self._core_cache, kwargs) - set_io_token_rate_limit_request_kwargs( - kwargs, - store_in_context=deployment_has_io_token_limits(deployment), - opens_responses_socket=opens_responses_socket, - metadata_channel=metadata_variable_name, - ) - - kwargs[metadata_variable_name].setdefault( - ROUTING_REQUEST_TAGS_METADATA_KEY, - tuple(get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name)), - ) - - ## DEPLOYMENT-LEVEL TAGS - deployment_tags: Final = deployment.get("litellm_params", {}).get("tags") - merged_tags: Final = tags_with_serving_deployment_tags(kwargs[metadata_variable_name], deployment_tags) - if merged_tags is not None: - kwargs[metadata_variable_name]["tags"] = list(merged_tags) - - baggage_stamp: Final = baggage_stamp_for_serving_deployment(kwargs[metadata_variable_name], deployment_tags) - if baggage_stamp is not None: - kwargs[metadata_variable_name][BAGGAGE_MAPPED_TAGS_METADATA_KEY] = baggage_stamp - - ## CREDENTIAL NAME AS TAG - credential_name: Final = deployment.get("litellm_params", {}).get("litellm_credential_name") - if credential_name: - credential_tag: Final = f"Credential: {credential_name}" - existing_tags = kwargs[metadata_variable_name].get("tags") or [] - if credential_tag not in existing_tags: - existing_tags.append(credential_tag) - kwargs[metadata_variable_name]["tags"] = existing_tags - - from litellm.llms.github_copilot.per_user_auth import ( - github_copilot_per_user_credential_name, - ) - - configured_per_user_name: Final = github_copilot_per_user_credential_name( - cast( # cast-ok: deployment litellm_params is a str-keyed dict - "Mapping[str, object]", deployment.get("litellm_params") - ) - ) - caller_credential_name: Final[object] = cast( # cast-ok: kwargs is the untyped request dict - "object", kwargs.get("litellm_credential_name") - ) - if ( - configured_per_user_name is not None - and isinstance(caller_credential_name, str) - and caller_credential_name != configured_per_user_name - ): - raise litellm_exceptions.BadRequestError( - message="litellm_credential_name cannot be overridden on a deployment that uses per-user GitHub OAuth", - model=deployment_model_name, - llm_provider="", - ) - - kwargs["model_info"] = model_info - - if function_name == "_ageneric_api_call_with_fallbacks": - from litellm.passthrough.timeout_utils import ( - resolve_llm_passthrough_timeout, - ) - - _router_timeout: Final = ( - self.request_timeout - if self.request_timeout is not None - else float(self._explicit_timeout) - if isinstance(self._explicit_timeout, (int, float)) - else None - ) - _router_stream_timeout: Final = ( - self.stream_timeout - if self.stream_timeout is not None - else self.request_timeout - if self.request_timeout is not None - else self._core_default_litellm_params.get("stream_timeout") - ) - passthrough_timeout: Final = partial( - resolve_llm_passthrough_timeout, - litellm_params=deployment["litellm_params"], - router_timeout=_router_timeout, - router_stream_timeout=_router_stream_timeout, - ) - timeout_request: Final = {**kwargs, "stream": True} if requests_streamed_retrieval(kwargs) else kwargs - kwargs.update( - CompletionTimeout.stream_request_kwargs( - passthrough_timeout(kwargs={**kwargs, "stream": False}), passthrough_timeout(kwargs=timeout_request) - ) - if splits_stream_timeout and timeout_request.get("stream") - else {"timeout": passthrough_timeout(kwargs=timeout_request)} - ) - elif splits_stream_timeout: - request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs) - deployment_params: Final = _DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) - kwargs.update( - CompletionTimeout.stream_request_kwargs( - self._get_non_stream_timeout(kwargs=request_kwargs, data=deployment_params), - self._get_stream_timeout(kwargs=request_kwargs, data=deployment_params) - if request_kwargs.get("stream") - else None, - ) - ) - else: - kwargs["timeout"] = self._get_timeout( - kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), - data=_DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]), - ) - - deployment_litellm_params: Final = _DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) - if router_config_requests_stream_usage( - _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), - ( - _LITELLM_PARAMS_ADAPTER.validate_python( - deployment_litellm_params - if isinstance(deployment_litellm_params, Mapping) - else deployment_litellm_params.model_dump(exclude_none=True) - ), - defaults, - ), - ): - kwargs[STREAM_USAGE_INJECTED_FIELD] = True - self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) - pinned_model_id: Final = deployment_model_id_pin(deployment_context) - if pinned_model_id is not None: - kwargs["model_id"] = pinned_model_id # rebind-ok: this method's contract is filling the caller's kwargs - hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) - if _is_fallback_hop(hop_kwargs): - EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( - self._core_ctx, - hop_kwargs.get("input"), - hop_kwargs.get("messages"), - (_FALLBACK_HOP_ADAPTER.validate_python(deployment),), - unmarked_origin=deployment_that_just_failed, - ) - if RESPONSES_CONTINUATION_KWARG in hop_kwargs: - del kwargs[RESPONSES_CONTINUATION_KWARG] # rebind-ok: this method prepares the call's kwargs in place - kwargs["input"] = responses_continuation_input_for( # rebind-ok: the deployment's continuation input - hop_kwargs, _FALLBACK_HOP_ADAPTER.validate_python(deployment) - ) - - def _get_async_openai_model_client(self, deployment: dict, kwargs: dict): - """ - Helper to get AsyncOpenAI or AsyncAzureOpenAI client that was created for the deployment - - The same OpenAI client is re-used to optimize latency / performance in production - - If dynamic api key is provided: - Do not re-use the client. Pass model_client=None. The OpenAI/ AzureOpenAI client will be recreated in the handler for the llm provider - """ - potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs, client_type="async") - - # check if provided keys == client keys # - dynamic_api_key: Final = kwargs.get("api_key", None) - if ( - dynamic_api_key is not None - and potential_model_client is not None - and dynamic_api_key != potential_model_client.api_key - ): - model_client = None - else: - model_client = potential_model_client - - return model_client - - def _get_stream_timeout( - self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams - ) -> object: - """Helper to get stream timeout from kwargs or deployment params""" - return ( - kwargs.get("stream_timeout", None) # the params dynamically set by user - or get_service_tier_timeout(kwargs.get("service_tier"), data, "stream_timeout_by_service_tier") - or data.get("stream_timeout", None) # timeout set on litellm_params for this deployment - or self.stream_timeout # timeout set on router - or self.request_timeout # litellm_settings.request_timeout (per-attempt) - or self._core_default_litellm_params.get("stream_timeout", None) - ) - - def _get_non_stream_timeout( - self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams - ) -> object: - """Helper to get non-stream timeout from kwargs or deployment params""" - return ( - kwargs.get("timeout", None) # the params dynamically set by user - or kwargs.get("request_timeout", None) # the params dynamically set by user - or get_service_tier_timeout(kwargs.get("service_tier"), data, "timeout_by_service_tier") - or data.get("timeout", None) # timeout set on litellm_params for this deployment - or data.get("request_timeout", None) # timeout set on litellm_params for this deployment - or self.request_timeout # litellm_settings.request_timeout (per-attempt) - or self.timeout # timeout set on router (router_settings.timeout) - or self._core_default_litellm_params.get("timeout", None) - ) - - def _get_timeout(self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams) -> object: - """Helper to get timeout from kwargs or deployment params""" - stream_timeout: Final = self._get_stream_timeout(kwargs=kwargs, data=data) if kwargs.get("stream") else None - return stream_timeout if stream_timeout is not None else self._get_non_stream_timeout(kwargs=kwargs, data=data) - - async def abatch_completion( - self, - models: list[str], - messages: list[dict[str, str]] | list[list[dict[str, str]]], - **kwargs, - ): - """ - Async Batch Completion. Used for 2 scenarios: - 1. Batch Process 1 request to N models on litellm.Router. Pass messages as List[Dict[str, str]] to use this - 2. Batch Process N requests to M models on litellm.Router. Pass messages as List[List[Dict[str, str]]] to use this - - Example Request for 1 request to N models: - ``` - response = await router.abatch_completion( - models=["gpt-3.5-turbo", "groq-llama"], - messages=[ - {"role": "user", "content": "is litellm becoming a better product ?"} - ], - max_tokens=15, - ) - ``` - - - Example Request for N requests to M models: - ``` - response = await router.abatch_completion( - models=["gpt-3.5-turbo", "groq-llama"], - messages=[ - [{"role": "user", "content": "is litellm becoming a better product ?"}], - [{"role": "user", "content": "who is this"}], - ], - ) - ``` - """ - ############## Helpers for async completion ################## - - async def _async_completion_no_exceptions(model: str, messages: list[AllMessageValues], **kwargs): - """ - Wrapper around self.async_completion that catches exceptions and returns them as a result - """ - try: - return await self.acompletion(model=model, messages=messages, **kwargs) - except Exception as e: - return e - - async def _async_completion_no_exceptions_return_idx( - model: str, - messages: list[AllMessageValues], - idx: int, # index of message this response corresponds to - **kwargs, - ): - """ - Wrapper around self.async_completion that catches exceptions and returns them as a result - """ - try: - return ( - await self.acompletion(model=model, messages=messages, **kwargs), - idx, - ) - except Exception as e: - return e, idx - - ############## Helpers for async completion ################## - - if isinstance(messages, list) and all(isinstance(m, dict) for m in messages): - _tasks = [] - for model in models: - # add each task but if the task fails - _tasks.append(_async_completion_no_exceptions(model=model, messages=messages, **kwargs)) - response = await asyncio.gather(*_tasks) - return response - elif isinstance(messages, list) and all(isinstance(m, list) for m in messages): - _tasks = [] - for idx, message in enumerate(messages): - for model in models: - # Request Number X, Model Number Y - _tasks.append( - _async_completion_no_exceptions_return_idx( - model=model, - idx=idx, - messages=message, - **kwargs, - ) - ) - responses: Final = await asyncio.gather(*_tasks) - final_responses: Final[list[list[object]]] = [[] for _ in range(len(messages))] - for response in responses: - if isinstance(response, tuple): - final_responses[response[1]].append(response[0]) - else: - final_responses[0].append(response) - return final_responses - - async def abatch_completion_one_model_multiple_requests( - self, model: str, messages: list[list[AllMessageValues]], **kwargs - ): - """ - Async Batch Completion - Batch Process multiple Messages to one model_group on litellm.Router - - Use this for sending multiple requests to 1 model - - Args: - model (List[str]): model group - messages (List[List[Dict[str, str]]]): list of messages. Each element in the list is one request - **kwargs: additional kwargs - Usage: - response = await self.abatch_completion_one_model_multiple_requests( - model="gpt-3.5-turbo", - messages=[ - [{"role": "user", "content": "hello"}, {"role": "user", "content": "tell me something funny"}], - [{"role": "user", "content": "hello good mornign"}], - ] - ) - """ - - async def _async_completion_no_exceptions(model: str, messages: list[AllMessageValues], **kwargs): - """ - Wrapper around self.async_completion that catches exceptions and returns them as a result - """ - try: - return await self.acompletion(model=model, messages=messages, **kwargs) - except Exception as e: - return e - - _tasks: Final = [] - for message_request in messages: - # add each task but if the task fails - _tasks.append(_async_completion_no_exceptions(model=model, messages=message_request, **kwargs)) - - response: Final = await asyncio.gather(*_tasks) - return response - - # fmt: off - - @overload - async def abatch_completion_fastest_response( - self, model: str, messages: list[dict[str, str]], stream: Literal[True], **kwargs - ) -> CustomStreamWrapper: - ... - - - - @overload - async def abatch_completion_fastest_response( - self, model: str, messages: list[dict[str, str]], stream: Literal[False] = False, **kwargs - ) -> ModelResponse: - ... - - # fmt: on - - async def abatch_completion_fastest_response( - self, - model: str, - messages: list[dict[str, str]], - stream: bool = False, - **kwargs, - ): - """ - model - List of comma-separated model names. E.g. model="gpt-4, gpt-3.5-turbo" - - Returns fastest response from list of model names. OpenAI-compatible endpoint. - """ - models: Final = [m.strip() for m in model.split(",")] - - async def _async_completion_no_exceptions( - model_name: str, messages: list[dict[str, str]], stream: bool, **kwargs: object - ) -> ModelResponse | CustomStreamWrapper | Exception: - """ - Wrapper around self.acompletion that catches exceptions and returns them as a result - """ - try: - result = await self.acompletion(model=model_name, messages=messages, stream=stream, **kwargs) - return result - except asyncio.CancelledError: - verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model_name) - raise - except Exception as e: - return e - - pending_tasks = [] - - async def check_response(task: asyncio.Task): - nonlocal pending_tasks - try: - result: Final = await task - if isinstance(result, (ModelResponse, CustomStreamWrapper)): - verbose_router_logger.debug("Received successful response. Cancelling other LLM API calls.") - # If a desired response is received, cancel all other pending tasks - for t in pending_tasks: - t.cancel() - return result - except Exception: - # Ignore exceptions, let the loop handle them - pass - finally: - # Remove the task from pending tasks if it finishes - try: - pending_tasks.remove(task) - except KeyError: - pass - - for model_name in models: - task = asyncio.create_task( - _async_completion_no_exceptions(model_name=model_name, messages=messages, stream=stream, **kwargs) - ) - pending_tasks.append(task) - - # Await the first task to complete successfully - while pending_tasks: - done, pending_tasks = await asyncio.wait(pending_tasks, return_when=asyncio.FIRST_COMPLETED) - for completed_task in done: - result = await check_response(completed_task) - - if result is not None: - # Return the first successful result - if (result_hidden_params := get_hidden_params(result)) is not None: - result_hidden_params["fastest_response_batch_completion"] = True - return result - - # If we exit the loop without returning, all tasks failed - raise Exception("All tasks failed") - - ### SCHEDULER ### - - # fmt: off - - @overload - async def schedule_acompletion( - self, - model: str, - messages: list[AllMessageValues], - priority: int | None = None, - stream: Literal[False] = False, - **kwargs, - ) -> ModelResponse: - ... - - @overload - async def schedule_acompletion( - self, model: str, messages: list[AllMessageValues], priority: int | None, stream: Literal[True], **kwargs - ) -> CustomStreamWrapper: - ... - - # fmt: on - - async def schedule_acompletion( - self, - model: str, - messages: list[AllMessageValues], - priority: int | None = None, - stream=False, - **kwargs, - ): - queued_priority, _ = self._core_split_request_priority(model=model, kwargs={**kwargs, "priority": priority}) - if queued_priority is None: - raise priority_error(model=model, message=unqueueable_priority_message(priority)) - return await self.acompletion(model=model, messages=messages, stream=stream, priority=queued_priority, **kwargs) - - def _core_split_request_priority( - self, model: str, kwargs: Mapping[str, object] - ) -> tuple[int | None, Mapping[str, object]]: - request_priority: Final = resolve_request_priority( - requested=kwargs.get("priority"), - default_priority=self.default_priority, - drops_params=lambda: self._drops_params(model=model, kwargs=kwargs), - ) - if isinstance(request_priority, InvalidPriority): - raise priority_error(model=model, message=request_priority.message) - return request_priority, {key: value for key, value in kwargs.items() if key != "priority"} - - def priority_header_update( - self, model: str, kwargs: Mapping[str, object], header: str | None - ) -> Mapping[str, int | str]: - return priority_header_fields(header, drops_params=lambda: self._drops_params(model=model, kwargs=kwargs)) - - def _drops_params(self, model: str, kwargs: Mapping[str, object]) -> bool: - return request_drops_params( - kwargs, self._core_default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs) - ) - - async def _schedule_factory( - self, - model: str, - priority: int, - original_function: Callable[..., Awaitable[object]], - args: tuple[object, ...], - kwargs: Mapping[str, object], - ): - await self._wait_for_scheduler_turn( - model=model, - priority=priority, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - request_kwargs=kwargs, - ) - try: - _response: Final = await original_function(*args, **kwargs) - response_hidden_params: Final = get_hidden_params(_response) - if response_hidden_params is not None: - additional_headers: Final = cast( # cast-ok: router headers are stored as a mutable mapping - dict[str, object], response_hidden_params.setdefault("additional_headers", {}) - ) - additional_headers.update({"x-litellm-request-prioritization-used": True}) - return _response - except Exception as e: - setattr(e, "priority", priority) - raise e - - async def _wait_for_scheduler_turn( - self, model: str, priority: int, parent_otel_span: Span | None, request_kwargs: Mapping[str, object] - ) -> None: - pinned: Final = self._pinned_deployment(model) - - async def healthy_deployments() -> Sequence[object]: - if pinned is not None: - return (pinned,) - deployments, _ = await self._async_get_healthy_deployments( - model=model, - parent_otel_span=parent_otel_span, - request_kwargs=share_holder_request_kwargs(request_kwargs, self._ptu_share_teams()), - ) - return deployments - - await self.scheduler.wait_for_turn( - request=FlowItem(priority=priority, request_id=str(uuid.uuid4()), model_name=model), - timeout=self.timeout, - get_healthy_deployments=healthy_deployments, - ) - - def _is_prompt_management_model(self, model: str) -> bool: - model_list: Final = self._core_get_model_list(model_name=model) - if model_list is None or len(model_list) != 1: - return False - - litellm_model: Final = model_list[0]["litellm_params"].get("model", None) - if litellm_model is None or "/" not in litellm_model: - return False - - split_litellm_model: Final = litellm_model.split("/")[0] - return split_litellm_model in self._host.dynamic().known_custom_logger_compatible_callbacks - - def _pinned_deployment(self, model: str) -> Deployment | None: - return self._core_get_deployment(model_id=model) if model not in self._core_model_names else None - - def _request_deployment_params(self, model: str, kwargs: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: - deployment: Final = self._pinned_deployment(model) - if deployment is not None: - return (deployment.litellm_params.model_dump(exclude_none=True),) - candidates: Final = ( - self._get_deployment_by_litellm_model(model=model) - if kwargs.get("specific_deployment") is True - else self._core_get_model_list(model_name=model, team_id=get_request_team_id(kwargs)) or () - ) - return tuple(candidate["litellm_params"] for candidate in candidates) - - def _prompt_management_acompletion( - self, - model: str, - messages: list[AllMessageValues], # mutable-ok: _prompt_management_factory takes the caller's list - ) -> Callable[..., Awaitable[object]]: - async def acompletion(**kwargs: object) -> object: # kwargs-ok: called with the request's own keyword payload - return await self._prompt_management_factory(model=model, messages=messages, kwargs=dict(kwargs)) - - return acompletion - - async def _prompt_management_factory( - self, - model: str, - messages: list[AllMessageValues], - kwargs: dict[str, Any], - ): - litellm_logging_object = kwargs.get("litellm_logging_obj", None) - if litellm_logging_object is None: - litellm_logging_object, kwargs = function_setup( - **{ - "original_function": "acompletion", - "rules_obj": Rules(), - "start_time": get_utc_datetime(), - **kwargs, - } - ) - litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) - caller_timeout: Final = kwargs.get("timeout") or kwargs.get("request_timeout") - caller_stream_timeout: Final = kwargs.get("stream_timeout") - specific_deployment: Final = kwargs.pop("specific_deployment", None) - prompt_management_deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=cast(list[dict[str, str]], messages), # cast-ok: selection reads messages structurally - specific_deployment=specific_deployment, - request_kwargs=kwargs, - ) - - self._update_kwargs_with_deployment( - deployment=prompt_management_deployment, kwargs=kwargs, splits_stream_timeout=True - ) - data: Final = prompt_management_deployment["litellm_params"].copy() - - litellm_model: Final = data.get("model", None) - - # litellm_agent/ prefix only strips the model name, no prompt_id needed - is_litellm_agent_model: Final = isinstance(litellm_model, str) and litellm_model.startswith("litellm_agent/") - - prompt_id = kwargs.get("prompt_id") or prompt_management_deployment["litellm_params"].get("prompt_id", None) - prompt_variables: Final = kwargs.get("prompt_variables") or prompt_management_deployment["litellm_params"].get( - "prompt_variables", None - ) - prompt_label: Final = kwargs.get("prompt_label", None) or prompt_management_deployment["litellm_params"].get( - "prompt_label", None - ) - - if not is_litellm_agent_model and (prompt_id is None or not isinstance(prompt_id, str)): - raise ValueError(f"Prompt ID is not set or not a string. Got={prompt_id}, type={type(prompt_id)}") - if prompt_variables is not None and not isinstance(prompt_variables, dict): - raise ValueError( - f"Prompt variables is set but not a dictionary. Got={prompt_variables}, type={type(prompt_variables)}" - ) - - ( - model, - messages, - optional_params, - ) = litellm_logging_object.get_chat_completion_prompt( - model=litellm_model, - messages=messages, - non_default_params=get_non_default_completion_params(kwargs=kwargs), - prompt_id=prompt_id, - prompt_variables=prompt_variables, - prompt_label=prompt_label, - request_kwargs=kwargs, - injected_for_every_deployment=True, - ) - - # Filter out prompt management specific parameters from data before merging - prompt_management_params: Final = { - "bitbucket_config", - "dotprompt_config", - "prompt_id", - "prompt_variables", - "prompt_label", - "prompt_version", - } - filtered_data: Final = {k: v for k, v in data.items() if k not in prompt_management_params} - - kwargs = {**filtered_data, **kwargs, **optional_params} - kwargs["model"] = model - kwargs["messages"] = messages - kwargs["litellm_logging_obj"] = litellm_logging_object - kwargs["prompt_id"] = prompt_id - kwargs["prompt_variables"] = prompt_variables - kwargs["prompt_label"] = prompt_label - - _model_list: Final = self._core_get_model_list(model_name=model) - if _model_list is None or len(_model_list) == 0: # if direct call to model - kwargs.pop("original_function") - return await self._host.invoke("acompletion")(**kwargs) - - target_group_timeouts: Final = { - key: value - for key, value in ( - ("timeout", caller_timeout if caller_timeout is not None else filtered_data.get("timeout")), - ( - "stream_timeout", - caller_stream_timeout if caller_stream_timeout is not None else filtered_data.get("stream_timeout"), - ), - ) - if value is not None - } - return await self.async_function_with_fallbacks( - **{key: value for key, value in kwargs.items() if key not in ("timeout", "stream_timeout")}, - **target_group_timeouts, - ) - - def image_generation(self, prompt: str, model: str, **kwargs): - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["original_function"] = self._image_generation - kwargs.setdefault("metadata", {}).update({"model_group": model}) - response: Final = self.function_with_fallbacks(**kwargs) - - return response - except Exception as e: - raise e - - def _image_generation(self, prompt: str, model: str, **kwargs): - model_name: Final = "" - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _image_generation()- model: %s", model) - with picks_reserve(True): - deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - - ### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit) - self.routing_strategy_pre_call_checks(deployment=deployment) - - response: Final = self._host.invoke("image_generation")( - **{ - **data, - "prompt": prompt, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.image_generation(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.image_generation(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def aimage_generation(self, prompt: str, model: str, **kwargs): - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["original_function"] = self._aimage_generation - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _aimage_generation(self, prompt: str, model: str, **kwargs): - model_name = model - start_time: Final = datetime.now() # noqa: DTZ005 # SDK logging uses naive local timestamps - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _image_generation()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - try: - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - except Exception as selection_error: - # Admitted attempts keep SDK-owned logging, including deployment callbacks and retry state - request_kwargs: Final = cast( # cast-ok: Python kwargs have string keys and opaque values - Mapping[str, object], kwargs - ) - if request_kwargs.get("litellm_logging_obj") is None: - refusal_kwargs: Final = {**request_kwargs, "model": model, "prompt": prompt} - refusal_kwargs.setdefault("litellm_call_id", str(uuid.uuid4())) - logging_obj, _ = function_setup( - "aimage_generation", - Rules(), - start_time, - **refusal_kwargs, # pyright: ignore[reportArgumentType] # preserve opaque SDK options, including is_async_call - ) - try: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=selection_error, - traceback_exception=traceback.format_exc(), - prefer_async_handlers=True, - ) - ) - finally: - _restore_correlation_context_if_supported(logging_obj) - raise - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - data: Final = deployment["litellm_params"].copy() - model_name = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("aimage_generation")( - **{ - **data, - "prompt": prompt, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.aimage_generation(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.aimage_generation(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def atranscription(self, file: FileTypes, model: str, **kwargs): - """ - Example Usage: - - ``` - from litellm import Router - client = Router(model_list = [ - { - "model_name": "whisper", - "litellm_params": { - "model": "whisper-1", - }, - }, - ]) - - audio_file = open("speech.mp3", "rb") - transcript = await client.atranscription( - model="whisper", - file=audio_file - ) - - ``` - """ - try: - kwargs["model"] = model - kwargs["file"] = file - kwargs["original_function"] = self._atranscription - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _atranscription(self, file: FileTypes, model: str, **kwargs): - model_name: Final = model - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _atranscription()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("atranscription")( - **{ - **data, - "file": file, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.atranscription(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.atranscription(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): - """ - Example Usage: - - ``` - from litellm import Router - client = Router(model_list = [ - { - "model_name": "tts", - "litellm_params": { - "model": "tts-1", - }, - }, - ]) - - async with client.aspeech( - model="tts", - voice="alloy", - input="the quick brown fox jumped over the lazy dogs", - api_base=None, - api_key=None, - organization=None, - project=None, - max_retries=1, - timeout=600, - client=None, - optional_params={}, - ) as response: - response.stream_to_file(speech_file_path) - - ``` - """ - try: - kwargs["model"] = model - kwargs["input"] = input - kwargs["voice"] = voice - kwargs["original_function"] = self._aspeech - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): - model_name: Final = model - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _aspeech()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("aspeech")( - **{ - **data, - "input": input, - "voice": data.get("voice") if voice is None else voice, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.aspeech(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.aspeech(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def arerank(self, model: str, **kwargs): - try: - kwargs["model"] = model - kwargs["input"] = input - kwargs["original_function"] = self._arerank - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _arerank(self, model: str, **kwargs): - model_name = None - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _rerank()- model: %s", model) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_name = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - response: Final = await self._host.invoke("arerank")( - **{ - **data, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.arerank(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.arerank(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - def text_completion( - self, - model: str, - prompt: str, - is_retry: bool | None = False, - is_fallback: bool | None = False, - is_async: bool | None = False, - **kwargs, - ): - messages: Final = [{"role": "user", "content": prompt}] - with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) - kwargs.setdefault("metadata", {}).update({"model_group": model}) - - # pick the one that is available (lowest TPM/RPM) - with endpoint_selection("text_completion", model=model), picks_reserve(True): - deployment: Final = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - data: Final = deployment["litellm_params"].copy() - for k, v in self._core_default_litellm_params.items(): - if k not in kwargs: # prioritize model-specific params > default router params - kwargs[k] = v - elif k == "metadata": - kwargs[k].update(v) - pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data)) - if pinned_model_id is not None: - kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults - - # call via litellm.completion() - return self._host.invoke("text_completion")( - **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs} - ) - except Exception as e: - raise e - - async def atext_completion( - self, - model: str, - prompt: str, - is_retry: bool | None = False, - is_fallback: bool | None = False, - is_async: bool | None = False, - **kwargs, - ): - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["original_function"] = self._atext_completion - - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - request_priority, request_kwargs = self._core_split_request_priority(model=model, kwargs=kwargs) - if request_priority is None: - return await self.async_function_with_fallbacks(**request_kwargs) - return await self._schedule_factory( - model=model, - priority=request_priority, - original_function=self.async_function_with_fallbacks, - args=(), - kwargs=request_kwargs, - ) - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _atext_completion(self, model: str, prompt: str, **kwargs): - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _atext_completion()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": prompt}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("atext_completion")( - **{ - **data, - "prompt": prompt, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.atext_completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.atext_completion(model=%s)\x1b[31m Exception %s\x1b[0m", model, e) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def aadapter_completion( - self, - adapter_id: str, - model: str, - is_retry: bool | None = False, - is_fallback: bool | None = False, - is_async: bool | None = False, - **kwargs, - ): - try: - kwargs["model"] = model - kwargs["adapter_id"] = adapter_id - kwargs["original_function"] = self._aadapter_completion - kwargs.setdefault("metadata", {}).update({"model_group": model}) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _aadapter_completion(self, adapter_id: str, model: str, **kwargs): - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _aadapter_completion()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "default text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("aadapter_completion")( - **{ - **data, - "adapter_id": adapter_id, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.aadapter_completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.aadapter_completion(model=%s)\x1b[31m Exception %s\x1b[0m", model, e) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def _asearch_with_fallbacks(self, original_function: Callable, **kwargs) -> object: - """ - Helper function to make a search API call through the router with load balancing and fallbacks. - Reuses the router's retry/fallback infrastructure. - """ - from litellm.router_utils.search_api_router import SearchAPIRouter - - return _OPAQUE_VALUE_ADAPTER.validate_python( - await SearchAPIRouter.async_search_with_fallbacks( - litellm_routing_ctx=self._core_ctx, - original_function=original_function, - **kwargs, - ) - ) - - async def _asearch_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): - """ - Helper function for search API calls - selects a search tool and calls the original function. - Called by async_function_with_fallbacks for each retry attempt. - """ - from litellm.router_utils.search_api_router import SearchAPIRouter - - return await SearchAPIRouter.async_search_with_fallbacks_helper( - litellm_routing_ctx=self._core_ctx, - model=model, - original_generic_function=original_generic_function, - **kwargs, - ) - - async def aguardrail( - self, - guardrail_name: str, - original_function: Callable, - **kwargs, - ): - """ - Execute a guardrail with load balancing and fallbacks. - - Args: - guardrail_name: Name of the guardrail to execute - original_function: The guardrail's execution function (e.g., async_pre_call_hook) - **kwargs: Additional arguments passed to the guardrail - - Returns: - Result from the guardrail execution - """ - kwargs["model"] = guardrail_name # For fallback system compatibility - kwargs["original_generic_function"] = original_function - kwargs["original_function"] = self._aguardrail_helper - self._update_kwargs_before_fallbacks( - model=guardrail_name, - kwargs=kwargs, - metadata_variable_name="litellm_metadata", - ) - verbose_router_logger.debug("Inside aguardrail() - guardrail_name: %s", guardrail_name) - response: Final = await self.async_function_with_fallbacks(**kwargs) - return response - - async def _aguardrail_helper( - self, - model: str, - original_generic_function: Callable, - **kwargs, - ): - """ - Helper for aguardrail - selects a guardrail deployment and executes it. - Called by async_function_with_fallbacks for each retry attempt. - - Args: - model: The guardrail_name (named 'model' for fallback system compatibility) - original_generic_function: The guardrail's execution function - **kwargs: Additional arguments - """ - guardrail_name: Final = model - selected_guardrail: Final = self._core_get_available_guardrail( - guardrail_name=guardrail_name, - ) - - verbose_router_logger.debug( - "Selected guardrail deployment: %s", selected_guardrail.get("litellm_params", {}).get("guardrail") - ) - - # Pass the selected guardrail config to the original function - kwargs["selected_guardrail"] = selected_guardrail - response: Final = await original_generic_function(**kwargs) - return response - - def _core_get_available_guardrail( - self, - guardrail_name: str, - ) -> "GuardrailTypedDict": - """ - Select a guardrail deployment using the router's load balancing strategy. - - Args: - guardrail_name: Name of the guardrail to select - - Returns: - Selected guardrail configuration dict - """ - from litellm.router_strategy.simple_shuffle import simple_shuffle - - healthy_deployments: Final = [g for g in self._core_guardrail_list if g.get("guardrail_name") == guardrail_name] - - if not healthy_deployments: - raise ValueError(f"No guardrail found with name: {guardrail_name}") - - if len(healthy_deployments) == 1: - return healthy_deployments[0] - - # Use simple_shuffle for weighted selection - return simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=healthy_deployments, - model=guardrail_name, - request_kwargs=None, - ) - - async def _ageneric_api_call_with_fallbacks( - self, model: str, original_function: Callable, attempt_function: Callable | None = None, **kwargs - ): - """ - Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router - - attempt_function runs every attempt of the chain instead of the plain helper, so a streaming - endpoint can wrap each attempt's stream with its own mid-stream fallback handling. - """ - try: - kwargs["model"] = model - kwargs["original_generic_function"] = original_function - kwargs["original_function"] = attempt_function or self._ageneric_api_call_with_fallbacks_helper - if attempt_function is not None: - controls: Final = per_request_fallback_controls(kwargs) - kwargs[MID_STREAM_FALLBACK_CONTROLS_KEY] = controls # rebind-ok: forwarded to every hop - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs, metadata_variable_name="litellm_metadata") - verbose_router_logger.debug("Inside ageneric_api_call_with_fallbacks() - model: %s", model) - response: Final = await self.async_function_with_fallbacks(**kwargs) - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - def _add_deployment_model_to_endpoint_for_llm_passthrough_route( - self, kwargs: dict[str, Any], model: str, model_name: str - ) -> dict[str, Any]: - """ - Add the deployment model to the endpoint for LLM passthrough route. - - e.g for bedrock invoke users can pass endpoint as /model/special-bedrock-model/invoke - it should be actually sent as /model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke - """ - if "endpoint" in kwargs and kwargs["endpoint"]: - # For provider-specific endpoints, strip the provider prefix from model_name - # e.g., "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0" -> "us.anthropic.claude-3-5-sonnet-20240620-v1:0" - from litellm import get_llm_provider - - try: - # get_llm_provider returns (model_without_prefix, provider, api_key, api_base) - stripped_model_name, _, _, _ = get_llm_provider( - model=model_name, - custom_llm_provider=kwargs.get("custom_llm_provider"), - api_base=kwargs.get("api_base"), - ) - replacement_model_name = stripped_model_name - except Exception: - # If get_llm_provider fails, fall back to using model_name as-is - replacement_model_name = model_name - - kwargs["endpoint"] = replace_path_segment(kwargs["endpoint"], model, replacement_model_name) - return kwargs - - async def _ageneric_api_call_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): - """ - Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router - """ - - passthrough_on_no_deployment: Final = kwargs.pop("passthrough_on_no_deployment", False) - function_name: Final = "_ageneric_api_call_with_fallbacks" - handler_name: Final[str | None] = getattr(original_generic_function, "__name__", None) - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - try: - with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): - deployment = ( # rebind-ok: set on success, see pre-init - await self._core_async_get_available_deployment( - model=model, - request_kwargs=kwargs, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - ) - ) - except Exception as e: - if passthrough_on_no_deployment: - return await original_generic_function(model=model, **kwargs) - raise e - - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name=function_name, - splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, - opens_responses_socket=handler_name == "_aresponses_websocket", - ) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self.total_calls[model_name] += 1 - - self._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, model=model, model_name=model_name - ) - - custom_llm_provider: Final = provider_for_generic_call(data) - generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "") - - response_kwargs: Final = { - **without_retrieve_stream_defaults(data, kwargs), - "caching": self._core_cache_responses, - **kwargs, - "model": model_name, - **_with_router_resolved_session_model(kwargs.get("session"), model_name), - **( - {"turn_admission": partial(self.async_admit_deployment_turn, model, deployment, kwargs)} - if generic_function_name == "_aresponses_websocket" - else {} - ), - } - # Only set custom_llm_provider if it's not None - if custom_llm_provider is not None: - response_kwargs["custom_llm_provider"] = custom_llm_provider - - compacted_input: Final = await compact_to_fit( - self._core_ctx, - deployment, - response_kwargs, - surface_for_call(generic_function_name), - ) - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await original_generic_function(**compacted_input) - - if self._should_raise_anthropic_refusal_error( - model=model, - original_generic_function=original_generic_function, - response=response, - kwargs=kwargs, - ): - from litellm.llms.anthropic.messages.native.utils import ( - safeguard_refusal_error, - ) - - refusal_details: Final = cast(dict, response["stop_details"]) # cast-ok: gate verified the shape - raise safeguard_refusal_error(model=model, stop_details=refusal_details) - - self.success_calls[model_name] += 1 - realtime_session_logging: Final = ( - _CALL_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") - if generic_function_name == "_arealtime" - else None - ) - realtime_session_deployment_id: Final = ( - _MODEL_INFO_ADAPTER.validate_python(kwargs["model_info"]).get("id") - if isinstance(realtime_session_logging, LiteLLMLogging) - and realtime_session_logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True - and realtime_session_logging.model_call_details.get(REALTIME_SESSION_ENDED_ABNORMALLY_KEY) is not True - else None - ) - if realtime_session_deployment_id is not None: - increment_deployment_successes_for_current_minute( - ctx=self._core_ctx, deployment_id=str(realtime_session_deployment_id) - ) - verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - except Exception as e: - verbose_router_logger.info( - "ageneric_api_call_with_fallbacks(model=%s)\x1b[31m Exception %s\x1b[0m", model, e - ) - if model is not None: - self.fail_calls[model] += 1 - if deployment is not None: - self._set_deployment_num_retries_on_exception(e, deployment) - self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) - raise e - - async def _aresponses_with_streaming_fallbacks( - self, original_function: Callable, **kwargs: Any - ) -> Union["ResponsesAPIResponse", "BaseResponsesAPIStreamingIterator"]: - """ - _ageneric_api_call_with_fallbacks for the Responses API, with every attempt's stream - carrying its own mid-stream fallback handling - (see _ageneric_api_call_with_fallbacks_responses_attempt). - """ - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - attempt_function=self._ageneric_api_call_with_fallbacks_responses_attempt, - **kwargs, - ) - - async def _ageneric_api_call_with_fallbacks_responses_attempt( - self, - model: str, - original_generic_function: Callable, - **kwargs: object, # kwargs-ok: forwarded verbatim to the per-attempt helper, shape varies per call site - ) -> Union["ResponsesAPIResponse", "BaseResponsesAPIStreamingIterator"]: - """ - One attempt of the Responses API fallback chain. A streaming result is wrapped with - _aresponses_streaming_iterator over this attempt's own kwargs, so a fallback hop that - fails mid-stream resumes the original group's chain instead of re-raising; the name keeps - _get_router_metadata_variable_name resolving to litellm_metadata for every hop. - """ - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - - controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None) - hop_kwargs: Final = mid_stream_fallback_hop_kwargs( - model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs - ) - shared_logging: Final = kwargs.get("litellm_logging_obj") - if isinstance(shared_logging, LiteLLMLogging): - shared_logging.carry_failed_attempt_usage() - response: Final = await self._ageneric_api_call_with_fallbacks_helper( - model=model, original_generic_function=original_generic_function, **kwargs - ) - carry_over_pre_routing_selection(live_kwargs=kwargs, snapshot=hop_kwargs) - carry_over_routed_deployment(live_kwargs=kwargs, snapshot=hop_kwargs) - if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator): - return await self._aresponses_streaming_iterator(response=response, initial_kwargs=hop_kwargs) - return response - - async def _aanthropic_messages_streaming_iterator( - self, - response: AsyncIterator[bytes], - initial_kwargs: dict[str, object], # mutable-ok: mutated in-place before re-entering the fallback chain - ) -> AsyncIterator[bytes]: - """ - Wrap an anthropic_messages (/v1/messages) streaming response so a - mid-stream provider error triggers the Router's fallback chain - (parity with _acompletion_streaming_iterator for the - chat-completions path). See #24004. - - anthropic_messages goes through _ageneric_api_call_with_fallbacks - rather than _acompletion, so the returned byte iterator is never - wrapped by the chat-completions fallback handler. Two failure - shapes land here: - - the completion-bridge path (deployments with no native - /v1/messages endpoint, via - LiteLLMMessagesToCompletionTransformationHandler) already - raises MidStreamFallbackError out of its underlying - CustomStreamWrapper; this wrapper only needs to catch it. - - a native Anthropic/Bedrock passthrough never raises anything - for a provider SSE `event: error` frame (e.g. `overloaded_error`, - `internal_server_error`) - it is forwarded to the client as-is - - so this wrapper detects it via parse_anthropic_error_event and - raises MidStreamFallbackError itself. - - Only an error before any real content (a content_block_delta frame) - has reached the caller triggers a fallback attempt, mirroring the - restriction _acompletion_streaming_iterator applies: once generated - output has already reached the caller, retrying would start a - second, overlapping Anthropic message lifecycle on the same SSE - stream, so the error is left to propagate instead of being retried - invisibly. A non-retriable client error (4xx other than 429) is - never worth a fallback attempt either, so it is also left to - propagate. - - Lifecycle/bookkeeping frames (message_start, content_block_start, - ping, ...) do not by themselves disqualify a fallback attempt - - Anthropic routinely sends message_start before an overload error. - When a fallback can still take over they are BUFFERED rather than - forwarded immediately, since forwarding one and then appending a - fallback attempt's own message_start would produce two overlapping - message lifecycles on one SSE stream; a `ping` carries no lifecycle, - so it is forwarded live even while lifecycle frames sit buffered, - keeping the connection alive during a long thinking pass. Buffered - frames are flushed, in order, the moment real content arrives (the - primary attempt has committed by then anyway) or once the stream - ends without ever producing content or an error. When no fallback - can take over the request is already committed, so every frame, - including pings and provider error frames, is forwarded live and - verbatim instead. - """ - from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( - FakeAnthropicMessagesStreamIterator, - ) - from litellm.llms.anthropic.messages.native.streaming_iterator import ( - aclose_if_supported, - aclose_superseded_attempt, - parse_anthropic_error_event, - parse_anthropic_refusal_stop_details, - ) - from litellm.llms.anthropic.messages.native.utils import ( - safeguard_refusal_error, - ) - - source_iterator: Final = response - replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator) - - async def close_source() -> None: - with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): - await aclose_if_supported(source_iterator) - - async def close_superseded() -> None: - with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): - await aclose_superseded_attempt(source_iterator) - - async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]: - from litellm.exceptions import MidStreamFallbackError - - # Lifecycle/bookkeeping frames (message_start, content_block_start, - # ...) are held back rather than forwarded immediately, but only - # while a fallback can still take over: Anthropic routinely sends - # message_start before an overload error, and once a byte reaches - # the client a fallback attempt can only append its OWN - # message_start, producing two overlapping message lifecycles on - # one SSE stream. A `ping` keepalive carries no lifecycle, so it - # is forwarded live even behind buffered frames, keeping the - # connection alive through a long thinking pass. Buffered frames - # are flushed the moment real content (content_block_delta) - # arrives - at that point the primary attempt has committed and a - # clean retry is no longer possible anyway - or once the primary - # stream ends without ever producing content. Hitting - # MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS forces the same early - # commit as real content arriving, so a hostile or pathological - # upstream can't grow the buffer forever. With no fallback able - # to take over there is nothing to buffer for, so every frame, - # including pings and provider error frames, is forwarded live. - model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group - has_generated_content = not ( # rebind-ok: set once real content is seen, the buffer cap is hit, or neither a retry nor a fallback can take over - self._anthropic_messages_stream_can_retry(initial_kwargs) - or self._stream_can_fall_back(model, initial_kwargs) - ) - buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline - # An error this loop raises off a frame, not the provider's. - raised_here: BaseException | None = None # rebind-ok: set when the loop raises off a frame - try: - async for chunk in source_iterator: - if _anthropic_stream_forwards_ping_live(chunk, has_generated_content): - yield chunk - continue - if _stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)): - has_generated_content = True - # A transport can split one SSE data line across byte chunks, so pre-content - # detection parses the accumulated buffer plus the current chunk, never the - # chunk alone; the buffer is already capped, which bounds this window too. - parse_window = ( - b"".join(c for c in (*buffered_lifecycle_chunks, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime - if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime - else chunk - ) - error_event = parse_anthropic_error_event(parse_window) - recoverable_frame_error = self._anthropic_messages_recoverable_frame_error( - error_event, chunk, has_generated_content, model, initial_kwargs - ) - refusal_stop_details = ( - parse_anthropic_refusal_stop_details(parse_window) - if not has_generated_content and error_event is None - else None - ) - if ( - refusal_stop_details is not None - and not replays_a_finished_agentic_turn - and self._refusal_fallback_available(model, initial_kwargs) - ): - refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) - raised_here = MidStreamFallbackError( - message=refusal_error.message, - model=model, - llm_provider="anthropic", - original_exception=refusal_error, - is_pre_first_chunk=True, - ) - raise raised_here - if not has_generated_content and error_event is None: - buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk) - continue - if recoverable_frame_error is not None: - assert error_event is not None - raised_here = MidStreamFallbackError( - message=error_event[1], - model=model, - llm_provider="anthropic", - original_exception=recoverable_frame_error, - is_pre_first_chunk=True, - ) - raise raised_here - for buffered_chunk in buffered_lifecycle_chunks: - yield buffered_chunk - buffered_lifecycle_chunks = () - yield chunk - for buffered_chunk in buffered_lifecycle_chunks: - yield buffered_chunk - except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate - if stream_error is not raised_here: - await self._areport_raised_messages_stream_failure(stream_error, initial_kwargs, source_iterator) - async for item in self._aanthropic_messages_recover_stream_error( - stream_error, - has_generated_content, - buffered_lifecycle_chunks, - model, - initial_kwargs, - wrapper, - close_superseded, - ): - yield item - finally: - await close_source() - - # Referenced by stream_with_fallbacks via closure - assigned here, before - # the generator body ever runs, so the reference resolves fine despite - # being defined textually after the function that captures it. - wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator) - return wrapper - - async def _areport_raised_messages_stream_failure( - self, stream_error: Exception, attempt_kwargs: Mapping[str, object], source: object - ) -> None: - """A /v1/messages stream whose provider iterator raised, counted against the deployment - that served it and cooled down by the rules, the same as a non-stream call's failure, so a - retry picks another deployment. A stream that logs itself (litellm's own stream logging - reaches the router's failure callback), or a failure the request's logging object already - logged, is not counted twice.""" - from litellm.exceptions import MidStreamFallbackError - from litellm.litellm_core_utils.litellm_logging import Logging - - if isinstance(getattr(source, "logging_obj", None), Logging): - return - logging_obj: Final = attempt_kwargs.get("litellm_logging_obj") - if isinstance(logging_obj, Logging) and logging_obj.model_call_details.get("has_logged_async_failure") is True: - return - bucket: Final = as_str_mapping(attempt_kwargs.get(get_metadata_variable_name_from_kwargs(attempt_kwargs))) or {} - model_info: Final = as_str_mapping(bucket.get("model_info")) or {} - deployment_id: Final = model_info.get("id") - if deployment_id is None: - return - deployment: Final = self.get_deployment(model_id=str(deployment_id)) - deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} - provider_error: Final = ( - stream_error.original_exception - if isinstance(stream_error, MidStreamFallbackError) and stream_error.original_exception is not None - else stream_error - ) - now: Final = datetime.now() # noqa: DTZ005 # naive, as the logging pipeline's times - logged: Final = { - "call_type": "anthropic_messages", - "exception": provider_error, - "litellm_params": { - **{ - name: deployment_params[name] for name in ("api_base", "cooldown_time") if name in deployment_params - }, - "metadata": dict(bucket), - "model_info": dict(model_info), - }, - } - try: - _ = self.deployment_callback_on_failure(logged, None, now, now) # pyright: ignore[reportUnknownMemberType] # an untyped callback - await self.async_deployment_callback_on_failure(logged, None, now, now) # pyright: ignore[reportUnknownMemberType] # see above - except Exception as error: # noqa: BLE001 # a recording never fails the stream - verbose_router_logger.warning("a stream's failure was not recorded: %s", error) - - async def _aanthropic_messages_recover_stream_error( - self, - stream_error: Exception, - has_generated_content: bool, - buffered_lifecycle_chunks: tuple[bytes, ...], - model: str, - initial_kwargs: dict[str, object], # mutable-ok: _aanthropic_messages_fallback_attempt mutates it - wrapper: "FallbackAwareAnthropicMessagesStream", - close_superseded: Callable[[], Awaitable[None]], - ) -> AsyncGenerator[bytes, None]: - """Turns a source-iterator failure into a fallback attempt or the error reaching the caller.""" - from litellm.exceptions import MidStreamFallbackError - - if isinstance(stream_error, MidStreamFallbackError) and _anthropic_stream_should_decline_fallback( - has_generated_content, stream_error - ): - for buffered_chunk in buffered_lifecycle_chunks: - yield buffered_chunk - if stream_error.original_exception is not None: - raise stream_error.original_exception from stream_error - raise stream_error - fallback_error: Final = ( - stream_error - if isinstance(stream_error, MidStreamFallbackError) - else _stream_fallback_error_for_raised(stream_error, model, has_generated_content) - ) - if fallback_error is None: - raise stream_error - await close_superseded() - outcome: Final = await self._astream_retry_same_group( - fallback_error, - initial_kwargs, - attempt_function=self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt, - ) - match outcome: - case _StreamRetryOpened(response=retried, attempted_retries=attempted, max_retries=budget): - async for item in self._aanthropic_messages_yield_recovered(retried, wrapper, (attempted, budget)): - yield item - case _StreamRetriesExhausted(error=last_error): - async for item in self._aanthropic_messages_fallback_attempt(last_error, initial_kwargs, wrapper): - yield item - case _: - assert_never(outcome) - - def _anthropic_messages_group_retry_policy(self, kwargs: Mapping[str, object]) -> dict[str, RetryPolicy] | None: - configured: Final = kwargs.get("model_group_retry_policy", self.model_group_retry_policy) - return cast("dict[str, RetryPolicy] | None", configured) # cast-ok: same type as the router attribute - - def _anthropic_messages_resolved_retry_policy(self, kwargs: Mapping[str, object]) -> RetryPolicy | None: - """ - The retry policy for this request's model group, unless the request opted out with num_retries=0. - A policy that does not resolve to a RetryPolicy governs nothing here, so the stream runs as it would - with none; the pre-stream retry loop still reports the malformed policy when an attempt fails. - """ - if kwargs.get("num_retries") == 0: - return None - model_group: Final = _request_model_group(kwargs) - try: - return _RESOLVED_RETRY_POLICY_ADAPTER.validate_python( - resolve_retry_policy( - retry_policy=self.retry_policy, - model_group=model_group, - model_group_retry_policy=self._anthropic_messages_group_retry_policy(kwargs), - ) - ) - except (TypeError, ValidationError) as malformed: - verbose_router_logger.warning( - "The retry policy for %s is not a RetryPolicy, streaming without one: %s", model_group, malformed - ) - return None - - def _anthropic_messages_plain_retry_budget(self, kwargs: Mapping[str, object]) -> int: - """ - The same precedence async_function_with_retries resolves for a failure raised before the - stream opened: the request's num_retries, then the routed deployment's, then the router's. - """ - request_num_retries: Final = kwargs.get("num_retries") - if isinstance(request_num_retries, int): - return request_num_retries - deployment_id: Final = routed_deployment_id(kwargs) - deployment: Final = self._core_get_deployment(deployment_id) if deployment_id is not None else None - deployment_num_retries: Final = _deployment_num_retries(deployment) - if deployment_num_retries is not None: - return deployment_num_retries - return self.num_retries if self.num_retries is not None else 0 - - def _anthropic_messages_policy_retries(self, trigger: Exception, kwargs: Mapping[str, object]) -> int | None: - policy: Final = self._anthropic_messages_resolved_retry_policy(kwargs) - if policy is None: - return None - return _get_num_retries_from_retry_policy(exception=trigger, retry_policy=policy) - - def _stream_retry_budget(self, trigger: Exception, kwargs: Mapping[str, object]) -> tuple[int, bool]: - """ - The budget an earlier retry of this request committed to, else the retry policy's grant when one - names this error, else the plain budget, with whether a policy governs the retry: a committed budget - is kept whichever deployment the retry lands on, as async_function_with_retries keeps its own. - """ - policy_retries: Final = self._anthropic_messages_policy_retries(trigger, kwargs) - committed_budget: Final = committed_retry_budget_for_request(kwargs) - if committed_budget is not None: - return committed_budget, policy_retries is not None - if policy_retries is None: - return self._anthropic_messages_plain_retry_budget(kwargs), False - return policy_retries, True - - def _anthropic_messages_stream_can_retry(self, kwargs: Mapping[str, object]) -> bool: - """ - Whether a pre-content failure of this stream would be retried within its own model group, - the other case where holding lifecycle frames back from the client buys a clean restart. - A retry policy names its budget per error class, so the largest budget it names bounds the - hold: holding frames one attempt too long is safe, forwarding them before a retry is not. - """ - attempted: Final = attempted_retries_for_request(kwargs) - committed_budget: Final = committed_retry_budget_for_request(kwargs) - if committed_budget is not None: - return committed_budget > attempted - plain_budget: Final = self._anthropic_messages_plain_retry_budget(kwargs) - policy: Final = self._anthropic_messages_resolved_retry_policy(kwargs) - ceiling: Final = plain_budget if policy is None else max(plain_budget, _retry_policy_ceiling(policy)) - return ceiling > attempted - - def _anthropic_messages_recoverable_frame_error( - self, - error_event: tuple[str, str, int] | None, - chunk: object, - has_generated_content: bool, - model_group: str, - kwargs: Mapping[str, object], - ) -> Exception | None: - """ - The exception a provider `event: error` frame before content recovers through when a retry of its - class or a fallback can still take over. A frame nothing can take over for (content already out, a - gateway verdict, or a class granted no retry with no fallback) reaches the client as the provider - sent it, the way the last exhausted attempt's does. - """ - if has_generated_content or error_event is None: - return None - error_type, message, status_code = error_event - if not _is_retriable_anthropic_status(status_code) or _anthropic_stream_error_is_gateway_verdict(chunk): - return None - frame_error: Final = anthropic_error_frame_exception(error_type, message, status_code, model_group) - budget, _ = self._stream_retry_budget(frame_error, kwargs) - if budget > attempted_retries_for_request(kwargs): - return frame_error - if self._stream_can_fall_back(model_group, kwargs): - return frame_error - return None - - def _stream_should_retry( - self, - trigger: Exception, - healthy_deployments: list[dict], # mutable-ok: should_retry_this_error's own parameter type - all_deployments: list[dict], # mutable-ok: should_retry_this_error's own parameter type - kwargs: Mapping[str, object], - ) -> bool: - try: - self.should_retry_this_error( - error=trigger, - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - context_window_fallbacks=_request_fallback_list( - kwargs, - "context_window_fallbacks", - cast( # cast-ok: untyped router attribute - "list[object] | None", self._core_context_window_fallbacks - ), - ), - content_policy_fallbacks=_request_fallback_list( - kwargs, - "content_policy_fallbacks", - cast( # cast-ok: untyped router attribute - "list[object] | None", self._core_content_policy_fallbacks - ), - ), - regular_fallbacks=_request_fallback_list( - kwargs, - "fallbacks", - cast("list[object] | None", self._core_fallbacks), # cast-ok: untyped router attribute - ), - ) - except Exception: # noqa: BLE001 # should_retry_this_error declines by raising the error it was given - return False - return True - - async def _astream_retry_same_group( - self, - e: "MidStreamFallbackError", - initial_kwargs: Mapping[str, object], - attempt_function: "Callable[..., Awaitable[object]]", - ) -> _StreamRetryOutcome: - """ - Re-runs the attempt within the request's own model group, the way async_function_with_retries - would have for a failure raised before the stream opened, until a retry opens a stream or the - budget runs out. Each retry's stream carries its own wrapper with the remaining budget, so a - retry that drops before content again continues the same count instead of starting over. - `attempt_function` is the request's own attempt (the /v1/messages attempt, or a chat completions one). - """ - model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group - retry_kwargs: Final = mid_stream_retry_kwargs(initial_kwargs) - handler: Final = retry_kwargs.get("original_generic_function") - healthy_deployments, all_deployments = await self._async_get_healthy_deployments( - model=model_group, - parent_otel_span=get_parent_otel_span_from_kwargs(retry_kwargs), - request_kwargs=share_holder_request_kwargs(initial_kwargs, self._ptu_share_teams()), - ) - budget, policy_applies = self._stream_retry_budget(_mid_stream_retry_trigger(e), initial_kwargs) - last_error = e # rebind-ok: the newest failure is what the fallback chain and the caller see - for attempt in range(attempted_retries_for_request(initial_kwargs), budget): - trigger = _mid_stream_retry_trigger(last_error) - if is_agentic_followup_failure(trigger): - return _StreamRetriesExhausted(last_error) - if not policy_applies and not self._stream_should_retry( - trigger, healthy_deployments, all_deployments, initial_kwargs - ): - return _StreamRetriesExhausted(last_error) - if get_metadata_variable_name_from_kwargs(retry_kwargs) in retry_kwargs: - self._core_log_retry(kwargs=retry_kwargs, e=trigger) - await self._host.sleep( - self._time_to_sleep_before_retry( - e=trigger, - remaining_retries=budget - attempt, - num_retries=budget, - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) - ) - record_retry_attempt(retry_kwargs, attempted_retries=attempt + 1, max_retries=budget) - verbose_router_logger.debug( - "Retrying stream that failed before content, attempt %s of %s", attempt + 1, budget - ) - carry_superseded_attempt_usage(retry_kwargs) - try: - with endpoint_selection( - getattr(handler, "__name__", None), getattr(handler, "__module__", None), model=model_group - ): - response = await attempt_function(**retry_kwargs) - except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry - wrapped = _stream_fallback_error_for_raised(retry_error, model_group, False, e.llm_provider) - if wrapped is None: - return _StreamRetriesExhausted(_stream_pre_content_error(retry_error, model_group, e.llm_provider)) - last_error = wrapped - continue - return _StreamRetryOpened(response, attempted_retries=attempt + 1, max_retries=budget) - return _StreamRetriesExhausted(last_error) - - async def _aanthropic_messages_yield_recovered( - self, - recovered: object, - wrapper: "FallbackAwareAnthropicMessagesStream", - retry_counters: tuple[int, int] | None = None, - ) -> AsyncGenerator[bytes, None]: - """ - Hands a retry's or a fallback's response to the client through the wrapper, closing it afterwards. - A retry stamps the retry headers async_function_with_retries would have for a pre-stream retry; - a later hop the recovered stream makes replaces them with its own, the way a fallback's do. - """ - from litellm.llms.anthropic.messages.native.streaming_iterator import ( - aclose_if_supported, - anthropic_messages_response_as_sse_events, - ) - - hidden_params, headers = PythonCore._prepare_fallback_hidden_params(recovered) - wrapper.merge_fallback_hidden_params(hidden_params, headers) - wrapper.adopt_fallback_source(recovered) - if retry_counters is not None: - add_retry_headers_to_response( - response=wrapper, attempted_retries=retry_counters[0], max_retries=retry_counters[1] - ) - try: - if hasattr(recovered, "__aiter__"): - async for item in cast("AsyncIterator[bytes]", recovered): # cast-ok: __aiter__ checked above - wrapper.follow_source_attribution() - yield item - return - # A recovery can resolve to a complete AnthropicMessagesResponse - # dict even for a streaming request (e.g. an agentic tool-use - # interception loop) - yielding it as-is would put a raw dict - # into a byte stream, so it's synthesized into the SSE - # lifecycle a real stream would have sent instead. - for event in anthropic_messages_response_as_sse_events( - cast("AnthropicMessagesResponse", recovered) # cast-ok: non-streaming shape by elimination - ): - yield event - finally: - with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): - await aclose_if_supported(recovered) - - async def _aanthropic_messages_fallback_attempt( - self, - e: "MidStreamFallbackError", - initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain - wrapper: "FallbackAwareAnthropicMessagesStream", - ) -> AsyncGenerator[bytes, None]: - """ - Re-enters the Router's fallback chain for a mid-stream - anthropic_messages error and yields whatever the fallback attempt - produces. Split out of _aanthropic_messages_streaming_iterator to - keep each function's cyclomatic complexity within the repo's C901 - budget. - """ - from litellm.exceptions import MidStreamFallbackError - - try: - model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: model group - fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the common_utils list|None param - "fallbacks", self._core_fallbacks - ) - context_window_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below - "context_window_fallbacks", self._core_context_window_fallbacks - ) - content_policy_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below - "content_policy_fallbacks", self._core_content_policy_fallbacks - ) - initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt - self._update_kwargs_before_fallbacks( - model=model_group, - kwargs=initial_kwargs, - metadata_variable_name="litellm_metadata", - ) - # The content-policy and context-window dispatch branches match on the trigger's own type, so - # such an error's MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. - fallback_trigger: Final[Exception] = ( - e.original_exception - if isinstance( - e.original_exception, - (litellm_exceptions.ContentPolicyViolationError, litellm_exceptions.ContextWindowExceededError), - ) - else e - ) - fallback_response: Final = _OPAQUE_VALUE_ADAPTER.validate_python( - await self.async_function_with_fallbacks_common_utils( - e=fallback_trigger, - disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, - ) - ) - async for fallback_item in self._aanthropic_messages_yield_recovered(fallback_response, wrapper): - yield fallback_item - except Exception as fallback_error: - verbose_router_logger.error("Anthropic messages streaming fallback also failed: %s", fallback_error) - if isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None: - raise fallback_error.original_exception from fallback_error - raise - - async def _aanthropic_messages_with_streaming_fallbacks( - self, - original_function: Callable, - **kwargs: object, # kwargs-ok: forwarded verbatim to original_function, shape varies per call site - ) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]: - """ - _ageneric_api_call_with_fallbacks for anthropic_messages, with every attempt's stream - carrying its own mid-stream fallback handling - (see _ageneric_api_call_with_fallbacks_anthropic_messages_attempt). Parity with - _aresponses_with_streaming_fallbacks for the Responses API. - """ - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - attempt_function=self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt, - **kwargs, - ) - - async def _ageneric_api_call_with_fallbacks_anthropic_messages_attempt( - self, - model: str, - original_generic_function: Callable[..., Awaitable[object]], - **kwargs: object, # kwargs-ok: forwarded verbatim to the per-attempt helper, shape varies per call site - ) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]: - """ - One attempt of the anthropic_messages fallback chain. A streaming result is wrapped with - _aanthropic_messages_streaming_iterator over this attempt's own kwargs, so a fallback hop - that fails mid-stream resumes the original group's chain instead of re-raising; the name - keeps _get_router_metadata_variable_name resolving to litellm_metadata for every hop. - """ - controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None) - hop_kwargs: Final = mid_stream_fallback_hop_kwargs( - model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs - ) - response: Final = await self._ageneric_api_call_with_fallbacks_helper( - model=model, original_generic_function=original_generic_function, **kwargs - ) - carry_over_pre_routing_selection(live_kwargs=kwargs, snapshot=hop_kwargs) - carry_over_routed_deployment(live_kwargs=kwargs, snapshot=hop_kwargs) - if kwargs.get("stream") and hasattr(response, "__aiter__"): - return await self._aanthropic_messages_streaming_iterator( - response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator - initial_kwargs=hop_kwargs, - ) - return response - - async def _dispatch_generic_call_type( - self, - call_type: str, - original_function: Callable, - **kwargs: object, # kwargs-ok: forwarded verbatim to the per-call-type helper, shape varies per call site - ): - """ - factory_function's shared dispatch for call types with no - call-specific handling, except anthropic_messages: kept out of - factory_function's own async_wrapper (already at the repo's C901 - complexity ceiling) so routing its mid-stream fallback handling - (#24004) doesn't add another branch there. - """ - if call_type == "anthropic_messages": - return await self._aanthropic_messages_with_streaming_fallbacks( - original_function=original_function, **kwargs - ) - return await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) - - @staticmethod - @contextlib.contextmanager - def _video_call_pin(handler_name: str | None, request_kwargs: Mapping[str, object]) -> Generator[None]: - """ - Scopes the deployment a video id names to the routing of one call on that video (status, content, remix, - edit, extension). Any other call routes as if no id were sent, whatever its body carries. - """ - video_id: Final = request_kwargs.get("video_id") - deployment_id: Final = ( - decode_video_id_with_provider(video_id).get("deployment_id") - if handler_name in VIDEO_ID_CALL_TYPES and isinstance(video_id, str) - else None - ) - token: Final = _video_call_deployment_id.set(deployment_id) - try: - yield - finally: - _video_call_deployment_id.reset(token) - - @staticmethod - @contextlib.contextmanager - def _route_capability_pin(handler_name: str | None) -> Generator[None]: - token: Final = _routed_call_type.set(handler_name) - try: - yield - finally: - _routed_call_type.reset(token) - - @staticmethod - def _narrow_to_video_deployment(healthy_deployments: Sequence[_DeploymentT]) -> Sequence[_DeploymentT]: - """ - A video lives on the provider account of the deployment that created it, so a call on an id naming a - deployment that every routing filter kept goes to that deployment, else to whichever one routing picks. - """ - deployment_id: Final = _video_call_deployment_id.get() - creating: Final = tuple( - deployment - for deployment in healthy_deployments - if deployment_id in PythonCore._deployment_ids((deployment,)) - ) - return creating or healthy_deployments - - def _generic_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs): - """ - Make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router - Args: - model: The model to use - original_function: The handler function to call (e.g., litellm.completion) - **kwargs: Additional arguments to pass to the handler function - Returns: - The response from the handler function - """ - handler_name: Final = original_function.__name__ - metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call") - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): - try: - verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) - self._update_kwargs_before_fallbacks( - model=model, - kwargs=kwargs, - metadata_variable_name=metadata_variable_name, - ) - with ( - self._video_call_pin(handler_name, kwargs), - self._route_capability_pin(handler_name), - picks_reserve(True), - ): - deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="generic_api_call", - splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, - ) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - - self.total_calls[model_name] += 1 - - # For passthrough routes, use the actual model from deployment - # and swap model name in endpoint if present - if "endpoint" in kwargs and kwargs["endpoint"]: - kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name) - kwargs["model"] = model_name - - # Perform pre-call checks for routing strategy - self.routing_strategy_pre_call_checks(deployment=deployment) - - custom_llm_provider: Final = provider_for_generic_call(data) - - response: Final = original_function( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name) - return response - except Exception as e: - verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - def embedding( - self, - model: str, - input: str | list, - is_async: bool | None = False, - **kwargs, - ) -> EmbeddingResponse: - try: - kwargs["model"] = model - kwargs["input"] = input - kwargs["original_function"] = self._embedding - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = self.function_with_fallbacks(**kwargs) - return response - except Exception as e: - raise e - - def _embedding(self, input: str | list[object], model: str, **kwargs): - model_name = None - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside embedding()- model: %s", model) - with embedding_deployment_selection(), picks_reserve(True): - deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - input=input, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_name = data["model"] - batch_error: Final = embedding_batch_too_large_error( - deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment), - input=input, - request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs), - model=model, - ) - logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") - if batch_error is not None: - if isinstance(logging_obj, LiteLLMLogging): - log_rejection_before_provider_call(logging_obj, batch_error) - raise batch_error - if isinstance(logging_obj, LiteLLMLogging): - clear_rejected_before_provider_call(logging_obj) - - potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs, client_type="sync") - # check if provided keys == client keys # - dynamic_api_key: Final = kwargs.get("api_key", None) - if ( - dynamic_api_key is not None - and potential_model_client is not None - and dynamic_api_key != potential_model_client.api_key - ): - model_client = None - else: - model_client = potential_model_client - - self.total_calls[model_name] += 1 - - ### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit) - self.routing_strategy_pre_call_checks(deployment=deployment) - - response: Final = self._host.invoke("embedding")( - **{ - **data, - "input": input, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.embedding(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.embedding(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def aembedding( - self, - model: str, - input: str | list, - is_async: bool | None = True, - **kwargs, - ) -> EmbeddingResponse: - try: - kwargs["model"] = model - kwargs["input"] = input - kwargs["original_function"] = self._aembedding - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _aembedding(self, input: str | list[object], model: str, **kwargs): - model_name = None - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _aembedding()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - with embedding_deployment_selection(): - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - input=input, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_name = data["model"] - batch_error: Final = embedding_batch_too_large_error( - deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment), - input=input, - request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs), - model=model, - ) - logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") - if batch_error is not None: - if isinstance(logging_obj, LiteLLMLogging): - await async_log_rejection_before_provider_call(logging_obj, batch_error) - raise batch_error - if isinstance(logging_obj, LiteLLMLogging): - clear_rejected_before_provider_call(logging_obj) - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("aembedding")( - **{ - **data, - "input": input, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.aembedding(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.aembedding(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - #### FILES API #### - async def acreate_file( - self, - model: str, - **kwargs, - ) -> OpenAIFileObject: - try: - kwargs["model"] = model - kwargs["original_function"] = self._acreate_file - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _acreate_file( - self, - model: str, - **kwargs, - ) -> OpenAIFileObject: - try: - from litellm.litellm_core_utils.routing_common_utils import add_model_file_id_mappings - - verbose_router_logger.debug("Inside _atext_completion()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - healthy_deployments: Final = await self._core_async_get_healthy_deployments( - model=model, - messages=[{"role": "user", "content": "files-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - parent_otel_span=parent_otel_span, - ) - - async def create_file_for_deployment(deployment: dict) -> OpenAIFileObject: - from litellm.litellm_core_utils.core_helpers import safe_deep_copy - - kwargs_copy: Final = safe_deep_copy(kwargs) - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs_copy, - function_name="acreate_file", - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs_copy, - ) - self.total_calls[model_name] += 1 - - ## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ## - # For DB/config deployments, use provider from deployment params - custom_llm_provider = data.get("custom_llm_provider") - stripped_model, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - # Preserve explicitly stored provider, fallback to inferred - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - ## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ## - purpose: Final = cast(OpenAIFilesPurpose | None, kwargs.get("purpose")) - file = cast(FileTypes | None, kwargs.get("file")) - if not file or not purpose: - raise Exception("file and file_purpose are required for create_file") - - replace_model_in_jsonl_bool: Final = should_replace_model_in_jsonl( - purpose=purpose, - passthrough=kwargs.get("passthrough") is True, - ) - if replace_model_in_jsonl_bool: - file = replace_model_in_jsonl( - file_content=file, - new_model_name=stripped_model, - keep_line_numbers=custom_llm_provider == LlmProviders.VERTEX_AI.value, - ) - - kwargs_copy["file"] = file - if custom_llm_provider == LlmProviders.LITELLM_PROXY.value: - kwargs_copy["extra_body"] = MappingProxyType( - { - **(kwargs_copy.get("extra_body") or MappingProxyType({})), - "target_model_names": stripped_model, - } - ) - if ( - "gcs_bucket_name" in data - ): # TODO: Remove this once we have a better way to handle GCS bucket name: Problem is that we need to pass the gcs_bucket_name to the router for the create_file call but it doesn't show up there - kwargs_copy.setdefault("litellm_metadata", {})["gcs_bucket_name"] = data["gcs_bucket_name"] - async with self._deployment_slot( - deployment=deployment, kwargs=kwargs_copy, parent_otel_span=parent_otel_span - ): - response = await self._host.invoke("acreate_file")( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs_copy, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acreate_file(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - - tasks: Final = [] - - if isinstance(healthy_deployments, dict): - tasks.append(create_file_for_deployment(healthy_deployments)) - else: - for deployment in healthy_deployments: - tasks.append(create_file_for_deployment(deployment)) - - responses: Final = await asyncio.gather(*tasks) - - if len(responses) == 0: - raise Exception("No healthy deployments found.") - - model_file_id_mapping: Final = add_model_file_id_mappings( - healthy_deployments=healthy_deployments, responses=responses - ) - returned_response: Final = cast(OpenAIFileObject, responses[0]) - returned_response_hidden_params: Final = get_hidden_params(returned_response) - if returned_response_hidden_params is not None: - returned_response_hidden_params["model_file_id_mapping"] = model_file_id_mapping - return returned_response - except Exception as e: - verbose_router_logger.exception( - "litellm.acreate_file(model=%s, %s)\x1b[31m Exception %s\x1b[0m", model, kwargs, e - ) - if model is not None: - self.fail_calls[model] += 1 - raise e - - #### VECTOR STORES API #### - async def avector_store_create( - self, - model: str | None, - **kwargs, - ): - """ - Create a vector store for a specific model. - - Args: - model: Model name from router config - **kwargs: Vector store creation parameters - - Returns: - VectorStoreCreateResponse - """ - with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): - try: - # If model is None, use the factory function approach (direct SDK call) - if model is None: - from litellm.vector_stores.main import acreate - - # Use the factory function to handle the call - factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") - return await factory_fn(**kwargs) - - from litellm.vector_stores import acreate as avector_store_create_sdk - - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "vector-store-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="avector_store_create", - ) - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - # Get custom provider from deployment params - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot( - deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span - ): - response = await avector_store_create_sdk( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - except Exception as e: - verbose_router_logger.exception( - "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e - ) - if model is not None: - self.fail_calls[model] += 1 - raise e - - def _override_vector_store_methods_for_router(self): - """ - Override factory-generated vector store methods with router-aware implementations. - This is called after _initialize_vector_store_endpoints() to ensure our custom - methods that handle deployment selection and credential injection are used instead - of the generic factory-generated ones. - """ - # Store references to the custom methods defined above - # These methods handle proper routing through deployments - # The methods are already defined as instance methods above - - async def acreate_batch( - self, - model: str, - **kwargs, - ) -> LiteLLMBatch: - try: - kwargs["model"] = model - kwargs["original_function"] = self._acreate_batch - metadata_variable_name: Final = get_router_metadata_variable_name(function_name="_acreate_batch") - self._update_kwargs_before_fallbacks( - model=model, - kwargs=kwargs, - metadata_variable_name=metadata_variable_name, - ) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _acreate_batch( - self, - model: str, - **kwargs, - ) -> LiteLLMBatch: - try: - verbose_router_logger.debug("Inside _acreate_batch()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - with batch_create_routing(): - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "files-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="_acreate_batch") - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ## - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("acreate_batch")( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acreate_batch(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - except Exception as e: - verbose_router_logger.exception( - "litellm._acreate_batch(model=%s, %s)\x1b[31m Exception %s\x1b[0m", model, kwargs, e - ) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - raise e - - async def aretrieve_batch( - self, - model: str | None = None, - **kwargs, - ) -> LiteLLMBatch: - """ - Iterate through all models in a model group to check for batch - - Future Improvement - cache the result. - """ - try: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - requested_model_group: Final = model - metadata_variable_name: Final = get_router_metadata_variable_name(function_name="aretrieve_batch") - if model is not None: - filtered_model_list: ( - list[DeploymentTypedDict] | list[dict] | dict | None - ) = await self._core_async_get_healthy_deployments( - model=model, - messages=[{"role": "user", "content": "retrieve-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - parent_otel_span=parent_otel_span, - ) - else: - filtered_model_list = self._core_get_model_list() - if filtered_model_list is None: - raise Exception("Router not yet initialized.") - - receieved_exceptions: Final = [] - - async def try_retrieve_batch(model_name: DeploymentTypedDict): - try: - from litellm.litellm_core_utils.core_helpers import safe_deep_copy - - model: Final = model_name["litellm_params"].get("model") - data: Final = model_name["litellm_params"].copy() - custom_llm_provider = data.get("custom_llm_provider") - if model is None: - raise Exception(f"Model not found in litellm_params for deployment: {model_name}") - # Update kwargs with the current model name or any other model-specific adjustments - ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ## - if not custom_llm_provider: - _, custom_llm_provider, _, _ = get_llm_provider(model=model) - new_kwargs: Final = safe_deep_copy(kwargs) - self._update_kwargs_with_deployment( - deployment=cast(dict, model_name), - kwargs=new_kwargs, - function_name="aretrieve_batch", - ) - model_group: Final = requested_model_group or model_name["model_name"] - if not new_kwargs[metadata_variable_name].get("model_group"): - new_kwargs[metadata_variable_name]["model_group"] = model_group - new_kwargs.pop("custom_llm_provider", None) - data.pop("custom_llm_provider", None) - return await self._host.invoke("aretrieve_batch")( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - **new_kwargs, - }, - ) - except Exception as e: - import traceback - - traceback.print_exc() - receieved_exceptions.append(e) - return None - - # Check all models in parallel - if ( - filtered_model_list is not None - and isinstance(filtered_model_list, list) - and len(filtered_model_list) > 0 - ): - results = await asyncio.gather( - *[try_retrieve_batch(cast(DeploymentTypedDict, model)) for model in filtered_model_list], - return_exceptions=True, - ) - elif filtered_model_list is not None and isinstance(filtered_model_list, dict): - results = await try_retrieve_batch(cast(DeploymentTypedDict, filtered_model_list)) - else: - raise Exception("No healthy deployments found.") - - # Check for successful responses and handle exceptions - if results is not None: - if isinstance(results, LiteLLMBatch): - return results - elif isinstance(results, list): - for result in results: - if isinstance(result, LiteLLMBatch): - return result - - # If no valid Batch response was found, raise the first encountered exception - if receieved_exceptions: - raise receieved_exceptions[0] # Raising the first exception encountered - - # If no exceptions were encountered, raise a generic exception - raise Exception(f"Unable to find batch in any model. Received errors - {receieved_exceptions}") - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def acancel_batch( - self, - model: str, - **kwargs, - ) -> LiteLLMBatch: - """ - Cancel a batch through the router with proper model-to-provider mapping. - """ - try: - kwargs["model"] = model - kwargs["original_function"] = self._acancel_batch - metadata_variable_name: Final = get_router_metadata_variable_name(function_name="_acancel_batch") - self._update_kwargs_before_fallbacks( - model=model, - kwargs=kwargs, - metadata_variable_name=metadata_variable_name, - ) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _acancel_batch( - self, - model: str, - **kwargs, - ) -> LiteLLMBatch: - try: - verbose_router_logger.debug("Inside _acancel_batch()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - with picks_reserve(False): - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "batch-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - selected_deployment_id: Final = (deployment.get("model_info") or {}).get("id") - data: Final = deployment["litellm_params"].copy() - resolved_credentials: Final = self._core_get_deployment_credentials_with_provider( - model_id=selected_deployment_id or model - ) - if resolved_credentials is not None: - data.update(resolved_credentials) - data.pop("litellm_credential_name", None) - model_name: Final = data["model"] - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="_acancel_batch") - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ## - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("acancel_batch")( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acancel_batch(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - except Exception as e: - verbose_router_logger.exception( - "litellm._acancel_batch(model=%s, %s)\x1b[31m Exception %s\x1b[0m", model, kwargs, e - ) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - raise e - - async def alist_batches( - self, - model: str, - **kwargs, - ): - """ - Return all the batches across all deployments of a model group. - """ - - filtered_model_list: Final = self._core_get_model_list(model_name=model) - if filtered_model_list is None: - raise Exception("Router not yet initialized.") - - async def try_retrieve_batch(model: DeploymentTypedDict): - try: - # Update kwargs with the current model name or any other model-specific adjustments - return await self._host.invoke("alist_batches")(**{**model["litellm_params"], **kwargs}) - except Exception: - return None - - # Check all models in parallel - results: Final = await asyncio.gather(*[try_retrieve_batch(model) for model in filtered_model_list]) - - final_results: Final[dict] = { - "object": "list", - "data": [], - "first_id": None, - "last_id": None, - "has_more": False, - } - - for result in results: - if result is not None: - ## check batch id - if final_results["first_id"] is None and hasattr(result, "first_id"): - final_results["first_id"] = getattr(result, "first_id") - final_results["last_id"] = getattr(result, "last_id") - final_results["data"].extend(result.data) - - ## check 'has_more' - if getattr(result, "has_more", False) is True: - final_results["has_more"] = True - - return final_results - - #### PASSTHROUGH API #### - - async def _pass_through_moderation_endpoint_factory( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - **kwargs, - ): - # update kwargs with model_group - self._update_kwargs_before_fallbacks( - model=kwargs.get("model", ""), - kwargs=kwargs, - ) - if kwargs.get("model") and self._core_get_model_list(model_name=kwargs["model"]): - deployment: Final = await self._core_async_get_available_deployment( - model=kwargs["model"], - request_kwargs=kwargs, - ) - kwargs["model"] = deployment["litellm_params"]["model"] - data: Final = deployment["litellm_params"].copy() - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - ) - kwargs.update(data) - - return await original_function(**kwargs) - - def factory_function( - self, - original_function: Callable, - call_type: Literal[ - "assistants", - "moderation", - "anthropic_messages", - "aresponses", - "acancel_responses", - "acompact_responses", - "responses", - "aget_responses", - "adelete_responses", - "afile_delete", - "afile_content", - "_arealtime", - "acreate_realtime_client_secret", - "arealtime_calls", - "acreate_realtime_transcription_session", - "acreate_realtime_translation_client_secret", - "arealtime_translation_calls", - "_aresponses_websocket", - "acreate_fine_tuning_job", - "acancel_fine_tuning_job", - "alist_fine_tuning_jobs", - "aretrieve_fine_tuning_job", - "alist_files", - "aimage_edit", - "allm_passthrough_route", - "alist_input_items", - "agenerate_content", - "generate_content", - "agenerate_content_stream", - "generate_content_stream", - "avector_store_search", - "avector_store_create", - "avector_store_retrieve", - "avector_store_list", - "avector_store_update", - "avector_store_delete", - "avector_store_file_create", - "avector_store_file_list", - "avector_store_file_retrieve", - "avector_store_file_content", - "avector_store_file_update", - "avector_store_file_delete", - "vector_store_search", - "vector_store_create", - "vector_store_retrieve", - "vector_store_list", - "vector_store_update", - "vector_store_delete", - "vector_store_file_create", - "vector_store_file_list", - "vector_store_file_retrieve", - "vector_store_file_content", - "vector_store_file_update", - "vector_store_file_delete", - "aocr", - "ocr", - "asearch", - "search", - "adecisions", - "decisions", - "asystemone", - "systemone", - "aextraction", - "extraction", - "aadapter_generate_content", - "avideo_generation", - "video_generation", - "avideo_list", - "video_list", - "avideo_status", - "video_status", - "avideo_content", - "video_content", - "avideo_remix", - "video_remix", - "avideo_create_character", - "video_create_character", - "avideo_get_character", - "video_get_character", - "avideo_edit", - "video_edit", - "avideo_extension", - "video_extension", - "acreate_container", - "create_container", - "alist_containers", - "list_containers", - "aretrieve_container", - "retrieve_container", - "adelete_container", - "delete_container", - "aupload_container_file", - "upload_container_file", - "alist_container_files", - "list_container_files", - "aretrieve_container_file", - "retrieve_container_file", - "adelete_container_file", - "delete_container_file", - "acreate_skill", - "alist_skills", - "aget_skill", - "adelete_skill", - "acreate_interaction", - "create_interaction", - "aget_interaction", - "get_interaction", - "adelete_interaction", - "delete_interaction", - "acancel_interaction", - "cancel_interaction", - "acreate_agent", - "create_agent", - "alist_agents", - "list_agents", - "aget_agent", - "get_agent", - "adelete_agent", - "delete_agent", - "alist_agent_versions", - "list_agent_versions", - ] = "assistants", - ): - """ - Creates appropriate wrapper functions for different API call types. - - Returns: - - A synchronous function for synchronous call types - - An asynchronous function for asynchronous call types - """ - # Handle synchronous call types - if call_type in ( - "responses", - "generate_content", - "generate_content_stream", - "ocr", - "search", - "decisions", - "systemone", - "extraction", - "video_generation", - "video_list", - "video_status", - "video_content", - "video_remix", - "create_container", - "list_containers", - "retrieve_container", - "delete_container", - ): - - def sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs) - - return sync_wrapper - - if call_type in ( - "vector_store_search", - "vector_store_create", - "vector_store_retrieve", - "vector_store_list", - "vector_store_update", - "vector_store_delete", - ): - - def vector_store_sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - provider_kwargs: Final = ( - MappingProxyType({**kwargs, "custom_llm_provider": custom_llm_provider}) - if custom_llm_provider and "custom_llm_provider" not in kwargs - else MappingProxyType(kwargs) - ) - search_kwargs: Final = ( - MappingProxyType( - { - **provider_kwargs, - "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( - router=self._core_public(), - metadata=self._vector_store_request_metadata(kwargs), - ), - } - ) - if call_type == "vector_store_search" - else provider_kwargs - ) - if search_kwargs.get("model"): - return self._generic_api_call_with_fallbacks(original_function=original_function, **search_kwargs) - if call_type == "vector_store_search": - return original_function(**MappingProxyType({**search_kwargs, "router": self._core_public()})) - return original_function(**search_kwargs) - - return vector_store_sync_wrapper - - if call_type in ( - "vector_store_file_create", - "vector_store_file_list", - "vector_store_file_retrieve", - "vector_store_file_content", - "vector_store_file_update", - "vector_store_file_delete", - ): - - def vector_store_file_sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - return original_function( - custom_llm_provider=custom_llm_provider, - client=client, - **kwargs, - ) - - return vector_store_file_sync_wrapper - - if call_type in ( - "create_agent", - "list_agents", - "get_agent", - "delete_agent", - "list_agent_versions", - ): - - def managed_agents_sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - if "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = "gemini" - return original_function(**kwargs) - - return managed_agents_sync_wrapper - - # Handle asynchronous call types - async def async_wrapper( - custom_llm_provider: str | None = None, - client: AsyncOpenAI | None = None, - **kwargs, - ): - if call_type == "assistants": - return await self._pass_through_assistants_endpoint_factory( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - client=client, - **kwargs, - ) - elif call_type == "moderation": - return await self._pass_through_moderation_endpoint_factory( - original_function=original_function, **kwargs - ) - elif call_type in ("asearch", "search"): - return await self._asearch_with_fallbacks( - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "avector_store_file_create", - "avector_store_file_list", - "avector_store_file_retrieve", - "avector_store_file_content", - "avector_store_file_update", - "avector_store_file_delete", - ): - return await self._init_vector_store_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - elif call_type == "aresponses": - return await self._aresponses_with_streaming_fallbacks( - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "acreate_realtime_client_secret", - "arealtime_calls", - "acreate_realtime_transcription_session", - "acreate_realtime_translation_client_secret", - "arealtime_translation_calls", - ): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - client=client, - **kwargs, - ) - elif call_type in ( - "anthropic_messages", - "_arealtime", - "_aresponses_websocket", - "acreate_fine_tuning_job", - "acancel_fine_tuning_job", - "alist_fine_tuning_jobs", - "aretrieve_fine_tuning_job", - "alist_files", - "aimage_edit", - "agenerate_content", - "agenerate_content_stream", - "aocr", - "ocr", - "adecisions", - "asystemone", - "aextraction", - "avideo_generation", - "avideo_list", - "avideo_status", - "avideo_content", - "avideo_remix", - "avideo_create_character", - "avideo_get_character", - "avideo_edit", - "avideo_extension", - "acreate_skill", - "alist_skills", - "aget_skill", - "adelete_skill", - ): - return await self._dispatch_generic_call_type( - call_type=call_type, - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "acreate_container", - "alist_containers", - "aretrieve_container", - "adelete_container", - "aupload_container_file", - "alist_container_files", - "aretrieve_container_file", - "adelete_container_file", - "aretrieve_container_file_content", - ): - return await self._init_containers_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - elif call_type == "allm_passthrough_route": - if client: - kwargs["client"] = client - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - passthrough_on_no_deployment=True, - **kwargs, - ) - elif call_type in ( - "aget_responses", - "acancel_responses", - "acompact_responses", - "adelete_responses", - "alist_input_items", - ): - return await self._init_responses_api_endpoints( - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "avector_store_search", - "avector_store_create", - "avector_store_retrieve", - "avector_store_list", - "avector_store_update", - "avector_store_delete", - ): - vector_store_kwargs: Final = ( - { - **kwargs, - "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( - router=self._core_public(), - metadata=self._vector_store_request_metadata(kwargs), - ), - } - if call_type == "avector_store_search" - else kwargs - ) - return await self._init_vector_store_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - call_type=call_type, - **vector_store_kwargs, - ) - elif call_type in ("afile_delete", "afile_content"): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - client=client, - **kwargs, - ) - elif call_type in ( - "acreate_interaction", - "create_interaction", - "aget_interaction", - "adelete_interaction", - "acancel_interaction", - ): - return await self._init_interactions_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - elif call_type in ( - "acreate_agent", - "alist_agents", - "aget_agent", - "adelete_agent", - "alist_agent_versions", - ): - return await self._init_managed_agents_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - - return async_wrapper - - @staticmethod - def _vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: - return vector_store_request_metadata(kwargs) - - async def _init_vector_store_api_endpoints( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - call_type: str | None = None, - **kwargs, - ): - """ - Initialize the Vector Store API endpoints on the router. - - If a model is provided in kwargs, use model-based routing to get - the deployment credentials. Otherwise, call the original function directly. - """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - - # If model is provided, use generic API call with fallbacks for proper routing - if kwargs.get("model"): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - **kwargs, - ) - - # For search, pass the router so provider transforms can resolve - # router-managed embedding models (e.g. S3 Vectors query embeddings). - # The merge also overrides any client-supplied `router` key. - if call_type == "avector_store_search": - search_kwargs: Final = MappingProxyType({**kwargs, "router": self._core_public()}) - return await original_function(**search_kwargs) - - # Otherwise, call the original function directly - return await original_function(**kwargs) - - async def _init_containers_api_endpoints( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - **kwargs, - ): - """ - Initialize the Containers API endpoints on the router. - - LiteLLM-managed container IDs (``cntr_...``) encode ``model_id`` and provider - metadata. When present, decode the ID, replace ``container_id`` with the - upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so - deployment credentials (e.g. regional ``api_base`` for Azure) match - :meth:`_init_responses_api_endpoints`. Create/list calls carry no container ID, so - they route through the deployment named by ``model`` when the caller passes one, - falling back to the direct call when no deployment matches. Otherwise call the - handler directly with global provider credentials. - """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - - from litellm.responses.utils import ResponsesAPIRequestUtils - - container_id: Final = kwargs.get("container_id") - _forwarded_model_id: Final = kwargs.get("model_id") - if isinstance(container_id, str): - decoded: Final = ResponsesAPIRequestUtils.decode_container_id(container_id) - original_id: Final = decoded.get("response_id", container_id) - if original_id != container_id: - kwargs["container_id"] = original_id - decoded_provider: Final = decoded.get("custom_llm_provider") - if decoded_provider and kwargs.get("custom_llm_provider") == "openai": - kwargs["custom_llm_provider"] = decoded_provider - # Fall back to the model_id forwarded by the proxy when the container_id - # is a native upstream ID (e.g. Azure hex cntr_) that carries no LiteLLM - # routing payload, so deployment credentials (api_base, api_key) are applied. - model_id: Final = decoded.get("model_id") or ( - _forwarded_model_id.strip() - if isinstance(_forwarded_model_id, str) and _forwarded_model_id.strip() - else None - ) - if model_id: - kwargs["model"] = model_id - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - **kwargs, - ) - - requested_model: Final = kwargs.get("model") - if isinstance(requested_model, str) and requested_model.strip(): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - passthrough_on_no_deployment=True, - **kwargs, - ) - - return await original_function(**kwargs) - - async def _init_responses_api_endpoints( - self, - original_function: Callable, - **kwargs, - ): - """ - Initialize the Responses API endpoints on the router. - - GET, DELETE, CANCEL Responses API Requests encode the model_id in the response_id, this function decodes the response_id and sets the model to the model_id. - """ - from litellm.responses.utils import ResponsesAPIRequestUtils - - model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(kwargs.get("response_id")) - if model_id is not None: - kwargs["model"] = model_id - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - **kwargs, - ) - - async def _init_interactions_api_endpoints( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - **kwargs, - ): - """ - Initialize the Interactions API endpoints on the router. - - GET, DELETE, CANCEL Interactions API Requests don't need model-based routing, - so we call the original function directly with the custom_llm_provider. - """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - # Default to gemini for interactions API - if "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = "gemini" - # If the proxy accidentally passed agent name as model, clear it - if kwargs.get("agent") and kwargs.get("model") == kwargs.get("agent"): - kwargs["model"] = None - # Model-based interactions use deployment routing + fallbacks; agent-only calls - # must not enter model-group lookup (agent name is not a LiteLLM deployment). - if kwargs.get("model"): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - **kwargs, - ) - return await original_function(**kwargs) - - async def _init_managed_agents_api_endpoints( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - **kwargs, - ): - """ - Initialize the Managed Agents API endpoints on the router (v1beta/agents). - - CRUD operations for Gemini managed agents don't need model-based routing, - so we call the original function directly with the custom_llm_provider. - """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - if "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = "gemini" - return await original_function(**kwargs) - - async def _pass_through_assistants_endpoint_factory( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - client: AsyncOpenAI | None = None, - **kwargs, - ): - """Internal helper function to pass through the assistants endpoint""" - if custom_llm_provider is None: - if self.assistants_config is not None: - custom_llm_provider = self.assistants_config["custom_llm_provider"] - kwargs.update(self.assistants_config["litellm_params"]) - else: - raise Exception( - "'custom_llm_provider' must be set. Either via:\n `Router(assistants_config={'custom_llm_provider': ..})` \nor\n `router.arun_thread(custom_llm_provider=..)`" - ) - return await original_function(custom_llm_provider=custom_llm_provider, client=client, **kwargs) - - #### [END] ASSISTANTS API #### - - async def _maybe_run_weighted_failover( - self, - exception: Exception, - original_model_group: str, - all_deployments: Sequence[DeploymentTypedDict], - args: tuple, - kwargs: dict, - input_kwargs: dict, - ) -> Any | None: - """Same-model-group retry after a failed deployment; returns None if not applicable.""" - strategy, _ = self._get_routing_context(original_model_group, kwargs) - if strategy != "simple-shuffle": - return None - - failed_id: Final[str | None] = getattr(exception, "failed_deployment_id", None) - if not failed_id: - return None - - metadata_variable_name: Final = self._get_metadata_variable_name_from_kwargs(kwargs) - meta = kwargs.get(metadata_variable_name) - if meta is None: - meta = {} - kwargs[metadata_variable_name] = meta - if not isinstance(meta, dict): - return None - configured_failed_id: Final[str | None] = getattr(exception, "failed_configured_deployment_id", None) - prev_excluded: Final = set(meta.get("_failover_excluded_ids") or []) - excluded: Final = prev_excluded | {configured_failed_id or failed_id} - - all_ids: Final = { - (d.get("model_info") or {}).get("id") - for d in all_deployments - if (d.get("model_info") or {}).get("id") is not None - } - # Only consider deployments that are currently healthy (not in cooldown). - # Using all_ids here would cause a wasteful run_async_fallback invocation - # that fails with RouterRateLimitError whenever the "remaining" entries - # are all in cooldown — the inner async_get_healthy_deployments call - # would find an empty list and raise immediately. - cooldown_ids = set(await async_get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=None)) - remaining: Final = (all_ids - cooldown_ids) - excluded - if not remaining: - return None - - verbose_router_logger.debug( - "Weighted failover: exclude=%r, remaining=%s for model_group=%r", - excluded, - len(remaining), - original_model_group, - ) - - meta["_failover_excluded_ids"] = list(excluded) - - entry: Final = { - "model": original_model_group, - "_excluded_deployment_ids": list(excluded), - } - # Build a local copy so the weighted-failover keys do not leak back to - # the caller's shared kwargs dict (any downstream fallback path reads - # the same dict and must not inherit our `_excluded_deployment_ids` - # entry). - failover_kwargs: Final = { - **input_kwargs, - "fallback_model_group": [entry], - "original_model_group": original_model_group, - } - try: - return await run_async_fallback(*args, **failover_kwargs) - except (openai.APIError, RouterRateLimitError, RouterRateLimitErrorBasic): - # Expected model-level failure on the retried deployment. All - # litellm provider errors derive from openai.APIError; if every - # remaining deployment in the group is in cooldown the router - # raises RouterRateLimitError (a ValueError, not an APIError). - # In either case defer to the regular fallback path. Programming - # errors (AttributeError, KeyError, TypeError, etc.) intentionally - # propagate so they remain visible. - return None - - async def async_function_with_fallbacks_common_utils( - self, - e: Exception, - disable_fallbacks: bool | None, - fallbacks: list | None, - context_window_fallbacks: list | None, - content_policy_fallbacks: list | None, - model_group: str | None, - args: tuple, - kwargs: dict, - include_fallback_errors: bool = False, - ): - """ - Common utilities for async_function_with_fallbacks - """ - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("Traceback%s", redact_string(traceback.format_exc())) - original_exception: Final = e - fallback_model_group = None - original_model_group: Final[str | None] = kwargs.get("model") - fallback_walk: Final = attempted_fallback_targets_from(kwargs) # pyright: ignore[reportUnknownArgumentType] # keep legacy kwargs untyped - # A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier - # behind the router name, and fallbacks are configured per tier, not per router. - lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group) - fallback_failure_exception_str = "" - no_fallback_group_explained = False - hop_depth: Final = kwargs.get("fallback_depth") - nested_fallback_hop: Final = isinstance(hop_depth, int) and hop_depth > 0 - - unsupported_params_without_fallback: Final = is_bad_request_without_fallback( - e, self.fallback_on_unsupported_params - ) - if ( - disable_fallbacks is True - or original_model_group is None - or is_guardrail_intervention(e) - or unsupported_params_without_fallback - ): - raise e - - input_kwargs: Final = { - "litellm_routing_ctx": self._core_ctx, - "original_exception": original_exception, - **kwargs, - "attempted_targets": fallback_walk, - } - - if "max_fallbacks" not in input_kwargs: - input_kwargs["max_fallbacks"] = self.max_fallbacks - if "fallback_depth" not in input_kwargs: - input_kwargs["fallback_depth"] = 0 - if include_fallback_errors: - input_kwargs["include_fallback_errors"] = True - - context_window_fallback_model_group: Final = ( - self._get_fallback_model_group_for_lookup_groups( - fallbacks=context_window_fallbacks, - lookup_groups=lookup_groups, - ) - if isinstance(e, litellm_exceptions.ContextWindowExceededError) and context_window_fallbacks is not None - else None - ) - _skip_order_fallback: Final = ( - isinstance(e, litellm_exceptions.ContentPolicyViolationError) - or context_window_fallback_model_group is not None - ) - _request_team_id: Final[str | None] = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id") - # Use wildcard-aware lookup so order-based fallback also works for model - # groups resolved via pattern routing (e.g. `openai/*` -> `openai/gpt-4.1-mini`). - order_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group - all_deployments: Final = self._core_get_model_list(model_name=order_model_group, team_id=_request_team_id) or () - _order_set: Final[set] = { - litellm_utils.get_deployment_order(d) - for d in all_deployments - if litellm_utils.get_deployment_order(d) is not None - } - order_values: Final[list] = sorted(_order_set) - if len(order_values) > 1 and not _skip_order_fallback: - # Determine which order levels have already been tried - current_target: Final = kwargs.get("_target_order") - skip_up_to: Final = current_target if current_target is not None else order_values[0] - # Build order-based fallback entries (skip already-tried levels) - order_fallback_entries: Final[list] = [ - {"model": order_model_group, "_target_order": o} for o in order_values if o > skip_up_to - ] - # Get external fallbacks — handle both standard and non-standard formats - external_fallback_group: list | None = None - if fallbacks is not None and lookup_groups: - if check_non_standard_fallback_format(fallbacks=fallbacks): - # Non-standard formats (e.g. ["claude-3-haiku"] or - # [{"model": "...", "messages": [...]}]) are passed through directly - external_fallback_group = fallbacks - else: - external_fallback_group, generic_idx = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, - lookup_groups=lookup_groups, - ) - if external_fallback_group is None and generic_idx is not None: - external_fallback_group = fallbacks[generic_idx]["*"] - - # Combined list: order fallbacks first, then external - combined_fallbacks: Final = order_fallback_entries + (external_fallback_group or []) - - if combined_fallbacks: - input_kwargs.update( - { - "fallback_model_group": combined_fallbacks, - "original_model_group": original_model_group, - } - ) - try: - return await run_async_fallback( - *args, - **input_kwargs, - ) - except Exception: - if isinstance(original_exception, litellm_exceptions.ContextWindowExceededError): - raise original_exception - raise - - # Weighted intra-group failover (simple-shuffle only); see _maybe_run_weighted_failover. - if ( - self.enable_weighted_failover - and not _skip_order_fallback - and not isinstance(e, litellm_exceptions.ContextWindowExceededError) - and original_model_group is not None - ): - response = await self._maybe_run_weighted_failover( - exception=e, - original_model_group=original_model_group, - all_deployments=all_deployments, - args=args, - kwargs=kwargs, - input_kwargs=input_kwargs, - ) - if response is not None: - return response - - try: - verbose_router_logger.info("Trying to fallback b/w models") - - # check if client-side fallbacks are used (e.g. fallbacks = ["gpt-3.5-turbo", "claude-3-haiku"] or fallbacks=[{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hey, how's it going?"}]}] - is_non_standard_fallback_format: Final = check_non_standard_fallback_format(fallbacks=fallbacks) - - if is_non_standard_fallback_format and context_window_fallback_model_group is None: - input_kwargs.update( - { - "fallback_model_group": fallbacks, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - - return response - - if isinstance(e, litellm_exceptions.ContextWindowExceededError): - if context_window_fallback_model_group is not None: - input_kwargs.update( - { - "fallback_model_group": context_window_fallback_model_group, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - return response - - elif context_window_fallbacks is None: - error_message = f"model={model_group}. context_window_fallbacks={mask_sensitive_structure(context_window_fallbacks)}. fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" - verbose_router_logger.info( - msg=f"Got 'ContextWindowExceededError'. No context_window_fallback set. Defaulting \ - to fallbacks, if available.{error_message}" - ) - - if self._host.dynamic().expose_router_debug_in_errors: - e.message += f"\n{error_message}" - elif isinstance(e, litellm_exceptions.ContentPolicyViolationError): - if content_policy_fallbacks is not None: - content_policy_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_for_lookup_groups( - fallbacks=content_policy_fallbacks, - lookup_groups=lookup_groups, - ) - ) - if content_policy_fallback_model_group is None: - raise original_exception - - input_kwargs.update( - { - "fallback_model_group": content_policy_fallback_model_group, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - return response - else: - error_message = format_content_policy_fallback_hint( - model_group, content_policy_fallbacks, fallbacks - ) - verbose_router_logger.info( - msg=f"Got 'ContentPolicyViolationError'. No content_policy_fallback set. Defaulting \ - to fallbacks, if available.{error_message}" - ) - - if self._host.dynamic().expose_router_debug_in_errors: - e.message += f"\n{error_message}" - if fallbacks is not None and lookup_groups: - verbose_router_logger.debug("inside model fallbacks: %s", mask_sensitive_structure(fallbacks)) - ( - fallback_model_group, - generic_fallback_idx, - ) = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}] - lookup_groups=lookup_groups, - ) - ## if none, check for generic fallback - if fallback_model_group is None and generic_fallback_idx is not None: - fallback_model_group = fallbacks[generic_fallback_idx]["*"] - - if fallback_model_group is None: - masked_fallbacks: Final = mask_sensitive_structure(fallbacks) - verbose_router_logger.info( - "No fallback model group found for lookup_groups=%s. Fallbacks=%s", - " -> ".join(lookup_groups), - masked_fallbacks, - ) - if ( - hasattr(original_exception, "message") - and self._host.dynamic().expose_router_debug_in_errors - and not nested_fallback_hop - ): - original_exception.message += format_no_fallback_group_message(lookup_groups, fallbacks) - no_fallback_group_explained = True - raise original_exception - - input_kwargs.update( - { - "fallback_model_group": fallback_model_group, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - - return response - except Exception as new_exception: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - fallback_failure_exception_str = truncate_fallback_error_detail(redact_string(str(new_exception))) - cooldown_info: Final = await async_get_cooldown_deployments_with_debug_info( - ctx=self._core_ctx, - parent_otel_span=parent_otel_span, - ) - verbose_router_logger.error( - "litellm.router.py::async_function_with_fallbacks() - " - "Error occurred while trying to do fallbacks - %s\n" - "Debug Information:\nCooldown Deployments=%s", - fallback_failure_exception_str, - cooldown_info, - ) - - attempted_fallback_group: Final = input_kwargs.get("fallback_model_group") - if ( - hasattr(original_exception, "message") - and self._host.dynamic().expose_router_debug_in_errors - and not no_fallback_group_explained - and not nested_fallback_hop - ): - original_exception.message += format_fallback_outcome_message( - model_group, - attempted_fallback_group, - fallback_failure_exception_str, - skipped=fallback_walk.skipped, - ) - - raise original_exception - - @tracer.wrap() - async def async_function_with_fallbacks(self, *args, **kwargs): - """ - Try calling the function_with_retries - If it fails after num_retries, fall back to another model group - """ - model_group: Final[str | None] = kwargs.get("model") - compaction_surface: Final = surface_for_call( - getattr(kwargs.get("original_generic_function") or kwargs.get("original_function"), "__name__", "") - ) - if compaction_surface is not None: - kwargs["_context_compaction_state"] = initialize_compaction_state(kwargs, compaction_surface) - clear_pre_routing_selection(kwargs) # pyright: ignore[reportUnknownArgumentType] # **kwargs is untyped at this boundary - if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets): - _fallback_metadata_key: Final = get_router_metadata_variable_name( - function_name=getattr(kwargs.get("original_function"), "__name__", None) - ) - _sibling_metadata_key: Final = ( - "metadata" if _fallback_metadata_key == "litellm_metadata" else "litellm_metadata" - ) - if isinstance(_sibling_metadata := kwargs.get(_sibling_metadata_key), dict): - # In place, like every other router bucket write: downstream resolves the bucket by - # key presence, so rebinding kwargs to a copy detaches the proxy's request_data write-backs - _sibling_metadata.pop("attempted_fallbacks", None) - _sibling_metadata.pop("original_model_group", None) - if isinstance(_fallback_metadata := kwargs.get(_fallback_metadata_key), dict): - _fallback_metadata["attempted_fallbacks"] = 0 - if model_group is not None: - _fallback_metadata["original_model_group"] = model_group - include_fallback_errors: Final = kwargs.get("include_fallback_errors", False) is True - disable_fallbacks: Final[bool | None] = kwargs.pop("disable_fallbacks", False) - record_disable_fallbacks(kwargs, disable_fallbacks is True) - fallbacks: Final[list | None] = kwargs.get("fallbacks", self._core_fallbacks) - context_window_fallbacks: list | None = kwargs.get( - "context_window_fallbacks", self._core_context_window_fallbacks - ) - content_policy_fallbacks: list | None = kwargs.get( - "content_policy_fallbacks", self._core_content_policy_fallbacks - ) - - mock_timeout: Final = kwargs.pop("mock_timeout", None) - - with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): - try: - self._handle_mock_testing_fallbacks( - kwargs=kwargs, - model_group=model_group, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) - - if mock_timeout is not None: - response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) - else: - response = await self.async_function_with_retries(*args, **kwargs) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("Async Response: %s", response) - response = add_fallback_headers_to_response( - response=response, - attempted_fallbacks=0, - ) - return response - except Exception as e: - return await self.async_function_with_fallbacks_common_utils( - e, - disable_fallbacks, - fallbacks, - context_window_fallbacks, - content_policy_fallbacks, - model_group, - args, - kwargs, - include_fallback_errors=include_fallback_errors, - ) - - def _handle_mock_testing_fallbacks( - self, - kwargs: dict, - model_group: str | None = None, - fallbacks: list | None = None, - context_window_fallbacks: list | None = None, - content_policy_fallbacks: list | None = None, - ): - """ - Helper function to raise a litellm Error for mock testing purposes. - - Raises: - litellm.InternalServerError: when `mock_testing_fallbacks=True` passed in request params - litellm_exceptions.ContextWindowExceededError: when `mock_testing_context_fallbacks=True` passed in request params - litellm.ContentPolicyViolationError: when `mock_testing_content_policy_fallbacks=True` passed in request params - """ - mock_testing_params: Final = MockRouterTestingParams.from_kwargs(kwargs) - if ( - mock_testing_params.mock_testing_fallbacks is not None - and mock_testing_params.mock_testing_fallbacks is True - ): - raise litellm_exceptions.InternalServerError( - model=model_group, - llm_provider="", - message=f"This is a mock exception for model={model_group}, to trigger a fallback. Fallbacks={fallbacks}", - ) - elif ( - mock_testing_params.mock_testing_context_fallbacks is not None - and mock_testing_params.mock_testing_context_fallbacks is True - ): - raise litellm_exceptions.ContextWindowExceededError( - model=model_group, - llm_provider="", - message=f"This is a mock exception for model={model_group}, to trigger a fallback. \ - Context_Window_Fallbacks={context_window_fallbacks}", - ) - elif ( - mock_testing_params.mock_testing_content_policy_fallbacks is not None - and mock_testing_params.mock_testing_content_policy_fallbacks is True - ): - raise litellm_exceptions.ContentPolicyViolationError( - model=model_group, - llm_provider="", - message=f"This is a mock exception for model={model_group}, to trigger a fallback. \ - Context_Policy_Fallbacks={content_policy_fallbacks}", - ) - - @staticmethod - def _deployment_ids_to_skip_on_retry(exception: Exception, already_skipped: object) -> tuple[str, ...]: - failed_deployment_id: Final[str | None] = getattr(exception, "retry_skip_deployment_id", None) or getattr( - exception, "failed_deployment_id", None - ) - status_code: Final = getattr(exception, "status_code", None) - if not failed_deployment_id or not isinstance(status_code, int): - return () - if python_host.invoke("should_retry")(status_code): - return () - already_skipped_ids: Final = as_retry_skipped_deployment_ids(already_skipped) - skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id)))) - verbose_router_logger.debug( - "Retry skips deployments that already answered %s to this request: %s", status_code, skipped - ) - return skipped - - @tracer.wrap() - async def async_function_with_retries(self, *args, **kwargs): - verbose_router_logger.debug("Inside async function with retries.") - original_function: Final = kwargs.pop("original_function") - fallbacks: Final = kwargs.pop("fallbacks", self._core_fallbacks) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - context_window_fallbacks: Final = kwargs.pop("context_window_fallbacks", self._core_context_window_fallbacks) - content_policy_fallbacks: Final = kwargs.pop("content_policy_fallbacks", self._core_content_policy_fallbacks) - # Support per-request model_group_retry_policy override (from key/team settings) - model_group_retry_policy: Final = kwargs.pop("model_group_retry_policy", self.model_group_retry_policy) - model_group: Final[str | None] = kwargs.get("model") - request_num_retries: Final[int | None] = kwargs.pop("num_retries", None) - num_retries = request_num_retries - if num_retries is None: - # Fall back to the router setting (then 0) so the comparisons below never - # hit `None > int`, which would mask the real upstream error with a TypeError. - num_retries = self.num_retries if self.num_retries is not None else 0 - - ## ADD MODEL GROUP SIZE TO METADATA - used for model_group_rate_limit_error tracking - _metadata: Final[dict] = kwargs.get("litellm_metadata", kwargs.get("metadata")) or {} - if "model_group" in _metadata and isinstance(_metadata["model_group"], str): - model_list: Final = self._core_get_model_list(model_name=_metadata["model_group"]) - if model_list is not None: - _metadata.update({"model_group_size": len(model_list)}) - - verbose_router_logger.debug( - "async function w/ retries: original_function - %s, num_retries - %s", original_function, num_retries - ) - ## ADD RETRY TRACKING TO METADATA - used for spend logs retry tracking - _metadata["attempted_retries"] = 0 - _metadata["max_retries"] = num_retries # Updated after overrides in exception handler - try: - self._handle_mock_testing_rate_limit_error(model_group=model_group, kwargs=kwargs) - # if the function call is successful, no exception will be raised and we'll break out of the loop - response = await self.make_call(original_function, *args, **kwargs) - response = add_retry_headers_to_response(response=response, attempted_retries=0, max_retries=None) - return response - except Exception as e: - if is_guardrail_intervention(e) or is_agentic_followup_failure(e): - raise - current_attempt = None - original_exception = e - deployment_num_retries: Final = getattr(e, "num_retries", None) - - if ( - request_num_retries is None - and deployment_num_retries is not None - and isinstance(deployment_num_retries, int) - ): - num_retries = deployment_num_retries - """ - Retry Logic - """ - ( - _healthy_deployments, - _all_deployments, - ) = await self._async_get_healthy_deployments( - model=kwargs.get("model") or "", - parent_otel_span=parent_otel_span, - request_kwargs=share_holder_request_kwargs( - cast(Mapping[str, object], kwargs), # cast-ok: kwargs keys are param names - self._ptu_share_teams(), - ), - ) - - # Check retry policy FIRST, before should_retry_this_error - # This allows retry policies to override the healthy deployments check - _retry_policy_applies = False - if request_num_retries != 0 and (self.retry_policy is not None or model_group_retry_policy is not None): - # get num_retries from retry policy - # Use the model_group captured at the start of the function, or get it from metadata - # kwargs.get("model") at this point is the deployment model, not the model_group - _model_group_for_retry_policy = model_group or _metadata.get("model_group") or kwargs.get("model") - # Use per-request model_group_retry_policy if provided, otherwise use self - _retry_policy_retries: Final = _get_num_retries_from_retry_policy( - exception=original_exception, - model_group=_model_group_for_retry_policy, - model_group_retry_policy=model_group_retry_policy, - retry_policy=self.retry_policy, - ) - if _retry_policy_retries is not None: - num_retries = _retry_policy_retries - _retry_policy_applies = True - - # raises an exception if this error should not be retries - # Skip this check if retry policy applies (retry policy takes precedence) - if not _retry_policy_applies: - self.should_retry_this_error( - error=e, - healthy_deployments=_healthy_deployments, - all_deployments=_all_deployments, - context_window_fallbacks=context_window_fallbacks, - regular_fallbacks=fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) - # Update max_retries after overrides (deployment_num_retries / retry_policy) - _metadata["max_retries"] = num_retries - - ## LOGGING - if num_retries > 0: - kwargs = self._core_log_retry(kwargs=kwargs, e=original_exception) - first_skipped_ids: Final = self._deployment_ids_to_skip_on_retry( - exception=original_exception, - already_skipped=kwargs.get("_retry_skipped_deployment_ids"), - ) - if first_skipped_ids: - kwargs["_retry_skipped_deployment_ids"] = first_skipped_ids # rebind-ok: the next attempt reads it - if isinstance( - original_exception, litellm_exceptions.UnsupportedEndpointError - ) and _every_deployment_refused( - _DEPLOYMENTS_ADAPTER.validate_python(_healthy_deployments), first_skipped_ids - ): - raise - else: - raise - - verbose_router_logger.debug("Retrying request with num_retries: %s", num_retries) - # decides how long to sleep before retry - retry_after: Final = self._time_to_sleep_before_retry( - e=original_exception, - remaining_retries=num_retries, - num_retries=num_retries, - healthy_deployments=_healthy_deployments, - all_deployments=_all_deployments, - ) - - await self._host.sleep(retry_after) - - for current_attempt in range(num_retries): - try: - # Update retry tracking metadata before each retry attempt - _metadata["attempted_retries"] = current_attempt + 1 - _metadata["max_retries"] = num_retries - carry_superseded_attempt_usage(kwargs) - # if the function call is successful, no exception will be raised and we'll break out of the loop - response = await self.make_call(original_function, *args, **kwargs) - if coroutine_checker.is_async_callable(response): # async errors are often returned as coroutines - response = await response - - response = add_retry_headers_to_response( - response=response, - attempted_retries=current_attempt + 1, - max_retries=num_retries, - ) - return response - - except Exception as e: - if is_agentic_followup_failure(e): - raise - # Always track the latest error so we raise the most - # recent exception instead of the first one. - original_exception = e - - ## LOGGING - kwargs = self._core_log_retry(kwargs=kwargs, e=e) - remaining_retries = num_retries - current_attempt - 1 - _model: str | None = kwargs.get("model") - if _model is not None: - ( - _healthy_deployments, - _, - ) = await self._async_get_healthy_deployments( - model=_model, - parent_otel_span=parent_otel_span, - request_kwargs=share_holder_request_kwargs( - cast(Mapping[str, object], kwargs), # cast-ok: kwargs keys are param names - self._ptu_share_teams(), - ), - ) - else: - _healthy_deployments = [] - - # Check if this error is non-retryable (e.g., 400 context - # window exceeded). If so, raise immediately instead of - # continuing the retry loop. Respect retry policy - # precedence - only check when no retry policy applies. - if not _retry_policy_applies: - try: - self.should_retry_this_error( - error=e, - healthy_deployments=_healthy_deployments, - all_deployments=_all_deployments, - context_window_fallbacks=context_window_fallbacks, - regular_fallbacks=fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) - except Exception: - raise e - - skipped_ids = self._deployment_ids_to_skip_on_retry( - exception=e, - already_skipped=kwargs.get("_retry_skipped_deployment_ids"), - ) - if skipped_ids: - kwargs["_retry_skipped_deployment_ids"] = skipped_ids # rebind-ok: the next attempt reads it - if isinstance(e, litellm_exceptions.UnsupportedEndpointError) and _every_deployment_refused( - _DEPLOYMENTS_ADAPTER.validate_python(_healthy_deployments), skipped_ids - ): - raise - _timeout = self._time_to_sleep_before_retry( - e=e, - remaining_retries=remaining_retries, - num_retries=num_retries, - healthy_deployments=_healthy_deployments, - all_deployments=_all_deployments, - ) - await self._host.sleep(_timeout) - - if type(original_exception) in self._host.dynamic().exception_types: - setattr(original_exception, "max_retries", num_retries) - # current_attempt is 0-indexed (0 to num_retries-1), so after loop completion - # it represents the last attempt index. The actual number of retries attempted - # is current_attempt + 1, which equals num_retries when all retries are exhausted. - # We've already verified num_retries > 0 before entering the loop, so current_attempt - # will always be set (never None) when we reach this point. - actual_retries_attempted: Final = current_attempt + 1 if current_attempt is not None else num_retries - setattr(original_exception, "num_retries", actual_retries_attempted) - - raise original_exception - - async def make_call(self, original_function: Any, *args: object, **kwargs: object): - """ - Handler for making a call to the .completion()/.embeddings()/etc. functions. - """ - model_group: Final = kwargs.get("model") - generic_function: Final = _OPAQUE_VALUE_ADAPTER.validate_python(kwargs.get("original_generic_function")) - endpoint_function: Final = ( - generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function) - ) - try: - with endpoint_selection( - getattr(endpoint_function, "__name__", None), - getattr(endpoint_function, "__module__", None), - model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group), - ): - response = original_function(*args, **kwargs) - if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): - response = await response - except Exception: - with contextlib.suppress(Exception): - await async_io_token_refund_failure( - self._core_cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs) - ) - raise - await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) - ## PROCESS RESPONSE HEADERS - response = await self.set_response_headers( - response=response, model_group=model_group if isinstance(model_group, str) else None, request_kwargs=kwargs - ) - - return response - - def _handle_mock_testing_rate_limit_error(self, kwargs: dict, model_group: str | None = None): - """ - Helper function to raise a mock litellm.RateLimitError error for testing purposes. - - Raises: - litellm.RateLimitError error when `mock_testing_rate_limit_error=True` passed in request params - """ - mock_testing_rate_limit_error: Final[bool | None] = kwargs.pop("mock_testing_rate_limit_error", None) - - available_models: Final = self._core_get_model_list(model_name=model_group) - num_retries: int | None = None - - if available_models is not None and len(available_models) == 1: - num_retries = available_models[0]["litellm_params"].get("num_retries") - - if mock_testing_rate_limit_error is not None and mock_testing_rate_limit_error is True: - verbose_router_logger.info( - "litellm.router.py::_mock_rate_limit_error() - Raising mock RateLimitError for model=%s", model_group - ) - raise litellm_exceptions.RateLimitError( - model=model_group, - llm_provider="", - message=f"This is a mock exception for model={model_group}, to trigger a rate limit error.", - num_retries=num_retries, - ) - - def should_retry_this_error( - self, - error: Exception, - healthy_deployments: list | None = None, - all_deployments: list | None = None, - context_window_fallbacks: list | None = None, - content_policy_fallbacks: list | None = None, - regular_fallbacks: list | None = None, - ): - """ - 1. raise an exception for ContextWindowExceededError if context_window_fallbacks is not None - 2. raise an exception for ContentPolicyViolationError if content_policy_fallbacks is not None - - 2. raise an exception for RateLimitError if - - there are no fallbacks - - there are no healthy deployments in the same model group - """ - _num_healthy_deployments = 0 - if healthy_deployments is not None and isinstance(healthy_deployments, list): - _num_healthy_deployments = len(healthy_deployments) - - _num_all_deployments = 0 - if all_deployments is not None and isinstance(all_deployments, list): - _num_all_deployments = len(all_deployments) - - ### CHECK IF RATE LIMIT / CONTEXT WINDOW ERROR / CONTENT POLICY VIOLATION ERROR w/ fallbacks available / Bad Request Error - if isinstance(error, litellm_exceptions.ContextWindowExceededError) and context_window_fallbacks is not None: - raise error - - if isinstance(error, litellm_exceptions.ContentPolicyViolationError) and content_policy_fallbacks is not None: - raise error - - status_code: Final = getattr(error, "status_code", None) - if status_code is not None and not self._host.invoke("should_retry")(status_code): - # 401/403 are special cases - allow retry if multiple deployments exist (handled below) - if status_code not in (401, 403) and not isinstance(error, litellm_exceptions.UnsupportedEndpointError): - raise error - - if isinstance(error, litellm_exceptions.NotFoundError): - raise error - # Error we should only retry if there are other deployments - if isinstance(error, openai.RateLimitError): - if ( - _num_healthy_deployments <= 0 # if no healthy deployments - and regular_fallbacks is not None # and fallbacks available - and len(regular_fallbacks) > 0 - ): - raise error # then raise the error - - if isinstance(error, (openai.AuthenticationError, openai.PermissionDeniedError)): - """ - - if other deployments available -> retry - - else -> raise error - """ - if _num_all_deployments <= 1: # if there is only 1 deployment for this model group then don't retry - raise error # then raise error - - # Do not retry if there are no healthy deployments - # just raise the error - if _num_healthy_deployments <= 0: # if no healthy deployments - raise error - - return True - - def function_with_fallbacks(self, *args, **kwargs): - """ - Sync wrapper for async_function_with_fallbacks - - Wrapped to reduce code duplication and prevent bugs. - """ - return run_async_function(self.async_function_with_fallbacks, *args, **kwargs) - - def _get_fallback_model_group_for_lookup_groups( - self, - fallbacks: list[dict[str, list[str]]], # mutable-ok: mirrors the shared resolver's contract - lookup_groups: tuple[str, ...], - ) -> list[str] | None: # mutable-ok: mirrors the shared resolver's contract - fallback_model_group, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, lookup_groups=lookup_groups - ) - return fallback_model_group - - def _get_first_default_fallback(self) -> str | None: - """ - Returns the first model from the default_fallbacks list, if it exists. - """ - if self._core_fallbacks is None: - return None - for fallback in self._core_fallbacks: - if isinstance(fallback, dict) and "*" in fallback: - default_list = fallback["*"] - if isinstance(default_list, list) and len(default_list) > 0: - return default_list[0] - return None - - def _time_to_sleep_before_retry( - self, - e: Exception, - remaining_retries: int, - num_retries: int, - healthy_deployments: list | None = None, - all_deployments: list | None = None, - ) -> int | float: - """ - Calculate back-off, then retry - - It should instantly retry only when: - 1. there are healthy deployments in the same model group - 2. there are fallbacks for the completion call - """ - - ## base case - single deployment - if all_deployments is not None and len(all_deployments) == 1: - pass - elif healthy_deployments is not None and isinstance(healthy_deployments, list) and len(healthy_deployments) > 0: - return 0 - - response_headers: httpx.Headers | None = None - if hasattr(e, "response") and hasattr(e.response, "headers"): - response_headers = e.response.headers - if hasattr(e, "litellm_response_headers"): - response_headers = e.litellm_response_headers - - if response_headers is not None: - timeout = self._host.invoke("calculate_retry_after")( - remaining_retries=remaining_retries, - max_retries=num_retries, - response_headers=response_headers, - min_timeout=self.retry_after, - ) - - else: - timeout = self._host.invoke("calculate_retry_after")( - remaining_retries=remaining_retries, - max_retries=num_retries, - min_timeout=self.retry_after, - ) - - return timeout - - ### HELPER FUNCTIONS - - async def deployment_callback_on_success( - self, - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, - end_time, # start/end time - ): - """ - Track remaining tpm/rpm quota for model in model_list - """ - try: - if kwargs.get("call_type") == "_arealtime": - return - if is_batch_retrieve_call_type(kwargs.get("call_type")): - return - standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) - if standard_logging_object is None: - raise ValueError("standard_logging_object is None") - litellm_params: Final = kwargs["litellm_params"] - metadata: Final = litellm_params.get("metadata") - if metadata is None: - return - model_group: Final = metadata.get("model_group", None) - model_info: Final = litellm_params.get("model_info", {}) or {} - deployment_id: Final = model_info.get("id", None) - if ( - model_group is None - or deployment_id is None - or self._core_get_deployment(model_id=str(deployment_id)) is None - ): - return - - if self.session_affinity_repin_on_failover: - await self._repin_session_after_success( - _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params), - _OPAQUE_VALUE_ADAPTER.validate_python(completion_response), - ) - - # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits - increment_deployment_successes_for_current_minute( - ctx=self._core_ctx, - deployment_id=str(deployment_id), - ) - - total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0) - counted_tokens: Final = get_counted_usage_tokens(litellm_params) - deployment_name: Final = metadata.get("deployment", None) - return await self._increment_deployment_usage( - deployment_id=str(deployment_id), - deployment_name=deployment_name if isinstance(deployment_name, str) else None, - model_group=model_group, - total_tokens=total_tokens if counted_tokens is None else max(0, total_tokens - counted_tokens), - rpm_increment=1 if counted_tokens is None else 0, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) - - except Exception as e: - verbose_router_logger.debug( - "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e - ) - - async def _repin_session_after_success(self, litellm_params: Mapping[str, object], response: object) -> None: - if not self.session_affinity_repin_on_failover: - return - try: - from litellm.router_utils.pre_call_checks.deployment_affinity_check import completed_affinity_response - - if not completed_affinity_response(response): - return - callback: Final = next( - (logger for logger in self._optional_loggers() if isinstance(logger, DeploymentAffinityCheck)), - None, - ) - if callback is None: - return - model_group: Final = callback.get_repin_model_group(litellm_params) - if model_group is None or self._is_priority_routing_group(model_group): - return - registered_ids: Final = self._core_get_candidate_model_ids_for_route( - model=model_group, team_id=get_request_team_id(litellm_params) - ) - cooldown_ids: Final = frozenset( - model_id - for model_id, _ in await self._core_cooldown_cache.async_get_active_cooldowns( - model_ids=list(registered_ids), parent_otel_span=None - ) - ) - await callback.async_repin_session_after_success(litellm_params, model_group, registered_ids - cooldown_ids) - except Exception as error: # noqa: BLE001 # Optional pin updates must not interrupt success accounting - verbose_router_logger.debug("Session re-pinning failed: %s", error) - - async def increment_deployment_usage_for_response( - self, - response: object, - request_kwargs: dict[str, object], - ) -> None: - if response is None: - return - try: - deployment_metadata: Final = find_deployment_metadata(request_kwargs) - model_group: Final = request_kwargs.get("model") - if deployment_metadata is None or not isinstance(model_group, str): - return - model_info: Final = deployment_metadata["model_info"] - deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None - if deployment_id is None: - return - total_tokens: Final = response_total_token_count(response) - deployment_name: Final = deployment_metadata.get("deployment") - deployment_metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = total_tokens - try: - await self._increment_deployment_usage( - deployment_id=str(deployment_id), - deployment_name=deployment_name if isinstance(deployment_name, str) else None, - model_group=model_group, - total_tokens=total_tokens, - rpm_increment=1, - parent_otel_span=get_parent_otel_span_from_kwargs(request_kwargs), - ) - except Exception: - deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None) - raise - except Exception as e: - verbose_router_logger.debug( - "litellm.router.Router::increment_deployment_usage_for_response(): Exception occured - %s", e - ) - - async def _increment_deployment_usage( - self, - *, - deployment_id: str, - deployment_name: str | None, - model_group: str, - total_tokens: float, - rpm_increment: int, - parent_otel_span: Span | None, - ) -> str | None: - from litellm.types.caching import RedisPipelineIncrementOperation - - if total_tokens <= 0 and rpm_increment <= 0: - return None - if not self._cached_deployment_tracks_usage(deployment_id, model_group): - return None - - current_minute: Final = get_utc_datetime().strftime("%H-%M") # use the same timezone regardless of system clock - tpm_key: Final = RouterCacheEnum.TPM.value.format( - id=deployment_id, current_minute=current_minute, model=deployment_name - ) - rpm_key: Final = RouterCacheEnum.RPM.value.format( - id=deployment_id, current_minute=current_minute, model=deployment_name - ) - pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [ - RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value) - for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment)) - ] - post_increment_values: Final = await self._core_cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - parent_otel_span=parent_otel_span, - ) - if post_increment_values is not None and self._core_cache.redis_cache is not None: - for operation, value in zip(pipeline_operations, post_increment_values): - await self._core_cache.async_set_cache( - operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value - ) - return tpm_key - - def _deployment_tracks_usage(self, deployment_id: str, model_group: str) -> bool: - deployment_info: Final = self._core_get_deployment(model_id=deployment_id) - if deployment_info is None: - return False - deployment_model_info: Final = self._core_get_router_model_info( - deployment=deployment_info, - received_model_name=model_group, - ) - configured_limits: Final = ( - deployment_info.get("tpm", None), - deployment_info.get("rpm", None), - deployment_info.litellm_params.tpm, - deployment_info.litellm_params.rpm, - deployment_model_info.get("tpm", None), - deployment_model_info.get("rpm", None), - ) - ## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are - ## set. IO deployments still record TPM/RPM usage here so TPM-aware - ## routing strategies see their real load in mixed model groups; their - ## itpm/otpm enforcement runs separately in ModelRateLimitingCheck. - return any(limit is not None for limit in configured_limits) or deployment_has_io_token_limits( - deployment_info.model_dump() - ) - - def sync_deployment_callback_on_success( - self, - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, - end_time, # start/end time - ) -> str | None: - """ - Tracks the number of successes for a deployment in the current minute (using in-memory cache) - - Returns: - - key: str - The key used to increment the cache - - None: if no key is found - """ - if is_batch_retrieve_call_type(kwargs.get("call_type")): - return None - id = None - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - model_info: Final = kwargs["litellm_params"].get("model_info", {}) or {} - id = model_info.get("id", None) - if model_group is None or id is None: - return None - elif isinstance(id, int): - id = str(id) - - if id is not None: - key: Final = increment_deployment_successes_for_current_minute( - ctx=self._core_ctx, - deployment_id=id, - ) - return key - - return None - - def deployment_callback_on_failure( - self, - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, - end_time, # start/end time - ) -> bool: - """ - 2 jobs: - - Tracks the number of failures for a deployment in the current minute (using in-memory cache) - - Puts the deployment in cooldown if it exceeds the allowed fails / minute - - Returns: - - True if the deployment should be put in cooldown - - False if the deployment should not be put in cooldown - """ - verbose_router_logger.debug("Router: Entering 'deployment_callback_on_failure'") - try: - exception: Final = kwargs.get("exception", None) - - if is_failure_outside_selected_deployment(exception): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Failure came from advisor orchestration or deployment selection, not the selected deployment." - ) - return False - - # Cache litellm_params to avoid repeated dict lookups - litellm_params: Final = kwargs.get("litellm_params", {}) - _model_info: Final = litellm_params.get("model_info", {}) - - if is_background_response_poll_not_found(exception, litellm_params): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Provider 404 came from the background response poll, not the deployment's health." - ) - return False - - if is_video_lookup_not_found(exception, kwargs): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Provider 404 says the video is not on this deployment's account, not that the deployment is down." - ) - return False - - exception_status: Final = getattr(exception, "status_code", "") - failure_call_details: Final = _CALL_KWARGS_ADAPTER.validate_python(kwargs) - - if is_caller_timeout_408(failure_call_details, exception_status): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "A timeout the caller set caused this 408, not the deployment's health." - ) - return False - - caller_failure_model_info: Final = ( - _MODEL_INFO_ADAPTER.validate_python(_model_info) if isinstance(_model_info, dict) else None - ) - caller_failure_model_id: Final = ( - caller_failure_model_info.get("id") if caller_failure_model_info is not None else None - ) - caller_failure_deployment: Final = ( - self._core_get_deployment(model_id=caller_failure_model_id) - if isinstance(caller_failure_model_id, str) - else None - ) - if is_caller_scoped_auth_failure(caller_failure_deployment, exception_status): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Caller-scoped OAuth authentication failed, not the deployment." - ) - return False - if is_caller_entra_credential_failure( - caller_failure_deployment, - self._active_default_litellm_params(), - failure_call_details, - exception_status, - exception, - ): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Entra ID rejected a credential the caller sent, not the deployment's own." - ) - return False - - exception_headers: Final = exception_mapping_utils._get_response_headers(original_exception=exception) - - # Determine cooldown time with priority: deployment config > response header > router default - deployment_cooldown: Final = _first_present( - _model_info if isinstance(_model_info, dict) else None, litellm_params, key="cooldown_time" - ) - - header_cooldown = None - if exception_headers is not None: - header_cooldown = litellm_utils._get_retry_after_from_exception_header( - response_headers=exception_headers - ) - ############################################## - # Logic to determine cooldown time - # 1. Check if a cooldown time is set in the deployment config - # 2. Check if a cooldown time is set in the response header - # 3. If no cooldown time is set, use the router default cooldown time - ############################################## - if deployment_cooldown is not None and deployment_cooldown >= 0: - _time_to_cooldown = deployment_cooldown - elif header_cooldown is not None and header_cooldown >= 0: - _time_to_cooldown = header_cooldown - else: - _time_to_cooldown = self._core_cooldown_time - - if isinstance(_model_info, dict): - deployment_id: Final[str | None] = _model_info.get("id") - if deployment_id is None: - return False - increment_deployment_failures_for_current_minute( - ctx=self._core_ctx, - deployment_id=deployment_id, - ) - result: Final = set_cooldown_deployments( - ctx=self._core_ctx, - exception_status=exception_status, - original_exception=exception, - deployment=deployment_id, - time_to_cooldown=_time_to_cooldown, - requested_model_group=( - get_litellm_metadata_from_kwargs( - cast("dict[str, object]", kwargs) # cast-ok: untyped kwargs dict - ) - or {} - ).get("model_group"), - request_kwargs=cast("dict[str, object]", kwargs), # cast-ok: kwargs is the untyped request dict - ) # setting deployment_id in cooldown deployments - - return result - else: - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. No model_info found." - ) - return False - - except Exception as e: - raise e - - async def async_deployment_callback_on_failure( - self, kwargs, completion_response: object | None, start_time, end_time - ): - """ - Update RPM usage for a deployment - """ - if is_batch_retrieve_call_type(kwargs.get("call_type")): - return - deployment_name: Final = kwargs["litellm_params"]["metadata"].get( - "deployment", None - ) # handles wildcard routes - by giving the original name sent to `litellm.completion` - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - model_info: Final = kwargs["litellm_params"].get("model_info", {}) or {} - id = model_info.get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - - ## RPM - rpm_key: Final = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name) - with service_target(ROUTER_USAGE_TARGET): - await self._core_cache.async_increment_cache( - key=rpm_key, - value=1, - parent_otel_span=parent_otel_span, - ttl=RoutingArgs.ttl.value, - ) - - def _get_metadata_variable_name_from_kwargs(self, kwargs: dict) -> Literal["metadata", "litellm_metadata"]: - """ - Helper to return what the "metadata" field should be called in the request data - - - New endpoints return `litellm_metadata` - - Old endpoints return `metadata` - - Context: - - LiteLLM used `metadata` as an internal field for storing metadata - - OpenAI then started using this field for their metadata - - LiteLLM is now moving to using `litellm_metadata` for our metadata - """ - return get_metadata_variable_name_from_kwargs(kwargs) - - def _core_log_retry(self, kwargs: dict, e: Exception) -> dict: - """ - When a retry or fallback happens, record which model group, deployment and attempt just failed and why, - and count it toward the request-wide num_retries_per_request cap - """ - from litellm.litellm_core_utils.get_litellm_params import get_base_model_from_litellm_call_metadata - from litellm.types.router import RetryAttemptRecord - - _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" - request_metadata: Final = cast("Mapping[str, object]", kwargs[_metadata_var]) # cast-ok: untyped kwargs - model_group: Final = kwargs.get("model") - model_info: Final = request_metadata.get("model_info") - deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None - base_model: Final = get_base_model_from_litellm_call_metadata( - metadata=request_metadata - ) or self._deployment_params_base_model(deployment_id if isinstance(deployment_id, str) else None) - deployment_model: Final = request_metadata.get("deployment") - attempted_retries: Final = request_metadata.get("attempted_retries") - attempt_record: Final[RetryAttemptRecord] = { - "model_group": model_group if isinstance(model_group, str) else None, - "deployment_id": deployment_id if isinstance(deployment_id, str) else None, - "model": deployment_model if isinstance(deployment_model, str) else None, - "base_model": base_model if isinstance(base_model, str) else None, - "exception_type": type(e).__name__, - "exception_string": str(e), - "attempted_retries": attempted_retries if type(attempted_retries) is int else None, - } - earlier_breadcrumbs: Final = request_metadata.get("previous_models") - kept_breadcrumbs: Final[tuple[object, ...]] = ( - tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :] - if isinstance(earlier_breadcrumbs, (list, tuple)) - else () - ) - breadcrumbs: Final = (*kept_breadcrumbs, attempt_record) - earlier: Final = request_metadata.get("request_retry_count") - request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1 - decision: Final = request_metadata.get("routing_decision") - router_name: Final = decision.get("router_model_name") if isinstance(decision, Mapping) else None - if isinstance(router_name, str): - phase_event( - "litellm.routing.retry", - { - "litellm.routing.router_model_name": router_name, - "litellm.retry.count": request_retry_count, - "error.type": type(e).__name__, - **({"litellm.deployment.model_group": model_group} if isinstance(model_group, str) else {}), - **({"litellm.deployment.id": deployment_id} if isinstance(deployment_id, str) else {}), - }, - ) - kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict - kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap - return kwargs - - def _deployment_params_base_model(self, deployment_id: str | None) -> str | None: - """The ``base_model`` a deployment declares under ``litellm_params`` (the Azure named-deployment - shape), the same key the success path's ``get_base_model_from_metadata`` reads.""" - deployment: Final = self._core_get_deployment(model_id=deployment_id) if deployment_id is not None else None - if deployment is None: - return None - extra_params: Final = cast( # cast-ok: pydantic extras are untyped, base_model is isinstance-checked below - "Mapping[str, object]", deployment.litellm_params.model_extra or {} - ) - base_model: Final = extra_params.get("base_model") - return base_model if isinstance(base_model, str) else None - - def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int: - """ - Update deployment rpm for that minute - - Returns: - - int: request count - """ - rpm_key: Final = deployment_id - - request_count = self._core_cache.get_cache(key=rpm_key, parent_otel_span=parent_otel_span, local_only=True) - if request_count is None: - request_count = 1 - self._core_cache.set_cache(key=rpm_key, value=request_count, local_only=True, ttl=60) # only store for 60s - else: - request_count += 1 - self._core_cache.set_cache(key=rpm_key, value=request_count, local_only=True) # don't change existing ttl - - return request_count - - def _has_default_fallbacks(self) -> bool: - if self._core_fallbacks is None: - return False - for fallback in self._core_fallbacks: - if isinstance(fallback, dict): - if "*" in fallback: - return True - return False - - def _has_content_policy_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: - """ - Whether a content-policy fallback would resolve for this request, keyed the same way - async_function_with_fallbacks_common_utils resolves it: the tier a pre-routing hook - selected wins over the requested group. Raising without this returning True would turn - a deliverable response into an error the fallback chain cannot recover from. - """ - content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self._core_content_policy_fallbacks) - if content_policy_fallbacks is not None: - return has_unattempted_fallback_target( - self._get_fallback_model_group_for_lookup_groups( - fallbacks=content_policy_fallbacks, - lookup_groups=fallback_lookup_groups(kwargs, model_group), - ), - kwargs, - ) - if self._has_default_fallbacks(): - return True - verbose_router_logger.debug( - "No content-policy fallback available. Returning original response. model=%s, content_policy_fallbacks=%s", - model_group, - content_policy_fallbacks, - ) - return False - - def _refusal_fallback_available(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: - """ - Whether a safeguard refusal can actually be recovered by the dispatcher. A configured - content-policy list is authoritative; with none configured at all, the dispatcher falls - through to the generic fallbacks lookup, so the gate mirrors that reachability and arms - on a resolving generic chain (tier first, then the requested group, then "*"). - """ - if fallbacks_disabled_for_request(kwargs): - return False - content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self._core_content_policy_fallbacks) - if content_policy_fallbacks is not None: - return self._has_content_policy_fallback(model_group, kwargs) - if self._has_default_fallbacks(): - return True - fallbacks: Final = kwargs.get("fallbacks", self._core_fallbacks) - if fallbacks is None: - return False - resolved, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, - lookup_groups=fallback_lookup_groups(kwargs, model_group), - ) - return has_unattempted_fallback_target(resolved, kwargs) - - def _stream_order_levels(self, model_group: str, kwargs: Mapping[str, Any]) -> tuple[int, ...]: - """ - The distinct deployment order levels the fallback dispatcher would see for this request, - computed the same way: the tier a pre-routing hook selected wins over the requested group. - """ - request_team_id: Final[str | None] = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id") - order_model_group: Final = get_pre_routing_selection(kwargs) or model_group - all_deployments: Final = self._core_get_model_list(model_name=order_model_group, team_id=request_team_id) or () - return tuple( - sorted( - { - litellm_utils.get_deployment_order(d) - for d in all_deployments - if litellm_utils.get_deployment_order(d) is not None - } - ) - ) - - def _stream_can_fall_back(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: - """ - Whether async_function_with_fallbacks_common_utils could still route a - MidStreamFallbackError somewhere for this request (order levels, weighted - failover, content-policy or generic fallbacks), which is the only case where - holding lifecycle frames back from the client buys a clean retry. Errs toward - True whenever a dispatcher path might reach a fallback. - """ - if fallbacks_disabled_for_request(kwargs): - return False - if self.enable_weighted_failover: - return True - order_levels: Final = self._stream_order_levels(model_group, kwargs) - if len(order_levels) > 1: - current_target: Final = kwargs.get("_target_order") - skip_up_to: Final = current_target if current_target is not None else order_levels[0] - if any(o > skip_up_to for o in order_levels): - return True - content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self._core_content_policy_fallbacks) - if content_policy_fallbacks is not None and self._has_content_policy_fallback(model_group, kwargs): - return True - fallbacks: Final = kwargs.get("fallbacks", self._core_fallbacks) - if not fallbacks: - return False - if check_non_standard_fallback_format(fallbacks=fallbacks): - return True - resolved, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, - lookup_groups=fallback_lookup_groups(kwargs, model_group), - ) - return has_unattempted_fallback_target(resolved, kwargs) - - def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool: - """ - Determines if a content policy error should be raised. - - Only raised if a fallback is available. - - Else, original response is returned. - """ - if response.choices and len(response.choices) > 0: - if response.choices[0].finish_reason != "content_filter": - return False - - return self._has_content_policy_fallback(model, kwargs) - - def _should_raise_anthropic_refusal_error( - self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, object] - ) -> bool: - """ - The /v1/messages twin of _should_raise_content_policy_error: an Anthropic safeguard - refusal (stop_reason "refusal" carrying stop_details) re-enters the fallback chain only - when a content-policy fallback is configured; a plain refusal without stop_details, or - any response with nothing configured, is returned to the client unchanged. - """ - from litellm.llms.anthropic.messages.native.utils import ( - get_safeguard_refusal_stop_details, - ) - - if getattr(original_generic_function, "__name__", "") != "anthropic_messages": - return False - if get_safeguard_refusal_stop_details(response) is None: - return False - return self._refusal_fallback_available(model, kwargs) - - def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None): - _all_deployments: list = [] - try: - _, _all_deployments = self._common_checks_available_deployment( - model=model, - ) - if isinstance(_all_deployments, dict): - return [] - except Exception: - pass - - unhealthy_deployments: Final = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - unhealthy_set: Final = set(unhealthy_deployments) - healthy_deployments: list = [d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_set] - healthy_deployments = self._filter_blocked_deployments(healthy_deployments) - - return healthy_deployments, _all_deployments - - async def _async_get_healthy_deployments( - self, model: str, parent_otel_span: Span | None, request_kwargs: Mapping[str, object] | None = None - ) -> tuple[list[dict], list[dict]]: - """ - Returns Tuple of: - - Tuple[List[Dict], List[Dict]]: - 1. healthy_deployments: list of healthy deployments the caller of ``request_kwargs`` may use - 2. all_deployments: list of all deployments that caller may use - """ - _all_deployments: list = [] - try: - _, _all_deployments = self._common_checks_available_deployment( - model=model, - request_kwargs=request_kwargs, - ) - if isinstance(_all_deployments, dict): - return [], _all_deployments - except Exception: - pass - - unhealthy_deployments: Final = await async_get_cooldown_deployments( - ctx=self._core_ctx, parent_otel_span=parent_otel_span - ) - # Convert to set for O(1) lookup instead of O(n) - unhealthy_deployments_set: Final = set(unhealthy_deployments) - healthy_deployments: list = [ - d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_deployments_set - ] - healthy_deployments = self._filter_blocked_deployments(healthy_deployments) - return healthy_deployments, _all_deployments - - def routing_strategy_pre_call_checks(self, deployment: dict): - """ - Mimics 'async_routing_strategy_pre_call_checks' - - Ensures consistent update rpm implementation for 'usage-based-routing-v2' - - Returns: - - None - - Raises: - - Rate Limit Exception - If the deployment is over it's tpm/rpm limits - """ - for _callback in self._routing_callbacks(): - if isinstance(_callback, CustomLogger): - _callback.pre_call_check(deployment) - - async def async_routing_strategy_pre_call_checks( - self, - deployment: dict, - parent_otel_span: Span | None, - logging_obj: LiteLLMLogging | None = None, - ): - """ - For usage-based-routing-v2, enables running rpm checks before the call is made, inside the semaphore. - - -> makes the calls concurrency-safe, when rpm limits are set for a deployment - - Returns: - - None - - Raises: - - Rate Limit Exception - If the deployment is over it's tpm/rpm limits - """ - for _callback in self._routing_callbacks(): - if isinstance(_callback, CustomLogger): - try: - await _callback.async_pre_call_check(deployment, parent_otel_span) - except litellm_exceptions.RateLimitError as e: - ## LOG FAILURE EVENT - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback.format_exc(), - prefer_async_handlers=True, - ) - ) - set_cooldown_deployments( - ctx=self._core_ctx, - exception_status=e.status_code, - original_exception=e, - deployment=deployment["model_info"]["id"], - time_to_cooldown=self._core_cooldown_time, - ) - raise e - except Exception as e: - ## LOG FAILURE EVENT - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback.format_exc(), - prefer_async_handlers=True, - ) - ) - raise e - - def _deployment_as_configured_now(self, deployment: Mapping[str, object]) -> Mapping[str, object]: - model_info: Final = deployment.get("model_info") - model_id: Final = ( - _MODEL_INFO_ADAPTER.validate_python(model_info).get("id") if isinstance(model_info, dict) else None - ) - current: Final = self._core_get_deployment(model_id) if isinstance(model_id, str) else None - return deployment if current is None else current.model_dump(exclude_none=True) - - async def async_admit_deployment_turn( - self, - model_group: str, - opened_on: Mapping[str, object], - request_kwargs: Mapping[str, object], - ) -> Mapping[str, object]: - deployment: Final = self._deployment_as_configured_now(opened_on) - rate_limit_check: Final = next( - (callback for callback in self._routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)), None - ) - io_token_admission: Final[Mapping[str, object]] = ( - MappingProxyType({}) - if rate_limit_check is None - else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) - ) - budget_limiter: Final = self._core_get_router_deployment_budget_limiter() - if budget_limiter is not None: - try: - await budget_limiter.async_filter_deployments( # pyright: ignore[reportUnknownMemberType] # legacy signature takes untyped deployment dicts - model=model_group, - healthy_deployments=[dict(deployment)], - messages=None, - request_kwargs=dict(request_kwargs), - ) - except ValueError as e: - raise litellm_exceptions.RateLimitError( - message=str(e), - llm_provider="", - model=model_group, - category=litellm_exceptions.RateLimitErrorCategory.LITELLM_RATE_LIMIT, - ) from e - await self._admit_turn_through_routing_strategy(model_group, deployment, request_kwargs) - await self.async_routing_strategy_pre_call_checks( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict - deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None - ) - return io_token_admission - - async def _admit_turn_through_routing_strategy( - self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] - ) -> None: - routing_kwargs: Final = {key: value for key, value in request_kwargs.items() if key != "litellm_logging_obj"} - strategy, selector = self._get_routing_context( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped request dict - model=model_group, request_kwargs=routing_kwargs - ) - if strategy == "simple-shuffle" or selector is None: - return - selected: Final = await self._select_deployment_async( # pyright: ignore[reportUnknownMemberType] # legacy signature takes untyped deployment dicts - strategy=strategy, - selector=selector, - model=model_group, - healthy_deployments=[dict(deployment)], - messages=None, - input=None, - request_kwargs=routing_kwargs, - reserve=False, - ) - if selected is None: - no_deployment: Final = await async_raise_no_deployment_exception( - ctx=self._core_ctx, model=model_group, parent_otel_span=None - ) - raise litellm_exceptions.RateLimitError( - message=str(no_deployment), - llm_provider="", - model=model_group, - category=litellm_exceptions.RateLimitErrorCategory.LITELLM_RATE_LIMIT, - ) from no_deployment - await self._async_override_selector_pre_call_check( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict - strategy, selector, dict(deployment), None - ) - - @contextlib.asynccontextmanager - async def _deployment_slot( - self, deployment: dict, kwargs: Mapping[str, object], parent_otel_span: Span | None - ) -> AsyncGenerator[None, None]: - """Holds the deployment's max_parallel_requests slot, if it has one, around the provider call. Routing - strategy pre-call checks run inside the slot so their rpm accounting stays concurrency-safe.""" - max_parallel_requests_limit: Final = self._get_client( - deployment=deployment, - kwargs=kwargs, - client_type="max_parallel_requests", - ) - async with contextlib.AsyncExitStack() as slot: - if isinstance(max_parallel_requests_limit, MaxParallelRequestsLimit): - slot.enter_context(max_parallel_requests_limit) - await self.async_routing_strategy_pre_call_checks(deployment=deployment, parent_otel_span=parent_otel_span) - yield - - async def async_callback_filter_deployments( - self, - model: str, - healthy_deployments: list[dict], - messages: list[AllMessageValues] | None, - parent_otel_span: Span | None, - request_kwargs: dict | None = None, - logging_obj: LiteLLMLogging | None = None, - ): - """ - For usage-based-routing-v2, enables running rpm checks before the call is made, inside the semaphore. - - -> makes the calls concurrency-safe, when rpm limits are set for a deployment - - Returns: - - None - - Raises: - - Rate Limit Exception - If the deployment is over it's tpm/rpm limits - """ - returned_healthy_deployments = healthy_deployments - for _callback in self._routing_callbacks(): - if isinstance(_callback, CustomLogger): - try: - returned_healthy_deployments = await _callback.async_filter_deployments( - model=model, - healthy_deployments=returned_healthy_deployments, - messages=messages, - request_kwargs=request_kwargs, - parent_otel_span=parent_otel_span, - ) - except Exception as e: - ## LOG FAILURE EVENT - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback.format_exc(), - prefer_async_handlers=True, - ) - ) - raise e - return returned_healthy_deployments - - @staticmethod - def _json_default_stable_id(value: object) -> str: - """json.dumps default= for generate_model_id: plain str() on an arbitrary - object (e.g. a RoutingPlugin instance) falls back to object.__repr__'s - ``, so the hash -- and deployment id -- would - change every restart. Use the class name instead, stable across restarts.""" - return f"{type(value).__module__}.{type(value).__qualname__}" - - @staticmethod - def _core_generate_model_id(model_group: str, litellm_params: dict) -> str: # mutable-ok: hashed read-only - """The stable id of a deployment without ``model_info.id``; see ``litellm.router_api.deployment_id``.""" - return _deployment_id(model_group, litellm_params) - - @staticmethod - def _inherit_builtin_cache_pricing(model_info: dict, backend_model: str, custom_llm_provider: str | None) -> None: - """Fill missing cache pricing on a custom-priced deployment entry from - the backend model's built-in cost map entry, so a deployment that - only spells out ``input_cost_per_token``/``output_cost_per_token`` - does not silently bill cache_read/cache_creation at 0. - - User-specified cache fields always win; only ``None``/missing entries - are inherited, and a user-set cache_read price also covers explicit - cache hits. No-op when the backend model has no canonical entry. - """ - cache_fields: Final = ( - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr", - "cache_creation_input_token_cost_above_200k_tokens", - "cache_read_input_token_cost", - "cache_read_input_token_cost_above_200k_tokens", - ) - inherited_fields: Final = ( - cache_fields - if model_info.get("cache_read_input_token_cost") is not None - else (*cache_fields, "explicit_cache_read_input_token_cost") - ) - missing_fields: Final = tuple(f for f in inherited_fields if model_info.get(f) is None) - if not missing_fields: - return - try: - backend_info: Final = python_host.invoke("get_model_info")( - model=backend_model, custom_llm_provider=custom_llm_provider - ) - except Exception: - return - for field in missing_fields: - backend_value = backend_info.get(field) - if backend_value is not None: - model_info[field] = backend_value - - @staticmethod - def _cost_map_backend_model(deployment: Deployment | DeploymentTypedDict) -> str: - model_info_base_model: Final = ( - deployment.model_info.base_model - if isinstance(deployment, Deployment) - else _configured_model_info(deployment).get("base_model") - ) - if isinstance(model_info_base_model, str) and model_info_base_model: - return model_info_base_model - params: Final = _LITELLM_PARAMS_ADAPTER.validate_python( - deployment.litellm_params.model_dump(exclude_none=True) - if isinstance(deployment, Deployment) - else deployment["litellm_params"] - ) - params_base_model: Final = params.get("base_model") - if isinstance(params_base_model, str) and params_base_model: - return params_base_model - model: Final = ( - deployment.litellm_params.model - if isinstance(deployment, Deployment) - else deployment["litellm_params"].get("model", "") - ) - invoked_model: Final = params.get("model_id") - if ( - isinstance(invoked_model, str) - and invoked_model - and (params.get("custom_llm_provider") == "bedrock" or model.startswith("bedrock/")) - ): - profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}" - profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock") - if not isinstance(profile_info, ModelInfoLookupMiss): - return profile_model - return model - - @staticmethod - def _inherit_builtin_service_tier_pricing( - model_info: dict, # mutable-ok: deployment cost-map entry filled in place - backend_model: str, - custom_llm_provider: str | None, - ) -> None: - """Inherit missing tier rates so a standalone entry does not fall back to custom standard rates.""" - if ptu_terms(model_info) is not None and is_ptu_cost_attribution_enabled(): - return - if all(model_info.get(field) is None for field in ("input_cost_per_token", "output_cost_per_token")): - return - try: - backend_info: Final = python_host.invoke("get_model_info")( - model=backend_model, custom_llm_provider=custom_llm_provider - ) - except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model - return - backend_entry: Final = python_host.dynamic().model_cost.get(backend_info.get("key") or "") - if not isinstance(backend_entry, dict): - return - for field, backend_value in backend_entry.items(): - if not field.endswith(SERVICE_TIER_COST_KEY_SUFFIXES): - continue - if model_info.get(field) is not None or backend_value is None: - continue - model_info[field] = copy.deepcopy(backend_value) - - @staticmethod - def _inherit_builtin_base_rates_for_off_peak( - model_info: dict, # mutable-ok: cost-map entry filled in place - backend_model: str, - custom_llm_provider: str | None, - ) -> None: - """Fill missing pricing fields on a deployment entry that only sets - ``off_peak_pricing``, from the backend model's built-in cost map entry. - - Cost lookup selects the deployment-scoped entry over the shared backend - entry only when the deployment entry carries a base pricing field, and - ``off_peak_pricing`` is deliberately kept off the shared entry, so a - deployment spelling out only its off-peak schedule would otherwise - never receive the discount. The backend model's entire canonical cost - map entry is copied, field by field, so threshold, tiered, - service-tier, cache, character, and per-second rates as well as - companion billing fields like ``web_search_billing_unit`` and the - regional uplift multipliers all carry over, and peak-hour billing - through the deployment entry matches the shared backend entry exactly. - The raw ``litellm.model_cost`` entry is the copy source rather than - ``get_model_info``'s view of it, since that view synthesizes zero flat - token rates for backends without one and storing those would mark a - tiered-only backend explicitly priced free. Values are deep-copied to - keep the builtin entry isolated. User-specified fields always win; - no-op when any base pricing field is already set or the backend model - has no canonical entry. - """ - if not model_info.get("off_peak_pricing"): - return - if any( - model_info.get(field) is not None - for field in ("input_cost_per_token", "input_cost_per_second", "cost_per_second", "tiered_pricing") - ): - return - try: - backend_info: Final = python_host.invoke("get_model_info")( - model=backend_model, custom_llm_provider=custom_llm_provider - ) - except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model - return - backend_entry: Final = python_host.dynamic().model_cost.get(backend_info.get("key") or "") - if not isinstance(backend_entry, dict): - return - for field, backend_value in backend_entry.items(): - if model_info.get(field) is not None or backend_value is None: - continue - model_info[field] = copy.deepcopy(backend_value) - - @staticmethod - def _inherit_builtin_tiered_output_rate( - model_info: dict, backend_model: str, custom_llm_provider: str | None - ) -> None: - """Fill a missing entry-level output rate on a deployment entry whose tier - table omits one, from the backend model's built-in cost map entry. - - A deployment's custom pricing is registered as its own standalone - ``litellm.model_cost`` entry holding only the supplied fields, and the - tiered-cost output fallback reads that same entry, so a tier table that - spells out only input-side rates would bill every completion at 0. - - A user-specified ``output_cost_per_token`` always wins. No-op without a - tier table, when every tier declares its own output rate, or when the - backend model has no canonical entry or no flat output rate: - ``get_model_info`` synthesizes a zero for tiered-only backends, and - storing that zero would mark the deployment as explicitly priced free. - """ - tiers: Final = model_info.get("tiered_pricing") - if not isinstance(tiers, list) or not tiers: - return - if model_info.get("output_cost_per_token") is not None: - return - if all(isinstance(tier, dict) and "output_cost_per_token" in tier for tier in tiers): - return - try: - backend_info: Final = python_host.invoke("get_model_info")( - model=backend_model, custom_llm_provider=custom_llm_provider - ) - except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model - return - backend_rate: Final = backend_info.get("output_cost_per_token") - if backend_rate: - model_info["output_cost_per_token"] = backend_rate - - def _create_deployment( - self, - deployment_info: dict, - _model_name: str, - _litellm_params: dict, - _model_info: dict, - *, - declared_id: str | None = None, - duplicate_ids: frozenset[str] = frozenset(), - ) -> Deployment | None: - """ - Create a deployment object and add it to the model list - - If the deployment is not active for the current environment, it is ignored - - Returns: - - Deployment: The deployment object - - None: If the deployment is not active for the current environment (if 'supported_environments' is set in litellm_params) - """ - try: - config_sourced: Final = _model_info.get("db_model") is not True - identity_error: Final = ( - ptu_identity_error( - declared_id=declared_id, - taken=declared_id in duplicate_ids, - current_id=_model_info.get("id"), - model_name=_model_name, - ) - if config_sourced and ptu_terms(_model_info) is not None - else None - ) - ptu_error: Final = ( - (ptu_config_error(_model_info, model_name=_model_name) or identity_error) if config_sourced else None - ) - declares_split: Final = declares_ptu_shares(_model_info) - if ptu_error is not None and (declares_split or is_ptu_cost_attribution_enabled()): - raise ValueError(ptu_error) - access_windows_error: Final = access_windows_config_error(_model_info, model_name=_model_name) - if access_windows_error is not None: - raise ValueError(access_windows_error) - capacity_warning: Final = ( - ptu_capacity_warning( - _model_name, - MappingProxyType( - { # pyright: ignore[reportUnknownArgumentType] # router deployment dicts are untyped - "model_info": _model_info, - "litellm_params": _litellm_params, - } - ), - ) - if is_ptu_cost_attribution_enabled() - else None - ) - if capacity_warning is not None: - verbose_router_logger.warning(capacity_warning) - zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None - merged_params: Final[Mapping[str, object]] = ( - _litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing}) - ) - litellm_params: Final[LiteLLM_Params] = validate_litellm_params_strictly( - LiteLLM_Params, dict(**merged_params) - ) - declared_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(_litellm_params) - warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=declared_params) - warn_on_ignored_model_capability_flags(model_name=_model_name, litellm_params=declared_params) - warn_on_unenforced_endpoint_flag(model_name=_model_name, model_info=_model_info) - deployment = without_model_capability_flags( - Deployment( - **deployment_info, - model_name=_model_name, - litellm_params=litellm_params, - model_info=_model_info, - ) - ) - refuse_unservable_numbers(deployment.to_json(exclude_none=True, warnings=False)) # pyright: ignore[reportUnknownMemberType] # pydantic's to_json - PythonCore._register_deployment_pricing(deployment) - - ## Check if LLM Deployment is allowed for this deployment - if self.deployment_is_active_for_environment(deployment=deployment) is not True: - verbose_router_logger.warning( - "Ignoring deployment %s as it is not active for environment %s", - deployment.model_name, - deployment.model_info["supported_environments"], - ) - return None - - # Validate tag_regex patterns BEFORE adding the deployment so we never - # have partially-initialised router state if a pattern is invalid. - _tag_regex: Final = deployment.litellm_params.get("tag_regex") or [] - for pattern in _tag_regex: - try: - re.compile(pattern) - except re.error as exc: - raise ValueError( - f"Invalid regex in tag_regex for model '{deployment.model_name}': {pattern!r} — {exc}" - ) from exc - - deployment = self._add_deployment(deployment=with_admin_api_key_secret_reference(deployment)) - - model: Final = deployment.to_json(exclude_none=True) - - self._add_model_to_list_and_index_map(model=model, model_id=deployment.model_info.id) - return deployment - except Exception as e: - if self.ignore_invalid_deployments: - if isinstance(e, litellm_exceptions.BadRequestError): - self._provider_unresolved_deployments = ( - *self._provider_unresolved_deployments, - partial( - self._create_deployment, - deployment_info=deployment_info, - _model_name=_model_name, - _litellm_params=_litellm_params, - _model_info=_model_info, - declared_id=declared_id, - duplicate_ids=duplicate_ids, - ), - ) - verbose_router_logger.exception( - "Error creating deployment: %s, ignoring and continuing with other deployments.", e - ) - return None - else: - raise e - - def _is_auto_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: - """ - Check if the deployment is an auto-router deployment (semantic router). - - Returns True if the litellm_params model starts with "auto_router/" - but NOT "auto_router/complexity_router" or "auto_router/adaptive_router" - (which use the complexity-router and adaptive-router strategies). - """ - return classify_strategy_router_model(litellm_params.model) == "semantic" - - @staticmethod - def _deployment_tags(deployment: Deployment) -> tuple[str, ...]: - """Deployment tags used to disambiguate strategy registries keyed by model_name.""" - return tuple(deployment.litellm_params.tags or ()) - - def init_auto_router_deployment(self, deployment: Deployment): - """ - Initialize the auto-router deployment. - - This will initialize the auto-router and add it to the auto-routers dictionary. - """ - from litellm.router_strategy.auto_router.auto_router import AutoRouter - - auto_router_config_path: Final[str | None] = deployment.litellm_params.auto_router_config_path - auto_router_config: Final[str | None] = deployment.litellm_params.auto_router_config - if auto_router_config_path is None and auto_router_config is None: - raise ValueError( - "auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params" - ) - - default_model: Final[str | None] = deployment.litellm_params.auto_router_default_model - if default_model is None: - raise ValueError( - "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params" - ) - - embedding_model: Final[str | None] = deployment.litellm_params.auto_router_embedding_model - if embedding_model is None: - raise ValueError( - "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params" - ) - - max_input_chars: Final[int | None] = deployment.litellm_params.auto_router_max_input_chars - - autor_router: Final[AutoRouter] = AutoRouter( - model_name=deployment.model_name, - auto_router_config_path=auto_router_config_path, - auto_router_config=auto_router_config, - default_model=default_model, - embedding_model=embedding_model, - ctx=self._core_ctx, - max_input_chars=(max_input_chars if max_input_chars is not None else DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS), - ) - self._register_pre_routing_strategy( - registry=self._core_auto_routers, - deployment=deployment, - strategy=autor_router, - strategy_label="Auto-router", - ) - - def _is_complexity_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: - """ - Check if the deployment is a complexity-router deployment. - - Returns True if the litellm_params model starts with "auto_router/complexity_router" - """ - return classify_strategy_router_model(litellm_params.model) == "complexity" - - def _core_config_deployments(self) -> Iterator[Mapping[str, object]]: - """The model_list rows that came from config.yaml rather than the DB (``model_info.db_model`` unset).""" - for deployment in self._core_model_list: - if not isinstance(deployment, Mapping): - continue - model_info = deployment.get("model_info") - if not (isinstance(model_info, Mapping) and model_info.get("db_model")): - yield deployment - - def auto_router_capability_violation(self, capability: GatedAutoRouterCapability) -> str | None: - """ - Why one more router claiming ``capability`` cannot join this router, or None when it can. - - Judged against every deployment currently on the model_list; an upsert pops the row being - edited first, so an edit of an existing gated router keeps its own slot. The limit is - resolved on every call through ``auto_router_capability_limit``; unset means unlimited, - which is the SDK default, and the proxy injects a resolver backed by its license. - """ - limit: Final = self.auto_router_capability_limit() if self.auto_router_capability_limit is not None else None - others: Final = count_capability_routers( - (deployment for deployment in self._core_model_list if isinstance(deployment, Mapping)), - capability=capability, - ) - return capability_limit_violation(capability=capability, held=others + 1, limit=limit) - - def init_complexity_router_deployment(self, deployment: Deployment): - """ - Initialize the complexity-router deployment. - - This will initialize the complexity-router and add it to the complexity-routers dictionary. - """ - # Import here to avoid circular imports — ComplexityRouter is a CustomLogger - # subclass that imports litellm internals which depend on router.py. - # This matches the AutoRouter pattern in init_auto_router_deployment above. - from litellm.router_strategy.complexity_router.complexity_router import ( - ComplexityRouter, - ) - from litellm.router_strategy.complexity_router.config import ( - ComplexityRouterConfig, - ) - - complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config - for capability in claimed_capabilities(complexity_router_config): - if (limit_violation := self.auto_router_capability_violation(capability)) is not None: - raise ValueError(limit_violation) - - default_model: Final = ( - ComplexityRouterConfig.model_validate(complexity_router_config).resolve_default_model( - deployment.litellm_params.complexity_router_default_model - ) - if complexity_router_config - else deployment.litellm_params.complexity_router_default_model - ) - - if default_model is None: - raise ValueError( - "complexity_router_default_model is required for complexity-router deployments, " - "or configure tiers in complexity_router_config. Please set it in the litellm_params" - ) - - complexity_router: Final[ComplexityRouter] = ComplexityRouter( - model_name=deployment.model_name, - default_model=default_model, - ctx=self._core_ctx, - complexity_router_config=complexity_router_config, - ) - self._register_pre_routing_strategy( - registry=self._core_complexity_routers, - deployment=deployment, - strategy=complexity_router, - strategy_label="Complexity-router", - ) - if complexity_router._uses_deployment_pin: - self._ensure_deployment_affinity_callback() - - def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: - """True when this deployment opts in via the `auto_router/adaptive_router` model prefix.""" - return classify_strategy_router_model(litellm_params.model) == "adaptive" - - def _deployment_participates_in_adaptive_routing(self, litellm_params: LiteLLM_Params) -> bool: - """True when this deployment owns an `adaptive_routers` entry once finalized: - a dedicated adaptive router, or a complexity router whose config enables the - adaptive companion. Mirrors the two arms of - `_finalize_adaptive_router_if_configured`, which is the registry's only writer.""" - if self._is_adaptive_router_deployment(litellm_params=litellm_params): - return True - if not self._is_complexity_router_deployment(litellm_params=litellm_params): - return False - config: Final = litellm_params.complexity_router_config - if not config: - return False - adaptive_flag: Final[object] = config.get("adaptive") - return bool(adaptive_flag) - - @staticmethod - def _has_registered_strategy( - registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], - model_name: str, - tags: tuple[str, ...], - ) -> bool: - """True when a strategy for this (model_name, tags) pair is already registered.""" - return any(existing.tags == tags for existing in registry.get(model_name, [])) - - def _register_pre_routing_strategy( - self, - registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], - deployment: Deployment, - strategy: _PreRoutingStrategyT, - strategy_label: str, - ) -> None: - """ - Register `strategy` under `deployment.model_name`, scoped by its tags. - Reusing a `model_name` is allowed when tags differ; a repeat of the same - (model_name, tags) pair is a misconfiguration and is rejected. - """ - tags: Final = self._deployment_tags(deployment) - if self._has_registered_strategy(registry, deployment.model_name, tags): - raise ValueError( - f"{strategy_label} deployment {deployment.model_name} with tags {list(tags)} already exists. " - "Please use a different model name or set different tags." - ) - registry[deployment.model_name] = [ - *registry.get(deployment.model_name, []), - TaggedPreRoutingStrategy( - tags=tags, - strategy=strategy, - definition_fingerprint=self._routing_definition_fingerprint(deployment), - deployment=deployment, - ), - ] - - @staticmethod - def _routing_definition_fingerprint(deployment: Deployment) -> str | None: - """Identify configured routing fields, excluding external files and live strategy state.""" - try: - definition: Final = json.dumps( - deployment.litellm_params.model_dump( - include=STRATEGY_ROUTER_PARAM_FIELDS | {"model", "tags"}, exclude_none=True, warnings=False - ), - sort_keys=True, - default=PythonCore._json_default_stable_id, - ) - except (TypeError, ValueError): - return None - return PythonCore._core_generate_model_id(deployment.model_name, {"routing_definition": definition}) - - @staticmethod - def _unregister_pre_routing_strategy( - registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], - model_name: str, - tags: tuple[str, ...], - ) -> bool: - """Drop the strategy registered for this exact (model_name, tags) pair, leaving - strategies registered under the same name with different tags in place. Returns - whether anything was actually dropped.""" - existing: Final = registry.get(model_name, []) - remaining: Final = [entry for entry in existing if entry.tags != tags] - if len(remaining) == len(existing): - return False - if remaining: - registry[model_name] = remaining - else: - registry.pop(model_name, None) - return True - - def _unregister_pre_routing_strategy_for_deployment(self, deployment: Deployment) -> None: - """ - Release the pre-routing strategy a deployment holds, so removing it from the - model_list also frees its (model_name, tags) slot. - - Without this, re-adding the deployment (an edit arriving via upsert_deployment, - or a router recreated under a name that was deleted earlier) hits the - "already exists" guard in `_register_pre_routing_strategy`, which - `ignore_invalid_deployments` swallows - the deployment then silently never - makes it back into the model_list. - - Released from every registry rather than the first match, because registration is - one-to-many: a complexity router configured with `adaptive` is also registered in - `adaptive_routers` under the same (model_name, tags) by the deferred finalize pass. - Guarded on the auto_router/ prefix so removing a *regular* deployment can't evict a - router that merely shares its model_name. - """ - if not deployment.litellm_params.model.startswith("auto_router/"): - return - model_name: Final = deployment.model_name - tags: Final = self._deployment_tags(deployment) - for registry in (self._core_auto_routers, self._core_complexity_routers, self._core_quality_routers): - self._unregister_pre_routing_strategy(registry, model_name, tags) - if self._unregister_pre_routing_strategy(self._core_adaptive_routers, model_name, tags): - self._sync_adaptive_router_hooks() - - def _finalize_adaptive_router_if_configured(self) -> None: - """Locate every adaptive-router deployment in the finalized model_list and - build an AdaptiveRouter for each. Safe no-op when none are configured. - Idempotent: skips any deployment whose (model_name, tags) pair is already - initialized, so hot-reloads don't rebuild routers that would lose state.""" - for entry in self._core_model_list or []: - lp = entry.get("litellm_params") if isinstance(entry, dict) else entry.litellm_params - lp_model = (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None - if not (lp_model and lp_model.startswith("auto_router/adaptive_router")): - continue - model_name = entry.get("model_name") if isinstance(entry, dict) else entry.model_name - if not model_name or not lp: - continue - deployment = Deployment( - model_name=model_name, - litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params.model_validate(lp)), - model_info=(entry.get("model_info") if isinstance(entry, dict) else entry.model_info), - ) - if self._has_registered_strategy( - self._core_adaptive_routers, model_name, self._deployment_tags(deployment) - ): - continue - self.init_adaptive_router_deployment(deployment=deployment) - - for model_name, tagged_complexity_routers in self._core_complexity_routers.items(): - for tagged in tagged_complexity_routers: - complexity_router = tagged.strategy - if not complexity_router.config.adaptive: - continue - if self._has_registered_strategy(self._core_adaptive_routers, model_name, tagged.tags): - continue - adaptive_router: AdaptiveRouter | None = complexity_router._ensure_adaptive_router() - if adaptive_router is not None: - self._core_adaptive_routers[model_name] = [ - *self._core_adaptive_routers.get(model_name, []), - TaggedPreRoutingStrategy( - tags=tagged.tags, - strategy=adaptive_router, - definition_fingerprint=tagged.definition_fingerprint, - deployment=tagged.deployment, - ), - ] - - self._sync_adaptive_router_hooks() - - def _sync_adaptive_router_hooks(self) -> None: - """Rebuild the AdaptiveRouterPostCallHook set so it is exactly one hook per - currently registered adaptive router. Run at every point the adaptive registry - changes, otherwise a released router keeps recording turns through its hook.""" - from litellm.router_strategy.adaptive_router.hooks import ( - AdaptiveRouterPostCallHook, - AdaptiveRouterResponseHeaderHook, - ) - - for callback in self._host.callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): - self._host.callbacks.remove_callback_from_all_lists(callback) - for tagged_adaptive_routers in self._core_adaptive_routers.values(): - for tagged in tagged_adaptive_routers: - self._host.callbacks.add_litellm_callback(AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy)) - if self._core_adaptive_routers: - self._host.callbacks.add_litellm_callback(AdaptiveRouterResponseHeaderHook()) - - def init_adaptive_router_deployment(self, deployment: Deployment) -> None: - """ - Build an AdaptiveRouter instance for this deployment and register its - post-call hook. Multiple adaptive routers can coexist on a single Router, - keyed by `deployment.model_name`. - - `model_to_prefs` and `model_to_cost` are derived from the OTHER models - already registered in `self._core_model_list` whose `model_name` appears in - `available_models`. Models not yet registered fall back to defaults. - """ - # Local import: AdaptiveRouter -> hooks -> classifier all import litellm - # internals which transitively import this module. (AGENTS.md exception clause.) - from litellm.router_strategy.adaptive_router.adaptive_router import ( - AdaptiveRouter, - ) - from litellm.router_strategy.adaptive_router.hooks import ( - AdaptiveRouterPostCallHook, - ) - from litellm.types.router import ( - AdaptiveRouterConfig, - AdaptiveRouterPreferences, - ) - - raw_config: Final = deployment.litellm_params.adaptive_router_config - if raw_config is None: - raise ValueError("adaptive_router_config is required for adaptive-router deployments.") - - config: Final = AdaptiveRouterConfig(**raw_config) - - model_to_prefs: Final[dict[str, AdaptiveRouterPreferences]] = {} - model_to_cost: Final[dict[str, float]] = {} - # O(k) via the name→indices map: only touch deployments whose name - # is listed in `available_models`, instead of scanning model_list. - for name in config.available_models: - deployments = self._core_deployments_named(name) - if not deployments: - continue - d = deployments[0] - mi_dict: dict[str, Any] = d.get("model_info") or {} - prefs_raw = mi_dict.get("adaptive_router_preferences") - if prefs_raw is not None: - model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) - - # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set. - cost = d.get("litellm_params", {}).get("input_cost_per_token") - if cost is None: - cost = mi_dict.get("input_cost_per_token") - if cost is not None: - model_to_cost[name] = float(cost) - - adaptive_router: Final = AdaptiveRouter( - router_name=deployment.model_name, - config=config, - model_to_prefs=model_to_prefs, - model_to_cost=model_to_cost, - ) - self._register_pre_routing_strategy( - registry=self._core_adaptive_routers, - deployment=deployment, - strategy=adaptive_router, - strategy_label="Adaptive-router", - ) - self._host.callbacks.add_litellm_callback(AdaptiveRouterPostCallHook(adaptive_router=adaptive_router)) - verbose_router_logger.info( - "AdaptiveRouter[%s] initialized with %d models", - deployment.model_name, - len(config.available_models), - ) - - def _is_quality_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: - """ - Check if the deployment is a quality-router deployment. - - Returns True if the litellm_params model starts with "auto_router/quality_router". - """ - return classify_strategy_router_model(litellm_params.model) == "quality" - - def init_quality_router_deployment(self, deployment: Deployment): - """ - Initialize the quality-router deployment. - - Resolves the default model from either `quality_router_default_model` or - `quality_router_config["default_model"]`, then instantiates the - QualityRouter and stores it in `self._core_quality_routers`. - """ - # Import here to mirror the AutoRouter / ComplexityRouter init pattern - # and avoid circular imports. - from litellm.router_strategy.quality_router.quality_router import ( - QualityRouter, - ) - - quality_router_config: Final[dict | None] = deployment.litellm_params.quality_router_config - - default_model: str | None = deployment.litellm_params.quality_router_default_model - if default_model is None and quality_router_config: - default_model = quality_router_config.get("default_model") - - if default_model is None: - raise ValueError( - "quality_router_default_model is required for quality-router deployments, " - "or set default_model in quality_router_config. Please configure it in the litellm_params" - ) - - quality_router: Final[QualityRouter] = QualityRouter( - model_name=deployment.model_name, - default_model=default_model, - ctx=self._core_ctx, - quality_router_config=quality_router_config, - ) - self._register_pre_routing_strategy( - registry=self._core_quality_routers, - deployment=deployment, - strategy=quality_router, - strategy_label="Quality-router", - ) - - def deployment_is_active_for_environment(self, deployment: Deployment) -> bool: - """ - Function to check if a llm deployment is active for a given environment. Allows using the same config.yaml across multople environments - - Requires `LITELLM_ENVIRONMENT` to be set in .env. Valid values for environment: - - development - - staging - - production - - Raises: - - ValueError: If LITELLM_ENVIRONMENT is not set in .env or not one of the valid values - - ValueError: If supported_environments is not set in model_info or not one of the valid values - """ - return model_info_is_active_for_environment(model_info=deployment.model_info) - - def set_model_list(self, model_list: list): - """Replace every deployment. All or nothing: when an entry raises, the deployments served - before the call are put back and the error is raised.""" - previous: Final = copy.deepcopy(list(self.__dict__.get("_core_model_list") or ())) - try: - self._replace_model_list(model_list) # pyright: ignore[reportUnknownMemberType] # the legacy untyped model list - except Exception: - self._replace_model_list(previous) # pyright: ignore[reportUnknownMemberType] # see above - raise - - def _replace_model_list(self, model_list: list) -> None: # pyright: ignore[reportUnknownParameterType, reportMissingTypeArgument] # the legacy untyped model list - original_model_list: Final = copy.deepcopy(model_list) - self._discovered_model_info_cache.flush_cache() - with self._deployment_index_lock: - self._core_model_list = [] - self.model_id_to_deployment_index_map = {} # Reset the index - self._core_model_name_to_deployment_indices = {} # Reset the model_name index - self._core_team_model_to_deployment_indices = {} # Reset the team_model index - self._core_team_pattern_routers = {} - self._core_team_public_model_names = frozenset() - # Reset per-strategy router registries so hot-reload doesn't leave - # stale routers pointing at the old model_list. - self._core_quality_routers = {} - self._core_complexity_routers = {} - self._core_auto_routers = {} - self._provider_unresolved_deployments = () - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works - - declared_ids: Final = tuple( - str(entry["model_info"]["id"]) - for entry in original_model_list - if isinstance(entry.get("model_info"), dict) and entry["model_info"].get("id") is not None - ) - duplicate_ids: Final = frozenset(model_id for model_id in declared_ids if declared_ids.count(model_id) > 1) - - ptu_declared: Final = tuple( - str(entry.get("model_name")) - for entry in original_model_list - if isinstance(entry.get("model_info"), dict) - and entry["model_info"].get("db_model") is not True - and declares_ptu(entry["model_info"]) - ) - if ptu_declared and not is_ptu_cost_attribution_enabled(): - verbose_router_logger.warning( - "PTU fields are set on config.yaml deployment(s) %s, but PTU cost attribution is disabled, so no " - "flat cost accrues and this traffic is billed per token. Set %s=True to enable it", - ", ".join(ptu_declared), - PTU_COST_ATTRIBUTION_ENV_VAR, - ) - - for model in original_model_list: - _model_name = model.pop("model_name") - _litellm_params = model.pop("litellm_params") - ## check if litellm params in os.environ - if isinstance(_litellm_params, dict): - for k, v in _litellm_params.items(): - if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k): - _litellm_params[k] = get_secret(v) - - _model_info: dict = model.pop("model_info", {}) - - declared_id = None if _model_info.get("id") is None else str(_model_info["id"]) - - # check if model info has id - if "id" not in _model_info: - _id = self._core_generate_model_id(_model_name, _litellm_params) - _model_info["id"] = _id - - if _litellm_params.get("organization", None) is not None and isinstance( - _litellm_params["organization"], list - ): # Addresses https://github.com/BerriAI/litellm/issues/3949 - for org in _litellm_params["organization"]: - _litellm_params["organization"] = org - self._create_deployment( - deployment_info=model, - _model_name=_model_name, - _litellm_params=_litellm_params, - _model_info=_model_info, - declared_id=declared_id, - duplicate_ids=duplicate_ids, - ) - else: - self._create_deployment( - deployment_info=model, - _model_name=_model_name, - _litellm_params=_litellm_params, - _model_info=_model_info, - declared_id=declared_id, - duplicate_ids=duplicate_ids, - ) - - verbose_router_logger.debug("\nInitialized Model List %s", self._core_get_model_names()) - self._core_model_names = {m["model_name"] for m in model_list} - - # Note: model_name_to_deployment_indices is already built incrementally - # by _create_deployment -> _add_model_to_list_and_index_map - - # Deferred: build the AdaptiveRouter strategy now that all underlying - # deployments have been registered. - self._finalize_adaptive_router_if_configured() - - def _add_deployment(self, deployment: Deployment) -> Deployment: - import os - - #### VALIDATE MODEL ######## - # Check if this is a prompt management model before validating as LLM provider - litellm_model: Final = deployment.litellm_params.model - if isinstance(deployment.litellm_params.drop_params, str): - verbose_router_logger.warning( - "model=%s drop_params=%r is not a flag value, treating it as unset", - deployment.model_name, - deployment.litellm_params.drop_params, - ) - is_prompt_management_model = False - - if "/" in litellm_model: - split_litellm_model: Final = litellm_model.split("/")[0] - if split_litellm_model in self._host.dynamic().known_custom_logger_compatible_callbacks: - is_prompt_management_model = True - - if is_prompt_management_model: - # For prompt management models, skip LLM provider validation - # The actual model will be resolved at runtime from the prompt file - _model = litellm_model - custom_llm_provider = None - dynamic_api_key = None - api_base = None - elif declared_authenticating_provider(litellm_model, deployment.litellm_params.custom_llm_provider) is not None: - _model, custom_llm_provider = ( # rebind-ok: each branch resolves the deployment provider - self._model_group_llm_provider(deployment.litellm_params) - ) - else: - # check if model provider in supported providers - ( - _model, - custom_llm_provider, - dynamic_api_key, - api_base, - ) = self._host.invoke("get_llm_provider")( - model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None), - api_base=deployment.litellm_params.api_base, - litellm_params=deployment.litellm_params, - ) - # done reading model["litellm_params"] - # Check if provider is supported: either in enum or JSON-configured - if ( - custom_llm_provider not in self._host.dynamic().provider_list - and not JSONProviderRegistry.exists(custom_llm_provider) - and not is_registered_custom_provider(custom_llm_provider) - ): - raise Exception(f"Unsupported provider - {custom_llm_provider}") - - #### DEPLOYMENT NAMES INIT ######## - self._core_deployment_names.append(deployment.litellm_params.model) - ############ Users can either pass tpm/rpm as a litellm_param or a router param ########### - # for get_available_deployment, we use the litellm_param["rpm"] - # in this snippet we also set rpm to be a litellm_param - if deployment.litellm_params.rpm is None and getattr(deployment, "rpm", None) is not None: - deployment.litellm_params.rpm = getattr(deployment, "rpm") - - if deployment.litellm_params.tpm is None and getattr(deployment, "tpm", None) is not None: - deployment.litellm_params.tpm = getattr(deployment, "tpm") - - # Check if user is trying to use model_name == "*" - # this is a catch all model for their specific api key - # if deployment.model_name == "*": - # if deployment.litellm_params.model == "*": - # # user wants to pass through all requests to litellm.acompletion for unknown deployments - # self._core_router_general_settings.pass_through_all_models = True - # else: - # self._core_default_deployment = deployment.to_json(exclude_none=True) - # Check if user is using provider specific wildcard routing - # example model_name = "databricks/*" or model_name = "anthropic/*" - if "*" in deployment.model_name: - # store this as a regex pattern - all deployments matching this pattern will be sent to this deployment - # Store deployment.model_name as a regex pattern - self._core_pattern_router.add_pattern(deployment.model_name, deployment.to_json(exclude_none=True)) - if deployment.model_info.id: - self._core_provider_default_deployment_ids.append(deployment.model_info.id) - - _team_id: Final = deployment.model_info.get("team_id") - _team_public_model_name: Final[str | None] = deployment.model_info.get("team_public_model_name") - if _team_id is not None and _team_public_model_name is not None and "*" in _team_public_model_name: - if _team_id not in self._core_team_pattern_routers: - self._core_team_pattern_routers[_team_id] = PatternMatchRouter() - self._core_team_pattern_routers[_team_id].add_pattern( - _team_public_model_name, deployment.to_json(exclude_none=True) - ) - - # Azure GPT-Vision Enhancements, users can pass os.environ/ - data_sources: Final = deployment.litellm_params.get("dataSources", []) or [] - - for data_source in data_sources: - params = data_source.get("parameters", {}) - for param_key in ["endpoint", "key"]: - # if endpoint or key set for Azure GPT Vision Enhancements, check if it's an env var - if param_key in params and params[param_key].startswith("os.environ/"): - env_name = params[param_key].replace("os.environ/", "") - params[param_key] = os.environ.get(env_name, "") - - # # init OpenAI, Azure clients - # InitalizeOpenAISDKClient.set_client( - # litellm_router_instance=self, model=deployment.to_json(exclude_none=True) - # ) - - if custom_llm_provider is not None: - self._initialize_deployment_for_pass_through( - deployment=deployment, - custom_llm_provider=custom_llm_provider, - ) - - ######################################################### - # Check if this is an auto-router deployment - ######################################################### - if self._is_auto_router_deployment(litellm_params=deployment.litellm_params): - self.init_auto_router_deployment(deployment=deployment) - - ######################################################### - # Check if this is a complexity-router deployment - ######################################################### - if self._is_complexity_router_deployment(litellm_params=deployment.litellm_params): - self.init_complexity_router_deployment(deployment=deployment) - - # NOTE: adaptive-router deployments are deferred to the end of - # set_model_list() because their init needs visibility into the OTHER - # deployments listed in `available_models` (which may not yet have - # been processed when this one is created). - ######################################################### - # Check if this is a quality-router deployment - ######################################################### - if self._is_quality_router_deployment(litellm_params=deployment.litellm_params): - self.init_quality_router_deployment(deployment=deployment) - - return deployment - - def _initialize_deployment_for_pass_through(self, deployment: Deployment, custom_llm_provider: str): - """ - Optional: Register vertex credentials for pass-through endpoints if `deployment.litellm_params.use_in_pass_through` is True - - Other providers need no registration here: PassthroughEndpointRouter.get_credentials resolves their credentials per-request from the live router deployments - """ - if deployment.litellm_params.use_in_pass_through is not True: - return - if custom_llm_provider != "vertex_ai": - return - credential_name: Final = deployment.litellm_params.litellm_credential_name - credential_values: Final = ( - CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else {} - ) - from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES - - vertex_project: Final = ( - cast("str | None", credential_values.get("vertex_project")) # cast-ok: vertex credential values are str - or deployment.litellm_params.vertex_project - ) - vertex_location: Final = ( - cast("str | None", credential_values.get("vertex_location")) # cast-ok: vertex credential values are str - or deployment.litellm_params.vertex_location - ) - vertex_credentials: Final = ( - cast( # cast-ok: vertex_credentials holds the typed credential union - "VERTEX_CREDENTIALS_TYPES | None", credential_values.get("vertex_credentials") - ) - or deployment.litellm_params.vertex_credentials - ) - - if vertex_project is None or vertex_location is None: - raise ValueError( - "vertex_project, and vertex_location must be set in litellm_params for pass-through endpoints." - ) - router_host().add_vertex_pass_through_credentials( - project_id=vertex_project, - location=vertex_location, - vertex_credentials=vertex_credentials, - ) - - def add_deployment(self, deployment: Deployment) -> Deployment | None: - """ - Parameters: - - deployment: Deployment - the deployment to be added to the Router - - Returns: - - The added deployment - - OR None (if deployment already exists) - """ - # check if deployment already exists - - _deployment_model_id: Final = deployment.model_info.id - if _deployment_model_id and self._core_has_model_id(_deployment_model_id): - return None - - declared_params: Final = deployment.litellm_params.model_dump(exclude_none=True) - warn_on_provider_credential_mismatch(model_name=deployment.model_name, litellm_params=declared_params) - warn_on_ignored_model_capability_flags(model_name=deployment.model_name, litellm_params=declared_params) - warn_on_unenforced_endpoint_flag( - model_name=deployment.model_name, model_info=deployment.model_info.model_dump() - ) - served: Final = with_admin_api_key_secret_reference(without_model_capability_flags(deployment)) - - # add to model list - _deployment: Final = served.to_json(exclude_none=True) - refuse_unservable_numbers(_deployment) - # initialize client - self._add_deployment(deployment=served) - - PythonCore._register_deployment_pricing(deployment=served) - - # add to model names - self._add_model_to_list_and_index_map(model=_deployment, model_id=served.model_info.id) - self._core_model_names.add(served.model_name) - self._sync_deployment_budget_config(deployment=served) - return served - - def _update_deployment_indices_after_removal(self, model_id: str, removal_idx: int) -> None: - """ - Helper method to update deployment indices after a deployment has been removed from model_list. - - Parameters: - - model_id: str - the id of the deployment that was removed - - removal_idx: int - the index where the deployment was removed from model_list - """ - self._discovered_model_info_cache.delete_cache(model_id) - # Update indices for all models after the removed one - for deployment_id, idx in self.model_id_to_deployment_index_map.items(): - if idx > removal_idx: - self.model_id_to_deployment_index_map[deployment_id] = idx - 1 - # Remove the deleted model from index - if model_id in self.model_id_to_deployment_index_map: - del self.model_id_to_deployment_index_map[model_id] - - # Update model_name_to_deployment_indices - for model_name, indices in list(self._core_model_name_to_deployment_indices.items()): - # Build new list without mutating the original - updated_indices = [] - for idx in indices: - if idx == removal_idx: - # Skip the removed index - continue - elif idx > removal_idx: - # Decrement indices after removal - updated_indices.append(idx - 1) - else: - # Keep indices before removal unchanged - updated_indices.append(idx) - - # Update or remove the entry - if len(updated_indices) > 0: - self._core_model_name_to_deployment_indices[model_name] = updated_indices - else: - del self._core_model_name_to_deployment_indices[model_name] - self._core_model_names.discard(model_name) - - # Update team_model_to_deployment_indices - for key, indices in list(self._core_team_model_to_deployment_indices.items()): - # Build new list without mutating the original - updated_indices = [] - for idx in indices: - if idx == removal_idx: - # Skip the removed index - continue - elif idx > removal_idx: - # Decrement indices after removal - updated_indices.append(idx - 1) - else: - # Keep indices before removal unchanged - updated_indices.append(idx) - - # Update or remove the entry - if len(updated_indices) > 0: - self._core_team_model_to_deployment_indices[key] = updated_indices - else: - del self._core_team_model_to_deployment_indices[key] - self._core_team_public_model_names = frozenset( - public_model_name for _, public_model_name in self._core_team_model_to_deployment_indices - ) - - self._core_pattern_router.remove_deployment(model_id) - for team_id in list(self._core_team_pattern_routers.keys()): - team_pattern_router = self._core_team_pattern_routers[team_id] - team_pattern_router.remove_deployment(model_id) - if not team_pattern_router.patterns: - del self._core_team_pattern_routers[team_id] - - def _update_team_model_index(self, model: dict, idx: int) -> None: - """ - Helper to update team_model_to_deployment_indices for a single deployment. - - Parameters: - - model: dict - the deployment to index - - idx: int - the index in model_list - """ - team_id: Final = (model.get("model_info") or {}).get("team_id") - team_public_model_name: Final = (model.get("model_info") or {}).get("team_public_model_name") - if team_id and team_public_model_name: - key: Final = (team_id, team_public_model_name) - self._core_team_public_model_names = self._core_team_public_model_names | frozenset( - {team_public_model_name} - ) - if key not in self._core_team_model_to_deployment_indices: - self._core_team_model_to_deployment_indices[key] = [] - if idx not in self._core_team_model_to_deployment_indices[key]: - self._core_team_model_to_deployment_indices[key].append(idx) - - def _add_model_to_list_and_index_map(self, model: dict, model_id: str | None = None) -> None: - """ - Helper method to add a model to the model_list and update both indices. - - Parameters: - - model: dict - the model to add to the list - - model_id: Optional[str] - the model ID to use for indexing. If None, will try to get from model["model_info"]["id"] - """ - with self._deployment_index_lock: - idx: Final = len(self._core_model_list) - self._core_model_list.append(model) - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - - # Update model_id index for O(1) lookup - if model_id is not None: - self.model_id_to_deployment_index_map[model_id] = idx - elif model.get("model_info", {}).get("id") is not None: - self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx - - # Update model_name index for O(1) lookup - model_name: Final = model.get("model_name") - if model_name: - if model_name not in self._core_model_name_to_deployment_indices: - self._core_model_name_to_deployment_indices[model_name] = [] - self._core_model_name_to_deployment_indices[model_name].append(idx) - - # Update team_model index for O(1) team-scoped lookup - self._update_team_model_index(model, idx) - - def _core_upsert_deployment(self, deployment: Deployment) -> Deployment | None: - """ - Add or update deployment - Parameters: - - deployment: Deployment - the deployment to be added to the Router - - Returns: - - The added/updated deployment - """ - _deployment_model_id: Final = deployment.model_info.id or "" - _deployment_on_router: Final[Deployment | None] = self._core_get_deployment(model_id=_deployment_model_id) - served: Final = without_model_capability_flags(deployment) - # `_create_deployment` lists the model name before it can still fail; a dropped deployment - # takes its name back off the list. - names_before: Final = len(self._core_deployment_names) - try: - if _deployment_on_router is not None: - # deployment with this model_id exists on the router - if ( - deployment.model_name == _deployment_on_router.model_name - and served.litellm_params == _deployment_on_router.litellm_params - and deployment.model_info == _deployment_on_router.model_info - ): - # No need to update - return None - - # if there is a new litellm param -> then update the deployment - # remove the previous deployment - with self._deployment_index_lock: - removal_idx: Final = self.model_id_to_deployment_index_map.get(_deployment_model_id) - if removal_idx is not None: - self._core_model_list.pop(removal_idx) - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - self._update_deployment_indices_after_removal( - model_id=_deployment_model_id, removal_idx=removal_idx - ) - - # Free the outgoing deployment's pre-routing strategy slot (keyed by the - # OLD model_name/tags) before the re-add below re-registers it. - self._unregister_pre_routing_strategy_for_deployment(deployment=_deployment_on_router) - - # if the model_id is not in router - self.add_deployment(deployment=deployment) - # add_deployment() builds every strategy EXCEPT the adaptive one, which - # set_model_list() defers until the whole model_list is visible. Re-run that - # deferred pass so an adaptive router whose slot was just released above is - # rebuilt rather than left unregistered. - if self._deployment_participates_in_adaptive_routing(litellm_params=deployment.litellm_params) or ( - _deployment_on_router is not None - and self._deployment_participates_in_adaptive_routing( - litellm_params=_deployment_on_router.litellm_params - ) - ): - self._finalize_adaptive_router_if_configured() - return deployment - except Exception as e: - if self.ignore_invalid_deployments: - verbose_router_logger.warning( - "Error upserting deployment %s (id=%s): %s. Dropping it and continuing with other deployments.", - deployment.model_name, - deployment.model_info.id, - e, - ) - del self._core_deployment_names[names_before:] - self._restore_deployment_after_failed_upsert( - previous_deployment=_deployment_on_router, model_id=_deployment_model_id - ) - return None - else: - raise e - - def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None: - """Put a deployment back the way it was before a failed upsert popped it. - - A rollback re-admits state that was already serving, so it does not go through the - capability ceiling a newcomer gets: with the ceiling tightened since the deployment first - registered, judging the rollback would drop a serving router over an unrelated failed edit. - """ - if previous_deployment is None or self._core_has_model_id(model_id): - return - limit_resolver: Final = self.auto_router_capability_limit - self.auto_router_capability_limit = None - try: - self.add_deployment(deployment=previous_deployment) - verbose_router_logger.info( - "Restored deployment %s (id=%s); it keeps serving its previous configuration.", - previous_deployment.model_name, - model_id, - ) - except Exception as restore_error: # noqa: BLE001 # best-effort restore: a second failure must not abort the reload - verbose_router_logger.warning( - "Could not restore previously served deployment %s (id=%s) after the failed upsert: %s", - previous_deployment.model_name, - model_id, - restore_error, - ) - finally: - self.auto_router_capability_limit = limit_resolver - - @staticmethod - def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: - """The ``litellm.model_cost`` keys a deployment's shared backend info is registered under.""" - backend_key: Final = model if custom_llm_provider is None else f"{custom_llm_provider}/{model}" - if "responses/" in backend_key: - return (backend_key, backend_key.replace("responses/", "")) - return (backend_key,) - - @staticmethod - def _catalog_entry_has_token_limit(catalog_entry: Mapping[str, object], field: str) -> bool: - return catalog_entry.get(field) is not None or ( - field == "max_tokens" and catalog_entry.get("max_output_tokens") is not None - ) - - @staticmethod - def _shared_backend_model_info_for_keys( - catalog: Mapping[str, Mapping[str, object]], - backend_keys: tuple[str, ...], - shared_model_info: Mapping[str, object], - ) -> Mapping[str, object]: - catalog_token_limit_fields: Final = frozenset( - field - for field in DEPLOYMENT_SCOPED_TOKEN_LIMIT_FIELDS - if any(PythonCore._catalog_entry_has_token_limit(catalog.get(key) or {}, field) for key in backend_keys) - ) - return {field: value for field, value in shared_model_info.items() if field not in catalog_token_limit_fields} - - @staticmethod - def _deployment_model_cost_payload(deployment: Deployment) -> dict[str, object]: # mutable-ok: cost-map entry - """The ``model_info`` a deployment contributes to ``litellm.model_cost``. - - Custom pricing lives on ``litellm_params`` rather than ``model_info``, and - the built-in cache-pricing inheritance is derived rather than stored, so - both are folded back in here. That keeps this reproducible from a - deployment alone, which is what lets a refresh rebuild the same entries. - """ - model_info: Final = dict( - _MODEL_INFO_ADAPTER.validate_python(deployment.model_info.model_dump(exclude_none=True)) - ) - for field in CustomPricingLiteLLMParams.model_fields: - field_value = deployment.litellm_params.get(field) - if field_value is not None: - model_info[field] = field_value - backend_model: Final = PythonCore._cost_map_backend_model(deployment) - pricing_provider: Final = ( - deployment.model_info.pricing_provider or deployment.litellm_params.custom_llm_provider - ) - params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) - if ( - params.get("model_id") - and backend_model != deployment.litellm_params.model - and ( - deployment.litellm_params.custom_llm_provider == "bedrock" - or deployment.litellm_params.model.startswith("bedrock/") - ) - ): - profile_info: Final = lookup_model_info_for_registration(model=backend_model, custom_llm_provider="bedrock") - if not isinstance(profile_info, ModelInfoLookupMiss): - explicit_info: Final = _MODEL_INFO_ADAPTER.validate_python(model_info) - priced_info: Final = model_info_with_pricing_overrides(profile_info, explicit_info) - model_info.update( - {key: value for key, value in priced_info.items() if is_server_derived_pricing_key(key)} - ) - PythonCore._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=pricing_provider, - ) - if model_info.get("input_cost_per_token") is not None: - PythonCore._inherit_builtin_cache_pricing( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=pricing_provider, - ) - PythonCore._inherit_builtin_service_tier_pricing( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=pricing_provider, - ) - PythonCore._inherit_builtin_tiered_output_rate( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=pricing_provider, - ) - return model_info - - @staticmethod - def _register_deployment_pricing(deployment: Deployment) -> None: - """Register a deployment's custom/inherited pricing in ``litellm.model_cost``. - - Takes only a ``Deployment``, so it registers pricing for a deployment that - is never added to ``self._core_model_list`` (a per-request client-side-credential - deployment) just as readily as one that is. - """ - PythonCore._register_deployment_in_model_cost( - model_id=deployment.model_info.id, - model_info=PythonCore._deployment_model_cost_payload(deployment), - model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - pricing_provider=deployment.model_info.pricing_provider, - ) - - @staticmethod - def _register_deployment_in_model_cost( - *, - model_id: str | None, - model_info: dict, # mutable-ok: cost-map entry - model: str, - custom_llm_provider: str | None, - pricing_provider: str | None = None, - ) -> None: - """Write a deployment's metadata into ``litellm.model_cost``. - - Runs when a deployment is added and again after a price data reload, so - the entries a refresh rebuilds are the ones a fresh boot would produce. - Nothing is recorded for replay: a refresh walks the live routers instead, - so a deleted, repointed or never-added deployment, and a discarded router, - drop out of the rebuild on their own. - - A strategy-router alias is never the deployment actually called or - billed, so custom pricing configured on it must not become a cost-map - price: an explicit zero would let ``_is_cost_explicitly_configured`` - treat the alias as a genuinely free model and waive budget checks for - requests that route to (and bill as) a real deployment. - """ - if classify_strategy_router_model(model) is not None: - model_info = {k: v for k, v in model_info.items() if k not in CustomPricingLiteLLMParams.model_fields} - - if model_id is not None: - python_host.invoke("register_model")( - model_cost={model_id: model_info}, - persist_across_reloads=False, - warning_display_name=model, - custom_llm_provider=pricing_provider or custom_llm_provider, - add_to_provider_model_sets=False, - ) - - ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes - backend_keys: Final = PythonCore._backend_cost_map_keys(model=model, custom_llm_provider=custom_llm_provider) - model_provider_prefix: Final = model.partition("/")[0] - catalog_provider: Final = ( - custom_llm_provider - if custom_llm_provider is not None - else model_provider_prefix - if model_provider_prefix in LlmProvidersSet - else None - ) - potential_names: Final = ( - get_potential_model_names(model, catalog_provider) if catalog_provider is not None else None - ) - catalog_backend_keys: Final = ( - ( - *backend_keys, - model, - potential_names["split_model"], - potential_names["stripped_model_name"], - potential_names["combined_model_name"], - potential_names["combined_stripped_model_name"], - potential_names["provider_prefixed_model_name"], - potential_names["region_free_combined_model_name"], - ) - if potential_names is not None - else (*backend_keys, model) - ) - - # For the shared backend key, keep only cost-map schema fields - # (minus custom pricing) so that one deployment's pricing overrides - # or custom metadata (id, access_via_team_ids, arbitrary keys) - # don't pollute another deployment sharing the same backend model - # name. Each deployment's full model_info is already stored under - # its unique model_id above. - catalog: Final = catalog_with_runtime_registrations() - shared_model_info: Final = PythonCore._shared_backend_model_info_for_keys( - catalog=catalog, - backend_keys=catalog_backend_keys, - shared_model_info=shared_backend_model_info(model_info), - ) - python_host.invoke("register_model")( - model_cost={_key: keep_registered_mode(_key, shared_model_info) for _key in backend_keys}, - persist_across_reloads=False, - ) - - def replay_model_cost_registrations(self) -> None: - """Re-assert this router's deployments onto a freshly fetched catalog. - - Reads ``model_list`` at call time, so only deployments the router still - serves are restored, plus any config deployment the fresh catalog now resolves. - """ - provider_unresolved: Final = self._provider_unresolved_deployments - self._provider_unresolved_deployments = () - for create_deployment in provider_unresolved: - create_deployment() - for entry in tuple(self._core_model_list): - try: - deployment = entry if isinstance(entry, Deployment) else Deployment(**entry) - except Exception: # noqa: BLE001 # a malformed entry must not abort the rest of the rebuild - verbose_router_logger.exception( - "Router: could not rebuild cost-map entry for a deployment during a price data reload" - ) - continue - PythonCore._register_deployment_in_model_cost( - model_id=deployment.model_info.id, - model_info=PythonCore._deployment_model_cost_payload(deployment), - model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - pricing_provider=deployment.model_info.pricing_provider, - ) - self._invalidate_model_group_info_cache() - - _replay_model_cost_registrations = replay_model_cost_registrations - - def _core_delete_deployment(self, id: str) -> Deployment | None: - """ - Parameters: - - id: str - the id of the deployment to be deleted - - Returns: - - The deleted deployment - - OR None (if deleted deployment not found) - """ - try: - with self._deployment_index_lock: - deployment_idx: Final = self.model_id_to_deployment_index_map.get(id) - if deployment_idx is None: - return None - # Pop the item from the list first - item: Final = self._core_model_list.pop(deployment_idx) - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - self._update_deployment_indices_after_removal(model_id=id, removal_idx=deployment_idx) - _budget_limiter: Final = self._core_get_router_deployment_budget_limiter() - if _budget_limiter is not None: - _budget_limiter.unregister_deployment_budget(model_id=id) - try: - self._unregister_pre_routing_strategy_for_deployment( - deployment=item if isinstance(item, Deployment) else Deployment(**item) - ) - except Exception: - verbose_router_logger.exception( - "delete_deployment: could not release pre-routing strategies for model_id=%s; " - "the deployment is out of the model_list and its indices are repaired", - id, - ) - except Exception: - return None - return item - - def _core_get_router_deployment_budget_limiter( - self, - ) -> RouterBudgetLimiting | None: - """ - Return the router's deployment-budget callback. - - Uses exact-type matching so proxy subclasses (e.g. virtual-key model budgets) - registered on litellm.callbacks are not mistaken for router deployment budgets. - """ - if self.router_budget_logger is not None: - return self.router_budget_logger - - if self.optional_callbacks: - for _cb in self.optional_callbacks: - if isinstance(_cb, RouterBudgetLimiting) and type(_cb) is RouterBudgetLimiting: - self.router_budget_logger = _cb - return _cb - return None - - def _deployment_has_budget_limits(self, deployment: Deployment) -> bool: - return ( - deployment.litellm_params.get("max_budget") is not None - and deployment.litellm_params.get("budget_duration") is not None - and deployment.model_info.id is not None - ) - - def _sync_deployment_budget_config(self, deployment: Deployment) -> None: - model_id: Final = deployment.model_info.id - if model_id is None: - return - - _budget_limiter = self._core_get_router_deployment_budget_limiter() - - if not self._deployment_has_budget_limits(deployment=deployment): - if _budget_limiter is not None: - _budget_limiter.unregister_deployment_budget(model_id=model_id) - return - - if _budget_limiter is None: - self.add_optional_pre_call_checks(optional_pre_call_checks=["router_budget_limiting"]) - _budget_limiter = self._core_get_router_deployment_budget_limiter() - - if _budget_limiter is not None: - _budget_limiter.register_deployment_budget(deployment=deployment.to_json(exclude_none=True)) - - def _core_deployments_named(self, model_name: str) -> tuple[DeploymentTypedDict, ...] | None: - """The deployments indexed under `model_name`, read under the index lock, or None when the name is not indexed.""" - with self._deployment_index_lock: - indices: Final = self._core_model_name_to_deployment_indices.get(model_name) - return None if indices is None else tuple(self._core_model_list[idx] for idx in indices) - - def _served_openai_models(self, model_name: str) -> tuple[str | None, ...]: - return served_openai_models(self._core_deployments_named(model_name) or ()) - - def team_public_model_names_of(self, team_id: str) -> frozenset[str]: - """The public names `team_id` serves its own deployments under.""" - with self._deployment_index_lock: - return frozenset(name for owner, name in self._core_team_model_to_deployment_indices if owner == team_id) - - def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[DeploymentTypedDict, ...] | None: - """The deployments `team_id` serves under `public_model_name`, or None when that pair is not indexed.""" - with self._deployment_index_lock: - indices: Final = self._core_team_model_to_deployment_indices.get((team_id, public_model_name)) - return None if indices is None else tuple(self._core_model_list[idx] for idx in indices) - - def _core_get_deployment(self, model_id: str) -> Deployment | None: - """ - Returns -> Deployment or None - - Raise Exception -> if model found in invalid format - """ - # Use O(1) lookup via model_id_to_deployment_index_map only - with self._deployment_index_lock: - idx: Final = self.model_id_to_deployment_index_map.get(model_id) - model: Final = None if idx is None else self._core_model_list[idx] - if model is not None: - if isinstance(model, dict): - return Deployment(**model) - elif isinstance(model, Deployment): - return model - else: - raise Exception(f"Model invalid format - {type(model)}") - - return None - - def _core_get_deployment_credentials(self, model_id: str) -> dict | None: - """ - Returns -> dict of credentials for a given model id. - - Returns None if the deployment is paused via `LiteLLM_ProxyModelTable.blocked`, - so file/batch/passthrough callers that resolve credentials directly cannot keep - using a paused deployment. - """ - deployment: Final = self._core_get_deployment(model_id=model_id) - if deployment is None or self._core_is_deployment_blocked(deployment): - return None - return CredentialLiteLLMParams.model_validate( - deployment.litellm_params.model_dump(exclude_none=True) - ).model_dump(exclude_none=True) - - def _core_get_deployment_by_model_group_name(self, model_group_name: str) -> Deployment | None: - """ - Returns -> Deployment or None - - Raise Exception -> if model found in invalid format - - Optimized with O(1) index lookup instead of O(n) linear scan. - """ - # O(1) lookup in model_name index - with self._deployment_index_lock: - indices: Final = self._core_model_name_to_deployment_indices.get(model_group_name) - # Return first deployment for this model_name - model: Final = self._core_model_list[indices[0]] if indices else None - if model is not None: - if isinstance(model, dict): - return Deployment(**model) - elif isinstance(model, Deployment): - return model - else: - raise Exception(f"Model Name invalid - {type(model)}") - return None - - @staticmethod - def _core_deployment_usable_by_team(model: Mapping[str, object] | Deployment, team_id: str | None) -> bool: - """ - A team-scoped deployment (``model_info.team_id`` set) is only usable by - callers from that same team; deployments without a team owner are shared. - """ - model_info: Final = model.get("model_info") if isinstance(model, dict) else model.model_info - owner_team_id: Final = cast(Mapping[str, object], model_info).get("team_id") if model_info is not None else None - return team_may_use_deployment(owner_team_id, team_id) - - def _get_model_group_deployment_usable_by_team( - self, model_group_name: str, team_id: str | None - ) -> Deployment | None: - """ - Like ``get_deployment_by_model_group_name``, but skips deployments owned - by other teams so a shared model name never resolves another team's - credentials. - """ - with self._deployment_index_lock: - indices: Final = self._core_model_name_to_deployment_indices.get(model_group_name) or () - usable: Final = ( - self._core_model_list[idx] - for idx in indices - if self._core_deployment_usable_by_team(self._core_model_list[idx], team_id) - ) - first_usable: Final = next(usable, None) - if first_usable is None: - return None - return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable - - async def _core_arefresh_model_info(self, *, client: AsyncHTTPHandler | None = None) -> None: - """Refresh token limits advertised by configured OpenAI-compatible deployments.""" - deployments: Final = iter(tuple(self._core_model_list)) - - async def refresh_worker() -> None: - for raw_deployment in deployments: - try: - await self._arefresh_deployment_model_info(raw_deployment, client=client) - except Exception: # noqa: BLE001 # one invalid deployment must not prevent refreshing the others - verbose_router_logger.debug("Could not refresh deployment model info") - - await asyncio.gather(*(refresh_worker() for _ in range(MODEL_INFO_REFRESH_CONCURRENCY))) - self._invalidate_model_group_info_cache() - - async def _arefresh_deployment_model_info( - self, raw_deployment: Mapping[str, object], *, client: AsyncHTTPHandler | None - ) -> None: - deployment: Final = Deployment.model_validate(raw_deployment) - params: Final = LiteLLM_Params.model_validate( - MappingProxyType( - { - **deployment.litellm_params.model_dump(exclude_none=True), - **( - self._core_get_deployment_credentials_with_provider(deployment.model_info.id or "") - or MappingProxyType({}) - ), - } - ) - ) - model, provider, dynamic_api_key, api_base = self._host.invoke("get_llm_provider")( - model=params.model, litellm_params=params - ) - if provider not in MODEL_INFO_DISCOVERY_PROVIDERS: - return - if api_base is None or "*" in model or params.get("use_clientside_credentials"): - return - api_key: Final = await aresolve_api_key_secret_reference( - params.api_key or dynamic_api_key, - model=model, - llm_provider=provider, - upstream=MappingProxyType({"api_base": params.api_base}), - ) - headers: Final = TypeAdapter(Mapping[str, str]).validate_python( - params.get("extra_headers") or params.get("headers") or MappingProxyType({}) - ) - auth_headers: Final = ( - MappingProxyType({"authorization": f"Bearer {api_key}"}) if api_key else MappingProxyType({}) - ) - limits: Final = await get_openai_compatible_model_info( - model=model, - api_base=api_base, - headers=MappingProxyType( - { - **auth_headers, - **MappingProxyType({key.lower(): value for key, value in headers.items()}), - } - ), - client=client or get_async_httpx_client(llm_provider=LlmProviders.OPENAI), - cache=self._core_cache.in_memory_cache, - ) - model_id: Final = deployment.model_info.id - if not limits or model_id is None or self._core_get_model_info(model_id) is not raw_deployment: - return - self._discovered_model_info_cache.max_size_in_memory = max(len(self._core_model_list), 1) - self._discovered_model_info_cache.delete_cache(model_id) - self._discovered_model_info_cache.set_cache( - model_id, DiscoveredDeploymentModelInfo(deployment=raw_deployment, limits=limits) - ) - self._invalidate_model_group_info_cache() - - def _core_get_discovered_model_info(self, model_id: str | None) -> Mapping[str, int]: - cached: Final[object] = self._discovered_model_info_cache.get_cache(model_id) - if ( - model_id is not None - and isinstance(cached, DiscoveredDeploymentModelInfo) - and cached.deployment is self._core_get_model_info(model_id) - ): - configured: Final = TypeAdapter(Mapping[str, object]).validate_python(cached.deployment["model_info"]) - return MappingProxyType({key: value for key, value in cached.limits.items() if configured.get(key) is None}) - return MappingProxyType({}) - - def _core_get_model_listing_info(self, model_name: str) -> DeploymentModelListingInfo | None: - """ - Return what the concrete deployments behind model_name contribute to its - /v1/models entry: the cost-map keys for their underlying models, plus the widest - configured or discovered token limits. Resolved via O(1) index - lookup. - - Returns None for wildcard-expanded or unknown names, where the listed name is the - real model name and no deployment-specific information exists, and treats a - malformed configured limit as absent rather than failing the listing. - - The whole group is read rather than just its first deployment, so a group that - mixes models does not advertise a window that depends on config order; the widest - one is reported, which is what get_model_group_info already shows the Admin UI. - Keys are deduplicated, so the ordinary group of interchangeable deployments of one - model still costs the caller a single cost-map lookup. Unlike get_model_group_info, - this never triggers pattern matching or deep copies, so it is safe to call per - listed model on the /v1/models hot path. - """ - deployments: Final = self._core_deployments_named(model_name) - if not deployments: - return None - - model_infos: Final = tuple( - MappingProxyType( - { - **self._core_get_discovered_model_info( - (deployment.get("model_info") or MappingProxyType({})).get("id") - ), - **MappingProxyType( - { - k: v - for k, v in (deployment.get("model_info") or MappingProxyType({})).items() - if v is not None - } - ), - } - ) - for deployment in deployments - ) - raw_deployments: Final[tuple[object, ...]] = deployments - pricing_keys: Final = (deployment_pricing_models(deployment) for deployment in raw_deployments) - cost_map_keys: Final = tuple(dict.fromkeys(keys[0] for keys in pricing_keys if keys)) - return DeploymentModelListingInfo( - cost_map_keys=cost_map_keys, - max_input_tokens=self._widest_configured_limit(model_infos, "max_input_tokens"), - max_output_tokens=self._widest_configured_limit(model_infos, "max_output_tokens"), - ) - - @staticmethod - def _widest_configured_limit(model_infos: Sequence[Mapping[str, object]], field: str) -> int | None: - """The largest usable value of ``field`` across a group's configured model_info blocks.""" - limits: Final = tuple( - limit - for limit in (coerce_token_limit(model_info.get(field)) for model_info in model_infos) - if limit is not None - ) - return max(limits) if limits else None - - def _core_get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]": - """ - Return (max_input_tokens, max_output_tokens) configured or discovered for a concrete - deployment of model_name, via O(1) index lookup. - - Returns (None, None) for wildcard-expanded or unknown names, and treats a - malformed configured value as absent rather than failing the caller. - - Deliberately reads one deployment rather than aggregating the group the way - get_model_listing_info does: its caller truncates an embedding input to this - value, so the widest window in a mixed group would be the wrong answer there. - """ - deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name) - if deployment is None: - return (None, None) - - model_info: Final = MappingProxyType( - { - **self._core_get_discovered_model_info(deployment.model_info.id), - **deployment.model_info.model_dump(exclude_none=True), - } - ) - return ( - coerce_token_limit(model_info.get("max_input_tokens")), - coerce_token_limit(model_info.get("max_output_tokens")), - ) - - def _core_get_configured_mode(self, model_name: str) -> "str | None": - """Return the mode explicitly configured for a concrete deployment.""" - deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name) - if deployment is None: - return None - - mode: Final = deployment.model_info.get("mode") - if isinstance(mode, str) and mode.strip(): - return mode - return None - - def _core_get_configured_display_name(self, model_name: str) -> "str | None": - """ - Return the display_name explicitly configured in a concrete deployment's - model_info for model_name, via O(1) index lookup. - - Returns None for wildcard-expanded or unknown names, and treats a - non-string or empty configured value as absent rather than failing the - listing. Like get_configured_token_limits, this never triggers pattern - matching or deep copies, so it is safe to call per listed model on the - /v1/models hot path. - """ - deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name) - if deployment is None: - return None - - display_name: Final = deployment.model_info.get("display_name") - if isinstance(display_name, str) and display_name.strip(): - return display_name - return None - - def _core_routable_model_group(self, model_name: str) -> str: - """ - The model group a request to model_name routes to: its target when - model_name is a `model_group_alias`, else model_name itself. - """ - target: Final = self._core_get_model_from_alias(model_name) - return target if target is not None else model_name - - def _routable_deployments(self, model_name: str, team_id: str | None) -> tuple[DeploymentTypedDict, ...]: - """ - The deployments a request from team_id to model_name can route to, in - model_list order and via the same O(1) index lookup as - get_configured_display_name, selected the way routing selects them: a - `model_group_alias` reads its target's deployments, another team's - deployment of the name is left out, a caller with no team reads the - deployments no team owns (every deployment of the name when a team owns - each one, which is what an admin's request routes to), and one an admin - paused via `LiteLLM_ProxyModelTable.blocked` is left out. - - Returns an empty tuple for wildcard-expanded or unknown names. - """ - named: Final = self._core_get_all_deployments( - model_name=self._core_routable_model_group(model_name), team_id=team_id - ) - usable: Final = tuple( - deployment for deployment in named if self._core_deployment_usable_by_team(deployment, team_id) - ) - return tuple( - deployment - for deployment in (usable or named) - if _configured_model_info(deployment).get("blocked") is not True - ) - - def _core_get_configured_service_tiers(self, model_name: str, team_id: str | None = None) -> tuple[object, ...]: - """ - Return the service_tiers value each deployment a request from team_id - to model_name can route to (see _routable_deployments) configures in - its model_info, unvalidated and in model_list order, None for a - deployment that sets none; the caller validates the shape it expects - and decides how the deployments combine. - - Returns an empty tuple for wildcard-expanded or unknown names. - """ - return tuple( - _configured_model_info(deployment).get("service_tiers") - for deployment in self._routable_deployments(model_name, team_id) - ) - - def _core_get_configured_max_input_tokens( - self, model_name: str, team_id: str | None = None - ) -> tuple[int | None, ...]: - """ - Return the max_input_tokens each deployment a request from team_id to - model_name can route to (see _routable_deployments) sets in its own - model_info, in model_list order, None for a deployment that sets none - or a malformed one. Unlike get_model_listing_info, neither the cost map - nor a discovered limit is read, so only an operator's own value counts, - and a value a saved /model/info response copied from the cost map is - skipped too. - - Returns an empty tuple for wildcard-expanded or unknown names. - """ - return tuple( - _operator_max_input_tokens(deployment) for deployment in self._routable_deployments(model_name, team_id) - ) - - def _core_get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: - """ - Whether each deployment a request from team_id to model_name can route - to (see _routable_deployments) answers Codex's client `tool_search`: - its model_info.supports_search_tool when that is a bool, else True when - its route builds an Anthropic Messages request, else None (unknown). - """ - return tuple(_search_tool_support(deployment) for deployment in self._routable_deployments(model_name, team_id)) - - def _core_get_routable_upstream_model(self, model_name: str, team_id: str | None = None) -> str | None: - """ - Return the `litellm_params.model` of the first deployment a request - from team_id to model_name can route to (see _routable_deployments). - - Returns None for wildcard-expanded or unknown names, and when the team - can route to no deployment of the name. - """ - deployment: Final = next(iter(self._routable_deployments(model_name, team_id)), None) - return deployment["litellm_params"].get("model") if deployment is not None else None - - def _core_get_credential_deployment(self, model_id: str, team_id: str | None = None) -> Deployment | None: - """ - The deployment a passthrough endpoint (files, batches, etc.) resolves for a - model id or model name: by deployment id first, then by model_name, then by - the team's exact public model name, then by wildcard pattern (team wildcards - before global ones, so a global "openai/*" never shadows the team's own - entry). Name and wildcard lookups never resolve another team's deployment. - - Returns None when nothing matches or the match is paused via - `LiteLLM_ProxyModelTable.blocked`, so callers cannot bypass an admin pause - by resolving the deployment directly. - """ - deployment: Final = ( - self._core_get_deployment(model_id=model_id) - or self._get_model_group_deployment_usable_by_team(model_group_name=model_id, team_id=team_id) - or self._get_team_public_name_deployment(model_id=model_id, team_id=team_id) - or self._get_wildcard_deployment_usable_by_team(model_id=model_id, team_id=team_id) - ) - if deployment is None or self._core_is_deployment_blocked(deployment): - return None - return deployment - - def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None: - if team_id is None: - return None - with self._deployment_index_lock: - team_indices: Final = self._core_team_model_to_deployment_indices.get((team_id, model_id)) - team_model: Final = self._core_model_list[team_indices[0]] if team_indices else None - if team_model is None: - return None - return Deployment(**team_model) if isinstance(team_model, dict) else team_model - - def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None: - team_pattern_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - team_wildcard_models: Final = team_pattern_router.route(model_id) if team_pattern_router else None - global_wildcard_models: Final = tuple( - wildcard_model - for wildcard_model in (self._core_pattern_router.route(model_id) or ()) - if self._core_deployment_usable_by_team(wildcard_model, team_id) - ) - potential_wildcard_models: Final = team_wildcard_models or global_wildcard_models - if not potential_wildcard_models: - return None - wildcard_deployment: Final = potential_wildcard_models[0] - if isinstance(wildcard_deployment, dict): - return Deployment(**wildcard_deployment) - if isinstance(wildcard_deployment, Deployment): - return wildcard_deployment - return None - - def _core_get_deployment_credentials_with_provider( - self, model_id: str, team_id: str | None = None - ) -> dict[str, Any] | None: - """ - Get API credentials and provider info from a model name in model_list. - Useful for passthrough endpoints (files, batches, etc.) that need credentials. - - Resolves the deployment with `get_credential_deployment` (by deployment id, - then model_name, team public model name, and wildcard pattern). - - Args: - model_id: Model ID or model name from model_list (e.g., "gpt-4o-litellm") - team_id: Optional team id of the caller. When set, team-scoped - deployments (indexed by team public model name, including team - wildcard models like "openai/*") are also considered. Name and - wildcard lookups never resolve a deployment owned by a - different team, so shared model names can't leak another - team's credentials. - - Returns: - Dictionary containing api_key, api_base, custom_llm_provider, etc. - Returns None if model not found, or if the resolved deployment is - paused via `LiteLLM_ProxyModelTable.blocked` (so passthrough callers - cannot bypass an admin pause by resolving credentials directly). - - Example: - credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm") - # Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", "model": "gpt-4o", ...} - """ - deployment: Final = self._core_get_credential_deployment(model_id=model_id, team_id=team_id) - if deployment is None: - return None - - # Get basic credentials - credentials: Final = CredentialLiteLLMParams.model_validate( - deployment.litellm_params.model_dump(exclude_none=True) - ).model_dump(exclude_none=True) - - # Resolve litellm_credential_name to actual credentials - if deployment.litellm_params.litellm_credential_name is not None: - credential_values: Final = CredentialAccessor.get_credential_values( - deployment.litellm_params.litellm_credential_name - ) - if not credential_values: - verbose_router_logger.warning( - "Credential '%s' not found in credential_list", deployment.litellm_params.litellm_credential_name - ) - credentials.update(credential_values) - # Remove the credential name since we've resolved it - credentials.pop("litellm_credential_name", None) - - credentials["model"] = deployment.litellm_params.model - - # Add custom_llm_provider - if deployment.litellm_params.custom_llm_provider: - credentials["custom_llm_provider"] = deployment.litellm_params.custom_llm_provider - elif "/" in deployment.litellm_params.model: - # Extract provider from "provider/model" format - credentials["custom_llm_provider"] = deployment.litellm_params.model.split("/")[0] - else: - credentials["custom_llm_provider"] = "openai" # default - - return credentials - - @overload - def _core_get_router_model_info( - self, - deployment: Union[dict, "Deployment"], - received_model_name: str, - id: None = None, - ) -> ModelMapInfo: - pass - - @overload - def _core_get_router_model_info(self, deployment: None, received_model_name: str, id: str) -> ModelMapInfo: - pass - - def _core_get_router_model_info( - self, - deployment: Union[dict, "Deployment"] | None, - received_model_name: str, - id: str | None = None, - ) -> ModelMapInfo: - """ - For a given model id, return the model info (max tokens, input cost, output cost, etc.). - - Augment litellm info with additional params set in `model_info`. - - For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set. - For every other provider, a base_model naming the same model as `model:` can only widen its context window. - - Returns - - ModelInfo - If found -> typed dict with max tokens, input cost, etc. - - Raises: - - ValueError -> If model is not mapped yet - """ - if id is not None: - _deployment: Final = self._core_get_deployment(model_id=id) - if _deployment is not None: - deployment = _deployment - - if deployment is None: - raise ValueError("Deployment not found") - - ## GET BASE MODEL - base_model = (deployment.get("model_info") or {}).get("base_model", None) - if base_model is None: - base_model = (deployment.get("litellm_params") or {}).get("base_model", None) - - model = base_model - - ## GET PROVIDER - reuse LiteLLM_Params if already constructed - litellm_params_data: Final = deployment.get("litellm_params") - litellm_params: LiteLLM_Params - if isinstance(litellm_params_data, LiteLLM_Params): - litellm_params = litellm_params_data - elif isinstance(litellm_params_data, dict) and "model" in litellm_params_data: - litellm_params = LiteLLM_Params.model_validate(litellm_params_data) - else: - raise ValueError( - f"Deployment missing valid litellm_params. " - f"Got: {type(litellm_params_data).__name__}, " - f"deployment_id: {(deployment.get('model_info') or {}).get('id', 'unknown')}" - ) - _model, custom_llm_provider, _, _ = self._host.invoke("get_llm_provider")( - model=litellm_params.model, - litellm_params=litellm_params, - ) - - ## SET MODEL TO 'model=' - if base_model is None + not azure - if custom_llm_provider == "azure" and base_model is None: - # Router init auto-registers every deployment name into - # litellm.model_cost as a zeroed stub, so membership alone can't - # tell a resolvable name apart; require usable limits/costs. - _azure_fallback_key = _model if _model.startswith("azure/") else f"azure/{_model}" - _fallback_entry = self._host.dynamic().model_cost.get(_azure_fallback_key) - _fallback_resolves = isinstance(_fallback_entry, Mapping) and ( - (_fallback_entry.get("max_input_tokens") or 0) > 0 - or (_fallback_entry.get("input_cost_per_token") or 0) > 0 - ) - if _fallback_resolves: - verbose_router_logger.debug( - "Azure deployment '%s' has no base_model set; using '%s' from the model cost map for max tokens, cost tracking, etc.", - _model, - _azure_fallback_key, - ) - else: - verbose_router_logger.error( - "Could not identify azure model '%s'. 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", - _model, - ) - elif custom_llm_provider != "azure": - model = _model - - if "*" in model: # only call pattern_router for wildcard models - potential_models: Final = self._core_pattern_router.route(received_model_name) - if potential_models is not None: - for potential_model in potential_models: - try: - if (potential_model.get("model_info") or {}).get("id") == ( - deployment.get("model_info") or {} - ).get("id"): - model = (potential_model.get("litellm_params") or {}).get("model") - break - except Exception: - pass - - ## GET LITELLM MODEL INFO - raises exception, if model is not mapped - if model is None: - # Handle case where base_model is None (e.g., Azure models without base_model set) - # Use the original model from litellm_params - model = _model - - if not model.startswith(f"{custom_llm_provider}/"): - model_info_name = f"{custom_llm_provider}/{model}" - else: - model_info_name = model - - model_info: Final = self._host.invoke("get_model_info")(model=model_info_name) - if model_info is None: - return model_info - - ## CHECK USER SET MODEL INFO - raw_user_model_info: Final = deployment.get("model_info") - user_model_info: Final = _MODEL_INFO_ADAPTER.validate_python( - raw_user_model_info.model_dump(exclude_none=True) - if isinstance(raw_user_model_info, BaseModel) - else raw_user_model_info or {} - ) - - # get_model_info() hands back an lru_cache'd dict, so merge into a copy; unset - # values are skipped or Deployment's None pricing defaults would erase the map's - merged_model_info: Final[ModelMapInfo] = { - **copy.deepcopy(model_info), - **_base_model_widened_window( - model_info, - _declared_base_model(user_model_info) or _declared_base_model(litellm_params.model_extra or {}), - custom_llm_provider, - ), - **self._core_get_discovered_model_info((deployment.get("model_info") or {}).get("id")), - **MappingProxyType({key: value for key, value in user_model_info.items() if value is not None}), - } - - return merged_model_info - - def _core_get_model_info(self, id: str) -> dict | None: - """ - For a given model id, return the model info - - Returns - - dict: the model in list with 'model_name', 'litellm_params', Optional['model_info'] - - None: could not find deployment in list - - Optimized with O(1) index lookup instead of O(n) linear scan. - """ - # O(1) lookup via model_id_to_deployment_index_map - with self._deployment_index_lock: - idx: Final = self.model_id_to_deployment_index_map.get(id) - return None if idx is None else self._core_model_list[idx] - - def _core_get_model_group(self, id: str) -> list | None: - """ - Return list of all models in the same model group as that model id - """ - - model_info: Final = self._core_get_model_info(id=id) - if model_info is None: - return None - - model_name: Final = model_info["model_name"] - return self._core_get_model_list(model_name=model_name) - - def _core_get_deployment_model_info(self, model_id: str, model_name: str) -> ModelInfo | None: - """ - For a given model id, return the model info - - 1. Check if model_id is in model info - 2. If not, check if litellm model name is in model info - 3. If not, return None - """ - from litellm.utils import cost_map_omits_token_price, update_dictionary - - model_info: ModelInfo | None = None - custom_model_info: dict | None = None - litellm_model_name_model_info: ModelInfo | None = None - base_model_key: str | None = None - - try: - if model_id in self._host.dynamic().model_cost: - cost_entry: Final = self._host.dynamic().model_cost.get(model_id) or MappingProxyType({}) - if not isinstance(cost_entry, Mapping): - raise TypeError(f"model_cost[{model_id!r}] is not a mapping") - custom_model_info = { - **copy.deepcopy(cost_entry), - **self._core_get_discovered_model_info(model_id), - } - except Exception: - pass - - try: - litellm_model_name_model_info = self._host.invoke("get_model_info")(model=model_name) - except Exception: - pass - - ## check for base model - try: - if custom_model_info is not None: - base_model: Final = custom_model_info.get("base_model", None) - if base_model is not None: - ## update litellm model info with base model info - base_model_info: Final = copy.deepcopy(self._host.invoke("get_model_info")(model=base_model)) - if base_model_info is not None: - base_model_key = base_model_info.get("key") - # Base model provides defaults, custom model info overrides - custom_model_info = update_dictionary( - cast(dict, base_model_info), - custom_model_info, - ) - except Exception: - pass - - # Three mutually exclusive scenarios for the model's metadata: - if custom_model_info is not None and litellm_model_name_model_info is not None: - # (1) It has both custom model_info set and exists in the built-in map - # merge with custom overriding built-in - model_info = cast( - ModelInfo, - update_dictionary( - copy.deepcopy(cast(dict, litellm_model_name_model_info)), - custom_model_info, - ), - ) - elif litellm_model_name_model_info is not None: - # (2) Built-in only — no custom pricing to merge - model_info = copy.deepcopy(litellm_model_name_model_info) - elif custom_model_info is not None: - # (3) Custom only — model not in built-in cost map yet - # custom_model_info already includes base_model defaults at this point, if applicable - model_info = cast(ModelInfo, custom_model_info) - - if model_info is None: - return None - builtin_key: Final = ( - litellm_model_name_model_info.get("key") if litellm_model_name_model_info is not None else None - ) - if cost_map_omits_token_price(model_id, builtin_key, base_model_key): - return cast( # cast-ok: TypedDict spread with overridden keys loses its type - ModelInfo, {**model_info, "input_cost_per_token": None, "output_cost_per_token": None} - ) - return model_info - - @staticmethod - def _model_group_llm_provider(litellm_params: LiteLLM_Params) -> tuple[str, str]: - try: - return resolve_model_and_provider_for_metadata( - model=litellm_params.model, - custom_llm_provider=litellm_params.custom_llm_provider, - ) - except litellm_exceptions.BadRequestError as e: - verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e) - return "", "" - - def _set_model_group_info( - self, model_group: str, user_facing_model_group_name: str, team_id: str | None = None - ) -> ModelGroupInfo | None: - """ - For a given model group name, return the combined model info - - Returns: - - ModelGroupInfo if able to construct a model group - - None if error constructing model group info - """ - model_group_info: ModelGroupInfo | None = None - - total_tpm: int | None = None - total_rpm: int | None = None - total_itpm: int | None = None - total_otpm: int | None = None - configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None - reasoning_efforts_initialized = False - reasoning_efforts_unknown = False - serves_team_deployments: Final = team_id is not None and bool( - self._team_deployments_named(team_id, model_group) - ) - model_list: Final = self._core_get_model_list_of_routed_group(model_group, team_id) - model_group_description: Final[str | None] = next( - (d for d in (_configured_model_info(m).get("description") for m in model_list) if isinstance(d, str)), - None, - ) - for model in model_list: - is_match = False - if ( - serves_team_deployments - or "model_name" in model - and model["model_name"] == model_group - or "model_name" in model - and self._core_pattern_router.route(model_group) is not None - ): # exact match - is_match = True - - if not is_match: - continue - # model in model group found # - litellm_params = LiteLLM_Params(**model["litellm_params"]) - # get configurable clientside auth params - configurable_clientside_auth_params = litellm_params.configurable_clientside_auth_params - - # Cache nested dict access to avoid repeated temporary dict allocations - model_litellm_params = model.get("litellm_params", {}) - model_info_dict = model.get("model_info", {}) - - # get model tpm - _deployment_tpm: int | None = None - if _deployment_tpm is None: - _deployment_tpm = model.get("tpm", None) - if _deployment_tpm is None: - _deployment_tpm = model_litellm_params.get("tpm", None) - if _deployment_tpm is None: - _deployment_tpm = model_info_dict.get("tpm", None) - - # get model rpm - _deployment_rpm: int | None = None - if _deployment_rpm is None: - _deployment_rpm = model.get("rpm", None) - if _deployment_rpm is None: - _deployment_rpm = model_litellm_params.get("rpm", None) - if _deployment_rpm is None: - _deployment_rpm = model_info_dict.get("rpm", None) - - _deployment_itpm: int | None = model.get("itpm") - if _deployment_itpm is None: - _deployment_itpm = model_litellm_params.get("itpm", None) - if _deployment_itpm is None: - _deployment_itpm = model_info_dict.get("itpm", None) - - _deployment_otpm: int | None = model.get("otpm") - if _deployment_otpm is None: - _deployment_otpm = model_litellm_params.get("otpm", None) - if _deployment_otpm is None: - _deployment_otpm = model_info_dict.get("otpm", None) - - # get model info - try: - model_id = model_info_dict.get("id", None) - if model_id is not None: - model_info = self._core_get_deployment_model_info( - model_id=model_id, model_name=litellm_params.model - ) - else: - model_info = None - except Exception: - model_info = None - - deployment_is_mapped = deployment_is_catalog_mapped(model_info, model_info_dict) - - # get llm provider - litellm_model, llm_provider = self._model_group_llm_provider(litellm_params) - - if model_info is None: - supported_openai_params = self._host.invoke("get_supported_openai_params")( - model=litellm_model, custom_llm_provider=llm_provider - ) - if supported_openai_params is None: - supported_openai_params = [] - - # Get mode from database model_info if available, otherwise default to "chat" - db_model_info = model.get("model_info", {}) - mode = db_model_info.get("mode", "chat") - input_cost_per_token = _cost_value_as_float(db_model_info.get("input_cost_per_token")) - output_cost_per_token = _cost_value_as_float(db_model_info.get("output_cost_per_token")) - - model_info = ModelMapInfo( - key=model_group, - max_tokens=None, - max_input_tokens=None, - max_output_tokens=None, - input_cost_per_token=input_cost_per_token, - output_cost_per_token=output_cost_per_token, - litellm_provider=llm_provider, - mode=mode, - supported_openai_params=supported_openai_params, - supports_system_messages=None, - ) - - if model_group_info is None: - model_group_info = ModelGroupInfo( - **{ - "model_group": user_facing_model_group_name, - "providers": [llm_provider], - "description": model_group_description, - **model_info, - "supports_continuation": model_info.get("supports_continuation") is True, - "supports_fast_mode": True, - "supported_reasoning_efforts": None, - "supported_endpoints": None, - } - ) - else: - # if max_input_tokens > curr - # if max_output_tokens > curr - # if input_cost_per_token > curr - # if output_cost_per_token > curr - # supports_parallel_function_calling == True - # supports_vision == True - # supports_function_calling == True - if llm_provider not in model_group_info.providers: - model_group_info.providers.append(llm_provider) - if ( - model_info.get("max_input_tokens", None) is not None - and model_info["max_input_tokens"] is not None - and ( - model_group_info.max_input_tokens is None - or model_info["max_input_tokens"] > model_group_info.max_input_tokens - ) - ): - model_group_info.max_input_tokens = model_info["max_input_tokens"] - if ( - model_info.get("max_output_tokens", None) is not None - and model_info["max_output_tokens"] is not None - and ( - model_group_info.max_output_tokens is None - or model_info["max_output_tokens"] > model_group_info.max_output_tokens - ) - ): - model_group_info.max_output_tokens = model_info["max_output_tokens"] - _input_cost_per_token = _cost_value_as_float(model_info.get("input_cost_per_token")) - if _input_cost_per_token is not None and ( - model_group_info.input_cost_per_token is None - or _input_cost_per_token > (model_group_info.input_cost_per_token or 0.0) - ): - model_group_info.input_cost_per_token = _input_cost_per_token - _output_cost_per_token = _cost_value_as_float(model_info.get("output_cost_per_token")) - if _output_cost_per_token is not None and ( - model_group_info.output_cost_per_token is None - or _output_cost_per_token > (model_group_info.output_cost_per_token or 0.0) - ): - model_group_info.output_cost_per_token = _output_cost_per_token - if ( - model_info.get("supports_parallel_function_calling", None) is not None - and model_info["supports_parallel_function_calling"] is True - ): - model_group_info.supports_parallel_function_calling = True - if model_info.get("supports_vision", None) is not None and model_info["supports_vision"] is True: - model_group_info.supports_vision = True - if ( - model_info.get("supports_function_calling", None) is not None - and model_info["supports_function_calling"] is True - ): - model_group_info.supports_function_calling = True - if ( - model_info.get("supports_web_search", None) is not None - and model_info["supports_web_search"] is True - ): - model_group_info.supports_web_search = True - if ( - model_info.get("supports_url_context", None) is not None - and model_info["supports_url_context"] is True - ): - model_group_info.supports_url_context = True - if model_info.get("supports_continuation") is True: - model_group_info.supports_continuation = True - - if model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True: - model_group_info.supports_reasoning = True - if ( - model_info.get("supported_openai_params", None) is not None - and model_info["supported_openai_params"] is not None - ): - model_group_info.supported_openai_params = model_info["supported_openai_params"] - if model_info.get("tpm", None) is not None and _deployment_tpm is None: - _deployment_tpm = model_info.get("tpm") - if model_info.get("rpm", None) is not None and _deployment_rpm is None: - _deployment_rpm = model_info.get("rpm") - - model_group_info.supported_endpoints = ( - tuple( - dict.fromkeys( - ( - *(model_group_info.supported_endpoints or ()), - *( - deployment_supported_endpoints( - model_info.get("supported_endpoints"), model, llm_provider - ) - if team_may_use_deployment(_configured_model_info(model).get("team_id"), team_id) - else () - ), - ) - ) - ) - or None - ) - model_group_info.supports_fast_mode = model_group_info.supports_fast_mode and ( - AnthropicModelInfo.supports_fast_mode(litellm_model, llm_provider) - ) - deployment_reasoning_efforts = resolve_supported_reasoning_efforts( - model_info, deployment_is_mapped=deployment_is_mapped - ) - if deployment_reasoning_efforts is None: - reasoning_efforts_unknown = True - model_group_info.supported_reasoning_efforts = None - elif not reasoning_efforts_initialized: - reasoning_efforts_initialized = True - if not reasoning_efforts_unknown: - model_group_info.supported_reasoning_efforts = deployment_reasoning_efforts - elif not reasoning_efforts_unknown: - model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts( - model_group_info.supported_reasoning_efforts, - deployment_reasoning_efforts, - ) - - if _deployment_tpm is not None: - if total_tpm is None: - total_tpm = 0 - total_tpm += _deployment_tpm - - if _deployment_rpm is not None: - if total_rpm is None: - total_rpm = 0 - total_rpm += _deployment_rpm - - if _deployment_itpm is not None: - if total_itpm is None: - total_itpm = 0 - total_itpm += _deployment_itpm - - if _deployment_otpm is not None: - if total_otpm is None: - total_otpm = 0 - total_otpm += _deployment_otpm - if model_group_info is not None: - ## UPDATE WITH TOTAL TPM/RPM FOR MODEL GROUP - if total_tpm is not None: - model_group_info.tpm = total_tpm - - if total_rpm is not None: - model_group_info.rpm = total_rpm - - if total_itpm is not None: - model_group_info.itpm = total_itpm - - if total_otpm is not None: - model_group_info.otpm = total_otpm - - ## UPDATE WITH CONFIGURABLE CLIENTSIDE AUTH PARAMS FOR MODEL GROUP - if configurable_clientside_auth_params is not None: - model_group_info.configurable_clientside_auth_params = configurable_clientside_auth_params - - return model_group_info - - def _core_get_model_group_info( - self, model_group: str, *, include_hidden: bool = False, team_id: str | None = None - ) -> ModelGroupInfo | None: - """ - For a given model group name, return the combined model info - - Returns: - - ModelGroupInfo if able to construct a model group - - None if error constructing model group info or hidden model group (unless include_hidden) - """ - ## Check if model group alias - if model_group in self._core_model_group_alias: - item: Final = self._core_model_group_alias[model_group] - if isinstance(item, str): - _router_model_group = item - elif isinstance(item, dict): - if item["hidden"] is True and not include_hidden: - return None - else: - _router_model_group = item["model"] - else: - return None - - return self._set_model_group_info( - model_group=_router_model_group, - user_facing_model_group_name=model_group, - team_id=team_id, - ) - - ## Check if actual model - return self._set_model_group_info( - model_group=model_group, user_facing_model_group_name=model_group, team_id=team_id - ) - - async def _core_get_model_group_usage(self, model_group: str) -> tuple[int | None, int | None]: - """ - Returns current tpm/rpm usage for model group - - Parameters: - - model_group: str - the received model name from the user (can be a wildcard route). - - Returns: - - usage: Tuple[tpm, rpm] - """ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - tpm_keys: Final[list[str]] = [] - rpm_keys: Final[list[str]] = [] - - model_list: Final = self._core_get_model_list(model_name=model_group) - if model_list is None: # no matching deployments - return None, None - - for model in model_list: - id: str | None = model.get("model_info", {}).get("id") - litellm_model: str | None = model["litellm_params"].get( - "model" - ) # USE THE MODEL SENT TO litellm.completion() - consistent with how global_router cache is written. - if id is None or litellm_model is None: - continue - tpm_keys.append( - RouterCacheEnum.TPM.value.format( - id=id, - model=litellm_model, - current_minute=current_minute, - ) - ) - rpm_keys.append( - RouterCacheEnum.RPM.value.format( - id=id, - model=litellm_model, - current_minute=current_minute, - ) - ) - combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys - - combined_tpm_rpm_values: Final = await self._core_cache.async_batch_get_cache(keys=combined_tpm_rpm_keys) - if combined_tpm_rpm_values is None: - return None, None - - tpm_usage_list: Final[list | None] = combined_tpm_rpm_values[: len(tpm_keys)] - rpm_usage_list: Final[list | None] = combined_tpm_rpm_values[len(tpm_keys) :] - - ## TPM - tpm_usage: int | None = None - if tpm_usage_list is not None: - for t in tpm_usage_list: - if isinstance(t, int): - if tpm_usage is None: - tpm_usage = 0 - tpm_usage += t - ## RPM - rpm_usage: int | None = None - if rpm_usage_list is not None: - for t in rpm_usage_list: - if isinstance(t, int): - if rpm_usage is None: - rpm_usage = 0 - rpm_usage += t - return tpm_usage, rpm_usage - - async def get_model_group_io_token_usage(self, model_group: str) -> tuple[int | None, int | None]: - """ - Returns current ITPM/OTPM usage for a model group (sum across deployments). - """ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") - itpm_keys: Final[list[str]] = [] - otpm_keys: Final[list[str]] = [] - - model_list: Final = self._core_get_model_list(model_name=model_group) - if model_list is None: - return None, None - - for model in model_list: - model_id: str | None = model.get("model_info", {}).get("id") - litellm_model: str | None = model["litellm_params"].get("model") - if model_id is None or litellm_model is None: - continue - itpm_keys.append( - RouterCacheEnum.ITPM.value.format( - id=model_id, - model=litellm_model, - current_minute=current_minute, - ) - ) - otpm_keys.append( - RouterCacheEnum.OTPM.value.format( - id=model_id, - model=litellm_model, - current_minute=current_minute, - ) - ) - - combined_values: Final = await self._core_cache.async_batch_get_cache(keys=itpm_keys + otpm_keys) - if combined_values is None: - return None, None - - itpm_values: Final = combined_values[: len(itpm_keys)] - otpm_values: Final = combined_values[len(itpm_keys) :] - - total_itpm: int | None = None - for value in itpm_values: - if isinstance(value, int): - total_itpm = (total_itpm or 0) + value - - total_otpm: int | None = None - for value in otpm_values: - if isinstance(value, int): - total_otpm = (total_otpm or 0) + value - - return total_itpm, total_otpm - - @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) - def _cached_get_model_group_info(self, model_group: str) -> ModelGroupInfo | None: - """ - Cached version of get_model_group_info, uses @lru_cache wrapper - - This is a speed optimization, since set_response_headers makes a call to get_model_group_info on every request - """ - return self._core_get_model_group_info(model_group) - - def _core_cached_model_group_info(self, model_group: str) -> ModelGroupInfo | None: - return self._cached_get_model_group_info(model_group) - - async def _core_get_remaining_model_group_usage(self, model_group: str) -> dict[str, int]: - model_group_info: Final = self._cached_get_model_group_info(model_group) - - returned_dict: Final[dict[str, int]] = {} - - # ITPM/OTPM groups emit input/output token headers, but they may also set - # tpm/rpm, so build both sets rather than returning early - clients and - # prometheus gauges that read the standard headers still get data. - if model_group_info is not None and (model_group_info.itpm is not None or model_group_info.otpm is not None): - current_itpm, current_otpm = await self.get_model_group_io_token_usage(model_group) - returned_dict.update( - build_io_token_rate_limit_headers( - itpm_limit=model_group_info.itpm, - otpm_limit=model_group_info.otpm, - current_itpm=current_itpm, - current_otpm=current_otpm, - ) - ) - - tpm_limit: Final = model_group_info.tpm if model_group_info is not None else None - rpm_limit: Final = model_group_info.rpm if model_group_info is not None else None - - if tpm_limit is None and rpm_limit is None: - return returned_dict - - current_tpm, current_rpm = await self._core_get_model_group_usage(model_group) - - if tpm_limit is not None: - returned_dict["x-ratelimit-remaining-tokens"] = tpm_limit - (current_tpm or 0) - returned_dict["x-ratelimit-limit-tokens"] = tpm_limit - if rpm_limit is not None: - returned_dict["x-ratelimit-remaining-requests"] = rpm_limit - (current_rpm or 0) - returned_dict["x-ratelimit-limit-requests"] = rpm_limit - - return returned_dict - - async def set_response_headers( - self, - response: object, - model_group: str | None = None, - request_kwargs: dict[str, object] | None = None, - ) -> Any: - """ - Add the most accurate rate limit headers for a given model response. - - ## TODO: add model group rate limit headers - # - if healthy_deployments > 1, return model group rate limit headers - # - else return the model's rate limit headers - """ - response = prepare_response_for_header_attachment(response) - if response is None: - return response - - additional_headers: Final = ensure_response_additional_headers(response) - apply_response_model_id( - response, find_deployment_metadata(request_kwargs) if request_kwargs is not None else None - ) - additional_headers["x-litellm-model-group"] = model_group - apply_quality_router_decision_headers(additional_headers, request_kwargs) - additional_headers.update(complexity_router_decision_headers(request_kwargs)) - - if model_group is not None: - remaining_usage: Final = await self._core_get_remaining_model_group_usage(model_group) - apply_remaining_usage_headers(additional_headers, remaining_usage) - return response - - def _build_model_name_index(self, model_list: list) -> None: - """ - Build model_name -> deployment indices mapping for O(1) lookups. - - This index allows us to find all deployments for a given model_name in O(1) time - instead of O(n) linear scan through the entire model_list. - """ - self._core_model_name_to_deployment_indices.clear() - self._core_team_model_to_deployment_indices.clear() - self._core_team_public_model_names = frozenset() - - for idx, model in enumerate(model_list): - model_name = model.get("model_name") - if model_name: - if model_name not in self._core_model_name_to_deployment_indices: - self._core_model_name_to_deployment_indices[model_name] = [] - self._core_model_name_to_deployment_indices[model_name].append(idx) - - self._update_team_model_index(model, idx) - - def _build_model_id_to_deployment_index_map(self, model_list: list): - """ - Build model index from model list to enable O(1) lookups immediately. - This is called during initialization to avoid the race condition where - requests arrive before model_id_to_deployment_index_map is populated. - """ - # First populate the model_list - self._core_model_list = [] - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - for _, model in enumerate(model_list): - # Extract model_info from the model dict - model_info = model.get("model_info", {}) - model_id = model_info.get("id") - - # If no ID exists, generate one using the same logic as set_model_list - if model_id is None: - model_name = model.get("model_name", "") - litellm_params = model.get("litellm_params", {}) - model_id = self._core_generate_model_id(model_name, litellm_params) - # Update the model_info in the original list - if "model_info" not in model: - model["model_info"] = {} - model["model_info"]["id"] = model_id - - self._add_model_to_list_and_index_map(model=model, model_id=model_id) - - def _core_get_model_ids(self, model_name: str | None = None, exclude_team_models: bool = False) -> list[str]: - """ - if 'model_name' is none, returns all. - - Returns list of model id's. - - Optimized with O(1) or O(k) index lookup when model_name provided, - instead of O(n) linear scan. - """ - ids: Final = [] - - if model_name is not None: - # O(1) lookup in model_name index, then O(k) iteration where k = deployments for this model_name - for model in self._core_deployments_named(model_name) or (): - if "model_info" in model and "id" in model["model_info"]: - if exclude_team_models and model["model_info"].get("team_id"): - continue - ids.append(model["model_info"]["id"]) - else: - # When model_name is None, return all model IDs - # Use the index map keys for O(n) where n = total deployments - with self._deployment_index_lock: - id_indices: Final = tuple(self.model_id_to_deployment_index_map.items()) - models: Final = tuple(self._core_model_list) - for model_id, idx in id_indices: - model = models[idx] - if "model_info" in model and "id" in model["model_info"]: - if exclude_team_models and model["model_info"].get("team_id"): - continue - ids.append(model_id) - - return ids - - def _core_get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: - """ - Deployment ids that could serve ``model`` for ``team_id``, following the same - precedence ``_common_checks_available_deployment`` uses to build a candidate pool: - ``model_group_alias``, then a routing group, then the first matching early-resolve - path for a name that is not a ``model_name`` (team route, wildcard pattern via - ``get_deployments_by_pattern``, team pattern router, default deployment), then the - ``model_name`` and team indexes. Delegating to the router's own resolvers keeps this - aligned with how a route actually resolves rather than re-deriving it, and unlike - ``_common_checks_available_deployment`` it is read-only: it does not apply request - fallbacks and (with ``include_team_models`` left off) does not raise. Lets a pre-call - check tell a genuine cross-group route from same-group unavailability without leaking - deployment ids into request kwargs bound for the provider. - """ - resolved: Final = self._core_get_model_from_alias(model=model) or model - routing_group_members: Final = self._get_routing_group_deployments(model=resolved, team_id=team_id) - if routing_group_members is not None: - return self._deployment_ids(routing_group_members) - early: Final = self._try_early_resolve_deployments_for_model_not_in_names( - model=resolved, request_team_id=team_id - ) - if early is not None: - early_deployments: Final = early[1] - return self._deployment_ids( - (early_deployments,) if isinstance(early_deployments, Mapping) else early_deployments - ) - return self._deployment_ids(self._core_get_all_deployments(model_name=resolved, team_id=team_id)) - - @staticmethod - def _deployment_ids(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]: - return frozenset( - str(model_info["id"]) - for deployment in deployments - for model_info in (deployment.get("model_info"),) - if isinstance(model_info, Mapping) and model_info.get("id") is not None - ) - - def _region_routable_deployments( - self, - deployments: Sequence[DeploymentTypedDict], - request_kwargs: Mapping[str, object] | None, - has_request_input: bool, - ) -> tuple[Mapping[str, Any], ...]: - """ - `deployments` without the ones `_pre_call_checks` keeps this request away from for being outside its allowed - region, typed for the tag routing filter - """ - allowed_model_region: Final = None if request_kwargs is None else request_kwargs.get("allowed_model_region") - if allowed_model_region is None or not self._core_enable_pre_call_checks or not has_request_input: - return tuple(deployments) - if not isinstance(allowed_model_region, str): - return () - return tuple( - deployment - for deployment in deployments - if is_region_allowed( - litellm_params=LiteLLM_Params.model_validate(deployment["litellm_params"]), - allowed_model_region=allowed_model_region, - ) - ) - - async def _async_routable_cooled_deployment_ids( - self, - model: str, - request_deployments: Sequence[DeploymentTypedDict], - cooldown_deployments: Iterable[str], - request_kwargs: Mapping[str, object], - has_request_input: bool, - ) -> frozenset[str]: - """ - Ids of the cooling deployments in `request_deployments` that region and tag routing would send this request - to, each tag routing pick counted along with the picks it falls back to once that one is unavailable, judged - without recording a tag match on the request - """ - cooled_ids: Final = frozenset(cooldown_deployments) - metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_kwargs) - metadata: Final = request_kwargs.get(metadata_variable_name) - if metadata_variable_name in request_kwargs and not isinstance(metadata, Mapping): - return frozenset() - tag_request: Final[Mapping[str, Any]] = ( - {**request_kwargs, metadata_variable_name: {**metadata}} - if isinstance(metadata, Mapping) - else request_kwargs - ) - routable_ids = frozenset[str]() # rebind-ok: gains each pass's tag routing picks - remaining = tuple( # rebind-ok: each pass drops what tag routing just picked - deployment - for deployment in self._region_routable_deployments(request_deployments, request_kwargs, has_request_input) - if self._deployment_ids((deployment,)) <= cooled_ids - ) - while remaining: - try: - picked = await get_deployments_for_tag( - ctx=self._core_ctx, - model=model, - healthy_deployments=remaining, - request_kwargs=tag_request, - metadata_variable_name=metadata_variable_name, - request_deployments=request_deployments, - ) - except ValueError: # tag routing sends this request to none of the rest - break - picked_identities = frozenset(id(deployment) for deployment in picked) - rest = tuple(deployment for deployment in remaining if id(deployment) not in picked_identities) - if len(rest) == len(remaining): - break - routable_ids |= self._deployment_ids( - tuple(deployment for deployment in remaining if id(deployment) in picked_identities) - ) - remaining = rest - return routable_ids - - def _core_has_model_id(self, candidate_id: str) -> bool: - """ - O(1) membership check for a deployment ID without allocating large lists. - - Note: Call sites may pass a variable named `model` when it actually - contains a deployment ID. This helper expects the deployment ID string. - - Uses the existing `model_id_to_deployment_index_map` which is kept - in sync by `_build_model_id_to_deployment_index_map` and model-list - mutation helpers. - """ - return candidate_id in self.model_id_to_deployment_index_map - - def _core_resolve_model_name_from_model_id( - self, model_id: str | None, custom_llm_provider: str | None = None - ) -> str | None: - """ - Resolve model_name from model_id. - - This method attempts to find the correct model_name to use with the router - so that litellm_params can be automatically injected from the model config. - - Strategy: - 1. First, check if model_id directly matches a model_name or deployment ID - 2. If not, search through router's model_list to find a match by litellm_params.model - 3. If not, keep model_id when the router routes that name itself (an alias, a routing group, or a wildcard - whose deployment serves `custom_llm_provider`), else name it `custom_llm_provider/model_id` when only that - form routes to such a wildcard - 4. Return the model_name if found, None otherwise - - Args: - model_id: The model_id extracted from decoded video_id - (could be model_name or litellm_params.model value) - custom_llm_provider: The provider extracted from the same id - - Returns: - model_name if found, None otherwise. If None, the request will fall through - to normal flow using environment variables. - """ - if not model_id: - return None - - # Strategy 1: Check if model_id directly matches a model_name or deployment ID - if model_id in self._core_model_names: - return model_id - if self._core_has_model_id(model_id): - deployment = self._core_get_deployment(model_id=model_id) - if deployment is not None and deployment.model_name: - return deployment.model_name - return model_id - - # Strategy 2: Search through router's model_list to find by litellm_params.model - all_models: Final = self._core_get_model_list(model_name=None) - if not all_models: - return None - - for deployment in all_models: - litellm_params = deployment.get("litellm_params", {}) - actual_model = litellm_params.get("model") - - # Match by exact match or by checking if actual_model ends with /model_id or :model_id - # e.g., model_id="veo-2.0-generate-001" matches actual_model="vertex_ai/veo-2.0-generate-001" - matches = ( - actual_model == model_id - or (actual_model and actual_model.endswith(f"/{model_id}")) - or (actual_model and actual_model.endswith(f":{model_id}")) - ) - - if matches: - model_name = deployment.get("model_name") - if model_name: - return model_name - - if self._core_is_recognized_model(model_id): - return model_id - if custom_llm_provider is None: - return None - if self._core_pattern_router.routes_to_provider(model_id, custom_llm_provider): - return model_id - provider_prefixed: Final = f"{custom_llm_provider}/{model_id}" - if not model_id.startswith(f"{custom_llm_provider}/") and self._core_pattern_router.routes_to_provider( - provider_prefixed, custom_llm_provider - ): - return provider_prefixed - return None - - def _core_resolve_served_model_name(self, model_name: str, served_model: str) -> str: - """ - `served_model` when `model_name` routes to it (the alias target, the routing group member, or the name - itself), so a request can return to the group that served an earlier one, else `model_name`. - """ - target: Final = self._core_get_model_from_alias(model_name) or model_name - routing_group: Final = self.get_routing_group(target) - members: Final = routing_group.models if routing_group is not None else () - return served_model if served_model == target or served_model in members else model_name - - def _core_map_team_model( - self, team_model_name: str | None, team_id: str, custom_llm_provider: str | None = None - ) -> str | None: - """ - Resolve explicit team models or select a deployment for model-less calls. - - Model-less calls with no provider use the first team deployment. With a provider, they - use the first team deployment of that provider, pinned by ID when its public name also - covers another provider. - """ - models: Final = self._core_get_model_list(model_name=team_model_name, team_id=team_id) - if not models: - return None - if team_model_name is None and custom_llm_provider is not None: - - def team_deployment_public_name(deployment: DeploymentTypedDict) -> str | None: - model_info: Final = _configured_model_info(deployment) - if model_info.get("team_id") != team_id: - return None - team_public_model_name: Final = model_info.get("team_public_model_name") - return ( - team_public_model_name - if isinstance(team_public_model_name, str) and team_public_model_name - else deployment["model_name"] - ) - - matching_deployment: Final = next( - ( - model - for model in models - if team_deployment_public_name(model) is not None - and provider_for_generic_call(model["litellm_params"]) == custom_llm_provider - ), - None, - ) - if matching_deployment is None: - return None - public_name: Final = team_deployment_public_name(matching_deployment) - if public_name is None: - return None - same_public_name_deployments: Final = tuple( - candidate for candidate in models if team_deployment_public_name(candidate) == public_name - ) - if all( - provider_for_generic_call(candidate["litellm_params"]) == custom_llm_provider - for candidate in same_public_name_deployments - ): - return public_name - deployment_id: Final = _configured_model_info(matching_deployment).get("id") - return deployment_id if isinstance(deployment_id, str) else None - for model in models: - if model.get("model_info", {}).get("team_id") == team_id: - if team_model_name is None: - # No model was specified (e.g. vector store endpoints). - # Return the deployment's public model name so the router - # can route to it and inject the BYOK API key. - return model.get("model_info", {}).get("team_public_model_name") or model.get("model_name") - return team_model_name - - # No team-scoped deployment found; wildcard/pattern routes are - # handled downstream by the pattern_router in _common_checks_available_deployment. - return None - - def _core_should_include_deployment( - self, model_name: str, model: DeploymentTypedDict | dict, team_id: str | None = None - ) -> bool: - """ - Get the team-specific model name if team_id matches the deployment. - """ - if ( - team_id is not None - and (model.get("model_info") or {}).get("team_id") == team_id - and model_name == (model.get("model_info") or {}).get("team_public_model_name") - ): - return True - elif model_name is not None and model["model_name"] == model_name: - # Fallback: check by internal model_name for non-team deployments - # or deployments that haven't been migrated to team_public_model_name yet - model_team_id: Final = (model.get("model_info") or {}).get("team_id") - if ( - team_id is None # requester has no team constraint - or model_team_id is None # global deployment - accessible to all teams - or model_team_id == team_id # deployment belongs to requester's team - ): - return True - # No match: deployment is for a different team or doesn't match the requested model - return False - - def _core_get_all_deployments( - self, - model_name: str, - model_alias: str | None = None, - team_id: str | None = None, - ) -> list[DeploymentTypedDict]: - """ - Return all deployments of a model name - - Used for accurate 'get_model_list'. - - if team_id specified, only return team-specific models - - Optimized with O(1) index lookup instead of O(n) linear scan. - - Note: when team_id is provided, O(1) lookup in - `team_model_to_deployment_indices` only applies when `model_name` is the - team public model name. If a caller passes an internal deployment model - name (for example, `model_name__`), this method falls back - to the standard model-name index / scan path. - """ - returned_models: Final[list[DeploymentTypedDict]] = [] - - # O(1) lookup in team_model index when team_id is provided - if team_id is not None: - team_models: Final = self._team_deployments_named(team_id, model_name) - if team_models is not None: - # O(k) where k = team deployments for this model_name (typically 1-10) - for model in team_models: - if not self._core_should_include_deployment(model_name=model_name, model=model, team_id=team_id): - continue - if model_alias is not None: - alias_model = model.copy() - alias_model["model_name"] = model_alias - returned_models.append(alias_model) - else: - returned_models.append(model) - if returned_models: - return returned_models - - # O(1) lookup in model_name index - named_models: Final = self._core_deployments_named(model_name) - if named_models is not None: - # O(k) where k = deployments for this model_name (typically 1-10) - for model in named_models: - if self._core_should_include_deployment(model_name=model_name, model=model, team_id=team_id): - if model_alias is not None: - # Optimized: Use shallow copy since we only modify top-level model_name - # This is much faster than deepcopy for nested dict structures - alias_model = model.copy() - alias_model["model_name"] = model_alias - returned_models.append(alias_model) - else: - returned_models.append(model) - elif team_id is not None: - # Fallback: if team_id is provided and model_name not in index, - # check if model_name matches any team_public_model_name - # O(n) scan but only when team_id lookup fails - for _idx, model in enumerate(self._core_model_list): - if self._core_should_include_deployment(model_name=model_name, model=model, team_id=team_id): - if model_alias is not None: - # Optimized: Use shallow copy since we only modify top-level model_name - alias_model = model.copy() - alias_model["model_name"] = model_alias - returned_models.append(alias_model) - else: - returned_models.append(model) - - return returned_models - - def _core_get_model_names(self, team_id: str | None = None) -> list[str]: - """ - Returns all possible model names for the router, including models defined via model_group_alias. - - If a team_id is provided, only deployments configured with that team_id (i.e. team‐specific models) - will yield their team public name. - """ - deployments: Final = self._core_get_model_list() or [] - model_names: Final = [] - - for deployment in deployments: - model_info = deployment.get("model_info") - if self._is_team_specific_model(model_info): - team_model_name = self._get_team_specific_model(deployment=deployment, team_id=team_id) - if team_model_name: - model_names.append(team_model_name) - else: - model_names.append(deployment.get("model_name", "")) - - return model_names - - def _core_get_fully_blocked_model_names(self) -> set[str]: - """ - Returns the set of model_names where every backing deployment has `blocked=True`. - - Used by `/v1/models` to hide paused models from client listings while still - surfacing them on admin endpoints (e.g. `/model/info`). A model with at least - one non-blocked deployment is still serviceable and remains visible. - """ - deployments: Final = self._core_get_model_list() or [] - blocked_by_name: Final[dict[str, bool]] = {} - for deployment in deployments: - name = deployment.get("model_name") or "" - if not name: - continue - is_blocked = (deployment.get("model_info") or {}).get("blocked") is True - if name in blocked_by_name: - blocked_by_name[name] = blocked_by_name[name] and is_blocked - else: - blocked_by_name[name] = is_blocked - return {name for name, fully_blocked in blocked_by_name.items() if fully_blocked} - - @staticmethod - def _core_are_all_deployments_blocked( - deployments: Sequence[DeploymentTypedDict], - ) -> bool: - return len(deployments) > 0 and all( - (deployment.get("model_info") or {}).get("blocked") is True for deployment in deployments - ) - - def _is_model_fully_blocked(self, model: str) -> bool: - deployments: Final = self._core_get_model_list(model_name=model) or [] - return self._core_are_all_deployments_blocked(deployments=deployments) - - async def _core_async_get_fully_unhealthy_model_names(self) -> set[str]: - """ - Returns the set of model names where every backing deployment is currently - marked unhealthy by background health checks (and the health state is not stale). - - Used by `/v1/models?healthy_only=true` to hide models that cannot serve any - request. A model with at least one healthy (or unknown-health) deployment - remains visible. Returns an empty set when no health state is available, so - callers fail open to the unfiltered listing. - - Notes: - - Mirrors `_async_filter_health_check_unhealthy_deployments`: when - `allowed_fails_policy` is set, cooldown is the sole routing exclusion - mechanism, so nothing is hidden here either. - - Team-specific public model names (`team_public_model_name`) are - aggregated alongside `model_name`, so team aliases of fully-unhealthy - deployments are hidden too (unlike `get_fully_blocked_model_names`, - which matches `model_name` only). - - Wildcard routes (e.g. `openai/*`) are matched by their literal - deployment name only; models expanded from a wildcard route are not - hidden (fail open). - - Intentionally diverges from the routing-time safety net (which - bypasses the health filter when every candidate is unhealthy and - still attempts the request): hiding here is presentation-only — - it answers "should this model be advertised?", not "should a - request for it still be attempted?". A hidden model can still be - called directly. - """ - if self._core_allowed_fails_policy is not None: - return set() - unhealthy_ids: Final = await self._core_health_state_cache.async_get_unhealthy_deployment_ids() - if not unhealthy_ids: - return set() - deployments: Final = self._core_get_model_list() or [] - unhealthy_by_name: Final[dict[str, bool]] = {} - for deployment in deployments: - model_info = deployment.get("model_info") or {} - names = [deployment.get("model_name") or ""] - team_public_model_name = model_info.get("team_public_model_name") - if team_public_model_name: - names.append(team_public_model_name) - is_unhealthy = model_info.get("id") in unhealthy_ids - for name in names: - if not name: - continue - if name in unhealthy_by_name: - unhealthy_by_name[name] = unhealthy_by_name[name] and is_unhealthy - else: - unhealthy_by_name[name] = is_unhealthy - return {name for name, fully_unhealthy in unhealthy_by_name.items() if fully_unhealthy} - - def _get_team_specific_model(self, deployment: DeploymentTypedDict, team_id: str | None = None) -> str | None: - """ - Get the team-specific model name if team_id matches the deployment. - - Args: - deployment: DeploymentTypedDict - The model deployment - team_id: Optional[str] - If passed, will return router models set with a `team_id` matching the passed `team_id`. - - Returns: - str: The `team_public_model_name` if team_id matches - None: If team_id doesn't match or no team info exists - """ - model_info: Final[dict | None] = deployment.get("model_info") or {} - if model_info is None: - return None - if team_id == model_info.get("team_id"): - return model_info.get("team_public_model_name") - return None - - def _is_team_specific_model(self, model_info: dict | None) -> bool: - """ - Check if model info contains team-specific configuration. - - Args: - model_info: Model information dictionary - - Returns: - bool: True if model has team-specific configuration - """ - return bool(model_info and model_info.get("team_id")) - - def _core_get_model_list_from_model_alias(self, model_name: str | None = None) -> list[DeploymentTypedDict]: - """ - Helper function to get model list from model alias. - - Used by `.get_model_list` to get model list from model alias. - """ - returned_models: Final[list[DeploymentTypedDict]] = [] - - if model_name is not None: - # Fast path: direct dict lookup avoids scanning all aliases for non-alias model names. - if model_name not in self._core_model_group_alias: - return returned_models - alias_items = [(model_name, self._core_model_group_alias[model_name])] - else: - alias_items = list(self._core_model_group_alias.items()) - - for model_alias, model_value in alias_items: - if isinstance(model_value, str): - _router_model_name: str = model_value - elif isinstance(model_value, dict): - _model_value = RouterModelGroupAliasItem(**model_value) - if _model_value["hidden"] is True and model_name is None: - continue - _router_model_name = _model_value["model"] - else: - continue - - if (alias_group := self.get_routing_group(_router_model_name)) is not None: - returned_models.extend( - {**row, "model_name": model_alias} for row in self._materialize_routing_group_rows((alias_group,)) - ) - else: - returned_models.extend( - self._core_get_all_deployments(model_name=_router_model_name, model_alias=model_alias) - ) - - return returned_models - - def _core_get_model_list_from_routing_groups(self, model_name: str | None = None) -> Sequence[DeploymentTypedDict]: - """ - Callable routing groups materialized as model-list rows, mirroring - `get_model_list_from_model_alias`: each member deployment is emitted - under the group's name (via `get_all_deployments`' `model_alias` - rewrite), which is what surfaces groups in `get_model_names`, - `/v1/models` discovery, `get_model_group_usage`, and the - blocked/unhealthy hiding that all read `get_model_list`. - """ - if model_name is not None: - group: Final = self.get_routing_group(model_name) - return self._materialize_routing_group_rows((group,)) if group is not None else () - cached: Final = self._routing_group_rows - if cached is not None: - return cached - rows: Final = self._materialize_routing_group_rows( - tuple( - callable_group - for name in self._core_routing_groups - if (callable_group := self.get_routing_group(name)) is not None - ) - ) - self._routing_group_rows = rows - return rows - - def _materialize_routing_group_rows(self, groups: tuple[RoutingGroup, ...]) -> tuple[DeploymentTypedDict, ...]: - return tuple( - self._as_routing_group_row(apply_routing_group_priority(group, member, deployment)) - for group in groups - for member in group.models - for deployment in self._core_get_all_deployments(model_name=member, model_alias=group.group_name) - ) - - @staticmethod - def _as_routing_group_row(deployment: DeploymentTypedDict) -> DeploymentTypedDict: - """ - A member deployment re-emitted under its group's name must not carry - the member's `access_groups`: access groups grant member names, never - the group, so inheriting them here would let a key holding a member's - access group list and call the whole group. - """ - model_info: Final = {k: v for k, v in (deployment.get("model_info") or {}).items() if k != "access_groups"} - return {**deployment, "model_info": model_info} - - TIER_PARAMS_NEVER_DROPPED: Final = frozenset(all_litellm_params) | frozenset( - { - "additional_drop_params", - "drop_params", - "messages", - "model", - "extra_headers", - "max_tokens", - "max_completion_tokens", - } - ) - - @staticmethod - def _declared_param_allowlist(params: Mapping[str, object]) -> frozenset[str]: - declared: Final = params.get("allowed_openai_params") - if not isinstance(declared, (list, tuple, set, frozenset)): - return frozenset() - return frozenset(entry for entry in declared if isinstance(entry, str)) - - @staticmethod - def _deployment_accepts_param(deployment: DeploymentTypedDict, group: str, param: str) -> bool: - deployment_params: Final = deployment.get("litellm_params") - if not deployment_params: - return True - if param in PythonCore._declared_param_allowlist(deployment_params): - return True - if declared_authenticating_provider( - str(deployment_params.get("model") or ""), deployment_params.get("custom_llm_provider") - ): - return True - return PythonCore._deployment_supports_param(deployment, group, param) - - @staticmethod - def _capability_model_and_provider(deployment_params: LiteLLMParamsTypedDict, group: str) -> tuple[str, str]: - model_label: Final = str(deployment_params.get("model") or group) - authenticating_provider: Final = declared_authenticating_provider( - model_label, deployment_params.get("custom_llm_provider") - ) - if authenticating_provider is not None: - return model_label.removeprefix(f"{authenticating_provider}/"), authenticating_provider - model, custom_llm_provider, _, _ = python_host.invoke("get_llm_provider")( - model=model_label, - custom_llm_provider=deployment_params.get("custom_llm_provider"), - litellm_params=LiteLLM_Params.model_validate(deployment_params), - ) - return model, custom_llm_provider - - @staticmethod - def _deployment_supports_param(deployment: DeploymentTypedDict, group: str, param: str) -> bool: - deployment_params: Final = deployment.get("litellm_params") - if not deployment_params: - return True - deployment_model_info: Final = deployment.get("model_info") - base_model: Final = ( - deployment_model_info.get("base_model") if deployment_model_info else None - ) or deployment_params.get("base_model") - try: - model, custom_llm_provider = PythonCore._capability_model_and_provider(deployment_params, group) - supported: Final = python_host.invoke("get_supported_openai_params")( - model=model, - custom_llm_provider=custom_llm_provider, - base_model=base_model if isinstance(base_model, str) else None, - ) - except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request - verbose_router_logger.debug( - "litellm.router.py::_deployment_supports_param: keeping %s for model=%s. Got - %s", param, group, e - ) - return True - return supported is None or param in supported - - def _tier_params_the_target_accepts( - self, model: str, tier_params: Mapping[str, object], request_kwargs: Mapping[str, object] - ) -> Mapping[str, object]: - """Drop an OpenAI param that no deployment behind ``model`` declares. - - A tier's litellm_params are an operator override applied to every request the tier routes, - so one the target cannot take turns that whole tier into a 400 raised before the request - leaves the proxy. The candidates are exactly what get_optional_params can reject, asked of - the module that raises, so credentials and endpoint controls are never at risk. - - TIER_PARAMS_NEVER_DROPPED is excluded on top of that, for two reasons. No provider lists a - litellm control among its supported params, so "no deployment declares it" means litellm - consumes it rather than that the target refuses it, and dropping one changes litellm's own - behavior: dropping drop_params or additional_drop_params silently disables the sanitization - the operator configured. Providers do list extra_headers, but it carries auth, tenancy and - routing information, so sending fewer headers than configured is worse than today's error. - Token ceilings stay for the same reason: a tier's max_tokens or max_completion_tokens is a - cost bound, and dropping it would let a caller's own larger value through where today the - mismatch fails loudly. - - The trade this filter makes is a param for a working request, which is right for one that - only shapes how the model answers and wrong for anything else. - - A param survives if ANY deployment could take it, because routing has not chosen one yet, - and it survives both an unresolvable provider and a group with no deployments, because a - best-effort filter must never narrow what the request already did. - - A github_copilot or chatgpt deployment counts as accepting everything, decided before any - lookup: resolving either provider runs its OAuth device flow, so a capability question - asked from the routing path can freeze the event loop for minutes waiting on a human. - - allowed_openai_params is the documented escape hatch for an outdated or incomplete - supported-params list: request-time validation extends the supported list with it before - comparing. The filter asks the same question, so a param named by the allowlist on the tier - overlay, the request, or a deployment's own litellm_params is never a drop candidate. - """ - deployments: Final = self._core_get_model_list(model_name=model) or () - if not deployments: - return tier_params - allowlisted: Final = self._declared_param_allowlist(tier_params) | self._declared_param_allowlist( - request_kwargs - ) - candidates: Final = provider_rejectable_params(tier_params) - self.TIER_PARAMS_NEVER_DROPPED - allowlisted - unsupported: Final = frozenset( - param - for param in candidates - if not any(self._deployment_accepts_param(deployment, model, param) for deployment in deployments) - ) - if not unsupported: - return tier_params - verbose_router_logger.warning( - "litellm.router.py: dropping tier params %s for model=%s, no deployment behind it declares them", - ", ".join(sorted(unsupported)), - model, - ) - return MappingProxyType({key: value for key, value in tier_params.items() if key not in unsupported}) - - def _core_get_model_list( - self, model_name: str | None = None, team_id: str | None = None - ) -> list[DeploymentTypedDict] | None: - """ - Includes router model_group_alias'es as well - - if team_id specified, returns matching team-specific models - """ - # Note: model_list and model_group_alias are always initialized in __init__ - # so hasattr checks are unnecessary - returned_models: list[DeploymentTypedDict] = [] - - if model_name is not None: - returned_models.extend(self._core_get_all_deployments(model_name=model_name, team_id=team_id)) - - returned_models.extend(self._core_get_model_list_from_model_alias(model_name=model_name)) - returned_models.extend(self._core_get_model_list_from_routing_groups(model_name=model_name)) - - if len(returned_models) == 0 and model_name is not None: - returned_models.extend(self._get_wildcard_deployments(model_name=model_name, team_id=team_id)) - - if model_name is None: - returned_models += self._core_model_list - - return returned_models - - def _get_wildcard_deployments(self, model_name: str, team_id: str | None = None) -> list[DeploymentTypedDict]: - """ - The deployments of the wildcard routes matching model_name (the proxy-wide - ones, plus team_id's own when given), each emitted under model_name. - """ - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - matches: Final = [ - *self._core_pattern_router.get_deployments_by_pattern(model=model_name), - *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), - ] - return [{**DeploymentTypedDict(**m), "model_name": model_name} for m in matches] - - def _core_wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: - """The deployments routing serves ``model`` from through a wildcard: ``team_id``'s own pattern - matching it, else the most specific proxy-wide pattern, never both.""" - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - team_owned: Final[Sequence[Mapping[str, object]]] = ( - team_router.get_deployments_by_pattern(model=model) if team_router is not None else () - ) - return team_owned or self._core_pattern_router.get_deployments_by_pattern(model=model) - - def _core_wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]: - """The models wildcard_route_deployments would route ``model`` to, without copying the deployments.""" - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - team_owned: Final = ( - team_router.get_models_by_pattern(model=model) if team_router is not None else frozenset[str]() - ) - return team_owned or self._core_pattern_router.get_models_by_pattern(model=model) - - def _core_team_ptu_ceiling(self, team_id: str, model_group: str) -> PTUTeamCeiling | None: - """The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``, - else None, kept until the model list next changes.""" - if team_id not in self._ptu_share_teams(): - return None - return self._cached_team_ptu_ceiling(team_id, model_group) - - def _collect_ptu_share_teams(self) -> frozenset[str]: - return ptu_share_teams(self._core_get_model_list() or ()) - - def _resolve_team_ptu_ceiling(self, team_id: str, model_group: str) -> PTUTeamCeiling | None: - return team_ptu_ceiling( - self._core_get_model_list() or (), - self._core_model_list, - team_id, - model_group, - lambda name: self._core_wildcard_route_deployments(name, team_id), - ) - - def _core_get_model_list_of_routed_group( - self, model_group: str, team_id: str | None = None - ) -> list[DeploymentTypedDict]: - """ - The deployments a request the router has already resolved to model_group is - served from: team_id's own deployments under that public name when the team has - any, else the ones named model_group, the routing group of that name, or the - wildcard route matching it when neither exists. - - Unlike get_model_list, model_group's own model_group_alias entry is not - followed. The router resolves an alias exactly once, so a group reached as - an alias target is served by its own deployments, never by a second hop: - in the chain X -> T -> U a request to X is served from T's deployments. - """ - if team_id is not None and self._team_deployments_named(team_id, model_group): - return self._core_get_all_deployments(model_name=model_group, team_id=team_id) - own: Final = self._get_all_deployments(model_name=model_group) - named: Final = [*own, *self._core_get_model_list_from_routing_groups(model_name=model_group)] - return named or self._get_wildcard_deployments(model_name=model_group) - - def _core_resolved_litellm_models( - self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False - ) -> tuple[str, ...]: - """The provider model strings `model_name` can actually be served by on this proxy. - - `get_model_list` composes every channel the request path itself uses (exact name, - model_group_alias, routing groups, wildcards), so this answers "which models will - answer a call to this name" rather than "what did the admin call it": the deployment - name is admin-arbitrary, and two names over one provider model are one model. - - Empty when the name resolves to no deployment. That is not the same fact as "the - call will fail" - a provider-qualified public name is served by the SDK with no - deployment behind it - so the fallback for an empty result is the caller's policy, - never this function's. - """ - if use_base_model: - return tuple( - backend_model - for deployment in self._core_get_model_list(model_name=model_name, team_id=team_id) or () - if (backend_model := self._cost_map_backend_model(deployment)) - ) - return tuple( - litellm_model - for deployment in self._core_get_model_list(model_name=model_name, team_id=team_id) or () - if isinstance(litellm_model := deployment.get("litellm_params", {}).get("model"), str) and litellm_model - ) - - def _invalidate_model_group_info_cache(self) -> None: - """Invalidate the cached model group info. - - Call this whenever self._core_model_list is modified to ensure the cache is rebuilt. - Also clears the auth-layer zero-cost cache, which depends on the same - ``ModelGroupInfo`` data — without this, an in-place pricing update on - an existing deployment (same model count) would keep a stale ``True`` - result and bypass budget enforcement. - """ - self._cached_get_model_group_info.cache_clear() - self._core_cached_deployment_model_info.cache_clear() - self._ptu_share_teams.cache_clear() - self._cached_team_ptu_ceiling.cache_clear() - self._cached_deployment_tracks_usage.cache_clear() - self.served_openai_models.cache_clear() - self._zero_cost_cache.clear() - self._routing_group_rows = None - - def _invalidate_access_groups_cache(self) -> None: - """Invalidate the cached access groups. - - Call this whenever self._core_model_list is modified to ensure the cache is rebuilt. It also - advances ``registry().version``, so every write path (``apply``, the old ``set_model_list`` - and ``add_deployment``) moves it. - """ - self._access_groups_cache = None - self._registry_version += 1 - - def _core_get_model_access_groups( - self, - model_name: str | None = None, - model_access_group: str | None = None, - team_id: str | None = None, - ) -> dict[str, list[str]]: - """ - If model_name is provided, only return access groups for that model. - - Parameters: - - model_name: Optional[str] - the received model name from the user (can be a wildcard route). If set, will only return access groups for that model. - - model_access_group: Optional[str] - the received model access group from the user. If set, will only return models for that access group. - - team_id: Optional[str] - the team id, to resolve team-specific models - """ - # Check if this is the no-args hot path (cacheable) - _use_cache: Final = model_name is None and model_access_group is None and team_id is None - - # Return cached result for the no-args hot path - if _use_cache and self._access_groups_cache is not None: - return self._access_groups_cache - - from collections import defaultdict - - access_groups: Final = defaultdict(list) - - model_list: Final = self._core_get_model_list(model_name=model_name, team_id=team_id) - if model_list: - for m in model_list: - _model_info = m.get("model_info") - if _model_info: - for group in _model_info.get("access_groups", []) or []: - if model_access_group is not None: - if group == model_access_group: - model_name = m["model_name"] - access_groups[group].append(model_name) - else: - model_name = m["model_name"] - access_groups[group].append(model_name) - - # Cache the result for the no-args hot path - if _use_cache: - self._access_groups_cache = dict(access_groups) - return self._access_groups_cache - - return access_groups - - def _core_is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: - """ - Return True if model access group is a wildcard route - """ - # GET ACCESS GROUPS - access_groups: Final = self._core_get_model_access_groups(model_access_group=model_access_group) - - if len(access_groups) == 0: - return False - - models: Final = access_groups.get(model_access_group, []) - - for model in models: - # CHECK IF MODEL ACCESS GROUP IS A WILDCARD ROUTE - if self._core_pattern_router.route(request=model) is not None: - return True - - return False - - def _core_get_settings(self) -> dict[str, object]: # mutable-ok: the settings snapshot is built fresh per call - """ - Get router settings method, returns a dictionary of the settings and their values. - For example get the set values for routing_strategy_args, routing_strategy, allowed_fails, cooldown_time, num_retries, timeout, max_retries, retry_after - """ - _all_vars: Final = vars(self) - _settings_to_return: Final = {} - vars_to_include: Final = [ - "routing_strategy_args", - "routing_strategy", - "allowed_fails", - "cooldown_time", - "num_retries", - "timeout", - "max_retries", - "retry_after", - "fallbacks", - "context_window_fallbacks", - "model_group_retry_policy", - "retry_policy", - "model_group_alias", - "enable_weighted_failover", - "enable_mid_stream_fallback_continuation", - "enable_tag_filtering", - "tag_routing_prefix", - "fallback_on_unsupported_params", - ] - - for var in vars_to_include: - stored = var if var in _all_vars else "_core_" + var - if stored in _all_vars: - _settings_to_return[var] = _all_vars[stored] - if ( - var == "routing_strategy_args" - and self._core_routing_strategy == "latency-based-routing" - and self.lowestlatency_logger is not None - ): - _settings_to_return[var] = self.lowestlatency_logger.routing_args.json() - - _settings_to_return["routing_groups"] = [ - group.model_dump(exclude=frozenset(("model_priorities",)) if group.model_priorities is None else None) - for group in self._core_routing_groups.values() - ] - return _settings_to_return - - def _switch_routing_strategy(self, routing_strategy: str | None, kwargs: Mapping[str, object]) -> None: - if routing_strategy == "lar1": - from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - - apply_lar1_routing_strategy(self._core_ctx, kwargs.get("routing_strategy_args")) - return - self.routing_strategy_init( - routing_strategy=routing_strategy, routing_strategy_args=kwargs.get("routing_strategy_args", {}) - ) - - def _core_update_settings(self, **kwargs) -> None: - """ - Update the router settings. - """ - _int_settings: Final = [ - "timeout", - "num_retries", - "retry_after", - "allowed_fails", - "cooldown_time", - ] - - _existing_router_settings: Final = self._core_get_settings() - model_group_alias_before: Final = self._core_model_group_alias - rebuild_routing_groups = False - routing_args_updated = False - for var in kwargs: - if var in RUNTIME_UPDATABLE_ROUTER_SETTINGS: - if var in _int_settings: - _casted_value = int(kwargs[var]) - setattr(self, var, _casted_value) - elif var == "routing_groups": - rebuild_routing_groups = True - elif var == "optional_pre_call_checks": - self.set_optional_pre_call_checks(kwargs[var]) - elif var == "enable_mid_stream_fallback_continuation": - self._set_mid_stream_fallback_continuation(kwargs[var]) - elif var == "retry_policy": - self._set_retry_policy(kwargs[var]) - else: - value = kwargs[var] - # only run routing strategy init if it has changed - if var == "routing_strategy": - value = self._normalize_strategy(value) - if _existing_router_settings["routing_strategy"] != value: - self._switch_routing_strategy(value, kwargs) - rebuild_routing_groups = True - elif var == "routing_strategy_args": - routing_args_updated = value != self.routing_strategy_args - setattr(self, var, value) - else: - verbose_router_logger.debug("Setting %s is not allowed", var) - - if routing_args_updated: - self._apply_updated_routing_strategy_args() - - if self._core_model_group_alias != model_group_alias_before: - self._invalidate_model_group_info_cache() - - if rebuild_routing_groups: - routing_groups_input: Final = kwargs.get("routing_groups", self._routing_groups_input) - self._init_routing_groups(routing_groups_input) - self._routing_groups_input = routing_groups_input - verbose_router_logger.debug("Updated Router settings: %s", self._core_get_settings()) - - @staticmethod - def _is_flag_enabled(value: object) -> bool: - return value is True or (isinstance(value, str) and str_to_bool(value) is True) - - def _set_mid_stream_fallback_continuation(self, enabled: object) -> None: - self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enabled) - if self.enable_mid_stream_fallback_continuation: - self._register_continuation_prefill_check() - - def _set_retry_policy(self, value: object) -> None: - policy: Final = RetryPolicy.model_validate(value) if isinstance(value, dict) else value - if policy is None or isinstance(policy, RetryPolicy): - self.retry_policy = policy - - def _get_client(self, deployment, kwargs, client_type=None): - """ - Returns the appropriate client based on the given deployment, kwargs, and client_type. - - Parameters: - deployment (dict): The deployment dictionary containing the clients. - kwargs (dict): The keyword arguments passed to the function. - client_type (str): The type of client to return. - - Returns: - The appropriate client based on the given client_type and kwargs. - """ - model_id: Final = deployment["model_info"]["id"] - parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs) - if client_type == "max_parallel_requests": - cache_key = f"{model_id}_max_parallel_requests_client" - client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) - if client is None: - InitalizeCachedClient.set_max_parallel_requests_client(ctx=self._core_ctx, model=deployment) - client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) - return client - elif client_type == "async": - if ( - deployment["litellm_params"].get( # pyright: ignore[reportUnknownMemberType] # legacy mapping - "kubernetes_pod_discovery" - ) - is True - ): - return None - if kwargs.get("stream") is True: - cache_key = f"{model_id}_stream_async_client" - client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) - return client - else: - cache_key = f"{model_id}_async_client" - client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) - return client - else: - if ( - deployment["litellm_params"].get( # pyright: ignore[reportUnknownMemberType] # legacy mapping - "kubernetes_pod_discovery" - ) - is True - ): - return None - if kwargs.get("stream") is True: - cache_key = f"{model_id}_stream_client" - client = self._core_cache.get_cache(key=cache_key, parent_otel_span=parent_otel_span) - return client - else: - cache_key = f"{model_id}_client" - client = self._core_cache.get_cache(key=cache_key, parent_otel_span=parent_otel_span) - return client - - def _count_pre_call_check_tokens( - self, - messages: Sequence[Mapping[str, object]] | None, - input: str | list[object] | None, - request_kwargs: Mapping[str, object] | None = None, - tokenizer_model: str = "", - deployments: Sequence[Mapping[str, object]] = (), - ) -> int: - """ - Count input tokens for context-window pre-call checks. - - Chat Completions send `messages`; the Responses API sends `input` (a string or - a list of Responses input items) plus an optional `instructions` system prompt. - The Responses payload is normalized to chat messages via the shared - LiteLLMCompletionResponsesConfig transform so the same token_counter path covers - both API surfaces and `instructions` tokens are included in the count. - - Prompt content the message list never carries is read from `request_kwargs`: - `tools` (Chat Completions, Responses and Anthropic Messages shapes) and the - Anthropic Messages top-level `system` block. - """ - from litellm.llms.anthropic.messages.native.utils import ( - anthropic_system_to_openai_message, - ) - - extras: Final = request_kwargs if request_kwargs is not None else MappingProxyType({}) - raw_instructions: Final = extras.get("instructions") - instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None - raw_tools: Final = extras.get("tools") - tools: Final = ( - cast(list[ChatCompletionToolParam], raw_tools) # cast-ok: token_counter formats any tool dict shape - if isinstance(raw_tools, list) and raw_tools - else None - ) - system_message: Final = anthropic_system_to_openai_message(extras.get("system")) - if messages is not None: - counted_messages: Final = (system_message, *messages) if system_message is not None else messages - return self._host.invoke("token_counter")(model=tokenizer_model, messages=counted_messages, tools=tools) - if input is not None: - from openai.types.responses.response_create_params import ResponseInputParam - - from litellm.litellm_core_utils.token_counter import split_embedding_input_tokens - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) - - text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) - typed_input: Final = cast( # cast-ok: str | list matches transform input - str | ResponseInputParam, text_input - ) - input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=typed_input, - responses_api_request={"instructions": instructions} if instructions is not None else {}, - only_since_compaction=True, - ) - return self._host.invoke("token_counter")( - model=tokenizer_model, - messages=cast(list, input_messages), # cast-ok: transformed chat messages - tools=tools, - ) - raise ValueError("Either messages or input must be provided to count tokens") - - def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None: - """The deployment's declared context window, or None when it declares none or cannot be resolved.""" - try: - model_info: Final = self._core_get_router_model_info( - deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts - received_model_name=model, - ) - except Exception as e: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others - verbose_router_logger.debug( - "litellm.router.py::_deployment_max_input_tokens: skipping deployment. Got - %s", e - ) - return None - max_input_tokens: Final = model_info.get("max_input_tokens") - return max_input_tokens if isinstance(max_input_tokens, int) else None - - def _pre_call_checks_need_token_count( - self, model: str, healthy_deployments: Sequence[Mapping[str, object]] - ) -> bool: - """Whether any healthy deployment declares a context window that a token count could exceed. - - Resolves each deployment the way ``_pre_call_checks`` does, so one unmappable deployment - cannot hide a later one that does declare a limit. - """ - return any( - self._deployment_max_input_tokens(model, deployment) is not None for deployment in healthy_deployments - ) - - async def _acount_pre_call_check_tokens( - self, - model: str, - healthy_deployments: Sequence[Mapping[str, object]], - messages: Sequence[Mapping[str, str]] | None, - input: str | Sequence[object] | None, - request_kwargs: Mapping[str, object] | None, - ) -> int | None: - """Count input tokens off the event loop, so a multi-MB prompt cannot stall the proxy. - - Returns None when no deployment limits its context window, and when counting fails. The - caller pairs this with ``skip_inline_token_count`` so neither case puts the count back on - the loop: a failed count only skips the context-window check - RPM, region and - invalid-params checks still run for every deployment. - """ - if messages is None and input is None: - return None - try: - if not self._pre_call_checks_need_token_count(model, healthy_deployments): - return None - return await offload_token_count(self._count_pre_call_check_tokens)( - messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter - input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter - request_kwargs=request_kwargs, - tokenizer_model=single_encoding_openai_model( - self.served_openai_models(model) or served_openai_models(healthy_deployments) - ) - or "", - deployments=healthy_deployments, - ) - except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request - verbose_router_logger.error( - "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e - ) - return None - - def _resolve_pre_call_check_input_tokens( - self, - messages: Sequence[Mapping[str, str]] | None, - input: str | Sequence[object] | None, - request_kwargs: Mapping[str, object] | None, - input_token_count: int | None, - count_inline: bool, - model: str, - deployments: Sequence[Mapping[str, object]], - ) -> int | None: - if input_token_count is not None: - return input_token_count - if not count_inline: - return None - try: - return self._count_pre_call_check_tokens( - messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter - input=cast(str | list[object] | None, input), # cast-ok: forwarded to the sync counter - request_kwargs=request_kwargs, - tokenizer_model=single_encoding_openai_model( - self.served_openai_models(model) or served_openai_models(deployments) - ) - or "", - deployments=deployments, - ) - except Exception as e: - verbose_router_logger.error( - "litellm.router.py::_pre_call_checks: failed to count tokens. " - "Skipping the context-window check only; RPM/region/param checks still apply. Got - %s", - e, - ) - return None - - @staticmethod - def _params_neither_dropped_nor_allowed(params: Mapping[str, object], candidates: frozenset[str]) -> frozenset[str]: - if params.get("drop_params") is True: - return frozenset() - dropped: Final = params.get("additional_drop_params") - dropped_names: Final = ( - frozenset(entry for entry in dropped if isinstance(entry, str)) - if isinstance(dropped, list) - else frozenset() - ) - return candidates - PythonCore._declared_param_allowlist(params) - dropped_names - - @staticmethod - def _deployment_lacks_params(deployment: DeploymentTypedDict, group: str, params: frozenset[str]) -> bool: - return any(not PythonCore._deployment_supports_param(deployment, group, param) for param in params) - - @staticmethod - def _deployment_rejects_params( - deployment: DeploymentTypedDict, group: str, params: frozenset[str], request_kwargs: Mapping[str, object] - ) -> bool: - call_settings: Final[Mapping[str, object]] = {**(deployment.get("litellm_params") or {}), **request_kwargs} - return PythonCore._deployment_lacks_params( - deployment, group, PythonCore._params_neither_dropped_nor_allowed(call_settings, params) - ) - - @staticmethod - def _deployment_declares_no_vision(deployment: DeploymentTypedDict, group: str) -> bool: - plain_deployment: Final = cast(Mapping[str, object], deployment) # cast-ok: router deployments are plain dicts - model_info: Final = cast(Mapping[str, object], plain_deployment.get("model_info") or {}) # cast-ok: plain dict - declared: Final = model_info.get("supports_vision") - if declared is not None: - return declared is False - raw_params: Final = plain_deployment.get("litellm_params") or {} - params: Final = cast(Mapping[str, object], raw_params) # cast-ok: router deployments are plain dicts - if not params: - return False - base_model: Final = model_info.get("base_model") or params.get("base_model") - try: - model, custom_llm_provider = PythonCore._capability_model_and_provider(deployment["litellm_params"], group) - return litellm_utils.is_vision_explicitly_disabled( - base_model if isinstance(base_model, str) else model, custom_llm_provider - ) - except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request - verbose_router_logger.debug( - "litellm.router.py::_deployment_declares_no_vision: keeping deployment for model=%s. Got - %s", group, e - ) - return False - - @staticmethod - def _deployment_serves_route( - deployment: DeploymentTypedDict, - group: str, - resolver: Callable[[str, LlmProviders], object | None], - ) -> bool: - if "litellm_params" not in deployment: - return True - - try: - model, resolved_custom_llm_provider = PythonCore._capability_model_and_provider( - deployment["litellm_params"], group - ) - provider: Final = next( - (member for member in LlmProviders if member.value == resolved_custom_llm_provider), - None, - ) - return provider is not None and resolver(model, provider) is not None - except Exception as e: # noqa: BLE001 # best-effort capability filter - verbose_router_logger.debug( - "litellm.router.py::_deployment_serves_route: keeping deployment for model=%s. Got - %s", group, e - ) - return True - - @staticmethod - def _without_deployments_that_cannot_serve_route( - healthy_deployments: _RouteDeploymentsT, group: str - ) -> _RouteDeploymentsT: - handler_name: Final = _routed_call_type.get() - resolver: Final = _ROUTE_CAPABILITY_RESOLVERS.get(handler_name) if handler_name is not None else None - if resolver is None or isinstance(healthy_deployments, Mapping): - return healthy_deployments - - filtered_deployments: Final = [ - deployment - for deployment in healthy_deployments - if PythonCore._deployment_serves_route( - cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape - group, - resolver, - ) - ] - if not filtered_deployments: - return healthy_deployments - return cast(_RouteDeploymentsT, filtered_deployments) # cast-ok: same list shape as the input candidates - - @staticmethod - def _decisions_deployment_refuses_media( - deployment: DeploymentTypedDict, group: str, request_kwargs: Mapping[str, object], *, audio: bool, videos: bool - ) -> bool: - from litellm.decisions.call import deployment_refuses_media - - plain_deployment: Final = cast(Mapping[str, object], deployment) # cast-ok: router deployments are plain dicts - model_info: Final = cast(Mapping[str, object], plain_deployment.get("model_info") or {}) # cast-ok: plain dict - deployment_params: Final = cast(Mapping[str, object], deployment["litellm_params"]) # cast-ok: plain dict - try: - model, custom_llm_provider = PythonCore._capability_model_and_provider(deployment["litellm_params"], group) - return deployment_refuses_media( - model, - custom_llm_provider, - model_info, - {**deployment_params, **request_kwargs}, - audio=audio, - videos=videos, - ) - except Exception as e: # noqa: BLE001 # best-effort filter: a deployment it cannot resolve answers for itself - verbose_router_logger.debug( - "litellm.router.py::_decisions_deployment_refuses_media: keeping a deployment of model=%s. Got - %s", - group, - e, - ) - return False - - def _without_decisions_deployments_refusing_media( - self, healthy_deployments: _RouteDeploymentsT, group: str, request_kwargs: Mapping[str, object] | None - ) -> _RouteDeploymentsT: - call_kwargs: Final = request_kwargs or {} - audio: Final = bool(call_kwargs.get("audio")) - videos: Final = bool(call_kwargs.get("videos")) - if ( - _routed_call_type.get() not in _DECISIONS_MEDIA_CALL_TYPES - or not (audio or videos) - or isinstance(healthy_deployments, Mapping) - ): - return healthy_deployments - - sent_kwargs: Final = {**self._active_default_litellm_params(), **call_kwargs} - deployments_taking_media: Final = [ - deployment - for deployment in healthy_deployments - if not PythonCore._decisions_deployment_refuses_media( - cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape - group, - sent_kwargs, - audio=audio, - videos=videos, - ) - ] - if not deployments_taking_media: - return healthy_deployments - return cast(_RouteDeploymentsT, deployments_taking_media) # cast-ok: same list shape as the input candidates - - @staticmethod - def _without_declared_text_only_deployments( - healthy_deployments: _RouteDeploymentsT, - group: str, - messages: list[dict[str, str]] | None, # mutable-ok: handed to _request_carries_image as-is - request_kwargs: Mapping[str, object] | None, - ) -> _RouteDeploymentsT: - if isinstance(healthy_deployments, Mapping) or not PythonCore._request_carries_image( - messages, request_kwargs or {} - ): - return healthy_deployments - - vision_capable_deployments: Final = [ - deployment - for deployment in healthy_deployments - if not PythonCore._deployment_declares_no_vision( - cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape - group, - ) - ] - if not vision_capable_deployments: - return healthy_deployments - return cast(_RouteDeploymentsT, vision_capable_deployments) # cast-ok: same list shape as the input candidates - - @staticmethod - def _request_carries_image( - messages: object, - request_kwargs: Mapping[str, object], - ) -> bool: - from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content - - return request_contains_image_content(messages or request_kwargs.get("messages") or request_kwargs.get("input")) - - @staticmethod - def _pre_call_check_requested_chat_params(request_kwargs: Mapping[str, object]) -> frozenset[str]: - requested: Final = litellm_utils.get_non_default_params(passed_params=dict(request_kwargs)) - return frozenset( - param - for param, value in requested.items() - if param in _PRE_CALL_CHECK_FILTERED_PARAMS and not isinstance(value, str) - ) - - def _pre_call_checks( - self, - model: str, - healthy_deployments: list, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - request_kwargs: dict | None = None, - input_token_count: int | None = None, - skip_inline_token_count: bool = False, - request_deployments: Sequence[DeploymentTypedDict] | None = None, - ): - """ - Filter out model in model group, if: - - - model context window < message length. For azure openai models, requires 'base_model' is set. - https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models - - if region given, filter out models not in that region / unknown region - - [TODO] function call and model doesn't support function calling - """ - - verbose_router_logger.debug("Starting Pre-call checks for deployments in model=%s", model) - - # Optimized: Use list() shallow copy instead of deepcopy - # We only pop from the list, not modify deployment dicts - 100x+ faster on hot path (every request) - _returned_deployments = list(healthy_deployments) - - invalid_model_indices: Final = set[int]() # Use set for O(1) membership checks - - _context_window_error = False - _potential_error_str = "" - - has_countable_input: Final = (messages is not None or input is not None) and not compaction_pending( - request_kwargs - ) - deployment_max_input_tokens: Final = tuple( - self._deployment_max_input_tokens(model, deployment) for deployment in _returned_deployments - ) - input_tokens: Final = self._resolve_pre_call_check_input_tokens( - messages=messages, - input=cast(str | list[object] | None, input), # cast-ok: forwarded to the resolver - request_kwargs=cast(Mapping[str, object] | None, request_kwargs), # cast-ok: forwarded to the resolver - input_token_count=input_token_count, - count_inline=has_countable_input - and not skip_inline_token_count - and any(limit is not None for limit in deployment_max_input_tokens), - model=model, - deployments=_returned_deployments, - ) - request_settings: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else {} - requested_params: Final[frozenset[str]] = ( - self._pre_call_check_requested_chat_params(request_settings) - if request_kwargs is not None and self._host.dynamic().drop_params is False - else frozenset() - ) - - for idx, deployment in enumerate(_returned_deployments): - # Cache nested dict access to avoid repeated temporary dict allocations - _litellm_params = deployment.get("litellm_params", {}) - _model_info = deployment.get("model_info", {}) - - max_input_tokens = deployment_max_input_tokens[idx] - if ( - max_input_tokens is not None - and has_countable_input - and input_tokens is not None - and input_tokens > max_input_tokens - ): - invalid_model_indices.add(idx) - _context_window_error = True - deployment_model = ( - _model_info.get("base_model") or _litellm_params.get("base_model") or _litellm_params.get("model") - ) - _potential_error_str += ( - f"Model={deployment_model}, Max Input Tokens={max_input_tokens}, Got={input_tokens}" - ) - continue - - ## REGION CHECK ## - if request_kwargs is not None and request_kwargs.get("allowed_model_region") is not None: - allowed_model_region = request_kwargs.get("allowed_model_region") - - if allowed_model_region is not None: - if not is_region_allowed( - litellm_params=LiteLLM_Params.model_validate(_litellm_params), - allowed_model_region=allowed_model_region, - ): - invalid_model_indices.add(idx) - continue - - typed_deployments: Final = tuple( - cast(DeploymentTypedDict, deployment) # cast-ok: router deployments are plain dicts - for deployment in _returned_deployments - ) - param_candidates: Final = tuple( - (idx, deployment) for idx, deployment in enumerate(typed_deployments) if idx not in invalid_model_indices - ) - candidates_lacking_params: Final = tuple( - self._deployment_lacks_params(deployment, model, requested_params) for _, deployment in param_candidates - ) - no_deployment_takes_params: Final = bool(candidates_lacking_params) and all(candidates_lacking_params) - candidates_failing_params: Final = ( - tuple( - self._deployment_rejects_params(deployment, model, requested_params, request_settings) - for _, deployment in param_candidates - ) - if no_deployment_takes_params - else candidates_lacking_params - ) - excluded_indices: Final = frozenset(invalid_model_indices) | frozenset( - idx for (idx, _), fails in zip(param_candidates, candidates_failing_params, strict=True) if fails - ) - - if len(excluded_indices) == len(_returned_deployments): - group_deployments: Final = ( - tuple(request_deployments) - if request_deployments is not None - else (tuple(self._core_get_all_deployments(model_name=model)) or typed_deployments) - ) - if ( - not _context_window_error - and group_deployments - and all( - self._deployment_rejects_params(deployment, model, requested_params, request_settings) - for deployment in group_deployments - ) - ): - raise litellm_exceptions.UnsupportedParamsError( - message=( - f"no deployment of model={model} supports {sorted(requested_params)}. " - "To drop these, set `litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, " - f"or send allowed_openai_params={sorted(requested_params)} to pass them through" - ), - model=model, - ) - if _context_window_error is True: - raise litellm_exceptions.ContextWindowExceededError( - message=f"litellm._pre_call_checks: Context Window exceeded for given call. No models have context window large enough for this call.\n{_potential_error_str}", - model=model, - llm_provider="", - ) - if len(excluded_indices) > 0: - # Single-pass filter using set for O(1) lookups (avoids O(n^2) from repeated pops) - _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in excluded_indices] - - return _returned_deployments - - def _core_get_model_from_alias(self, model: str) -> str | None: - """ - Get the model from the alias. - - Returns: - - str, the litellm model name - - None, if model is not in model group alias - """ - return resolve_model_group_alias(self._core_model_group_alias, model) - - def _get_deployment_by_litellm_model(self, model: str) -> list[DeploymentTypedDict]: - """ - Get the deployment by litellm model. - """ - return [m for m in self._core_model_list if m["litellm_params"]["model"] == model] - - def _core_get_routable_deployments( - self, - model: str, - request_kwargs: Mapping[str, object] = MappingProxyType({}), - call_type: str | None = None, - ) -> tuple[DeploymentTypedDict, ...]: - try: - with endpoint_selection(call_type, model=model): - _, candidates = self._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # its signature leaves input and the result untyped - except (litellm_exceptions.BadRequestError, ValueError): - return () - if isinstance(candidates, dict): - return (cast(DeploymentTypedDict, candidates),) # cast-ok: the router returns one deployment as a dict - return tuple(cast(list[DeploymentTypedDict], candidates)) # cast-ok: the router returns deployments as a list - - def _try_early_resolve_deployments_for_model_not_in_names( - self, - model: str, - request_team_id: str | None, - include_team_models: bool = False, - ) -> tuple[str, list | dict] | None: - """ - When ``model`` is not in ``self._core_model_names``, try team routes, pattern routes, - team pattern routers, then default deployment. Returns None if none apply. - """ - if model in self._core_model_names: - return None - # Check for team-specific deployments by team_public_model_name. - # This intentionally takes priority over team pattern routers below, - # so that named team deployments shadow wildcard/pattern routes. - if request_team_id is not None: - team_deployments = self._core_get_all_deployments(model_name=model, team_id=request_team_id) - if team_deployments: - return model, team_deployments - elif include_team_models: - team_deployments = self._team_deployments_across_teams(model) - if team_deployments: - return model, team_deployments - - wildcard_deployments: Final = self._core_wildcard_route_deployments(model, request_team_id) - if wildcard_deployments: - return model, list(wildcard_deployments) - - if self._core_default_deployment is not None: - # Shallow copy with nested litellm_params copy (100x+ faster than deepcopy) - updated_deployment: Final = self._core_default_deployment.copy() - updated_deployment["litellm_params"] = self._core_default_deployment["litellm_params"].copy() - updated_deployment["litellm_params"]["model"] = model - return model, updated_deployment - - return None - - def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]: - """Every team's deployments under public name `model`, for a proxy admin calling without a team.""" - with self._deployment_index_lock: - team_deployments: Final = [ - self._core_model_list[index] - for (_, public_model_name), indices in self._core_team_model_to_deployment_indices.items() - if public_model_name == model - for index in indices - ] - team_ids: Final = { - team_id - for deployment in team_deployments - for team_id in [(deployment.get("model_info") or {}).get("team_id")] - if team_id is not None - } - if len(team_ids) > 1: - raise litellm_exceptions.BadRequestError( - message=( - f"Model name '{model}' matches deployments from multiple teams. " - "Specify the deployment ID directly to disambiguate." - ), - model=model, - llm_provider="", - ) - return team_deployments - - def _core_deployments_for_request( - self, model: str, request_kwargs: Mapping[str, object] - ) -> Sequence[DeploymentTypedDict]: - """The deployments `model` names for this caller, through the same alias, then team-first, then - global, then admin-across-teams resolution `_common_checks_available_deployment` applies, so - strategy selection and compression policy can never disagree with deployment selection about - which marker a name means.""" - registered_name: Final = self._core_get_model_from_alias(model=model) or model - team_id: Final = get_request_team_id(request_kwargs) - deployments: Final = self._core_get_all_deployments(model_name=registered_name, team_id=team_id) - if deployments or team_id is not None or not is_proxy_admin_request(request_kwargs): - return deployments - return self._team_deployments_across_teams(registered_name) - - @staticmethod - def _is_strategy_marker_deployment(deployment: Mapping[str, object]) -> bool: - litellm_params: Final = deployment.get("litellm_params") - if not isinstance(litellm_params, Mapping): - return False - deployment_model: Final = litellm_params.get("model") - return isinstance(deployment_model, str) and classify_strategy_router_model(deployment_model) is not None - - def _common_checks_available_deployment( - self, - model: str, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - request_kwargs: Mapping[str, object] | None = None, - ) -> tuple[str, list | dict]: - """ - Common checks for 'get_available_deployment' across sync + async call. - - If 'healthy_deployments' returned is None, this means the user chose a specific deployment - - Returns - - str, the litellm model name - - List, if multiple models chosen - - Dict, if specific model chosen - """ - - request_team_id: Final = get_request_team_id(request_kwargs) - # check if aliases set on litellm model alias map - if specific_deployment is True: - return model, self._drop_strategy_markers( - model, - self._filter_reserved_deployments( - model=model, - healthy_deployments=self._get_deployment_by_litellm_model(model=model), - request_team_id=request_team_id, - ), - ) - elif model not in self._core_model_names and self._core_has_model_id(model): - deployment: Final = self._core_get_deployment(model_id=model) - if deployment is not None: - deployment_model: Final = deployment.litellm_params.model - return deployment_model, cast( # cast-ok: contract requires a plain dict for a single deployment - dict, - self._filter_reserved_deployments( - model=deployment_model, - healthy_deployments=( - cast( # cast-ok: model_dump of a router deployment - DeploymentTypedDict, deployment.model_dump(exclude_none=True) - ), - ), - request_team_id=request_team_id, - )[0], - ) - raise ValueError( - f"LiteLLM Router: Trying to call specific deployment, but Model ID :{model} does not exist in Model ID map" - ) - - _model_from_alias: Final = self._core_get_model_from_alias(model=model) - if _model_from_alias is not None: - model = _model_from_alias - - _routing_group_deployments: Final = self._get_routing_group_deployments(model=model, team_id=request_team_id) - if _routing_group_deployments is None: - early: Final = self._try_early_resolve_deployments_for_model_not_in_names( - model=model, - request_team_id=request_team_id, - include_team_models=is_proxy_admin_request(request_kwargs), - ) - if early is not None: - if not isinstance(early[1], list): - return early[0], cast( # cast-ok: contract requires a plain dict for a single deployment - dict, - self._filter_reserved_deployments( - model=early[0], - healthy_deployments=( - cast( # cast-ok: early resolve returns a router deployment - DeploymentTypedDict, early[1] - ), - ), - request_team_id=request_team_id, - )[0], - ) - return early[0], self._drop_strategy_markers( - early[0], - self._filter_reserved_deployments( - model=early[0], - healthy_deployments=early[1], - request_team_id=request_team_id, - ), - ) - - ## get healthy deployments - ### get all deployments - healthy_deployments = ( - _routing_group_deployments - if _routing_group_deployments is not None - else self._core_get_all_deployments(model_name=model, team_id=request_team_id) - ) - _pre_model_access_group_filter_len: Final = len(healthy_deployments) - healthy_deployments = self._filter_reserved_deployments( - model=model, - healthy_deployments=self._filter_deployments_by_model_access_groups( - model=model, - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - request_team_id=request_team_id, - ), - request_team_id=request_team_id, - ) - _access_group_filter_emptied_candidates = ( - _pre_model_access_group_filter_len > 0 and len(healthy_deployments) == 0 - ) - - if len(healthy_deployments) == 0: - # check if the user sent in a deployment name instead - # Do not fall back when access-group filtering removed every candidate; - # _get_deployment_by_litellm_model does not re-apply that filter. - if _pre_model_access_group_filter_len == 0: - _litellm_model_deployments: Final = self._get_deployment_by_litellm_model(model=model) - healthy_deployments = self._filter_reserved_deployments( - model=model, - healthy_deployments=self._filter_deployments_by_model_access_groups( - model=model, - healthy_deployments=_litellm_model_deployments, - request_kwargs=request_kwargs, - request_team_id=request_team_id, - ), - request_team_id=request_team_id, - ) - # If the litellm-model lookup produced candidates that access-group - # filtering then removed, treat this the same as the by-name path - # being emptied: prevent default-model fallback from bypassing the - # restriction (the fallback model may have no access_groups and - # would short-circuit the filter). - if len(_litellm_model_deployments) > 0 and len(healthy_deployments) == 0: - _access_group_filter_emptied_candidates = True - - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("initial list of deployments: %s", healthy_deployments) - - if len(healthy_deployments) == 0: - # Check for default fallbacks if no deployments are found for the requested model - # Do not fall back to another model when access-group filtering removed every - # candidate for the requested name: re-filtering the fallback model can be a - # no-op when it has no access_groups, incorrectly serving a different model. - if self._has_default_fallbacks() and not _access_group_filter_emptied_candidates: - fallback_model: Final = self._get_first_default_fallback() - if fallback_model: - verbose_router_logger.info( - "Model '%s' not found. Attempting to use default fallback model '%s'.", model, fallback_model - ) - # Re-assign model to the fallback and try to get deployments again - model = fallback_model - healthy_deployments = self._core_get_all_deployments(model_name=model, team_id=request_team_id) - healthy_deployments = self._filter_reserved_deployments( - model=model, - healthy_deployments=self._filter_deployments_by_model_access_groups( - model=model, - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - request_team_id=request_team_id, - ), - request_team_id=request_team_id, - ) - - # If still no deployments after checking for fallbacks, raise an error - if len(healthy_deployments) == 0: - raise litellm_exceptions.BadRequestError( - message=f"You passed in model={model}. {RouterErrors.no_healthy_deployments.value}", - model=model, - llm_provider="", - ) - - if self._host.dynamic().model_alias_map and model in self._host.dynamic().model_alias_map: - model = self._host.dynamic().model_alias_map[ - model - ] # update the model to the actual value if an alias has been passed in - - return model, self._drop_strategy_markers(model, healthy_deployments) - - def _drop_strategy_markers( - self, model: str, deployments: Sequence[DeploymentTypedDict] - ) -> list[DeploymentTypedDict]: - """A strategy marker is never a callable deployment, whichever resolution arm produced it.""" - selectable: Final = [d for d in deployments if not self._is_strategy_marker_deployment(d)] - if deployments and not selectable: - raise litellm_exceptions.BadRequestError( - message=f"You passed in model={model}. {RouterErrors.only_strategy_marker_deployments.value}", - model=model, - llm_provider="", - ) - eligible: Final = eligible_deployments(model, selectable) - if isinstance(eligible, tuple): - return list(eligible) - raise eligible - - def _filter_reserved_deployments( - self, - model: str, - healthy_deployments: Sequence[DeploymentTypedDict], - request_team_id: str | None, - ) -> tuple[DeploymentTypedDict, ...]: - result: Final = filter_reserved_deployments( - self._drop_strategy_markers(model, healthy_deployments), request_team_id, now=datetime.now(timezone.utc) - ) - if result.blocking_window is not None and len(result.deployments) == 0: - raise litellm_exceptions.BadRequestError( - message=f"Deployment {model} is reserved for another team until " - f"{result.blocking_window.end:%H:%M} {result.blocking_window.timezone}", - model=model, - llm_provider="", - ) - shared: Final = filter_ptu_shared_deployments(result.deployments, request_team_id) - if shared.withheld and len(shared.deployments) == 0: - raise litellm_exceptions.BadRequestError( - message=f"Deployment {model} is reserved for the teams holding a PTU share of it", - model=model, - llm_provider="", - ) - return shared.deployments - - def _filter_deployments_by_model_access_groups( - self, - model: str, - healthy_deployments: list[DeploymentTypedDict], - request_kwargs: Mapping[str, object] | None, - request_team_id: str | None, - ) -> list[DeploymentTypedDict]: - """ - Restrict candidate deployments to caller-authorized model access groups. - - This is only applied when: - - request metadata includes `user_api_key_auth`, and - - caller permissions for this model are access-group-only - (no explicit model, wildcard, or all-proxy grants). - """ - if not healthy_deployments or request_kwargs is None: - return healthy_deployments - - user_api_key_auth: Final = request_user_api_key_auth(request_kwargs) - if user_api_key_auth is None: - return healthy_deployments - - object_models: Final = set(getattr(user_api_key_auth, "models", []) or []) - object_team_models: Final = set(getattr(user_api_key_auth, "team_models", []) or []) - allowed_models: Final = object_models | object_team_models - if not allowed_models: - return healthy_deployments - - # If caller has direct model/wildcard/all-proxy access, do not constrain - # deployment choice by access group. - if model in allowed_models or "*" in allowed_models or "all-proxy-models" in allowed_models: - return healthy_deployments - - access_groups_for_model: Final = self._core_get_model_access_groups(model_name=model, team_id=request_team_id) - if len(access_groups_for_model) == 0: - return healthy_deployments - - allowed_access_groups: Final = set(access_groups_for_model.keys()) & allowed_models - if not allowed_access_groups: - # No overlap means this request was not authorized via model access - # group membership for this model, so do not force group filtering. - return healthy_deployments - - filtered_deployments: Final = [] - for deployment in healthy_deployments: - deployment_model_info = deployment.get("model_info") or {} - deployment_access_groups = set(deployment_model_info.get("access_groups", []) or []) - if deployment_access_groups & allowed_access_groups: - filtered_deployments.append(deployment) - - return filtered_deployments - - async def _core_async_get_healthy_deployments( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ) -> list[dict] | dict: - """ - Get the healthy deployments for a model. - - Returns: - - List[Dict], if multiple models chosen - *OR* - - Dict, if specific model chosen - """ - - model, healthy_deployments = self._common_checks_available_deployment( - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - - # IF TEAM ID SPECIFIED ON MODEL, AND REQUEST CONTAINS USER_API_KEY_TEAM_ID, FILTER OUT MODELS THAT ARE NOT IN THE TEAM - ## THIS PREVENTS WRITING FILES OF OTHER TEAMS TO MODELS THAT ARE TEAM-ONLY MODELS - by_id: Final = healthy_deployments if isinstance(healthy_deployments, dict) else None # pyright: ignore[reportUnknownVariableType] # the legacy deployment dict - healthy_deployments = filter_team_based_models( - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - ) - if by_id is not None and not healthy_deployments: - # Another team's deployment asked for by id: refused under its group's name. - group_name: Final = str(by_id.get("model_name")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above - model = group_name # rebind-ok: the error names the group - - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("healthy_deployments after team filter: %s", healthy_deployments) - - healthy_deployments = filter_web_search_deployments( - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - ) - - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("healthy_deployments after web search filter: %s", healthy_deployments) - - healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows - self._without_deployments_that_cannot_serve_route( - healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type - model, - ) - ) - - if isinstance(healthy_deployments, dict): - if (healthy_deployments.get("model_info") or {}).get("blocked") is True: - raise litellm_exceptions.ServiceUnavailableError( - message=f"Model '{model}' is currently paused and cannot accept requests.", - model=model, - llm_provider="", - ) - return healthy_deployments - - request_deployments: Final = cast( # cast-ok: router deployments are plain dicts - Sequence[DeploymentTypedDict], healthy_deployments - ) - - # Health-check-based filtering (before cooldown) - healthy_deployments = ( # rebind-ok: filtered candidates enter health checks - await self._async_filter_health_check_unhealthy_deployments( - healthy_deployments=healthy_deployments, - parent_otel_span=parent_otel_span, - health_check_probe=health_check_probe, - ) - ) - - routing_read_batch: Final = RoutingReadBatch.active() - cooldown_deployments: Final = ( - await async_get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - if routing_read_batch is None - else await routing_read_batch.async_get_cooldown_deployments( - ctx=self._core_ctx, - healthy_deployments=healthy_deployments, - parent_otel_span=parent_otel_span, - ) - ) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("cooldown deployments: %s", cooldown_deployments) - _pre_cooldown_deployments: Final = healthy_deployments - healthy_deployments = self._filter_cooldown_deployments( - healthy_deployments=healthy_deployments, - cooldown_deployments=cooldown_deployments, - ) - # Safety net: only bypass cooldown filter when health-check routing is - # driving cooldown (i.e. allowed_fails_policy is set). Without a policy, - # cooldowns are from real request failures and must not be bypassed. - if ( - not healthy_deployments - and self._core_enable_health_check_routing - and self._core_allowed_fails_policy is not None - ): - verbose_router_logger.warning( - "All deployments in cooldown via health-check routing, bypassing cooldown filter" - ) - healthy_deployments = _pre_cooldown_deployments - - healthy_deployments = self._filter_blocked_deployments(healthy_deployments) - - healthy_deployments = await self.async_callback_filter_deployments( - model=model, - healthy_deployments=healthy_deployments, - messages=(cast(list[AllMessageValues], messages) if messages is not None else None), - request_kwargs=request_kwargs, - parent_otel_span=parent_otel_span, - ) - - if self._core_enable_pre_call_checks and (messages is not None or input is not None): - deployments_to_check: Final = healthy_deployments - healthy_deployments = self._pre_call_checks( - model=model, - healthy_deployments=deployments_to_check, - messages=messages, - input=input, - request_kwargs=request_kwargs, - input_token_count=await self._acount_pre_call_check_tokens( - model=model, - healthy_deployments=deployments_to_check, - messages=messages, - input=input, - request_kwargs=request_kwargs, - ), - skip_inline_token_count=True, - request_deployments=request_deployments, - ) - # check if user wants to do tag based routing - healthy_deployments = await get_deployments_for_tag( - ctx=self._core_ctx, - model=model, - request_kwargs=request_kwargs, - healthy_deployments=healthy_deployments, - metadata_variable_name=self._get_metadata_variable_name_from_kwargs(request_kwargs), - request_deployments=request_deployments, - ) - - # narrow to whatever `self._core_routing_plugins` left in candidate_models - healthy_deployments = self._filter_by_routing_plugin_candidates( - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - ) - - ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) - _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = self._without_decisions_deployments_refusing_media( - self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts - model, - messages, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ), - model, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ) - embedding_batch_admitting: Final = deployments_admitting_embedding_batch( - batch_candidates, - _OPAQUE_VALUE_ADAPTER.validate_python(input), - _OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs), - ) - embedding_batch_routed: Final = ( - embedding_batch_admitting - if embedding_batch_admitting is not None - else deployments_when_none_admits_embedding_batch( - batch_candidates, - cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts - await self._async_routable_cooled_deployment_ids( - model=model, - request_deployments=request_deployments, - cooldown_deployments=cooldown_deployments, - request_kwargs=cast(Mapping[str, object], request_kwargs), # cast-ok: str-keyed request kwargs - has_request_input=True, - ), - _OPAQUE_VALUE_ADAPTER.validate_python(input), - _OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs), - ) - ) - healthy_deployments = litellm_utils.get_order_filtered_deployments( - list( - self._narrow_to_video_deployment( - without_deployments_that_cannot_batch( - embedding_batch_routed, caller_deployments=lambda: batch_candidates - ) - ) - ), - target_order=_target_order, - ) - - ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in - ## this request via weighted-failover. Always honored, regardless of the - ## router-level flag, so a stale exclusion key on kwargs cannot escape. - _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm_utils.get_excluded_filtered_deployments( - cast(list[dict], healthy_deployments), - excluded_deployment_ids=_excluded_deployment_ids, - ) - - ## RETRY SKIP ## -> drop deployments that already refused this request with a - ## non-retryable status, unless that leaves nothing, so the caller still gets - ## the provider's own error instead of a no-deployments error. - _retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids( - request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None - ) - healthy_deployments = ( - litellm_utils.get_excluded_filtered_deployments( - healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids - ) - or healthy_deployments - ) - - if len(healthy_deployments) == 0: - exception: Final = await async_raise_no_deployment_exception( - ctx=self._core_ctx, - model=model, - parent_otel_span=parent_otel_span, - route_model_ids=await self._async_routable_cooled_deployment_ids( - model=model, - request_deployments=request_deployments, - cooldown_deployments=cooldown_deployments, - request_kwargs=cast(Mapping[str, object], request_kwargs), # cast-ok: str-keyed request kwargs - has_request_input=messages is not None or input is not None, - ), - ) - raise exception - - return healthy_deployments - - @staticmethod - def _pop_effort_from_nested_carrier(request_kwargs: dict[str, object], carrier: str) -> None: - nested: Final = request_kwargs.get(carrier) - if not isinstance(nested, Mapping): - return - sanitized: Final = {key: value for key, value in nested.items() if key != "effort"} - if sanitized: - request_kwargs[carrier] = sanitized # rebind-ok: copy-on-write, so a shared nested carrier is never edited - else: - request_kwargs.pop(carrier, None) - - @staticmethod - def _tier_ceiling_under_the_surface_name( - tier_litellm_params: Mapping[str, object], responses_call: bool - ) -> Mapping[str, object]: - """``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are one - ceiling under three names, and each surface reads exactly one of them: the - Responses bridge builds its internal ``max_tokens`` from ``max_output_tokens`` - and would overwrite the tier's, chat and /v1/messages never read - ``max_output_tokens``, and litellm already renames ``max_tokens`` to - ``max_completion_tokens`` for the OpenAI models that require it. Collapse - whatever the tier carries onto the surface's own name, preferring a value the - operator already wrote under that name.""" - surface_key: Final = "max_output_tokens" if responses_call else "max_tokens" - carried: Final = tuple( - key - for key in (surface_key, "max_tokens", "max_completion_tokens", "max_output_tokens") - if key in tier_litellm_params - ) - if not carried: - return tier_litellm_params - return MappingProxyType( - { - **{k: v for k, v in tier_litellm_params.items() if k not in OUTPUT_TOKEN_CEILING_PARAMS}, - surface_key: tier_litellm_params[carried[0]], - } - ) - - def _pin_tier_params_onto_request( - self, - model: str, - tier_litellm_params: Mapping[str, object] | None, - request_kwargs: dict, - responses_call: bool, - ) -> bool: - """Apply per-tier litellm_params and report whether they pinned an output ceiling.""" - if not tier_litellm_params: - return False - reject_api_key_secret_reference(tier_litellm_params) - accepted_tier_params: Final = self._tier_params_the_target_accepts(model, tier_litellm_params, request_kwargs) - surface_tier_params: Final = self._tier_ceiling_under_the_surface_name( - accepted_tier_params, responses_call=responses_call - ) - self._drop_client_carriers_a_tier_pin_supersedes(request_kwargs, surface_tier_params) - request_kwargs.update(surface_tier_params) - return not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(surface_tier_params) - - @staticmethod - def _restore_client_output_ceiling(request_kwargs: MutableMapping[str, object]) -> None: - """Restore caller intent before routing, so neither a previous tier pin nor alias defaults - become the next attempt's budget. The proxy strips caller-written stamps at ingestion.""" - callers_ceiling: Final = stamped_client_ceiling(request_kwargs) - if callers_ceiling is None: - return - for carrier in OUTPUT_TOKEN_CEILING_PARAMS: - request_kwargs.pop(carrier, None) - request_kwargs.update(callers_ceiling) - - @staticmethod - def _stamp_client_output_ceiling( - request_kwargs: dict[str, object], # mutable-ok: preserve request-owned metadata across routing attempts - ) -> None: - _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket.setdefault( - CLIENT_OUTPUT_CEILING_METADATA_KEY, - {carrier: request_kwargs[carrier] for carrier in OUTPUT_TOKEN_CEILING_PARAMS if carrier in request_kwargs}, - ) - - @staticmethod - def _drop_client_carriers_a_tier_pin_supersedes( - request_kwargs: dict[str, object], - tier_litellm_params: Mapping[str, object], - ) -> None: - """Tier litellm_params are deliberate operator overrides, but provider - translations let a caller-supplied carrier of the same setting - (``thinking``, ``output_config.effort``, ``reasoning.effort``) outrank - the ``reasoning_effort`` alias, so a pinned effort only reaches the wire - if the client's other encodings are removed before the merge. Non-effort - fields a carrier also holds (``output_config.format``, - ``reasoning.summary``) are kept. An output ceiling has the same shape: - ``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are - one setting under three names, and a provider handed two of them either - rejects the request or picks one by iteration order.""" - if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params): - PythonCore._stamp_client_output_ceiling(request_kwargs) - for carrier in OUTPUT_TOKEN_CEILING_PARAMS: - request_kwargs.pop(carrier, None) - if "reasoning_effort" not in tier_litellm_params: - return - request_kwargs.pop("thinking", None) - PythonCore._pop_effort_from_nested_carrier(request_kwargs, "output_config") - PythonCore._pop_effort_from_nested_carrier(request_kwargs, "reasoning") - - async def _async_get_available_deployment_unresolved( # noqa: ANN202 # legacy selector result remains inferred - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - ): - """ - Async implementation of 'get_available_deployments'. - - Allows all cache calls to be made async => 10x perf impact (8rps -> 100 rps). - """ - with phase_span(f"route {model}"): - return await self._async_get_available_deployment( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - reserve=pick_reserves(default=True), - ) - - async def _async_get_available_deployment( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None, - input: str | list | None, - specific_deployment: bool | None, - reserve: bool, - ): - if ( - self._core_routing_strategy != "usage-based-routing-v2" - and self._core_routing_strategy != "simple-shuffle" - and self._core_routing_strategy != "cost-based-routing" - and self._core_routing_strategy != "latency-based-routing" - and self._core_routing_strategy != "least-busy" - ): # prevent regressions for other routing strategies, that don't have async get available deployments implemented. - return self._get_available_deployment_unresolved( - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - try: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(request_kwargs) - - ######################################################### - # Execute Pre-Routing Hooks - # this hook can modify the model, messages before the routing decision is made - ######################################################### - responses_call: Final = input is not None and messages is None - pre_routing_hook_response: Final = await self._core_async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - if pre_routing_hook_response is not None: - model = pre_routing_hook_response.model - messages = pre_routing_hook_response.messages - record_pre_routing_selection(request_kwargs, model) - self._pin_tier_params_onto_request( - model=model, - tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, - request_kwargs=request_kwargs, - responses_call=responses_call, - ) - ######################################################### - - # Resolve the strategy and logger AFTER the pre-routing hook, since - # the hook can replace `model` and routing-group lookup must key - # off the final model name. - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - pick_attributes: Final = _deployment_pick_attributes( - model, _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) - ) - routing_read_batch: Final = RoutingReadBatch.for_strategy(strategy, strategy_selector) - - with RoutingReadBatch.scoped(routing_read_batch): - healthy_deployments: Final = await self._core_async_get_healthy_deployments( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - parent_otel_span=parent_otel_span, - ) - if isinstance(healthy_deployments, dict): - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, healthy_deployments, parent_otel_span - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return healthy_deployments - - # When encrypted content affinity pins to a specific deployment, - if request_kwargs.get("_encrypted_content_affinity_pinned") and len(healthy_deployments) == 1: - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, healthy_deployments[0], parent_otel_span - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return healthy_deployments[0] - - start_time: Final = time.time() - if strategy == "simple-shuffle": - shuffled: Final = simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=healthy_deployments, - model=model, - request_kwargs=request_kwargs, - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return shuffled - with PrefetchedUsage.scoped( - routing_read_batch.prefetched_usage if routing_read_batch is not None else None - ): - deployment: Final = await self._select_deployment_async( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - reserve=reserve and request_kwargs not in (None, {}), - ) - if deployment is None: - exception: Final = await async_raise_no_deployment_exception( - ctx=self._core_ctx, - model=model, - parent_otel_span=parent_otel_span, - ) - raise exception - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, deployment, parent_otel_span - ) - verbose_router_logger.info( - "get_available_deployment for model: %s, Selected deployment: %s for model: %s", - model, - self.print_deployment(deployment), - model, - ) - - end_time: Final = time.time() - _duration: Final = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.ROUTER, - duration=_duration, - call_type=".async_get_available_deployments", - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) - ) - - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return deployment - except Exception as e: - mark_deployment_selection_failure(e) - traceback_exception: Final = traceback.format_exc() - # if router rejects call -> log to langfuse/otel/etc. - if request_kwargs is not None: - logging_obj: Final = request_kwargs.get("litellm_logging_obj", None) - - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback_exception, - prefer_async_handlers=True, - ) - ) - raise e - - _core_async_get_available_deployment = ( # pyright: ignore[reportUnknownVariableType] # legacy selector return - async_resolve_pods_after( - _async_get_available_deployment_unresolved # pyright: ignore[reportUnknownArgumentType] # legacy type - ) - ) - - async def async_get_available_deployment_for_pass_through( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - ): - """ - Async version of get_available_deployment_for_pass_through - - Only returns deployments configured with use_in_pass_through=True - """ - with phase_span(f"route {model}"): - return await self._async_get_available_deployment_for_pass_through( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - - async def _async_get_available_deployment_for_pass_through( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None, - input: str | list | None, - specific_deployment: bool | None, - ): - try: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(request_kwargs) - - # 1. Execute pre-routing hook - responses_call: Final = input is not None and messages is None - pre_routing_hook_response: Final = await self._core_async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - if pre_routing_hook_response is not None: - model = pre_routing_hook_response.model - messages = pre_routing_hook_response.messages - record_pre_routing_selection(request_kwargs, model) - self._pin_tier_params_onto_request( - model=model, - tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, - request_kwargs=request_kwargs, - responses_call=responses_call, - ) - - # 2. Get healthy deployments - healthy_deployments: Final = await self._core_async_get_healthy_deployments( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - parent_otel_span=parent_otel_span, - ) - - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - - # 3. If specific deployment returned, verify if it supports pass-through - if isinstance(healthy_deployments, dict): - if (healthy_deployments.get("model_info") or {}).get("blocked") is True: - raise litellm_exceptions.ServiceUnavailableError( - message=f"Model '{model}' is currently paused and cannot accept requests.", - model=model, - llm_provider="", - ) - litellm_params: Final = healthy_deployments.get("litellm_params", {}) - if litellm_params.get("use_in_pass_through"): - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, healthy_deployments, parent_otel_span - ) - return healthy_deployments - else: - raise litellm_exceptions.BadRequestError( - message=f"Deployment {healthy_deployments.get('model_info', {}).get('id')} does not support pass-through endpoint (use_in_pass_through=False)", - model=model, - llm_provider="", - ) - - # 4. Filter deployments that support pass-through - pass_through_deployments = self._filter_pass_through_deployments(healthy_deployments=healthy_deployments) - - if len(pass_through_deployments) == 0: - raise litellm_exceptions.BadRequestError( - message=f"Model {model} has no deployments configured with use_in_pass_through=True. Please add use_in_pass_through: true to the deployment configuration", - model=model, - llm_provider="", - ) - - # 5. Apply load balancing strategy - start_time: Final = time.perf_counter() - if strategy == "simple-shuffle": - return simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=pass_through_deployments, - model=model, - request_kwargs=request_kwargs, - ) - deployment: Final = await self._select_deployment_async( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=pass_through_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), - ) - - if deployment is None: - exception: Final = await async_raise_no_deployment_exception( - ctx=self._core_ctx, - model=model, - parent_otel_span=parent_otel_span, - ) - raise exception - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, deployment, parent_otel_span - ) - - verbose_router_logger.info( - "async_get_available_deployment_for_pass_through model: %s, selected deployment: %s", - model, - self.print_deployment(deployment), - ) - - end_time: Final = time.perf_counter() - _duration: Final = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.ROUTER, - duration=_duration, - call_type=".async_get_available_deployments", - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) - ) - - return deployment - except Exception as e: - mark_deployment_selection_failure(e) - traceback_exception: Final = traceback.format_exc() - if request_kwargs is not None: - logging_obj: Final = request_kwargs.get("litellm_logging_obj", None) - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback_exception, - prefer_async_handlers=True, - ) - ) - raise e - - async def _run_routing_plugins( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, object]] | None, - ) -> RoutingContext: - """ - Build a RoutingContext for `model`, run it through `self._core_routing_plugins` - in order, then stash the narrowed candidate list and accumulated signals - onto `request_kwargs["metadata"]` so `_filter_by_routing_plugin_candidates` - (called later, during healthy-deployment filtering) can consume them. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - resolve_structured_messages, - ) - - candidate_models: Final = list(self._core_resolved_litellm_models(model)) - - metadata_key: Final = self._get_metadata_variable_name_from_kwargs(request_kwargs) - metadata: Final = request_kwargs.setdefault(metadata_key, {}) - - context = RoutingContext( - raw_messages=messages or [], - structured_messages=resolve_structured_messages(messages=messages, request_kwargs=request_kwargs) or [], - candidate_models=candidate_models, - metadata=metadata, - ) - - for plugin in self._core_routing_plugins: - context = await plugin.run(context) - - metadata["routing_plugin_signals"] = context.signals - if len(context.candidate_models) < len(candidate_models): - metadata["_routing_plugin_candidate_models"] = context.candidate_models - - return context - - def _filter_by_routing_plugin_candidates( - self, - healthy_deployments: list[dict] | dict, - request_kwargs: dict, - ) -> list[dict] | dict: - """ - Narrow `healthy_deployments` to whatever `self._core_routing_plugins` left in - `context.candidate_models`. Raises rather than silently falling back to - the unfiltered pool -- a plugin narrowing to nothing is a policy decision - (e.g. no model this tenant's budget allows), not something to bypass. - """ - if not self._core_routing_plugins or not isinstance(healthy_deployments, list): - return healthy_deployments - - metadata_key: Final = self._get_metadata_variable_name_from_kwargs(request_kwargs) - candidate_models: Final = (request_kwargs.get(metadata_key) or {}).get("_routing_plugin_candidate_models") - # `is None` (not falsy-check): a plugin narrowing to an empty list must - # still hit the "no deployments left" raise below, not be treated the - # same as "no plugin ever set this key". - if candidate_models is None: - return healthy_deployments - - candidate_set: Final = set(candidate_models) - filtered: Final = [d for d in healthy_deployments if d.get("litellm_params", {}).get("model") in candidate_set] - - if not filtered: - raise ValueError(f"No deployments left after routing-plugin filtering. candidate_models={candidate_models}") - - return filtered - - def _select_pre_routing_strategy( - self, model: str, request_kwargs: Mapping[str, object] - ) -> "TaggedPreRoutingStrategy[PreRoutingStrategy] | None": - """ - Resolve the pre-routing strategy for `model`, disambiguating deployments - that share a `model_name` by matching the request's tags against each - registered strategy's tags before falling back to the first registered. - Returns the tagged registry entry so the caller can tell whether the - request's tags were what selected it, and can locate the marker - deployment the strategy was registered from via its (model_name, tags) - pair. - - The registries are keyed by the marker deployment's own `model_name`, which - for a team-scoped router is the internal `model_name_{team}_{uuid}` while - the caller sends the team's public name. So the names looked up are the - `model_name`s of whatever deployments this caller's request resolves `model` - to, and `model` itself when it resolves to none. - - With tag filtering enabled, router-wide or by the request's - enable_tag_filtering (which the proxy sets from key/team - router_settings), strategies that all carry real tags matching none of - the request's do not capture it when the name also has plain - deployments: returning None hands the request to ordinary tag-aware - deployment selection. - """ - registries: Final = ( - self._core_auto_routers, - self._core_complexity_routers, - self._core_adaptive_routers, - self._core_quality_routers, - ) - if not any(registries): - return None - deployments: Final = self._core_deployments_for_request(model, request_kwargs) - registered_names: Final = tuple(dict.fromkeys(str(d["model_name"]) for d in deployments)) or (model,) - candidates: Final = tuple( - tagged for registry in registries for name in registered_names for tagged in registry.get(name, []) - ) - if not candidates: - return None - - request_tags: Final = get_routing_tags_from_request_kwargs(request_kwargs) - if request_tags: - for tagged in candidates: - if tagged.tags and is_valid_deployment_tag( - list(tagged.tags), request_tags, self._core_tag_filtering_match_any - ): - return tagged - for tagged in candidates: - if "default" in tagged.tags: - return tagged - request_scoped_filtering: Final = request_kwargs.get("enable_tag_filtering") is True - if ( - (self._core_enable_tag_filtering or request_scoped_filtering) - and all(tagged.tags for tagged in candidates) - and any(not self._is_strategy_marker_deployment(d) for d in deployments) - ): - return None - return candidates[0] - - @staticmethod - def _request_header(request_kwargs: Mapping[str, object], header_name: str) -> str | None: - proxy_server_request: Final = request_kwargs.get("proxy_server_request") - if not isinstance(proxy_server_request, Mapping): - return None - headers: Final = proxy_server_request.get("headers") - if not isinstance(headers, Mapping): - return None - return next( - ( - value - for key, value in headers.items() - if isinstance(key, str) and key.lower() == header_name and isinstance(value, str) - ), - None, - ) - - def _claude_code_session_router_cache_key(self, request_kwargs: Mapping[str, object]) -> str | None: - session_id: Final = self._request_header(request_kwargs, "x-claude-code-session-id") - if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None: - return None - metadata_name: Final = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" - metadata: Final = request_kwargs.get(metadata_name) - if not isinstance(metadata, Mapping): - return None - caller_scope: Final = metadata.get("user_api_key_hash") - if not isinstance(caller_scope, str) or not caller_scope: - return None - return f"claude_code_session_router:v1:{caller_scope}:{session_id}" - - @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: - try: - await self._claude_code_session_router_cache.async_delete_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request - verbose_router_logger.warning( - "Failed to delete Claude Code session router binding; the binding may remain until its TTL expires: %s", - e, - ) - - @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: - session_cache: Final = self._claude_code_session_router_cache - try: - if session_cache.redis_cache is None: - return await session_cache.async_get_cache(key=cache_key) - return await session_cache.redis_cache.async_get_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis - log_redis_failure( - verbose_router_logger, - logging.WARNING, - "Failed to read Claude Code session router binding; using the requested model", - e, - ) - return None - - @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _resolve_claude_code_session_router( - self, - model: str, - registered_model_name: str, - request_kwargs: Mapping[str, object], - ) -> str: - if is_native_compaction_call(): - return registered_model_name - if not any( - ( - self._core_auto_routers, - self._core_complexity_routers, - self._core_adaptive_routers, - self._core_quality_routers, - ) - ): - return registered_model_name - cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) - if cache_key is None or not isinstance(request_kwargs, dict): - return registered_model_name - if request_kwargs.get("fallback_depth") not in (None, 0): - return registered_model_name - - agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") - if agent_id is not None: - bound_model: Final = await self._get_claude_code_session_router_binding(cache_key) - if not isinstance(bound_model, str): - return registered_model_name - bound_registered_model: Final = self._core_get_model_from_alias(model=bound_model) or bound_model - if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) - return registered_model_name - await self._claude_code_session_router_cache.async_set_cache( - key=cache_key, - value=bound_model, - ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, - ) - self._stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model) - return bound_registered_model - - if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - return registered_model_name - await self._claude_code_session_router_cache.async_set_cache( - key=cache_key, - value=model, - ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, - ) - return registered_model_name - - async def _core_async_pre_routing_hook( - self, - model: str, - request_kwargs: dict[str, object], - messages: list[dict[str, Any]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - ) -> PreRoutingHookResponse | None: - """ - This hook is called before the routing decision is made. - - Used for the litellm auto-router to modify the request before the routing decision is made. - - `model` is whatever the caller asked for, which may be a `model_group_alias` key or a team's - public model name, while the strategy registries and the marker deployment are keyed by the - marker's own `model_name`, so every lookup below resolves the alias first and the team name - through the deployment path. Only the lookups: the caller-facing name stays the alias, since - spend metadata is stamped before routing and the response carries the tier group the - strategy picked. - """ - self._restore_client_output_ceiling(request_kwargs) - self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) - requested_registered_model_name: Final = self._core_get_model_from_alias(model=model) or model - registered_model_name: Final = await self._resolve_claude_code_session_router( - model=model, - registered_model_name=requested_registered_model_name, - request_kwargs=request_kwargs, - ) - - ######################################################### - # Run the routing-plugin pipeline, if any plugins are configured. - # Plugins narrow the candidate deployment pool (consumed later by - # `_filter_by_routing_plugin_candidates`) and may attach signals for - # downstream strategies (auto-router, complexity-router, ...) to read. - ######################################################### - if self._core_routing_plugins: - await self._run_routing_plugins( - model=registered_model_name, request_kwargs=request_kwargs, messages=messages - ) - - selected_strategy: Final = self._select_pre_routing_strategy( - model=registered_model_name, request_kwargs=request_kwargs - ) - if selected_strategy is None: - await arm_compaction(request_kwargs, None) - self._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None - ) - self._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, key=CONSUMED_REQUEST_TAGS_METADATA_KEY, value=None - ) - return None - - from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter - - marker: Final = self._selected_strategy_marker_deployment( - model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs - ) - definition: Final = selected_strategy.deployment - config_id: Final = definition.model_info.id if definition else None - updated_at: Final = definition.model_info.updated_at if definition else None - router_kind: Final = classify_strategy_router_model(definition.litellm_params.model) if definition else None - provenance: Final[StandardLoggingRoutingDecision] = { - "router_model_name": definition.model_name if definition else registered_model_name, - **({"router_type": router_kind} if router_kind is not None else {}), - **({"router_config_id": config_id} if config_id is not None else {}), - **({"router_config_updated_at": updated_at.isoformat()} if updated_at is not None else {}), - **( - {"router_config_fingerprint": selected_strategy.definition_fingerprint} - if selected_strategy.definition_fingerprint is not None - else {} - ), - } - phase_attributes(routing_decision_attributes(provenance)) - - reject_recursive_compactor(registered_model_name) - await arm_compaction( - request_kwargs, - selected_strategy.strategy.config.context_compaction - if isinstance(selected_strategy.strategy, ComplexityRouter) - else None, - tuple( - dict.fromkeys( - member - for pool in selected_strategy.strategy.config.tiers.values() - for member in ((pool,) if isinstance(pool, str) else pool) - ) - ) - if isinstance(selected_strategy.strategy, ComplexityRouter) - else (), - parent_model=model, - ctx=self._core_ctx, - allow_escalation=isinstance(selected_strategy.strategy, ComplexityRouter) - and selected_strategy.strategy.config.enable_context_window_escalation, - messages=messages, - ) - - await router_host().authorize_strategy_inference( - deployment=marker, - request_kwargs=request_kwargs, - llm_router=self._core_public(), - ) - - # Same tag-aware lookup the proxy's pre-call arming used, so an alias with - # several tag-scoped markers cannot suppress under one and route under another. - compression_policy: Final = router_host().compression_policy( - llm_router=self._core_public(), - model_alias=registered_model_name, - request_kwargs=request_kwargs, - request_tags=get_routing_tags_from_request_kwargs(request_kwargs), - ) - # Shared compression already ran in the pre-call hook, so reuse it rather than - # compressing twice. Conditional on arming having actually happened: only the - # proxy arms, and on the SDK path the shortcut would skip both hops entirely. - needs_independent_routing_compression: Final = compression_policy is not None and not ( - compression_policy.is_same - and compression_policy.model is not None - and router_host().model_hop_compression_armed() - ) - routing_messages: Final = ( - await router_host().messages_for_routing( - policy=compression_policy, messages=messages, request_kwargs=request_kwargs - ) - if needs_independent_routing_compression - else None - ) - - routed: Final = await selected_strategy.strategy.async_pre_routing_hook( - model=registered_model_name, - request_kwargs=request_kwargs, - messages=routing_messages if routing_messages is not None else messages, - input=input, - specific_deployment=specific_deployment, - ) - traced_routed: Final = ( - routed.model_copy(update={"routing_decision": {**provenance, **routed.routing_decision}}) - if routed is not None and routed.routing_decision is not None - else routed - ) - # Routing-only compression must not leak into the response: the model call and - # deployment-context filtering key off this field. Compared by value, since - # pydantic rebuilds the list rather than keeping the object passed in. - pre_routing_hook_response: Final = ( - traced_routed.model_copy(update={"messages": messages}) - if traced_routed is not None and routing_messages is not None and traced_routed.messages == routing_messages - else traced_routed - ) - self._record_routing_decision( - request_kwargs=request_kwargs, - routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), - ) - self._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, - key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, - value=(pre_routing_hook_response.session_affinity_ttl_seconds if pre_routing_hook_response else None), - ) - self._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, - key=CONSUMED_REQUEST_TAGS_METADATA_KEY, - value=self._consumed_request_tags_stamp( - selected_strategy=selected_strategy, - pre_routing_hook_response=pre_routing_hook_response, - request_tags=get_routing_tags_from_request_kwargs(request_kwargs), - ), - ) - - # `model` (the alias, e.g. "smart-router") is never the deployment actually - # called - apply the router marker's own litellm_params to the request, - # since the tier/route deployment the hook selected won't have them. The - # marker entry is looked up by its `auto_router/` model prefix and the - # selected strategy's tags, never by list position: plain deployments may - # share the alias `model_name` and must not leak their params (`api_base`, - # `api_key`, ...) onto the routed call. Router-only fields (tpm, rpm, - # weight, complexity_router_config, ...) are excluded from the actual - # outbound LLM call downstream by litellm.types.utils.all_litellm_params, - # not here. Custom pricing fields ARE call params, so they must be - # excluded here: they price the alias, not the deployment the hook - # selected, and forwarding them re-registers the routed deployment at - # the alias's price (an explicit 0 makes every alias request bill $0). - # Forwarded params only fill gaps: the keys inserted here ride along on the - # request (top level, so sibling requests sharing a `metadata` dict never see - # them) until `_update_kwargs_with_deployment` drops any the selected - # deployment sets itself (its own `aws_region_name` beats the marker's). - # Per-tier `litellm_params` on the hook response are deliberate overrides - # the caller applies on top, so those keys and equivalent output-limit - # names are never tracked as alias defaults here. - marker_params: Final = ( - self._forwardable_alias_marker_params( - model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs - ) - if pre_routing_hook_response is not None - else () - ) - tier_param_keys: Final = ( - tuple(pre_routing_hook_response.litellm_params or ()) if pre_routing_hook_response is not None else () - ) - overridden_marker_keys: Final = frozenset(tier_param_keys) | ( - OUTPUT_TOKEN_CEILING_PARAMS if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_param_keys) else frozenset() - ) - newly_forwarded: Final = tuple( - (key, value) - for key, value in marker_params - if key not in request_kwargs and key not in overridden_marker_keys - ) - if pre_routing_hook_response is not None: - self._stamp_client_output_ceiling(request_kwargs) - request_kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, None) - request_kwargs.update(newly_forwarded) - if newly_forwarded: - request_kwargs.update(((_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, tuple(key for key, _ in newly_forwarded)),)) - - return pre_routing_hook_response - - def _selected_strategy_marker_deployment( - self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] - ) -> DeploymentTypedDict | None: - markers: Final = tuple( - deployment - for deployment in self._core_deployments_for_request(model, request_kwargs) - if "model" in deployment["litellm_params"] - and str(deployment["litellm_params"]["model"]).startswith(AUTO_ROUTER_MODEL_PREFIX) - ) - tag_matched: Final = tuple( - deployment - for deployment in markers - if (tuple(deployment["litellm_params"]["tags"] or ()) if "tags" in deployment["litellm_params"] else ()) - == strategy_tags - ) - return tag_matched[0] if tag_matched else (markers[0] if markers else None) - - def _forwardable_alias_marker_params( - self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] - ) -> tuple[tuple[str, object], ...]: - marker: Final = self._selected_strategy_marker_deployment( - model=model, strategy_tags=strategy_tags, request_kwargs=request_kwargs - ) - if marker is None: - return () - return tuple( - (key, value) - for key, value in marker["litellm_params"].items() - if key not in _ALIAS_PARAMS_NEVER_FORWARDED - and key not in CustomPricingLiteLLMParams.model_fields - and value is not None - ) - - @staticmethod - def _forwarded_alias_marker_keys_the_deployment_sets( - deployment: Mapping[str, object], forwarded_keys: object - ) -> tuple[str, ...]: - deployment_litellm_params: Final = deployment.get("litellm_params") - if not isinstance(deployment_litellm_params, Mapping) or not isinstance(forwarded_keys, tuple): - return () - return tuple( - key - for key in forwarded_keys - if isinstance(key, str) and PythonCore._deployment_sets_litellm_param(deployment_litellm_params, key) - ) - - @staticmethod - def _deployment_sets_litellm_param(deployment_litellm_params: Mapping[str, object], key: str) -> bool: - value: Final = deployment_litellm_params.get(key) - if value is None: - return False - field: Final = LiteLLM_Params.model_fields.get(key) - return field is None or value != field.default - - def _consumed_request_tags_stamp( - self, - selected_strategy: "TaggedPreRoutingStrategy[PreRoutingStrategy]", - pre_routing_hook_response: PreRoutingHookResponse | None, - request_tags: Sequence[str], - ) -> ConsumedRequestTagsStamp | None: - """Record which tags picked the router and which model group it rewrote to, or None. - - A request whose tags matched the selected strategy's tags has already spent those - tags on picking the router; re-applying them to the routed tier's model group would - empty the pool unless every tier deployment repeats the marker's tag. Only the - strategy's own tags are spent: the request's other tags keep constraining - deployment selection inside the routed group, and key/team policy tags are - untouched because tag filtering separately re-applies whatever - `metadata.inherited_tags` carries for the stamped group. - """ - if pre_routing_hook_response is None or not selected_strategy.tags or not request_tags: - return None - if not is_valid_deployment_tag(selected_strategy.tags, request_tags, self._core_tag_filtering_match_any): - return None - return ConsumedRequestTagsStamp(model_group=pre_routing_hook_response.model, tags=selected_strategy.tags) - - @staticmethod - def _record_routing_decision( - request_kwargs: dict, - routing_decision: StandardLoggingRoutingDecision | None, - ) -> None: - """Make the request's metadata describe THIS routing attempt, and only this one. - - Fallbacks re-enter the hook with the same `request_kwargs`, so an attempt that - picks a plain model group after an auto-router group failed must clear the - earlier decision; leaving it would attribute the first router's tier and cause - to the deployment that actually served the request. Every attempt therefore - writes or clears, never just writes. - """ - from litellm.router_strategy.auto_router.baseline_request import capture_baseline_parameters - from litellm.types.router import BaselineRouteStamp - - phase_attributes(routing_decision_attributes(routing_decision)) - baseline_model: Final = routing_decision.get("savings_baseline_model") if routing_decision else None - baseline_id: Final = routing_decision.get("savings_baseline_deployment_id") if routing_decision else None - router_name: Final = routing_decision.get("router_model_name") if routing_decision else None - caller_parameters: Final = ( - capture_baseline_parameters(request_kwargs) if router_name and baseline_model else None - ) - PythonCore._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, - key="_autorouter_baseline_route", - value=( - BaselineRouteStamp(router_name, baseline_model, baseline_id, caller_parameters) - if router_name and baseline_model and baseline_id - else None - ), - ) - PythonCore._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, - key="routing_decision", - value=( - None - if routing_decision is None - else PythonCore._redact_prompt_text_if_needed( - request_kwargs=request_kwargs, routing_decision=routing_decision - ) - ), - ) - - @staticmethod - def _stamp_or_clear_metadata_key(request_kwargs: dict, key: str, value: object | None) -> None: - """Write a proxy-internal metadata key for THIS routing attempt, or clear it. - - Fallbacks and retries re-enter the pre-routing hook with the same - `request_kwargs`, so every attempt must write or clear, never just write; - a value left behind by an earlier attempt would be attributed to this one. - `get_or_create_metadata_bucket` is the single owner of "which dict holds - proxy-internal metadata": it picks `litellm_metadata` when present (so the - value never lands in the `metadata` dict that routes like /v1/messages - forward to the provider) and replaces a non-dict value rather than silently - skipping the write. Clearing pops from BOTH buckets so a request whose - bucket resolution changed between attempts cannot resurrect a stale value. - """ - if value is None: - for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")): - if isinstance(bucket, dict): - bucket.pop(key, None) - return - - _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket[key] = value - - @staticmethod - def _redact_prompt_text_if_needed( - request_kwargs: Mapping[str, object], - routing_decision: StandardLoggingRoutingDecision, - ) -> StandardLoggingRoutingDecision: - """Drop verbatim prompt text from the record when message logging is redacted. - - An operator who turns message logging off has said prompt content must not reach - the logs, so the fields that quote the prompt (the matched keywords, and the - signals that name them) are omitted. Derived values are kept, because a tier, a - cause, a score or an escalation flag aggregates the prompt rather than - reproducing any of it, and dropping them would leave the row unexplainable for - no privacy gain. Applied here rather than in each strategy so a strategy added - later cannot bypass it. - """ - from litellm.litellm_core_utils.redact_messages import ( - should_redact_message_logging, - ) - - if not should_redact_message_logging( - { - "litellm_params": request_kwargs, - "standard_callback_dynamic_params": request_kwargs.get("standard_callback_dynamic_params"), - } - ): - return routing_decision - kept: Final = { - field: value - for field, value in routing_decision.items() - if field not in PROMPT_QUOTING_ROUTING_DECISION_FIELDS - } - return cast(StandardLoggingRoutingDecision, kept) # cast-ok: dropping optional keys preserves the type - - def _get_available_deployment_unresolved( # noqa: ANN202 # legacy selector result remains inferred - self, - model: str, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - request_kwargs: dict | None = None, - ): - """ - Returns the deployment based on routing strategy - """ - if self._core_routing_plugins: - raise ValueError( - "Router(plugins=[...]) is configured, but this call resolved to the synchronous " - "deployment-selection path, which never runs the routing-plugin pipeline. This " - "happens for sync Router methods (e.g. Router.completion()) and for async calls " - "with a routing_strategy that has no async-native selector (e.g. legacy " - "'usage-based-routing', v1). Silently skipping " - "configured plugins would let a policy plugin (e.g. a deny-all rule) be bypassed. " - "Use an async Router method with a supported routing_strategy (simple-shuffle, " - "usage-based-routing-v2, cost-based-routing, latency-based-routing, least-busy), " - "or remove `plugins` from the Router config." - ) - # users need to explicitly call a specific deployment, by setting `specific_deployment = True` as completion()/embedding() kwarg - # When this was no explicit we had several issues with fallbacks timing out - - model, healthy_deployments = self._common_checks_available_deployment( - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows - self._without_deployments_that_cannot_serve_route( - healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type - model, - ) - ) - # As the async pick: a team-owned deployment serves only its own team's callers, also when - # asked for by id (then refused under its group's name). - by_id: Final = healthy_deployments if isinstance(healthy_deployments, dict) else None # pyright: ignore[reportUnknownVariableType] # the legacy deployment dict - healthy_deployments = filter_team_based_models( # rebind-ok: the team filter narrows # pyright: ignore[reportUnknownVariableType] # legacy selector type - healthy_deployments=healthy_deployments, # pyright: ignore[reportArgumentType] # legacy selector type - request_kwargs=request_kwargs, - ) - if by_id is not None and not healthy_deployments: - group_name: Final = str(by_id.get("model_name")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above - model = group_name # rebind-ok: the error names the group - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - routing_kwargs: Final = _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) - pick_attributes: Final = _deployment_pick_attributes(model, routing_kwargs) - - if isinstance(healthy_deployments, dict): - if (healthy_deployments.get("model_info") or {}).get("blocked") is True: - raise litellm_exceptions.ServiceUnavailableError( - message=f"Model '{model}' is currently paused and cannot accept requests.", - model=model, - llm_provider="", - ) - self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return healthy_deployments - - parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs) - request_deployments: Final = cast( # cast-ok: router deployments are plain dicts - Sequence[DeploymentTypedDict], healthy_deployments - ) - - # Health-check-based filtering (before cooldown) - healthy_deployments = self._filter_health_check_unhealthy_deployments( - healthy_deployments=healthy_deployments, - parent_otel_span=parent_otel_span, - ) - - cooldown_deployments: Final = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - _pre_cooldown_deployments: Final = healthy_deployments - healthy_deployments = self._filter_cooldown_deployments( - healthy_deployments=healthy_deployments, - cooldown_deployments=cooldown_deployments, - ) - if ( - not healthy_deployments - and self._core_enable_health_check_routing - and self._core_allowed_fails_policy is not None - ): - verbose_router_logger.warning( - "All deployments in cooldown via health-check routing, bypassing cooldown filter" - ) - healthy_deployments = _pre_cooldown_deployments - - healthy_deployments = continuation_deployments( - self._filter_blocked_deployments(healthy_deployments), routing_kwargs - ) - - # filter pre-call checks - if self._core_enable_pre_call_checks and (messages is not None or input is not None): - healthy_deployments = self._pre_call_checks( - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - request_deployments=request_deployments, - ) - - ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) - _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = self._without_decisions_deployments_refusing_media( - self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts - model, - messages, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ), - model, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ) - embedding_batch_admitting: Final = deployments_admitting_embedding_batch( - batch_candidates, - _OPAQUE_VALUE_ADAPTER.validate_python(input), - routing_kwargs, - ) - embedding_batch_routed: Final = ( - embedding_batch_admitting - if embedding_batch_admitting is not None - else deployments_when_none_admits_embedding_batch( - batch_candidates, - cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts - self._deployment_ids( - self._region_routable_deployments( - request_deployments, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: str-keyed request kwargs - True, - ) - ) - & frozenset(cooldown_deployments), - _OPAQUE_VALUE_ADAPTER.validate_python(input), - routing_kwargs, - ) - ) - healthy_deployments = litellm_utils.get_order_filtered_deployments( - list( - self._narrow_to_video_deployment( - without_deployments_that_cannot_batch( - embedding_batch_routed, - caller_deployments=lambda: deployments_the_caller_may_use(batch_candidates, routing_kwargs), - ) - ) - ), - target_order=_target_order, - ) - - ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in - ## this request via weighted-failover. See async counterpart in - ## async_get_healthy_deployments for details. - _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm_utils.get_excluded_filtered_deployments( - healthy_deployments, - excluded_deployment_ids=_excluded_deployment_ids, - ) - - ## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments. - _retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids( - request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None - ) - healthy_deployments = ( - litellm_utils.get_excluded_filtered_deployments( - healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids - ) - or healthy_deployments - ) - - if len(healthy_deployments) == 0: - model_ids = self._core_get_model_ids(model_name=model) - _cooldown_time = self._core_cooldown_cache.get_min_cooldown( - model_ids=model_ids, parent_otel_span=parent_otel_span - ) - _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - raise RouterRateLimitError( - model=model, - cooldown_time=_cooldown_time, - enable_pre_call_checks=self._core_enable_pre_call_checks, - cooldown_list=_cooldown_list, - model_ids=model_ids, - cooldown_deployments=get_cooldown_deployment_info( - self._core_ctx, - self._deployment_ids( - self._region_routable_deployments( - request_deployments, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: str-keyed request kwargs - messages is not None or input is not None, - ) - ), - _cooldown_list, - ), - ) - - if strategy == "simple-shuffle": - # if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm - ############## Check 'weight' param set for weighted pick ################# - shuffled: Final = simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=healthy_deployments, - model=model, - request_kwargs=request_kwargs, - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return shuffled - deployment: Final = self._select_deployment_sync( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), - ) - - if deployment is None: - verbose_router_logger.info("get_available_deployment for model: %s, No deployment available", model) - model_ids = self._core_get_model_ids(model_name=model) - _cooldown_time = self._core_cooldown_cache.get_min_cooldown( - model_ids=model_ids, parent_otel_span=parent_otel_span - ) - _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - raise RouterRateLimitError( - model=model, - cooldown_time=_cooldown_time, - enable_pre_call_checks=self._core_enable_pre_call_checks, - cooldown_list=_cooldown_list, - model_ids=model_ids, - ) - self._override_selector_pre_call_check(strategy, strategy_selector, deployment) - verbose_router_logger.info( - "get_available_deployment for model: %s, Selected deployment: %s for model: %s", - model, - self.print_deployment(deployment), - model, - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return deployment - - get_available_deployment = ( # pyright: ignore[reportUnknownVariableType] # legacy selector return - resolve_pods_after( - _get_available_deployment_unresolved # pyright: ignore[reportUnknownArgumentType] # legacy annotations - ) - ) - - def _core_get_available_deployment_for_pass_through( # noqa: ANN202 # unannotated before the rename too; its inferred return mixes untyped deployment dicts - self, - model: str, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - request_kwargs: dict | None = None, - ): - """ - Returns deployments available for pass-through endpoints (based on load balancing strategy) - - Similar to get_available_deployment, but only returns deployments with use_in_pass_through=True - - Args: - model: Model name - messages: Optional list of messages - input: Optional input data - specific_deployment: Whether to find a specific deployment - request_kwargs: Optional request parameters - - Returns: - Dict: Selected deployment configuration - - Raises: - BadRequestError: If no deployment is configured with use_in_pass_through=True - RouterRateLimitError: If no pass-through deployments are available - """ - # 1. Perform common checks to get healthy deployments list - model, healthy_deployments = self._common_checks_available_deployment( - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - - # 2. If the returned is a specific deployment (Dict), verify and return directly - if isinstance(healthy_deployments, dict): - if (healthy_deployments.get("model_info") or {}).get("blocked") is True: - raise litellm_exceptions.ServiceUnavailableError( - message=f"Model '{model}' is currently paused and cannot accept requests.", - model=model, - llm_provider="", - ) - litellm_params: Final = healthy_deployments.get("litellm_params", {}) - if litellm_params.get("use_in_pass_through"): - self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments) - return healthy_deployments - else: - # Specific deployment does not support pass-through - raise litellm_exceptions.BadRequestError( - message=f"Deployment {healthy_deployments.get('model_info', {}).get('id')} does not support pass-through endpoint (use_in_pass_through=False)", - model=model, - llm_provider="", - ) - - # 3. Filter deployments that support pass-through - pass_through_deployments = self._filter_pass_through_deployments(healthy_deployments=healthy_deployments) - - if len(pass_through_deployments) == 0: - # No deployments support pass-through - raise litellm_exceptions.BadRequestError( - message=f"Model {model} has no deployment configured with use_in_pass_through=True. Please add use_in_pass_through: true in the deployment configuration", - model=model, - llm_provider="", - ) - pass_through_model_ids: Final = tuple( - deployment["model_info"]["id"] - for deployment in pass_through_deployments - if "id" in deployment.get("model_info", {}) - ) - - # 4. Apply health-check and cooldown filtering - parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs) - request_deployments: Final = cast( # cast-ok: router deployments are plain dicts - Sequence[DeploymentTypedDict], pass_through_deployments - ) - pass_through_deployments = self._filter_health_check_unhealthy_deployments( - healthy_deployments=pass_through_deployments, - parent_otel_span=parent_otel_span, - ) - cooldown_deployments: Final = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - pass_through_deployments = self._filter_cooldown_deployments( - healthy_deployments=pass_through_deployments, - cooldown_deployments=cooldown_deployments, - ) - pass_through_deployments = self._filter_blocked_deployments(pass_through_deployments) - - # 5. Apply pre-call checks (if enabled) - if self._core_enable_pre_call_checks and (messages is not None or input is not None): - pass_through_deployments = self._pre_call_checks( - model=model, - healthy_deployments=pass_through_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - request_deployments=request_deployments, - ) - - if self._is_priority_routing_group(model): - pass_through_deployments = litellm_utils.get_order_filtered_deployments( - pass_through_deployments, - target_order=request_kwargs.pop("_target_order", None) if request_kwargs is not None else None, - ) - - if len(pass_through_deployments) == 0: - model_ids = self._core_get_model_ids(model_name=model) - _cooldown_time = self._core_cooldown_cache.get_min_cooldown( - model_ids=model_ids, parent_otel_span=parent_otel_span - ) - _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - raise RouterRateLimitError( - model=model, - cooldown_time=_cooldown_time, - enable_pre_call_checks=self._core_enable_pre_call_checks, - cooldown_list=_cooldown_list, - model_ids=pass_through_model_ids, - cooldown_deployments=get_cooldown_deployment_info( - self._core_ctx, - self._deployment_ids( - self._region_routable_deployments( - request_deployments, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: str-keyed request kwargs - messages is not None or input is not None, - ) - ), - _cooldown_list, - ), - ) - - # 6. Apply load balancing strategy - if strategy == "simple-shuffle": - return simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=pass_through_deployments, - model=model, - request_kwargs=request_kwargs, - ) - deployment: Final = self._select_deployment_sync( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=pass_through_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), - ) - - if deployment is None: - verbose_router_logger.info( - "get_available_deployment_for_pass_through model: %s, no available deployments", model - ) - model_ids = self._core_get_model_ids(model_name=model) - _cooldown_time = self._core_cooldown_cache.get_min_cooldown( - model_ids=model_ids, parent_otel_span=parent_otel_span - ) - _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - raise RouterRateLimitError( - model=model, - cooldown_time=_cooldown_time, - enable_pre_call_checks=self._core_enable_pre_call_checks, - cooldown_list=_cooldown_list, - model_ids=model_ids, - ) - self._override_selector_pre_call_check(strategy, strategy_selector, deployment) - - verbose_router_logger.info( - "get_available_deployment_for_pass_through model: %s, selected deployment: %s", - model, - self.print_deployment(deployment), - ) - return deployment - - def _filter_cooldown_deployments( - self, healthy_deployments: list[dict], cooldown_deployments: list[str] - ) -> list[dict]: - """ - Filters out the deployments currently cooling down from the list of healthy deployments - - Args: - healthy_deployments: List of healthy deployments - cooldown_deployments: List of model_ids cooling down. cooldown_deployments is a list of model_id's cooling down, cooldown_deployments = ["16700539-b3cd-42f4-b426-6a12a1bb706a", "16700539-b3cd-42f4-b426-7899"] - - Returns: - List of healthy deployments - """ - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("cooldown deployments: %s", cooldown_deployments) - # Convert to set for O(1) lookup and use list comprehension for O(n) filtering - cooldown_set: Final = set(cooldown_deployments) - return [deployment for deployment in healthy_deployments if deployment["model_info"]["id"] not in cooldown_set] - - def _filter_blocked_deployments(self, healthy_deployments: list[dict]) -> list[dict]: - """ - Filters out deployments that an admin has paused via `LiteLLM_ProxyModelTable.blocked`. - - Applied alongside the cooldown filter on every routing entry point that calls - `_common_checks_available_deployment` directly — the primary sync/async path, - the sync pass-through path, and the retry / health-check helpers — so paused - deployments never serve a request. The async pass-through path inherits this - filter through its delegation to `async_get_healthy_deployments`. - """ - return [ - deployment - for deployment in healthy_deployments - if (deployment.get("model_info") or {}).get("blocked") is not True - ] - - @staticmethod - def _core_is_deployment_blocked(deployment: "Deployment") -> bool: - """ - Returns True when a `Deployment` Pydantic instance carries the admin-paused - flag. Used by credential-lookup helpers so passthrough file / batch endpoints - cannot bypass the pause by resolving credentials directly. - """ - model_info: Final[object | None] = getattr(deployment, "model_info", None) - if model_info is None: - return False - return getattr(model_info, "blocked", None) is True - - async def _async_filter_health_check_unhealthy_deployments( - self, - healthy_deployments: list[dict], - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ) -> list[dict]: - """ - Filter out deployments marked unhealthy by background health checks. - No-op when enable_health_check_routing is False. - When background_health_check_model_groups is set, only deployments in the - listed model groups are filtered; every other group keeps its configured - routing strategy untouched, and a router-level allowed_fails_policy no - longer disables the filter for the listed groups. - Returns all deployments if health state is unavailable, stale, or would - exclude every candidate (safety net). - """ - if not self._core_enable_health_check_routing: - return healthy_deployments - - # When allowed_fails_policy is set, cooldown is the sole routing exclusion - # mechanism -- skip the binary health check filter so the policy threshold - # is respected before any deployment is excluded. With a model-group - # allowlist the filter is already scoped, so listed groups keep it. - scoped_groups: Final = self._core_background_health_check_model_groups - if self._core_allowed_fails_policy is not None and scoped_groups is None: - return healthy_deployments - - unhealthy_ids: Final = await self._core_health_state_cache.async_get_unhealthy_deployment_ids( - parent_otel_span=parent_otel_span - ) - if not unhealthy_ids: - return healthy_deployments - - filtered: Final = [ - d - for d in healthy_deployments - if d["model_info"]["id"] not in unhealthy_ids - or (scoped_groups is not None and d["model_name"] not in scoped_groups) - ] - - if not filtered: - return [] if health_check_probe else healthy_deployments - - return filtered - - def _filter_health_check_unhealthy_deployments( - self, - healthy_deployments: list[dict], - parent_otel_span: Span | None = None, - ) -> list[dict]: - """Sync version of _async_filter_health_check_unhealthy_deployments.""" - if not self._core_enable_health_check_routing: - return healthy_deployments - - scoped_groups: Final = self._core_background_health_check_model_groups - if self._core_allowed_fails_policy is not None and scoped_groups is None: - return healthy_deployments - - unhealthy_ids: Final = self._core_health_state_cache.get_unhealthy_deployment_ids( - parent_otel_span=parent_otel_span - ) - if not unhealthy_ids: - return healthy_deployments - - filtered: Final = [ - d - for d in healthy_deployments - if d["model_info"]["id"] not in unhealthy_ids - or (scoped_groups is not None and d["model_name"] not in scoped_groups) - ] - - if not filtered: - verbose_router_logger.warning("All deployments marked unhealthy by health checks, bypassing health filter") - return healthy_deployments - - return filtered - - def _filter_pass_through_deployments(self, healthy_deployments: list[dict]) -> list[dict]: - """ - Filter out deployments configured with use_in_pass_through=True - - Args: - healthy_deployments: List of healthy deployments - - Returns: - List[Dict]: Only includes a list of deployments that support pass-through - """ - verbose_router_logger.debug( - "Filter pass-through deployments from %s healthy deployments", len(healthy_deployments) - ) - - pass_through_deployments: Final = [ - deployment - for deployment in healthy_deployments - if deployment.get("litellm_params", {}).get("use_in_pass_through", False) - ] - - verbose_router_logger.debug("Found %s deployments with pass-through enabled", len(pass_through_deployments)) - - return pass_through_deployments - - def _track_deployment_metrics(self, deployment, parent_otel_span: Span | None, response=None): - """ - Tracks successful requests rpm usage. - """ - try: - model_id: Final = deployment.get("model_info", {}).get("id", None) - if response is None: - # update self.deployment_stats - if model_id is not None: - self._update_usage(model_id, parent_otel_span) # update in-memory cache for tracking - except Exception as e: - verbose_router_logger.error("Error in _track_deployment_metrics: %s", e) - - def get_num_retries_from_retry_policy(self, exception: Exception, model_group: str | None = None): - return _get_num_retries_from_retry_policy( - exception=exception, - model_group=model_group, - model_group_retry_policy=self.model_group_retry_policy, - retry_policy=self.retry_policy, - ) - - def _core_get_allowed_fails_from_policy(self, exception: Exception) -> int | None: - """ - BadRequestErrorRetries: Optional[int] = None - AuthenticationErrorRetries: Optional[int] = None - TimeoutErrorRetries: Optional[int] = None - RateLimitErrorRetries: Optional[int] = None - ContentPolicyViolationErrorRetries: Optional[int] = None - """ - # if we can find the exception then in the retry policy -> return the number of retries - allowed_fails_policy: Final[AllowedFailsPolicy | None] = self._core_allowed_fails_policy - - if allowed_fails_policy is None: - return None - - if ( - isinstance(exception, litellm_exceptions.AuthenticationError) - and allowed_fails_policy.AuthenticationErrorAllowedFails is not None - ): - return allowed_fails_policy.AuthenticationErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.Timeout) - and allowed_fails_policy.TimeoutErrorAllowedFails is not None - ): - return allowed_fails_policy.TimeoutErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.RateLimitError) - and allowed_fails_policy.RateLimitErrorAllowedFails is not None - ): - return allowed_fails_policy.RateLimitErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.ContentPolicyViolationError) - and allowed_fails_policy.ContentPolicyViolationErrorAllowedFails is not None - ): - return allowed_fails_policy.ContentPolicyViolationErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.BadRequestError) - and allowed_fails_policy.BadRequestErrorAllowedFails is not None - ): - return allowed_fails_policy.BadRequestErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.InternalServerError) - and allowed_fails_policy.InternalServerErrorAllowedFails is not None - ): - return allowed_fails_policy.InternalServerErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.ServiceUnavailableError) - and allowed_fails_policy.ServiceUnavailableErrorAllowedFails is not None - ): - return allowed_fails_policy.ServiceUnavailableErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.BadGatewayError) - and allowed_fails_policy.BadGatewayErrorAllowedFails is not None - ): - return allowed_fails_policy.BadGatewayErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.NotFoundError) - and allowed_fails_policy.NotFoundErrorAllowedFails is not None - ): - return allowed_fails_policy.NotFoundErrorAllowedFails - - def _initialize_alerting(self): - from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting - - if self.alerting_config is None: - return - - router_alerting_config: Final[AlertingConfig] = self.alerting_config - - _slack_alerting_logger: Final = SlackAlerting( - alerting_threshold=router_alerting_config.alerting_threshold, - alerting=["slack"], - default_webhook_url=router_alerting_config.webhook_url, - ) - - self._core_slack_alerting_logger = _slack_alerting_logger - - self._host.callbacks.add_litellm_callback(_slack_alerting_logger) - self._host.callbacks.add_litellm_success_callback(_slack_alerting_logger.response_taking_too_long_callback) - verbose_router_logger.info("\033[94m\nInitialized Alerting for litellm.Router\033[0m\n") - - def _core_set_custom_routing_strategy(self, CustomRoutingStrategy: CustomRoutingStrategyBase) -> None: - """ - Sets get_available_deployment and async_get_available_deployment on an instanced of litellm.Router - - Use this to set your custom routing strategy - - Args: - CustomRoutingStrategy: litellm.router.CustomRoutingStrategyBase - """ - - strategy: Final = CustomRoutingStrategy - setattr( - self, - "get_available_deployment", - resolve_pods_after_bound( - strategy.get_available_deployment, # pyright: ignore[reportUnknownArgumentType] # legacy stub - self._core_kubernetes_pod_discovery, - ), - ) - setattr( - self, - "_core_async_get_available_deployment", - async_resolve_pods_after_bound( - strategy.async_get_available_deployment, # pyright: ignore[reportUnknownArgumentType] # legacy stub - self._core_kubernetes_pod_discovery, - ), - ) - - def _reset_custom_routing_strategy(self) -> None: - for attr in ("get_available_deployment", "_core_async_get_available_deployment"): - if attr in self.__dict__: - delattr(self, attr) - - def flush_cache(self): - self._host.set_global_cache(None) - self._core_cache.flush_cache() - session_in_memory_cache: Final = self._claude_code_session_router_cache.in_memory_cache - if session_in_memory_cache is not None: - session_in_memory_cache.flush_cache() - self._core_prompt_caching_pin_cache.in_memory_cache.flush_cache() - - def reset(self): - ## clean up on close - self._host.set_callback_list("success_callback", []) - self._host.set_callback_list("_async_success_callback", []) - self._host.set_callback_list("failure_callback", []) - self._host.set_callback_list("_async_failure_callback", []) - self.retry_policy = None - self.flush_cache() - - -def _name_core_members_after_old_names() -> None: - # An old member that moved to ``_core_`` can still be reached as the function itself - # (``cached_deployment_model_info`` is an lru_cache over one), so errors that name the method - # keep the name it had before. - for name, member in vars(PythonCore).items(): - old_name = name.removeprefix("_core_") - if old_name != name and inspect.isfunction(member) and hasattr(DeprecatedRouterMembers, old_name): - member.__qualname__ = f"Router.{old_name}" - - -_name_core_members_after_old_names() diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 395b580d9b2..c6ca4a1f399 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -11,11 +11,9 @@ itself only through its host, and the host never calls back into a Python router Python router would have running for the request reach (``OptionalChecks.outcome_callbacks``) - sleeps, fallback access and budget checks, and cache reads and writes are host steps -``RustCore`` does not build, wrap or read a ``PythonCore``. Settings are resolved from the -``RouterConfig`` and the host's process settings; the views (``registry()``, ``settings()``, -``health()``) are read from the Rust handle. A member the Rust core does not serve raises -``RustCoreUnsupported`` (an ``AttributeError``) instead of reaching a Python core, so deleting -``router_python_core.py`` cannot change what this core does. +Settings are resolved from the ``RouterConfig`` and the host's process settings; the views +(``registry()``, ``settings()``, ``health()``) are read from the Rust handle. A member the Rust +core does not serve raises ``RustCoreUnsupported`` (an ``AttributeError``). """ import contextlib @@ -23,13 +21,14 @@ import dataclasses import inspect import itertools import json -from collections.abc import Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterable, Mapping, Sequence +from contextvars import ContextVar from datetime import date, datetime, time from functools import partial from time import perf_counter from time import time as unix_time from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, final +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, TypeVar, final, overload from litellm._internal_context import service_target from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, HEALTH_CHECKS_TARGET @@ -46,7 +45,12 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import ( ) from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type -from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, without +from litellm.litellm_core_utils.core_helpers import ( + as_object_list, + as_str_mapping, + get_parent_otel_span_from_kwargs, + without, +) from litellm.litellm_core_utils.routing_common_utils import ( format_content_policy_fallback_hint, format_fallback_outcome_message, @@ -79,6 +83,7 @@ from litellm.router_rust.invoke.endpoints import ( attempt_call, clean_realtime_session_deployment, compacted_attempt, + endpoint_names, failover_fields, generic_plan, log_refused_call, @@ -208,9 +213,11 @@ from litellm.router_utils.kubernetes_pod_discovery import ( resolve_pods_after_bound, ) from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_io_token_refund_failure from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingReadBatch from litellm.types.router import ( + AllowedFailsPolicy, Deployment, RouterErrors, RouterGeneralSettings, @@ -218,7 +225,7 @@ from litellm.types.router import ( RoutingGroup, RoutingStrategy, ) -from litellm.types.utils import LlmProviders +from litellm.types.utils import CooldownDeploymentInfo, LlmProviders if TYPE_CHECKING: from opentelemetry.trace import Span @@ -444,6 +451,23 @@ _NUMBER_SETTINGS: Final = frozenset( _UPDATE_INT_SETTINGS: Final = frozenset({"timeout", "num_retries", "retry_after", "allowed_fails", "cooldown_time"}) +def _raised_before_provider(error: BaseException) -> bool: + """A request litellm refused before it reached the provider (a parameter the provider does not + take, a malformed priority): neither a failure nor a request of the deployment it picked.""" + from litellm.exceptions import UnsupportedParamsError + + body: Final = getattr(error, "body", None) + return isinstance(error, UnsupportedParamsError) or (isinstance(body, Mapping) and body.get("param") == "priority") + + +def _shares_logging(request: Mapping[str, object]) -> bool: + """The caller passed the logging object every attempt of the request logs on, which logs one + success a request (``should_run_logging``).""" + from litellm.litellm_core_utils.litellm_logging import Logging + + return isinstance(request.get("litellm_logging_obj"), Logging) + + class RustCoreUnsupported(AttributeError): """A member or request the Rust core does not serve. It is never forwarded to a Python core.""" @@ -541,6 +565,14 @@ def _strategy_name(value: object) -> str | None: return None if value is None else normalize_strategy(value) or "simple-shuffle" +_U32_MAX: Final = 2**32 - 1 + + +def _crate_count(value: int | None) -> int | None: + """A count setting as the crate's u32 reads it: a negative one is 0, a huge one the most it holds.""" + return None if value is None else min(max(value, 0), _U32_MAX) + + def _lenient_int(value: object, default: int) -> int: """A count setting kept as given, as the crate reads it: an int, or a number spelled another way (``1.5``, ``"3"``) truncated to one; ``default`` when it is no number. The Python core stores the @@ -763,9 +795,9 @@ class Resolved: if name in {"model_group_alias", "default_litellm_params"}: updates[name] = MappingProxyType(dict(as_str_mapping(value) or {})) elif name in _NUMBER_SETTINGS: + # `allowed_fails` saved later is no explicit one: only the constructor's argument is + # (`explicit_allowed_fails`), so a value equal to the default keeps the rate-based default. updates[name] = self._number_setting(name, value) - if name == "allowed_fails": - updates["explicit_allowed_fails"] = updates[name] elif name in _SELECTION_SETTINGS: updates[name] = self._selection_setting(name, value) elif name == "disable_cooldowns": @@ -823,7 +855,9 @@ class Resolved: "routing_strategy": self.routing_strategy, "cooldown_time": self.cooldown_time, "disable_cooldowns": self.disable_cooldowns, - "allowed_fails_policy": _plain(self.allowed_fails_policy), + # A policy that is no policy or mapping (an attribute write of a list) is none, as the + # Python core's `get_allowed_fails_from_policy` reads no field off it. + "allowed_fails_policy": _plain(self.allowed_fails_policy) if self._has_fails_policy() else None, "retry_after": self.retry_after, "timeout": self.timeout, "routing_groups": group_documents(self.routing_groups) if self.routing_groups else None, @@ -842,6 +876,9 @@ class Resolved: # An explicit `routing_strategy=None` reaches the Rust core as "None": no strategy picks. return kept if self.routing_strategy is not None else {**kept, "routing_strategy": "None"} + def _has_fails_policy(self) -> bool: + return isinstance(self.allowed_fails_policy, AllowedFailsPolicy | Mapping) + def dynamic_document(self, version: int, model_names: Sequence[str]) -> Mapping[str, object]: values: Final = { "version": version, @@ -849,9 +886,11 @@ class Resolved: "fallbacks": _plain(self.fallbacks), "context_window_fallbacks": _plain(self.context_window_fallbacks), "content_policy_fallbacks": _plain(self.content_policy_fallbacks), - "num_retries": self.num_retries, - "max_fallbacks": self.max_fallbacks, - "allowed_fails": self.allowed_fails, + # A negative count reads as none at all in the Python core (`range(num_retries)`, + # `fallback_depth >= max_fallbacks`, `fails > allowed_fails`); the crate reads a u32. + "num_retries": _crate_count(self.num_retries), + "max_fallbacks": _crate_count(self.max_fallbacks), + "allowed_fails": _crate_count(self.allowed_fails), "allowed_fails_is_default": True if self.allowed_fails_is_default() else None, "fallback_keys": self._fallback_keys(model_names), "default_deployment": _plain(served_default_deployment(self.default_deployment)[0]), @@ -878,7 +917,9 @@ class Resolved: for entry in _as_list(entries) if (mapping := as_str_mapping(entry)) ) - names: Final = {*model_names, *self.model_group_alias} + # A bare name a wildcard serves is in no list of names: the keys' own bare spellings stand for it. + bare: Final = (key.split("/", 1)[1] for key in keys if "/" in key) + names: Final = {*model_names, *self.model_group_alias, *bare} prefixed: Final = { name: spelled for name in names if (spelled := _provider_prefixed_model_group(name, keys)) is not None } @@ -1094,6 +1135,14 @@ class RustHealth(RouterHealth): core._resolved.allowed_fails_policy, # pyright: ignore[reportPrivateUsage] # see store ) + @overload + async def usage( + self, model_group: str, *, kind: Literal["current"] = "current" + ) -> tuple[int | None, int | None]: ... + + @overload + async def usage(self, model_group: str, *, kind: Literal["remaining"]) -> Mapping[str, int]: ... + async def usage( # pyright: ignore[reportIncompatibleMethodOverride] # the same two kinds self, model_group: str, *, kind: Literal["current", "remaining"] = "current" ) -> tuple[int | None, int | None] | Mapping[str, int]: @@ -1167,6 +1216,21 @@ _EndpointImpl: TypeAlias = Callable[..., Awaitable[object]] _StrategyInputs: TypeAlias = tuple[str, object, str, list[dict[str, object]]] +def _pick_answer( + step: Mapping[str, object], candidates: Sequence[object], picked: object +) -> list[object]: # mutable-ok: a JSON list + """A strategy step's answer: the picked id, or its position when the candidates share ids (the + step then carries ``records``), so the pick serves the entry it chose as Python's list does.""" + found: Final = picked_id(picked) + if step.get("records") is None or found is None: + return [found] if found is not None else [] + # The entry the selector returned, else (a selector that returned a copy) the first with its id. + position: Final = next((index for index, record in enumerate(candidates) if record is picked), None) + if position is not None: + return [position] + return [next((index for index, record in enumerate(candidates) if picked_id(record) == found), 0)] + + @final class _PickCounts(Protocol): """The router cache as ``track_deployment_pick`` uses it: each deployment id holds its pick count.""" @@ -1193,6 +1257,14 @@ class _EndpointPort: def router(self) -> "RustCore": return self._core + @property + def public(self) -> object: + """The ``Router`` that owns the core, which litellm functions taking ``router=`` get; the core + itself when built alone.""" + owner: Final[object] = self._core.__dict__.get("_core_owner") + public: Final[object] = owner() if callable(owner) else None + return self._core if public is None else public + @property def host(self) -> PythonHost: return self._core._host # pyright: ignore[reportPrivateUsage] # see the class docstring @@ -1343,6 +1415,7 @@ class RustCore( ) self._hosted._core_cache = self._cache # pyright: ignore[reportPrivateUsage] # usage reads come from the cache the Rust core writes ctx: Final = RoutingContext(self) # pyright: ignore[reportArgumentType] # the strategies read registry(), settings(), health() and select() + self._ctx: Final = ctx strategies: Final = PreRoutingStrategies( self._hosted, host, ctx, config.get("auto_router_capability_limit"), self._cache ) @@ -1373,6 +1446,7 @@ class RustCore( self._checks: Final = OptionalChecks(self, host, self._cache, self._hosted, ctx, config.arguments) self._hosted.checks = self._checks strategies.ensure_affinity = self._checks.ensure_deployment_affinity_callback + strategies.hand_over_affinity(self._checks.adopt) self._checks.build(config.get("optional_pre_call_checks"), config.get("model_list")) self._slack_alerting_logger: SlackAlerting | None = initialize_alerting(host, config.get("alerting_config")) # pyright: ignore[reportArgumentType] # the argument as passed if self._slack_alerting_logger is not None: @@ -1384,6 +1458,11 @@ class RustCore( apply_lar1_routing_strategy(ctx, resolved.routing_strategy_args) # pyright: ignore[reportArgumentType] # the core's RoutingContext + @property + def _core_ctx(self) -> RoutingContext: + """The router's ``RoutingContext``: what a routing helper gets where it once took the router.""" + return self._ctx + @property def _registry_state(self) -> HostRegistry: return self._hosted @@ -1417,6 +1496,9 @@ class RustCore( if enabled: register_continuation_prefill_check(self._host) + def __dir__(self) -> list[str]: # mutable-ok: dir() returns a list + return sorted({*super().__dir__(), *endpoint_names()}) + def __getattr__(self, name: str) -> object: if name in self.__dict__.get("_deleted", ()): raise self._no_attribute(name) @@ -1550,6 +1632,27 @@ class RustCore( str(e).replace("\r", "").replace("\n", ""), ) + @contextlib.contextmanager + def _usage_read_scope(self, model: str, request: Mapping[str, object]) -> Generator[None, None, None]: + """``RoutingReadBatch.for_strategy``: a usage-based-routing-v2 pick reads its tpm/rpm counters + in the cooldown read's MGET (see ``_batch_read``), and the strategy reads them from there.""" + from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 + + try: + selector: Final = self._prefetch_usage_selector(model, request) + deployments: Final = self._hosted._core_get_model_list(model_name=model) or [] # pyright: ignore[reportPrivateUsage] # see __init__ + except Exception: # noqa: BLE001 # the batched read is an optimisation: the reads stay apart + yield + return + if not isinstance(selector, LowestTPMLoggingHandler_v2) or not deployments: + yield + return + token: Final = _USAGE_READ.set((RoutingReadBatch(usage_selector=selector), list(deployments))) + try: + yield + finally: + _USAGE_READ.reset(token) + def _prefetch_usage_selector(self, model: str, request_kwargs: Mapping[str, object] | None) -> object: """``_get_routing_context``'s selector when it is a usage-based-routing-v2 one: the request's strategy override, else the routing group ``model`` resolves into, else the router's own.""" @@ -1913,6 +2016,7 @@ class RustCore( # `_core_update_redis_cache`: the routing state's cache gains the Redis cache. self._cache.attach_redis_cache(changes.redis_cache) self._checks.prompt_caching_pin_cache.attach_redis_cache(changes.redis_cache) + self._prerouter.attach_redis_cache(changes.redis_cache) return ApplyResult( version=self._hosted._registry_version, # pyright: ignore[reportPrivateUsage] # see __init__ upserted=upserted, @@ -1967,6 +2071,8 @@ class RustCore( self._settings = None # the snapshot holds settings fields, not the plain attributes if "model_group_alias" in values: self._hosted._core_model_group_alias = dict(resolved.model_group_alias) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # see __init__ + # A group's info reads through the alias map. + self._hosted._invalidate_model_group_info_cache() # pyright: ignore[reportPrivateUsage] # see __init__ self._apply_strategy(previous, resolved, given, settled=settled) if "routing_groups" in values: self._settings = None @@ -2162,11 +2268,15 @@ class RustCore( def __delattr__(self, name: str) -> None: # `del router.` unsets it: the settings view reports it unset from then on. - from litellm.router_api import SETTING_FIELDS + from litellm.router_api import PLAIN_SETTING_ATTRIBUTES, SETTING_FIELDS # The deprecated attribute deletes the core's own name (``_core_``). setting: Final = name.removeprefix("_core_") - if setting not in SETTING_FIELDS and setting not in _DELETABLE_ATTRIBUTES: + if ( + setting not in SETTING_FIELDS + and setting not in PLAIN_SETTING_ATTRIBUTES + and setting not in _DELETABLE_ATTRIBUTES + ): if name in self.__dict__ or name in self._unset_attributes: object.__delattr__(self, name) return @@ -2220,6 +2330,10 @@ class RustCore( elif name in PLAIN_DEFAULTS or name == "max_retries": # Plain attributes on the Python core, outside the setting lists: stored as given. self._store_plain({name: value}) + elif name == "default_priority": + # Queued calls read it from the next call on, kept as given, as the Python core's attribute is. + self._endpoints.assign_setting(name, value) + self._deleted = self._deleted - {name} elif name == "scheduler": self._endpoints.set_scheduler(value) self._deleted = self._deleted - {name} @@ -2392,7 +2506,7 @@ class RustCore( return await self._endpoints.generic(plan_for("aanthropic_messages"), kwargs) async def anthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: endpoint payload - """The proxy's ``/v1/messages`` route name for ``aanthropic_messages``; async, as in PythonCore.""" + """The proxy's ``/v1/messages`` route name for ``aanthropic_messages``; async.""" return await self._endpoints.generic(plan_for("anthropic_messages"), kwargs) async def aimage_generation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload @@ -2446,18 +2560,24 @@ class RustCore( self._refuse_deleted_reads({**request, **(overrides or {})}) self._sync_routing_state() model: Final = request.get("model") - if model is None and "model" in request: - raise self._no_model_error(plan, request) if "model" not in request: raise TypeError(f"{plan.endpoint}() needs a model name") - name: Final = model if isinstance(model, str) else self.routed_name(model) + name: Final = ( + model if isinstance(model, str) else self._no_model_name() if model is None else self.routed_name(model) + ) field: Final = plan.prepare_field or "metadata" if request.get(field) is not None and not isinstance(request.get(field), Mapping): raise TypeError(f"{field} must be a mapping") prepared: Final = dict(request) if prepared_already else self._endpoints.prepare(plan, request, model) # pyright: ignore[reportArgumentType] # the name as given, which the metadata keeps - return self._handle.start( - plan.endpoint, self._routing_inputs(plan, request, prepared, name, overrides) - ), prepared + if not prepared_already: + # `async_function_with_fallbacks`: a selection the caller sent is not one the router made, + # so it never exempts a fallback target from the access gate. + from litellm.router_utils.fallback_event_handlers import clear_pre_routing_selection + + clear_pre_routing_selection(prepared) + # `async_function_with_fallbacks` raises a pick for no model group as it is, with no fallback. + routed: Final = {**(overrides or {}), "disable_fallbacks": True} if model is None else overrides + return self._handle.start(plan.endpoint, self._routing_inputs(plan, request, prepared, name, routed)), prepared def service_logged(self, call_type: str, start_time: float, end_time: float, request: Mapping[str, object]) -> None: """The ``ServiceTypes.ROUTER`` success event of a timed step, as a task of the running loop.""" @@ -2523,28 +2643,17 @@ class RustCore( _ = inferred_provider(model) # pyright: ignore[reportArgumentType] # see above return str(model) - def _no_model_error(self, plan: Plan, request: Mapping[str, object]) -> BaseException: - """``model=None``: the pick fails in ``_common_checks_available_deployment``, naming the - first default fallback when there is one, and with no group the fallback wrapper raises - that error as it is (``async_function_with_fallbacks_common_utils``).""" - import litellm.exceptions as litellm_exceptions - - named: Final = next( + def _no_model_name(self) -> str: + """``model=None``: ``_common_checks_available_deployment`` picks for the first default + fallback when there is one, and else for the name ``None``, which only a wildcard matches.""" + return next( ( - targets[0] + str(targets[0]) for entry in _as_list(self._resolved.fallbacks) if (targets := _as_list((as_str_mapping(entry) or {}).get("*"))) ), - None, + "None", ) - error: Final = litellm_exceptions.BadRequestError( - message=f"You passed in model={named}. {RouterErrors.no_healthy_deployments.value}", - model=named, # pyright: ignore[reportArgumentType] # None, as the Python router passes it - llm_provider="", - ) - if not plan.sync: - _Driver(self, plan, request)._log_pick_failure(error) # pyright: ignore[reportPrivateUsage] # the async pick's own logging - return error def _routing_inputs( self, @@ -2579,7 +2688,10 @@ class RustCore( ``make_call``, and a recovered stream is async.""" sync: Final = plan is not None and plan.sync make_call: Final = plan is None or not plan.single_attempt - callbacks: Final = self.outcome_callbacks(request, sync=sync, make_call=make_call) + callbacks: Final = { + **self.outcome_callbacks(request, sync=sync, make_call=make_call), + **({"shared_logging": True} if _shares_logging(request) else {}), + } return {"outcome_callbacks": callbacks} if callbacks else {} def _selection_fields(self, plan: Plan, request: Mapping[str, object]) -> Mapping[str, object]: @@ -2600,6 +2712,10 @@ class RustCore( **({"web_search": True} if web_search else {}), **({"route_handler": handler} if serves_route_capability(handler) else {}), **({"decisions_media": True} if decisions_media else {}), + # `batch_create_routing`: a batch create keeps to the deployments that can run batches. + **({"batch_create": True} if plan.kind == "batch" else {}), + # `deployments_admitting_embedding_batch`: an embedding keeps to the deployments admitting its input. + **({"embedding_batch": True} if plan.kind == "embedding" else {}), } def _picks_async(self) -> bool: @@ -2607,12 +2723,14 @@ class RustCore( def _process_renamed(self, model: str, group: str) -> str: """``model`` as the error names it (``_common_checks_available_deployment``): the group a name - resolved to by name, alias or routing group, renamed through the process ``model_alias_map``.""" + resolved to by name, alias or routing group, or the group of the deployment an id names, + renamed through the process ``model_alias_map``.""" resolution: Final = self._handle.resolve(model) - if resolution is None or resolution["kind"] not in ("Group", "Alias", "RoutingGroup"): + if resolution is None or resolution["kind"] not in ("Group", "Alias", "RoutingGroup", "Id"): return model + named: Final = str(resolution["model_group"]) if resolution["kind"] == "Id" else group renames: Final = self._host.dynamic().model_alias_map - return renames[group] if group in renames else group # noqa: SIM401 # `renames` is any process mapping, read as Python reads it + return renames[named] if named in renames else named # noqa: SIM401 # `renames` is any process mapping, read as Python reads it def _max_parallel_requests_client(self, deployment: Mapping[str, object], kwargs: Mapping[str, object]) -> object: """``_get_client(client_type="max_parallel_requests")``: the deployment's slot limit.""" @@ -2719,7 +2837,10 @@ class RustCore( with owning_picks(self._selectors.reservations, call_id_of(CALL_IDENTITY.validate_python(request))): call, prepared = self._start(plan, request, routed, prepared_already=prepared_already) arm_compaction_state(self._port, plan, prepared) # pyright: ignore[reportArgumentType] # the call's own dict - with picks_reserve(plan.endpoint not in _UNRESERVED_OPS): + with ( + picks_reserve(plan.endpoint not in _UNRESERVED_OPS), + self._usage_read_scope(str(prepared.get("model")), prepared), + ): driver, step = await self._run(plan, call, prepared, fresh=True, passthrough=passthrough) if step["kind"] == "failed": raise driver.failed(step) @@ -2744,6 +2865,16 @@ class RustCore( driver: Final = _Driver( self, plan, prepared, cut_deployment, stream_failure, fresh=fresh, passthrough=passthrough ) + try: + return await self._run_steps(driver, call) + except BaseException: + driver.close_route_phase() + raise + + async def _run_steps( + self, driver: "_Driver", call: "RouterCall[_CallSteps]" + ) -> "tuple[_Driver, Mapping[str, object]]": + """``_run``'s loop over the call's steps.""" outcome: object = None # rebind-ok: the answer to the last step while True: step = call.next(outcome) @@ -2769,6 +2900,7 @@ class RustCore( case "pre_route": outcome = await driver.pre_route(step) case "done" | "failed": + driver.close_route_phase() return driver, step case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list raise RuntimeError(f"router-core returned an unknown step {other!r}") @@ -3003,10 +3135,17 @@ class RustCore( if not isinstance(request, dict): return None depth: Final = step.get("fallback_depth") + if isinstance(depth, int) and depth > 0: + # A fallback hop routes on its own copy of the metadata buckets (`run_async_fallback`'s + # kwargs): what its hook clears and stamps leaves the caller's metadata as the first hop left it. + for name in ("metadata", "litellm_metadata"): + bucket = request.get(name) + if isinstance(bucket, dict): + request[name] = dict(bucket) # pyright: ignore[reportUnknownArgumentType] # the request's metadata dict r: Final = self._resolved owner: Final[object] = self.__dict__.get("_core_owner") public: Final[object] = owner() if callable(owner) else None - return await self._prerouter.pre_route( + routed: Final[object] = await self._prerouter.pre_route( str(step.get("model")), request, # pyright: ignore[reportUnknownArgumentType] # the request kwargs the hook stamps messages=request.get("messages") if messages is _NO_INPUT else messages, @@ -3023,6 +3162,7 @@ class RustCore( if self._hosted._core_routing_plugins # pyright: ignore[reportPrivateUsage] # the live plugin list else None, ) + return _served_as(routed, request) async def _run_plugins( self, @@ -3075,9 +3215,15 @@ class RustCore( selector: Final = self._selectors.selector_for( strategy, str(step.get("scope")), str(group) if group is not None else None, request ) - candidates: Final = [ - dict(as_str_mapping(self._handle.deployment(str(id))) or {}) for id in _as_list(step.get("candidates")) - ] + # Candidates sharing an id come with their own records, so the pick can land on any of them. + records: Final = step.get("records") + candidates: Final = ( + [dict(as_str_mapping(record) or {}) for record in _as_list(records)] + if records is not None + else [ + dict(as_str_mapping(self._handle.deployment(str(id))) or {}) for id in _as_list(step.get("candidates")) + ] + ) return strategy, selector, str(step.get("model")), candidates def _records_of(self, ids: object) -> list[dict[str, object]]: # mutable-ok: the deployment dicts filters take @@ -3144,6 +3290,45 @@ class RustCore( ) ] + def _override_selector(self, step: Mapping[str, object], request: Mapping[str, object]) -> object: + """The request's own strategy selector when it overrides the router's (``_override_selectors``), + whose pre-call check is in no callback list.""" + strategy: Final = str(step.get("strategy")) + selector: Final = self._selectors.selector_for(strategy, "override", None, request) + return selector if selector is not None and self._selectors.is_override(strategy, selector) else None + + def _embedding_batch( + self, step: Mapping[str, object], request: Mapping[str, object], input: object, *, tags: bool + ) -> list[str]: # mutable-ok: the step's answer + """``deployments_admitting_embedding_batch`` over the candidates; when none admits the input, + nothing while a routable, unpaused cooling deployment does (the pick waits out its cooldown), + else the candidates with the largest limit, so the 400 names the most the group accepts.""" + from litellm.router_utils.embedding_batch_size_check import ( + deployments_admitting_embedding_batch, + deployments_when_none_admits_embedding_batch, + embedding_deployment_selection, + ) + + sent: Final = request.get("input", input) + candidates: Final = self._records_of(step.get("candidates")) + with embedding_deployment_selection(): + admitting: Final = deployments_admitting_embedding_batch(candidates, sent, dict(request)) + if admitting is not None: + return [found for record in admitting if (found := picked_id(record)) is not None] + pool: Final = self._records_of(step.get("pool")) + pool_ids: Final = [found for record in pool if (found := picked_id(record)) is not None] + cooling: Final = [id for id, _ in self._cooldowns.get_active_cooldowns(pool_ids, None)] # pyright: ignore[reportArgumentType] # a list of ids + routable: Final = _region_routable(pool, request, True, self._resolved.enable_pre_call_checks) + routable_cooling: Final = ( + self._tag_routable_ids(str(step.get("model")), pool, routable, frozenset(cooling), request) + if tags + else frozenset(id for record in routable if (id := picked_id(record)) in cooling) + ) + kept: Final = deployments_when_none_admits_embedding_batch( + candidates, pool, routable_cooling, sent, dict(request) + ) + return [found for record in kept if (found := picked_id(record)) is not None] + async def _answer_select( self, step: Mapping[str, object], @@ -3167,6 +3352,14 @@ class RustCore( return self._route_incapable(step) if step.get("stage") == "media_refusing": return self._media_refusing(step, request) + if step.get("stage") == "embedding_batch": + return self._embedding_batch(step, request, input, tags=True) + if step.get("stage") == "override_check": + selector = self._override_selector(step, request) + if selector is not None: + for record in self._records_of(step.get("candidates")): + await selector.async_pre_call_check(dict(record), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + return step["candidates"] if step.get("stage") == "provider_request": return _provider_request(step) inputs: Final = self._strategy_inputs(step, request) @@ -3185,7 +3378,7 @@ class RustCore( picked: Final = ( candidates[0] if pinned - else await pick_async( + else await _pick_with_prefetched_usage( strategy, selector, # pyright: ignore[reportArgumentType] # a selector from selector_for model, @@ -3202,8 +3395,7 @@ class RustCore( await selector.async_pre_call_check(dict(as_str_mapping(picked) or {}), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger if picked is not None and timed: self.service_logged(".async_get_available_deployments", start_time, clock(), request) - found: Final = picked_id(picked) - return [found] if found is not None else [] + return _pick_answer(step, candidates, picked) def _answer_select_sync( self, step: Mapping[str, object], request: Mapping[str, object], messages: object, input: object @@ -3221,6 +3413,14 @@ class RustCore( return self._route_incapable(step) if step.get("stage") == "media_refusing": return self._media_refusing(step, request) + if step.get("stage") == "embedding_batch": + return self._embedding_batch(step, request, input, tags=False) + if step.get("stage") == "override_check": + selector = self._override_selector(step, request) + if selector is not None: + for record in self._records_of(step.get("candidates")): + selector.pre_call_check(dict(record)) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + return step["candidates"] if step.get("stage") == "provider_request": return _provider_request(step) inputs: Final = self._strategy_inputs(step, request) @@ -3242,8 +3442,7 @@ class RustCore( ) if picked is not None and self._selectors.is_override(strategy, selector): selector.pre_call_check(dict(as_str_mapping(picked) or {})) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # see above - found: Final = picked_id(picked) - return [found] if found is not None else [] + return _pick_answer(step, candidates, picked) # -- cache steps ----------------------------------------------------------------------------- @@ -3259,6 +3458,9 @@ class RustCore( read: Final[object] = await self._cache.async_batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above store: Final = self._cooldowns.cooldown_store + usage_read: Final = _USAGE_READ.get() + if usage_read is not None: + return self._corrected_cooldowns(keys, await self._batch_read_with_usage(keys, *usage_read)) with service_target(ROUTER_COOLDOWNS_TARGET): prefetched: Final = await RoutingReadBatch._read_prefetched(((store, keys),)) # pyright: ignore[reportPrivateUsage] # the batch's own prefetch read results: Final[object] = ( @@ -3266,6 +3468,37 @@ class RustCore( ) return self._corrected_cooldowns(keys, results) + async def _batch_read_with_usage( + self, + keys: list[str], + batch: RoutingReadBatch, + deployments: list[Mapping[str, object]], # mutable-ok: the read's key and deployment lists + ) -> object: + """The cooldown read with the usage-based-routing-v2 counters of ``deployments`` in one MGET + (``RoutingReadBatch.async_get_cooldown_deployments``), the counters kept on ``batch``.""" + from litellm.caching.dual_cache import DualCache + from litellm.router_strategy.lowest_tpm_rpm_v2 import PrefetchedUsage + from litellm.router_utils.routing_read_batch import ( + _routing_read_target, # pyright: ignore[reportPrivateUsage] # the batch's own target name + ) + + selector: Final = batch.usage_selector + assert selector is not None # set by `_usage_read_scope` + usage_keys: Final = list(itertools.chain(*selector.usage_counter_keys(deployments))) # pyright: ignore[reportArgumentType] # deployment dicts + reads: Final = ((self._cooldowns.cooldown_store, keys), (selector.router_cache, usage_keys)) + with service_target(_routing_read_target(keys, usage_keys)): + results: Final = await RoutingReadBatch._read_prefetched( + reads + ) or await DualCache.async_batch_get_cache_shared( # pyright: ignore[reportPrivateUsage] # see above + reads, parent_otel_span=None + ) + usage_values: Final = results[1] + batch.prefetched_usage = PrefetchedUsage( + keys=frozenset(usage_keys), + values=None if usage_values is None else MappingProxyType(dict(zip(usage_keys, usage_values))), + ) + return results[0] + def _batch_read_sync(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list """``_batch_read`` for a sync pick: ``get_active_cooldowns``' read.""" if not keys or not keys[0].endswith(":cooldown"): @@ -3295,7 +3528,13 @@ class RustCore( ) for key, amount in zip(keys, amounts) ] - values: Final = await self._cache.async_increment_cache_pipeline(increment_list=operations) + try: + values: Final = await self._cache.async_increment_cache_pipeline(increment_list=operations) + except Exception as e: # noqa: BLE001 # usage that cannot be written never fails the call, as on the Python router + from litellm._logging import verbose_router_logger + + verbose_router_logger.debug("_increment_deployment_usage: usage was not recorded. Got - %s", e) + return if values is None or self._cache.redis_cache is None: return for operation, value in zip(operations, values): @@ -3397,6 +3636,81 @@ class RustCore( # -- errors the core raises when no attempt failed ------------------------------------------ + def _routable_cooled_info( + self, + step: Mapping[str, object], + group: str, + cooling: Sequence[str], + request: Mapping[str, object], + *, + has_input: bool, + tags: bool, + ) -> tuple[CooldownDeploymentInfo, ...] | None: + """``RouterRateLimitError.cooldown_deployments``: the cooling deployments the request could + route to (region and, on an async pick, tag routing over the deployments the core listed + before its cooldown filter); ``None`` when the strategy refused every candidate.""" + request_ids: Final = step.get("error_request_ids") + if request_ids is None: + return None + records: Final = self._records_of(request_ids) + routable: Final = _region_routable(records, request, has_input, self._resolved.enable_pre_call_checks) + route: Final = ( + self._tag_routable_ids(group, records, routable, frozenset(cooling), request) + if tags + else frozenset(id for record in routable if (id := picked_id(record)) is not None) + ) + return get_cooldown_deployment_info(RoutingContext(self), route, cooling) # pyright: ignore[reportArgumentType] # reads registry() + + def _tag_routable_ids( + self, + group: str, + pool: Sequence[Mapping[str, object]], + routable: Sequence[Mapping[str, object]], + cooled: frozenset[str], + request: Mapping[str, object], + ) -> frozenset[str]: + """The cooling deployments in ``routable`` that tag routing would send the request to, each + pick counted along with the picks it falls back to once that one is unavailable, judged + without recording a tag match on the request.""" + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + from litellm.router_strategy.tag_based_routing import get_deployments_for_tag + + field: Final = get_metadata_variable_name_from_kwargs(dict(request)) + metadata: Final = request.get(field) + if field in request and not isinstance(metadata, Mapping): + return frozenset() + tag_request: Final = {**request, field: {**metadata}} if isinstance(metadata, Mapping) else dict(request) + found = frozenset[str]() # rebind-ok: gains each pass's tag routing picks + remaining = [ # rebind-ok: each pass drops what tag routing just picked + record for record in routable if picked_id(record) in cooled + ] + while remaining: + try: + # Tag routing awaits nothing: its coroutine finishes on the first send. + picked = _finished( + get_deployments_for_tag( + ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # reads registry() and settings() + model=group, + healthy_deployments=remaining, # pyright: ignore[reportArgumentType] # deployment dicts + request_kwargs=tag_request, # pyright: ignore[reportArgumentType] # the request kwargs + metadata_variable_name=field, + request_deployments=pool, # pyright: ignore[reportArgumentType] # deployment dicts + ) + ) + except ValueError: # tag routing sends this request to none of the rest + break + picked_ids = frozenset( + id + for record in (picked if isinstance(picked, list | tuple) else ()) + if (id := picked_id(record)) is not None + ) # pyright: ignore[reportArgumentType] # deployment dicts + rest = [record for record in remaining if picked_id(record) not in picked_ids] + if len(rest) == len(remaining): + break + found |= picked_ids + remaining = rest + return found + def _failure_error( # noqa: C901 # one branch per failure Python tells apart self, step: Mapping[str, object], @@ -3405,6 +3719,8 @@ class RustCore( *, pass_through: str = "", request: Mapping[str, object] | None = None, + has_input: bool = True, + tags: bool = True, ) -> BaseException: """The error a call or pick that failed before any attempt raises; ``pass_through`` names the selection kind (``pass_through``/``async_pass_through``) whose errors it words.""" @@ -3534,10 +3850,8 @@ class RustCore( enable_pre_call_checks=self._resolved.enable_pre_call_checks, cooldown_list=cooling, model_ids=list(ids), # pyright: ignore[reportArgumentType] # a list of ids - cooldown_deployments=get_cooldown_deployment_info( - RoutingContext(self), # pyright: ignore[reportArgumentType] # reads registry() - ids, - cooling, + cooldown_deployments=self._routable_cooled_info( + step, str(group), cooling, request or {}, has_input=has_input, tags=tags ), ) if last is not None: @@ -3581,6 +3895,88 @@ def _sleep_seconds(step: Mapping[str, object]) -> float: return 0 if seconds == 0 else seconds +def _served_as(routed: object, request: Mapping[str, object]) -> object: + """A pre-routed group, paired with the group the hook stamped the request as served for + (``fallback_lookup_groups``' routed group) when that differs: the level's fallbacks look it up.""" + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + + if not isinstance(routed, str): + return routed + metadata: Final = request.get(get_metadata_variable_name_from_kwargs(dict(request))) + stamped: Final = metadata.get("model_group") if isinstance(metadata, Mapping) else None + return (routed, stamped) if isinstance(stamped, str) and stamped and stamped != routed else routed + + +_USAGE_READ: Final[ContextVar["tuple[RoutingReadBatch, list[Mapping[str, object]]] | None"]] = ContextVar( + "rust_core_usage_read", default=None +) + + +async def _pick_with_prefetched_usage(*args: object, **kwargs: object) -> object: + """``pick_async`` with the counters the pick's cooldown read fetched (``PrefetchedUsage``).""" + from litellm.router_strategy.lowest_tpm_rpm_v2 import PrefetchedUsage + + usage_read: Final = _USAGE_READ.get() + with PrefetchedUsage.scoped(usage_read[0].prefetched_usage if usage_read is not None else None): + return await pick_async(*args, **kwargs) # pyright: ignore[reportArgumentType, reportCallIssue] # pick_async's own arguments + + +def _deployment_pick_attributes(model: str, request: Mapping[str, object]) -> Mapping[str, str | int]: + """Bounded attributes for one deployment pick; the attempt is 1-based within its model group.""" + metadata: Final = request.get("litellm_metadata", request.get("metadata")) + attempted: Final = metadata.get("attempted_retries") if isinstance(metadata, Mapping) else None + retries: Final = attempted if isinstance(attempted, int) else 0 + depth: Final = request.get("fallback_depth") + hop: Final = (isinstance(depth, int) and depth > 0) or model != request.get("model") + reason: Final = "retry" if retries > 0 else "fallback" if hop else "initial" + return MappingProxyType( + { + "litellm.deployment.attempt": retries + 1, + "litellm.deployment.reason": reason, + "litellm.deployment.model_group": model, + } + ) + + +def _region_routable( + records: Sequence[Mapping[str, object]], + request: Mapping[str, object], + has_input: bool, + pre_call_checks: bool, +) -> Sequence[Mapping[str, object]]: + """``records`` without the ones ``_pre_call_checks`` keeps the request away from for being + outside its ``allowed_model_region``.""" + from litellm.types.router import LiteLLM_Params + from litellm.utils import is_region_allowed + + allowed: Final = request.get("allowed_model_region") + if allowed is None or not pre_call_checks or not has_input: + return records + if not isinstance(allowed, str): + return () + return [ + record + for record in records + if is_region_allowed( + litellm_params=LiteLLM_Params.model_validate(record.get("litellm_params") or {}), + allowed_model_region=allowed, + ) + ] + + +_Value = TypeVar("_Value") + + +def _finished(coroutine: Coroutine[object, object, _Value]) -> _Value: + """The value of a coroutine that awaits nothing, run to its end without an event loop.""" + try: + coroutine.send(None) + except StopIteration as finished: + return finished.value # pyright: ignore[reportAny] # the coroutine's return value + coroutine.close() + raise RuntimeError("the coroutine awaited something") + + def _uses_pass_through(record: object) -> bool: params: Final = as_str_mapping((as_str_mapping(record) or {}).get("litellm_params")) or {} return bool(params.get("use_in_pass_through")) @@ -3676,28 +4072,46 @@ class _Selection: if step.get("single"): # A deployment named by its id is listed alone, as a dict. (record,) = _as_list(step.get("deployments")) - return for_requested_model(typed_deployment(as_str_mapping(record) or {}), self._model) + return self._served(record) # The deployments as served: a priority group's members carry their priority as order. - return [ - for_requested_model(typed_deployment(as_str_mapping(record) or {}), self._model) - for record in _as_list(step.get("deployments")) - ] + return [self._served(record) for record in _as_list(step.get("deployments"))] if step.get("deployment_id") in (DEFAULT_DEPLOYMENT_FAILING_ID, DEFAULT_DEPLOYMENT_COPIED_ID): # A default deployment every call to fails, or whose id is not a string, is picked as Python copies it. picked: Final = default_deployment_copy(self._core._resolved.default_deployment, self._model) # pyright: ignore[reportPrivateUsage] # the selection serves its own core return picked # pyright: ignore[reportReturnType] # the assigned value's copy, as the Python core returns it - return for_requested_model(typed_deployment(as_str_mapping(step.get("deployment")) or {}), self._model) + return self._served(step.get("deployment")) - def _failed(self, step: Mapping[str, object]) -> BaseException: + def _served(self, record: object) -> Mapping[str, object]: + """A deployment as served: its record, with the litellm_params values JSON carries as text.""" + return _with_live_params( + self._core, for_requested_model(typed_deployment(as_str_mapping(record) or {}), self._model) + ) + + def _failed(self, step: Mapping[str, object], *, sync: bool) -> BaseException: core: Final = self._core error: Final = core._failure_error( # pyright: ignore[reportPrivateUsage] # see above - step, self._model, self._filter_error, pass_through=self._kind, request=self._request + step, + self._model, + self._filter_error, + pass_through=self._kind, + request=self._request, + has_input=self._messages is not None or self._input is not None, + tags=not sync, ) if self._kind not in ("pass_through", "sync_pick"): _log_refused_pick(self._request, error) return error async def run(self) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + from litellm.integrations.otel.runtime import phase_span + + with ( + phase_span(f"route {self._model}"), + self._core._usage_read_scope(str(self._model), self._request), # pyright: ignore[reportPrivateUsage] # the selection serves its own core + ): + return await self._run() + + async def _run(self) -> Mapping[str, object] | Sequence[Mapping[str, object]]: core: Final = self._core outcome: object = None # rebind-ok: the answer to the last step while True: @@ -3733,7 +4147,7 @@ class _Selection: # `_core_async_get_available_deployment` resolves its pick to a pod. return await self._resolved_pod(core, done) case "failed": - raise self._failed(step) + raise self._failed(step, sync=False) case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list raise RuntimeError(f"router-core returned an unknown step {other!r}") @@ -3769,7 +4183,7 @@ class _Selection: # `get_available_deployment` resolves its pick to a pod. return self._resolved_pod_sync(core, done) case "failed": - raise self._failed(step) + raise self._failed(step, sync=True) case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list raise RuntimeError(f"router-core returned an unknown step {other!r}") @@ -3809,6 +4223,23 @@ def _log_refused_image_pick(request: Mapping[str, object], error: BaseException) _restore_correlation_context_if_supported(logging_obj) +def _with_live_params(core: "RustCore", deployment: Mapping[str, object]) -> Mapping[str, object]: + """``deployment`` with the values its JSON record carries as text: litellm_params holding an + Exception ``mock_response``, a callable or a str subclass that marks the value (an admin's secret + reference), and model_info values with keys that are not strings.""" + info: Final = as_str_mapping(deployment.get("model_info")) + model_id: Final = (info or {}).get("id") + live: Final = core._hosted.live_params(model_id) # pyright: ignore[reportPrivateUsage] # the core's own registry + live_info: Final = core._hosted.live_info(model_id) # pyright: ignore[reportPrivateUsage] # same + if not live and not live_info: + return deployment + return { + **deployment, + **({"litellm_params": {**(as_str_mapping(deployment.get("litellm_params")) or {}), **live}} if live else {}), + **({"model_info": {**(info or {}), **live_info}} if live_info else {}), + } + + def _custom_pick(picked: object) -> Mapping[str, object]: """A custom strategy's answer as the deployment document the Rust core serves.""" document: Final = as_str_mapping(picked) @@ -3863,6 +4294,65 @@ _NO_ROUTING_PRE_CALL_CHECKS: Final = frozenset({"search", "rerank", "moderation" _UNPICKED_KINDS: Final = frozenset({"search", "create_file"}) +# The info lines the Python router's endpoint helpers logged for an attempt: its success, and for +# these its failure too. +_ATTEMPT_LOG_NAMES: Final = MappingProxyType( + { + "acompletion": "litellm.acompletion", + "completion": "litellm.completion", + "aembedding": "litellm.aembedding", + "embedding": "litellm.embedding", + "aimage_generation": "litellm.aimage_generation", + "image_generation": "litellm.image_generation", + "atranscription": "litellm.atranscription", + "aspeech": "litellm.aspeech", + "arerank": "litellm.arerank", + "atext_completion": "litellm.atext_completion", + "aadapter_completion": "litellm.aadapter_completion", + } +) +_SUCCESS_ONLY_LOG_NAMES: Final = MappingProxyType( + { + "acancel_batch": "litellm.acancel_batch", + "acreate_batch": "litellm.acreate_batch", + "acreate_file": "litellm.acreate_file", + "avector_store_create": "litellm.avector_store_create", + } +) + + +def _log_level_traceback(error: BaseException | None) -> None: + """``async_function_with_fallbacks_common_utils``' DEBUG traceback of the error a level hands + its fallbacks: the stack the ERROR line leaves out.""" + import logging + import traceback + + from litellm._logging import redact_string, verbose_router_logger + + if error is None or not verbose_router_logger.isEnabledFor(logging.DEBUG): + return + formatted: Final = "".join(traceback.format_exception(error)) + verbose_router_logger.debug("Traceback%s", redact_string(formatted)) + + +def _log_attempt(plan: Plan, model_name: object, error: BaseException | None = None) -> None: + """``litellm.(model=...) 200 OK`` or ``... Exception ...``, as the Python router logged each attempt.""" + from litellm._logging import verbose_router_logger + + name: Final = _ATTEMPT_LOG_NAMES.get(plan.endpoint) + if error is not None: + if name is not None: + verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", name, model_name, str(error)) + return + served_name: Final = ( + name + or _SUCCESS_ONLY_LOG_NAMES.get(plan.endpoint) + or ("ageneric_api_call_with_fallbacks" if plan.kind == "generic" else None) + ) + if served_name is not None: + verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", served_name, model_name) + + class _Driver: """Answers the attempt and fallback steps of one call, and shapes its result. @@ -3897,6 +4387,8 @@ class _Driver: # the level that stream served, on the failure ``async_function_with_fallbacks_common_utils`` # was handed, which is what it raises when no fallback answers. self._base_level: Final = stream_failure[0] if stream_failure is not None else 0 + # How many fallback hops deep the walk is: the "Successful fallback" lines a served hop logs. + self._fallback_depth: int = 0 self._level_errors: Final[dict[int, BaseException]] = dict( # mutable-ok: filled as the call runs (stream_failure,) if stream_failure else () ) @@ -3916,6 +4408,9 @@ class _Driver: # When the latest attempt started: Python reads it off the logging object, when one is set up. self._attempt_started: datetime | None = None # A single-attempt endpoint runs neither of the wrappers that stamp the walk on the request. + # The open ``route {model_group}`` phase span of the pick in progress (async calls only). + self._route_phase: contextlib.ExitStack | None = None + self._route_model: str = "" self._walk: Final = WalkMetadata( request, plan.metadata_field, self._params_base_model, live=fresh and not plan.single_attempt ) @@ -4025,7 +4520,9 @@ class _Driver: def _picked_error(self, event: Mapping[str, object]) -> BaseException: """The error a pick an event names raised, as a failed step with that error raises it.""" model: Final = str(event.get("failed_model") or event.get("model")) - return self._core._failure_error(event, model, None, request=self._request) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + return self._core._failure_error( # pyright: ignore[reportPrivateUsage] # the driver serves its own core + event, model, None, request=self._request, has_input=self._has_input(), tags=not self._sync + ) def _pick_raised(self, error: Exception, *, logged: bool) -> Mapping[str, object]: """A pick raised on the host (a strategy, a tag filter, a pre-routing hook): the try fails with @@ -4097,6 +4594,16 @@ class _Driver: self._chain_to_hop_first(picked, at_level) self._level_errors[at_level] = picked + def _log_hop(self, event: Mapping[str, object]) -> None: + """``run_async_fallback``'s "Falling back to model_group" line, for a hop to ``event``'s target.""" + from litellm._logging import verbose_router_logger + from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure + + target: Final = event["unroutable"] if "unroutable" in event else event.get("to") + verbose_router_logger.info("Falling back to model_group = %s", mask_sensitive_structure(target)) + depth: Final = event.get("depth") + self._fallback_depth = depth if isinstance(depth, int) else self._fallback_depth + 1 + def _carry_superseded_usage(self) -> None: """``carry_superseded_attempt_usage``, before each same-group retry and each fallback hop.""" from litellm.router_utils.fallback_event_handlers import carry_superseded_attempt_usage @@ -4198,11 +4705,39 @@ class _Driver: primary.message += self._explanation(primary, str(self._request.get("model"))) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the message is what this explains self._explains_failure = False + def open_route_phase(self, model: object) -> None: + """``_async_get_available_deployment_unresolved``'s ``route {model}`` phase: the pick's reads + nest in it, until the pick ends.""" + from litellm.integrations.otel.runtime import phase_span + + if self._sync or self._route_phase is not None: + return + stack: Final = contextlib.ExitStack() + stack.enter_context(phase_span(f"route {model}")) + self._route_phase = stack + self._route_model = str(model) + + def close_route_phase(self, selected_in: object = None) -> None: + """The pick ended: a served pick marks ``litellm.request.deployment_selected`` inside it first.""" + phase: Final = self._route_phase + if phase is None: + return + if isinstance(selected_in, str): + from litellm.integrations.otel.runtime import phase_event + + phase_event("litellm.request.deployment_selected", _deployment_pick_attributes(selected_in, self._request)) + self._route_phase = None + phase.close() + def take_events(self, events: object) -> None: # noqa: C901 # one branch per router-core event for raw_event in _as_list(events): event = as_str_mapping(raw_event) if event is None: continue + if event.get("kind") == "pick_started": + self.close_route_phase() + self.open_route_phase(event.get("model")) + continue if event.get("kind") == "retry": self._carry_superseded_usage() continue @@ -4211,6 +4746,7 @@ class _Driver: continue if event.get("kind") == "selected": self._selected(event) + self.close_route_phase(self._route_model) continue if event.get("kind") == "retries_opened": self._walk.retries_opened(self._stamped_budget(event)) @@ -4219,6 +4755,7 @@ class _Driver: self._walk.retry_budget(event) continue if event.get("kind") == "pick_failed": + self.close_route_phase() self._pick_failed(event) continue if event.get("kind") == "order_chain_failed": @@ -4254,6 +4791,7 @@ class _Driver: self._fallbacks_failed(event) continue if event.get("kind") == "fallback": + self._log_hop(event) hop_to = event.get("level") self._hop_first.pop(hop_to if isinstance(hop_to, int) else 0, None) self._walk.hop(event) @@ -4273,6 +4811,7 @@ class _Driver: if event.get("kind") == "falling_back": level = event.get("level") at = level if isinstance(level, int) else 0 + _log_level_traceback(self._level_errors.get(at, self._last_error)) self._explains_failure = self._explains_failure or at == self._base_level self._hint_level_failure(self._level_errors.get(at, self._last_error), str(event.get("model"))) continue @@ -4320,12 +4859,14 @@ class _Driver: def _call(self, step: Mapping[str, object]) -> "tuple[Mapping[str, object], AttemptCall]": name: Final = step.get("model") requested: Final = str(name) if name is not None else str(self._request.get("model")) - deployment: Final = ( + deployment: Mapping[str, object] = ( # rebind-ok: given its live params below # A default deployment whose id is not a string is called as Python copies it, id and all. as_str_mapping(default_deployment_copy(self._core._resolved.default_deployment, requested)) or {} # pyright: ignore[reportPrivateUsage] # the driver serves its own core if step.get("deployment_id") == DEFAULT_DEPLOYMENT_COPIED_ID else for_requested_model(typed_deployment(as_str_mapping(step.get("deployment")) or {}), requested) ) + # The values its JSON record carries as text (an Exception `mock_response`) called as configured. + deployment = _with_live_params(self._core, deployment) self._attempts += 1 self._attempt_started = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline's api_call_start_time level: Final = step.get("fallback_level") @@ -4343,7 +4884,7 @@ class _Driver: self._strip_undecryptable_reasoning(step, deployment, call.kwargs) shaped: Final = shape_attempt_kwargs(call.kwargs, deployment) # The walk's metadata for the attempt's level (retries, fallback hops) and its hop kwargs. - call.kwargs = self._walk.stamp(step, shaped) + call.kwargs = self._walk.stamp(step, shaped, deployment) if self._plan.stamps_walk else shaped field: Final = self._plan.metadata_field if isinstance(call.served, dict) and field in call.kwargs: call.served[field] = call.kwargs[field] @@ -4387,6 +4928,7 @@ class _Driver: order: Final = step.get("target_order") deployment_id: Final = step.get("deployment_id") + routed: Final = step.get("routed") return StreamAttempt( core=self._core, op=self._op, @@ -4398,6 +4940,7 @@ class _Driver: level=number("fallback_level"), deployment_id=deployment_id if isinstance(deployment_id, str) else None, attempted_retries=number("retry"), + routed=routed if isinstance(routed, str) else None, slot=slot, ) @@ -4410,6 +4953,8 @@ class _Driver: ) -> Mapping[str, object]: # What litellm's failure logging read, before the attempt helper's `except` wrote on it. detail: Final = str(error) + if call is not None: + _log_attempt(self._plan, call.model_name, error) served: Final = call.served if call is not None else {} # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values stamp_error(self._plan, error, deployment, served) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values self._core._endpoints.attempt_failed(self._plan, error, deployment, call, str(step.get("model"))) # pyright: ignore[reportPrivateUsage] # see above @@ -4425,14 +4970,41 @@ class _Driver: self._attempt_started, ) model_info: Final = as_str_mapping(deployment.get("model_info")) or {} + exits: Final = self._core.failure_exits(error, logged, model_info.get("id")) + # A call under credentials the request supplied ran as a request-scoped deployment + # (`_handle_clientside_credential`): its failure is never the shared deployment's. + ran_as: Final = (as_str_mapping(as_str_mapping(served).get("model_info") if served else None) or {}).get("id") # pyright: ignore[reportUnknownArgumentType] # see above + request_scoped: Final = self._core.ran_as_request_deployment( + ran_as, model_info.get("id") + ) or self._caller_entra_failure(error, served, model_info.get("id")) # pyright: ignore[reportUnknownArgumentType] # see above return { **attempt_failure(error), "detail": detail, - **self._core.failure_exits(error, logged, model_info.get("id")), + **exits, + **({"outside_deployment": True} if request_scoped or _raised_before_provider(error) else {}), # `async_deployment_callback_on_failure` does not run for a sync call, and skips a batch retrieve. - "uncounted_request": self._plan.sync or is_batch_retrieve_call_type(logged.get("call_type")), + "uncounted_request": self._plan.sync + or is_batch_retrieve_call_type(logged.get("call_type")) + or _raised_before_provider(error), } + def _caller_entra_failure(self, error: BaseException, sent: Mapping[str, object], deployment_id: object) -> bool: + """``is_caller_entra_credential_failure``: Entra ID rejected a credential the caller sent, not + the deployment's own, so the failure is the caller's.""" + from litellm.router_utils.cooldown_handlers import is_caller_entra_credential_failure + + if not isinstance(deployment_id, str): + return False + hosted: Final = self._core._hosted # pyright: ignore[reportPrivateUsage] # the driver serves its own core + defaults: Final = {k: v for k, v in self._core._resolved.default_litellm_params.items() if v is not None} # pyright: ignore[reportPrivateUsage] # see above + return is_caller_entra_credential_failure( + hosted._core_get_deployment(deployment_id), # pyright: ignore[reportPrivateUsage] # the host registry's id lookup + defaults, + {"litellm_params": sent}, + getattr(error, "status_code", ""), + error, + ) + def _chain_to_hop_first(self, error: BaseException, level: int) -> None: """A retry's failure is raised while the hop's first failure is handled (the retry loop runs in its ``except`` block), so it chains to it.""" @@ -4444,6 +5016,13 @@ class _Driver: self, call: "AttemptCall", response: object, host_op: str, stream: "StreamAttempt | None" = None ) -> Mapping[str, object]: self._last_response = response + _log_attempt(self._plan, call.model_name) + if self._fallback_depth > 0: + from litellm._logging import verbose_router_logger + + # One line per `run_async_fallback` the served hop returns through. + for _ in range(self._fallback_depth): + verbose_router_logger.info("Successful fallback b/w models.") tokens: Final = total_tokens(response) # A wrapped stream reports its success or failure when it ends; the attempt counts the # request and the tokens read so far, which the end's report leaves out. @@ -4625,6 +5204,15 @@ class _Driver: # plain dict fails the attempt. self._core._endpoints.counters.succeeded(self._plan, call) # pyright: ignore[reportPrivateUsage] # see above except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop + if call is not None: + # `make_call`: a failed call hands back the io tokens it reserved, whatever its + # logging object (shared across attempts, it logs only the first failure). + with contextlib.suppress(Exception): + await async_io_token_refund_failure( + self._core._cache, # pyright: ignore[reportPrivateUsage] # the driver serves its own core + call.kwargs, # pyright: ignore[reportArgumentType] # the call's kwargs dict + parent_otel_span=get_parent_otel_span_from_kwargs(call.kwargs), + ) failed: Final = self._failed(step, deployment, call, error) if reading_stream is None or not logs_failure: return {**failed, "recorded": not logs_failure} @@ -4724,6 +5312,10 @@ class _Driver: ctx: Final = RoutingContext(self._core) # pyright: ignore[reportArgumentType] # the checks read registry() and settings() request: Final = dict(self._request) + if self._plan.kind in _GENERIC_API_KINDS: + # The generic helpers' kwargs carry the function they call, which the checks read (a + # search fallback runs the search tool grants). + request["original_generic_function"] = self._core._host.invoke(self._plan.host_op) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # the driver serves its own core try: if target != get_pre_routing_selection(request) and not await ctx.fallback_access_allowed( model=target, request_kwargs=request @@ -4825,11 +5417,17 @@ class _Driver: await await_all(self.outcomes.take_pending()) return counted + def _has_input(self) -> bool: + """The request carries messages or an input, which region routing reads.""" + return self._request.get("messages") is not None or self._request.get("input") is not None + def failed(self, step: Mapping[str, object]) -> BaseException: model: Final = str(self._request.get("model")) primary: Final = self._level_errors.get(self._base_level) if primary is None: - return self._core._failure_error(step, model, self._last_error, request=self._request) # pyright: ignore[reportPrivateUsage] # see above + return self._core._failure_error( # pyright: ignore[reportPrivateUsage] # see above + step, model, self._last_error, request=self._request, has_input=self._has_input(), tags=not self._sync + ) if ( self._base_level == 0 and self._explains_failure diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index c1dc5162ccf..e2c621ccb7d 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -73,6 +73,11 @@ class EndpointCore(Protocol): """The core itself, for the provider handlers and helpers that take a router.""" ... + @property + def public(self) -> object: + """The ``Router`` that owns the core, for litellm functions that take ``router=``.""" + ... + @property def host(self) -> "PythonHost": ... @@ -232,6 +237,9 @@ class Plan: fail_key: CountKey | None = "deployment" # One attempt, no retries and no fallbacks. single_attempt: bool = False + # Whether the attempt's metadata carries the walk's keys (``original_model_group``, the retry + # counts): not for an endpoint whose metadata goes to the provider as the caller sent it. + stamps_walk: bool = True # The named parameter the endpoint takes beside **kwargs, sent between the deployment params and # the request (`{**data, name: value, **kwargs}`). named: tuple[str, ...] = () @@ -267,8 +275,9 @@ class Plan: def stamps_failed_deployment(self) -> bool: """Whether the attempt stamps ``failed_deployment_id`` on its error, which weighted failover reads: every endpoint but file and batch creation (whose result belongs to the deployment - that served it), moderation's pass-through and search.""" - return self.kind not in ("batch", "cancel_batch", "create_file", "moderation", "search") + that served it), moderation's pass-through, search and vector store creation (whose + handler calls the provider without the attempt wrapper).""" + return self.kind not in ("batch", "cancel_batch", "create_file", "moderation", "search", "vector_store_create") @property def metadata_field(self) -> str: @@ -450,6 +459,7 @@ PLANS: Final[Mapping[str, Plan]] = MappingProxyType( stamp="none", fail_key="group", single_attempt=True, + stamps_walk=False, ), "amoderation": Plan( "amoderation", @@ -636,6 +646,44 @@ _SYNC_VECTOR_STORE_FILES: Final[Mapping[str, str]] = MappingProxyType( } ) +# The endpoints ``RustEndpoints`` serves with a method of its own, by the same name. +_DIRECT_ENDPOINTS: Final = frozenset( + { + "avector_store_create", + "acreate_file", + "acreate_batch", + "acancel_batch", + "aretrieve_batch", + "alist_batches", + "aadapter_completion", + "abatch_completion", + "abatch_completion_one_model_multiple_requests", + "abatch_completion_fastest_response", + "schedule_acompletion", + "priority_header_update", + "aguardrail", + "asearch", + } +) + + +def endpoint_names() -> frozenset[str]: + """Every endpoint name ``RustEndpoints.member`` may serve: what ``dir()`` of a router lists.""" + return frozenset( + { + *PLANS, + *_HOST_GENERIC, + *_VIDEO_DIRECT, + *_SYNC_GENERIC, + *_UNBRANCHED, + *_ASSISTANTS, + *_MANAGED_AGENTS, + *_SYNC_VECTOR_STORE_FILES, + *_DIRECT_ENDPOINTS, + "create_interaction", + } + ) + def plan_for(endpoint: str) -> Plan: plan: Final = generic_plan(endpoint) @@ -968,8 +1016,7 @@ def _search_attempt( kwargs: dict[str, object], # mutable-ok: the attempt's own kwargs group: str, ) -> AttemptCall: - """`SearchAPIRouter.async_search_with_fallbacks_helper`: one of the search tools named - ``group`` at random, called with its provider, credentials and params under the request's.""" + """One of the search tools named ``group`` at random, called with its provider, credentials and params under the request's.""" import random tools: Final = [tool for tool in core.endpoints.search_tools if as_dict(tool).get("search_tool_name") == group] @@ -1359,7 +1406,15 @@ class RustEndpoints: self._search_tools: list[object] = tools # mutable-ok: the router's live list self._guardrails: list[object] = list(_as_items(config.get("guardrail_list"))) # mutable-ok: same parsed: Final = parse_default_priority(config.get("default_priority")) - self._default_priority: Final = None if isinstance(parsed, InvalidPriority) else parsed + if isinstance(parsed, InvalidPriority): + from litellm._logging import verbose_router_logger + + verbose_router_logger.warning( + "router_settings.default_priority is ignored, no request is queued by default: %s", + parsed.message, + ) + # apply() assigns it after construction, kept as given, as the Python core's attribute is. + self._default_priority: int | None = None if isinstance(parsed, InvalidPriority) else parsed self._polling_interval: Final = config.get("polling_interval") self._caching_groups: Final = config.get("caching_groups") self._default_max_parallel_requests: object = config.get("default_max_parallel_requests") @@ -1375,15 +1430,17 @@ class RustEndpoints: self._cache_responses = value def assign_setting(self, name: str, value: object) -> None: - """``search_tools``, ``guardrail_list`` or ``default_max_parallel_requests`` assigned after - construction, kept as given: later searches, guardrail picks and slot limits built from - then on read it, as the Python core's attributes are read.""" + """``search_tools``, ``guardrail_list``, ``default_max_parallel_requests`` or + ``default_priority`` assigned after construction, kept as given: later searches, guardrail + picks, slot limits and queued calls read it, as the Python core's attributes are read.""" if name == "search_tools": self._search_tools = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it elif name == "guardrail_list": self._guardrails = value # pyright: ignore[reportAttributeAccessIssue] # same elif name == "default_max_parallel_requests": self._default_max_parallel_requests = value + elif name == "default_priority": + self._default_priority = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it else: raise ValueError(f"not an endpoint setting: {name}") @@ -1490,23 +1547,8 @@ class RustEndpoints: return self._sync_vector_store_files(name) if name == "create_interaction": return self._factory("create_interaction", "interactions.create") - direct: Final = { - "avector_store_create": self.avector_store_create, - "acreate_file": self.acreate_file, - "acreate_batch": self.acreate_batch, - "acancel_batch": self.acancel_batch, - "aretrieve_batch": self.aretrieve_batch, - "alist_batches": self.alist_batches, - "aadapter_completion": self.aadapter_completion, - "abatch_completion": self.abatch_completion, - "abatch_completion_one_model_multiple_requests": self.abatch_completion_one_model_multiple_requests, - "abatch_completion_fastest_response": self.abatch_completion_fastest_response, - "schedule_acompletion": self.schedule_acompletion, - "priority_header_update": self.priority_header_update, - "aguardrail": self.aguardrail, - "asearch": self.asearch, - }.get(name) - if direct is not None: + if name in _DIRECT_ENDPOINTS: + direct: Final[Callable[..., object]] = getattr(self, name) # pyright: ignore[reportAny] # one of this class's own endpoint methods return direct if name == "responses": return self._sync_factory(name, PLANS["responses"]) @@ -1538,7 +1580,10 @@ class RustEndpoints: client: object, kwargs: dict[str, object], # mutable-ok: the wrapper's own kwargs ) -> object: - plan: Final = generic_plan(name) or _generic(name, host_op) + listed: Final = generic_plan(name) + # The endpoint's own plan, unless the wrapper calls another litellm function for it + # (``avector_store_create`` without a model calls ``vector_stores.main.acreate``). + plan: Final = listed if listed is not None and listed.host_op == host_op else _generic(name, host_op) if name in _CONTAINER_ENDPOINTS: return await self._containers(plan, custom_llm_provider, kwargs) if name in _RESPONSE_ID_ENDPOINTS: @@ -1618,7 +1663,7 @@ class RustEndpoints: if provider_kwargs.get("model"): return self._core.call_sync(plan, provider_kwargs, _SINGLE_ATTEMPT) if name == "vector_store_search": - return self._invoke(plan.host_op)(**{**provider_kwargs, "router": self._core.router}) + return self._invoke(plan.host_op)(**{**provider_kwargs, "router": self._core.public}) return self._invoke(plan.host_op)(**provider_kwargs) sync_wrapper.__name__ = name @@ -1695,7 +1740,7 @@ class RustEndpoints: if kwargs.get("model"): return await self._ageneric(plan, kwargs) if name == "avector_store_search": - return await self._direct(plan.host_op, {**kwargs, "router": self._core.router}) + return await self._direct(plan.host_op, {**kwargs, "router": self._core.public}) return await self._direct(plan.host_op, kwargs) async def _interactions( @@ -1862,10 +1907,9 @@ class RustEndpoints: if self._scheduler is None: from litellm.scheduler import Scheduler - self._scheduler = ( - Scheduler(polling_interval=self._polling_interval) # pyright: ignore[reportArgumentType] # as configured - if self._polling_interval is not None - else Scheduler() + self._scheduler = Scheduler( + polling_interval=self._polling_interval, # pyright: ignore[reportArgumentType] # as configured + redis_cache=self._core.routing_cache.redis_cache, ) return self._scheduler @@ -1880,7 +1924,12 @@ class RustEndpoints: assert isinstance(scheduler, Scheduler) async def healthy_deployments() -> Sequence[object]: - healthy: Final = await self._core.select(model, dict(request), kind="healthy") + try: + healthy: Final = await self._core.select(model, dict(request), kind="healthy") + except Exception: # noqa: BLE001 # `_async_get_healthy_deployments` lists none on any lookup failure + # It lists none while every deployment cools down, or for a group it cannot look up: + # the request keeps its place in the queue, and the call itself raises the error. + return [] return list(healthy) if isinstance(healthy, Sequence) else [healthy] await scheduler.wait_for_turn( @@ -2440,6 +2489,7 @@ __all__ = [ "attempt_call", "clean_realtime_session_deployment", "compacted_attempt", + "endpoint_names", "failover_fields", "generic_plan", "log_refused_call", diff --git a/litellm/router_rust/invoke/streaming.py b/litellm/router_rust/invoke/streaming.py index ca70143072f..d8b4f87501a 100644 --- a/litellm/router_rust/invoke/streaming.py +++ b/litellm/router_rust/invoke/streaming.py @@ -108,11 +108,16 @@ class StreamAttempt: deployment_id: str | None attempted_retries: int committed_budget: int | None = None + # The group a pre-routing hook selected for ``group``, which a fallback lookup keys on first. + routed: str | None = None # The deployment's max_parallel_requests slot the attempt holds; an open chat stream takes it. slot: contextlib.ExitStack | None = None # Set by the wrapper that reports the stream's outcome when it ends; the attempt's answer is then # deferred and counts only the request and the tokens read so far (``counted_tokens``). reports_outcome: bool = False + # Whether the logging object had logged a failure when the wrapper took the stream: an earlier + # attempt's failure blocks this one's report, the stream's own log of its failure does not. + failure_logged_before: frozenset[str] = frozenset() counted_tokens: int = 0 outcome_reported: bool = False # The background recording of the stream's failure, which a same-group retry waits for so it @@ -132,6 +137,7 @@ class StreamAttempt: "deployment_id": self.deployment_id, "attempted_retries": self.attempted_retries, "committed_budget": self.committed_budget, + "routed": self.routed, } def policy(self, trigger: BaseException | None = None) -> Mapping[str, object]: @@ -192,6 +198,13 @@ class StreamAttempt: """The outcome callbacks that run for this attempt when its stream ends.""" return self.core.outcome_callbacks(self.kwargs, sync=sync) + def take_outcome(self) -> None: + """The wrapper reports this attempt's outcome when its stream ends.""" + self.reports_outcome = True + self.failure_logged_before = frozenset( + event for event in ("async_failure", "sync_failure") if failure_already_logged(self.kwargs, event) + ) + def _claim(self) -> bool: if not self.reports_outcome or self.outcome_reported: return False @@ -209,9 +222,9 @@ class StreamAttempt: def report_failure(self, error: BaseException) -> None: """The stream failed after it opened: count the failure, cool down by the rules and count one more request (``async_deployment_callback_on_failure``), in the background. A logging object - shared with an earlier failure blocks it: the stream logs through ``attempt_scoped_copy``, read - here, when the stream fails.""" - if self._claim() and not failure_already_logged(self.kwargs, "async_failure"): + shared with an earlier failure blocks it, read when the wrapper took the stream: the stream's own + failure log marks the same object before this runs.""" + if self._claim() and "async_failure" not in self.failure_logged_before: self.failure_recording = _record_in_background( self.core.report_stream_failure(self._failure_logged(error), self._callbacks(sync=False)) ) @@ -235,7 +248,7 @@ class StreamAttempt: def report_failure_sync(self, error: BaseException) -> None: """A sync stream failed after it opened: counted and cooled down, with no request counted.""" - if self._claim() and not failure_already_logged(self.kwargs, "sync_failure"): + if self._claim() and "sync_failure" not in self.failure_logged_before: logged: Final = self._failure_logged(error) _record_now(lambda: self.core.report_stream_failure_sync(logged, self._callbacks(sync=True))) @@ -284,20 +297,28 @@ async def _areported( items: AsyncIterable[_ItemT], total_tokens: Callable[[object], int], observe: Callable[[object], None] | None = None, + *, + fails_before_first_item: bool = True, ) -> AsyncGenerator[_ItemT, None]: """``items``, reporting the attempt's stream once they end: a success with the tokens - ``total_tokens`` reads off the last item, or the failure they raise. Closed early, nothing is - reported. The wrapper marks the attempt (``reports_outcome``) when it wraps the stream, before - the attempt answers.""" + ``total_tokens`` reads off the last item, or the failure they raise, unless they raise it before + their first item and ``fails_before_first_item`` is off. Closed early, nothing is reported. The + wrapper marks the attempt (``reports_outcome``) when it wraps the stream, before the attempt + answers.""" last: object = None # rebind-ok: the newest item + started = False # rebind-ok: whether an item came try: async for item in items: last = item + started = True if observe is not None: observe(item) yield item except Exception as error: - attempt.report_failure(error) + if started or fails_before_first_item: + attempt.report_failure(error) + else: + attempt.outcome_reported = True # rebind-ok: a stream that never opened counts nothing raise attempt.report_success(total_tokens(last), last) @@ -918,7 +939,7 @@ def wrap_chat_stream( # noqa: C901 # one function per Python wrapper it ports after content, the provider's error is raised. ``held_slot`` (the deployment's max_parallel_requests slot) is released when the stream is exhausted, closed, or falls back.""" slot: Final = held_slot if held_slot is not None else contextlib.ExitStack() - attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + attempt.take_outcome() async def close_model_response() -> None: if not hasattr(model_response, "aclose"): @@ -1013,7 +1034,7 @@ def wrap_chat_stream_sync( attempt: StreamAttempt, model_response: CustomStreamWrapper, first_chunk: _FirstChunk[Iterator[object]] ) -> FallbackAwareStreamWrapper: """``_completion_streaming_iterator``: the sync twin, which never continues after content.""" - attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + attempt.take_outcome() def stream_with_fallbacks() -> Generator[object, None, None]: fallback_response: object = None # rebind-ok: tracked for cleanup @@ -1253,7 +1274,7 @@ def wrap_anthropic_stream( replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator) model: Final = attempt.group usage: Final = _AnthropicStreamUsage() - attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + attempt.take_outcome() async def close_source() -> None: with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): @@ -1527,7 +1548,7 @@ def wrap_responses_stream( # noqa: C901 # one function per Python wrapper it p from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, get_openai_response_types source_iterator: Final = response - attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + attempt.take_outcome() openai_types: Final = get_openai_response_types() terminal_event_types: Final = ( openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, @@ -1664,7 +1685,7 @@ def wrap_responses_stream( # noqa: C901 # one function per Python wrapper it p else cast("MidStreamFallbackError", stream_error) # cast-ok: not a context window error ) if response_id_exposed: - _raise_original_mid_stream_error(e) + _raise_released_stream_error(attempt, e) async with contextlib.aclosing( _responses_fallback_attempt(attempt, e, source_iterator, wrapper.adopt_fallback_headers, held) ) as fallback_stream: @@ -1686,6 +1707,33 @@ def wrap_responses_stream( # noqa: C901 # one function per Python wrapper it p return wrapper +def _raise_released_stream_error(attempt: StreamAttempt, e: MidStreamFallbackError) -> NoReturn: + """A background stream whose response id is already out raises its provider error as is; a content + policy refusal before output reads as the router's error would after the walk found no fallback.""" + from litellm.litellm_core_utils.routing_common_utils import format_content_policy_no_fallback_note + from litellm.router_utils.fallback_event_handlers import fallback_lookup_groups, fallbacks_disabled_for_request + + original: Final = e.original_exception + request: Final = attempt.request + model_group: Final = request.get("model") + if ( + attempt.core._host.dynamic().expose_router_debug_in_errors # pyright: ignore[reportPrivateUsage] # the attempt's own core + and isinstance(original, litellm_exceptions.ContentPolicyViolationError) + and isinstance(model_group, str) + and not fallbacks_disabled_for_request(request) + and _mid_stream_error_before_output(e) + ): + walked: Final = attempt.core._walked_fallbacks # pyright: ignore[reportPrivateUsage] # the attempt's own core + original.message += format_content_policy_no_fallback_note( + model_group, + fallback_lookup_groups(request, model_group), + request.get("content_policy_fallbacks", walked.get("content_policy_fallbacks")), + request.get("fallbacks", walked.get("fallbacks")), + attempt.depth > 0, + ) + _raise_original_mid_stream_error(e) + + def _responses_continuation_request( attempt: StreamAttempt, e: MidStreamFallbackError, source_iterator: "BaseResponsesAPIStreamingIterator" ) -> Mapping[str, object] | None: @@ -1782,6 +1830,18 @@ async def _responses_fallback_attempt( verbose_router_logger.debug("stream_with_fallbacks(aresponses): error closing fallback: %s", exc) +def _reported_transcription_stream(attempt: StreamAttempt, response: object) -> object: + """A transcription stream reports its attempt's outcome when it ends: a stream that fails after it + opened cools its deployment down by the rules, as the router's failure callback does.""" + from litellm.litellm_core_utils.audio_utils.transcription_streaming import LoggingAsyncTranscriptionStream + + if not isinstance(response, LoggingAsyncTranscriptionStream): + return response + attempt.take_outcome() + response._iterator = _areported(attempt, response._iterator, lambda _last: 0) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the stream's own iterator, wrapped + return response + + # --- the attempt hook --------------------------------------------------------------------------- @@ -1798,13 +1858,37 @@ async def open_attempt_stream(attempt: StreamAttempt, response: object) -> objec return wrap_chat_stream(attempt, response, first_chunk, held_slot) if not attempt.kwargs.get("stream"): return response + if attempt.op == "atranscription": + return _reported_transcription_stream(attempt, response) + if attempt.op == "atext_completion": + return _reported_text_completion_stream(attempt, response) if attempt.op == "aresponses" and isinstance(response, BaseResponsesAPIStreamingIterator): + # `_aresponses_streaming_iterator` wraps a call that asked for a stream, not one the + # deployment's own `stream` setting forced. + if not attempt.request.get("stream"): + return response return wrap_responses_stream(attempt, response) if attempt.op in ("aanthropic_messages", "anthropic_messages") and hasattr(response, "__aiter__"): return wrap_anthropic_stream(attempt, cast("AsyncIterator[bytes]", response)) # cast-ok: a byte stream return response +def _reported_text_completion_stream(attempt: StreamAttempt, response: object) -> object: + """A text-completion stream reports its attempt's outcome when it ends, as litellm's logging + does for it: a success with the stream's tokens, or a failure after the first chunk. One that + fails before its first chunk counts nothing, as its logging carries no deployment then.""" + from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper + + if not isinstance(response, TextCompletionStreamWrapper): + return response + attempt.take_outcome() + chunks: Final = cast("AsyncIterable[object]", response.completion_stream) # cast-ok: an untyped attribute + response.completion_stream = _areported( # rebind-ok: the stream's own iterator, wrapped + attempt, chunks, _chat_total_tokens, fails_before_first_item=False + ) + return response + + def open_attempt_stream_sync(attempt: StreamAttempt, response: object) -> object: if attempt.op == "completion" and isinstance(response, CustomStreamWrapper): if response.completion_stream is None and response.make_call is not None: # pyright: ignore[reportUnknownMemberType] # untyped attributes diff --git a/litellm/router_rust/invoke/walk.py b/litellm/router_rust/invoke/walk.py index 4dab8edc07f..17e20119f43 100644 --- a/litellm/router_rust/invoke/walk.py +++ b/litellm/router_rust/invoke/walk.py @@ -41,6 +41,17 @@ _SERVED_KEYS: Final = ("model_info", "deployment") _HOP_KWARGS: Final = ("fallback_depth", "max_fallbacks", "attempted_targets", "include_fallback_errors") +def _reentrant_hop_stamps(request: Mapping[str, object], field: str) -> Mapping[str, object]: + """A hop ``run_async_fallback`` re-entered with its own ``attempted_targets`` keeps the fallback + stamps it carries in place of the root level's (``async_function_with_fallbacks`` skips them).""" + from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets + + if not isinstance(request.get("attempted_targets"), AttemptedFallbackTargets): + return {} + bucket: Final = as_str_mapping(request.get(field)) or {} + return {key: bucket[key] for key in ("attempted_fallbacks", "original_model_group") if key in bucket} + + class _Scope: __slots__ = ("kwargs", "metadata") @@ -146,6 +157,7 @@ class WalkMetadata: root: Final[dict[str, object]] = {"attempted_fallbacks": 0} # mutable-ok: the root level's metadata if model is not None: root["original_model_group"] = model + root.update(_reentrant_hop_stamps(request, field)) self._request: Final = request self._field: Final = field self._params_base_model: Final = params_base_model @@ -252,10 +264,18 @@ class WalkMetadata: def hop_kwargs(self, event: Mapping[str, object]) -> dict[str, object]: # mutable-ok: a fresh dict for the loggers """The kwargs the level held after a hop (``kwargs["model"] = target``), as the - fallback loggers see them.""" + fallback loggers see them: the request's own metadata (the proxy's ``user_api_key_*`` + labels) with the level's walk stamps laid over it, as ``run_async_fallback``'s + ``_hop_metadata`` copies it.""" scope: Final = self._scope(event.get("scope")) model: Final = scope.metadata.get("model_group", event.get("to")) - return {**self._request, **scope.kwargs, "model": model, self._field: dict(scope.metadata)} + request_bucket: Final = as_str_mapping(self._request.get(self._field)) or {} + return { + **self._request, + **scope.kwargs, + "model": model, + self._field: {**request_bucket, **scope.metadata}, + } def log_retry(self, event: Mapping[str, object], error: BaseException | None) -> None: """A ``retry_logged`` event: the failure joins the scope's ``previous_models``.""" @@ -289,8 +309,11 @@ class WalkMetadata: """The level's ``request_retry_count``, as the retry loop has counted it so far.""" return self._scope(step.get("scope")).metadata.get("request_retry_count") - def stamp(self, step: Mapping[str, object], kwargs: Mapping[str, object]) -> Mapping[str, object]: - """The attempt's kwargs with its level's walk metadata and hop kwargs laid over them.""" + def stamp( + self, step: Mapping[str, object], kwargs: Mapping[str, object], deployment: Mapping[str, object] | None = None + ) -> Mapping[str, object]: + """The attempt's kwargs with its level's walk metadata and hop kwargs laid over them. + ``deployment``: the deployment the attempt calls, whose own ``mock_timeout`` a hop keeps.""" scope: Final = self._scope(step.get("scope")) metadata: Final = scope.metadata metadata["attempted_retries"] = step.get("retry") @@ -299,10 +322,31 @@ class WalkMetadata: self._write_through( scope, {"attempted_retries": metadata["attempted_retries"], "max_retries": metadata["max_retries"]} ) - served: Final = dict(as_str_mapping(kwargs.get(self._field)) or {}) + # `include_fallback_errors` is the walk's own flag: the fallback loggers see it, the provider never does. + hop: Final = { + key: scope.kwargs[key] for key in _HOP_KWARGS if key in scope.kwargs and key != "include_fallback_errors" + } + depth: Final = hop.get("fallback_depth") + given: Final = kwargs.get(self._field) + # The requested group's attempts stamp the caller's own dict in place, as + # `_update_kwargs_with_deployment` stamps it: the proxy's post-call write-backs land in the dict + # its spend row reads. A fallback hop stamps its own copy (`run_async_fallback`'s kwargs). + served: Final = ( + given # pyright: ignore[reportUnknownVariableType] # the caller's metadata dict + if isinstance(given, dict) and not (isinstance(depth, int) and depth > 0) + else dict(as_str_mapping(given) or {}) + ) served.update({key: metadata[key] for key in _WALK_KEYS if key in metadata}) metadata.update({key: served.get(key) for key in _SERVED_KEYS}) - return {**kwargs, **{key: scope.kwargs[key] for key in _HOP_KWARGS if key in scope.kwargs}, self._field: served} + # `async_function_with_fallbacks` pops the request's `mock_timeout` for the requested group's own + # retries: no hop sees it. A deployment's own `mock_timeout` is merged in per call, so it stays. + own: Final = as_str_mapping((deployment or {}).get("litellm_params")) or {} + sent: Final = ( + {key: value for key, value in kwargs.items() if key != "mock_timeout"} + if isinstance(depth, int) and depth > 0 and "mock_timeout" not in own + else kwargs + ) + return {**sent, **hop, self._field: served} async def _log_failed_fallbacks(ctx: "RoutingContext", detail: str, parent_otel_span: "Span | None") -> None: diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index b906a496550..75afe38548e 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -218,6 +218,12 @@ class OptionalChecks: self.optional_callbacks.append(callback) self._host.callbacks.add_litellm_callback(callback) + def adopt(self, callback: object) -> None: + """Own a check already registered with the host: list it in ``optional_callbacks``.""" + if self.optional_callbacks is None: + self.optional_callbacks = [] + self.optional_callbacks.append(callback) + def ensure_deployment_affinity_callback(self) -> None: """``_ensure_deployment_affinity_callback``: the affinity check with every global flag off.""" if any(isinstance(cb, DeploymentAffinityCheck) for cb in self.optional_callbacks or ()): @@ -323,7 +329,9 @@ class OptionalChecks: if self.optional_callbacks is None or not any(type(cb) is callback_cls for cb in self.optional_callbacks): return self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls] - others: Final = tuple(owner for owner in live_check_owners if id(owner) != id(self._owner)) + others: Final = tuple( + owner for owner in live_check_owners if id(owner) != id(self._owner) and _router_alive(owner) + ) if any(any(type(cb) is callback_cls for cb in (owner.optional_callbacks or ())) for owner in others): return callbacks: Final = _callbacks(self._host) @@ -552,22 +560,42 @@ class RustOptionalChecks: def optional_callbacks(self) -> "list[object] | None": # mutable-ok: the router's live list return self._opt.optional_callbacks + @optional_callbacks.setter + def optional_callbacks(self, value: "list[object] | None") -> None: # mutable-ok: the list a caller assigns + """Assigning the list replaces it, as on the Python router: the checks read the new one.""" + self._opt.optional_callbacks = value + @property def router_budget_logger(self) -> RouterBudgetLimiting | None: return self._opt.router_budget_logger + # The affinity settings are plain attributes on the Python router: an assignment is what an + # affinity check added afterwards (`add_optional_pre_call_checks`) is built with. + @property def deployment_affinity_ttl_seconds(self) -> object: return self._opt.deployment_affinity_ttl_seconds + @deployment_affinity_ttl_seconds.setter + def deployment_affinity_ttl_seconds(self, value: object) -> None: + self._opt.deployment_affinity_ttl_seconds = value + @property def session_affinity_repin_on_failover(self) -> object: return self._opt.session_affinity_repin_on_failover + @session_affinity_repin_on_failover.setter + def session_affinity_repin_on_failover(self, value: object) -> None: + self._opt.session_affinity_repin_on_failover = value + @property def model_group_affinity_config(self) -> "dict[str, list[str]] | None": # mutable-ok: the argument as passed return self._opt.model_group_affinity_config + @model_group_affinity_config.setter + def model_group_affinity_config(self, value: "dict[str, list[str]] | None") -> None: + self._opt.model_group_affinity_config = value + def add_optional_pre_call_checks(self, optional_pre_call_checks: "Sequence[object] | None") -> None: self._opt.add(optional_pre_call_checks) @@ -648,3 +676,9 @@ __all__ = [ "filtered_ids", "register_router_selector", ] + + +def _router_alive(owner: object) -> bool: + """A core whose public ``Router`` was dropped holds no check: it only waits for the collector.""" + handle: Final = vars(owner).get("_core_owner") if hasattr(owner, "__dict__") else None + return not callable(handle) or handle() is not None diff --git a/litellm/router_rust/steps/precall.py b/litellm/router_rust/steps/precall.py index 156c7b5e20f..227bb8de66f 100644 --- a/litellm/router_rust/steps/precall.py +++ b/litellm/router_rust/steps/precall.py @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.token_counter import ( single_encoding_openai_model, split_embedding_input_tokens, ) +from litellm.router_attempt import for_requested_model from litellm.router_python_host import python_host from litellm.router_strategy.complexity_router.context_compaction import compaction_pending from litellm.types.router import LiteLLM_Params @@ -229,10 +230,15 @@ class PreCallFacts: return self._has_input and any(limit is not None for limit in self._windows.values()) def _tokenizer_model(self) -> str: - named: Final = self._registry._core_deployments_named(self._model) or () # pyright: ignore[reportPrivateUsage] # the host registry's index + # A wildcard deployment serves the requested model (`PatternMatchRouter.route`). + named: Final = tuple( + for_requested_model(record, self._name) + for record in self._registry._core_deployments_named(self._model) or () # pyright: ignore[reportPrivateUsage] # the host registry's index + ) return ( single_encoding_openai_model( - served_openai_models(named) or served_openai_models(self._record(id) for id in self._candidates) + served_openai_models(named) + or served_openai_models(for_requested_model(self._record(id), self._name) for id in self._candidates) ) or "" ) diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py index a83fbe29996..63d15a1a80a 100644 --- a/litellm/router_rust/steps/prerouting.py +++ b/litellm/router_rust/steps/prerouting.py @@ -65,6 +65,7 @@ from litellm.types.utils import ( from litellm.utils import provider_rejectable_params if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache from litellm.router_python_host import PythonHost from litellm.router_rust.surface.registry import HostRegistry from litellm.router_strategy.auto_router.auto_router import PreRoutingHookResponse @@ -406,6 +407,13 @@ class PreRoutingStrategies: if router._uses_deployment_pin: # pyright: ignore[reportPrivateUsage] # the Python router reads the same flag self._ensure_deployment_affinity_callback() + def hand_over_affinity(self, adopt: Callable[[object], None]) -> None: + """Give the affinity check registered while the deployments were built, before the optional + checks existed, to those checks: they own it from then on, as ``optional_callbacks``.""" + if self._affinity is not None: + adopt(self._affinity) + self._affinity = None + def _ensure_deployment_affinity_callback(self) -> None: """``_ensure_deployment_affinity_callback``: the affinity check (every global flag off) a complexity router's session pin needs, registered once.""" @@ -568,7 +576,9 @@ def _record_routing_decision( ) -def _as_output_cap(value: object) -> int | None: +def as_output_cap(value: object) -> int | None: + """A client-sent output cap coerced to an int: ints, floats and numeric strings, never bools + or negatives.""" if isinstance(value, bool) or not isinstance(value, int | float | str): return None try: @@ -578,7 +588,8 @@ def _as_output_cap(value: object) -> int | None: return cap if cap >= 0 else None -def _stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None: +def stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None: + """None means unstamped; an empty mapping means the original caller supplied no limit.""" stamped: Final = next( ( bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) @@ -593,7 +604,7 @@ def _stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str { str(carrier): cap for carrier, value in stamped.items() - if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := _as_output_cap(value)) is not None + if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None } ) @@ -601,7 +612,7 @@ def _stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str def _restore_client_output_ceiling( request_kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs, restored in place ) -> None: - callers_ceiling: Final = _stamped_client_ceiling(request_kwargs) + callers_ceiling: Final = stamped_client_ceiling(request_kwargs) if callers_ceiling is None: return for carrier in OUTPUT_TOKEN_CEILING_PARAMS: @@ -710,6 +721,10 @@ class PreRouter: self._ctx: Final = ctx self._session_cache: Final = DualCache(redis_cache=routing_cache.redis_cache, in_memory_cache=InMemoryCache()) + def attach_redis_cache(self, redis_cache: "RedisCache") -> None: + """``update_redis_cache``: the Claude Code session bindings are shared through it from now on.""" + self._session_cache.attach_redis_cache(redis_cache) + # -- the step -------------------------------------------------------------------------------- async def pre_route( @@ -1104,4 +1119,4 @@ def _allowlist(params: Mapping[str, object]) -> frozenset[str]: return frozenset(entry for entry in declared if isinstance(entry, str)) -__all__ = ["PreRouter", "PreRoutingStrategies", "is_strategy_marker"] +__all__ = ["PreRouter", "PreRoutingStrategies", "as_output_cap", "is_strategy_marker", "stamped_client_ceiling"] diff --git a/litellm/router_rust/surface/health.py b/litellm/router_rust/surface/health.py index 5547471855a..0653518d0d7 100644 --- a/litellm/router_rust/surface/health.py +++ b/litellm/router_rust/surface/health.py @@ -255,6 +255,17 @@ class RustHealthCallbacks: deployment: Final = self._deployment_of(deployment_id) return deployment is not None and _authenticates_per_caller(deployment.litellm_params) + def ran_as_request_deployment(self, logged_id: object, deployment_id: object) -> bool: + """Whether a call against ``deployment_id`` ran under credentials the request supplied + (``_handle_clientside_credential``): as ``logged_id``, a request-scoped deployment the router + does not hold, so its failure is that id's, never the shared deployment's.""" + return ( + logged_id is not None + and deployment_id is not None + and str(logged_id) != str(deployment_id) + and self._deployment_of(logged_id) is None + ) + def failure_exits( self, exception: object, logged: Mapping[str, object], deployment_id: object ) -> Mapping[str, bool]: @@ -268,6 +279,7 @@ class RustHealthCallbacks: is_caller_scoped_auth_failure, is_caller_timeout_408, is_failure_outside_selected_deployment, + is_response_lookup_not_found, is_video_lookup_not_found, ) @@ -277,6 +289,7 @@ class RustHealthCallbacks: is_failure_outside_selected_deployment(exception if isinstance(exception, BaseException) else None) or (isinstance(exception, Exception) and is_background_response_poll_not_found(exception, litellm_params)) or is_video_lookup_not_found(exception, logged) + or is_response_lookup_not_found(exception, logged) or is_caller_timeout_408(logged, status) # pyright: ignore[reportArgumentType] # the status as raised or is_caller_scoped_auth_failure(self._deployment_of(deployment_id), status) # pyright: ignore[reportArgumentType] # see above ) @@ -498,16 +511,17 @@ class RustHealthCallbacks: completion_response: object, start_time: object, end_time: object, - ) -> None: - """Count a success and the call's tokens and request against a tracked deployment.""" - await self.report_stream_success(kwargs, completion_response) + ) -> str | None: + """Count a success and the call's tokens and request against a tracked deployment; returns + the counter's key, as the sync callback does.""" + return await self.report_stream_success(kwargs, completion_response) async def report_stream_success( self, kwargs: Mapping[str, object], completion_response: object, callbacks: Mapping[str, bool] | None = None, - ) -> None: + ) -> str | None: """``deployment_callback_on_success``, for a stream that ended. The re-pin is the host's half of that callback, so it reads ``callbacks`` here.""" from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type @@ -528,7 +542,7 @@ class RustHealthCallbacks: total: Final = float(tokens) if isinstance(tokens, int | float) else 0.0 counted: Final = get_counted_usage_tokens(as_str_mapping(kwargs.get("litellm_params")) or {}) usage_model: Final = target[1].get("deployment") - await self._drive_cache_steps( + finished: Final = await self._drive_cache_steps( self._handle.report_success( target[0], total if counted is None else max(0.0, total - counted), @@ -538,6 +552,8 @@ class RustHealthCallbacks: callbacks, ) ) + key: Final = finished.get("key") + return key if isinstance(key, str) else None async def count_realtime_session_success(self, deployment_id: str) -> None: """``increment_deployment_successes_for_current_minute`` for a realtime session that ended diff --git a/litellm/router_rust/surface/model_list.py b/litellm/router_rust/surface/model_list.py index 405a4b47d7b..bb7fb53acea 100644 --- a/litellm/router_rust/surface/model_list.py +++ b/litellm/router_rust/surface/model_list.py @@ -187,10 +187,16 @@ class HostedModelList(list[Entry]): self, handle: Callable[[], "RouterCoreHandle | None"], apply: Callable[[Mapping[str, object]], None], + live: Callable[[object], Mapping[str, object]] = lambda _: {}, + live_info: Callable[[object], Mapping[str, object]] = lambda _: {}, ) -> None: super().__init__() self._handle: Final = handle self._apply: Final = apply + # A deployment's litellm_params values its record carries as text, by id (``live_params``). + self._live: Final = live + # A deployment's model_info values its record reads back changed (non-string keys), by id. + self._live_info: Final = live_info # The handle and registry version the contents were last read at. self._read_at: tuple[object, int] | None = None self._stable: _Stable = {} @@ -225,7 +231,7 @@ class HostedModelList(list[Entry]): keys, counts = _keys(records, self._id_counts if appended else {}) entries: Final[list[Entry]] = [] # mutable-ok: filled here for key, record in zip(keys, records, strict=True): - typed = typed_deployment(_record(record)) + typed = self._with_live(typed_deployment(_record(record))) held = self._stable.get(key) if key is not None else None kept = held is not None and held[0] == typed entry = held[1] if held is not None and kept else _HostedEntry(typed, self._entry_written) @@ -236,6 +242,18 @@ class HostedModelList(list[Entry]): self._id_counts = counts return entries + def _with_live(self, record: Mapping[str, object]) -> Mapping[str, object]: + params: Final = as_str_mapping(record.get("litellm_params")) + info: Final = as_str_mapping(record.get("model_info")) + model_id: Final = (info or {}).get("id") + live: Final = self._live(model_id) if params is not None else {} + live_info: Final = self._live_info(model_id) if info is not None else {} + if live: + record = {**record, "litellm_params": {**(params or {}), **live}} # rebind-ok: given its live values + if live_info: + record = {**record, "model_info": {**(info or {}), **live_info}} # rebind-ok: see above + return record + # -- a build ---------------------------------------------------------------------------------- def begin_build(self, *, replacing: bool) -> None: diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index a55a1cfc6d6..75d30ebb2be 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -42,6 +42,7 @@ router and every ``registry()`` query dispatches the way it does on the Python c import copy import dataclasses +import json import os import re import threading @@ -50,7 +51,7 @@ from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast # noqa: TID251 # entries as a caller wrote them -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from litellm import exceptions as litellm_exceptions from litellm import get_secret_str @@ -93,7 +94,7 @@ from litellm.litellm_core_utils.routing_groups import apply_routing_group_priori from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.responses.bridge_check import keep_registered_mode -from litellm.router_api import RouterRegistry +from litellm.router_api import RouterRegistry, deployment_id from litellm.router_rust.invoke.shaping import ROUTER_CLIENTS_TARGET from litellm.router_rust.surface.indexes import HostedIndexes from litellm.router_rust.surface.model_info import HostModelInfo @@ -579,6 +580,34 @@ def _index_key(value: object) -> str | None: return value if isinstance(value, str) else None +def _round_trips(value: object) -> bool: + """Whether JSON reads ``value`` back equal to itself.""" + try: + loaded: Final[object] = json.loads(json.dumps(value, default=str)) # pyright: ignore[reportAny] # plain JSON + except (TypeError, ValueError): + return False + return loaded == value + + +def _has_non_str_keys(value: object) -> bool: + """Whether ``value`` holds a dict key that is not a string, which a JSON record reads back as one.""" + pending: list[object] = [value] + seen: set[int] = set() + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + if isinstance(current, Mapping): + for key, item in current.items(): + if not isinstance(key, str): + return True + pending.append(item) + elif isinstance(current, list | tuple): + pending.extend(current) + return False + + class HostRegistry(HostModelInfo): """The live model_list and its indexes, kept the way the Python router core keeps them.""" @@ -597,7 +626,13 @@ class HostRegistry(HostModelInfo): self._deployment_index_lock: Final = threading.RLock() self._core_model_group_alias: dict[str, str | RouterModelGroupAliasItem] = model_group_alias # mutable-ok: the live alias map # fmt: skip # The live model_list: router-core's deployments read back, every write a changeset. - self._core_model_list: Final[HostedModelList] = HostedModelList(self._live_handle, self._hand_changes) + self._core_model_list: Final[HostedModelList] = HostedModelList( + self._live_handle, self._hand_changes, self.live_params, self.live_info + ) + # The litellm_params values a deployment's JSON cannot carry (an Exception or model object as + # ``mock_response``, a callable ``azure_ad_token_provider``), by deployment id: an attempt is called with them, as Python's list held them. + self._live_params: Final[dict[str, Mapping[str, object]]] = {} # mutable-ok: written as deployments register + self._live_info: Final[dict[str, Mapping[str, object]]] = {} # mutable-ok: written as deployments register # The indexes the registry hands out: router-core's, read back, every write a changeset. self._indexes: Final = HostedIndexes(self._live_handle, self._hand_changes) self._core_deployment_latency_map: dict[str, float] = {} # mutable-ok: the live latency map @@ -854,6 +889,7 @@ class HostRegistry(HostModelInfo): self.build_runner(entries, build.path, resolve_secrets, install, partial(self._answer_shown, build)) finally: self._core_model_list.end_build() + self._drop_removed_live_values() if build.failure is not None: # router-core served what it built before the entry that raised (`set_model_list` # cleared the list first), so the router is left as the Python core leaves it. @@ -872,9 +908,11 @@ class HostRegistry(HostModelInfo): """The answer to one of router-core's build steps. A ``set_model_list`` deployment that fails is left out when the router ignores invalid deployments (``_create_deployment``), and retried on the next price data reload when its provider did not resolve.""" - if build.failure is not None: - return {"invalid": True} kind: Final = step["kind"] + if build.failure is not None: + # The build already failed: every later step is skipped, a secrets step unresolved. + refs_left: Final = step.get("refs") + return [None] * len(refs_left) if kind == "secrets" and isinstance(refs_left, list) else {"invalid": True} # pyright: ignore[reportUnknownArgumentType] # only the count is read index: Final = int(str(step["index"])) refs: Final = step.get("refs") if kind == "secrets" and isinstance(refs, list) and not any(_ref_env_only(ref) for ref in refs): @@ -969,7 +1007,13 @@ class HostRegistry(HostModelInfo): routed: Final = as_str_mapping(step["entry"]) or {} routed_info: Final = as_str_mapping(routed.get("model_info")) or {} if "id" not in entry.model_info: - entry.model_info["id"] = routed_info.get("id") + # router-core hashes the JSON form; a value JSON reads back changed hashes as the + # Python core hashes it (`deployment_id` over the params as given). + entry.model_info["id"] = ( + routed_info.get("id") + if _round_trips(entry.litellm_params) + else deployment_id(entry.model_name, entry.litellm_params) + ) if entry.per_organization: entry.litellm_params["organization"] = (as_str_mapping(routed.get("litellm_params")) or {}).get( "organization" @@ -1159,6 +1203,40 @@ class HostRegistry(HostModelInfo): deployment.litellm_params.rpm = getattr(deployment, "rpm") # noqa: B009 # an extra field of the entry # rebind-ok: the core sets it on the deployment it adds if deployment.litellm_params.tpm is None and getattr(deployment, "tpm", None) is not None: deployment.litellm_params.tpm = getattr(deployment, "tpm") # noqa: B009 # an extra field of the entry # rebind-ok: see above + live: Final = { + key: value + for key, value in { + **deployment.litellm_params.__dict__, + **(deployment.litellm_params.__pydantic_extra__ or {}), + }.items() + # Values JSON cannot carry: errors, callables, a str subclass that marks the value (an + # admin's secret reference, which only the deployment's own key may resolve), and any + # other value JSON reads back changed (a set, a tuple, a datetime), which the Python + # core keeps as given. A model is read back as its dump, as on the Python core. + if isinstance(value, BaseException) + or callable(value) + or (isinstance(value, str) and type(value) is not str) + or not (isinstance(value, BaseModel) or _round_trips(value)) + } + # model_info values JSON reads back changed: a dict with keys that are not strings. + live_info: Final = { + key: value + for key, value in { + **deployment.model_info.__dict__, + **(deployment.model_info.__pydantic_extra__ or {}), + }.items() + if _has_non_str_keys(value) + } + model_id: Final = deployment.model_info.id + if model_id is not None: + if live: + self._live_params[model_id] = live + else: + self._live_params.pop(model_id, None) + if live_info: + self._live_info[model_id] = live_info + else: + self._live_info.pop(model_id, None) data_sources: Final = deployment.litellm_params.get("dataSources", []) or [] for source, key, value in data_source_values: data_sources[source]["parameters"][key] = value @@ -1270,6 +1348,17 @@ class HostRegistry(HostModelInfo): raise RuntimeError("the host registry writes deployments through its Rust core") self.changes_runner(changes) self._core_model_list.refresh() + self._drop_removed_live_values() + + def _drop_removed_live_values(self) -> None: + """Forget the live values of ids no deployment carries any more, so a router that rotates + ids does not keep the removed deployments' errors and callables.""" + handle: Final = self._live_handle() + if handle is None or not (self._live_params or self._live_info): + return + for held in (self._live_params, self._live_info): + for model_id in [id for id in held if handle.deployment(id) is None]: + del held[model_id] def refresh(self) -> None: """Read router-core's list again after a write (only what was appended, when that is all @@ -1447,6 +1536,14 @@ class HostRegistry(HostModelInfo): idx: Final[object] = None if handle is None or key is None else handle.id_position(key) return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it + def live_params(self, model_id: object) -> Mapping[str, object]: + """The litellm_params values of deployment ``model_id`` that its JSON record cannot carry.""" + return self._live_params.get(model_id, {}) if isinstance(model_id, str) else {} + + def live_info(self, model_id: object) -> Mapping[str, object]: + """The model_info values a deployment's JSON record reads back changed, by id.""" + return self._live_info.get(model_id, {}) if isinstance(model_id, str) else {} + def _core_get_deployment(self, model_id: str) -> Deployment | None: entry: Final = self._core_get_model_info(model_id) return None if entry is None else Deployment(**entry) # pyright: ignore[reportArgumentType] # a stored entry @@ -1541,8 +1638,14 @@ class HostRegistry(HostModelInfo): for name in (self._core_routing_groups or {}) if (group := self.get_routing_group(name)) is not None ] + # The provider is inferred only for a name no deployment or alias serves: inferring it for an + # authenticating provider's name would start its sign-in (`_get_wildcard_deployments`). wildcard: Final = ( - [model_name, f"{inferred_provider(model_name)}/{model_name}"] if model_name is not None else [] + [] + if model_name is None + else [model_name] + if model_name in alias_map or model_name in self._core_model_names + else [model_name, f"{inferred_provider(model_name)}/{model_name}"] ) return {"aliases": aliases, "groups": groups, "wildcard": wildcard, "only": only} @@ -1561,7 +1664,8 @@ class HostRegistry(HostModelInfo): info: Final = _model_info_of(model) if team_id is not None and info.get("team_id") == team_id and model_name == info.get("team_public_model_name"): return True - if model_name is not None and model.get("model_name") == model_name: # pyright: ignore[reportUnnecessaryComparison] # the core checks it + # A record with no model_name raises KeyError, as the Python core's lookup does. + if model_name is not None and model["model_name"] == model_name: # pyright: ignore[reportUnnecessaryComparison] # the core checks it model_team_id: Final = info.get("team_id") return team_id is None or model_team_id is None or model_team_id == team_id return False diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 03485ebd5f8..4c918776223 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -108,6 +108,26 @@ def is_video_lookup_not_found(exception: object, failure_kwargs: Mapping[str, ob return getattr(exception, "status_code", None) == 404 and failure_kwargs.get("call_type") in VIDEO_ID_CALL_TYPES +# The calls on an existing response id. ``compact_responses`` is left out: it sends a model and input, +# so its 404 can mean a missing model rather than a missing response. +RESPONSE_ID_CALL_TYPES: Final = frozenset( + { + "aget_responses", + "get_responses", + "adelete_responses", + "delete_responses", + "acancel_responses", + "cancel_responses", + "alist_input_items", + } +) + + +def is_response_lookup_not_found(exception: object, failure_kwargs: Mapping[str, object]) -> bool: + """Whether a call on an existing response id failed with a provider 404: the response is not on that account.""" + return getattr(exception, "status_code", None) == 404 and failure_kwargs.get("call_type") in RESPONSE_ID_CALL_TYPES + + _EXCEPTION_POLICY_FIELDS: Final[tuple[tuple[type, str], ...]] = ( # ContentPolicyViolationError subclasses BadRequestError, so it must be checked first. (litellm.ContentPolicyViolationError, "ContentPolicyViolationErrorAllowedFails"), diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 81012d38b65..e5d944b3aef 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -163,6 +163,9 @@ class PromptCachingDeploymentCheck(CustomLogger): return healthy_deployments + def exempt_from_message_redaction(self) -> bool: + return True + async def async_log_success_event( self, kwargs: _PromptCachingSuccessEventKwargs, diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py deleted file mode 100644 index 17a56e62e41..00000000000 --- a/litellm/router_utils/search_api_router.py +++ /dev/null @@ -1,252 +0,0 @@ -""" -Router utilities for Search API integration. - -Handles search tool selection, load balancing, and fallback logic for search requests. -""" - -import asyncio -import random -import traceback -from collections.abc import Callable -from functools import partial -from types import MappingProxyType -from typing import TYPE_CHECKING, Final - -from litellm._logging import verbose_router_logger -from litellm.router_api import ChangeSet -from litellm.router_deprecated import renamed_router_param -from litellm.types.router import SearchToolLiteLLMParams, SearchToolTypedDict - -if TYPE_CHECKING: - from litellm.router_api import RoutingContext - - -class SearchAPIRouter: - """ - Static utility class for routing search API calls through the LiteLLM router. - - Provides methods for search tool selection, load balancing, and fallback handling. - """ - - @staticmethod - def _resolve_search_provider_credentials( - *, - tool_litellm_params: SearchToolLiteLLMParams, - ) -> tuple[str | None, str | None]: - """ - Resolve search provider credentials from tool configuration ONLY. - - Credentials are stored only in search_tool.litellm_params, never in team/key metadata. - This ensures secrets are not exposed in team/key API responses. - - Args: - tool_litellm_params: Search tool litellm_params with credentials - - Returns: - Tuple of (api_key, api_base) from tool configuration - """ - resolved_api_key: Final[str | None] = tool_litellm_params.get("api_key") - resolved_api_base: Final[str | None] = tool_litellm_params.get("api_base") - - return resolved_api_key, resolved_api_base - - @staticmethod - @renamed_router_param("router_instance") - async def update_router_search_tools(ctx: "RoutingContext", search_tools: list) -> None: - """ - Update the router with search tools from the database. - - This method is called by a cron job to sync search tools from DB to router. - - Args: - ctx: The routing context of the router to update - search_tools: List of search tool configurations from the database - """ - try: - verbose_router_logger.debug("Adding %s search tools to router", len(search_tools)) - - # Convert search tools to the format expected by the router - router_search_tools: Final[list] = [] - for tool in search_tools: - # Create dict that matches SearchToolTypedDict structure - router_search_tool: SearchToolTypedDict = { - "search_tool_id": tool.get("search_tool_id"), - "search_tool_name": tool.get("search_tool_name"), - "litellm_params": tool.get("litellm_params", {}), - "search_tool_info": tool.get("search_tool_info"), - } - router_search_tools.append(router_search_tool) - - # Update the router's search_tools list - ctx.apply(ChangeSet(assign={"search_tools": router_search_tools})) - - verbose_router_logger.info("Successfully updated router with %s search tool(s)", len(router_search_tools)) - - except Exception as e: - verbose_router_logger.exception("Error updating router with search tools: %s", e) - raise e - - @staticmethod - @renamed_router_param("router_instance") - def get_matching_search_tools( - ctx: "RoutingContext", - search_tool_name: str, - ) -> list: - """ - Get all search tools matching the given name. - - Args: - ctx: The routing context of the router - search_tool_name: Name of the search tool to find - - Returns: - List of matching search tool configurations - - Raises: - ValueError: If no matching search tools are found - """ - matching_tools: Final = [ - tool for tool in ctx.settings().search_tools if tool.get("search_tool_name") == search_tool_name - ] - - if not matching_tools: - raise ValueError(f"Search tool '{search_tool_name}' not found in router.search_tools") - - return matching_tools - - @staticmethod - @renamed_router_param("router_instance", new="litellm_routing_ctx") - async def async_search_with_fallbacks( - litellm_routing_ctx: "RoutingContext", - original_function: Callable, - **kwargs, - ): - """ - Helper function to make a search API call through the router with load balancing and fallbacks. - Reuses the router's retry/fallback infrastructure. - - Args: - litellm_routing_ctx: The routing context of the router - original_function: The original litellm.asearch function - **kwargs: Search parameters including search_tool_name, query, etc. - - Returns: - SearchResponse from the search API - """ - ctx: Final = litellm_routing_ctx - try: - search_tool_name: Final = kwargs.get("search_tool_name", kwargs.get("model")) - - if not search_tool_name: - raise ValueError("search_tool_name or model parameter is required for search") - - # Set up kwargs for the fallback system - kwargs["model"] = search_tool_name # Use model field for compatibility with fallback system - kwargs["original_generic_function"] = original_function - kwargs["original_function"] = partial( - SearchAPIRouter.async_search_with_fallbacks_helper, - litellm_routing_ctx=ctx, - ) - - # Update kwargs before fallbacks (for logging, metadata, etc) - ctx.update_kwargs_before_fallbacks( - model=search_tool_name, - kwargs=kwargs, - metadata_variable_name="litellm_metadata", - ) - - available_search_tool_names: Final = [tool.get("search_tool_name") for tool in ctx.settings().search_tools] - verbose_router_logger.debug( - "Inside SearchAPIRouter.async_search_with_fallbacks() - search_tool_name: %s, Available Search Tools: %s", - search_tool_name, - available_search_tool_names, - ) - - # Use the existing retry/fallback infrastructure - response: Final = await ctx.with_fallbacks(**kwargs) - return response - - except Exception as e: - from litellm.router_utils.handle_error import send_llm_exception_alert - - asyncio.create_task( - send_llm_exception_alert( - ctx=ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - @staticmethod - @renamed_router_param("router_instance", new="litellm_routing_ctx") - async def async_search_with_fallbacks_helper( - litellm_routing_ctx: "RoutingContext", - model: str, - original_generic_function: Callable, - **kwargs, - ): - """ - Helper function for search API calls - selects a search tool and calls the original function. - Called by async_function_with_fallbacks for each retry attempt. - - Args: - litellm_routing_ctx: The routing context of the router - model: The search tool name (passed as model for compatibility) - original_generic_function: The original litellm.asearch function - **kwargs: Search parameters - - Returns: - SearchResponse from the selected search provider - """ - ctx: Final = litellm_routing_ctx - search_tool_name: Final = model # model field contains the search_tool_name - - try: - # Find matching search tools - matching_tools: Final = SearchAPIRouter.get_matching_search_tools( - ctx=ctx, - search_tool_name=search_tool_name, - ) - - # Simple random selection for load balancing across multiple providers with same name - # For search tools, we use simple random choice since they don't have TPM/RPM constraints - selected_tool: Final = random.choice(matching_tools) - - # Extract search provider and other params from litellm_params - litellm_params: Final = selected_tool.get("litellm_params", {}) - search_provider: Final = litellm_params.get("search_provider") - if not search_provider: - raise ValueError(f"search_provider not found in litellm_params for search tool '{search_tool_name}'") - - api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( - tool_litellm_params=litellm_params, - ) - protected_params: Final = frozenset(("search_provider", "api_key", "api_base")) - search_params: Final = MappingProxyType( - { - key: value - for params in (litellm_params, kwargs) - for key, value in params.items() - if key not in protected_params and value is not None - } - ) - - verbose_router_logger.debug("Selected search tool with provider: %s", search_provider) - - # Call the original search function with the provider config - response: Final = await original_generic_function( - search_provider=search_provider, - api_key=api_key, - api_base=api_base, - **search_params, - ) - - return response - - except Exception as e: - verbose_router_logger.error( - "Error in SearchAPIRouter.async_search_with_fallbacks_helper for %s: %s", search_tool_name, e - ) - raise e diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index a12f79c6c8a..9ba8cbc35a7 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -23,7 +23,6 @@ class Route(str, Enum): OCR = "ocr" TOKEN_COUNTER = "token_counter" TOKENIZER = "tokenizer" - ROUTER = "router" class Target(str, Enum): @@ -225,7 +224,6 @@ RULES: Final[Rules] = ( RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY), - RouteRule(Route.ROUTER, Rollout.RUST_OPT_IN), RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), ) diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index c4f7a2d8a7b..4d808742434 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -176,11 +176,16 @@ class Status1(Enum): incomplete = "incomplete" budget_exceeded = "budget_exceeded" queued = "queued" + requires_continuation = "requires_continuation" class InteractionStatusUpdate(LiteLLMBaseModel): interaction_id: str | None = None status: Status1 | None = None + continuation_token: str | None = Field( + None, + description="Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + ) event_type: Literal["interaction.status_update"] = "interaction.status_update" event_id: str | None = Field( None, @@ -343,6 +348,7 @@ class Status3(Enum): INCOMPLETE = "incomplete" BUDGET_EXCEEDED = "budget_exceeded" QUEUED = "queued" + REQUIRES_CONTINUATION = "requires_continuation" class ModelOption(RootModel[str]): @@ -604,6 +610,9 @@ class ToolChoice(RootModel[ToolChoiceType | ToolChoiceConfig]): class Usage(LiteLLMBaseModel): total_input_tokens: int | None = Field(None, description="Number of tokens in the prompt (context).") + initial_input_tokens: int | None = Field( + None, description="Number of tokens in the prompt of the first hop of a long decode." + ) input_tokens_by_modality: list[ModalityTokens] | None = Field( None, description="A breakdown of input token usage by modality." ) @@ -625,6 +634,7 @@ class Usage(LiteLLMBaseModel): None, description="A breakdown of tool-use token usage by modality." ) total_reasoning_tokens: int | None = Field(None, description="Number of tokens of thoughts for thinking models.") + total_thought_tokens: int | None = Field(None, description="Number of tokens of thoughts for thinking models.") total_tokens: int | None = Field( None, description="Total token count for the interaction request (prompt + responses + other\ninternal tokens).", @@ -787,6 +797,10 @@ class Interaction(LiteLLMBaseModel): description="The mime type of the response. This is required if response_format is set.", ) previous_interaction_id: str | None = Field(None, description="The ID of the previous interaction, if any.") + continuation_token: str | None = Field( + None, + description="Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + ) input: str | list[Content] | list[Turn] | Content | None = Field( None, description="The inputs for the interaction." ) @@ -843,6 +857,10 @@ class CreateModelInteractionParams(LiteLLMBaseModel): description="The mime type of the response. This is required if response_format is set.", ) previous_interaction_id: str | None = Field(None, description="The ID of the previous interaction, if any.") + continuation_token: str | None = Field( + None, + description="Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + ) input: str | list[Content] | list[Turn] | Content = Field(..., description="The inputs for the interaction.") generation_config: GenerationConfig | None = Field( None, @@ -894,6 +912,10 @@ class CreateAgentInteractionParams(LiteLLMBaseModel): description="The mime type of the response. This is required if response_format is set.", ) previous_interaction_id: str | None = Field(None, description="The ID of the previous interaction, if any.") + continuation_token: str | None = Field( + None, + description="Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + ) input: str | list[Content] | list[Turn] | Content = Field(..., description="The inputs for the interaction.") agent_config: DynamicAgentConfig | DeepResearchAgentConfig | None = Field( None, description="Configuration for the agent.", discriminator="type" @@ -1080,6 +1102,7 @@ class InteractionsAPIResponse(BaseLiteLLMOpenAIResponseObject): # New schema field (Api-Revision: 2026-05-20). steps: list[dict[str, Any]] | None = None usage: dict[str, Any] | None = None + continuation_token: str | None = None _hidden_params: dict = PrivateAttr(default_factory=dict) @@ -1126,6 +1149,7 @@ class InteractionsAPIStreamingResponse(BaseLiteLLMOpenAIResponseObject): step: dict[str, Any] | None = None interaction_id: str | None = None interaction: dict[str, Any] | None = None + continuation_token: str | None = None _hidden_params: dict = PrivateAttr(default_factory=dict) diff --git a/litellm/types/proxy/claude_code_endpoints.py b/litellm/types/proxy/claude_code_endpoints.py index 26c610a4f11..5d8479b8ce1 100644 --- a/litellm/types/proxy/claude_code_endpoints.py +++ b/litellm/types/proxy/claude_code_endpoints.py @@ -2,6 +2,7 @@ Claude Code Marketplace endpoint types for LiteLLM Proxy """ +from datetime import datetime from typing import Final, Literal, TypeAlias from pydantic import AliasChoices, ConfigDict, Field, TypeAdapter, model_validator @@ -32,7 +33,7 @@ class PluginMetadata(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - version: str | None = Field("1.0.0", description="Semantic version") + version: str | None = Field("1.0.0", description="Semantic version, defaults to 1.0.0 on creation") description: str | None = Field(None, description="Plugin description") author: PluginAuthor | None = Field(None, description="Plugin author") homepage: str | None = Field(None, description="Plugin homepage URL") @@ -105,11 +106,14 @@ class UpdatePluginRequest(PluginSpec): The plugin name is the resource identity and is supplied as the path parameter, so it cannot be changed here. This is a full replace: omitted - fields reset to their defaults, so version is cleared rather than - defaulting to the create-time "1.0.0". + fields reset to their defaults, except version which automatically increments + the published patch version when omitted or null. """ - version: str | None = Field(None, description="Semantic version; cleared if omitted") + version: str | None = Field(None, description="Explicit higher SemVer, or omit/null for an automatic patch bump") + expected_updated_at: datetime | None = Field( + None, description="Published updated_at read by the editor. A stale value returns 409", exclude=True + ) class PluginResponse(LiteLLMBaseModel): diff --git a/litellm/types/proxy/skill_hub.py b/litellm/types/proxy/skill_hub.py index bf28ce592b3..d575defcda7 100644 --- a/litellm/types/proxy/skill_hub.py +++ b/litellm/types/proxy/skill_hub.py @@ -1,8 +1,10 @@ from collections.abc import Mapping -from typing import Literal +from datetime import datetime +from typing import Literal, TypeAlias from pydantic import ConfigDict, Field +from litellm.models.skill_submission import SubmissionOperation, SubmissionStatus from litellm.types.llms.base import LiteLLMBaseModel @@ -13,10 +15,17 @@ class SkillHubSkill(LiteLLMBaseModel): name: str description: str | None = None version: str | None = None + updated_at: datetime | None = None visibility: Literal["public", "private"] source: Mapping[str, str] download_url: str | None = None content_url: str | None = None + domain: str | None = None + namespace: str | None = None + keywords: tuple[str, ...] | None = None + download_count: int | None = Field( + None, description="Gateway ZIP downloads, included in proxy-admin catalog listings" + ) class SkillHubList(LiteLLMBaseModel): @@ -34,14 +43,17 @@ class ImportSkillRequest(LiteLLMBaseModel): skill: str | None = Field(None, description="Skill name to select when a repository contains several skills") ref: str = Field("HEAD", description="Git commit, tag or branch to import") path: str | None = Field(None, description="Repository-relative skill directory") - version: str | None = Field(None, description="Version label for the imported snapshot") + version: str | None = Field(None, description="Semantic version, or omit/null to start at 1.0.0") class SkillHubMetadata(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, extra="forbid") description: str | None = Field(None, max_length=1024) - version: str | None = None + version: str | None = Field(None, description="Explicit higher SemVer, or omit/null for an automatic patch bump") + expected_updated_at: datetime | None = Field( + None, description="Published updated_at read by the editor. A stale value returns 409", exclude=True + ) domain: str | None = None namespace: str | None = None keywords: tuple[str, ...] | None = None @@ -49,3 +61,105 @@ class SkillHubMetadata(LiteLLMBaseModel): class SkillHubPublication(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, extra="forbid") + + +class SkillHubVersion(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + version: str | None + description: str | None + manifest: Mapping[str, object] + published_at: datetime | None + published_by: str | None + content_kind: Literal["stored_archive", "source_reference"] + download_url: str | None + + +class SkillHubVersions(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + versions: tuple[SkillHubVersion, ...] + next_offset: int | None + + +class SkillHubAgentSkill(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + name: str + description: str + digest: str + origin: str | None + + +class SkillHubProposal(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + name: str + description: str | None + version: str | None + domain: str | None = None + namespace: str | None = None + keywords: tuple[str, ...] | None = None + source: Mapping[str, str] + agent_skill: SkillHubAgentSkill | None = None + + +class SkillHubPendingSkill(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + status: Literal["pending_review"] = "pending_review" + action: Literal["submitted"] = "submitted" + submission_id: str + proposal: SkillHubProposal + + +SkillHubMutationResponse: TypeAlias = SkillHubSkill | SkillHubPendingSkill + + +class SkillHubSubmission(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + skill_name: str + operation: SubmissionOperation + status: SubmissionStatus + proposal: SkillHubProposal + has_archive: bool + submitted_by: str + created_at: datetime + base_skill_id: str | None + base_updated_at: datetime | None + reviewed_by: str | None + reviewed_at: datetime | None + rejection_reason: str | None + + +class SkillHubSubmissionList(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + submissions: tuple[SkillHubSubmission, ...] + count: int + has_more: bool + + +class SkillHubSubmissionDetail(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + submission: SkillHubSubmission + current_skill: SkillHubSkill | None + + +class SkillHubReviewResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + status: Literal["success"] = "success" + action: Literal["approved", "rejected"] + submission: SkillHubSubmission + skill: SkillHubSkill | None + + +class SkillHubRejection(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + reason: str | None = Field(None, max_length=2000) diff --git a/litellm/types/proxy/skill_scanning.py b/litellm/types/proxy/skill_scanning.py new file mode 100644 index 00000000000..19577c8a9ac --- /dev/null +++ b/litellm/types/proxy/skill_scanning.py @@ -0,0 +1,91 @@ +from datetime import datetime +from typing import Annotated, Final, Literal +from urllib.parse import urlsplit + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class SkillScannerConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + endpoint: str = Field(max_length=2048, description="External scanner's authenticated POST endpoint") + api_key_env: str = Field( + pattern=r"^[A-Za-z_][A-Za-z0-9_]*$", description="Environment secret name for scanner auth" + ) + use_llm: bool = False + timeout_seconds: int = Field(default=180, ge=10, le=600) + + @field_validator("endpoint") + @classmethod + def validate_endpoint(cls, value: str) -> str: + parsed: Final = urlsplit(value) + if ( + not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.query + or parsed.fragment + or not ( + parsed.scheme == "https" + or (parsed.scheme == "http" and parsed.hostname in {"127.0.0.1", "::1", "localhost"}) + ) + ): + raise ValueError( + "Scanner endpoint must use HTTPS without credentials, query or fragment (HTTP loopback allowed)" + ) + return value + + +class SkillScanFinding(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + rule_id: str = Field(max_length=200) + severity: Literal["LOW", "MEDIUM", "HIGH", "CRITICAL"] + message: str = Field(max_length=8000) + file: str = Field(max_length=2048) + start_line: int | None = Field(default=None, ge=1) + explanation: str | None = Field(default=None, max_length=8000) + remediation: str | None = Field(default=None, max_length=8000) + + +class SkillScanReport(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + scanner_version: str = Field(min_length=1, max_length=100) + complete: bool + scan_mode: Literal["static-only", "static+llm"] + risk_score: int = Field(ge=0, le=100) + findings: tuple[SkillScanFinding, ...] = Field(max_length=500) + limitations: tuple[Annotated[str, Field(max_length=2000)], ...] = Field(default=(), max_length=100) + + +class SkillScan(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + status: Literal["running", "completed", "failed"] + requested_by: str + started_at: datetime + expires_at: datetime + completed_at: datetime | None = None + archive_sha256: str | None = None + report: SkillScanReport | None = None + error: str | None = None + + def at(self, now: datetime) -> "SkillScan": + if self.status == "running" and self.expires_at <= now: + return self.model_copy( + update={"status": "failed", "error": "Scan timed out or was interrupted. Try scanning again."} + ) + return self + + +class SkillScanRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class SkillScanResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + configured: bool + scan: SkillScan | None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ba29a67a36c..ca434154f68 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -703,6 +703,8 @@ class CallTypes(str, Enum): adelete_responses = "adelete_responses" cancel_responses = "cancel_responses" acancel_responses = "acancel_responses" + compact_responses = "compact_responses" + acompact_responses = "acompact_responses" alist_input_items = "alist_input_items" llm_passthrough_route = "llm_passthrough_route" allm_passthrough_route = "allm_passthrough_route" @@ -806,6 +808,8 @@ CallTypesLiteral = Literal[ "send_message", "aresponses", "responses", + "compact_responses", + "acompact_responses", "acreate_skill", "acreate_realtime_client_secret", "arealtime_calls", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 76fd85d001b..a29f97ee9ea 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4046,6 +4046,8 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_1hr": 2e-05, "cache_read_input_token_cost": 2.5e-07, + "deprecation_date": "2027-09-01", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "litellm_provider": "azure_ai", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -4068,7 +4070,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, "prompt_cache_min_tokens": 512, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models" + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" }, "azure_ai/claude-opus-5": { "deprecation_date": "2027-07-08", @@ -13458,13 +13460,16 @@ "supports_reasoning": true }, "azure_ai/mistral-document-ai-2512": { + "deprecation_date": "2027-02-15", + "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.003, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" ], - "source": "https://ai.azure.com/catalog/models/mistral-document-ai-2512" + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", + "successor_model": "azure_ai/mistral-ocr-4-0" }, "azure_ai/mistral-ocr-4-0": { "litellm_provider": "azure_ai", @@ -32320,7 +32325,9 @@ "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", "successor_model": "gemini/gemini-3.8-flash", "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, - "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 + "input_cost_per_token_above_200k_tokens_flex": 1.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, + "output_cost_per_token_above_200k_tokens_flex": 7.5e-06 }, "gemini/gemini-2.5-computer-use-preview-10-2025": { "deprecation_date": "2026-07-28", @@ -32811,7 +32818,7 @@ "output_cost_per_token": 9e-06, "output_cost_per_video_token": 1.75e-05, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_endpoints": [ "/v1beta/interactions" ], @@ -32833,7 +32840,7 @@ "tpm": 800000, "deprecation_date": "2026-10-22", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", - "successor_model": "gemini/gemini-omni-1.1-flash" + "successor_model": "gemini/gemini-omni-1.2-flash" }, "gemini/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, @@ -32899,7 +32906,9 @@ "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, - "output_cost_per_token_above_200k_tokens_batches": 9e-06 + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -32965,7 +32974,9 @@ "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, - "output_cost_per_token_above_200k_tokens_batches": 9e-06 + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -33613,14 +33624,14 @@ "gemini/veo-3.1-fast-generate-preview": { "deprecation_date": "2026-10-22", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", - "successor_model": "gemini/gemini-omni-1.1-flash", + "successor_model": "gemini/gemini-omni-1.2-flash", "litellm_provider": "gemini", "max_input_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.1, "output_cost_per_second_1080p": 0.12, "output_cost_per_second_4k": 0.3, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_modalities": [ "text" ], @@ -33631,13 +33642,13 @@ "gemini/veo-3.1-generate-preview": { "deprecation_date": "2026-10-22", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", - "successor_model": "gemini/gemini-omni-1.1-flash", + "successor_model": "gemini/gemini-omni-1.2-flash", "litellm_provider": "gemini", "max_input_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, "output_cost_per_second_4k": 0.6, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_modalities": [ "text" ], @@ -33648,13 +33659,13 @@ "gemini/veo-3.1-lite-generate-preview": { "deprecation_date": "2026-10-22", "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", - "successor_model": "gemini/gemini-omni-1.1-flash", + "successor_model": "gemini/gemini-omni-1.2-flash", "litellm_provider": "gemini", "max_input_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.05, "output_cost_per_second_1080p": 0.08, - "source": "https://ai.google.dev/gemini-api/docs/video", + "source": "https://ai.google.dev/gemini-api/docs/deprecations", "supported_modalities": [ "text" ], @@ -40887,6 +40898,22 @@ "output_cost_per_token": 4e-07, "supports_tool_choice": true }, + "jina-reranker-v3.5": { + "input_cost_per_token": 5e-08, + "litellm_provider": "jina_ai", + "max_input_tokens": 131072, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" + }, + "jina-reranker-m0": { + "input_cost_per_token": 5e-08, + "litellm_provider": "jina_ai", + "max_input_tokens": 10240, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://api.jina.ai/v1/models" + }, "jina-reranker-v2-base-multilingual": { "input_cost_per_token": 5e-08, "litellm_provider": "jina_ai", @@ -43765,7 +43792,8 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/zai-glm-5-3": { "cache_read_input_token_cost": 1.4e-07, @@ -45264,6 +45292,18 @@ "source": "https://tokenfactory.nebius.com/api/public/models_info", "cache_read_input_token_cost": 1.5e-07 }, + "nebius/moonshotai/Kimi-K3.1-Pro": { + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "nebius", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true, + "supports_prompt_caching": true, + "source": "https://tokenfactory.nebius.com/api/public/models_info" + }, "nebius/moonshotai/Kimi-K3": { "max_input_tokens": 1048576, "max_output_tokens": 1024000, @@ -60260,6 +60300,7 @@ "rpm": 10, "gemini_audio_only_live": true, "input_cost_per_second": 8.33333333333e-05, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_response_schema": false, "supports_reasoning": true }, @@ -67653,17 +67694,20 @@ "source": "https://deepinfra.com/pricing" }, "gemini/gemini-omni-1.1-flash": { + "deprecation_date": "2026-12-03", + "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", - "max_input_tokens": 1048576, + "max_input_tokens": 131072, "max_output_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, "output_cost_per_video_token": 1.75e-05, "rpm": 2000, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://ai.google.dev/gemini-api/docs/models/gemini-omni-flash", + "successor_model": "gemini/gemini-omni-1.2-flash", "supported_endpoints": [ "/v1beta/interactions" ], @@ -71179,6 +71223,8 @@ }, "gpt-5.5-cyber-preview": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-02-01", + "deprecation_source": "https://developers.openai.com/api/docs/deprecations", "input_cost_per_token": 1.25e-05, "litellm_provider": "openai", "mode": "chat", @@ -72138,8 +72184,10 @@ "output_cost_per_token": 4.5e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_audio_output": true, "tpm": 250000, "rpm": 10, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -72159,8 +72207,10 @@ "output_cost_per_token": 4.5e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supports_audio_input": true, + "supports_audio_output": true, "tpm": 250000, "rpm": 10, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -76457,6 +76507,33 @@ "supports_tool_choice": true, "supports_vision": true }, + "us-gov.xai.grok-4.7": { + "supports_regex_lookaround": false, + "supported_endpoints": ["/v1/chat/completions"], + "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, + "supports_bedrock_runtime_chat_completions_response_format": true, + "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_converse", + "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", + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "baseten/deepseek-ai/DeepSeek-V4.1-Flash-Fast": { "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 6e-07, @@ -78425,7 +78502,7 @@ "cache_creation_input_token_cost": 1e-05, "cache_creation_input_token_cost_above_1hr": 1.6e-05, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2027-10-01", + "deprecation_date": "2027-01-15", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "input_cost_per_token": 8e-06, "litellm_provider": "azure_ai", @@ -78435,7 +78512,7 @@ "output_cost_per_token": 4e-05, "prompt_cache_min_tokens": 512.0, "search_context_cost_per_query": {"search_context_size_high":0.02,"search_context_size_low":0.02,"search_context_size_medium":0.02}, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models", + "source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, @@ -82770,6 +82847,39 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true }, + "gpt-6-cyber": { + "cache_creation_input_token_cost": 2.5e-05, + "cache_creation_input_token_cost_above_12hr": 5e-05, + "cache_read_input_token_cost": 2e-06, + "input_cost_per_token": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.0001, + "source": "https://developers.openai.com/api/docs/models/gpt-6-cyber", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true + }, "gpt-6-astra-law": { "cache_read_input_token_cost": 1e-06, "input_cost_per_token": 1e-05, @@ -83130,6 +83240,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_prompt_caching": true, "supports_function_calling": true, "supports_response_schema": true, @@ -83146,6 +83261,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_prompt_caching": true, "supports_function_calling": true, "supports_response_schema": true, @@ -83155,6 +83275,7 @@ "litellm_provider": "empiriolabs", "mode": "chat", "max_input_tokens": 202745, + "max_output_tokens": 128000, "input_cost_per_token": 8.25e-07, "output_cost_per_token": 3.301e-06, "input_cost_per_token_above_32k_tokens": 1.1e-06, @@ -83178,6 +83299,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-2" @@ -83230,6 +83356,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-3" @@ -83245,6 +83376,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-3-flash" @@ -83342,6 +83478,28 @@ "supports_response_schema": true, "source": "https://empiriolabs.ai/models/kimi-k3" }, + "empiriolabs/kimi-k3-1-pro": { + "litellm_provider": "empiriolabs", + "mode": "chat", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 2e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1.8e-06, + "supports_system_messages": true, + "supports_vision": true, + "supports_reasoning": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_function_calling": true, + "source": "https://empiriolabs.ai/models/kimi-k3-1-pro" + }, "empiriolabs/mimo-v2-6-flash": { "litellm_provider": "empiriolabs", "mode": "chat", @@ -83721,6 +83879,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" }, @@ -83737,6 +83900,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-27b" }, @@ -83753,6 +83921,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" }, @@ -83769,6 +83942,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" }, @@ -83830,6 +84008,10 @@ "max_output_tokens": 65536, "input_cost_per_token": 2.9e-08, "output_cost_per_token": 2.87e-07, + "input_cost_per_token_above_128k_tokens": 1.15e-07, + "output_cost_per_token_above_128k_tokens": 1.147e-06, + "input_cost_per_token_above_256k_tokens": 1.72e-07, + "output_cost_per_token_above_256k_tokens": 1.72e-06, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -83852,6 +84034,11 @@ "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, "supports_audio_input": true, "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" }, @@ -83865,6 +84052,11 @@ "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, "supports_response_schema": true, "supports_audio_input": true, "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" @@ -83924,6 +84116,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-6-27b" }, @@ -84421,7 +84618,7 @@ "max_output_tokens": 131072, "input_cost_per_token": 1e-06, "output_cost_per_token": 2.7e-06, - "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost": 1e-07, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -85310,6 +85507,24 @@ "output_cost_per_token": 6.09e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/XiaomiMiMo/MiMo-V2.6-Flash": { + "cache_read_input_token_cost": 2.24e-09, + "input_cost_per_token": 1.12e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 2.24e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.6-Pro": { + "cache_read_input_token_cost": 2.88e-09, + "input_cost_per_token": 3.48e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6.96e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/deepseek-ai/DeepSeek-Prover-V2-671B": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -85434,6 +85649,15 @@ "output_cost_per_token": 3.168e-06, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/deepseek-ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 4.4e-07, @@ -85443,6 +85667,24 @@ "output_cost_per_token": 1.32e-06, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/deepseek-ai/deepseek-v4.1-flash-1017": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/gmicloud/Ripple-65B-Preview": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/google/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, @@ -85759,6 +86001,14 @@ "output_cost_per_token": 2.5e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/openai/gpt-oss-20b": { + "input_cost_per_token": 4e-08, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/stepfun-ai/Step-3.7-Flash": { "cache_read_input_token_cost": 4e-08, "input_cost_per_token": 2e-07, @@ -85803,6 +86053,15 @@ "output_cost_per_token": 6e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/tencent/hy4": { + "cache_read_input_token_cost": 2.1e-08, + "input_cost_per_token": 4.17e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.2505e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/tencent/hy4-preview": { "cache_read_input_token_cost": 4.2e-08, "input_cost_per_token": 8.34e-07, @@ -86176,5 +86435,100 @@ "output_cost_per_token": 6e-07, "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true + }, + "vertex_ai/gemini-omni-1.2-flash": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 260640, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": false, + "supports_audio_output": false, + "supports_computer_use": false, + "supports_function_calling": false, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_tool_choice": false, + "supports_url_context": false, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": false + }, + "gemini/gemini-omni-1.2-flash": { + "input_cost_per_image_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "input_cost_per_video_token": 1.5e-06, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.2e-05, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "video", + "audio" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_reasoning": true, + "supports_system_messages": false, + "supports_video_input": true, + "supports_vision": true + }, + "databricks/databricks-kimi-k3-1-pro": { + "cache_creation_input_token_cost": 1.99997e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "metadata": {"notes":"Input/output cost per token is dbu cost * $0.070. Billing reads the per-token dollar fields; the '*_dbu_cost_per_token' fields are the published Databricks rates, kept for reference."}, + "mode": "chat", + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.000142857, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "token_count_framing": 100.0 } } diff --git a/pyproject.toml b/pyproject.toml index 2cc2d0f70a3..ee6cb6c1dcc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,6 +95,7 @@ proxy = [ "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", "expression>=5.6.0,<6.0", "backports-zstd>=1.0.0,<2.0; python_version < '3.14'", + "semver>=3.0.4,<4", ] # Thin client install for the `lite` CLI on developer laptops. The CLI's heavy # imports are all guarded, so it runs on the base SDK plus these packages, and @@ -269,6 +270,7 @@ proxy-dev = [ "a2a-sdk==1.1.0", "backports-zstd==1.8.0; python_version < '3.14'", "httptools>=0.7.1,<0.8", + "semver>=3.0.4,<4", ] ci = [ "psutil==7.2.2", diff --git a/schema.prisma b/schema.prisma index 78f16d9f1fd..818f2dc4858 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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 diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 42ee6178bf0..f6e9c2157d7 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -76,6 +76,7 @@ ignored_function_names = [ "_acreate_batch", "_acreate_file", "__init__", + "_router_init", # Never run: it carries Router.__init__'s signature and docstring "avector_store_create", # Tested via proxy vector_store_endpoints (files lack "router" in name) "_override_vector_store_methods_for_router", # No-op placeholder, called during Router init "_merge_tools_from_deployment", # Tested indirectly via _update_kwargs_with_deployment (test files lack "router" in name) diff --git a/tests/code_coverage_tests/router_litellm_globals_check.py b/tests/code_coverage_tests/router_litellm_globals_check.py index 172a3754000..ec38507a3bc 100644 --- a/tests/code_coverage_tests/router_litellm_globals_check.py +++ b/tests/code_coverage_tests/router_litellm_globals_check.py @@ -20,7 +20,6 @@ CHECKED = ( "litellm/router_api.py", "litellm/router_core.py", "litellm/router_deprecated.py", - "litellm/router_python_core.py", "litellm/router_rust/core.py", ) diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py index 1c41526101a..8eb7651f51c 100644 --- a/tests/code_coverage_tests/router_surface_check.py +++ b/tests/code_coverage_tests/router_surface_check.py @@ -20,7 +20,7 @@ Member uses are found the same way the boundary survey found them: attribute access on a receiver whose name looks like a router (llm_router, litellm_router_instance, router, ...) or on the Router class, plus getattr/setattr/hasattr with a constant name on such a receiver. Members -include the container file endpoints Router.__init__ registers by name. +include the endpoints the Rust core serves by name, container file endpoints included. Run: python tests/code_coverage_tests/router_surface_check.py [--write] --write regenerates the allowlist from the current tree (for shrinking it). @@ -41,7 +41,6 @@ ROUTER_FILES = { "litellm/router_api.py", "litellm/router_config.py", "litellm/router_core.py", - "litellm/router_python_core.py", "litellm/router_python_host.py", "litellm/router_rust/core.py", "litellm/router_deprecated.py", @@ -49,8 +48,15 @@ ROUTER_FILES = { } # The Rust core's host-side glue modules are part of the router. ROUTER_FILE_PREFIXES = ("litellm/router_rust/",) -ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_python_core.py", "litellm/router_deprecated.py"] -ROUTER_CLASS_NAMES = {"Router", "PythonCore", "DeprecatedRouterMembers"} +ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_rust/core.py", "litellm/router_deprecated.py"] +ROUTER_CLASS_NAMES = {"Router", "RustCore", "DeprecatedRouterMembers"} +# The Rust core serves its endpoint methods by name (``RustCore.__getattr__``); these name lists hold them. +ROUTER_ENDPOINT_LISTS = { + "litellm/router_rust/surface/legacy_members.py": {"_ROUTER_ENDPOINTS", "_ASSISTANTS_ENDPOINTS"}, + "litellm/router_rust/invoke/endpoints.py": {"_DIRECT_ENDPOINTS"}, +} +# ``RustEndpoints.member`` serves these methods by name, from its ``direct`` table. +ROUTER_ENDPOINT_TABLE = "litellm/router_rust/invoke/endpoints.py" SCAN_ROOTS = ["litellm", "enterprise"] SKIP_DIRS = { @@ -70,7 +76,7 @@ RECEIVER = re.compile( r"(^|[._])(llm_router|router|litellm_router|litellm_router_instance|router_instance" r"|_router|proxy_router|llm_router_instance|router_obj|mock_router|Router)$" ) -# Router.__init__ registers these with setattr(self, name, ...), so they never appear as self.. +# The Rust core serves these by name too, so they never appear as self.. CONTAINER_FILE_ENDPOINTS = "litellm/containers/endpoints.json" DYNAMIC_FUNCS = {"getattr", "setattr", "hasattr"} @@ -109,6 +115,24 @@ def router_members() -> Set[str]: and isinstance(node.ctx, ast.Store) ): members.add(node.attr) + for rel, names in ROUTER_ENDPOINT_LISTS.items(): + with open(os.path.join(ROOT, rel), encoding="utf-8") as f: + tree = ast.parse(f.read()) + for node in tree.body: + target = node.target if isinstance(node, ast.AnnAssign) else None + if isinstance(target, ast.Name) and target.id in names and node.value is not None: + members.update( + c.value for c in ast.walk(node.value) if isinstance(c, ast.Constant) and isinstance(c.value, str) + ) + with open(os.path.join(ROOT, ROUTER_ENDPOINT_TABLE), encoding="utf-8") as f: + tree = ast.parse(f.read()) + for node in ast.walk(tree): + if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "direct": + for table in ast.walk(node.value): + if isinstance(table, ast.Dict): + members.update( + k.value for k in table.keys if isinstance(k, ast.Constant) and isinstance(k.value, str) + ) with open(os.path.join(ROOT, CONTAINER_FILE_ENDPOINTS), encoding="utf-8") as f: for endpoint in json.load(f)["endpoints"]: members.update((endpoint["name"], endpoint["async_name"])) diff --git a/tests/integration/_support/daily_activity.py b/tests/integration/_support/daily_activity.py index 346fc156ea9..cedcb80843a 100644 --- a/tests/integration/_support/daily_activity.py +++ b/tests/integration/_support/daily_activity.py @@ -272,16 +272,32 @@ def assert_key_reported( date: str, metadata: Mapping[str, JsonValue], metrics: Mapping[str, float], + *, + entity_rows: Mapping[str, int] | None = None, ) -> None: assert response.status_code == 200, response.text body: Final = object_value(response.json()) records: Final = tuple(object_value(record) for record in records_of_key(body, api_key)) assert records, response.text assert all(record["metadata"] == metadata for record in records), response.text - assert all(counted(record["metrics"]) == pytest.approx(metrics) for record in records), response.text days: Final = body["results"] assert isinstance(days, list) and len(days) == 1, response.text day: Final = object_value(days[0]) + breakdown: Final = object_value(day["breakdown"]) + if entity_rows is None: + assert all(counted(record["metrics"]) == pytest.approx(metrics) for record in records), response.text + else: + totals: Final = records_of_key({name: value for name, value in breakdown.items() if name != "entities"}, api_key) + assert totals, response.text + assert all(counted(object_value(record)["metrics"]) == pytest.approx(metrics) for record in totals), response.text + entities: Final = object_value(breakdown["entities"]) + assert set(entities) == set(entity_rows), response.text + for entity, rows in entity_rows.items(): + bucket: Final = object_value(entities[entity]) + assert counted(bucket["metrics"]) == pytest.approx(seeded_metrics(rows)), response.text + keys: Final = object_value(bucket["api_key_breakdown"]) + assert set(keys) == {api_key}, response.text + assert counted(object_value(keys[api_key])["metrics"]) == pytest.approx(seeded_metrics(rows)), response.text assert day["date"] == date, response.text assert counted(day["metrics"]) == pytest.approx(metrics), response.text assert object_value(body["metadata"])["total_spend"] == pytest.approx(metrics["spend"]), response.text diff --git a/tests/integration/_support/device_login.py b/tests/integration/_support/device_login.py index 55519678f40..55eb5a753bb 100644 --- a/tests/integration/_support/device_login.py +++ b/tests/integration/_support/device_login.py @@ -25,6 +25,7 @@ from pydantic import JsonValue CHATGPT_AUTH_HOST: Final = "auth.openai.com" GITHUB_HOST: Final = "github.com" GITHUB_API_HOST: Final = "api.github.com" +COPILOT_API_HOST: Final = "api.githubcopilot.com" AUTH_HOSTS: Final = (CHATGPT_AUTH_HOST, GITHUB_HOST, GITHUB_API_HOST) CHATGPT_ACCOUNT: Final = "acct-device-login" _FAR_FUTURE: Final = 4102444800 @@ -79,7 +80,17 @@ COPILOT_MINTED_KEY: Final = "tid=device-login;minted" COPILOT_REJECTED_KEY: Final = "tid=device-login;rejected" CONTROL_KEY: Final = "device-login-control-key" ACCEPTED_BEARERS: Final = frozenset( - {CHATGPT_STORED, CHATGPT_REFRESHED, CHATGPT_FIRST_LOGIN, COPILOT_STORED_KEY, COPILOT_MINTED_KEY, CONTROL_KEY} + { + CHATGPT_STORED, + CHATGPT_REFRESHED, + CHATGPT_FIRST_LOGIN, + COPILOT_ACCESS, + COPILOT_FIRST_LOGIN_ACCESS, + COPILOT_UNSAVED_KEY_ACCESS, + COPILOT_STORED_KEY, + COPILOT_MINTED_KEY, + CONTROL_KEY, + } ) REJECTED_DETAIL: Final = "the scripted provider rejected this bearer token" _DEVICE_AUTH_ID: Final = "device-auth-synthetic" @@ -180,6 +191,10 @@ def _auth_hosts(api_url: str, grant: threading.Event) -> Callable[[Request], Rep return _json(200, {"access_token": COPILOT_FIRST_LOGIN_ACCESS, "token_type": "bearer"}) case ("api.github.com", "GET", "/copilot_internal/v2/token"): return _copilot_api_key(request, api_url) + case ("api.github.com", "GET", "/copilot_internal/user"): + if request.headers.get("authorization") == f"token {COPILOT_REJECTED_ACCESS}": + return _json(401, {"message": "Bad credentials"}) + return _json(200, {"endpoints": {"api": f"https://{COPILOT_API_HOST}"}}) case _: return _switched_off() @@ -204,9 +219,22 @@ class Switches: switch.clear() +@dataclass(frozen=True, slots=True) +class _ProviderAPI: + plain: Wire + secure: Wire + + @property + def url(self) -> str: + return self.plain.url + + def drain(self) -> tuple[Request, ...]: + return (*self.plain.drain(), *self.secure.drain()) + + @dataclass(frozen=True, slots=True) class Peers: - api: Wire + api: _ProviderAPI auth: Wire tunnel: str cert: Path @@ -228,6 +256,14 @@ class Peers: def auth_connections(self) -> tuple[str, ...]: return tuple(self.authorities.get_nowait() for _ in range(self.authorities.qsize())) + def assert_only_copilot_discovery(self, access_token: str) -> None: + assert set(self.auth_connections()) <= {f"{GITHUB_API_HOST}:443"} + requests: Final = self.auth.drain() + assert { + (request.method, request.target, request.headers.get("host"), request.headers.get("authorization")) + for request in requests + } <= {("GET", "/copilot_internal/user", GITHUB_API_HOST, f"token {access_token}")}, requests + def dropped(self) -> tuple[Hangup, ...]: return tuple(self.hangups.get_nowait() for _ in range(self.hangups.qsize())) @@ -279,9 +315,8 @@ class _TunnelServer(socketserver.ThreadingTCPServer): @contextmanager def _tunnel( - destination: Wire, switches: Switches, authorities: SimpleQueue[str], hangups: SimpleQueue[Hangup] + destination: Wire, provider: Wire, switches: Switches, authorities: SimpleQueue[str], hangups: SimpleQueue[Hangup] ) -> Generator[str]: - destination_port: Final = int(destination.url.rsplit(":", 1)[1]) class Tunnel(socketserver.StreamRequestHandler): rbufsize = 0 @@ -306,20 +341,24 @@ def _tunnel( if len(request_line) < 2: return authority: Final = request_line[1] - if authority.rsplit(":", 1)[0] not in AUTH_HOSTS: + host: Final = authority.rsplit(":", 1)[0] + if host not in (*AUTH_HOSTS, COPILOT_API_HOST): self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n") return - authorities.put(authority) - if switches.refuse_connect.is_set(): + is_auth: Final = host in AUTH_HOSTS + if is_auth: + authorities.put(authority) + if is_auth and switches.refuse_connect.is_set(): self.wfile.write(b"HTTP/1.1 502 Bad Gateway\r\ncontent-length: 0\r\n\r\n") return - if self.hold(authority, opened, switches.hold_connect) is None: + if is_auth and self.hold(authority, opened, switches.hold_connect) is None: return self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n") - client_hello: Final = self.hold(authority, opened, switches.stall_tls) + client_hello: Final = self.hold(authority, opened, switches.stall_tls) if is_auth else b"" if client_hello is None: return self.request.settimeout(10) + destination_port: Final = int((destination if is_auth else provider).url.rsplit(":", 1)[1]) with socket.create_connection(("127.0.0.1", destination_port), timeout=10) as upstream: upstream.sendall(client_hello) outbound: Final = threading.Thread(target=_pipe, args=(self.request, upstream)) @@ -340,16 +379,17 @@ def _tunnel( @contextmanager def device_login_peers(directory: Path) -> Generator[Peers]: - cert, key = write_self_signed_cert(directory, AUTH_HOSTS) + cert, key = write_self_signed_cert(directory, (*AUTH_HOSTS, COPILOT_API_HOST)) switches: Final = Switches(threading.Event(), threading.Event(), threading.Event(), threading.Event()) authorities: Final[SimpleQueue[str]] = SimpleQueue() hangups: Final[SimpleQueue[Hangup]] = SimpleQueue() with ( wire_server(_provider_api(rv.ResponsesVendor())) as api, + wire_server(_provider_api(rv.ResponsesVendor()), tls=server_context(cert, key)) as copilot_api, wire_server(_auth_hosts(api.url, switches.grant), tls=server_context(cert, key)) as auth, - _tunnel(auth, switches, authorities, hangups) as tunnel, + _tunnel(auth, copilot_api, switches, authorities, hangups) as tunnel, ): - yield Peers(api, auth, tunnel, cert, authorities, hangups, switches) + yield Peers(_ProviderAPI(api, copilot_api), auth, tunnel, cert, authorities, hangups, switches) def chatgpt_record( diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index e278a796146..d4cbccfc43e 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -1,16 +1,25 @@ from __future__ import annotations +import io +import selectors +import socket import ssl import threading import time from collections.abc import Callable, Generator, Iterator, Mapping -from contextlib import contextmanager +from contextlib import contextmanager, suppress from dataclasses import dataclass from functools import partial from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from queue import SimpleQueue from types import MappingProxyType -from typing import BinaryIO, Final +from typing import TYPE_CHECKING, BinaryIO, Final + +if TYPE_CHECKING: + from collections.abc import Buffer + +_POLL_SECONDS: Final = 0.05 +_WRITE_PIECE_BYTES: Final = 64 * 1024 @dataclass(frozen=True, slots=True) @@ -97,12 +106,108 @@ def _await_release( return gate.is_set() +class _RequestReader(io.BufferedReader): + cut = False + + def readline(self, size: int | None = -1, /) -> bytes: + line: Final = super().readline(size) + self.cut = self.cut or not line.endswith(b"\n") + return line + + +@dataclass(frozen=True, slots=True) +class _Blocked: + since: float + writing: bool + + +class _Connection(io.BufferedIOBase): + blocked: _Blocked | None = None + + def __init__(self, connection: socket.socket, closing: threading.Event) -> None: + super().__init__() + self.connection = connection + self.closing = closing + self.selector = selectors.DefaultSelector() + self.selector.register(connection, selectors.EVENT_READ) + + def readable(self) -> bool: + return True + + def writable(self) -> bool: + return True + + def readinto(self, buffer: Buffer, /) -> int: + if not self._wait_for_bytes(): + return 0 + with self._blocking(writing=False): + return self.connection.recv_into(buffer) + + def write(self, buffer: Buffer, /) -> int: + view: Final = memoryview(buffer).cast("B") + for start in range(0, view.nbytes, _WRITE_PIECE_BYTES): + with self._blocking(writing=True): + self.connection.sendall(view[start : start + _WRITE_PIECE_BYTES]) + return view.nbytes + + def close(self) -> None: + self.selector.close() + super().close() + + def stalled(self, grace: float, teardown_started: float) -> _Blocked | None: + blocked: Final = self.blocked + if blocked is None or time.monotonic() - max(blocked.since, teardown_started) < grace: + return None + return blocked + + def shut(self) -> None: + with suppress(OSError): + self.connection.shutdown(socket.SHUT_RDWR) + + def _wait_for_bytes(self) -> bool: + while not self.closing.is_set(): + if self._has_bytes(_POLL_SECONDS): + return True + return self._has_bytes(0) + + def _has_bytes(self, timeout: float) -> bool: + decrypted: Final = isinstance(self.connection, ssl.SSLSocket) and self.connection.pending() > 0 + return decrypted or bool(self.selector.select(timeout)) + + @contextmanager + def _blocking(self, writing: bool) -> Iterator[None]: + self.blocked = _Blocked(time.monotonic(), writing) + try: + yield + finally: + self.blocked = None + + +def _close_unless_stalled( + close: Callable[[], None], connections: SimpleQueue[_Connection], grace: float, errors: SimpleQueue[Exception] +) -> None: + teardown_started: Final = time.monotonic() + closer: Final = threading.Thread(target=close) + closer.start() + while closer.is_alive(): + for connection in tuple(connections.get_nowait() for _ in range(connections.qsize())): + stalled: Final = connection.stalled(grace, teardown_started) + if stalled is None: + connections.put(connection) + continue + if stalled.writing: + errors.put(TimeoutError(f"a reply write made no progress for {grace} s")) + connection.shut() + closer.join(timeout=_POLL_SECONDS) + + @contextmanager def wire_server( respond: Callable[[Request], Reply], tls: ssl.SSLContext | None = None, port: int = 0, keep_alive: bool = False, + teardown_grace: float = 10, ) -> Generator[Wire, None, None]: """Owned TCP peer; requests traverse the real HTTP client and serialization. With `keep_alive` the peer honours HTTP/1.1 persistent connections so `connections()` counts the client's TCP sessions.""" @@ -111,14 +216,20 @@ def wire_server( disconnected: Final[SimpleQueue[str]] = SimpleQueue() connected: Final[SimpleQueue[str]] = SimpleQueue() closed: Final[SimpleQueue[str]] = SimpleQueue() + streams: Final[SimpleQueue[_Connection]] = SimpleQueue() closing: Final = threading.Event() class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" - timeout = 5 + request: socket.socket + reader: _RequestReader def setup(self) -> None: - super().setup() + stream: Final = _Connection(self.request, closing) + self.reader = _RequestReader(stream) + self.rfile = self.reader + self.wfile = stream + streams.put(stream) connected.put(f"{self.client_address[0]}:{self.client_address[1]}") def finish(self) -> None: @@ -130,6 +241,9 @@ def wire_server( self.wfile.flush() def respond(self) -> None: + if self.reader.cut: + self.close_connection = True + return headers: Final = {name.lower(): value for name, value in self.headers.items()} try: body: Final = read_body(headers, self.rfile) @@ -221,7 +335,7 @@ def wire_server( self.socket = tls.wrap_socket(self.socket, server_side=True) with OwnedHTTPServer(("127.0.0.1", port), Handler) as server: - thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": _POLL_SECONDS}) thread.start() try: yield Wire( @@ -236,6 +350,6 @@ def wire_server( server.shutdown() thread.join(timeout=6) assert not thread.is_alive(), "Owned HTTP server survived cleanup" - server.server_close() + _close_unless_stalled(server.server_close, streams, teardown_grace, errors) failure: Final = None if errors.empty() else errors.get_nowait() assert failure is None, f"Owned HTTP peer failed: {failure!r}" diff --git a/tests/integration/authorization/test_responses_lifecycle_during_call_guardrail.py b/tests/integration/authorization/test_responses_lifecycle_during_call_guardrail.py new file mode 100644 index 00000000000..3227caaa7f5 --- /dev/null +++ b/tests/integration/authorization/test_responses_lifecycle_during_call_guardrail.py @@ -0,0 +1,114 @@ +"""GET, compact, DELETE and cancel on /v1/responses answer 200 when a during_call guardrail is on by default +(LIT-9770): the hook used to raise `'acompact_responses' is not a valid CallTypes` and the proxy returned 500""" + +import json +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, string_value +from integration._support.process import owned_proxy +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.shares_worker_state + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MODEL: Final = "openai/gpt-5.4-mini" + + +def _input() -> list[JsonValue]: + return [ + {"role": "user", "content": f"Remember the number {uuid.uuid4().hex}."}, + {"role": "assistant", "content": "Got it."}, + ] + + +def _response(response_id: str) -> Reply: + body: Final[dict[str, JsonValue]] = { + "id": response_id, + "object": "response", + "created_at": 1, + "status": "completed", + "model": _MODEL.removeprefix("openai/"), + "output": [ + { + "type": "message", + "id": f"msg_{uuid.uuid4().hex}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Got it.", "annotations": []}], + } + ], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}, + } + return Reply(body=json.dumps(body).encode()) + + +def _deleted(response_id: str) -> Reply: + return Reply(body=json.dumps({"id": response_id, "object": "response.deleted", "deleted": True}).encode()) + + +def _guarded_config(path: Path) -> Path: + base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + guardrail: Final[dict[str, JsonValue]] = { + "guardrail_name": f"content-filter-{uuid.uuid4().hex}", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "during_call", + "default_on": True, + "blocked_words": [{"keyword": "forbiddenword", "action": "BLOCK"}], + }, + } + path.write_text(yaml.safe_dump({**base, "guardrails": [guardrail]})) + return path + + +@pytest.fixture(scope="module") +def guarded(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("during-call-guardrail") + with owned_proxy(session_gateway, directory, {}, config=_guarded_config(directory / "config.yaml")) as candidate: + yield candidate + + +def _create(guarded: Gateway, provider: SharedProvider, upstream_id: str) -> str: + provider.expect(_response(upstream_id)) + created: Final = guarded.request("POST", "/v1/responses", {"model": _MODEL, "input": _input(), "store": True}) + assert created.status_code == 200, created.text + (upstream,) = provider.received() + assert upstream.target.endswith("/responses"), upstream.target + return string_value(_JSON_OBJECT.validate_json(created.content)["id"]) + + +def _lifecycle(guarded: Gateway, provider: SharedProvider, operation: str) -> tuple[httpx.Response, str]: + if operation == "compact": + provider.expect(_response(f"resp_{uuid.uuid4().hex}")) + return guarded.request( + "POST", "/v1/responses/compact", {"model": _MODEL, "input": _input()} + ), "/responses/compact" + upstream_id: Final = f"resp_{uuid.uuid4().hex}" + managed_id: Final = _create(guarded, provider, upstream_id) + provider.expect(_deleted(upstream_id) if operation == "delete" else _response(upstream_id)) + if operation == "get": + return guarded.request("GET", f"/v1/responses/{managed_id}"), f"/responses/{upstream_id}" + if operation == "delete": + return guarded.request("DELETE", f"/v1/responses/{managed_id}"), f"/responses/{upstream_id}" + return guarded.request("POST", f"/v1/responses/{managed_id}/cancel"), f"/responses/{upstream_id}/cancel" + + +@pytest.mark.parametrize("operation", ["get", "compact", "delete", "cancel"]) +def test_responses_lifecycle_answers_200_with_a_during_call_guardrail_on( + guarded: Gateway, provider: SharedProvider, operation: str +) -> None: + response, upstream_target = _lifecycle(guarded, provider, operation) + assert response.status_code == 200, response.text + (upstream,) = provider.received() + assert upstream.target.endswith(upstream_target), upstream.target diff --git a/tests/integration/database/test_skill_publication.py b/tests/integration/database/test_skill_publication.py new file mode 100644 index 00000000000..476c6bd4399 --- /dev/null +++ b/tests/integration/database/test_skill_publication.py @@ -0,0 +1,201 @@ +import io +import json +import os +import stat +import uuid +import zipfile +from concurrent.futures import ThreadPoolExecutor +from contextlib import AsyncExitStack, ExitStack +from typing import Final + +import psycopg +import pytest +from psycopg import sql + +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.plugin_upload_repository import PluginUploadRepository +from tests.integration._support.client import Gateway, list_value, object_value, string_value +from tests.integration._support.database import read_rows + + +def _archive(name: str) -> bytes: + buffer: Final = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr( + "SKILL.md", f"---\nname: {name}\ndescription: Published instructions\n---\nRetain these exact bytes.\n", + ) + return buffer.getvalue() + + +def _delete(gateway: Gateway, name: str) -> None: + response: Final = gateway.request("DELETE", f"/claude-code/plugins/{name}") + assert response.status_code == 200, response.text + + +@pytest.mark.asyncio +async def test_snapshot_failure_rolls_back_catalog_archive_and_history_then_source_replacement_preserves_bytes( + gateway: Gateway, +) -> None: + name: Final = f"integration-version-{uuid.uuid4().hex}" + path: Final = f"/claude-code/plugins/{name}" + history_path: Final = f"/skill-hub/skills/{name}/versions" + content: Final = _archive(name) + async with AsyncExitStack() as cleanup: + created: Final = gateway.request_multipart( + "/claude-code/plugins/upload", + {"metadata": json.dumps({"name": name})}, + {"file": ("skill.zip", content, "application/zip")}, + ) + assert created.status_code == 200, created.text + cleanup.callback(_delete, gateway, name) + client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache())) + await client.connect() + cleanup.push_async_callback(client.disconnect) + repository: Final = PluginUploadRepository(client) + metadata: Final = await repository.get_metadata(name) + assert metadata is not None + snapshot: Final = (await repository.versions(metadata, 100, 0))[0] + snapshot_bytes: Final = await repository.version_content(metadata, snapshot.id) + assert snapshot_bytes is not None + before: Final = gateway.get(path) + history: Final = gateway.get(history_path) + original: Final = object_value(list_value(history["versions"])[0]) + archive_path: Final = string_value(original["download_url"]) + original_response: Final = gateway.request("GET", archive_path) + assert original_response.status_code == 200, original_response.text + published_content: Final = original_response.content + with zipfile.ZipFile(io.BytesIO(published_content)) as published, zipfile.ZipFile(io.BytesIO(content)) as uploaded: + assert published.read("SKILL.md") == uploaded.read("SKILL.md") + current_path: Final = string_value(object_value(before["source"])["url"]) + original_current: Final = gateway.request("GET", current_path) + assert original_current.status_code == 200, original_current.text + witness: Final = f"publication_{uuid.uuid4().hex}" + trigger: Final = f"{witness}_trigger" + function: Final = f"{witness}_function" + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection, ExitStack() as failure: + connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(witness))) + failure.callback(connection.execute, sql.SQL("DROP SEQUENCE {}").format(sql.Identifier(witness))) + connection.execute(sql.SQL("GRANT USAGE ON SEQUENCE {} TO PUBLIC").format(sql.Identifier(witness))) + connection.execute( + sql.SQL( + "CREATE FUNCTION {}() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN " + "IF EXISTS (SELECT 1 FROM \"LiteLLM_ClaudeCodePluginTable\" " + "WHERE id = NEW.plugin_id AND name = {}) THEN PERFORM nextval({}); " + "RAISE EXCEPTION 'publication snapshot failure'; END IF; RETURN NEW; END $$" + ).format(sql.Identifier(function), sql.Literal(name), sql.Literal(witness)) + ) + failure.callback(connection.execute, sql.SQL("DROP FUNCTION {}()").format(sql.Identifier(function))) + connection.execute( + sql.SQL('CREATE TRIGGER {} BEFORE INSERT ON "LiteLLM_ClaudeCodePluginVersion" ' + "FOR EACH ROW EXECUTE FUNCTION {}()").format(sql.Identifier(trigger), sql.Identifier(function)) + ) + failure.callback( + connection.execute, + sql.SQL('DROP TRIGGER {} ON "LiteLLM_ClaudeCodePluginVersion"').format(sql.Identifier(trigger)), + ) + rejected: Final = gateway.request( + "PUT", path, {"source": {"source": "github", "repo": "example/changed"}, "description": "must roll back"}, + ) + assert rejected.status_code == 500, rejected.text + assert connection.execute( + sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness)) + ).fetchone() == (True,) + assert gateway.get(path) == before + assert gateway.get(history_path) == history + current_archive: Final = gateway.request("GET", current_path) + assert current_archive.status_code == 200 and current_archive.content == original_current.content + updated: Final = gateway.request( + "PUT", path, {"source": {"source": "github", "repo": "example/changed"}}, + ) + assert updated.status_code == 200, updated.text + after: Final = gateway.get(history_path) + versions: Final = tuple(object_value(value) for value in list_value(after["versions"])) + assert [version["version"] for version in versions] == ["1.0.1", "1.0.0"] + assert versions[0]["content_kind"] == "source_reference" + assert versions[1]["content_kind"] == "stored_archive" + retained: Final = gateway.request("GET", archive_path) + assert retained.status_code == 200 and retained.content == published_content + assert read_rows( + 'SELECT plugin_name FROM "LiteLLM_ClaudeCodePluginUpload" WHERE plugin_name=%s', (name,) + ) == [] + assert snapshot in await repository.versions(metadata, 100, 0) + assert await repository.version(metadata, snapshot.id) == snapshot + assert await repository.version_content(metadata, snapshot.id) == snapshot_bytes + hidden: Final = gateway.request("POST", f"{path}/disable") + assert hidden.status_code == 200, hidden.text + assert await repository.versions(metadata, 100, 0) == () + assert await repository.version(metadata, snapshot.id) is None + assert await repository.version_content(metadata, snapshot.id) is None + + +def test_neutral_publication_preserves_portable_files_and_atomic_metadata_history(gateway: Gateway) -> None: + name: Final = f"integration-neutral-{uuid.uuid4().hex}" + path: Final = f"/skill-hub/skills/{name}" + files: Final = { + "SKILL.md": f"---\nname: {name}\ndescription: Portable integration fixture.\n---\n\nRead references/notes.md.\n", + "references/notes.md": "Preserve the original supporting file.\n", + "scripts/check.sh": "#!/bin/sh\nprintf 'portable\\n'\n", + } + buffer: Final = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + for filename, content in files.items(): + entry: Final = zipfile.ZipInfo(filename) + entry.create_system = 3 + entry.external_attr = (stat.S_IFREG | (0o755 if filename.endswith(".sh") else 0o644)) << 16 + archive.writestr(entry, content) + with ExitStack() as cleanup: + created: Final = gateway.request_multipart( + "/skill-hub/skills/upload", {}, {"file": ("skill.zip", buffer.getvalue(), "application/zip")}, + ) + assert created.status_code == 200, created.text + cleanup.callback(_delete, gateway, name) + assert created.json()["version"] == "1.0.0" + before: Final = gateway.get(path) + assert string_value(before["updated_at"]) == created.json()["updated_at"] + history_before: Final = gateway.get(f"{path}/versions") + first: Final = object_value(list_value(history_before["versions"])[0]) + version_id: Final = string_value(first["id"]) + archive_path: Final = f"{path}/versions/{version_id}/archive" + downloaded: Final = gateway.request("GET", archive_path) + assert downloaded.status_code == 200, downloaded.text + with zipfile.ZipFile(io.BytesIO(downloaded.content)) as portable: + assert {member: portable.read(member).decode() for member in portable.namelist()} == files + assert stat.S_IMODE(portable.getinfo("scripts/check.sh").external_attr >> 16) == 0o755 + assert gateway.get(f"{path}/versions/{version_id}") == first + + def publish(index: int) -> int: + response: Final = gateway.request( + "PUT", path, + {"description": f"Registry edit {index}", "expected_updated_at": before["updated_at"]}, + ) + if response.status_code == 200: + assert string_value(response.json()["updated_at"]) != before["updated_at"] + return response.status_code + + with ThreadPoolExecutor(max_workers=2) as pool: + statuses: Final = tuple(pool.map(publish, range(2))) + assert sorted(statuses) == [200, 409] + assert publish(2) == 409 + current: Final = gateway.get(path) + history_after: Final = gateway.get(f"{path}/versions") + versions: Final = tuple(object_value(value) for value in list_value(history_after["versions"])) + assert [version["version"] for version in versions] == ["1.0.1", "1.0.0"] + latest: Final = versions[0] + manifest: Final = object_value(latest["manifest"]) + assert current["version"] == manifest["version"] == latest["version"] + assert current["description"] == manifest["description"] == latest["description"] + assert manifest["agent_skill"] == object_value(first["manifest"])["agent_skill"] + stored: Final = read_rows( + 'SELECT version, description, manifest_json FROM "LiteLLM_ClaudeCodePluginTable" WHERE name=%s', (name,) + ) + assert len(stored) == 1 + assert stored[0]["version"] == current["version"] + assert json.loads(string_value(stored[0]["manifest_json"])) == manifest + retained: Final = gateway.request("GET", archive_path) + latest_archive: Final = gateway.request("GET", f"{path}/versions/{latest['id']}/archive") + assert retained.content == latest_archive.content == downloaded.content + rejected: Final = gateway.request("PUT", path, {"version": "1.0.0", "description": "rejected"}) + assert rejected.status_code == 400, rejected.text + assert gateway.get(path) == current + assert gateway.get(f"{path}/versions") == history_after diff --git a/tests/integration/database/test_skill_submissions.py b/tests/integration/database/test_skill_submissions.py new file mode 100644 index 00000000000..cc00bf6e359 --- /dev/null +++ b/tests/integration/database/test_skill_submissions.py @@ -0,0 +1,454 @@ +import asyncio +import hashlib +import io +import os +import uuid +import zipfile +from collections.abc import AsyncIterator, Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal + +import psycopg +import pytest +import pytest_asyncio +from pydantic import JsonValue, TypeAdapter + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PluginUploadStore, + PublishedPlugin, +) +from litellm.models.skill_submission import ( + SkillSubmission, + SkillSubmissionStore, + SubmissionConflict, + SubmissionLimitExceeded, + SubmissionReviewed, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.plugin_upload_repository import PluginUploadRepository +from litellm.repositories.skill_submission_repository import SkillSubmissionRepository +from tests.integration._support.client import eventually +from tests.integration._support.database import read_rows, write_rows + + +@dataclass(frozen=True, slots=True) +class PendingUpdate: + publications: PluginUploadStore + submissions: SkillSubmissionStore + original: PluginPublication + proposal: PluginPublication + base: PublishedPlugin + submission: SkillSubmission + + +def _publication(name: str, description: str, base_url: str = "https://skills.example.test") -> PluginPublication: + buffer: Final = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("SKILL.md", f"---\nname: {name}\ndescription: {description}\n---\n{description}\n") + content: Final = buffer.getvalue() + digest: Final = hashlib.sha256(content).hexdigest() + return PluginPublication( + name=name, + version=None, + description=description, + manifest={ + "name": name, + "description": description, + "extension": {"retained": True}, + "source": { + "source": "archive", + "url": f"{base_url}/claude-code/plugins/{name}/archive/{digest}.zip", + "sha256": digest, + }, + }, + content=content, + created_by=f"{name}-owner", + ) + + +def _published_rows(name: str) -> Sequence[Mapping[str, JsonValue]]: + return read_rows( + """SELECT 'catalog' AS storage, to_jsonb(p) AS value + FROM "LiteLLM_ClaudeCodePluginTable" p WHERE name=%s + UNION ALL + SELECT 'archive', to_jsonb(a) FROM "LiteLLM_ClaudeCodePluginUpload" a WHERE plugin_name=%s + UNION ALL + SELECT 'history', to_jsonb(v) FROM "LiteLLM_ClaudeCodePluginVersion" v + JOIN "LiteLLM_ClaudeCodePluginTable" p ON p.id=v.plugin_id WHERE p.name=%s + ORDER BY storage, value""", + (name, name, name), + ) + + +@pytest_asyncio.fixture(loop_scope="function") +async def pending_update() -> AsyncIterator[PendingUpdate]: + name: Final = f"integration-submission-{uuid.uuid4().hex}" + client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache())) + await client.connect() + try: + publications: Final = PluginUploadRepository(client) + submissions: Final = SkillSubmissionRepository(client) + original: Final = _publication(name, "Published instructions") + base: Final = await publications.publish(original) + assert isinstance(base, PublishedPlugin) + proposal: Final = _publication(name, "Proposed instructions", "https://member.example.test") + submission: Final = await submissions.submit( + proposal, f"{name}-owner", PluginPublicationBase(base.id, base.updated_at) + ) + assert isinstance(submission, SkillSubmission) + yield PendingUpdate(publications, submissions, original, proposal, base, submission) + finally: + try: + write_rows('DELETE FROM "LiteLLM_SkillSubmission" WHERE plugin_name=%s', (name,)) + write_rows('DELETE FROM "LiteLLM_ClaudeCodePluginTable" WHERE name=%s', (name,)) + finally: + await client.disconnect() + + +async def test_approval_commits_version_archive_history_and_review(pending_update: PendingUpdate) -> None: + case: Final = pending_update + before: Final = await case.publications.get_metadata(case.original.name) + assert before is not None + assert await case.publications.get_content(before) == case.original.content + assert case.submission.proposal.version is None + reviewed: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", True, None, publication_base_url="https://skills.example.test" + ) + assert isinstance(reviewed, SubmissionReviewed) + assert reviewed.plugin is not None and reviewed.plugin.version == "1.0.1" + stored: Final = await case.submissions.get(case.submission.id, None) + assert stored is not None + assert (stored.status, stored.reviewed_by, stored.rejection_reason) == ( + "approved", + "integration-reviewer", + None, + ) + assert stored.reviewed_at is not None + assert stored.proposal == case.submission.proposal + assert await case.submissions.get_archive(stored.id, None) == case.proposal.content + current: Final = await case.publications.get_plugin(case.original.name) + assert current is not None + published_manifest: Final = { + **case.proposal.manifest, + "source": { + **TypeAdapter(dict[str, str]).validate_python(case.submission.proposal.manifest["source"]), + "url": f"https://skills.example.test/claude-code/plugins/{case.original.name}/archive/" + f"{hashlib.sha256(case.proposal.content or b'').hexdigest()}.zip", + }, + } + assert (current.id, current.version, current.description, current.manifest) == ( + case.base.id, + reviewed.plugin.version, + case.proposal.description, + {**published_manifest, "version": reviewed.plugin.version}, + ) + metadata: Final = await case.publications.get_metadata(case.original.name) + assert metadata is not None + assert await case.publications.get_content(metadata) == case.proposal.content + history: Final = await case.publications.versions(metadata, 10, 0) + assert tuple((version.version, version.description, version.archive_available) for version in history) == ( + ("1.0.1", case.proposal.description, True), + ("1.0.0", case.original.description, True), + ) + for version, publication, manifest in zip( + history, (case.proposal, case.original), (published_manifest, case.original.manifest), strict=True + ): + assert version.manifest == {**manifest, "version": version.version} + assert version.published_by == publication.created_by + assert await case.publications.version_content(metadata, version.id) == publication.content + + +async def test_simultaneous_approve_reject_has_one_winner(pending_update: PendingUpdate) -> None: + case: Final = pending_update + before: Final = _published_rows(case.original.name) + with psycopg.connect(os.environ["DATABASE_URL"]) as blocker: + blocker.execute('SELECT id FROM "LiteLLM_SkillSubmission" WHERE id=%s FOR UPDATE', (case.submission.id,)) + blocker_pid: Final = str(blocker.info.backend_pid) + reviews: Final = ( + asyncio.create_task( + case.submissions.review( + case.submission.id, + "integration-approver", + True, + None, + publication_base_url="https://skills.example.test", + ) + ), + asyncio.create_task( + case.submissions.review(case.submission.id, "integration-rejector", False, "Needs revision") + ), + ) + try: + await asyncio.to_thread( + eventually, + lambda: read_rows( + """WITH RECURSIVE blocked AS ( + SELECT %s::integer AS pid + UNION + SELECT a.pid FROM pg_stat_activity a + JOIN blocked b ON b.pid = ANY(pg_blocking_pids(a.pid)) + ) SELECT pid FROM blocked WHERE pid<>%s::integer""", + (blocker_pid, blocker_pid), + ), + lambda rows: len(rows) == 2, + seconds=3, + ) + finally: + blocker.rollback() + results: Final = await asyncio.gather(*reviews) + winners: Final = tuple(result for result in results if isinstance(result, SubmissionReviewed)) + assert len(winners) == 1 + assert sum(isinstance(result, SubmissionConflict) for result in results) == 1 + stored: Final = await case.submissions.get(case.submission.id, None) + assert stored is not None and stored.reviewed_at is not None + winner: Final = winners[0] + assert (stored.status, stored.reviewed_by, stored.rejection_reason) == ( + winner.submission.status, + winner.submission.reviewed_by, + winner.submission.rejection_reason, + ) + if isinstance(results[0], SubmissionConflict): + assert (stored.status, stored.reviewed_by, stored.rejection_reason) == ( + "rejected", + "integration-rejector", + "Needs revision", + ) + assert _published_rows(case.original.name) == before + return + assert (stored.status, stored.reviewed_by, stored.rejection_reason) == ( + "approved", + "integration-approver", + None, + ) + metadata: Final = await case.publications.get_metadata(case.original.name) + assert metadata is not None + assert await case.publications.get_content(metadata) == case.proposal.content + assert tuple(version.version for version in await case.publications.versions(metadata, 10, 0)) == ( + "1.0.1", + "1.0.0", + ) + + +async def test_publication_failure_rolls_back_and_leaves_submission_pending(pending_update: PendingUpdate) -> None: + case: Final = pending_update + collision_id: Final = str(uuid.uuid4()) + write_rows( + """INSERT INTO "LiteLLM_ClaudeCodePluginVersion" (id, plugin_id, version) + VALUES (%s, %s, '1.0.1')""", + (collision_id, case.base.id), + ) + before: Final = _published_rows(case.original.name) + failed: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", True, None, publication_base_url="https://skills.example.test" + ) + assert isinstance(failed, SubmissionConflict) + assert _published_rows(case.original.name) == before + assert await case.submissions.get(case.submission.id, None) == case.submission + assert await case.submissions.get_archive(case.submission.id, None) == case.proposal.content + write_rows('DELETE FROM "LiteLLM_ClaudeCodePluginVersion" WHERE id=%s', (collision_id,)) + retried: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", True, None, publication_base_url="https://skills.example.test" + ) + assert isinstance(retried, SubmissionReviewed) + assert retried.plugin is not None and retried.plugin.version == "1.0.1" + stored: Final = await case.submissions.get(case.submission.id, None) + assert stored is not None and stored.status == "approved" + + +@pytest.mark.parametrize("change", ("update", "recreate")) +async def test_stale_target_cannot_publish( + pending_update: PendingUpdate, change: Literal["update", "recreate"] +) -> None: + case: Final = pending_update + intervening: Final = _publication(case.original.name, "Intervening publication") + if change == "recreate": + write_rows('DELETE FROM "LiteLLM_ClaudeCodePluginTable" WHERE name=%s', (case.original.name,)) + replacement: Final = await case.publications.publish( + intervening, + PluginPublicationBase(case.base.id, case.base.updated_at) if change == "update" else None, + ) + assert isinstance(replacement, PublishedPlugin) + before: Final = _published_rows(case.original.name) + reviewed: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", True, None, publication_base_url="https://skills.example.test" + ) + assert isinstance(reviewed, PluginPublicationError) and reviewed.kind == "conflict" + assert _published_rows(case.original.name) == before + assert await case.submissions.get(case.submission.id, None) == case.submission + + +async def test_concurrent_submissions_share_user_limit_and_review_releases_capacity( + pending_update: PendingUpdate, +) -> None: + case: Final = pending_update + owner: Final = case.submission.submitted_by + base: Final = PluginPublicationBase(case.base.id, case.base.updated_at) + for _ in range(18): + assert isinstance(await case.submissions.submit(case.proposal, owner, base), SkillSubmission) + peer: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache())) + await peer.connect() + try: + with psycopg.connect(os.environ["DATABASE_URL"]) as blocker: + blocker.execute( + "SELECT pg_advisory_xact_lock(hashtext('litellm.skill_submissions'), hashtext(%s))", (owner,) + ) + blocker_pid: Final = str(blocker.info.backend_pid) + requests: Final = ( + asyncio.create_task(case.submissions.submit(case.proposal, owner, base)), + asyncio.create_task(SkillSubmissionRepository(peer).submit(case.proposal, owner, base)), + ) + try: + await asyncio.to_thread( + eventually, + lambda: read_rows( + "SELECT pid FROM pg_stat_activity WHERE %s::integer = ANY(pg_blocking_pids(pid))", + (blocker_pid,), + ), + lambda rows: len(rows) == 2, + seconds=3, + ) + finally: + blocker.rollback() + results: Final = await asyncio.gather(*requests) + finally: + await peer.disconnect() + assert sum(isinstance(result, SkillSubmission) for result in results) == 1 + assert sum(isinstance(result, SubmissionLimitExceeded) for result in results) == 1 + full: Final = await case.submissions.list(owner, "pending_review", 20, 0) + assert (len(full.submissions), full.has_more) == (20, False) + other: Final = await case.submissions.submit(case.proposal, f"{owner}-other", base) + assert isinstance(other, SkillSubmission) + rejected: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", False, "Review complete" + ) + assert isinstance(rejected, SubmissionReviewed) + assert await case.submissions.get_archive(case.submission.id, owner) == case.proposal.content + replacement: Final = await case.submissions.submit(case.proposal, owner, base) + assert isinstance(replacement, SkillSubmission) + after: Final = await case.submissions.list(owner, "pending_review", 20, 0) + assert len(after.submissions) == 20 + assert all(row.id != case.submission.id and row.submitted_by == owner for row in after.submissions) + + +async def test_submission_pages_filter_owner_and_status_before_limit(pending_update: PendingUpdate) -> None: + case: Final = pending_update + owner: Final = case.submission.submitted_by + other: Final = f"{owner}-other" + generated: Final = await asyncio.gather( + *( + case.submissions.submit(case.proposal, identity, PluginPublicationBase(case.base.id, case.base.updated_at)) + for identity in (owner, other, owner, other) + ) + ) + submitted: Final = tuple(result for result in generated if isinstance(result, SkillSubmission)) + assert len(submitted) == 4 + rejected: Final = await case.submissions.review(submitted[0].id, "integration-reviewer", False, None) + assert isinstance(rejected, SubmissionReviewed) + write_rows( + """UPDATE "LiteLLM_SkillSubmission" SET created_at=TIMESTAMP '9999-01-01' + WHERE plugin_name=%s""", + (case.original.name,), + ) + own_pending: Final = tuple(sorted((case.submission.id, submitted[2].id), reverse=True)) + first: Final = await case.submissions.list(owner, "pending_review", 1, 0) + second: Final = await case.submissions.list(owner, "pending_review", 1, 1) + empty: Final = await case.submissions.list(owner, "pending_review", 1, 2) + assert (tuple(row.id for row in first.submissions), first.has_more) == (own_pending[:1], True) + assert (tuple(row.id for row in second.submissions), second.has_more) == (own_pending[1:], False) + assert (empty.submissions, empty.has_more) == ((), False) + reviewed: Final = await case.submissions.list(owner, "rejected", 20, 0) + assert (tuple(row.id for row in reviewed.submissions), reviewed.has_more) == ((submitted[0].id,), False) + own_all: Final = await case.submissions.list(owner, None, 20, 0) + assert tuple(row.id for row in own_all.submissions) == tuple(sorted((*own_pending, submitted[0].id), reverse=True)) + other_page: Final = await case.submissions.list(other, None, 20, 0) + assert tuple(row.id for row in other_page.submissions) == tuple( + sorted((submitted[1].id, submitted[3].id), reverse=True) + ) + global_pending: Final = await case.submissions.list(None, "pending_review", 3, 0) + all_pending: Final = tuple(sorted((*own_pending, submitted[1].id, submitted[3].id), reverse=True)) + assert (tuple(row.id for row in global_pending.submissions), global_pending.has_more) == (all_pending[:3], True) + + +@pytest.mark.parametrize("completion_seconds", (13, 21), ids=("before-deadline", "at-deadline")) +async def test_scan_claim_retry_and_late_result_preserve_proposal( + pending_update: PendingUpdate, completion_seconds: int, +) -> None: + from datetime import datetime, timedelta, timezone + + from litellm.models.skill_submission import SkillScanInput + from litellm.types.proxy.skill_scanning import SkillScan, SkillScanReport + + case: Final = pending_update + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + first: Final = SkillScan( + id="first", status="running", requested_by="reviewer", started_at=now, + expires_at=now + timedelta(seconds=10), + ) + started: Final = await case.submissions.start_scan(case.submission.id, first) + assert isinstance(started, SkillScanInput) + assert started.content == case.proposal.content + assert started.scan.archive_sha256 == hashlib.sha256(case.proposal.content or b"").hexdigest() + duplicate: Final = await case.submissions.start_scan(case.submission.id, first.model_copy(update={"id": "duplicate"})) + assert isinstance(duplicate, SubmissionConflict) + second: Final = first.model_copy(update={ + "id": "second", "started_at": now + timedelta(seconds=11), "expires_at": now + timedelta(seconds=21), + }) + retry: Final = await case.submissions.start_scan(case.submission.id, second) + assert isinstance(retry, SkillScanInput) + report: Final = SkillScanReport( + scanner_version="test-version", complete=False, scan_mode="static-only", risk_score=90, findings=(), + ) + await case.submissions.finish_scan(case.submission.id, started.scan.model_copy(update={ + "status": "completed", "report": report, "completed_at": now + timedelta(seconds=12), + })) + assert await case.submissions.get_scan(case.submission.id) == retry.scan + completed: Final = retry.scan.model_copy(update={ + "status": "completed", "report": report, "completed_at": now + timedelta(seconds=completion_seconds), + }) + expected: Final = completed if completion_seconds < 21 else retry.scan.at(now + timedelta(seconds=21)) + await case.submissions.finish_scan(case.submission.id, completed) + assert await case.submissions.get_scan(case.submission.id) == expected + assert await case.submissions.get(case.submission.id, None) == case.submission + approved: Final = await case.submissions.review(case.submission.id, "reviewer", True, None, publication_base_url="https://gateway.test") + assert isinstance(approved, SubmissionReviewed) + assert await case.submissions.get_scan(case.submission.id) == expected + denied: Final = await case.submissions.start_scan(case.submission.id, second) + assert isinstance(denied, SubmissionConflict) + + +@pytest.mark.parametrize("available, remote", ((True, False), (False, False), (False, True))) +async def test_metadata_scan_requires_the_matching_managed_archive( + pending_update: PendingUpdate, available: bool, remote: bool +) -> None: + from datetime import datetime, timedelta, timezone + + from litellm.models.skill_submission import SkillScanInput + from litellm.types.proxy.skill_scanning import SkillScan + + case: Final = pending_update + submission: Final = await case.submissions.submit( + replace(case.original, content=None, manifest={ + **case.original.manifest, + **({"source": {"source": "archive", "url": "https://example.org/skill.zip"}} if remote else {}), + }), "member", PluginPublicationBase(case.base.id, case.base.updated_at) + ) + if not available: + write_rows('DELETE FROM "LiteLLM_ClaudeCodePluginUpload" WHERE plugin_name=%s', (case.original.name,)) + now: Final = datetime.now(timezone.utc) + result: Final = await case.submissions.start_scan(submission.id, SkillScan( + id="metadata-scan", status="running", requested_by="admin", started_at=now, + expires_at=now + timedelta(seconds=10), + )) + if available or remote: + assert isinstance(result, SkillScanInput) + assert result.content == (None if remote else case.original.content) + if remote: + assert result.source == {"source": "archive", "url": "https://example.org/skill.zip"} + assert result.scan.archive_sha256 is None + else: + assert isinstance(result, SubmissionConflict) + assert await case.submissions.get_scan(submission.id) is None diff --git a/tests/integration/database/test_spend_lock_order.py b/tests/integration/database/test_spend_lock_order.py new file mode 100644 index 00000000000..a7e00b7fa87 --- /dev/null +++ b/tests/integration/database/test_spend_lock_order.py @@ -0,0 +1,197 @@ +import asyncio +import subprocess +import sys +from collections.abc import AsyncIterator, Callable, Iterator, Sequence +from pathlib import Path +from types import SimpleNamespace +from typing import Final + +import psycopg +import pytest +import pytest_asyncio +from prisma import Prisma + +from litellm.proxy._types import DBSpendUpdateTransactions +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model +from litellm.repositories import unit_of_work +from tests.integration._support.database import scratch_database + +_ROOT: Final = Path(__file__).resolve().parents[3] +_KEY: Final = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683" +_CHARGES: Final = DBSpendUpdateTransactions( + user_list_transactions={"user": 1.0}, + end_user_list_transactions={}, + key_list_transactions={_KEY: 1.0}, + team_list_transactions={"team": 1.0}, + team_member_list_transactions={"team_id::team::user_id::user": 1.0}, + org_list_transactions={"org": 1.0}, + org_member_list_transactions={"org_id::org::user_id::user": 1.0}, + project_list_transactions={"project": 1.0}, + tag_list_transactions={"tag": 1.0}, + agent_list_transactions={}, + model_access_group_list_transactions={"group": 1.0}, +) +_SEED: Final = ( + "INSERT INTO \"LiteLLM_BudgetTable\" (budget_id, created_by, updated_by) VALUES ('budget', 't', 't')", + 'INSERT INTO "LiteLLM_OrganizationTable" (organization_id, organization_alias, budget_id, created_by, updated_by) ' + "VALUES ('org', 'org', 'budget', 't', 't')", + "INSERT INTO \"LiteLLM_UserTable\" (user_id) VALUES ('user')", + 'INSERT INTO "LiteLLM_TeamTable" (team_id, members_with_roles) ' + """VALUES ('team', '[{"user_id": "user", "role": "user"}]')""", + "INSERT INTO \"LiteLLM_TeamMembership\" (team_id, user_id, budget_id) VALUES ('team', 'user', 'budget')", + "INSERT INTO \"LiteLLM_OrganizationMembership\" (organization_id, user_id) VALUES ('org', 'user')", + f"INSERT INTO \"LiteLLM_VerificationToken\" (token, budget_id) VALUES ('{_KEY}', 'budget')", + "INSERT INTO \"LiteLLM_ProjectTable\" (project_id, budget_id, created_by, updated_by) " + "VALUES ('project', 'budget', 't', 't')", + "INSERT INTO \"LiteLLM_TagTable\" (tag_name, budget_id) VALUES ('tag', 'budget')", + "INSERT INTO \"LiteLLM_ModelAccessGroupBudgetTable\" (access_group_name, budget_id) VALUES ('group', 'budget')", +) + + +@pytest.fixture(scope="module") +def schema_sql() -> str: + return subprocess.run( + [sys.executable, "-I", "-m", "prisma", "migrate", "diff", "--from-empty", "--to-schema-datamodel", + str(_ROOT / "schema.prisma"), "--script"], + check=True, + capture_output=True, + text=True, + cwd=_ROOT, + ).stdout + + +@pytest.fixture +def database(schema_sql: str) -> Iterator[str]: + with scratch_database("spend_lock_order") as url, psycopg.connect(url, autocommit=True) as connection: + connection.execute(schema_sql) + for statement in _SEED: + connection.execute(statement) + yield url + + +@pytest_asyncio.fixture(loop_scope="function") +async def ledger(database: str) -> AsyncIterator[Prisma]: + client: Final = Prisma(datasource={"url": database}) + await client.connect() + try: + yield client + finally: + await client.disconnect() + + +async def _apply_charges(client: Prisma) -> None: + async with client.tx() as transaction: + await DBSpendUpdateWriter().apply_accounting_entities(client, None, transaction, _CHARGES) + + +async def _ledger_waits_for_a_lock(url: str) -> None: + with psycopg.connect(url, autocommit=True) as observer: + for _ in range(200): + if observer.execute("SELECT 1 FROM pg_locks WHERE NOT granted").fetchone(): + return + await asyncio.sleep(0.05) + raise AssertionError("the ledger never waited on the other writer") + + +def _spend(url: str) -> Sequence[object]: + with psycopg.connect(url) as connection: + return connection.execute( + 'SELECT (SELECT spend FROM "LiteLLM_UserTable"), (SELECT spend FROM "LiteLLM_TeamTable"), ' + '(SELECT spend FROM "LiteLLM_TeamMembership"), (SELECT spend FROM "LiteLLM_VerificationToken"), ' + '(SELECT spend FROM "LiteLLM_OrganizationTable"), (SELECT spend FROM "LiteLLM_ProjectTable")' + ).fetchone() or () + + +_Writer = Callable[[psycopg.Connection[tuple[object, ...]]], None] + + +def _roster_first(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute("SELECT pg_advisory_xact_lock(hashtext('team'))") + + +def _roster_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute("SELECT 1 FROM \"LiteLLM_UserTable\" WHERE user_id = 'user' FOR UPDATE") + connection.execute("UPDATE \"LiteLLM_TeamTable\" SET members_with_roles = members_with_roles WHERE team_id = 'team'") + + +def _access_group_first(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute("UPDATE \"LiteLLM_TeamTable\" SET access_group_ids = ARRAY['ag'] WHERE team_id = 'team'") + + +def _access_group_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute(f"UPDATE \"LiteLLM_VerificationToken\" SET access_group_ids = ARRAY['ag'] WHERE token = '{_KEY}'") + + +def _cascade_first(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute(unit_of_work._LOCK_BUDGET_ORGS_SQL.replace("$1", "%s"), (["budget"],)) + + +def _cascade_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None: + for statement in ( + unit_of_work._LOCK_BUDGET_MEMBERSHIPS_SQL, + unit_of_work._LOCK_BUDGET_KEYS_SQL, + unit_of_work._LOCK_BUDGET_PROJECTS_SQL, + unit_of_work._LOCK_BUDGET_TAGS_SQL, + unit_of_work._LOCK_BUDGET_ACCESS_GROUPS_SQL, + ): + connection.execute(statement.replace("$1", "%s"), (["budget"],)) + connection.execute("UPDATE \"LiteLLM_TagTable\" SET spend = 0 WHERE budget_id = 'budget'") + + +@pytest.mark.parametrize( + ("first", "rest"), + [(_roster_first, _roster_rest), (_access_group_first, _access_group_rest), (_cascade_first, _cascade_rest)], + ids=["team_roster", "access_group", "budget_cascade"], +) +@pytest.mark.asyncio(loop_scope="function") +async def test_charges_wait_for_a_concurrent_writer_instead_of_deadlocking( + database: str, ledger: Prisma, first: _Writer, rest: _Writer +) -> None: + """The other writer takes its first lock, the ledger starts and must block holding nothing the other writer + needs next, so the other writer finishes and the charge lands once after it.""" + with psycopg.connect(database) as writer: + writer.execute("SET lock_timeout = '10s'") + first(writer) + charges: Final = asyncio.create_task(_apply_charges(ledger)) + await _ledger_waits_for_a_lock(database) + await asyncio.to_thread(rest, writer) + writer.commit() + await asyncio.wait_for(charges, timeout=20) + + assert _spend(database) == (1.0, 1.0, 1.0, 1.0, 1.0, 1.0) + + +async def _rename(client: Prisma) -> None: + await sync_model_allowlists_for_renamed_model( + SimpleNamespace(db=client), + model_id="deployment", + old_name="old-model", + new_name="new-model", + llm_router=None, + user_api_key_cache=UserApiKeyCache(), + ) + + +@pytest.mark.asyncio(loop_scope="function") +async def test_a_model_rename_waits_for_a_roster_change_instead_of_deadlocking(database: str, ledger: Prisma) -> None: + """The roster change holds the team lock and the user row, so the rename must wait before it locks the team + row the roster change updates next.""" + with psycopg.connect(database, autocommit=True) as seed: + seed.execute("UPDATE \"LiteLLM_TeamTable\" SET models = ARRAY['old-model']") + seed.execute("UPDATE \"LiteLLM_UserTable\" SET models = ARRAY['old-model']") + with psycopg.connect(database) as writer: + writer.execute("SET lock_timeout = '10s'") + _roster_first(writer) + writer.execute("SELECT 1 FROM \"LiteLLM_UserTable\" WHERE user_id = 'user' FOR UPDATE") + rename: Final = asyncio.create_task(_rename(ledger)) + await _ledger_waits_for_a_lock(database) + await asyncio.to_thread(_roster_rest, writer) + writer.commit() + await asyncio.wait_for(rename, timeout=20) + + with psycopg.connect(database) as connection: + assert connection.execute( + 'SELECT (SELECT models FROM "LiteLLM_TeamTable"), (SELECT models FROM "LiteLLM_UserTable")' + ).fetchone() == (["new-model"], ["new-model"]) diff --git a/tests/integration/database/test_spendlogs_delivery.py b/tests/integration/database/test_spendlogs_delivery.py index 19f01a6b902..d1227e85565 100644 --- a/tests/integration/database/test_spendlogs_delivery.py +++ b/tests/integration/database/test_spendlogs_delivery.py @@ -131,9 +131,9 @@ async def test_duplicate_rows_in_one_flush_keep_first_accounting_payload(deliver "accounting_log_disabled": mode == "disabled", } assert await repository.insert_with_deliveries((first, later), write_spend_logs=mode != "url") == 1 - assert delivery_database.execute( - 'SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"' - ).fetchone() == (first["accounting_payload"],) + assert delivery_database.execute('SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"').fetchone() == ( + first["accounting_payload"], + ) assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == ( [(first["spend"],)] if mode == "local" else [] ) @@ -415,11 +415,28 @@ def test_delivery_migration_preserves_unused_old_receipts(state): assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) +@pytest.mark.parametrize( + ("payload", "huge_state"), + [ + ( + "jsonb_set(accounting_payload, '{metadata}', ('{\"tokens_saved\": ' || repeat('9', 401) || '}')::jsonb)", + "applied", + ), + ( + "jsonb_set(accounting_payload, '{metadata}', ('{\"tokens_saved\": ' || repeat('9', 5000) || '}')::jsonb)", + "parked", + ), + ("'[]'::jsonb", "parked"), + ], + ids=["beyond_float", "beyond_python_int_parse", "not_an_object"], +) @pytest.mark.asyncio -async def test_malformed_payload_is_parked_without_losing_other_deliveries(delivery_database, delivery_client): - _insert(delivery_database, (_row("bad"), _row("good"))) +async def test_a_payload_the_engine_cannot_decode_never_blocks_other_deliveries( + delivery_database, delivery_client, payload, huge_state +): + _insert(delivery_database, (_row("huge"), _row("good"))) delivery_database.execute( - 'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload=$$ [] $$::jsonb WHERE request_id=$$bad$$' + f'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload = {payload} WHERE request_id = $$huge$$' ) async def apply(transaction, deliveries): @@ -428,12 +445,16 @@ async def test_malformed_payload_is_parked_without_losing_other_deliveries(deliv ) repository = AccountingDeliveryRepository(delivery_client) - assert (await repository.apply_pending(apply, max_attempts=2)).failed == 2 - assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,) - delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP') - assert (await repository.apply_pending(apply, max_attempts=2)).parked == ("bad",) - assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1 - assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + results = [await repository.apply_pending(apply, max_attempts=1) for _ in range(3)] + _insert(delivery_database, (_row("later"),)) + results.append(await repository.apply_pending(apply, max_attempts=1)) + + applied: Final = 3 if huge_state == "applied" else 2 + assert sum(result.applied for result in results) == applied + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5 * applied,) + assert delivery_database.execute( + 'SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id' + ).fetchall() == [("good", "applied"), ("huge", huge_state), ("later", "applied")] @pytest.mark.asyncio @@ -463,19 +484,29 @@ async def test_ledger_batch_retrieves_charge_once_in_every_logging_mode( from litellm.proxy.utils import jsonify_object completed = { - **_row("batch_completed"), "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + **_row("batch_completed"), + "call_type": "aretrieve_batch", + "status": "success", + "spend": 0.5, "metadata": '{"status":"success","private":"request detail"}', "accounting_payload": { - "request_id": "batch_completed", "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + "request_id": "batch_completed", + "call_type": "aretrieve_batch", + "status": "success", + "spend": 0.5, "metadata": {"accounting_savings": {"compression": 0.2}}, }, } monkeypatch.setattr(proxy_server, "general_settings", {"spend_logs_metadata_fields": {"include": ["status"]}}) client = SimpleNamespace(db=delivery_client, writer_db=delivery_client, jsonify_object=jsonify_object) writer = DBSpendUpdateWriter() - first = await writer._claim_settle_once_cost_row(payload=completed, prisma_client=client, disable_spend_logs=disabled) + first = await writer._claim_settle_once_cost_row( + payload=completed, prisma_client=client, disable_spend_logs=disabled + ) again = await writer._claim_settle_once_cost_row( - payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, prisma_client=client, disable_spend_logs=disabled + payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, + prisma_client=client, + disable_spend_logs=disabled, ) assert (first, again) == (True, False) assert delivery_database.execute('SELECT spend, metadata FROM "LiteLLM_SpendLogs"').fetchall() == ( diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index e1894aa66bf..60f3c074cca 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -63,10 +63,10 @@ _DAILY_TABLES: Final = ( ) -@pytest.mark.parametrize("batch_size,row_count,expected,distinct_models", [(None, 251, 251, False), ("200", 251, 200, False), (None, 10001, 10000, False), (None, 3001, 3001, True)]) +@pytest.mark.parametrize("batch_size,row_count,transactions,distinct_models", [(None, 251, 1, False), (None, 1001, 3, False), ("200", 251, 2, False), ("3001", 3001, 1, True)]) @pytest.mark.asyncio -async def test_worker_applies_the_configured_backlog_in_one_transaction( - parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, expected: int, distinct_models: bool +async def test_worker_drains_a_backlog_in_one_run_one_configured_batch_per_transaction( + parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, transactions: int, distinct_models: bool ): from litellm.accounting.worker import flush_accounting_deliveries @@ -92,13 +92,37 @@ async def test_worker_applies_the_configured_backlog_in_one_transaction( )) == row_count - 1 started = time.monotonic() await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) - print(f"Ledger batch applied={expected} queued={row_count} seconds={time.monotonic() - started:.3f}") + print(f"Ledger backlog queued={row_count} seconds={time.monotonic() - started:.3f}") with psycopg.connect(parity_database) as connection: - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (expected,) - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (row_count - expected,) - assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (1,) - assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (expected * 0.25,) - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (expected if distinct_models else 1,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (row_count,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (0,) + assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (transactions,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (row_count * 0.25,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (row_count if distinct_models else 1,) + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_worker_removes_more_expired_receipts_than_one_cleanup_batch_in_one_run( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "100") + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, state, applied_at) ' + "SELECT 'expired-' || n, '{}'::jsonb, false, 'applied', CURRENT_TIMESTAMP - INTERVAL '2 days' " + "FROM generate_series(1, 501) AS n" + ) + await client.db.connect() + try: + await flush_accounting_deliveries(DBSpendUpdateWriter(), client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) finally: await client.db.disconnect() @@ -404,6 +428,7 @@ async def _assert_legacy_database_parity( assert len(rows) == 1 if path == "ledger" and retention == "lean": assert rows[0]["accounting_log_disabled"] is True + assert not {"messages", "response", "proxy_server_request", "metadata"}.intersection(rows[0]) frozen = rows[0]["accounting_payload"] if isinstance(frozen, str): frozen = json.loads(frozen) diff --git a/tests/integration/management/test_claude_code_plugin_installation_preference.py b/tests/integration/management/test_claude_code_plugin_installation_preference.py index aaf6f02ff4e..6d41c690709 100644 --- a/tests/integration/management/test_claude_code_plugin_installation_preference.py +++ b/tests/integration/management/test_claude_code_plugin_installation_preference.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from functools import partial from pathlib import Path +from threading import Lock from typing import Final import httpx @@ -248,7 +249,7 @@ def test_duplicate_json_key_keeps_the_last_value(gateway: Gateway) -> None: assert stored_manifest(name)["installation_preference"] == "available" -def test_repeated_identical_update_is_idempotent(gateway: Gateway) -> None: +def test_repeated_identical_updates_publish_patch_versions_without_changing_preference(gateway: Gateway) -> None: name: Final = plugin_name() with gateway.scenario() as scenario: registered(scenario, name, installation_preference="required") @@ -257,7 +258,9 @@ def test_repeated_identical_update_is_idempotent(gateway: Gateway) -> None: replaced(gateway, name, installation_preference="required") second: Final = stored_rows(name) assert len(first) == 1, first - assert second == first + first_manifest: Final = JSON_OBJECT.validate_json(string_value(first[0]["manifest_json"])) + assert first_manifest["version"] == "1.0.1", first + assert second == [{**first[0], "manifest_json": json.dumps({**first_manifest, "version": "1.0.2"})}] assert_preference_served(gateway, name, "required") @@ -442,14 +445,16 @@ def test_burst_reads_across_processes_while_the_preference_flips(gateway: Gatewa name: Final = plugin_name() targets: Final = (gateway, peer) flips: Final = tuple(PREFERENCES[(index + 1) % len(PREFERENCES)] for index in range(13)) + writer: Final = Lock() with gateway.scenario() as scenario: registered(scenario, name, installation_preference="available") def job(index: int) -> tuple[str, JsonValue]: if index % 4 == 3: - return "flip", replace( - gateway, name, replacement(name, installation_preference=flips[index // 4]) - ).status_code + with writer: + return "flip", replace( + gateway, name, replacement(name, installation_preference=flips[index // 4]) + ).status_code return "read", marketplace_entry(targets[index % 2], name).get("installationPreference") with ThreadPoolExecutor(max_workers=8) as pool: diff --git a/tests/integration/management/test_skill_download_history.py b/tests/integration/management/test_skill_download_history.py new file mode 100644 index 00000000000..00445dd7e27 --- /dev/null +++ b/tests/integration/management/test_skill_download_history.py @@ -0,0 +1,163 @@ +import hashlib +import io +import os +import uuid +import zipfile +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from functools import partial +from pathlib import Path +from typing import Final + +import psycopg +import pytest +from pydantic import TypeAdapter + +from litellm.models.plugin_upload import PluginDownloadsResponse +from litellm.types.proxy.skill_hub import SkillHubList +from tests.integration._support.actors import GatewayClient, owned_settings_client +from tests.integration._support.client import Gateway, Scenario, object_value, string_value +from tests.integration._support.database import write_rows + +PLUGINS: Final = "/claude-code/plugins" +SKILLS: Final = "/skill-hub/skills" +HISTORY: Final = TypeAdapter(PluginDownloadsResponse) + + +@pytest.fixture(scope="module") +def gateway(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("skill-downloads") + with owned_settings_client(GatewayClient(session_gateway, {}), directory, {"public_skills_index": True}) as client: + with client.gateway.scenario() as publisher: + yield replace(client.gateway, key=publisher.key(user_id=publisher.user(user_role="proxy_admin"))) + + +def _delete(gateway: Gateway, name: str) -> None: + assert gateway.request("DELETE", f"{PLUGINS}/{name}").status_code in {200, 404} + + +def _upload(scenario: Scenario, name: str) -> str: + scenario.cleanups.callback(_delete, scenario.gateway, name) + content: Final = io.BytesIO() + with zipfile.ZipFile(content, "w") as archive: + archive.writestr("SKILL.md", f"---\nname: {name}\ndescription: Synthetic download regression\n---\n") + response: Final = scenario.gateway.request_multipart( + f"{SKILLS}/upload", {"version": "1.2.3"}, {"file": ("skill.zip", content.getvalue(), "application/zip")} + ) + assert response.status_code == 200, response.text + return string_value(object_value(scenario.gateway.get(f"{PLUGINS}/{name}")["source"])["url"]) + + +def _download(gateway: Gateway, key: str | None, url: str) -> str: + response: Final = gateway.client.get( + url, headers={"Authorization": f"Bearer {key}"} if key else {}, params={"user_id": "forged-user"} + ) + assert response.status_code == 200, response.text + return hashlib.sha256(response.content).hexdigest() + + +def _history(gateway: Gateway, name: str, **params: str) -> PluginDownloadsResponse: + history: Final = gateway.get(f"{SKILLS}/{name}/downloads", params) + assert gateway.get(f"{PLUGINS}/{name}/downloads", params) == history + return HISTORY.validate_python(history) + + +def test_attribution_survives_actor_deletion_and_name_reuse_starts_empty(gateway: Gateway) -> None: + with gateway.scenario() as skills: + name: Final = f"download-{uuid.uuid4().hex}" + legacy: Final = _upload(skills, name) + portable: Final = f"{SKILLS}/{name}/archive" + original_id: Final = string_value(gateway.get(f"{SKILLS}/{name}")["id"]) + assert gateway.client.get(f"{SKILLS}/{name}/content").status_code == 200 + assert _history(gateway, name).total_downloads == 0 + assert SkillHubList.model_validate(gateway.get(SKILLS, {"query": name})).skills[0].download_count == 0 + with gateway.scenario() as actors: + team: Final = actors.team(team_alias=name) + email: Final = f"{name}@example.test" + user: Final = actors.user(user_role="internal_user", user_email=email) + gateway.post("/team/member_add", {"team_id": team, "member": {"role": "user", "user_id": user}}) + key: Final = actors.key(user_id=user, team_id=team, key_alias=name) + digests: Final = tuple(_download(gateway, key, url) for url in (legacy, portable)) + for url in (legacy, portable): + _download(gateway, None, url) + assert gateway.request("GET", url, key="sk-invalid").status_code == 401 + recorded: Final = _history(gateway, name) + assert (recorded.total_downloads, recorded.anonymous_downloads) == (4, 2) + assert SkillHubList.model_validate(gateway.get(SKILLS, {"query": name})).skills[0].download_count == 4 + assert tuple((row.version, row.authenticated, row.user_id, row.user_email, row.key_alias, + row.team_id, row.team_alias) for row in recorded.downloads) == ( + ("1.2.3", True, user, email, name, team, name), + ) * 2 + assert tuple(row.digest for row in recorded.downloads) == digests[::-1] + for prefix in (PLUGINS, SKILLS): + assert gateway.request("GET", f"{prefix}/{name}/downloads", key=key).status_code == 403 + assert gateway.client.get(f"{prefix}/{name}/downloads").status_code == 401 + assert _history(gateway, name) == recorded + _delete(gateway, name) + assert gateway.rows('SELECT id FROM "LiteLLM_ClaudeCodePluginDownload" WHERE plugin_id = %s', (original_id,)) == [] + _upload(skills, name) + assert gateway.get(f"{SKILLS}/{name}")["id"] != original_id + assert _history(gateway, name) == PluginDownloadsResponse(0, 0, (), None) + + +def test_concurrent_counts_and_keyset_pages_preserve_scope_and_deleted_cursors(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + name: Final = f"download-{uuid.uuid4().hex}" + legacy: Final = _upload(scenario, name) + portable: Final = f"{SKILLS}/{name}/archive" + key: Final = scenario.key() + with ThreadPoolExecutor(max_workers=4) as pool: + tuple(pool.map(partial(_download, gateway, key), (legacy, portable) * 4)) + initial: Final = _history(gateway, name) + ordered: Final = tuple(row.id for row in initial.downloads) + assert initial.total_downloads == len(frozenset(ordered)) == 8 + assert ordered == tuple(sorted(ordered, key=int, reverse=True)) + timestamp: Final = gateway.get(f"{PLUGINS}/{name}")["updated_at"] + with ThreadPoolExecutor(max_workers=4) as pool: + tuple(pool.map(partial(_download, gateway, None), (legacy, portable) * 8)) + assert gateway.get(f"{PLUGINS}/{name}")["updated_at"] == timestamp + assert _history(gateway, name) == replace(initial, total_downloads=24, anonymous_downloads=16) + first: Final = _history(gateway, name, limit="3") + assert tuple(row.id for row in first.downloads) == ordered[:3] + assert first.next_cursor == ordered[2] + _download(gateway, key, portable) + other: Final = f"{name}-other" + _download(gateway, key, _upload(scenario, other)) + write_rows( + 'DELETE FROM "LiteLLM_ClaudeCodePluginDownload" WHERE id IN (%s::bigint, %s::bigint)', + (ordered[2], ordered[4]), database_url=gateway.database_url, + ) + assert first.next_cursor is not None + second: Final = _history(gateway, name, limit="3", cursor=first.next_cursor) + assert tuple(row.id for row in second.downloads) == (ordered[3], ordered[5], ordered[6]) + assert (second.total_downloads, second.anonymous_downloads) == (23, 16) + assert second.next_cursor is not None + last: Final = _history(gateway, name, limit="3", cursor=second.next_cursor) + assert tuple(row.id for row in last.downloads) == (ordered[7],) + assert last.next_cursor is None + assert _history(gateway, other, cursor=first.next_cursor) == PluginDownloadsResponse(1, 0, (), None) + catalog: Final = SkillHubList.model_validate(gateway.get(SKILLS, {"query": name})) + assert {skill.name: skill.download_count for skill in catalog.skills} == {name: 23, other: 1} + + +def test_repeat_migration_is_catalog_only_while_download_tables_are_locked(gateway: Gateway) -> None: + migration: Final = (Path(__file__).resolve().parents[3] / + "litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql" + ).read_bytes() + with gateway.scenario() as scenario: + name: Final = f"download-{uuid.uuid4().hex}" + url: Final = _upload(scenario, name) + for key in (None, gateway.key): + _download(gateway, key, url) + before: Final = _history(gateway, name) + database_url: Final = gateway.database_url or os.environ["DATABASE_URL"] + with psycopg.connect(database_url) as locker, psycopg.connect(database_url, autocommit=True) as migrator: + locker.execute( + 'LOCK TABLE "LiteLLM_ClaudeCodePluginTable", "LiteLLM_ClaudeCodePluginDownload" IN ACCESS EXCLUSIVE MODE' + ) + for _ in range(2): + with migrator.transaction(): + migrator.execute("SET LOCAL statement_timeout = '500ms'") + migrator.execute(migration) + assert _history(gateway, name) == before diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index f6743a8751a..0c97ae62d07 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -42,7 +42,6 @@ from openai.types.chat import ChatCompletionMessageParam from openai.types.responses.tool_param import Mcp from pydantic import JsonValue, TypeAdapter - pytestmark: Final = [pytest.mark.xdist_group("mcp_llm_private_gateway"), pytest.mark.shares_worker_state] @@ -589,7 +588,11 @@ def _upstream_model(surface: Surface) -> str: }[surface] -PRICES: Final[Mapping[str, JsonValue]] = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002, "num_retries": 0} +PRICES: Final[Mapping[str, JsonValue]] = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "num_retries": 0, +} @contextmanager @@ -917,7 +920,9 @@ def test_an_open_server_in_another_gateway_cannot_widen_this_gateways_tools(gate register_mcp(scenario, local, "open" + uuid.uuid4().hex[:8], allow_all_keys=True) listed, called = official_client_outcomes(shared, key, "/mcp", f"{foreign_alias}-add", ADD) assert set(listed.tools) == { - f"{foreign_alias}-add", f"{foreign_alias}-multiply", f"{foreign_alias}-fail" + f"{foreign_alias}-add", + f"{foreign_alias}-multiply", + f"{foreign_alias}-fail", }, listed.tools assert called.ok and called.text == str(ADD["a"] + ADD["b"]), called assert [call["body"]["params"]["name"] for call in _peer_add_calls(foreign)] == ["add"] @@ -1528,6 +1533,7 @@ def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_ one_round_rows: Final = _model_rows( one_round.headers["x-litellm-call-id"], _logs_a_failure_billing(_price(FIRST_ROUND) + FIXED_MARGIN) ) + assert _statuses(one_round_rows) == ["failure"], one_round_rows rig.wire.drain() response: Final = rig.send(key, [AUTO], stream=True, guardrails=[STREAM_BLOCK]) assert STREAM_BLOCK_DETAIL in response.text, response.text diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index 91f502179f3..384d6635bb0 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -67,9 +67,9 @@ class _TokenDirs: def environment(self) -> Mapping[str, str]: return MappingProxyType({"CHATGPT_TOKEN_DIR": str(self.chatgpt), "GITHUB_COPILOT_TOKEN_DIR": str(self.copilot)}) - def store_valid_tokens(self, api_url: str) -> None: + def store_valid_tokens(self) -> None: dl.write_chatgpt(self.chatgpt, dl.chatgpt_record(dl.CHATGPT_STORED)) - dl.write_copilot_key(self.copilot, dl.COPILOT_STORED_KEY, api_url) + dl.write_copilot_access(self.copilot, dl.COPILOT_ACCESS) def _token_dirs(directory: Path) -> _TokenDirs: @@ -242,16 +242,16 @@ def test_booting_without_tokens_lists_login_deployments_and_serves_only_after_to assert peers.auth_connections() == () assert dl.bearer(peers.api.drain()[-1]) == dl.CONTROL_KEY - dirs.store_valid_tokens(peers.api.url) + dirs.store_valid_tokens() _reload(owned) _every_worker_lists(owned, workers) for served_logins, name in enumerate(_LOGIN_NAMES, start=1): served: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert served.status == 200, served.text _successes_logged(key, at_least=_CONTROL_CALLS_BEFORE_THE_RELOAD + served_logins) - assert peers.auth_connections() == () + peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) forwarded: Final = peers.api.drain() - assert {dl.bearer(request) for request in forwarded} == {dl.CHATGPT_STORED, dl.COPILOT_STORED_KEY} + assert {dl.bearer(request) for request in forwarded} == {dl.CHATGPT_STORED, dl.COPILOT_ACCESS} dl.write_chatgpt(dirs.chatgpt, dl.chatgpt_record(dl.CHATGPT_EXPIRED, refresh_token=dl.GOOD_REFRESH)) peers.switches.hold_connect.set() @@ -287,7 +287,7 @@ def test_a_token_removed_after_boot_turns_every_request_into_the_refusal( ) -> None: dirs: Final = _token_dirs(tmp_path) with dl.device_login_peers(tmp_path) as peers: - dirs.store_valid_tokens(peers.api.url) + dirs.store_valid_tokens() with owned_proxy_process( gateway, tmp_path, @@ -310,7 +310,7 @@ def test_a_token_removed_after_boot_turns_every_request_into_the_refusal( refused: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert refused.status == 400, refused.text assert refusal in refused.message(), refused.text - assert peers.auth_connections() == () + peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) tail: Final = owned.log.read_text() assert dl.CHATGPT_USER_CODE not in tail and dl.COPILOT_USER_CODE not in tail, tail[-3000:] control: Final = _chat(_base_url(owned), key, _CONTROL, uuid.uuid4().hex) diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index 39e36fd95fa..9074ea83189 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -2,7 +2,6 @@ import asyncio import hashlib import itertools import json -import re import threading import time import uuid @@ -26,9 +25,10 @@ from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.timing import assert_cut_short, timed from integration._support.wire import Request -from litellm.llms.chatgpt.authenticator import DEVICE_CODE_COOLDOWN_SECONDS from pydantic import JsonValue +from litellm.llms.chatgpt.authenticator import DEVICE_CODE_COOLDOWN_SECONDS + pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state] Provider: TypeAlias = Literal["chatgpt", "copilot"] @@ -52,7 +52,7 @@ _REFUSALS: Final[Mapping[Provider, str]] = MappingProxyType( {"chatgpt": dl.CHATGPT_REFUSAL, "copilot": dl.COPILOT_REFUSAL} ) _STORED_BEARERS: Final[Mapping[Provider, str]] = MappingProxyType( - {"chatgpt": dl.CHATGPT_STORED, "copilot": dl.COPILOT_STORED_KEY} + {"chatgpt": dl.CHATGPT_STORED, "copilot": dl.COPILOT_ACCESS} ) _CONTROL_MODEL: Final = "device-login-control" _EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({"num_retries": 0, "cache": {"no-cache": True}}) @@ -154,7 +154,7 @@ class _Scene: if provider == "chatgpt": dl.write_chatgpt(self.rig.chatgpt_dir, dl.chatgpt_record(dl.CHATGPT_STORED)) return - dl.write_copilot_key(self.rig.copilot_dir, dl.COPILOT_STORED_KEY, self.rig.peers.api.url) + dl.write_copilot_access(self.rig.copilot_dir, dl.COPILOT_ACCESS) def _config(peers: dl.Peers, directory: Path) -> Path: @@ -403,7 +403,7 @@ def test_missing_token_is_refused_with_the_login_instructions_before_any_auth_ho @pytest.mark.parametrize("cell", _CELLS, ids=_Cell.name) -def test_stored_token_serves_the_request_without_an_auth_host_call(scene: _Scene, cell: _Cell) -> None: +def test_stored_token_serves_the_request_without_a_device_login(scene: _Scene, cell: _Cell) -> None: scene.store_valid_token(cell.provider) call: Final = scene.call(cell) served: Final = _serve(call) @@ -412,7 +412,10 @@ def test_stored_token_serves_the_request_without_an_auth_host_call(scene: _Scene received: Final = scene.rig.peers.api.drain() (forwarded,) = _carrying(received, call.marker) assert dl.bearer(forwarded) == _STORED_BEARERS[cell.provider], _describe(received) - assert scene.rig.peers.auth_connections() == () + if cell.provider == "copilot": + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) + else: + assert scene.rig.peers.auth_connections() == () (row,) = _spend_rows(scene.key) assert row["status"] == "success", row @@ -487,7 +490,7 @@ def test_missing_token_is_refused_on_the_other_model_routes(scene: _Scene, route assert _describe(scene.rig.peers.api.drain()) == () -def test_copilot_embeddings_are_served_with_the_stored_api_key(scene: _Scene) -> None: +def test_copilot_embeddings_are_served_with_the_stored_access_token(scene: _Scene) -> None: scene.store_valid_token("copilot") served: Final = scene.post( "/v1/embeddings", {"model": "github_copilot/text-embedding-3-small", "input": "login", **_EXTRA} @@ -496,9 +499,9 @@ def test_copilot_embeddings_are_served_with_the_stored_api_key(scene: _Scene) -> (item,) = rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_json(served.text)["data"]) assert item["embedding"] == [0.25, 0.5, 0.75], served.text (forwarded,) = scene.rig.peers.api.drain() - assert (forwarded.method, dl.bearer(forwarded)) == ("POST", dl.COPILOT_STORED_KEY), forwarded.target + assert (forwarded.method, dl.bearer(forwarded)) == ("POST", dl.COPILOT_ACCESS), forwarded.target assert forwarded.target.endswith("/embeddings"), forwarded.target - assert scene.rig.peers.auth_connections() == () + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) @pytest.mark.parametrize("provider", _PROVIDERS) @@ -562,71 +565,67 @@ def test_stored_token_the_provider_rejects_surfaces_its_401_without_a_login( if provider == "chatgpt": dl.write_chatgpt(scene.rig.chatgpt_dir, dl.chatgpt_record(dl.CHATGPT_REJECTED)) else: - dl.write_copilot_key(scene.rig.copilot_dir, dl.COPILOT_REJECTED_KEY, scene.rig.peers.api.url) + dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_REJECTED_ACCESS) call: Final = scene.call(_Cell(provider, endpoint, False, "httpx")) served: Final = _serve(call) assert served.status == 401, served.text assert dl.REJECTED_DETAIL in served.text, served.text received: Final = scene.rig.peers.api.drain() assert {dl.bearer(request) for request in _carrying(received, call.marker)} == { - dl.CHATGPT_REJECTED if provider == "chatgpt" else dl.COPILOT_REJECTED_KEY + dl.CHATGPT_REJECTED if provider == "chatgpt" else dl.COPILOT_REJECTED_ACCESS }, _describe(received) - assert scene.rig.peers.auth_connections() == () + if provider == "copilot": + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_REJECTED_ACCESS) + else: + assert scene.rig.peers.auth_connections() == () @pytest.mark.parametrize("stale_key", (False, True), ids=("missing-api-key", "expired-api-key")) -def test_copilot_api_key_is_minted_once_from_the_stored_access_token(scene: _Scene, stale_key: bool) -> None: +def test_copilot_access_token_is_forwarded_without_minting_an_api_key(scene: _Scene, stale_key: bool) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_ACCESS) if stale_key: dl.write_copilot_key(scene.rig.copilot_dir, dl.COPILOT_REJECTED_KEY, scene.rig.peers.api.url, -3600) + legacy_key: Final = scene.rig.copilot_dir / "api-key.json" + stored_key: Final = legacy_key.read_bytes() if stale_key else None call: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) served: Final = _serve(call) assert served.status == 200, served.text assert set(rv.MARKER.findall(served.text)) == {call.marker}, served.text - (minted,) = scene.rig.peers.auth.drain() - assert (minted.method, minted.target, minted.headers.get("host"), minted.headers.get("authorization")) == ( - "GET", - "/copilot_internal/v2/token", - dl.GITHUB_API_HOST, - f"token {dl.COPILOT_ACCESS}", - ) + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) (forwarded,) = _carrying(scene.rig.peers.api.drain(), call.marker) - assert dl.bearer(forwarded) == dl.COPILOT_MINTED_KEY - assert json.loads((scene.rig.copilot_dir / "api-key.json").read_text())["token"] == dl.COPILOT_MINTED_KEY + assert dl.bearer(forwarded) == dl.COPILOT_ACCESS + assert (legacy_key.read_bytes() if legacy_key.exists() else None) == stored_key -def test_copilot_api_key_that_cannot_be_saved_is_minted_at_most_once_per_worker(scene: _Scene) -> None: +def test_copilot_access_token_serves_when_the_legacy_api_key_path_is_a_directory(scene: _Scene) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_UNSAVED_KEY_ACCESS) unwritable_key_file: Final = scene.rig.copilot_dir / "api-key.json" unwritable_key_file.mkdir() first: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) first_served: Final = _serve(first) assert first_served.status == 200, first_served.text - (minted,) = scene.rig.peers.auth.drain() - assert minted.headers.get("authorization") == f"token {dl.COPILOT_UNSAVED_KEY_ACCESS}" later: Final = tuple(scene.call(_Cell("copilot", "chat", False, "httpx")) for _ in range(3)) later_served: Final = tuple(_serve(call) for call in later) assert [served.status for served in later_served] == [200, 200, 200], [served.text for served in later_served] - assert len(scene.rig.peers.auth.drain()) <= 1 + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_UNSAVED_KEY_ACCESS) received: Final = scene.rig.peers.api.drain() forwarded: Final = tuple(_carrying(received, call.marker) for call in (first, *later)) assert [tuple(dl.bearer(request) for request in requests) for requests in forwarded] == [ - (dl.COPILOT_MINTED_KEY,) + (dl.COPILOT_UNSAVED_KEY_ACCESS,) ] * 4, _describe(received) assert unwritable_key_file.is_dir() def test_copilot_access_token_github_rejects_fails_the_request_without_a_device_login(scene: _Scene) -> None: dl.write_copilot_access(scene.rig.copilot_dir, dl.COPILOT_REJECTED_ACCESS) - served: Final = _serve(scene.call(_Cell("copilot", "chat", False, "httpx"))) - assert served.status == 400, served.text - assert "Failed to refresh API key" in _error_message(served.text), served.text + call: Final = scene.call(_Cell("copilot", "chat", False, "httpx")) + served: Final = _serve(call) + assert served.status == 401, served.text + assert dl.REJECTED_DETAIL in _error_message(served.text), served.text _failure_logged(served) - reached: Final = scene.rig.peers.auth.drain() - assert {(request.method, request.target) for request in reached} == {("GET", "/copilot_internal/v2/token")}, ( - _describe(reached) - ) - assert _describe(scene.rig.peers.api.drain()) == () + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_REJECTED_ACCESS) + (forwarded,) = _carrying(scene.rig.peers.api.drain(), call.marker) + assert (forwarded.method, dl.bearer(forwarded)) == ("POST", dl.COPILOT_REJECTED_ACCESS) def test_expired_copilot_api_key_with_no_access_token_is_refused_before_any_auth_host_call(scene: _Scene) -> None: @@ -727,7 +726,10 @@ def test_a_token_mounted_after_a_refusal_serves_without_a_restart_and_removing_i after: Final = _serve(scene.call(cell)) assert after.status == 400, after.text assert _REFUSALS[provider] in _error_message(after.text), after.text - assert scene.rig.peers.auth_connections() == () + if provider == "copilot": + scene.rig.peers.assert_only_copilot_discovery(dl.COPILOT_ACCESS) + else: + assert scene.rig.peers.auth_connections() == () _COOLDOWN_SECONDS_LEFT: Final = 45 diff --git a/tests/integration/providers/test_gemini_interactions_wire.py b/tests/integration/providers/test_gemini_interactions_wire.py index 58a7ccb31c3..ded0718bd2b 100644 --- a/tests/integration/providers/test_gemini_interactions_wire.py +++ b/tests/integration/providers/test_gemini_interactions_wire.py @@ -1,5 +1,4 @@ import json -import uuid from pathlib import Path from typing import Final @@ -147,6 +146,7 @@ def test_interaction_forwards_native_body_to_gemini(gateway: Gateway, prefix: st "agent": None, "updated": None, "outputs": None, + "continuation_token": None, }, response.text assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions")], response.text @@ -179,6 +179,7 @@ def test_interaction_authenticates_google_key_header(gateway: Gateway, prefix: s "agent": None, "updated": None, "outputs": None, + "continuation_token": None, }, response.text assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions")], response.text @@ -218,6 +219,7 @@ def test_agent_interaction_uses_env_gemini_credentials(gateway: Gateway, tmp_pat "agent": "deep-research-pro-preview-12-2025", "updated": None, "outputs": None, + "continuation_token": None, }, response.text assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions")], response.text @@ -277,3 +279,78 @@ def test_stream_interaction_ends_without_openai_done_terminator(gateway: Gateway expected: Final = b"".join(_proxy_stream_event(payload, model) for _event_type, payload in _STREAM_EVENTS) assert response.content == expected, response.text assert [(r.method, r.target) for r in wire.drain()] == [("POST", "/v1beta/interactions?alt=sse")] + + +_LONG_DECODE_MODEL: Final = "gemini/gemini-4-argon" +_LONG_DECODE_BODY: Final = { + "input": "prove the prime number theorem", + "generation_config": {"max_output_tokens": 120000, "thinking_level": "high"}, +} + + +def _hop(hop_id: str, status: str, text: str, token: str | None, input_tokens: int, output_tokens: int) -> dict: + return { + "id": hop_id, + "object": "interaction", + "model": "gemini-4-argon", + "status": status, + "created": "2026-11-05T00:00:00Z", + "steps": [{"type": "model_output", "content": [{"type": "text", "text": text}]}], + "usage": { + "total_input_tokens": input_tokens, + "total_output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + "initial_input_tokens": 79, + }, + "continuation_token": token, + } + + +_FIRST_HOP: Final = _hop("interaction-hop-1", "requires_continuation", "Lemma 1.", "tok-1", 79, 32768) +_SECOND_HOP: Final = _hop("interaction-hop-2", "completed", " QED.", None, 32847, 1533) + + +def test_interaction_keeps_requires_continuation_status_and_token_for_the_caller(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.target == "/v1beta/interactions" + return Reply(body=json.dumps(_FIRST_HOP).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_LONG_DECODE_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY) + key: Final = scenario.key(models=[model]) + response: Final = gateway.request( + "POST", "/v1beta/interactions", {"model": model, **_LONG_DECODE_BODY}, key=key + ) + assert response.status_code == 200, response.text + sent: Final = [json.loads(r.body) for r in wire.drain()] + + assert sent == [{"model": "gemini-4-argon", **_LONG_DECODE_BODY}], sent + payload: Final = response.json() + assert payload["id"] == "interaction-hop-1", response.text + assert payload["status"] == "requires_continuation", response.text + assert payload["continuation_token"] == "tok-1", response.text + assert payload["usage"]["initial_input_tokens"] == 79, response.text + + +def test_interaction_forwards_a_caller_continuation_token_to_resume_a_long_decode(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.target == "/v1beta/interactions" + return Reply(body=json.dumps(_SECOND_HOP).encode()) + + resume: Final = { + "generation_config": _LONG_DECODE_BODY["generation_config"], + "previous_interaction_id": "interaction-hop-1", + "continuation_token": "tok-1", + } + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_LONG_DECODE_MODEL, api_base=wire.url, api_key=_GEMINI_DEPLOYMENT_KEY) + key: Final = scenario.key(models=[model]) + response: Final = gateway.request("POST", "/v1beta/interactions", {"model": model, **resume}, key=key) + assert response.status_code == 200, response.text + sent: Final = [json.loads(r.body) for r in wire.drain()] + + assert sent == [{"model": "gemini-4-argon", **resume}], sent + payload: Final = response.json() + assert payload["id"] == "interaction-hop-2", response.text + assert payload["status"] == "completed", response.text + assert payload["continuation_token"] is None, response.text diff --git a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py index 018c6a653d1..696253f6a98 100644 --- a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py +++ b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py @@ -299,14 +299,14 @@ def cooldown_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Cooldown pytest.param("chat", _BRIDGED, "pause_after_content", id="chat-bridged-to-responses"), ), ) -def test_a_deployment_whose_stream_stalls_stays_out_of_cooldown( +def test_a_deployment_whose_stream_stalls_is_cooled_down( cooldown_rig: _CooldownRig, endpoint: Endpoint, backend: str, shape: Shape ) -> None: with _peer(shape) as wire, cooldown_rig.gateway.scenario() as scenario: group, identity = _deployment(scenario, wire, backend) response: Final = _streamed(cooldown_rig.gateway, endpoint, group) error: Final = _failure_row(cooldown_rig.gateway, _ROWS, (group,)) - assert not cooled(cooldown_rig.cache, identity), (error, response.text) + assert eventually(lambda: cooled(cooldown_rig.cache, identity), bool, _ROW_SECONDS), (error, response.text) @_OWNED_PROXY diff --git a/tests/integration/routing/test_include_fallback_errors_wire.py b/tests/integration/routing/test_include_fallback_errors_wire.py index 9075f5e41d3..2f48bc43a22 100644 --- a/tests/integration/routing/test_include_fallback_errors_wire.py +++ b/tests/integration/routing/test_include_fallback_errors_wire.py @@ -61,10 +61,10 @@ _PROVIDERS: Final = MappingProxyType( ) _STREAMING: Final = (pytest.param(False, id="non-stream"), pytest.param(True, id="stream")) _NON_BOOLEAN_FLAGS: Final = ( - pytest.param(1, True, id="int"), - pytest.param("", False, id="empty-string"), - pytest.param([], False, id="list"), - pytest.param("x" * 5120, True, id="five-kilobytes"), + pytest.param(1, id="int"), + pytest.param("", id="empty-string"), + pytest.param([], id="list"), + pytest.param("x" * 5120, id="five-kilobytes"), ) _OPENAI_UNAUTHORIZED: Final = Reply( status=401, @@ -326,10 +326,8 @@ def test_default_gateway_keeps_the_flag_off_the_wire_on_a_fallback( assert observed.attempted == "1", observed -@pytest.mark.parametrize(("value", "errors_reported"), _NON_BOOLEAN_FLAGS) -def test_default_gateway_keeps_a_non_boolean_flag_off_the_wire( - registered: _Registered, value: JsonValue, errors_reported: bool -) -> None: +@pytest.mark.parametrize("value", _NON_BOOLEAN_FLAGS) +def test_default_gateway_keeps_a_non_boolean_flag_off_the_wire(registered: _Registered, value: JsonValue) -> None: sent: Final = _send( registered.gateway, _PATHS["chat"], @@ -497,9 +495,9 @@ def test_exposed_chat_without_a_fallback_keeps_the_flag_off_the_wire_and_the_err assert observed.errors_header is None, observed -@pytest.mark.parametrize(("value", "errors_reported"), _NON_BOOLEAN_FLAGS) -def test_exposed_chat_keeps_a_non_boolean_flag_off_the_wire_and_the_errors_header_follows_its_truthiness( - exposed: _Exposed, value: JsonValue, errors_reported: bool +@pytest.mark.parametrize("value", _NON_BOOLEAN_FLAGS) +def test_exposed_chat_keeps_a_non_boolean_flag_off_the_wire_and_reports_no_errors( + exposed: _Exposed, value: JsonValue ) -> None: sent: Final = _send( exposed.proxy, @@ -515,7 +513,7 @@ def test_exposed_chat_keeps_a_non_boolean_flag_off_the_wire_and_the_errors_heade assert observed.leaks == (), observed assert observed.hit == (_EXPOSED_OPENAI_PRIMARY, _EXPOSED_OPENAI_BACKUP), observed assert observed.attempted == "1", observed - assert _errors_mention_the_outage(observed.errors_header) is errors_reported, observed + assert observed.errors_header is None, observed def test_exposed_proxy_rejects_an_unauthenticated_flagged_request_before_the_wire(exposed: _Exposed) -> None: diff --git a/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py new file mode 100644 index 00000000000..13bf67582b3 --- /dev/null +++ b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py @@ -0,0 +1,95 @@ +import collections +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import eventually, gateway_from_environment, object_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.wire import Reply, Wire, wire_server +from integration.routing._prompt_cache_pin_support import ( + DEFAULT_MARK, + anthropic_deployment, + ask_turn, + chat_body, + claude_peer, + new_marker, + one_new_pin, + pin_config, + post_chat, + read_pins, + system_turn, +) +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +from litellm.constants import REDACTED_BY_LITELLM + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +_GROUP: Final = "pin-logging-off" +_LEFT: Final = "pin-logging-off-left" +_RIGHT: Final = "pin-logging-off-right" +_FOLLOW_UPS: Final = 20 +_REDACTED_MESSAGES: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}] +_LOGGED_EVENT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + redis: "Redis[bytes]" + logger_sink: Wire + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("prompt-cache-pin-message-logging-off") + with ( + gateway_from_environment() as gateway, + owned_redis(directory) as cache, + wire_server(claude_peer) as left, + wire_server(claude_peer) as right, + wire_server(lambda request: Reply()) as logger_sink, + ): + config: Final = pin_config( + directory, + [anthropic_deployment(_GROUP, _LEFT, left.url), anthropic_deployment(_GROUP, _RIGHT, right.url)], + litellm_settings={ + "turn_off_message_logging": True, + "callbacks": ["integration._support.redaction_probe.probe"], + "cache": False, + "enable_redis_auth_cache": False, + }, + ) + environment: Final = { + "REDIS_HOST": cache.host, + "REDIS_PORT": str(cache.port), + "REDACTION_PROBE_URL": logger_sink.url, + } + with ( + owned_proxy_process(gateway, directory, environment, config=config, workers=1) as owned, + Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis, + ): + yield _Rig(owned, redis, logger_sink) + + +def test_follow_ups_stay_on_the_first_deployment_while_the_logger_sees_redacted_messages(rig: _Rig) -> None: + before: Final = frozenset(read_pins(rig.redis)) + body: Final = chat_body(_GROUP, [system_turn(new_marker(), DEFAULT_MARK), ask_turn(new_marker())]) + + _, first_served = post_chat(rig.owned.gateway, body) + _, pin = one_new_pin(rig.redis, before) + assert pin.model_id == first_served, (pin, first_served) + + served: Final = tuple(post_chat(rig.owned.gateway, body)[1] for _ in range(_FOLLOW_UPS)) + assert served == (first_served,) * _FOLLOW_UPS, collections.Counter(served) + + eventually(rig.logger_sink.received.qsize, lambda count: count >= 1 + _FOLLOW_UPS, seconds=30) + logged: Final = tuple( + object_value(_LOGGED_EVENT.validate_json(request.body)["standard"])["messages"] + for request in rig.logger_sink.drain() + ) + assert len(logged) == 1 + _FOLLOW_UPS, len(logged) + assert all(messages == _REDACTED_MESSAGES for messages in logged), logged diff --git a/tests/integration/sdk/test_device_code_login_guard_sdk.py b/tests/integration/sdk/test_device_code_login_guard_sdk.py index 7640fe13d5d..256871faf7c 100644 --- a/tests/integration/sdk/test_device_code_login_guard_sdk.py +++ b/tests/integration/sdk/test_device_code_login_guard_sdk.py @@ -184,7 +184,7 @@ def clean(bench: _Bench) -> Iterator[_Bench]: _FIRST_LOGIN_BEARERS: Final[Mapping[Provider, str]] = MappingProxyType( - {"chatgpt": dl.CHATGPT_FIRST_LOGIN, "copilot": dl.COPILOT_MINTED_KEY} + {"chatgpt": dl.CHATGPT_FIRST_LOGIN, "copilot": dl.COPILOT_FIRST_LOGIN_ACCESS} ) @@ -205,7 +205,7 @@ def test_sync_call_on_the_main_thread_keeps_the_interactive_first_login( assert {dl.bearer(request) for request in forwarded} == {_FIRST_LOGIN_BEARERS[provider]}, [ (request.method, request.target) for request in forwarded ] - expected_files: Final = ("auth.json",) if provider == "chatgpt" else ("access-token", "api-key.json") + expected_files: Final = ("auth.json",) if provider == "chatgpt" else ("access-token",) assert clean.token_files() == expected_files diff --git a/tests/integration/sdk/test_router_include_fallback_errors_wire.py b/tests/integration/sdk/test_router_include_fallback_errors_wire.py index fb17e05b02a..a8952c8c2b5 100644 --- a/tests/integration/sdk/test_router_include_fallback_errors_wire.py +++ b/tests/integration/sdk/test_router_include_fallback_errors_wire.py @@ -36,11 +36,11 @@ _BURST: Final = 12 _CALLBACK_WINDOW: Final = 15.0 _STREAMING: Final = (pytest.param(False, id="non-stream"), pytest.param(True, id="stream")) _NON_BOOLEAN_FLAGS: Final = ( - pytest.param(1, True, id="int"), - pytest.param("", False, id="empty-string"), - pytest.param([], False, id="list"), - pytest.param("x" * 5120, True, id="five-kilobytes"), - pytest.param(False, False, id="false"), + pytest.param(1, id="int"), + pytest.param("", id="empty-string"), + pytest.param([], id="list"), + pytest.param("x" * 5120, id="five-kilobytes"), + pytest.param(False, id="false"), ) _UNAUTHORIZED: Final = Reply( status=401, @@ -263,10 +263,8 @@ def test_without_a_fallback_the_flag_still_stays_off_the_wire() -> None: assert outcome.errors == (), outcome -@pytest.mark.parametrize(("value", "errors_reported"), _NON_BOOLEAN_FLAGS) -def test_a_non_boolean_flag_stays_off_the_wire_and_the_errors_header_follows_its_truthiness( - value: object, errors_reported: bool -) -> None: +@pytest.mark.parametrize("value", _NON_BOOLEAN_FLAGS) +def test_a_flag_other_than_true_stays_off_the_wire_and_reports_no_errors(value: object) -> None: with wire_server(_peer) as wire: router: Final = _router(wire) outcome: Final = _outcome(wire, _complete(router, "primary", stream=False, include_fallback_errors=value)) @@ -274,7 +272,7 @@ def test_a_non_boolean_flag_stays_off_the_wire_and_the_errors_header_follows_its assert outcome.hit == ("primary", "backup"), outcome assert outcome.text == f"{_ANSWER} backup", outcome assert outcome.attempted == 1, outcome - assert (len(outcome.errors) == 1) is errors_reported, outcome + assert outcome.errors == (), outcome class _Recorder(CustomLogger): diff --git a/tests/integration/security/test_skill_scanner.py b/tests/integration/security/test_skill_scanner.py new file mode 100644 index 00000000000..e18378463ee --- /dev/null +++ b/tests/integration/security/test_skill_scanner.py @@ -0,0 +1,58 @@ +import asyncio +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Final + +import psutil +import pytest +from fastapi import HTTPException + +from cookbook.skill_scanner import server +from tests.integration._support.process import _wait_until_gone, group_members, stop_processes + + +@pytest.mark.parametrize("mode", ("deadline", "cap", "complete")) +async def test_adapter_bounds_worker_lifetime_and_reaps_its_process_group( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, mode: str, +) -> None: + script: Final = ( + "import subprocess, sys, time\n" + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(30)'], " + "stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)\n" + "print(child.pid, flush=True)\n" + + ("pass\n" if mode == "complete" else "time.sleep(30)\n") + ) + child: Final = await asyncio.create_subprocess_exec( + sys.executable, "-c", script, stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, start_new_session=True, + ) + try: + assert child.stdout is not None + descendant: Final = psutil.Process(int(await asyncio.wait_for(child.stdout.readline(), 5))) + if mode == "complete": + await asyncio.wait_for(child.wait(), 5) + output: Final = tmp_path / "report.json" + output.write_text('{"complete": true}') + + async def launch(*args: object, **kwargs: object) -> asyncio.subprocess.Process: + return child + + monkeypatch.setattr(server.asyncio, "create_subprocess_exec", launch) + monkeypatch.setattr(server, "SCAN_TIMEOUT", 0.05 if mode == "cap" else 30) + deadline: Final = ( + None if mode == "cap" + else datetime.now(timezone.utc) + timedelta(seconds=0.05 if mode == "deadline" else 30) + ) + result: Final = server.run_worker("unused.zip", output, False, deadline) + if mode == "complete": + assert await asyncio.wait_for(result, 5) == {"complete": True} + else: + with pytest.raises(HTTPException) as raised: + await asyncio.wait_for(result, 5) + assert raised.value.status_code == 504 + assert child.returncode is not None + _wait_until_gone(descendant, 3) + finally: + await asyncio.to_thread(stop_processes, group_members(child.pid), 0.1) + await child.wait() diff --git a/tests/integration/spend/fixtures/daily_activity_user.json b/tests/integration/spend/fixtures/daily_activity_user.json index ae05e661071..058e12a5044 100644 --- a/tests/integration/spend/fixtures/daily_activity_user.json +++ b/tests/integration/spend/fixtures/daily_activity_user.json @@ -1 +1 @@ -{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"user-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}},"guardrails":{}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"user-1":5}}} +{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"user-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{"user_email":"user@example.com","user_alias":null},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}},"guardrails":{}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"user-1":5}}} diff --git a/tests/integration/spend/golden/daily_activity_user_aggregated.json b/tests/integration/spend/golden/daily_activity_user_aggregated.json index 185d333965a..3189c000035 100644 --- a/tests/integration/spend/golden/daily_activity_user_aggregated.json +++ b/tests/integration/spend/golden/daily_activity_user_aggregated.json @@ -1,5 +1,5 @@ { - "metadata": {"entity_total_api_keys":null, + "metadata": {"entity_total_api_keys":{"user-1": 5}, "api_key_limit": 100, "has_more": false, "page": 1, @@ -346,7 +346,181 @@ } } }, - "entities": {}, + "entities": { + "user-1": { + "api_key_breakdown": { + "key-a": { + "metadata": { + "key_alias": "alias-a", + "key_exists": true, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 1, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 2, + "spend": 100.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 4 + } + }, + "key-b": { + "metadata": { + "key_alias": "alias-b", + "key_exists": true, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 1, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 3, + "spend": 90.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 5 + } + }, + "key-c": { + "metadata": { + "key_alias": "alias-c", + "key_exists": true, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 1, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 4, + "spend": 80.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 6 + } + }, + "key-cache": { + "metadata": { + "key_alias": "alias-cache", + "key_exists": true, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 1, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 1000, + "spend": 2.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 1002 + } + }, + "key-target": { + "metadata": { + "key_alias": "deleted-target", + "key_exists": false, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 2, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 5, + "spend": 1.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 7 + } + } + }, + "metadata": { + "user_alias": null, + "user_email": "user@example.com" + }, + "metrics": { + "api_requests": 5, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 6, + "completion_tokens": 10, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 1014, + "spend": 1273.0, + "successful_requests": 5, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 1024 + } + } + }, "mcp_servers": {}, "model_groups": { "model-cache": { diff --git a/tests/integration/spend/test_daily_activity_key_owner.py b/tests/integration/spend/test_daily_activity_key_owner.py index cec19ce5ea0..5a9d8c7467c 100644 --- a/tests/integration/spend/test_daily_activity_key_owner.py +++ b/tests/integration/spend/test_daily_activity_key_owner.py @@ -57,6 +57,7 @@ def test_key_whose_daily_spend_names_two_users_is_reported_with_no_owner(gateway DAY, key_metadata(), seeded_metrics(2), + entity_rows={first: 1, second: 1}, ) @@ -74,6 +75,7 @@ def test_daily_spend_rows_naming_no_user_do_not_hide_the_one_user_the_others_nam DAY, key_metadata(user=owner, email=email), seeded_metrics(2), + entity_rows={owner: 1, "Unassigned": 1}, ) diff --git a/tests/integration/spend/test_daily_activity_key_owner_faults.py b/tests/integration/spend/test_daily_activity_key_owner_faults.py index 1dba97d6ae1..2eeca3c1588 100644 --- a/tests/integration/spend/test_daily_activity_key_owner_faults.py +++ b/tests/integration/spend/test_daily_activity_key_owner_faults.py @@ -199,7 +199,7 @@ def test_key_stops_being_reported_with_an_owner_once_a_second_user_spends_with_i with daily_rows((user_row(second, api_key, DAY),)): shared: Final = activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key) assert_key_reported(alone, api_key, DAY, key_metadata(user=first, email=email), seeded_metrics(1)) - assert_key_reported(shared, api_key, DAY, key_metadata(), seeded_metrics(2)) + assert_key_reported(shared, api_key, DAY, key_metadata(), seeded_metrics(2), entity_rows={first: 1, second: 1}) @pytest.mark.timeout(300) diff --git a/tests/unit/accounting/test_metrics.py b/tests/unit/accounting/test_metrics.py index 49b46eb36c2..054ceaba967 100644 --- a/tests/unit/accounting/test_metrics.py +++ b/tests/unit/accounting/test_metrics.py @@ -1,9 +1,17 @@ +import importlib.util +import sys + +import pytest from prometheus_client.metrics import MetricWrapperBase from litellm.accounting import metrics from litellm.accounting.metrics import record_deliveries from litellm.repositories.accounting_delivery_repository import DeliveryMetrics +_DELIVERY_METRICS = DeliveryMetrics( + pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192 +) + def _sample(metric: MetricWrapperBase, suffix: str = "") -> float: family = metric.collect()[0] @@ -11,18 +19,31 @@ def _sample(metric: MetricWrapperBase, suffix: str = "") -> float: def test_delivery_metrics_publish_backlog_and_clamp_clock_skew() -> None: - batches_before = _sample(metrics._BATCH_TIME, "_count") + instruments = metrics._INSTRUMENTS + assert instruments is not None + batches_before = _sample(instruments.batch_time, "_count") - record_deliveries( - DeliveryMetrics(pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192), - 0.5, - ) + record_deliveries(_DELIVERY_METRICS, 0.5) assert ( - _sample(metrics._PENDING), - _sample(metrics._PARKED), - _sample(metrics._OLDEST), - _sample(metrics._BATCH_COUNT), - _sample(metrics._TABLE_BYTES), + _sample(instruments.pending), + _sample(instruments.parked), + _sample(instruments.oldest), + _sample(instruments.batch_count), + _sample(instruments.table_bytes), ) == (4.0, 1.0, 0.0, 7.0, 8192.0) - assert _sample(metrics._BATCH_TIME, "_count") == batches_before + 1 + assert _sample(instruments.batch_time, "_count") == batches_before + 1 + + +def test_delivery_metrics_are_skipped_when_prometheus_client_is_not_installed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem(sys.modules, "prometheus_client", None) + spec = importlib.util.spec_from_file_location("accounting_metrics_without_prometheus", metrics.__file__) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + + spec.loader.exec_module(module) + + assert module.record_deliveries(_DELIVERY_METRICS, 0.5) is None diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py index 17186a89922..c78d6bec757 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -150,6 +150,7 @@ async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj) monkeypatch.setattr(proxy_server, "prisma_client", client) monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False) monkeypatch.setattr(client, "spend_log_transactions", []) monkeypatch.setattr(client, "model_usage_transactions", []) writer: Final = DBSpendUpdateWriter() @@ -521,7 +522,9 @@ def test_stored_payload_keeps_guardrail_billing_but_not_guardrail_content() -> N assert dict(daily_user.increments)["compression_saved_tokens"] == 7 -def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(monkeypatch: pytest.MonkeyPatch) -> None: +def test_caller_supplied_savings_beyond_json_range_are_clamped_and_never_drop_the_charge( + monkeypatch: pytest.MonkeyPatch, +) -> None: import litellm monkeypatch.setitem( @@ -552,6 +555,8 @@ def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(m ) restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str)) assert dict(restored["metadata"])["accounting_savings"]["compression"] == 0.0 - assert [dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"] == [ - 0.5 - ] + assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1} + assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1} + assert [ + dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken" + ] == [0.5] diff --git a/tests/unit/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py index 85422819926..fa0b7631955 100644 --- a/tests/unit/caching/test_redis_cluster_cache.py +++ b/tests/unit/caching/test_redis_cluster_cache.py @@ -155,20 +155,15 @@ def test_cache_init_creates_redis_cache_without_cluster_config( ) def test_router_create_redis_cache_cluster_detection(startup_nodes, env_var, expected_cache_type, monkeypatch): """ - Test that Router._create_redis_cache() creates RedisClusterCache when - either startup_nodes is in config or REDIS_CLUSTER_NODES env var is set. + Test that a Router given Redis settings builds a RedisClusterCache when + either startup_nodes is in cache_kwargs or REDIS_CLUSTER_NODES env var is set. Also verifies that explicit startup_nodes take precedence over env var. Regression test for https://github.com/BerriAI/litellm/issues/22748 """ from litellm import Router - cache_config = dict( - host="mockhost", - port=6379, - password="mock-password", - startup_nodes=startup_nodes, - ) + cache_kwargs = {"startup_nodes": startup_nodes} if startup_nodes is not None else None if env_var is not None: monkeypatch.setenv("REDIS_CLUSTER_NODES", env_var) @@ -178,8 +173,15 @@ def test_router_create_redis_cache_cluster_detection(startup_nodes, env_var, exp def _mock_redis_cache_init(*args, **kwargs): ... with patch.object(RedisCache, "__init__", _mock_redis_cache_init): - redis_cache = Router._create_redis_cache(cache_config) - assert isinstance(redis_cache, expected_cache_type) + router = Router( + model_list=[{"model_name": "m", "litellm_params": {"model": "gpt-4o", "api_key": "x"}}], + redis_host="mockhost", + redis_port=6379, + redis_password="mock-password", + cache_kwargs=cache_kwargs, + ) + redis_cache = router.health().store.redis_cache + assert type(redis_cache) is expected_cache_type def _isolated_redis_cache(host: str) -> RedisCache: diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 4a07c97b709..7a224232f4c 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -44,20 +44,33 @@ os.environ["PYTHON_DOTENV_DISABLED"] = "1" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import -import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency +from litellm.anthropic_beta_headers_manager import ( + reload_beta_headers_config, # noqa: E402 # same import-time dependency +) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER # noqa: E402 # same import-time dependency -from litellm.anthropic_beta_headers_manager import reload_beta_headers_config # noqa: E402 # same import-time dependency -from litellm.litellm_core_utils.prompt_templates import factory as prompt_factory_module # noqa: E402 # same import-time dependency +from litellm.litellm_core_utils.prompt_templates import ( + factory as prompt_factory_module, # noqa: E402 # same import-time dependency +) from litellm.litellm_core_utils.prompt_templates import ( # noqa: E402 # same import-time dependency image_handling as image_handling_module, ) -from litellm.llms.gemini.chat import transformation as gemini_chat_transformation_module # noqa: E402 # same import-time dependency from litellm.llms.custom_httpx.async_client_cleanup import ( # noqa: E402 # same import-time dependency close_litellm_async_clients, ) -from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module # noqa: E402 # same import-time dependency +from litellm.llms.gemini.chat import ( + transformation as gemini_chat_transformation_module, # noqa: E402 # same import-time dependency +) +from litellm.proxy.db import ( + tool_registry_writer as tool_registry_writer_module, # noqa: E402 # same import-time dependency +) +from litellm.router_rust.surface.replay import ( # noqa: E402 # same import-time dependency + _live_registries as live_router_registries, # pyright: ignore[reportPrivateUsage] # the routers the model-cost replay reaches +) +from litellm.router_utils.global_callback_holders import ( # noqa: E402 # same import-time dependency + routers_holding_global_callbacks, +) LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"] AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( @@ -252,12 +265,6 @@ def socket_tempdir(socket_dir: Path, monkeypatch: pytest.MonkeyPatch) -> Path: return socket_dir -@pytest.fixture -def python_router_core(monkeypatch: pytest.MonkeyPatch) -> None: - """Keep ``Router`` on ``PythonCore``, for a test that patches a member the Rust core never reads.""" - monkeypatch.setenv("LITELLM_RUST", "0") - - @pytest.fixture(autouse=True, scope="session") def bundled_tiktoken_cache() -> None: importlib.import_module("litellm.litellm_core_utils.default_encoding") @@ -293,10 +300,12 @@ def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> It @pytest.fixture(autouse=True) def isolate_litellm_globals() -> Iterator[None]: original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS} - original_callback_holders: Final = frozenset(litellm_router_module._routers_holding_global_callbacks) + original_callback_holders: Final = frozenset(routers_holding_global_callbacks) original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS} original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)} - original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__} + original_clients: Final = { + name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__ + } original_loggers: Final = { logger: (logger.level, logger.disabled, logger.propagate, list(logger.handlers), list(logger.filters)) for logger in ALL_LOGGERS @@ -305,7 +314,7 @@ def isolate_litellm_globals() -> Iterator[None]: _flush_client_caches() for name in CALLBACK_LISTS: setattr(litellm, name, []) - litellm_router_module._routers_holding_global_callbacks.clear() + routers_holding_global_callbacks.clear() for name in RESET_TO_NONE_GLOBALS: setattr(litellm, name, None) for name in MODULE_LEVEL_CLIENTS: @@ -320,8 +329,8 @@ def isolate_litellm_globals() -> Iterator[None]: litellm.__dict__.update(original_clients) for name, value in (original_callbacks | original_reset | original_restored).items(): setattr(litellm, name, value) - litellm_router_module._routers_holding_global_callbacks.clear() - litellm_router_module._routers_holding_global_callbacks.update(original_callback_holders) + routers_holding_global_callbacks.clear() + routers_holding_global_callbacks.update(original_callback_holders) for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items(): logger.setLevel(level) logger.disabled = disabled @@ -333,19 +342,17 @@ def isolate_litellm_globals() -> Iterator[None]: @pytest.fixture(autouse=True) def isolate_router_model_cost_state() -> Iterator[None]: - original_live_routers: Final = frozenset(litellm_router_module._live_routers) + original_live_registries: Final = frozenset(live_router_registries) original_runtime_registered_model_cost: Final = { model_key: dict(model_value) for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() } litellm_utils_module._invalidate_model_cost_lowercase_map() - litellm_router_module._live_routers.clear() + live_router_registries.clear() litellm_utils_module._runtime_registered_model_cost.clear() yield - for router in tuple(litellm_router_module._live_routers): - litellm_router_module._live_routers.discard(router) - for router in original_live_routers: - litellm_router_module._live_routers.add(router) + live_router_registries.clear() + live_router_registries.update(original_live_registries) litellm_utils_module._runtime_registered_model_cost.clear() litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) litellm_utils_module._invalidate_model_cost_lowercase_map() diff --git a/tests/unit/containers/test_azure_container_transformation.py b/tests/unit/containers/test_azure_container_transformation.py index 38ea9323017..87b3def2a0b 100644 --- a/tests/unit/containers/test_azure_container_transformation.py +++ b/tests/unit/containers/test_azure_container_transformation.py @@ -558,13 +558,14 @@ class TestAzureContainerKnownFailureRegressions: self, monkeypatch ): """Operations without container_id (create, list) must NOT route via - _ageneric_api_call_with_fallbacks using a caller-supplied model_id. + a deployment using a caller-supplied model_id. Security boundary: only the path that holds a validated container_id is trusted to fall back to the forwarded model_id. A caller setting model_id without container_id on POST /v1/containers must not gain access to an arbitrary deployment UUID. """ + import litellm.containers from litellm.router import Router router = Router( @@ -582,32 +583,27 @@ class TestAzureContainerKnownFailureRegressions: ] ) - fallback_called = {"called": False} + called_with: dict = {} - async def _mock_fallback(original_function, **kwargs): - fallback_called["called"] = True + async def _capture(**kwargs): + called_with.update(kwargs) return {} - monkeypatch.setattr(router, "_ageneric_api_call_with_fallbacks", _mock_fallback) + monkeypatch.setattr(litellm.containers, "acreate_container", _capture) - original_called = {"called": False} - - async def _noop(**kwargs): - original_called["called"] = True - return {} - - # No container_id — simulates create/list; caller injects a model_id - await router._init_containers_api_endpoints( - original_function=_noop, + # No container_id — simulates create; caller injects a model_id + await router.acreate_container( + name="c", model_id="deployment-uuid-123", custom_llm_provider="azure", ) - assert not fallback_called["called"], ( - "_ageneric_api_call_with_fallbacks must NOT be called when " + assert called_with, "the container call must still run directly" + assert called_with.get("api_key") is None, ( + "the deployment's credentials must NOT be applied when " "container_id is absent, even if model_id is supplied" ) - assert original_called["called"], "original_function must be called directly" + assert called_with.get("api_base") is None def test_regression_httpx_empty_params_strips_query_string(self): """httpx erases the URL query-string when params={} (empty dict) is passed. @@ -974,9 +970,10 @@ class TestAzureContainerKnownFailureRegressions: self, monkeypatch ): """Native Azure container IDs (cntr_ + hex, no LiteLLM payload) must - still route through _ageneric_api_call_with_fallbacks using the - model_id forwarded from the proxy ownership check so that deployment - credentials (api_base) are applied.""" + still route through the deployment named by the model_id forwarded + from the proxy ownership check so that deployment credentials + (api_base) are applied.""" + import litellm.containers from litellm.router import Router router = Router( @@ -996,25 +993,22 @@ class TestAzureContainerKnownFailureRegressions: called_with: dict = {} - async def _mock_fallback(original_function, **kwargs): + async def _capture(**kwargs): called_with.update(kwargs) return {} - monkeypatch.setattr(router, "_ageneric_api_call_with_fallbacks", _mock_fallback) + monkeypatch.setattr(litellm.containers, "aretrieve_container", _capture) native_azure_id = "cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df" - async def _noop(**kwargs): - return {} - - await router._init_containers_api_endpoints( - original_function=_noop, + await router.aretrieve_container( container_id=native_azure_id, model_id="deployment-uuid-123", custom_llm_provider="azure", ) - assert called_with.get("model") == "deployment-uuid-123", ( - "_ageneric_api_call_with_fallbacks must be called with the forwarded " - "model_id when the container_id carries no LiteLLM routing payload" + assert called_with.get("api_base") == "https://my-resource.cognitiveservices.azure.com", ( + "the forwarded model_id's deployment credentials must be applied " + "when the container_id carries no LiteLLM routing payload" ) + assert called_with.get("container_id") == native_azure_id diff --git a/tests/unit/integration_support/test_wire.py b/tests/unit/integration_support/test_wire.py new file mode 100644 index 00000000000..c786895720f --- /dev/null +++ b/tests/unit/integration_support/test_wire.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +import datetime +import io +import socket +import ssl +import time +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from http.client import HTTPConnection, HTTPResponse, HTTPSConnection +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import NameOID + +from tests.integration._support.wire import Reply, Request, wire_server + +PAUSE_SECONDS: Final = 6 +HANDLER_READ_SECONDS: Final = 0.3 +FIRST: Final = b'{"call": "first"}' +LATE: Final = b'{"call": "late"}' +UNREAD: Final = b"x" * (32 * 1024 * 1024) + + +def _self_signed_server_context(directory: Path) -> ssl.SSLContext: + key: Final = ec.generate_private_key(ec.SECP256R1()) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "wire peer")]) + now: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + certificate: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(1) + .not_valid_before(now) + .not_valid_after(now + datetime.timedelta(days=36500)) + .sign(key, hashes.SHA256()) + ) + (directory / "peer.pem").write_bytes( + certificate.public_bytes(serialization.Encoding.PEM) + + key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + ) + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(directory / "peer.pem") + return context + + +def _trusting() -> ssl.SSLContext: + context: Final = ssl.create_default_context() + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + return context + + +def _client(url: str) -> HTTPConnection: + address: Final = urlsplit(url) + if address.scheme == "http": + return HTTPConnection(address.hostname or "", address.port) + return HTTPSConnection(address.hostname or "", address.port, context=_trusting()) + + +def _tls_socket(url: str) -> ssl.SSLSocket: + address: Final = urlsplit(url) + return _trusting().wrap_socket(socket.create_connection((address.hostname or "", address.port or 0))) + + +def _echo(request: Request) -> Reply: + return Reply(body=request.body) + + +def _answer_slowly(request: Request) -> Reply: + time.sleep(HANDLER_READ_SECONDS) + return Reply(body=request.body) + + +def _answer_in_paused_chunks(request: Request) -> Reply: + return Reply(chunks=(request.body[:1], request.body[1:]), pause_between_chunks=HANDLER_READ_SECONDS) + + +def _hold_first(request: Request) -> Reply: + if request.target == "/first": + time.sleep(HANDLER_READ_SECONDS) + return Reply(body=request.body) + + +def _paced_pieces(client: socket.socket) -> Iterator[bytes]: + while piece := client.recv(256 * 1024): + yield piece + time.sleep(0.02) + + +def _read_slowly(client: socket.socket) -> bytes: + return b"".join(_paced_pieces(client)) + + +def _read_after(client: socket.socket, delay: float) -> bytes: + time.sleep(delay) + with client.makefile("rb") as stream: + return stream.read() + + +def _post(target: str, body: bytes) -> bytes: + return b"POST %s HTTP/1.1\r\nhost: peer\r\ncontent-length: %d\r\n\r\n%s" % (target.encode(), len(body), body) + + +def _pause_then_send(connection: HTTPConnection, body: bytes) -> HTTPResponse: + time.sleep(PAUSE_SECONDS) + connection.request("POST", "/late", body) + return connection.getresponse() + + +def _send_headers_then_pause_before_the_body(connection: HTTPConnection, body: bytes) -> HTTPResponse: + connection.putrequest("POST", "/late") + connection.putheader("content-length", str(len(body))) + connection.endheaders() + time.sleep(PAUSE_SECONDS) + connection.send(body) + return connection.getresponse() + + +@pytest.mark.parametrize( + "send_late", + (_pause_then_send, _send_headers_then_pause_before_the_body), + ids=("pause-before-the-request", "pause-before-the-body"), +) +def test_a_kept_alive_connection_answers_a_request_whose_bytes_arrive_after_a_long_pause( + send_late: Callable[[HTTPConnection, bytes], HTTPResponse], +) -> None: + with wire_server(_echo, keep_alive=True) as wire: + address: Final = urlsplit(wire.url) + connection: Final = HTTPConnection(address.hostname or "", address.port) + connection.request("POST", "/first", FIRST) + first: Final = connection.getresponse().read() + late: Final = send_late(connection, LATE) + assert (first, late.status, late.read()) == (FIRST, 200, LATE) + assert tuple(request.body for request in wire.drain()) == (FIRST, LATE) + assert wire.connections() == 1 + assert connection.sock.recv(1) == b"" + connection.close() + + +def test_a_request_cut_off_by_teardown_never_reaches_respond() -> None: + with wire_server(_echo, keep_alive=True) as wire: + address: Final = urlsplit(wire.url) + connection: Final = HTTPConnection(address.hostname or "", address.port) + connection.request("POST", "/first", FIRST) + first: Final = connection.getresponse().read() + connection.sock.sendall(b"POST /cut HTTP/1.1\r\ncontent-type: application/json\r\ncontent-len") + time.sleep(HANDLER_READ_SECONDS) + assert (first, tuple(request.target for request in wire.drain())) == (FIRST, ("/first",)) + connection.close() + + +@pytest.mark.timeout(60) +def test_a_reply_its_client_stopped_reading_fails_teardown_instead_of_hanging() -> None: + peer: Final = ExitStack() + wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=0.5)) + address: Final = urlsplit(wire.url) + with socket.create_connection((address.hostname or "", address.port or 0)) as client: + client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n") + while wire.received.empty(): + time.sleep(0.01) + with pytest.raises(AssertionError, match="no progress"): + peer.close() + + +@pytest.mark.timeout(60) +def test_a_long_reply_its_client_is_still_reading_at_teardown_is_not_cut() -> None: + peer: Final = ExitStack() + wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=0.5)) + address: Final = urlsplit(wire.url) + with socket.create_connection((address.hostname or "", address.port or 0)) as client, ThreadPoolExecutor(1) as pool: + client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n") + while wire.received.empty(): + time.sleep(0.01) + reply: Final = pool.submit(_read_slowly, client) + time.sleep(1) + peer.close() + received: Final = reply.result() + assert (received[: len(b"HTTP/1.1 200")], received[-len(UNREAD) :] == UNREAD) == (b"HTTP/1.1 200", True) + + +@pytest.mark.timeout(60) +def test_a_reply_its_client_starts_reading_after_teardown_starts_still_arrives_whole() -> None: + peer: Final = ExitStack() + wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=2)) + address: Final = urlsplit(wire.url) + with socket.create_connection((address.hostname or "", address.port or 0)) as client, ThreadPoolExecutor(1) as pool: + client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n") + while wire.received.empty(): + time.sleep(0.01) + time.sleep(2.5) + reply: Final = pool.submit(_read_after, client, 0.3) + peer.close() + received: Final = reply.result() + assert (received[: len(b"HTTP/1.1 200")], received[-len(UNREAD) :] == UNREAD) == (b"HTTP/1.1 200", True) + + +def test_a_complete_request_sent_as_teardown_starts_is_still_answered() -> None: + with wire_server(_echo, keep_alive=True) as wire: + connection: Final = _client(wire.url) + connection.request("POST", "/first", FIRST) + first: Final = connection.getresponse().read() + connection.request("POST", "/complete", LATE) + late: Final = connection.getresponse() + assert (first, late.status, late.read(), tuple(request.target for request in wire.drain())) == ( + FIRST, + 200, + LATE, + ("/first", "/complete"), + ) + connection.close() + + +@pytest.mark.parametrize("secure", (False, True), ids=("plain", "tls")) +def test_a_reply_still_being_written_at_teardown_reaches_its_client_whole(secure: bool, tmp_path: Path) -> None: + slow: Final = Reply(chunks=(b"first", b"second"), pause_between_chunks=HANDLER_READ_SECONDS) + tls: Final = _self_signed_server_context(tmp_path) if secure else None + with wire_server(lambda _: slow, tls=tls, keep_alive=True) as wire: + connection: Final = _client(wire.url) + connection.request("GET", "/slow") + response: Final = connection.getresponse() + head: Final = response.read(len(b"first")) + assert head + response.read() == b"firstsecond" + connection.close() + + +def test_a_complete_request_queued_behind_a_busy_handler_is_still_answered() -> None: + with wire_server(_hold_first, keep_alive=True) as wire: + address: Final = urlsplit(wire.url) + client: Final = socket.create_connection((address.hostname or "", address.port or 0)) + client.sendall(_post("/first", FIRST)) + while wire.received.empty(): + time.sleep(0.01) + client.sendall(_post("/second", LATE)) + with client, client.makefile("rb") as stream: + replies: Final = stream.read() + assert ( + tuple(request.target for request in wire.drain()), + replies.count(b"HTTP/1.1 200"), + replies[-len(LATE) :], + ) == ( + ("/first", "/second"), + 2, + LATE, + ) + + +@pytest.mark.parametrize("respond", (_answer_slowly, _answer_in_paused_chunks), ids=("slow-respond", "paused-chunks")) +def test_a_reply_that_takes_longer_than_the_teardown_grace_still_reaches_its_client( + respond: Callable[[Request], Reply], +) -> None: + with wire_server(respond, keep_alive=True, teardown_grace=HANDLER_READ_SECONDS / 2) as wire: + connection: Final = _client(wire.url) + connection.request("POST", "/slow", FIRST) + while wire.received.empty(): + time.sleep(0.01) + response: Final = connection.getresponse() + assert (response.status, response.read()) == (200, FIRST) + connection.close() + + +@pytest.mark.timeout(60) +def test_a_tls_request_larger_than_the_read_buffer_is_answered(tmp_path: Path) -> None: + large: Final = b"x" * (3 * io.DEFAULT_BUFFER_SIZE // 2) + with wire_server(_echo, tls=_self_signed_server_context(tmp_path), keep_alive=True) as wire: + client: Final = _tls_socket(wire.url) + client.sendall(_post("/large", large)) + response: Final = HTTPResponse(client) + response.begin() + assert (response.status, response.read()) == (200, large) + client.close() + + +@pytest.mark.timeout(60) +def test_a_tls_record_cut_off_by_teardown_closes_its_connection_instead_of_hanging(tmp_path: Path) -> None: + tls: Final = _self_signed_server_context(tmp_path) + with wire_server(_echo, tls=tls, keep_alive=True, teardown_grace=HANDLER_READ_SECONDS) as wire: + client: Final = _tls_socket(wire.url) + socket.socket.sendall(client, b"\x17\x03\x03\x40\x00" + b"x" * 16) + time.sleep(HANDLER_READ_SECONDS) + assert wire.drain() == () + client.close() diff --git a/tests/unit/interactions/interactions.openapi.json b/tests/unit/interactions/interactions.openapi.json index 3830f5093cb..b67ad0aceaa 100644 --- a/tests/unit/interactions/interactions.openapi.json +++ b/tests/unit/interactions/interactions.openapi.json @@ -756,7 +756,7 @@ { "label": "list", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n res, err := client.Environments.ListEnvironments(ctx, operations.ListEnvironmentsRequest{})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.ListEnvironmentsResponse != nil {\n for _, env := range res.ListEnvironmentsResponse.Environments {\n if env.ID != \"\" {\n fmt.Println(env.ID)\n }\n }\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n res, err := client.Environments.ListEnvironments(ctx, operations.ListEnvironmentsRequest{})\n if err != nil {\n log.Fatal(err)\n }\n\n if res.ListEnvironmentsResponse != nil {\n for _, env := range res.ListEnvironmentsResponse.Environments {\n if env.ID != \"\" {\n fmt.Println(env.ID)\n }\n }\n }\n}\n" } ] }, @@ -834,12 +834,12 @@ { "label": "create", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n req := operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n Sources: []gaos_interactions.Source{\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeInline),\n Target: genai.Ptr(\"main.py\"),\n Content: genai.Ptr(\"print('Hello, World!')\"),\n },\n },\n },\n }\n\n res, err := client.Environments.CreateEnvironment(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n req := operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n Sources: []gaos_interactions.Source{\n {\n Type: genai.Ptr(gaos_interactions.SourceTypeInline),\n Target: genai.Ptr(\"main.py\"),\n Content: genai.Ptr(\"print('Hello, World!')\"),\n },\n },\n },\n }\n\n res, err := client.Environments.CreateEnvironment(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" }, { "label": "copy", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.CreateEnvironment(ctx, operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n FromEnvironment: genai.Ptr(sourceEnvID),\n },\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_environments \"google.golang.org/genai/interactions/models/environments\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n res, err := client.Environments.CreateEnvironment(ctx, operations.CreateEnvironmentRequest{\n Body: gaos_environments.CreateEnvironmentRequest{\n FromEnvironment: genai.Ptr(sourceEnvID),\n },\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n\n }\n}\n" } ] } @@ -967,12 +967,12 @@ { "label": "list_files", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"src\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n fileType := \"\"\n if file.Type != nil {\n fileType = string(*file.Type)\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, fileType, size)\n }\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"src\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n fileType := \"\"\n if file.Type != nil {\n fileType = string(*file.Type)\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, fileType, size)\n }\n }\n}\n" }, { "label": "get_file", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"main.py\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, size)\n }\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n res, err := client.Environments.Files.List(ctx, operations.GetEnvironmentFilesRequest{\n Environment: envID,\n Path: \"main.py\",\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.GetEnvironmentFilesResponse != nil {\n for _, file := range res.GetEnvironmentFilesResponse.Files {\n name := \"\"\n if file.Name != nil {\n name = *file.Name\n }\n var size int64\n if file.SizeBytes != nil {\n size = *file.SizeBytes\n }\n fmt.Println(name, size)\n }\n }\n}\n" } ] } @@ -1036,7 +1036,7 @@ { "label": "delete", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n _, err = client.Environments.DeleteEnvironment(ctx, operations.DeleteEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n fmt.Println(\"Environment deleted successfully.\")\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n _, err = client.Environments.DeleteEnvironment(ctx, operations.DeleteEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n fmt.Println(\"Environment deleted successfully.\")\n}\n" } ] }, @@ -1093,7 +1093,7 @@ { "label": "get", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client := interactions.New()\n\n\n res, err := client.Environments.GetEnvironment(ctx, operations.GetEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n\n res, err := client.Environments.GetEnvironment(ctx, operations.GetEnvironmentRequest{\n ID: envID,\n })\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Environment != nil && res.Environment.ID != \"\" {\n fmt.Println(res.Environment.ID)\n }\n}\n" } ] } @@ -1171,7 +1171,7 @@ "deep_research": { "summary": "Deep Research", "value": { - "agent": "deep-research-pro-preview-12-2025", + "agent": "deep-research-preview-04-2026", "created": "2025-11-26T12:22:47Z", "id": "v1_ChdPU0F4YWFtNkFwS2kxZThQZ05lbXdROBIXT1NBeGFhbTZBcEtpMWU4UGdOZW13UTg", "object": "interaction", @@ -1472,7 +1472,7 @@ { "label": "deep_research", "lang": "sh", - "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": \"deep-research-pro-preview-12-2025\",\n \"input\": \"Find a cure to cancer\",\n \"background\": true\n }'\n" + "source": "curl -X POST https://generativelanguage.googleapis.com/v1beta/interactions \\\n -H \"x-goog-api-key: $GEMINI_API_KEY\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": \"deep-research-preview-04-2026\",\n \"input\": \"Find a cure to cancer\",\n \"background\": true\n }'\n" }, { "label": "antigravity", @@ -1517,7 +1517,7 @@ { "label": "deep_research", "lang": "python", - "source": "from google import genai\n\nclient = genai.Client()\ninteraction = client.interactions.create(\n agent=\"deep-research-pro-preview-12-2025\",\n input=\"find a cure to cancer\",\n background=True,\n)\nprint(interaction.status)\n" + "source": "from google import genai\n\nclient = genai.Client()\ninteraction = client.interactions.create(\n agent=\"deep-research-preview-04-2026\",\n input=\"find a cure to cancer\",\n background=True,\n)\nprint(interaction.status)\n" }, { "label": "antigravity", @@ -1562,7 +1562,7 @@ { "label": "deep_research", "lang": "javascript", - "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n agent: 'deep-research-pro-preview-12-2025',\n input: 'find a cure to cancer',\n background: true,\n});\nconsole.log(interaction.status);\n" + "source": "import {GoogleGenAI} from '@google/genai';\n\nconst ai = new GoogleGenAI({});\nconst interaction = await ai.interactions.create({\n agent: 'deep-research-preview-04-2026',\n input: 'find a cure to cancer',\n background: true,\n});\nconsole.log(interaction.status);\n" }, { "label": "antigravity", @@ -1607,7 +1607,7 @@ { "label": "deep_research", "lang": "java", - "source": "import com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionStatus;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateAgentInteraction params =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"deep-research-pro-preview-12-2025\"))\n .input(InteractionsInput.of(\"find a cure to cancer\"))\n .background(true)\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.status().map(InteractionStatus::value).orElse(\"\"));\n" + "source": "import com.google.genai.gaos.models.interactions.AgentOption;\nimport com.google.genai.Client;\nimport com.google.genai.gaos.models.interactions.CreateAgentInteraction;\nimport com.google.genai.gaos.models.operations.CreateInteractionRequestBody;\nimport com.google.genai.gaos.models.operations.CreateInteractionResponse;\nimport com.google.genai.gaos.models.interactions.Interaction;\nimport com.google.genai.gaos.models.interactions.InteractionStatus;\nimport com.google.genai.gaos.models.interactions.InteractionsInput;\n\nClient client = new Client();\nCreateAgentInteraction params =\n CreateAgentInteraction.builder()\n .agent(AgentOption.of(\"deep-research-preview-04-2026\"))\n .input(InteractionsInput.of(\"find a cure to cancer\"))\n .background(true)\n .build();\nCreateInteractionResponse response =\n client.interactions.create(CreateInteractionRequestBody.of(params));\nInteraction interaction =\n response.interaction().orElseThrow(() -> new RuntimeException(\"No interaction returned\"));\nSystem.out.println(interaction.status().map(InteractionStatus::value).orElse(\"\"));\n" }, { "label": "antigravity", @@ -1652,7 +1652,7 @@ { "label": "deep_research", "lang": "go", - "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"deep-research-pro-preview-12-2025\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"find a cure to cancer\")),\n Background: genai.Ptr(true),\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil {\n fmt.Println(res.Interaction.Status)\n }\n}\n" + "source": "package main\n\nimport (\n \"context\"\n \"fmt\"\n \"log\"\n \"google.golang.org/genai\"\n gaos_interactions \"google.golang.org/genai/interactions/models/interactions\"\n \"google.golang.org/genai/interactions/models/operations\"\n)\n\nfunc main() {\n ctx := context.Background()\n client, err := genai.NewClient(ctx, nil)\n if err != nil {\n log.Fatal(err)\n }\n\n body := operations.NewCreateInteractionRequestBody(gaos_interactions.CreateAgentInteraction{\n Agent: gaos_interactions.AgentOption(\"deep-research-preview-04-2026\"),\n Input: genai.Ptr(gaos_interactions.NewInteractionsInput(\"find a cure to cancer\")),\n Background: genai.Ptr(true),\n })\n\n req := operations.CreateInteractionRequest{\n Body: body,\n }\n\n res, err := client.Interactions.Create(ctx, req)\n if err != nil {\n log.Fatal(err)\n }\n\n if res.Interaction != nil {\n fmt.Println(res.Interaction.Status)\n }\n}\n" }, { "label": "antigravity", @@ -3041,6 +3041,22 @@ "description": "An agent definition for the CreateAgent API.\nThis message is the target for annotation-parser-based JSON parsing.\nNew format:\n {\n \"id\": \"customer-sentinel\",\n \"base_agent\": \"\",\n \"system_instruction\": \"...\",\n \"base_environment\": { \"type\": \"remote\", \"sources\": [...] },\n \"tools\": [ {\"type\": \"code_execution\"} ]\n }", "type": "object", "properties": { + "id": { + "description": "The unique identifier for the agent.", + "type": "string" + }, + "base_agent": { + "description": "The base agent to extend.", + "type": "string" + }, + "system_instruction": { + "description": "System instruction for the agent.", + "type": "string" + }, + "description": { + "description": "Agent description for developers to quickly read and understand.", + "type": "string" + }, "agent_config": { "description": "Configuration parameters for the agent.", "oneOf": [ @@ -3060,22 +3076,6 @@ } ] }, - "id": { - "description": "The unique identifier for the agent.", - "type": "string" - }, - "base_agent": { - "description": "The base agent to extend.", - "type": "string" - }, - "system_instruction": { - "description": "System instruction for the agent.", - "type": "string" - }, - "description": { - "description": "Agent description for developers to quickly read and understand.", - "type": "string" - }, "tools": { "description": "The tools available to the agent.", "type": "array", @@ -3091,13 +3091,11 @@ "description": "The agent to interact with.", "type": "string", "enum": [ - "deep-research-pro-preview-12-2025", "deep-research-preview-04-2026", "deep-research-max-preview-04-2026", "antigravity-preview-05-2026" ], "x-speakeasy-enum-descriptions": [ - "Gemini Deep Research Agent", "Gemini Deep Research Agent", "Gemini Deep Research Max Agent", "Use the Antigravity managed agent to perform multi-step tasks that require reasoning, file operations, and tool use." @@ -3330,14 +3328,6 @@ "type": { "const": "audio" }, - "data": { - "type": "string", - "format": "byte", - "x-speakeasy-base64-input-mode": "file" - }, - "uri": { - "type": "string" - }, "mime_type": { "type": "string", "enum": [ @@ -3371,6 +3361,14 @@ "WEBM audio format" ] }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, "rate": { "description": "Deprecated. Use sample_rate instead. The value is ignored.", "type": "integer", @@ -3817,32 +3815,6 @@ } ] }, - "response_format": { - "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", - "oneOf": [ - { - "$ref": "#/components/schemas/ResponseFormat" - }, - { - "title": "ResponseFormatList", - "type": "array", - "items": { - "$ref": "#/components/schemas/ResponseFormat" - } - } - ] - }, - "environment": { - "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", - "oneOf": [ - { - "$ref": "#/components/schemas/EnvironmentConfig" - }, - { - "type": "string" - } - ] - }, "id": { "description": "Required. Output only. A unique identifier for the interaction completion.", "type": "string", @@ -3858,7 +3830,8 @@ "failed", "cancelled", "incomplete", - "budget_exceeded" + "budget_exceeded", + "requires_continuation" ], "readOnly": true, "x-google-enum-deprecated": [ @@ -3868,7 +3841,8 @@ false, false, false, - true + true, + false ], "x-google-enum-descriptions": [ "The interaction is in progress.", @@ -3877,7 +3851,8 @@ "The interaction failed.", "The interaction was cancelled.", "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", - "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11).", + "The interaction requires at least one more request for it to be\ncompleted. The next request should set the continuation_token to the one\nreturned by this request." ] }, "created": { @@ -3909,6 +3884,21 @@ "$ref": "#/components/schemas/ResponseModality" } }, + "response_format": { + "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", + "oneOf": [ + { + "$ref": "#/components/schemas/ResponseFormat" + }, + { + "title": "ResponseFormatList", + "type": "array", + "items": { + "$ref": "#/components/schemas/ResponseFormat" + } + } + ] + }, "response_mime_type": { "description": "The mime type of the response. This is required if response_format is set.", "type": "string", @@ -3918,6 +3908,17 @@ "description": "The ID of the previous interaction, if any.", "type": "string" }, + "environment": { + "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentConfig" + }, + { + "type": "string" + } + ] + }, "environment_id": { "description": "Output only. The environment ID for the interaction. Only populated if environment\nconfig is set in the request.", "type": "string", @@ -3945,6 +3946,12 @@ "type": "string" } }, + "continuation_token": { + "description": "Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, "stream": { "description": "Input only. Whether the interaction will be streamed.", "type": "boolean", @@ -4018,32 +4025,6 @@ "$ref": "#/components/schemas/GenerationConfig", "writeOnly": true }, - "response_format": { - "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", - "oneOf": [ - { - "$ref": "#/components/schemas/ResponseFormat" - }, - { - "title": "ResponseFormatList", - "type": "array", - "items": { - "$ref": "#/components/schemas/ResponseFormat" - } - } - ] - }, - "environment": { - "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", - "oneOf": [ - { - "$ref": "#/components/schemas/EnvironmentConfig" - }, - { - "type": "string" - } - ] - }, "id": { "description": "Required. Output only. A unique identifier for the interaction completion.", "type": "string", @@ -4059,7 +4040,8 @@ "failed", "cancelled", "incomplete", - "budget_exceeded" + "budget_exceeded", + "requires_continuation" ], "readOnly": true, "x-google-enum-deprecated": [ @@ -4069,7 +4051,8 @@ false, false, false, - true + true, + false ], "x-google-enum-descriptions": [ "The interaction is in progress.", @@ -4078,7 +4061,8 @@ "The interaction failed.", "The interaction was cancelled.", "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", - "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11).", + "The interaction requires at least one more request for it to be\ncompleted. The next request should set the continuation_token to the one\nreturned by this request." ] }, "created": { @@ -4110,6 +4094,21 @@ "$ref": "#/components/schemas/ResponseModality" } }, + "response_format": { + "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", + "oneOf": [ + { + "$ref": "#/components/schemas/ResponseFormat" + }, + { + "title": "ResponseFormatList", + "type": "array", + "items": { + "$ref": "#/components/schemas/ResponseFormat" + } + } + ] + }, "response_mime_type": { "description": "The mime type of the response. This is required if response_format is set.", "type": "string", @@ -4119,6 +4118,17 @@ "description": "The ID of the previous interaction, if any.", "type": "string" }, + "environment": { + "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentConfig" + }, + { + "type": "string" + } + ] + }, "environment_id": { "description": "Output only. The environment ID for the interaction. Only populated if environment\nconfig is set in the request.", "type": "string", @@ -4146,6 +4156,12 @@ "type": "string" } }, + "continuation_token": { + "description": "Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, "stream": { "description": "Input only. Whether the interaction will be streamed.", "type": "boolean", @@ -4348,11 +4364,31 @@ "type": "string", "enum": [ "application/pdf", - "text/csv" + "text/csv", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + "application/vnd.ms-powerpoint", + "application/vnd.oasis.opendocument.presentation", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "application/msword", + "application/vnd.openxmlformats-officedocument.wordprocessingml.template", + "application/vnd.oasis.opendocument.text", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "application/vnd.ms-excel", + "application/vnd.oasis.opendocument.spreadsheet" ], "x-google-enum-descriptions": [ "PDF document format", - "CSV document format" + "CSV document format", + "PPTX presentation format", + "PPT presentation format", + "ODP presentation format", + "DOCX document format", + "DOC document format", + "DOTX document template format", + "ODT document format", + "XLSX spreadsheet format", + "XLS spreadsheet format", + "ODS spreadsheet format" ] } }, @@ -4379,6 +4415,37 @@ "type": { "const": "document" }, + "mime_type": { + "type": "string", + "enum": [ + "application/pdf", + "text/csv", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + "application/vnd.ms-powerpoint", + "application/vnd.oasis.opendocument.presentation", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "application/msword", + "application/vnd.openxmlformats-officedocument.wordprocessingml.template", + "application/vnd.oasis.opendocument.text", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "application/vnd.ms-excel", + "application/vnd.oasis.opendocument.spreadsheet" + ], + "x-google-enum-descriptions": [ + "PDF document format", + "CSV document format", + "PPTX presentation format", + "PPT presentation format", + "ODP presentation format", + "DOCX document format", + "DOC document format", + "DOTX document template format", + "ODT document format", + "XLSX spreadsheet format", + "XLS spreadsheet format", + "ODS spreadsheet format" + ] + }, "data": { "type": "string", "format": "byte", @@ -4386,17 +4453,6 @@ }, "uri": { "type": "string" - }, - "mime_type": { - "type": "string", - "enum": [ - "application/pdf", - "text/csv" - ], - "x-google-enum-descriptions": [ - "PDF document format", - "CSV document format" - ] } }, "required": [ @@ -4474,23 +4530,6 @@ "type": "string", "readOnly": true }, - "network": { - "description": "Network configuration for the environment.", - "oneOf": [ - { - "$ref": "#/components/schemas/EnvironmentNetworkEgressAllowlist" - }, - { - "type": "string", - "enum": [ - "disabled" - ], - "x-google-enum-descriptions": [ - "All network egress is blocked." - ] - } - ] - }, "sources": { "description": "Sources to be mounted into the environment.", "type": "array", @@ -4526,6 +4565,23 @@ "" ] }, + "network": { + "description": "Network configuration for the environment.", + "oneOf": [ + { + "$ref": "#/components/schemas/EnvironmentNetworkEgressAllowlist" + }, + { + "type": "string", + "enum": [ + "disabled" + ], + "x-google-enum-descriptions": [ + "All network egress is blocked." + ] + } + ] + }, "file_count": { "description": "Output only. The number of files in the environment, output only.", "type": "string", @@ -4551,6 +4607,12 @@ "type": { "const": "remote" }, + "sources": { + "type": "array", + "items": { + "$ref": "#/components/schemas/Source" + } + }, "network": { "description": "Network configuration for the environment.", "oneOf": [ @@ -4568,12 +4630,6 @@ } ] }, - "sources": { - "type": "array", - "items": { - "$ref": "#/components/schemas/Source" - } - }, "environment_id": { "description": "Optional. The environment ID for the interaction. If specified, the request will\nupdate the existing environment instead of creating a new one.", "type": "string" @@ -5347,30 +5403,6 @@ "description": "Configuration parameters for model interactions.", "type": "object", "properties": { - "tool_choice": { - "description": "The tool choice configuration.", - "oneOf": [ - { - "$ref": "#/components/schemas/ToolChoiceConfig" - }, - { - "type": "string", - "enum": [ - "auto", - "any", - "none", - "validated" - ], - "x-google-enum-descriptions": [ - "Auto tool choice.", - "Any tool choice.", - "No tool choice.", - "Validated tool choice." - ], - "x-speakeasy-model-namespace": "interactions" - } - ] - }, "temperature": { "description": "Controls the randomness of the output.", "type": "number", @@ -5395,6 +5427,30 @@ "type": "string" } }, + "tool_choice": { + "description": "The tool choice configuration.", + "oneOf": [ + { + "$ref": "#/components/schemas/ToolChoiceConfig" + }, + { + "type": "string", + "enum": [ + "auto", + "any", + "none", + "validated" + ], + "x-google-enum-descriptions": [ + "Auto tool choice.", + "Any tool choice.", + "No tool choice.", + "Validated tool choice." + ], + "x-speakeasy-model-namespace": "interactions" + } + ] + }, "thinking_level": { "description": "The level of thought tokens that the model should generate.", "$ref": "#/components/schemas/ThinkingLevel" @@ -6211,14 +6267,6 @@ "type": { "const": "image" }, - "data": { - "type": "string", - "format": "byte", - "x-speakeasy-base64-input-mode": "file" - }, - "uri": { - "type": "string" - }, "mime_type": { "type": "string", "enum": [ @@ -6242,6 +6290,14 @@ "TIFF image format" ] }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, "resolution": { "description": "The resolution of the media.", "$ref": "#/components/schemas/MediaResolution" @@ -6385,7 +6441,8 @@ "failed", "cancelled", "incomplete", - "budget_exceeded" + "budget_exceeded", + "requires_continuation" ], "readOnly": true, "x-google-enum-deprecated": [ @@ -6395,7 +6452,8 @@ false, false, false, - true + true, + false ], "x-google-enum-descriptions": [ "The interaction is in progress.", @@ -6404,7 +6462,8 @@ "The interaction failed.", "The interaction was cancelled.", "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", - "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11).", + "The interaction requires at least one more request for it to be\ncompleted. The next request should set the continuation_token to the one\nreturned by this request." ] }, "created": { @@ -6434,9 +6493,33 @@ "$ref": "#/components/schemas/GenerationConfig", "writeOnly": true }, + "system_instruction": { + "description": "System instruction for the interaction.", + "type": "string" + }, + "tools": { + "description": "A list of tool declarations the model may call during interaction.", + "type": "array", + "items": { + "$ref": "#/components/schemas/Tool" + } + }, "input": { "$ref": "#/components/schemas/InteractionsInput" }, + "usage": { + "description": "Output only. Statistics on the interaction request's token usage.", + "$ref": "#/components/schemas/Usage", + "readOnly": true + }, + "response_modalities": { + "description": "The requested modalities of the response (TEXT, IMAGE, AUDIO).", + "type": "array", + "deprecated": true, + "items": { + "$ref": "#/components/schemas/ResponseModality" + } + }, "response_format": { "description": "Enforces that the generated response is a JSON object that complies with\nthe JSON schema specified in this field.", "oneOf": [ @@ -6452,6 +6535,15 @@ } ] }, + "response_mime_type": { + "description": "The mime type of the response. This is required if response_format is set.", + "type": "string", + "deprecated": true + }, + "previous_interaction_id": { + "description": "The ID of the previous interaction, if any.", + "type": "string" + }, "environment": { "description": "The environment configuration for the interaction. Can be an object\nspecifying remote environment sources or a string referencing an existing\nenvironment ID.", "oneOf": [ @@ -6463,39 +6555,6 @@ } ] }, - "system_instruction": { - "description": "System instruction for the interaction.", - "type": "string" - }, - "tools": { - "description": "A list of tool declarations the model may call during interaction.", - "type": "array", - "items": { - "$ref": "#/components/schemas/Tool" - } - }, - "usage": { - "description": "Output only. Statistics on the interaction request's token usage.", - "$ref": "#/components/schemas/Usage", - "readOnly": true - }, - "response_modalities": { - "description": "The requested modalities of the response (TEXT, IMAGE, AUDIO).", - "type": "array", - "deprecated": true, - "items": { - "$ref": "#/components/schemas/ResponseModality" - } - }, - "response_mime_type": { - "description": "The mime type of the response. This is required if response_format is set.", - "type": "string", - "deprecated": true - }, - "previous_interaction_id": { - "description": "The ID of the previous interaction, if any.", - "type": "string" - }, "environment_id": { "description": "Output only. The environment ID for the interaction. Only populated if environment\nconfig is set in the request.", "type": "string", @@ -6531,6 +6590,12 @@ }, "readOnly": true }, + "continuation_token": { + "description": "Opaque token to resume a long decode. Output: set when status is\nREQUIRES_CONTINUATION and decoding can be resumed. Input: pass the latest\ntoken back unchanged in CreateInteraction to continue decoding.", + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, "stream": { "description": "Input only. Whether the interaction will be streamed.", "type": "boolean", @@ -6813,16 +6878,16 @@ }, "readOnly": true }, - "agent": { - "description": "The agent to interact with.", - "type": "string" - }, "continuation_token": { "description": "Opaque token to resume a long decode.\nOutput: set when status is INCOMPLETE and decoding can be resumed.\nInput: pass the latest token back unchanged in CreateInteraction to\ncontinue decoding.", "type": "string", "format": "byte", "x-speakeasy-base64-input-mode": "file" }, + "agent": { + "description": "The agent to interact with.", + "type": "string" + }, "object": { "description": "Output only. The resource type.", "type": "string", @@ -6863,7 +6928,8 @@ "failed", "cancelled", "incomplete", - "budget_exceeded" + "budget_exceeded", + "requires_continuation" ], "x-google-enum-deprecated": [ false, @@ -6872,7 +6938,8 @@ false, false, false, - true + true, + false ], "x-google-enum-descriptions": [ "The interaction is in progress.", @@ -6881,7 +6948,8 @@ "The interaction failed.", "The interaction was cancelled.", "The interaction is completed, but contains incomplete results (e.g.\nhitting max_tokens).", - "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11)." + "Deprecated: Token and execution budget exhaustion returns INCOMPLETE\n(11).", + "The interaction requires at least one more request for it to be\ncompleted. The next request should set the continuation_token to the one\nreturned by this request." ] }, "event_type": { @@ -7076,14 +7144,14 @@ "type": { "const": "mcp_server" }, - "url": { - "description": "The full URL for the MCPServer endpoint.\nExample: \"https://api.example.com/mcp\"", - "type": "string" - }, "name": { "description": "The name of the MCPServer.", "type": "string" }, + "url": { + "description": "The full URL for the MCPServer endpoint.\nExample: \"https://api.example.com/mcp\"", + "type": "string" + }, "headers": { "description": "Optional: Fields for authentication headers, timeouts, etc., if needed.", "type": "object", @@ -7362,7 +7430,7 @@ }, "ModelOption": { "title": "Model", - "description": "The model that will complete your prompt.\\n\\nSee [models](https://ai.google.dev/gemini-api/docs/models) for additional details.", + "description": "The model that will complete your prompt.\\n\\nSee [models](https://aistudio.google.com/docs/models) for additional details.", "type": "string", "enum": [ "gemma-4-26b-a4b-it", @@ -8695,6 +8763,20 @@ "description": "Configuration for speech recognition (transcription).", "type": "object", "properties": { + "language_codes": { + "description": "Optional. BCP-47 language codes providing hints about the languages present in the\naudio. If omitted or empty, defaults to automatic language detection.", + "type": "array", + "items": { + "type": "string" + } + }, + "custom_vocabulary": { + "description": "Optional. A list of custom vocabulary phrases to bias the speech recognition model\ntoward recognizing specific terms.", + "type": "array", + "items": { + "type": "string" + } + }, "mode": { "description": "Discriminated transcription mode options or enum.", "oneOf": [ @@ -8714,20 +8796,6 @@ } ] }, - "language_codes": { - "description": "Optional. BCP-47 language codes providing hints about the languages present in the\naudio. If omitted or empty, defaults to automatic language detection.", - "type": "array", - "items": { - "type": "string" - } - }, - "custom_vocabulary": { - "description": "Optional. A list of custom vocabulary phrases to bias the speech recognition model\ntoward recognizing specific terms.", - "type": "array", - "items": { - "type": "string" - } - }, "timestamp_granularities": { "description": "Optional. The granularity of timestamps to include in the transcription output.\nSupported values: \"word\". If empty, no timestamps are generated.", "type": "array", @@ -9332,6 +9400,11 @@ "type": "integer", "format": "int32" }, + "initial_input_tokens": { + "description": "Number of input tokens used in the first turn of this interaction.", + "type": "integer", + "format": "int32" + }, "input_tokens_by_modality": { "description": "A breakdown of input token usage by modality.", "type": "array", @@ -9497,6 +9570,10 @@ "description": "The URI of the video.", "type": "string" }, + "resolution": { + "description": "The resolution of the media.", + "$ref": "#/components/schemas/MediaResolution" + }, "processing": { "description": "How the model processes this video for understanding.", "oneOf": [ @@ -9516,10 +9593,6 @@ } ] }, - "resolution": { - "description": "The resolution of the media.", - "$ref": "#/components/schemas/MediaResolution" - }, "name": { "description": "A user-defined name for this content block. Can be referenced by the model\nin the final response.", "type": "string" @@ -9573,14 +9646,6 @@ "type": { "const": "video" }, - "data": { - "type": "string", - "format": "byte", - "x-speakeasy-base64-input-mode": "file" - }, - "uri": { - "type": "string" - }, "mime_type": { "type": "string", "enum": [ @@ -9608,6 +9673,14 @@ "JPEG 2000 video format" ] }, + "data": { + "type": "string", + "format": "byte", + "x-speakeasy-base64-input-mode": "file" + }, + "uri": { + "type": "string" + }, "resolution": { "description": "The resolution of the media.", "$ref": "#/components/schemas/MediaResolution" diff --git a/tests/unit/interactions/test_main.py b/tests/unit/interactions/test_main.py index 5af5dd7d82a..cd32eb11485 100644 --- a/tests/unit/interactions/test_main.py +++ b/tests/unit/interactions/test_main.py @@ -559,28 +559,34 @@ class TestInteractionsAcreateOffline: @pytest.mark.usefixtures("fake_provider_credentials") @pytest.mark.asyncio - async def test_acreate_gemini_keeps_continuation_token_out_of_the_request(self, respx_mock: MockRouter) -> None: + async def test_acreate_gemini_sends_a_caller_continuation_token_to_resume_a_decode( + self, respx_mock: MockRouter + ) -> None: route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( return_value=httpx.Response( 200, json={ "id": "interaction-offline", "object": "interaction", - "model": "gemini-3.8-flash", + "model": "gemini-4-argon", "status": "completed", "steps": [{"type": "model_output", "content": [{"type": "text", "text": "Hi"}]}], - "usage": {"input_tokens": 6, "output_tokens": 3}, + "usage": {"total_input_tokens": 6, "total_output_tokens": 3}, }, ) ) response: Final = await interactions.acreate( - model="gemini/gemini-3.8-flash", - input="Prove it", + model="gemini/gemini-4-argon", + previous_interaction_id="interaction-hop-1", continuation_token="dG9rLTE=", api_key="gemini-offline", ) - # The Interactions API answered 400 "Unknown parameter 'continuation_token'." on 2026-10-13 - assert json.loads(route.calls.last.request.content) == {"model": "gemini-3.8-flash", "input": "Prove it"} + # Google's spec (2026-11-05) takes continuation_token on CreateModelInteractionParams; a 2026-10-13 400 is gone + assert json.loads(route.calls.last.request.content) == { + "model": "gemini-4-argon", + "previous_interaction_id": "interaction-hop-1", + "continuation_token": "dG9rLTE=", + } assert response.status == "completed" @pytest.mark.usefixtures("fake_provider_credentials") @@ -746,9 +752,7 @@ def test_create_with_system_instruction_for_google_and_responses( def test_create_streaming_returns_interaction_events( respx_mock: MockRouter, ) -> None: - route: Final = respx_mock.post( - "https://generativelanguage.googleapis.com/v1beta/interactions?alt=sse" - ).mock( + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions?alt=sse").mock( return_value=httpx.Response( 200, headers={"content-type": "text/event-stream"}, @@ -856,3 +860,65 @@ def test_create_streaming_hands_completed_interaction_to_success_callback( ) assert recorder.logged == (_GEMINI_COMPLETED_INTERACTION,) + + +_GEMINI_URL: Final = "https://generativelanguage.googleapis.com/v1beta/interactions" +_LONG_DECODE_REQUEST: Final = { + "model": "gemini/gemini-4-argon", + "input": "Prove the prime number theorem.", + "generation_config": {"max_output_tokens": 120000, "thinking_level": "high"}, +} +_FIRST_HOP: Final = { + "id": "hop-1", + "object": "interaction", + "model": "gemini-4-argon", + "status": "requires_continuation", + "steps": [{"type": "model_output", "content": [{"type": "text", "text": "Lemma 1."}]}], + "usage": { + "total_input_tokens": 79, + "total_output_tokens": 32668, + "total_thought_tokens": 100, + "total_tokens": 32847, + "initial_input_tokens": 79, + }, + "continuation_token": "tok-1", +} + + +@pytest.mark.usefixtures("fake_provider_credentials") +def test_create_keeps_requires_continuation_status_and_token_for_the_caller(respx_mock: MockRouter) -> None: + route: Final = respx_mock.post(_GEMINI_URL).mock(return_value=httpx.Response(200, json=_FIRST_HOP)) + + response: Final = interactions.create(**_LONG_DECODE_REQUEST, api_key="gemini-offline") + + assert route.call_count == 1 + assert response.status == "requires_continuation" + assert response.continuation_token == "tok-1" + assert response.usage == _FIRST_HOP["usage"] + + +@pytest.mark.usefixtures("fake_provider_credentials") +def test_create_sends_a_caller_continuation_token_to_resume_a_long_decode(respx_mock: MockRouter) -> None: + route: Final = respx_mock.post(_GEMINI_URL).mock( + return_value=httpx.Response( + 200, json={**_FIRST_HOP, "id": "hop-2", "status": "completed", "continuation_token": None} + ) + ) + + response: Final = interactions.create( + model="gemini/gemini-4-argon", + previous_interaction_id="hop-1", + continuation_token="tok-1", + generation_config=_LONG_DECODE_REQUEST["generation_config"], + api_key="gemini-offline", + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert sent == { + "model": "gemini-4-argon", + "previous_interaction_id": "hop-1", + "continuation_token": "tok-1", + "generation_config": _LONG_DECODE_REQUEST["generation_config"], + }, sent + assert response.status == "completed" + assert response.continuation_token is None diff --git a/tests/unit/interactions/test_openapi_compliance.py b/tests/unit/interactions/test_openapi_compliance.py index b6cc92e0adc..71db2cd24e3 100644 --- a/tests/unit/interactions/test_openapi_compliance.py +++ b/tests/unit/interactions/test_openapi_compliance.py @@ -2,7 +2,7 @@ OpenAPI compliance tests for Google Interactions API. Validates that our SDK requests/responses match the OpenAPI spec snapshot in -interactions.openapi.json, taken from https://ai.google.dev/static/api/interactions.openapi.json +interactions.openapi.json, taken from https://aistudio.google.com/api-reference/interactions.openapi.json (litellm/types/interactions/README.md says how to refresh it). The scheduled interactions-spec-canary workflow runs the same tests against the live spec by passing --interactions-spec-url, so a change Google makes shows up there, not in every PR. @@ -23,7 +23,7 @@ import pytest from openapi_core import OpenAPI from pydantic import BaseModel, TypeAdapter -OPENAPI_SPEC_URL = "https://ai.google.dev/static/api/interactions.openapi.json" +OPENAPI_SPEC_URL = "https://aistudio.google.com/api-reference/interactions.openapi.json" OPENAPI_SPEC_SNAPSHOT: Final = Path(__file__).with_name("interactions.openapi.json") SPEC_URL_OPTION: Final = TypeAdapter(str | None) @@ -47,10 +47,10 @@ def _requested_spec_url(request: pytest.FixtureRequest) -> str | None: def _load_openapi_spec(live_url: str | None) -> LoadedSpec: if live_url is None: return LoadedSpec(source=str(OPENAPI_SPEC_SNAPSHOT), spec=json.loads(OPENAPI_SPEC_SNAPSHOT.read_text())) - with httpx.Client(transport=httpx.HTTPTransport(retries=3), timeout=30.0) as client: + with httpx.Client(transport=httpx.HTTPTransport(retries=3), timeout=30.0, follow_redirects=True) as client: response: Final = client.get(live_url) response.raise_for_status() - return LoadedSpec(source=str(response.url), spec=response.json()) + return LoadedSpec(source=live_url, spec=response.json()) def _model_create_request_schema(spec_dict: Dict[str, Any]) -> Dict[str, Any]: diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_get_api_base.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_api_base.py index 90fd5ba6c88..07e153f537b 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_get_api_base.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_get_api_base.py @@ -1,6 +1,5 @@ -import json - import pytest +import respx import litellm from litellm.litellm_core_utils.llm_response_utils import get_api_base as get_api_base_module @@ -53,17 +52,20 @@ class TestDeclaredAuthenticatingProvider: assert resolution_lookups == [] assert api_base == expected - def test_copilot_keeps_the_enterprise_endpoint_from_disk(self, isolated_token_dirs, resolution_lookups): + def test_copilot_uses_the_user_endpoint_for_an_enterprise_api_base(self, isolated_token_dirs, resolution_lookups): token_dir = isolated_token_dirs / "github_copilot" token_dir.mkdir() - (token_dir / "api-key.json").write_text( - json.dumps({"endpoints": {"api": "https://api.enterprise.githubcopilot.com"}}) - ) - - api_base = litellm.get_api_base(model="github_copilot/gpt-4o", optional_params={}) + token = f"gho_{isolated_token_dirs.name}" + (token_dir / "access-token").write_text(token) + with respx.mock(assert_all_called=True) as router: + route = router.get("https://api.github.com/copilot_internal/user").respond( + json={"endpoints": {"api": "https://api.enterprise.githubcopilot.com"}} + ) + api_base = litellm.get_api_base(model="github_copilot/gpt-4o", optional_params={}) assert resolution_lookups == [] assert api_base == "https://api.enterprise.githubcopilot.com" + assert route.calls.last.request.headers["authorization"] == f"token {token}" def test_explicit_api_base_still_wins(self, isolated_token_dirs, resolution_lookups): api_base = litellm.get_api_base( diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index 25de5ef60d9..5779250950d 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -796,7 +796,6 @@ def test_metadata_callers_never_log_in_for_a_declared_chatgpt_deployment( from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.proxy.spend_tracking.savings import _resolve_model from litellm.router import Router - from litellm.types.router import LiteLLM_Params from litellm.utils import supports_native_streaming monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") @@ -805,4 +804,7 @@ def test_metadata_callers_never_log_in_for_a_declared_chatgpt_deployment( assert get_api_base("chatgpt/gpt-5.4", {"model": "chatgpt/gpt-5.4"}) == "https://chatgpt.com/backend-api/codex" identity = _resolve_model("chatgpt/gpt-5.4", None) assert identity is not None and identity.provider == "chatgpt" - assert Router._model_group_llm_provider(LiteLLM_Params(model="chatgpt/gpt-5.4")) == ("gpt-5.4", "chatgpt") + group = Router(model_list=[{"model_name": "codex", "litellm_params": {"model": "chatgpt/gpt-5.4"}}]).registry().group( + "codex" + ) + assert group is not None and group.providers == ["chatgpt"] diff --git a/tests/unit/litellm_core_utils/test_safe_json_dumps.py b/tests/unit/litellm_core_utils/test_safe_json_dumps.py index 1b8bdd21326..740ce4f433a 100644 --- a/tests/unit/litellm_core_utils/test_safe_json_dumps.py +++ b/tests/unit/litellm_core_utils/test_safe_json_dumps.py @@ -230,9 +230,7 @@ def test_pydantic_base_model(): inner: InnerModel tags: list - outer = OuterModel( - name="test", inner=InnerModel(value=42, label="hello"), tags=["a", "b"] - ) + outer = OuterModel(name="test", inner=InnerModel(value=42, label="hello"), tags=["a", "b"]) # Test a pydantic model at the top level result = json.loads(safe_dumps(outer)) @@ -261,3 +259,9 @@ def test_safe_json_structure_keeps_tuples_and_drops_non_string_keys(): assert structure == {"models": ("A", "B"), "tags": ["X", "Y"], "nested": {"deep": ("C",)}} assert type(structure["models"]) is tuple assert json.loads(safe_dumps(data)) == {"models": ["a", "b"], "tags": ["x", "y"], "nested": {"deep": ["c"]}} + + +def test_integer_limit_clamps_both_signs_and_leaves_bools_and_floats_alone() -> None: + assert safe_json_structure( + {"big": 10**400, "small": -(10**400), "inside": 7, "flag": True, "ratio": 1e300}, integer_limit=100 + ) == {"big": 100, "small": -100, "inside": 7, "flag": True, "ratio": 1e300} diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 80845b1f12d..3a4231443dd 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -11047,3 +11047,24 @@ def test_claude_cache_control_cap_does_not_count_a_message_level_hint_that_list_ assert sorted(_cached_texts(body)) == ["system prompt", "turn 1", "turn 2", "turn 3"] assert json.dumps(body).count('"cache_control"') == 4 + + +@pytest.mark.parametrize("model", ["claude-opus-4-7", "claude-haiku-4-5-20251001"]) +def test_get_supported_params_reports_sampling_params_only_when_the_request_keeps_them( + local_model_cost_map, model: str +): + """What get_supported_openai_params says about temperature/top_p must match what + map_openai_params does with them under drop_params, so a caller that pre-flights + from the list is never told a param is fine that the request path then drops.""" + config = AnthropicConfig() + + supported = config.get_supported_openai_params(model=model) + sent = config.map_openai_params( + non_default_params={"temperature": 0.1, "top_p": 0.5}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert ("temperature" in supported) == ("temperature" in sent), (model, supported, sent) + assert ("top_p" in supported) == ("top_p" in sent), (model, supported, sent) diff --git a/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py b/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py index 810dfffaf8b..4da4c16dff9 100644 --- a/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py +++ b/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py @@ -1053,32 +1053,30 @@ async def test_executor_failure_is_not_tagged(): def _router_with_advisor_deployment( - recorder, advisor_model="claude-opus-4-8", deployment_model=None, model_group_alias=None + advisor_model="claude-opus-4-8", deployment_model=None, model_group_alias=None ): """Build a Router whose only deployment is the advisor model on Foundry. - The recorder replaces ``litellm.anthropic_messages`` before construction - because Router binds it at init time, so the returned Router exercises the - real deployment-resolution path and records what it dispatched. + The Router resolves ``litellm.anthropic_messages`` per call, so each test + patches it with a recorder around the handler call; the Router then + exercises the real deployment-resolution path and records what it dispatched. """ - import litellm from litellm.router import Router - with patch("litellm.anthropic_messages", new=recorder): - return Router( - model_list=[ - { - "model_name": advisor_model, - "litellm_params": { - "model": deployment_model or f"azure_ai/{advisor_model}", - "api_base": "http://127.0.0.1:1/foundry", - "api_key": "fake-foundry-key", - }, - } - ], - model_group_alias=model_group_alias, - num_retries=0, - ) + return Router( + model_list=[ + { + "model_name": advisor_model, + "litellm_params": { + "model": deployment_model or f"azure_ai/{advisor_model}", + "api_base": "http://127.0.0.1:1/foundry", + "api_key": "fake-foundry-key", + }, + } + ], + model_group_alias=model_group_alias, + num_retries=0, + ) @pytest.mark.asyncio @@ -1094,7 +1092,7 @@ async def test_advisor_sub_call_routes_through_proxy_router(): router_calls.append(kwargs) return _make_text_response("Use trial division.", model="claude-opus-4-8") - router = _router_with_advisor_deployment(recorder) + router = _router_with_advisor_deployment() call_count = 0 @@ -1111,6 +1109,7 @@ async def test_advisor_sub_call_routes_through_proxy_router(): side_effect=mock_call, ), patch.object(proxy_server, "llm_router", router), + patch("litellm.anthropic_messages", new=recorder), ): h = AdvisorOrchestrationHandler() result = await h.handle( @@ -1155,7 +1154,7 @@ async def test_advisor_sub_call_routes_through_router_for_alias_and_wildcard(rou router_calls.append(kwargs) return _make_text_response("Use trial division.", model="claude-opus-4-8") - router = _router_with_advisor_deployment(recorder, **router_kwargs) + router = _router_with_advisor_deployment(**router_kwargs) call_count = 0 @@ -1172,6 +1171,7 @@ async def test_advisor_sub_call_routes_through_router_for_alias_and_wildcard(rou side_effect=mock_call, ), patch.object(proxy_server, "llm_router", router), + patch("litellm.anthropic_messages", new=recorder), ): h = AdvisorOrchestrationHandler() await h.handle( @@ -1204,7 +1204,7 @@ async def test_advisor_sub_call_bypasses_router_for_unconfigured_model(): router_calls.append(kwargs) return _make_text_response("should not be used") - router = _router_with_advisor_deployment(recorder, advisor_model="some-other-model") + router = _router_with_advisor_deployment(advisor_model="some-other-model") call_count = 0 @@ -1223,6 +1223,7 @@ async def test_advisor_sub_call_bypasses_router_for_unconfigured_model(): side_effect=mock_call, ), patch.object(proxy_server, "llm_router", router), + patch("litellm.anthropic_messages", new=recorder), ): h = AdvisorOrchestrationHandler() await h.handle( @@ -1253,7 +1254,7 @@ async def test_advisor_sub_call_client_override_bypasses_router(): router_calls.append(kwargs) return _make_text_response("should not be used") - router = _router_with_advisor_deployment(recorder) + router = _router_with_advisor_deployment() sub_calls = [] @@ -1278,6 +1279,7 @@ async def test_advisor_sub_call_client_override_bypasses_router(): side_effect=mock_call, ), patch.object(proxy_server, "llm_router", router), + patch("litellm.anthropic_messages", new=recorder), patch.dict(proxy_server.general_settings, {"allow_client_side_credentials": True}), patch.object(litellm, "user_url_validation", False), ): diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index d66ba3f324e..65f8400e53c 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -477,6 +477,7 @@ def test_default_max_retries_env_var_reaches_azure_sdk_client(): "aget_responses", "adelete_responses", "acancel_responses", + "acompact_responses", "aresponses_websocket", "alist_input_items", "acreate_fine_tuning_job", diff --git a/tests/unit/llms/base_llm/responses/test_transformation.py b/tests/unit/llms/base_llm/responses/test_transformation.py index 50d918fa5be..68d9e79dd81 100644 --- a/tests/unit/llms/base_llm/responses/test_transformation.py +++ b/tests/unit/llms/base_llm/responses/test_transformation.py @@ -27,8 +27,8 @@ from litellm.types.llms.openai import( ResponsesAPIStreamEvents, ) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome -from typing import Any, AsyncIterator, Final, List -from unittest.mock import AsyncMock, MagicMock, Mock, patch +from typing import Any, Final, List +from unittest.mock import Mock, patch @pytest.mark.asyncio @@ -751,586 +751,6 @@ class TestBaseResponsesAPIStreamingIterator: mock_logging_obj.async_failure_handler.assert_not_called() mock_logging_obj.failure_handler.assert_not_called() -@pytest.fixture() -def _vcr_outcome_gate_router_unit(request, vcr): - install_live_call_probe(request, vcr) - yield - record_vcr_outcome(request, vcr) - -@pytest.fixture(scope="function") -def setup_and_teardown_router_unit(): - """ - This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. - """ - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - - asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) - importlib.reload(litellm) - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - importlib.reload(litellm.proxy.proxy_server) - except Exception: - pass - loop = asyncio.get_event_loop_policy().new_event_loop() - asyncio.set_event_loop(loop) - yield - loop.close() - asyncio.set_event_loop(None) - -def _make_router() -> Router: - return Router( - model_list=[ - { - "model_name": "primary", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-test", - }, - }, - { - "model_name": "fallback", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "sk-test", - }, - }, - ] - ) - -def _make_completed_event(input_tokens: int, output_tokens: int, total_tokens: int) -> ResponseCompletedEvent: - response = ResponsesAPIResponse.model_construct( - usage=ResponseAPIUsage( - input_tokens=input_tokens, - output_tokens=output_tokens, - total_tokens=total_tokens, - ) - ) - return ResponseCompletedEvent.model_construct( - type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - response=response, - ) - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_extract_partial_responses_usage_native_completed(): - """Native path: completed_response carries usage → returned as-is.""" - completed = _make_completed_event(11, 7, 18) - source = MagicMock() - source.completed_response = completed - - usage = Router._extract_partial_responses_usage(source) - assert usage is not None - assert usage.input_tokens == 11 - assert usage.output_tokens == 7 - assert usage.total_tokens == 18 - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_extract_partial_responses_usage_no_completed_response(): - """Native path: no completed_response → returns None.""" - source = MagicMock() - source.completed_response = None - - usage = Router._extract_partial_responses_usage(source) - assert usage is None - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_extract_partial_responses_usage_bridge_iterator_no_completed_response(): - """ - Regression for #35411: the bridge iterator - (LiteLLMCompletionStreamingIterator) overrides __init__ without calling - super().__init__(), so completed_response was never set until the stream - reached RESPONSE_COMPLETED. On a mid-stream provider error (before - completion) the fallback recovery path read source_iterator.completed_response - and raised AttributeError, masking the real provider error and bypassing - fallbacks. The attribute must always exist and default to None. - """ - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - - wrapper = MagicMock() - wrapper.logging_obj = MagicMock() - iterator = LiteLLMCompletionStreamingIterator( - model="anthropic/claude-sonnet-4-5", - litellm_custom_stream_wrapper=wrapper, - request_input="hi", - responses_api_request={}, - ) - - assert iterator.completed_response is None - # No chat chunks collected yet and no completed_response → must return - # None instead of raising AttributeError. - assert Router._extract_partial_responses_usage(iterator) is None - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_combine_responses_fallback_usage_sums_completed_event(): - """Partial-stream usage is summed into the fallback event's usage.""" - fallback_event = _make_completed_event(5, 3, 8) - partial = ResponseAPIUsage(input_tokens=11, output_tokens=7, total_tokens=18) - - Router._combine_responses_fallback_usage(fallback_event, partial) - - combined = fallback_event.response.usage - assert combined is not None - assert combined.input_tokens == 16 - assert combined.output_tokens == 10 - assert combined.total_tokens == 26 - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_combine_responses_fallback_usage_passthrough_for_unknown_event(): - """Events that are not completed/failed/incomplete are not mutated.""" - other = MagicMock() # not a ResponseCompletedEvent etc. → isinstance false - partial = ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2) - Router._combine_responses_fallback_usage(other, partial) - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_passthrough(): - """ - Without MidStreamFallbackError, the wrapper yields source events - unchanged and returns a BaseResponsesAPIStreamingIterator subclass. - """ - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - - events = [_make_completed_event(1, 1, 2)] - - class _FakeSource: - """Minimal source iterator. Provides every attribute the wrapper - constructor reads from source_iterator.""" - - def __init__(self) -> None: - self._i = 0 - self.completed_response = None - self.response = MagicMock() - self.model = "openai/gpt-4o-mini" - self.logging_obj = MagicMock() - self.responses_api_provider_config = MagicMock() - self.start_time = 0.0 - self.litellm_metadata = {} - self.custom_llm_provider = "openai" - self.request_data = {} - self.call_type = "aresponses" - self._hidden_params: dict = {} - - def __aiter__(self) -> AsyncIterator[Any]: - return self - - async def __anext__(self): - if self._i >= len(events): - raise StopAsyncIteration - ev = events[self._i] - self._i += 1 - return ev - - async def aclose(self): - return None - - router = _make_router() - source = _FakeSource() - - wrapper = await router._aresponses_streaming_iterator(source, initial_kwargs={"model": "primary"}) - assert isinstance(wrapper, BaseResponsesAPIStreamingIterator) - - collected = [ev async for ev in wrapper] - assert len(collected) == 1 - assert collected[0].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_aresponses_with_streaming_fallbacks_non_streaming_passthrough(): - """Non-streaming response is returned unchanged, no wrap.""" - router = _make_router() - plain_response = MagicMock() - - async def fake_original(**_kwargs): - return plain_response - - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_helper", - new=AsyncMock(return_value=plain_response), - ): - out = await router._aresponses_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=False, - ) - assert out is plain_response - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_aresponses_with_streaming_fallbacks_wraps_streaming_iterator(): - """Streaming response is wrapped via _aresponses_streaming_iterator.""" - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - - router = _make_router() - streaming_iter = MagicMock(spec=BaseResponsesAPIStreamingIterator) - wrapped = MagicMock(spec=BaseResponsesAPIStreamingIterator) - - async def fake_original(**_kwargs): - return streaming_iter - - with ( - patch.object( - router, - "_ageneric_api_call_with_fallbacks_helper", - new=AsyncMock(return_value=streaming_iter), - ), - patch.object( - router, - "_aresponses_streaming_iterator", - new=AsyncMock(return_value=wrapped), - ) as mock_wrap, - ): - out = await router._aresponses_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - ) - assert out is wrapped - mock_wrap.assert_awaited_once() - -def _make_three_tier_router(**router_kwargs) -> Router: - return Router( - model_list=[ - {"model_name": "primary", "litellm_params": {"model": "openai/primary-model", "api_key": "sk-test"}}, - {"model_name": "fb1", "litellm_params": {"model": "openai/fb1-model", "api_key": "sk-test"}}, - {"model_name": "fb2", "litellm_params": {"model": "openai/fb2-model", "api_key": "sk-test"}}, - ], - num_retries=0, - **router_kwargs, - ) - -def _mid_stream_failure(model: str): - import litellm - from litellm.exceptions import MidStreamFallbackError - - return MidStreamFallbackError( - message="stream dropped", - model=model, - llm_provider="openai", - original_exception=litellm.InternalServerError(message="stream dropped", llm_provider="openai", model=model), - is_pre_first_chunk=True, - ) - -def _scripted_responses_stream(events: list, error: Exception | None = None): - from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator - - class _ScriptedStream(BaseResponsesAPIStreamingIterator): - def __init__(self) -> None: - self._events = list(events) - self._hidden_params: dict = {} - self.completed_response = None - - def __aiter__(self): - return self - - async def __anext__(self): - if self._events: - return self._events.pop(0) - if error is not None: - raise error - raise StopAsyncIteration - - async def aclose(self) -> None: - return None - - return _ScriptedStream() - -def _three_tier_original(calls: list, primary_fails_pre_stream: bool): - import litellm - - completed_event = _make_completed_event(1, 1, 2) - - async def fake_original(**kwargs): - model = kwargs["model"] - calls.append(model) - if model == "openai/primary-model": - if primary_fails_pre_stream: - raise litellm.InternalServerError(message="primary down", llm_provider="openai", model=model) - return _scripted_responses_stream([], _mid_stream_failure(model)) - if model == "openai/fb1-model": - return _scripted_responses_stream([], _mid_stream_failure(model)) - return _scripted_responses_stream([completed_event]) - - return fake_original, completed_event - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_aresponses_pre_stream_primary_failure_then_hop_stream_failure_reaches_second_entry(): - """Regression: fallbacks=[{"primary": ["fb1", "fb2"]}]. The primary fails before streaming, - fb1 is reached through the regular fallback chain and then fails mid-stream. Only the - primary's stream used to be wrapped, so fb1's mid-stream failure either re-raised or - re-tried fb1 itself; fb2 was unreachable.""" - router = _make_three_tier_router(fallbacks=[{"primary": ["fb1", "fb2"]}]) - calls: list = [] - fake_original, completed_event = _three_tier_original(calls, primary_fails_pre_stream=True) - - stream = await router._aresponses_with_streaming_fallbacks( - original_function=fake_original, model="primary", stream=True, input="hi" - ) - collected = [event async for event in stream] - - assert calls == ["openai/primary-model", "openai/fb1-model", "openai/fb2-model"] - assert collected == [completed_event] - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_aresponses_two_consecutive_mid_stream_failures_reach_second_entry(): - """Regression: the primary and fb1 both fail mid-stream; fb2 must still be tried.""" - router = _make_three_tier_router(fallbacks=[{"primary": ["fb1", "fb2"]}]) - calls: list = [] - fake_original, completed_event = _three_tier_original(calls, primary_fails_pre_stream=False) - - stream = await router._aresponses_with_streaming_fallbacks( - original_function=fake_original, model="primary", stream=True, input="hi" - ) - collected = [event async for event in stream] - - assert calls == ["openai/primary-model", "openai/fb1-model", "openai/fb2-model"] - assert collected == [completed_event] - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_aresponses_per_request_fallbacks_survive_into_hop_streams(): - """Regression: a request-level fallbacks list (key or team router_settings) is popped - before each attempt runs, so a hop's mid-stream re-entry used to see only the router's - own (empty) list and gave up after fb1.""" - router = _make_three_tier_router() - calls: list = [] - fake_original, completed_event = _three_tier_original(calls, primary_fails_pre_stream=False) - - stream = await router._aresponses_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - input="hi", - fallbacks=[{"primary": ["fb1", "fb2"]}], - ) - collected = [event async for event in stream] - - assert calls == ["openai/primary-model", "openai/fb1-model", "openai/fb2-model"] - assert collected == [completed_event] - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_aresponses_attempt_strips_the_controls_carrier_and_wraps_every_hop_stream(): - """Each attempt of the chain, not only the primary's, comes back wrapped for mid-stream - failover, and the per-request controls carrier rides into the wrapper's re-entry kwargs - without ever reaching the provider call.""" - from types import MappingProxyType - - from litellm.router_utils.fallback_event_handlers import ( - MID_STREAM_FALLBACK_CONTROLS_KEY, - MidStreamFallbackControls, - ) - - router = _make_three_tier_router() - completed_event = _make_completed_event(1, 1, 2) - hop_stream = _scripted_responses_stream([completed_event]) - seen: dict = {} - - async def fake_original(**kwargs): - seen.update(kwargs) - return hop_stream - - controls = MidStreamFallbackControls(MappingProxyType({"fallbacks": [{"primary": ["fb1", "fb2"]}]})) - stream = await router._ageneric_api_call_with_fallbacks_responses_attempt( - model="fb1", - original_generic_function=fake_original, - stream=True, - input="hi", - **{MID_STREAM_FALLBACK_CONTROLS_KEY: controls}, - ) - collected = [event async for event in stream] - - assert seen["model"] == "openai/fb1-model" - assert MID_STREAM_FALLBACK_CONTROLS_KEY not in seen - assert "fallbacks" not in seen - assert stream is not hop_stream - assert collected == [completed_event] - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_aresponses_fallback_on_in_stream_error_event(): - """A retriable in-stream error event (429) must trigger the router's mid-stream - fallback path: the wrapper catches MidStreamFallbackError raised by the source - iterator and yields the fallback stream instead of surfacing the error.""" - import json - from unittest.mock import Mock - - import litellm - from litellm.exceptions import MidStreamFallbackError - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - from litellm.types.llms.openai import ErrorEvent, ErrorEventError - - router = _make_router() - - error_payload = { - "type": "error", - "error": {"type": "tokens", "code": "rate_limit_exceeded", "message": "rate limited"}, - } - sse_bytes = f"data: {json.dumps(error_payload)}\n\n".encode() - - async def mock_aiter_bytes(): - yield sse_bytes - - mock_response = Mock() - mock_response.headers = {} - mock_response.aiter_bytes = mock_aiter_bytes - mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) - mock_logging_obj.model_call_details = {"litellm_params": {}} - mock_logging_obj.completion_start_time = None - mock_config = Mock(spec=BaseResponsesAPIConfig) - mock_config.transform_streaming_response.return_value = ErrorEvent( - type=ResponsesAPIStreamEvents.ERROR, - sequence_number=0, - error=ErrorEventError(type="tokens", code="rate_limit_exceeded", message="rate limited"), - ) - - source = ResponsesAPIStreamingIterator( - response=mock_response, - model="gpt-5", - responses_api_provider_config=mock_config, - logging_obj=mock_logging_obj, - custom_llm_provider="openai", - ) - - fallback_event = _make_completed_event(1, 1, 2) - - class _FallbackStream: - def __init__(self) -> None: - self._done = False - - def __aiter__(self): - return self - - async def __anext__(self): - if self._done: - raise StopAsyncIteration - self._done = True - return fallback_event - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FallbackStream()), - ) as mock_fallback: - wrapped = await router._aresponses_streaming_iterator( - response=source, - initial_kwargs={"model": "primary", "input": "original question"}, - ) - collected = [ev async for ev in wrapped] - - assert collected == [fallback_event] - mock_fallback.assert_awaited_once() - raised = mock_fallback.await_args.kwargs["e"] - assert isinstance(raised, MidStreamFallbackError) - assert raised.status_code == 429 - assert isinstance(raised.original_exception, litellm.RateLimitError) - assert raised.original_exception.status_code == 429 - assert mock_fallback.await_args.kwargs["kwargs"]["input"] == "original question" - -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_aresponses_fallback_uses_continuation_input_after_partial_content(): - """When output text was already streamed before the error, the fallback re-entry - must carry a continuation input with the partial assistant text instead of - retrying the original input from scratch (which would duplicate streamed content).""" - import json - from unittest.mock import Mock - - from litellm.exceptions import MidStreamFallbackError - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - from litellm.types.llms.openai import ErrorEvent, ErrorEventError - - router = _make_router() - - events = [ - {"type": "response.output_text.delta", "delta": "partial answer"}, - {"type": "error", "error": {"type": "server_error", "code": "internal_error", "message": "boom"}}, - ] - sse_payload = b"".join(f"data: {json.dumps(event)}\n\n".encode() for event in events) - - async def mock_aiter_bytes(): - yield sse_payload - - mock_response = Mock() - mock_response.headers = {} - mock_response.aiter_bytes = mock_aiter_bytes - mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) - mock_logging_obj.model_call_details = {"litellm_params": {}} - mock_logging_obj.completion_start_time = None - mock_config = Mock(spec=BaseResponsesAPIConfig) - - def transform(model, parsed_chunk, logging_obj): - if parsed_chunk.get("type") == "error": - return ErrorEvent( - type=ResponsesAPIStreamEvents.ERROR, - sequence_number=0, - error=ErrorEventError(**parsed_chunk["error"]), - ) - delta_event = Mock() - delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA - delta_event.delta = parsed_chunk["delta"] - return delta_event - - mock_config.transform_streaming_response.side_effect = transform - - source = ResponsesAPIStreamingIterator( - response=mock_response, - model="gpt-5", - responses_api_provider_config=mock_config, - logging_obj=mock_logging_obj, - custom_llm_provider="openai", - ) - - fallback_event = _make_completed_event(1, 1, 2) - - class _FallbackStream: - def __init__(self) -> None: - self._done = False - - def __aiter__(self): - return self - - async def __anext__(self): - if self._done: - raise StopAsyncIteration - self._done = True - return fallback_event - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FallbackStream()), - ) as mock_fallback: - wrapped = await router._aresponses_streaming_iterator( - response=source, - initial_kwargs={"model": "primary", "input": "original question"}, - ) - collected = [ev async for ev in wrapped] - - assert collected[-1] == fallback_event - raised = mock_fallback.await_args.kwargs["e"] - assert isinstance(raised, MidStreamFallbackError) - assert raised.is_pre_first_chunk is False - assert raised.generated_content == "partial answer" - continuation = mock_fallback.await_args.kwargs["kwargs"]["input"] - assert isinstance(continuation, list) - assert continuation[0]["content"][0]["text"] == "original question" - assert continuation[-2]["role"] == "assistant" - assert continuation[-2]["content"][0]["text"] == "partial answer" - assert continuation[-1]["role"] == "user" - - def _sse(*events: Mapping[str, object]) -> httpx.Response: return httpx.Response( 200, @@ -1894,42 +1314,109 @@ async def test_bridged_responses_fallback_continues_after_a_server_run_tool_sear router.discard() -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") +def _openai_completed(model: str) -> Mapping[str, object]: + return { + "type": "response.completed", "sequence_number": 1, + "response": { + "id": f"resp_{model}", "object": "response", "created_at": 0, "status": "completed", "model": model, + "output": [], "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + } + + @pytest.mark.asyncio -async def test_aresponses_client_error_event_skips_fallback(): - """A 400-mapped in-stream error (raised as APIError, not MidStreamFallbackError) - must surface to the caller without invoking the router's fallback path.""" - import litellm +@pytest.mark.parametrize( + ("primary_fails_pre_stream", "request_level_fallbacks"), + [(True, False), (False, False), (False, True)], + ids=["pre_stream_then_hop_stream_failure", "two_mid_stream_failures", "request_level_fallbacks"], +) +async def test_responses_fallback_chain_reaches_the_second_entry_after_a_hop_stream_fails( + primary_fails_pre_stream: bool, + request_level_fallbacks: bool, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression: fallbacks=[{"primary": ["fb1", "fb2"]}]. fb1 is reached (after the primary failed before or + during streaming) and then fails mid-stream; fb2 must still be tried, including when the fallbacks list + came with the request rather than the router.""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + calls: Final[list[str]] = [] - router = _make_router() + def upstream(request: httpx.Request) -> httpx.Response: + model: Final = json.loads(request.content)["model"] + calls.append(model) + if model == "primary-model" and primary_fails_pre_stream: + return httpx.Response(500, json={"error": {"message": "primary down", "type": "server_error"}}) + if model in ("primary-model", "fb1-model"): + return _openai_stream_cut_after("") + return _sse(_openai_completed(model)) - class _ClientErrorSource: - completed_response = None - - def __aiter__(self): - return self - - async def __anext__(self): - raise litellm.APIError( - status_code=400, - message="bad request", - llm_provider="openai", - model="gpt-5", - ) - - wrapped = await router._aresponses_streaming_iterator( - response=_ClientErrorSource(), - initial_kwargs={"model": "primary"}, + respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=upstream) + fallbacks: Final = [{"primary": ["fb1", "fb2"]}] + router: Final = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": f"openai/{name}-model", "api_key": "sk-test"}} + for name in ("primary", "fb1", "fb2") + ], + num_retries=0, + **({} if request_level_fallbacks else {"fallbacks": fallbacks}), ) + try: + stream: Final = await router.aresponses( + model="primary", input="hi", stream=True, **({"fallbacks": fallbacks} if request_level_fallbacks else {}) + ) + events: Final = [event async for event in stream] + finally: + router.discard() - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - with pytest.raises(litellm.APIError) as exc_info: - async for _ in wrapped: - pass + assert calls == ["primary-model", "fb1-model", "fb2-model"] + assert [event.type for event in events] == [ResponsesAPIStreamEvents.RESPONSE_COMPLETED] - assert exc_info.value.status_code == 400 - mock_fallback.assert_not_awaited() + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("error", "falls_back"), + [ + ({"type": "tokens", "code": "rate_limit_exceeded", "message": "rate limited"}, True), + ({"type": "invalid_request_error", "code": "invalid_request_error", "message": "bad request"}, False), + ], + ids=["retriable_429_falls_back", "client_400_surfaces"], +) +async def test_responses_in_stream_error_event_falls_back_only_when_retriable( + error: Mapping[str, object], falls_back: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + """A retriable in-stream error event (429) takes the router's mid-stream fallback path; a 400-mapped one + surfaces to the caller without invoking a fallback.""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + calls: Final[list[str]] = [] + + def upstream(request: httpx.Request) -> httpx.Response: + model: Final = json.loads(request.content)["model"] + calls.append(model) + if model == "primary-model": + return _sse({"type": "error", "sequence_number": 0, "error": error}) + return _sse(_openai_completed(model)) + + respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=upstream) + router: Final = Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": f"openai/{name}-model", "api_key": "sk-test"}} + for name in ("primary", "fallback") + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + ) + try: + stream: Final = await router.aresponses(model="primary", input="hi", stream=True) + if falls_back: + events: Final = [event async for event in stream] + assert [event.type for event in events] == [ResponsesAPIStreamEvents.RESPONSE_COMPLETED] + assert calls == ["primary-model", "fallback-model"] + else: + with pytest.raises(litellm.BadRequestError) as exc_info: + async for _ in stream: + pass + assert exc_info.value.status_code == 400 + assert calls == ["primary-model"] + finally: + router.discard() diff --git a/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py b/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py index 88fc8ba6e6a..d4dd8489927 100644 --- a/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py +++ b/tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py @@ -1,15 +1,14 @@ -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import httpx import pytest from pydantic import ValidationError - from litellm.exceptions import AuthenticationError +from litellm.llms.github_copilot.common_utils import GetAccessTokenError from litellm.llms.github_copilot.embedding.transformation import ( GithubCopilotEmbeddingConfig, ) -from litellm.llms.github_copilot.common_utils import GetAPIKeyError from litellm.types.utils import EmbeddingResponse @@ -41,7 +40,7 @@ def test_github_copilot_embedding_config_validate_environment(): assert "x-request-id" in validated_headers # Test with authentication failure - config.authenticator.get_api_key.side_effect = GetAPIKeyError( + config.authenticator.get_api_key.side_effect = GetAccessTokenError( message="Failed to get API key", status_code=401, ) diff --git a/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py index 20ff65f11d9..fbdf1ba4591 100644 --- a/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py +++ b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py @@ -4,9 +4,8 @@ from unittest.mock import MagicMock import httpx import pytest - from litellm.exceptions import AuthenticationError -from litellm.llms.github_copilot.common_utils import GetAPIKeyError +from litellm.llms.github_copilot.common_utils import GetAccessTokenError from litellm.llms.github_copilot.messages.transformation import ( GithubCopilotAnthropicMessagesConfig, ) @@ -216,7 +215,9 @@ def test_github_copilot_anthropic_messages_validate_environment_auth_error(): # Mock the authenticator to raise an error config.authenticator = MagicMock() - config.authenticator.get_api_key.side_effect = GetAPIKeyError(status_code=401, message="No valid API key found") + config.authenticator.get_api_key.side_effect = GetAccessTokenError( + status_code=401, message="No valid API key found" + ) with pytest.raises(AuthenticationError): config.validate_anthropic_messages_environment( diff --git a/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py b/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py index 55fd02bc4dd..8eb70d366b1 100644 --- a/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py +++ b/tests/unit/llms/github_copilot/test_github_copilot_authenticator.py @@ -1,559 +1,384 @@ -import json -import logging import os -import time from concurrent.futures import ThreadPoolExecutor -from datetime import datetime, timedelta from pathlib import Path -from types import MappingProxyType from typing import Final -from unittest.mock import MagicMock, mock_open, patch +from unittest.mock import MagicMock, patch import pytest import respx -from litellm.llms.github_copilot.authenticator import Authenticator -from litellm.llms.github_copilot.common_utils import ( - APIKeyExpiredError, - GetAccessTokenError, - GetAPIKeyError, - GetDeviceCodeError, - RefreshAPIKeyError, -) +from litellm.llms.github_copilot.authenticator import Authenticator, github_api_headers +from litellm.llms.github_copilot.common_utils import GetAccessTokenError, GetDeviceCodeError -API_KEY_URL: Final = "https://copilot-key.test/copilot_internal/v2/token" -REFRESHED_KEY: Final = MappingProxyType({"token": "new-api-key", "expires_at": 4102444800}) +COPILOT_USER_URL: Final = "https://api.github.com/copilot_internal/user" +DEVICE_CODE_URL: Final = "https://github.com/login/device/code" TENANT_API: Final = "https://api.tenant.githubcopilot.com" -REPAIRED_API: Final = "https://api.repaired.githubcopilot.com" TRUNCATED_UTF_16: Final = b"\xff\xfe\x00\xd8" @pytest.fixture -def logged_in_token_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: +def logged_in_token_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, str]: monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) - monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", API_KEY_URL) - (tmp_path / "access-token").write_text("gho_logged_in") - return tmp_path + token: Final = f"gho_{tmp_path.name}" + (tmp_path / "access-token").write_text(token) + return tmp_path, token -class TestGitHubCopilotAuthenticator: - @pytest.fixture - def authenticator(self): - with ( - patch("os.path.exists", return_value=False), - patch("os.makedirs") as mock_makedirs, - ): - auth = Authenticator() - mock_makedirs.assert_called_once() - return auth +def test_init_uses_access_token_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + auth: Final = Authenticator() + assert auth.token_dir == str(tmp_path) + assert auth.access_token_file == str(tmp_path / "access-token") - @pytest.fixture - def mock_http_client(self): - mock_client = MagicMock() - mock_response = MagicMock() - mock_client.get.return_value = mock_response - mock_client.post.return_value = mock_response - mock_response.raise_for_status.return_value = None - return mock_client, mock_response - def test_init(self): - """Test the initialization of the authenticator.""" - with ( - patch("os.path.exists", return_value=False), - patch("os.makedirs") as mock_makedirs, - ): - auth = Authenticator() - assert auth.token_dir.endswith("/github_copilot") - assert auth.access_token_file.endswith("/access-token") - assert auth.api_key_file.endswith("/api-key.json") - mock_makedirs.assert_called_once() +def test_init_uses_default_access_token_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.delenv("GITHUB_COPILOT_TOKEN_DIR", raising=False) - def test_ensure_token_dir(self): - """Test that the token directory is created if it doesn't exist.""" - with ( - patch("os.path.exists", return_value=False), - patch("os.makedirs") as mock_makedirs, - ): - auth = Authenticator() - mock_makedirs.assert_called_once_with(auth.token_dir, exist_ok=True) + auth: Final = Authenticator() - def test_get_github_headers(self, authenticator): - """Test that GitHub headers are correctly generated.""" - headers = authenticator._get_github_headers() - assert "accept" in headers - assert "editor-version" in headers - assert "user-agent" in headers - assert "content-type" in headers + assert auth.token_dir == str(tmp_path / ".config/litellm/github_copilot") + assert auth.access_token_file == str(tmp_path / ".config/litellm/github_copilot/access-token") - headers_with_token = authenticator._get_github_headers("test-token") - assert headers_with_token["authorization"] == "token test-token" - def test_get_access_token_from_file(self, logged_in_token_dir: Path) -> None: - """Test retrieving an access token from a file.""" - assert Authenticator().get_access_token() == "gho_logged_in" +def test_ensure_token_dir_creates_missing_directory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + token_dir: Final = str(tmp_path / "github_copilot") + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", token_dir) - def test_get_access_token_login(self, authenticator): - """Test logging in to get an access token.""" - mock_token = "mock-access-token" + auth: Final = Authenticator() - with ( - patch.object(authenticator, "_login", return_value=mock_token), - patch("builtins.open", mock_open()), - patch("builtins.open", side_effect=IOError) as mock_read, - ): - token = authenticator.get_access_token() - assert token == mock_token - authenticator._login.assert_called_once() + assert Path(auth.token_dir).is_dir() - @pytest.mark.asyncio - async def test_get_access_token_refuses_device_code_login_in_event_loop(self, authenticator): - with ( - patch("builtins.open", side_effect=FileNotFoundError), - patch.object(authenticator, "_login") as mock_login, - ): - with pytest.raises(GetAccessTokenError) as exc: - authenticator.get_access_token() - assert exc.value.status_code == 401 - assert "event loop" in str(exc.value) - assert authenticator.access_token_file not in str(exc.value) - mock_login.assert_not_called() +def test_get_github_headers_include_access_token() -> None: + headers: Final = github_api_headers("github-token") - def test_get_access_token_refuses_device_code_login_in_worker_thread(self, authenticator): - with ( - patch("builtins.open", side_effect=FileNotFoundError), - patch.object(authenticator, "_login") as mock_login, - ): - with ThreadPoolExecutor(max_workers=1) as pool: - with pytest.raises(GetAccessTokenError) as exc: - pool.submit(authenticator.get_access_token).result() + assert headers == { + "accept": "application/json", + "editor-version": "vscode/1.85.1", + "editor-plugin-version": "copilot/1.155.0", + "user-agent": "GithubCopilot/1.155.0", + "accept-encoding": "gzip,deflate,br", + "content-type": "application/json", + "authorization": "token github-token", + } - assert exc.value.status_code == 401 - assert "worker thread" in str(exc.value) - assert authenticator.access_token_file not in str(exc.value) - mock_login.assert_not_called() - @pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"]) - def test_saved_access_token_in_another_encoding_is_sent_on_refresh( - self, logged_in_token_dir: Path, encoding: str - ) -> None: - """An access-token file a Windows editor saved as UTF-16 or with a BOM still refreshes the key.""" - (logged_in_token_dir / "access-token").write_bytes("gho_logged_in\n".encode(encoding)) +def test_get_access_token_reads_saved_token(logged_in_token_dir: tuple[Path, str]) -> None: + _, token = logged_in_token_dir + assert Authenticator().get_access_token() == token - with respx.mock(assert_all_called=True) as router: - route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) - assert Authenticator().get_api_key() == "new-api-key" - assert route.calls.last.request.headers["authorization"] == "token gho_logged_in" +def test_get_access_token_logs_in_and_saves_token( + logged_in_token_dir: tuple[Path, str], +) -> None: + token_dir, _ = logged_in_token_dir + access_token_file: Final = token_dir / "access-token" + access_token_file.unlink() + auth: Final = Authenticator() - def test_undecodable_access_token_is_refused_like_a_missing_one(self, logged_in_token_dir: Path) -> None: - """A proxy worker gets the login hint for an access-token file it cannot decode, and the file is kept.""" - access_token_file: Final = logged_in_token_dir / "access-token" - access_token_file.write_bytes(TRUNCATED_UTF_16) + with patch.object(auth, "_login", return_value="github-access-token") as login: + assert auth.get_access_token() == "github-access-token" - with ThreadPoolExecutor(max_workers=1) as pool, pytest.raises(GetAccessTokenError) as exc: - pool.submit(Authenticator().get_access_token).result() + login.assert_called_once() + assert access_token_file.read_text() == "github-access-token" - assert "worker thread" in str(exc.value) - assert access_token_file.read_bytes() == TRUNCATED_UTF_16 - def test_get_access_token_failure(self, authenticator): - """Test that an exception is raised after multiple login failures.""" - with ( - patch.object( - authenticator, - "_login", - side_effect=GetDeviceCodeError(message="Test error", status_code=400), - ), - patch("builtins.open", side_effect=IOError), - ): - with pytest.raises(GetAccessTokenError): - authenticator.get_access_token() - assert authenticator._login.call_count == 3 +def test_get_api_key_returns_github_token_without_http( + logged_in_token_dir: tuple[Path, str], +) -> None: + _, token = logged_in_token_dir + with patch("litellm.llms.github_copilot.authenticator.get_httpx_client") as get_client: + assert Authenticator().get_api_key() == token + get_client.assert_not_called() - def test_get_api_key_from_file(self, authenticator): - """Test retrieving an API key from a file.""" - future_time = (datetime.now() + timedelta(hours=1)).timestamp() - mock_api_key_data = json.dumps( - {"token": "mock-api-key", "expires_at": future_time} - ).encode() - with patch("builtins.open", mock_open(read_data=mock_api_key_data)): - api_key = authenticator.get_api_key() - assert api_key == "mock-api-key" - - def test_get_api_key_expired(self, authenticator): - """Test refreshing an expired API key.""" - past_time = (datetime.now() - timedelta(hours=1)).timestamp() - mock_expired_data = json.dumps( - {"token": "expired-api-key", "expires_at": past_time} - ).encode() - mock_new_data = { - "token": "new-api-key", - "expires_at": (datetime.now() + timedelta(hours=1)).timestamp(), - } - - with ( - patch("builtins.open", mock_open(read_data=mock_expired_data)), - patch.object(authenticator, "_refresh_api_key", return_value=mock_new_data), - patch("json.dump") as mock_json_dump, - ): - api_key = authenticator.get_api_key() - assert api_key == "new-api-key" - authenticator._refresh_api_key.assert_called_once() - - @pytest.mark.parametrize( - "malformed", - [ - b"[]", - json.dumps({"token": "tid=abc", "expires_at": "soon"}).encode(), - TRUNCATED_UTF_16, - json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": []}).encode(), - ], - ids=["json-list", "text-expires-at", "truncated-utf-16", "list-endpoints-beside-live-key"], - ) - def test_get_api_key_malformed_file_refreshes_like_missing( - self, logged_in_token_dir: Path, malformed: bytes - ) -> None: - """A file that is not a Copilot key object is ignored, refreshed over, never answered with a parse error.""" - api_key_file: Final = logged_in_token_dir / "api-key.json" - api_key_file.write_bytes(malformed) - - with respx.mock(assert_all_called=True) as router: - route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) - assert Authenticator().get_api_key() == "new-api-key" - - assert route.call_count == 1 - assert route.calls.last.request.headers["authorization"] == "token gho_logged_in" - assert json.loads(api_key_file.read_text()) == REFRESHED_KEY - - @pytest.mark.parametrize( - "malformed", - [ - b"[]", - json.dumps({"token": "tid=abc", "expires_at": "soon"}).encode(), - TRUNCATED_UTF_16, - json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": []}).encode(), - ], - ids=["json-list", "text-expires-at", "truncated-utf-16", "list-endpoints-beside-live-key"], - ) - def test_get_api_key_malformed_file_is_kept_when_refresh_fails( - self, logged_in_token_dir: Path, malformed: bytes - ) -> None: - """Ignoring the file never deletes it, so a failed refresh leaves the operator's file for inspection.""" - api_key_file: Final = logged_in_token_dir / "api-key.json" - api_key_file.write_bytes(malformed) - - with respx.mock(assert_all_called=True) as router: - router.get(API_KEY_URL).respond(401) - with pytest.raises(GetAPIKeyError): - Authenticator().get_api_key() - - assert api_key_file.read_bytes() == malformed - - def test_refreshed_key_that_cannot_be_saved_is_reused_with_its_endpoint(self, logged_in_token_dir: Path) -> None: - unwritable_key_file: Final = logged_in_token_dir / "api-key.json" - unwritable_key_file.mkdir() - first: Final = Authenticator() - second: Final = Authenticator() - - with respx.mock(assert_all_called=True) as router: - route: Final = router.get(API_KEY_URL).respond(json={**REFRESHED_KEY, "endpoints": {"api": TENANT_API}}) - served: Final = (first.get_api_key(), second.get_api_key(), first.get_api_key()) - - assert served == ("new-api-key",) * 3 - assert (first.get_api_base(), second.get_api_base()) == (TENANT_API, TENANT_API) - assert route.call_count == 1 - assert unwritable_key_file.is_dir() - - def test_refreshed_key_that_cannot_be_saved_and_has_expired_is_refreshed_again( - self, logged_in_token_dir: Path - ) -> None: - (logged_in_token_dir / "api-key.json").mkdir() +def test_get_api_base_uses_user_endpoint_and_memoizes_per_access_token( + logged_in_token_dir: tuple[Path, str], +) -> None: + _, token = logged_in_token_dir + with respx.mock(assert_all_called=True) as router: + route: Final = router.get(COPILOT_USER_URL).respond(json={"endpoints": {"api": TENANT_API}}) auth: Final = Authenticator() + assert auth.get_api_base() == TENANT_API + assert Authenticator().get_api_base() == TENANT_API - with respx.mock(assert_all_called=True) as router: - route: Final = router.get(API_KEY_URL).respond(json={"token": "expired-api-key", "expires_at": 1}) - served: Final = (auth.get_api_key(), auth.get_api_key()) + assert route.call_count == 1 + assert route.calls.last.request.headers["authorization"] == f"token {token}" - assert served == ("expired-api-key",) * 2 - assert route.call_count == 2 - - def test_kept_key_and_its_endpoint_stay_paired_after_the_key_file_is_repaired( - self, logged_in_token_dir: Path - ) -> None: - key_file: Final = logged_in_token_dir / "api-key.json" - key_file.mkdir() - auth: Final = Authenticator() - - with respx.mock(assert_all_called=True) as router: - router.get(API_KEY_URL).respond(json={**REFRESHED_KEY, "endpoints": {"api": TENANT_API}}) - auth.get_api_key() - key_file.rmdir() - key_file.write_text( - json.dumps({"token": "repaired-api-key", "expires_at": 9999999999, "endpoints": {"api": REPAIRED_API}}) - ) - - assert (auth.get_api_key(), auth.get_api_base()) in ( - ("new-api-key", TENANT_API), - ("repaired-api-key", REPAIRED_API), - ) - - def test_unsaved_key_is_not_reused_after_the_access_token_changes(self, logged_in_token_dir: Path) -> None: - (logged_in_token_dir / "api-key.json").mkdir() - - with respx.mock(assert_all_called=True) as router: - route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) - Authenticator().get_api_key() - (logged_in_token_dir / "access-token").write_text("gho_logged_in_again") - Authenticator().get_api_key() - - assert [call.request.headers["authorization"] for call in route.calls] == [ - "token gho_logged_in", - "token gho_logged_in_again", - ] - - @pytest.mark.parametrize( - "malformed", - [ - b"[]", - json.dumps({"endpoints": {"api": 5}}).encode(), - json.dumps({"token": "t", "expires_at": 1, "endpoints": []}).encode(), - TRUNCATED_UTF_16, - ], - ids=["json-list", "number-api", "list-endpoints", "truncated-utf-16"], - ) - def test_get_api_base_malformed_file_returns_none( - self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, malformed: bytes - ) -> None: - """get_api_base runs before get_api_key on every request, so it must survive malformed shapes too.""" - monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) - (tmp_path / "api-key.json").write_bytes(malformed) +@pytest.mark.parametrize("status_code", [401, 403, 500]) +def test_get_api_base_returns_none_when_user_lookup_fails( + logged_in_token_dir: tuple[Path, str], status_code: int +) -> None: + with respx.mock(assert_all_called=True) as router: + router.get(COPILOT_USER_URL).respond(status_code=status_code) assert Authenticator().get_api_base() is None - def test_get_api_base_keeps_endpoints_when_the_key_is_malformed(self, logged_in_token_dir: Path) -> None: - """The endpoint is read on its own, so a bad key never sends the refreshed key to the default host.""" - (logged_in_token_dir / "api-key.json").write_text( - json.dumps({"token": "tid=abc", "expires_at": "soon", "endpoints": {"api": TENANT_API}}) - ) - auth: Final = Authenticator() +def test_get_api_base_retries_after_failed_user_lookup(logged_in_token_dir: tuple[Path, str]) -> None: + auth: Final = Authenticator() + with respx.mock(assert_all_called=True) as router: + failed_lookup: Final = router.get(COPILOT_USER_URL).respond(status_code=503) + assert auth.get_api_base() is None + assert failed_lookup.call_count == 1 + + with respx.mock(assert_all_called=True) as router: + successful_lookup: Final = router.get(COPILOT_USER_URL).respond(json={"endpoints": {"api": TENANT_API}}) assert auth.get_api_base() == TENANT_API - with respx.mock(assert_all_called=True) as router: - router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) - assert auth.get_api_key() == "new-api-key" + assert successful_lookup.call_count == 1 - @pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"]) - def test_live_key_file_in_another_json_encoding_is_used_without_a_refresh( - self, logged_in_token_dir: Path, encoding: str - ) -> None: - """A key file a Windows editor saved as UTF-16 or with a BOM still holds a live key and tenant.""" - live_key: Final = json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": {"api": TENANT_API}}) - api_key_file: Final = logged_in_token_dir / "api-key.json" - api_key_file.write_bytes(live_key.encode(encoding)) - auth: Final = Authenticator() - with respx.mock: - assert auth.get_api_base() == TENANT_API - assert auth.get_api_key() == "tid=abc" +def test_get_api_base_returns_none_when_user_lookup_response_has_no_endpoint( + logged_in_token_dir: tuple[Path, str], +) -> None: + with respx.mock(assert_all_called=True) as router: + router.get(COPILOT_USER_URL).respond(json={"endpoints": {}}) + assert Authenticator().get_api_base() is None - assert api_key_file.read_bytes() == live_key.encode(encoding) - def test_ignored_key_file_is_logged_without_its_contents( - self, logged_in_token_dir: Path, caplog: pytest.LogCaptureFixture - ) -> None: - """The warning names the file but never echoes a token saved in the wrong shape.""" - (logged_in_token_dir / "api-key.json").write_text( - json.dumps({"token": ["tid=secret-value"], "expires_at": 4102444800}) - ) +def test_get_api_base_returns_none_when_user_lookup_endpoints_are_not_a_mapping( + logged_in_token_dir: tuple[Path, str], +) -> None: + with respx.mock(assert_all_called=True) as router: + router.get(COPILOT_USER_URL).respond(json={"endpoints": []}) + assert Authenticator().get_api_base() is None - with caplog.at_level(logging.WARNING, logger="LiteLLM"), respx.mock(assert_all_called=True) as router: - router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY)) - assert Authenticator().get_api_key() == "new-api-key" - assert "not a Copilot API key file" in caplog.text - assert "tid=secret-value" not in caplog.text +def test_get_api_base_rejects_untrusted_user_endpoint( + logged_in_token_dir: tuple[Path, str], +) -> None: + with respx.mock(assert_all_called=True) as router: + router.get(COPILOT_USER_URL).respond(json={"endpoints": {"api": "https://attacker.example"}}) + assert Authenticator().get_api_base() is None - def test_refresh_api_key(self, authenticator, mock_http_client): - """Test refreshing an API key.""" - mock_client, mock_response = mock_http_client - mock_token = "mock-access-token" - mock_api_key_data = {"token": "new-api-key", "expires_at": 12345} - with ( - patch.object(authenticator, "get_access_token", return_value=mock_token), - patch( - "litellm.llms.github_copilot.authenticator.get_httpx_client", - return_value=mock_client, - ), - patch.object(mock_response, "json", return_value=mock_api_key_data), - ): - result = authenticator._refresh_api_key() - assert result == mock_api_key_data - mock_client.get.assert_called_once() - authenticator.get_access_token.assert_called_once() +def test_get_api_base_does_not_trigger_device_login_without_saved_token( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + with patch("litellm.llms.github_copilot.authenticator.get_httpx_client") as get_client: + assert Authenticator().get_api_base() is None + get_client.assert_not_called() - def test_refresh_api_key_failure(self, authenticator, mock_http_client): - """Test failure to refresh an API key.""" - mock_client, mock_response = mock_http_client - mock_token = "mock-access-token" - with ( - patch.object(authenticator, "get_access_token", return_value=mock_token), - patch( - "litellm.llms.github_copilot.authenticator.get_httpx_client", - return_value=mock_client, - ), - patch.object(mock_response, "json", return_value={}), - ): - with pytest.raises(RefreshAPIKeyError): - authenticator._refresh_api_key() - assert mock_client.get.call_count == 3 +@pytest.mark.asyncio +async def test_get_access_token_refuses_device_code_login_in_event_loop() -> None: + auth: Final = Authenticator() + with ( + patch("builtins.open", side_effect=FileNotFoundError), + patch.object(auth, "_login") as login, + ): + with pytest.raises(GetAccessTokenError) as exc: + auth.get_access_token() - def test_get_device_code(self, authenticator, mock_http_client): - """Test getting a device code.""" - mock_client, mock_response = mock_http_client - mock_device_code_data = { - "device_code": "mock-device-code", - "user_code": "ABCD-EFGH", - "verification_uri": "https://github.com/login/device", - } + assert exc.value.status_code == 401 + assert "event loop" in str(exc.value) + login.assert_not_called() - with ( - patch( - "litellm.llms.github_copilot.authenticator.get_httpx_client", - return_value=mock_client, - ), - patch.object(mock_response, "json", return_value=mock_device_code_data), - ): - result = authenticator._get_device_code() - assert result == mock_device_code_data - mock_client.post.assert_called_once() - def test_poll_for_access_token(self, authenticator, mock_http_client): - """Test polling for an access token.""" - mock_client, mock_response = mock_http_client - mock_token_data = {"access_token": "mock-access-token"} +def test_get_access_token_refuses_device_code_login_in_worker_thread() -> None: + auth: Final = Authenticator() + with ( + patch("builtins.open", side_effect=FileNotFoundError), + patch.object(auth, "_login") as login, + ): + with ThreadPoolExecutor(max_workers=1) as pool: + with pytest.raises(GetAccessTokenError) as exc: + pool.submit(auth.get_access_token).result() - with ( - patch( - "litellm.llms.github_copilot.authenticator.get_httpx_client", - return_value=mock_client, - ), - patch.object(mock_response, "json", return_value=mock_token_data), - patch("time.sleep"), - ): - result = authenticator._poll_for_access_token("mock-device-code") - assert result == "mock-access-token" - mock_client.post.assert_called_once() + assert exc.value.status_code == 401 + assert "worker thread" in str(exc.value) + login.assert_not_called() - def test_login(self, authenticator): - """Test the login process.""" - mock_device_code_data = { - "device_code": "mock-device-code", - "user_code": "ABCD-EFGH", - "verification_uri": "https://github.com/login/device", - } - mock_token = "mock-access-token" - with ( - patch.object( - authenticator, "_get_device_code", return_value=mock_device_code_data - ), - patch.object( - authenticator, "_poll_for_access_token", return_value=mock_token - ), - patch("builtins.print") as mock_print, - ): - result = authenticator._login() - assert result == mock_token - authenticator._get_device_code.assert_called_once() - authenticator._poll_for_access_token.assert_called_once_with( - "mock-device-code" - ) - mock_print.assert_called_once() +@pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"]) +def test_get_access_token_reads_saved_token_with_supported_encoding( + logged_in_token_dir: tuple[Path, str], encoding: str +) -> None: + token_dir, token = logged_in_token_dir + (token_dir / "access-token").write_bytes(f"{token}\n".encode(encoding)) - def test_get_api_base_from_file(self, authenticator): - """Test retrieving the API base endpoint from a file.""" - mock_api_key_data = json.dumps( + assert Authenticator().get_access_token() == token + + +def test_get_access_token_refuses_undecodable_saved_token( + logged_in_token_dir: tuple[Path, str], +) -> None: + token_dir, _ = logged_in_token_dir + access_token_file: Final = token_dir / "access-token" + access_token_file.write_bytes(TRUNCATED_UTF_16) + + with ThreadPoolExecutor(max_workers=1) as pool, pytest.raises(GetAccessTokenError) as exc: + pool.submit(Authenticator().get_access_token).result() + + assert exc.value.status_code == 401 + assert "worker thread" in str(exc.value) + assert access_token_file.read_bytes() == TRUNCATED_UTF_16 + + +def test_get_access_token_retries_device_flow_three_times( + logged_in_token_dir: tuple[Path, str], +) -> None: + token_dir, _ = logged_in_token_dir + (token_dir / "access-token").unlink() + auth: Final = Authenticator() + + with ( + patch.object( + auth, + "_login", + side_effect=GetDeviceCodeError(message="device flow failed", status_code=400), + ) as login, + pytest.raises(GetAccessTokenError, match="after 3 attempts"), + ): + auth.get_access_token() + + assert login.call_count == 3 + + +def test_get_access_token_returns_login_token_when_file_write_fails( + logged_in_token_dir: tuple[Path, str], +) -> None: + token_dir, _ = logged_in_token_dir + (token_dir / "access-token").unlink() + auth: Final = Authenticator() + + with ( + patch.object(auth, "_login", return_value="github-access-token"), + patch("builtins.open", side_effect=(FileNotFoundError, PermissionError)), + ): + assert auth.get_access_token() == "github-access-token" + + +@pytest.mark.parametrize( + ("environment", "url", "body"), + [ + ( + {"GITHUB_COPILOT_DEVICE_CODE_URL": "https://custom.example/device"}, + "https://custom.example/device", + {"client_id": "Iv1.b507a08c87ecfe98", "scope": "read:user"}, + ), + ( + {"GITHUB_COPILOT_CLIENT_ID": "custom-client"}, + DEVICE_CODE_URL, + {"client_id": "custom-client", "scope": "read:user"}, + ), + ], +) +def test_get_device_code_uses_runtime_url_and_client_id( + environment: dict[str, str], url: str, body: dict[str, str] +) -> None: + client: Final = MagicMock() + response: Final = MagicMock() + response.json.return_value = { + "device_code": "dc", + "user_code": "UC", + "verification_uri": "https://github.com/login/device", + } + response.raise_for_status.return_value = None + client.post.return_value = response + with ( + patch.dict(os.environ, environment), + patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=client), + ): + assert Authenticator()._get_device_code()["device_code"] == "dc" + client.post.assert_called_once() + assert client.post.call_args.args[0] == url + assert client.post.call_args.kwargs["json"] == body + + +def test_get_device_code_sends_device_flow_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("GITHUB_COPILOT_DEVICE_CODE_URL", raising=False) + monkeypatch.delenv("GITHUB_COPILOT_CLIENT_ID", raising=False) + client: Final = MagicMock() + response: Final = MagicMock() + device_code: Final = { + "device_code": "device-code", + "user_code": "ABCD-EFGH", + "verification_uri": "https://github.com/login/device", + } + response.json.return_value = device_code + response.raise_for_status.return_value = None + client.post.return_value = response + + with patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=client): + assert Authenticator()._get_device_code() == device_code + + client.post.assert_called_once_with( + DEVICE_CODE_URL, + headers=github_api_headers(), + json={"client_id": "Iv1.b507a08c87ecfe98", "scope": "read:user"}, + ) + + +def test_poll_for_access_token_sends_device_code_and_grant_type( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", raising=False) + monkeypatch.delenv("GITHUB_COPILOT_CLIENT_ID", raising=False) + client: Final = MagicMock() + response: Final = MagicMock() + response.json.return_value = {"access_token": "github-access-token"} + response.raise_for_status.return_value = None + client.post.return_value = response + + with patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=client): + assert Authenticator()._poll_for_access_token("device-code") == "github-access-token" + + client.post.assert_called_once_with( + "https://github.com/login/oauth/access_token", + headers=github_api_headers(), + json={ + "client_id": "Iv1.b507a08c87ecfe98", + "device_code": "device-code", + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + }, + ) + + +def test_poll_for_access_token_uses_runtime_url_and_client_id() -> None: + client: Final = MagicMock() + response: Final = MagicMock() + response.json.return_value = {"access_token": "github-token"} + response.raise_for_status.return_value = None + client.post.return_value = response + with ( + patch.dict( + os.environ, { - "token": "mock-api-key", - "expires_at": (datetime.now() + timedelta(hours=1)).timestamp(), - "endpoints": {"api": "https://api.enterprise.githubcopilot.com"}, - } - ).encode() - with patch("builtins.open", mock_open(read_data=mock_api_key_data)): - api_base = authenticator.get_api_base() - assert api_base == "https://api.enterprise.githubcopilot.com" + "GITHUB_COPILOT_ACCESS_TOKEN_URL": "https://custom.example/token", + "GITHUB_COPILOT_CLIENT_ID": "custom-client", + }, + ), + patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=client), + patch("time.sleep"), + ): + assert Authenticator()._poll_for_access_token("dc") == "github-token" + assert client.post.call_args.args[0] == "https://custom.example/token" + assert client.post.call_args.kwargs["json"]["client_id"] == "custom-client" - def test_get_device_code_with_custom_url(self, authenticator, mock_http_client): - """GITHUB_COPILOT_DEVICE_CODE_URL env var must be used by _get_device_code at call time.""" - mock_client, mock_response = mock_http_client - custom_url = "https://custom.example.com/device" - mock_response.json.return_value = { - "device_code": "dc", - "user_code": "UC", - "verification_uri": "https://example.com", - } - with ( - patch.dict(os.environ, {"GITHUB_COPILOT_DEVICE_CODE_URL": custom_url}), - patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client), - ): - authenticator._get_device_code() - assert mock_client.post.call_args[0][0] == custom_url - def test_get_device_code_with_custom_client_id(self, authenticator, mock_http_client): - """GITHUB_COPILOT_CLIENT_ID env var must appear as client_id in the device-code request body.""" - mock_client, mock_response = mock_http_client - custom_id = "custom_client_id" - mock_response.json.return_value = { - "device_code": "dc", - "user_code": "UC", - "verification_uri": "https://example.com", - } - with ( - patch.dict(os.environ, {"GITHUB_COPILOT_CLIENT_ID": custom_id}), - patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client), - ): - authenticator._get_device_code() - assert mock_client.post.call_args[1]["json"]["client_id"] == custom_id +def test_login_requests_device_code_and_polls_for_access_token() -> None: + auth: Final = Authenticator() + device_code: Final = { + "device_code": "device-code", + "user_code": "ABCD-EFGH", + "verification_uri": "https://github.com/login/device", + } - def test_poll_for_access_token_with_custom_url(self, authenticator, mock_http_client): - """GITHUB_COPILOT_ACCESS_TOKEN_URL env var must be used by _poll_for_access_token at call time.""" - mock_client, mock_response = mock_http_client - custom_url = "https://custom.example.com/token" - mock_response.json.return_value = {"access_token": "tok"} - with ( - patch.dict(os.environ, {"GITHUB_COPILOT_ACCESS_TOKEN_URL": custom_url}), - patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client), - patch("time.sleep"), - ): - authenticator._poll_for_access_token("dc") - assert mock_client.post.call_args[0][0] == custom_url + with ( + patch.object(auth, "_get_device_code", return_value=device_code) as get_device_code, + patch.object(auth, "_poll_for_access_token", return_value="github-access-token") as poll, + patch("builtins.print") as print_message, + ): + assert auth._login() == "github-access-token" - def test_poll_for_access_token_with_custom_client_id(self, authenticator, mock_http_client): - """GITHUB_COPILOT_CLIENT_ID env var must appear as client_id in the polling request body.""" - mock_client, mock_response = mock_http_client - custom_id = "custom_client_id" - mock_response.json.return_value = {"access_token": "tok"} - with ( - patch.dict(os.environ, {"GITHUB_COPILOT_CLIENT_ID": custom_id}), - patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client), - patch("time.sleep"), - ): - authenticator._poll_for_access_token("dc") - assert mock_client.post.call_args[1]["json"]["client_id"] == custom_id - - def test_refresh_api_key_with_custom_url(self, authenticator, mock_http_client): - """GITHUB_COPILOT_API_KEY_URL env var must be used by _refresh_api_key at call time.""" - mock_client, mock_response = mock_http_client - custom_url = "https://custom.example.com/api-key" - mock_response.json.return_value = {"token": "api-tok", "expires_at": 9999999999} - with ( - patch.dict(os.environ, {"GITHUB_COPILOT_API_KEY_URL": custom_url}), - patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client), - patch.object(authenticator, "get_access_token", return_value="access-tok"), - ): - authenticator._refresh_api_key() - assert mock_client.get.call_args[0][0] == custom_url + get_device_code.assert_called_once() + poll.assert_called_once_with("device-code") + print_message.assert_called_once_with( + "Please visit https://github.com/login/device and enter code ABCD-EFGH to authenticate.", + flush=True, + ) diff --git a/tests/unit/llms/github_copilot/test_github_copilot_transformation.py b/tests/unit/llms/github_copilot/test_github_copilot_transformation.py index 2dbfae07526..d453d7c6622 100644 --- a/tests/unit/llms/github_copilot/test_github_copilot_transformation.py +++ b/tests/unit/llms/github_copilot/test_github_copilot_transformation.py @@ -1,29 +1,17 @@ -import asyncio import json -from datetime import datetime, timedelta -from typing import AsyncGenerator -from unittest.mock import AsyncMock, MagicMock, mock_open, patch - -import pytest - +from unittest.mock import MagicMock, patch import httpx -from respx import MockRouter +import pytest import litellm # Import at the top to make the patch work correctly -import litellm.llms.github_copilot.chat.transformation -from litellm import Choices, Message, ModelResponse, Usage, acompletion, completion +from litellm import ModelResponse, completion from litellm.exceptions import AuthenticationError -from litellm.llms.github_copilot.authenticator import Authenticator from litellm.llms.github_copilot.chat.transformation import GithubCopilotConfig from litellm.llms.github_copilot.common_utils import ( - APIKeyExpiredError, GetAccessTokenError, - GetAPIKeyError, - GetDeviceCodeError, - RefreshAPIKeyError, ) @@ -37,9 +25,7 @@ def test_github_copilot_config_get_openai_compatible_provider_info(): config.authenticator = MagicMock() config.authenticator.get_api_key.return_value = mock_api_key # Test with dynamic endpoint - config.authenticator.get_api_base.return_value = ( - "https://api.enterprise.githubcopilot.com" - ) + config.authenticator.get_api_base.return_value = "https://api.enterprise.githubcopilot.com" # Test with default values model = "github_copilot/gpt-4" @@ -73,7 +59,7 @@ def test_github_copilot_config_get_openai_compatible_provider_info(): assert api_base == "https://api.githubcopilot.com" # Test with authentication failure - config.authenticator.get_api_key.side_effect = GetAPIKeyError( + config.authenticator.get_api_key.side_effect = GetAccessTokenError( message="Failed to get API key", status_code=401, ) @@ -145,7 +131,6 @@ def test_completion_github_copilot_mock_response( def test_transform_messages_disable_copilot_system_to_assistant(monkeypatch): """Test that system messages are converted to assistant unless disable_copilot_system_to_assistant is True.""" - import litellm from litellm.llms.github_copilot.chat.transformation import GithubCopilotConfig # Save original value @@ -365,10 +350,6 @@ def test_x_initiator_header_system_only_messages(): assert headers["X-Initiator"] == "user" - - - - def test_copilot_vision_request_header_with_image(): """Test that Copilot-Vision-Request header is added when messages contain images""" config = GithubCopilotConfig() @@ -693,13 +674,8 @@ class TestGithubCopilotTransformResponse: assert result.choices[0].message.tool_calls is not None assert len(result.choices[0].message.tool_calls) == 1 assert result.choices[0].message.tool_calls[0]["id"] == "toolu_01ABC" - assert ( - result.choices[0].message.tool_calls[0]["function"]["name"] == "get_weather" - ) - assert ( - '"Boston, MA"' - in result.choices[0].message.tool_calls[0]["function"]["arguments"] - ) + assert result.choices[0].message.tool_calls[0]["function"]["name"] == "get_weather" + assert '"Boston, MA"' in result.choices[0].message.tool_calls[0]["function"]["arguments"] def test_transform_response_anthropic_native_multiple_text_blocks(self): """All text blocks must be concatenated, not only the first.""" @@ -867,12 +843,8 @@ class TestGithubCopilotTransformParsedResponseDict: @patch("litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client") -@patch( - "litellm.llms.openai.openai.OpenAIChatCompletion.make_sync_openai_chat_completion_request" -) -def test_openai_handler_repairs_github_copilot_empty_choices( - mock_request, mock_get_client -): +@patch("litellm.llms.openai.openai.OpenAIChatCompletion.make_sync_openai_chat_completion_request") +def test_openai_handler_repairs_github_copilot_empty_choices(mock_request, mock_get_client): """ The OpenAI SDK handler calls convert_to_model_response_object directly on the SDK's parsed output, bypassing transform_response. convert raises APIError on @@ -929,7 +901,10 @@ def test_validate_environment_uses_per_user_session_and_skips_authenticator(): session = GithubCopilotUserSession(token="user-copilot-token", api_base="https://api.githubcopilot.com") headers = config.validate_environment( - headers={}, model="github_copilot/gpt-4o", messages=[], optional_params={}, + headers={}, + model="github_copilot/gpt-4o", + messages=[], + optional_params={}, litellm_params={"github_copilot_user_session": session}, ) assert headers["Authorization"] == "Bearer user-copilot-token" @@ -946,9 +921,7 @@ def test_transform_response_carries_upstream_usage(): "object": "chat.completion", "created": 1, "model": "github_copilot/gpt-4o", - "choices": [ - {"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"} - ], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18}, }, ) diff --git a/tests/unit/llms/github_copilot/test_github_host.py b/tests/unit/llms/github_copilot/test_github_host.py index daf5bfaab80..cee02dc072f 100644 --- a/tests/unit/llms/github_copilot/test_github_host.py +++ b/tests/unit/llms/github_copilot/test_github_host.py @@ -32,28 +32,25 @@ def test_normalized_github_host_accepts_only_bare_supported_hosts(raw: object, e def test_github_com_endpoints_read_url_overrides_at_call_time(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://device.example/code") monkeypatch.setenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", "https://device.example/token") - monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", "https://api.example/copilot") monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", "custom-client") endpoints: Final = github_host_endpoints("github.com") assert endpoints.device_code_url == "https://device.example/code" assert endpoints.access_token_url == "https://device.example/token" - assert endpoints.api_key_url == "https://api.example/copilot" + assert endpoints.copilot_user_url == "https://api.github.com/copilot_internal/user" assert endpoints.client_id == "custom-client" def test_ghe_endpoints_ignore_github_com_url_overrides(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://device.example/code") monkeypatch.setenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", "https://device.example/token") - monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", "https://api.example/copilot") monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", "custom-client") endpoints: Final = github_host_endpoints("octo.ghe.com") assert endpoints.device_code_url == "https://octo.ghe.com/login/device/code" assert endpoints.access_token_url == "https://octo.ghe.com/login/oauth/access_token" - assert endpoints.api_key_url == "https://api.octo.ghe.com/copilot_internal/v2/token" assert endpoints.copilot_user_url == "https://api.octo.ghe.com/copilot_internal/user" assert endpoints.user_url == "https://api.octo.ghe.com/user" assert endpoints.default_copilot_api_base == "https://copilot-api.octo.ghe.com" diff --git a/tests/unit/llms/github_copilot/test_per_user_auth.py b/tests/unit/llms/github_copilot/test_per_user_auth.py index 5274f0c1cb8..1c3f32aa5f4 100644 --- a/tests/unit/llms/github_copilot/test_per_user_auth.py +++ b/tests/unit/llms/github_copilot/test_per_user_auth.py @@ -12,8 +12,9 @@ import httpx import pytest import litellm -from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST +from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST, GITHUB_COPILOT_USER_SESSION_TTL_SECONDS from litellm.exceptions import ( + APIConnectionError, AuthenticationError, BadRequestError, CallerCredentialAuthenticationError, @@ -21,27 +22,27 @@ from litellm.exceptions import ( ServiceUnavailableError, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.github_copilot.github_host import github_host_endpoints from litellm.llms.github_copilot.per_user_auth import ( _SESSION_CACHE, GITHUB_COPILOT_USER_SESSION_KWARG_KEY, GithubCopilotUserSession, _session_cache_key, aattach_github_copilot_user_session, - aexchange_github_token, + acheck_copilot_seat, apoll_device_flow, + aresolve_copilot_user_session, astart_device_flow, attach_github_copilot_user_session, evict_copilot_user_session, - exchange_github_token, - github_copilot_host_endpoints, github_copilot_auth_mode, github_copilot_user_session_from, + resolve_copilot_user_session, validated_copilot_api_base, ) -from litellm.llms.github_copilot.github_host import github_host_endpoints from litellm.types.utils import CredentialItem, ModelResponse -EXCHANGE_URL = "https://api.github.com/copilot_internal/v2/token" +COPILOT_USER_URL = "https://api.github.com/copilot_internal/user" DEVICE_CODE_URL = "https://github.com/login/device/code" ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token" _GITHUB_ENDPOINTS: Final = github_host_endpoints("github.com") @@ -62,10 +63,8 @@ def _credential( ) -def _exchange_payload(token="copilot-token", expires_at=None, api_base="https://api.githubcopilot.com"): +def _user_payload(api_base="https://api.githubcopilot.com"): return { - "token": token, - "expires_at": expires_at if expires_at is not None else int(time.time()) + 3600, "endpoints": {"api": api_base}, } @@ -92,10 +91,10 @@ def github_http(respond: Callable[[httpx.Request], httpx.Response]): yield requests -def _exchange_responder(payload=None, status=200): +def _user_responder(payload=None, status=200): def respond(request: httpx.Request) -> httpx.Response: - if str(request.url) == EXCHANGE_URL: - return httpx.Response(status, json=payload or _exchange_payload(), request=request) + if str(request.url) == COPILOT_USER_URL: + return httpx.Response(status, json=payload or _user_payload(), request=request) pytest.fail(f"unexpected request to {request.url}") return respond @@ -170,17 +169,17 @@ def test_validated_copilot_api_base(raw, expected): "https://copilot-api.octo.ghe.com:8443", ], ) -def test_ghe_exchange_response_must_stay_on_its_tenant_host(raw: object) -> None: +def test_ghe_user_response_must_stay_on_its_tenant_host(raw: object) -> None: endpoints: Final = github_host_endpoints("octo.ghe.com") assert validated_copilot_api_base(raw, endpoints) == "https://copilot-api.octo.ghe.com" -def test_github_com_exchange_rejects_a_ghe_copilot_host() -> None: +def test_github_com_user_response_rejects_a_ghe_copilot_host() -> None: endpoints: Final = github_host_endpoints("github.com") assert validated_copilot_api_base("https://copilot-api.octo.ghe.com", endpoints) == "https://api.githubcopilot.com" -def test_ghe_exchange_uses_stored_host_not_request_kwargs(): +def test_ghe_user_lookup_uses_stored_host_not_request_kwargs(): credential: Final = CredentialItem( credential_name="copilot-cred", credential_values={"github_copilot_auth_type": "per_user_oauth", "github_copilot_host": "octo.ghe.com"}, @@ -188,8 +187,8 @@ def test_ghe_exchange_uses_stored_host_not_request_kwargs(): ) def respond(request: httpx.Request) -> httpx.Response: - assert str(request.url) == "https://api.octo.ghe.com/copilot_internal/v2/token" - return httpx.Response(200, json=_exchange_payload(api_base="https://copilot-api.octo.ghe.com"), request=request) + assert str(request.url) == "https://api.octo.ghe.com/copilot_internal/user" + return httpx.Response(200, json=_user_payload(api_base="https://copilot-api.octo.ghe.com"), request=request) with patch.object(litellm, "credential_list", [credential]), github_http(respond) as requests: kwargs: Final = _kwargs_with_connection(github_host="octo.ghe.com") @@ -201,7 +200,7 @@ def test_ghe_exchange_uses_stored_host_not_request_kwargs(): assert len(requests) == 1 -def test_invalid_stored_host_fails_before_exchange(): +def test_invalid_stored_host_fails_before_user_lookup(): credential: Final = CredentialItem( credential_name="copilot-cred", credential_values={"github_copilot_auth_type": "per_user_oauth", "github_copilot_host": "https://octo.ghe.com"}, @@ -225,35 +224,36 @@ def test_session_repr_and_str_hide_the_token(caplog): assert "secret-copilot-token" not in caplog.text -def test_exchange_returns_session_and_caches_per_user(): - with github_http(_exchange_responder()) as requests: - s1 = exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) +def test_user_lookup_returns_github_token_and_caches_per_user(): + api_base: Final = "https://tenant.githubcopilot.com" + with github_http(_user_responder(payload=_user_payload(api_base=api_base))) as requests: + s1 = resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 1 + assert str(requests[0].url) == COPILOT_USER_URL assert requests[0].headers["Authorization"] == "token gh-token" - s2 = exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + s2 = resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 1 - assert s2.token == s1.token == "copilot-token" - # different user -> separate exchange - exchange_github_token("user-b", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + assert s2.token == s1.token == "gh-token" + assert s1.api_base == api_base + resolve_copilot_user_session("user-b", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 2 - # different github token -> separate exchange - exchange_github_token("user-a", "gh-token-2", "copilot-cred", _GITHUB_ENDPOINTS) + resolve_copilot_user_session("user-a", "gh-token-2", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 3 def test_cached_session_is_not_reused_after_github_token_changes(): - with github_http(_exchange_responder()) as requests: - exchange_github_token("user-a", "gh-token-a", "copilot-cred", _GITHUB_ENDPOINTS) - exchange_github_token("user-a", "gh-token-b", "copilot-cred", _GITHUB_ENDPOINTS) + with github_http(_user_responder()) as requests: + resolve_copilot_user_session("user-a", "gh-token-a", "copilot-cred", _GITHUB_ENDPOINTS) + resolve_copilot_user_session("user-a", "gh-token-b", "copilot-cred", _GITHUB_ENDPOINTS) assert [request.headers["Authorization"] for request in requests] == ["token gh-token-a", "token gh-token-b"] @pytest.mark.parametrize("status", [401, 403, 404]) -def test_exchange_auth_failures_raise_caller_auth_error_and_evict(status): - with github_http(_exchange_responder(status=status)): +def test_user_lookup_auth_failures_raise_caller_auth_error_and_evict(status): + with github_http(_user_responder(status=status)): with pytest.raises(CallerCredentialAuthenticationError) as exc: - exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert "Reconnect GitHub Copilot for credential 'copilot-cred' in the LiteLLM UI (LLM Credentials)" in str( exc.value ) @@ -262,52 +262,135 @@ def test_exchange_auth_failures_raise_caller_auth_error_and_evict(status): ) -def test_exchange_429_raises_caller_rate_limit(): - with github_http(_exchange_responder(status=429)): +def test_user_lookup_429_raises_caller_rate_limit(): + with github_http(_user_responder(status=429)): with pytest.raises(CallerCredentialRateLimitError): - exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) -def test_exchange_other_failure_raises_service_unavailable_without_body(): - with github_http(_exchange_responder(status=500, payload={"secret": "body"})): +def test_user_lookup_other_failure_raises_service_unavailable_without_body(): + with github_http(_user_responder(status=500, payload={"secret": "body"})): with pytest.raises(ServiceUnavailableError) as exc: - exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert "body" not in str(exc.value) -def test_exchange_expiry_margin_refetches_near_expiry(): - with github_http(_exchange_responder(_exchange_payload(expires_at=int(time.time()) + 30))) as requests: - exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) - # ttl under the 60s safety margin -> not cached, second call refetches - exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) +def test_session_cache_uses_fixed_ttl(): + with ( + patch("litellm.llms.github_copilot.per_user_auth.time.time", return_value=1000) as clock, + github_http(_user_responder()) as requests, + ): + resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + clock.return_value = 1000 + GITHUB_COPILOT_USER_SESSION_TTL_SECONDS - 1 + resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + assert len(requests) == 1 + clock.return_value = 1000 + GITHUB_COPILOT_USER_SESSION_TTL_SECONDS + 1 + resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 2 @pytest.mark.asyncio -async def test_aexchange_single_flight_one_http_call(): +async def test_acheck_copilot_seat_rejects_models_403(per_user_credential): def respond(request: httpx.Request) -> httpx.Response: - return httpx.Response(200, json=_exchange_payload(), request=request) + if str(request.url) == COPILOT_USER_URL: + return httpx.Response(200, json=_user_payload(), request=request) + if str(request.url) == "https://api.githubcopilot.com/models": + return httpx.Response(403, request=request) + pytest.fail(f"unexpected request to {request.url}") with github_http(respond) as requests: - s1, s2 = await asyncio.gather( - aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS), - aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS), - ) - assert len(requests) == 1 - assert s1.token == s2.token == "copilot-token" + with pytest.raises(CallerCredentialAuthenticationError): + await acheck_copilot_seat("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + + assert [str(request.url) for request in requests] == [ + COPILOT_USER_URL, + "https://api.githubcopilot.com/models", + ] + assert requests[1].headers["Authorization"] == "Bearer gh-token" @pytest.mark.asyncio -async def test_aexchange_401_raises_and_evicts(): - with github_http(_exchange_responder(status=401)): +@pytest.mark.parametrize( + ("status", "error_type"), + [(429, CallerCredentialRateLimitError), (500, ServiceUnavailableError)], +) +async def test_acheck_copilot_seat_maps_models_failures( + per_user_credential, + status: int, + error_type: type[Exception], +) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if str(request.url) == COPILOT_USER_URL: + return httpx.Response(200, json=_user_payload(), request=request) + if str(request.url) == "https://api.githubcopilot.com/models": + return httpx.Response(status, request=request) + pytest.fail(f"unexpected request to {request.url}") + + with github_http(respond): + with pytest.raises(error_type): + await acheck_copilot_seat("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + + +@pytest.mark.asyncio +async def test_acheck_copilot_seat_maps_models_request_errors(per_user_credential) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if str(request.url) == COPILOT_USER_URL: + return httpx.Response(200, json=_user_payload(), request=request) + if str(request.url) == "https://api.githubcopilot.com/models": + raise httpx.ConnectError("connection failed", request=request) + pytest.fail(f"unexpected request to {request.url}") + + with github_http(respond): + with pytest.raises(APIConnectionError, match="models check request failed"): + await acheck_copilot_seat("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + + +@pytest.mark.asyncio +async def test_acheck_copilot_seat_passes_when_models_are_usable(per_user_credential): + def respond(request: httpx.Request) -> httpx.Response: + if str(request.url) == COPILOT_USER_URL: + return httpx.Response(200, json=_user_payload(), request=request) + if str(request.url) == "https://api.githubcopilot.com/models": + return httpx.Response(200, json={"data": []}, request=request) + pytest.fail(f"unexpected request to {request.url}") + + with github_http(respond) as requests: + session = await acheck_copilot_seat("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + + assert session.token == "gh-token" + assert session.api_base == "https://api.githubcopilot.com" + assert [str(request.url) for request in requests] == [ + COPILOT_USER_URL, + "https://api.githubcopilot.com/models", + ] + + +@pytest.mark.asyncio +async def test_aresolve_single_flight_one_http_call(): + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=_user_payload(), request=request) + + with github_http(respond) as requests: + s1, s2 = await asyncio.gather( + aresolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS), + aresolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS), + ) + assert len(requests) == 1 + assert s1.token == s2.token == "gh-token" + + +@pytest.mark.asyncio +async def test_aresolve_401_raises_and_evicts(): + with github_http(_user_responder(status=401)): with pytest.raises(CallerCredentialAuthenticationError): - await aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + await aresolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert ( _SESSION_CACHE.get_cache(_session_cache_key("user-a", "copilot-cred", GITHUB_COPILOT_DEFAULT_HOST)) is None ) ACCOUNT_URL = "https://api.github.com/copilot_internal/user" +EXCHANGE_URL: Final = "https://api.github.com/copilot_internal/v2/token" _CUSTOM_CLIENT_ID: Final = "Ov23liCustomOrgApp" _ACCOUNT_API_BASE: Final = "https://api.business.githubcopilot.com" @@ -327,17 +410,20 @@ def _account_responder(payload=None, status=200, model_payload=None): return respond -@pytest.fixture -def custom_client_id(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", _CUSTOM_CLIENT_ID) +@pytest.fixture(params=[None, _CUSTOM_CLIENT_ID], ids=["default-client-id", "custom-client-id"]) +def copilot_client_id(request, monkeypatch: pytest.MonkeyPatch): + if request.param is None: + monkeypatch.delenv("GITHUB_COPILOT_CLIENT_ID", raising=False) + return + monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", request.param) @pytest.mark.asyncio -async def test_custom_client_id_uses_the_github_token_without_the_copilot_token_swap(custom_client_id): +async def test_user_lookup_uses_the_github_token_for_any_client_id(copilot_client_id): endpoints: Final = github_host_endpoints("github.com") with github_http(_account_responder()) as requests: - first = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints) - second = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints) + first = await aresolve_copilot_user_session("user-a", "gho_org_app", "copilot-cred", endpoints) + second = await aresolve_copilot_user_session("user-a", "gho_org_app", "copilot-cred", endpoints) assert [str(r.url) for r in requests] == [ACCOUNT_URL] assert requests[0].headers["Authorization"] == "token gho_org_app" @@ -353,22 +439,24 @@ async def test_custom_client_id_uses_the_github_token_without_the_copilot_token_ _account_payload(access_type_sku="no_access"), ], ) -def test_custom_client_id_without_a_copilot_seat_raises_reconnect(custom_client_id, payload): +def test_user_lookup_without_a_copilot_seat_raises_reconnect(copilot_client_id, payload): with github_http(_account_responder(payload)): with pytest.raises(CallerCredentialAuthenticationError): - exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) + resolve_copilot_user_session("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) @pytest.mark.parametrize("status", [401, 403, 404]) -def test_custom_client_id_rejected_account_lookup_raises_reconnect(custom_client_id, status): +def test_rejected_user_lookup_raises_reconnect(copilot_client_id, status): with github_http(_account_responder(status=status)): with pytest.raises(CallerCredentialAuthenticationError): - exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) + resolve_copilot_user_session("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) -def test_custom_client_id_account_api_base_must_stay_on_copilot_hosts(custom_client_id): +def test_user_lookup_api_base_must_stay_on_copilot_hosts(copilot_client_id): with github_http(_account_responder(_account_payload(api_base="https://evil.example"))): - session = exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")) + session = resolve_copilot_user_session( + "user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com") + ) assert session.api_base == "https://api.githubcopilot.com" @@ -392,16 +480,16 @@ def _kwargs_with_connection( @pytest.mark.asyncio async def test_aattach_injects_session_into_kwargs(per_user_credential): - with github_http(_exchange_responder()): + with github_http(_user_responder()): kwargs = _kwargs_with_connection() await aattach_github_copilot_user_session(kwargs) session = github_copilot_user_session_from(kwargs) assert session is not None - assert session.token == "copilot-token" + assert session.token == "gho_user_token" def test_attach_sync_injects_session(per_user_credential): - with github_http(_exchange_responder()): + with github_http(_user_responder()): kwargs = _kwargs_with_connection() attach_github_copilot_user_session(kwargs) assert isinstance(kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY], GithubCopilotUserSession) @@ -424,7 +512,7 @@ def test_attach_rejects_connection_from_a_different_github_host( ): with patch.object(litellm, "credential_list", [_credential(github_host=credential_host)]): kwargs: Final = _kwargs_with_connection(github_host=connection_host) - with github_http(lambda request: pytest.fail(f"unexpected token exchange at {request.url}")) as requests: + with github_http(lambda request: pytest.fail(f"unexpected user lookup at {request.url}")) as requests: with pytest.raises(CallerCredentialAuthenticationError) as exc: attach_github_copilot_user_session(kwargs) assert "Connect GitHub Copilot for credential 'copilot-cred'" in str(exc.value) @@ -432,11 +520,11 @@ def test_attach_rejects_connection_from_a_different_github_host( def test_legacy_connection_without_host_uses_github_com(per_user_credential): - with github_http(_exchange_responder()) as requests: + with github_http(_user_responder()) as requests: kwargs: Final = _kwargs_with_connection(github_host=None) attach_github_copilot_user_session(kwargs) assert len(requests) == 1 - assert str(requests[0].url) == EXCHANGE_URL + assert str(requests[0].url) == COPILOT_USER_URL def test_attach_shared_mode_does_nothing_without_connection(): @@ -447,10 +535,10 @@ def test_attach_shared_mode_does_nothing_without_connection(): def test_evict_drops_cached_session(): - with github_http(_exchange_responder()) as requests: - exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + with github_http(_user_responder()) as requests: + resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) evict_copilot_user_session("user-a", "copilot-cred", _GITHUB_ENDPOINTS.host) - exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) assert len(requests) == 2 @@ -559,7 +647,7 @@ def test_two_per_user_calls_with_identical_bodies_both_reach_upstream( # test-q monkeypatch.setattr(litellm, "cache", Cache()) with ( - github_http(_exchange_responder()), + github_http(_user_responder()), patch( "litellm.main.openai_chat_completions.completion", return_value=_chat_completion_response(), @@ -584,7 +672,7 @@ async def test_two_per_user_async_calls_with_identical_bodies_both_reach_upstrea monkeypatch.setattr(litellm, "cache", Cache()) with ( - github_http(_exchange_responder()), + github_http(_user_responder()), patch( "litellm.main.openai_chat_completions.acompletion", new=AsyncMock(return_value=_chat_completion_response()), @@ -605,27 +693,19 @@ def test_shared_mode_second_identical_call_hits_response_cache( # test-quality- ): """Control: shared device-login github_copilot keeps caching exactly as before, so a second identical call does not touch the upstream.""" - import json - from litellm.caching import Cache - # a mounted shared login: api-key.json with a far-future token, no device flow needed - (tmp_path / "api-key.json").write_text( - json.dumps( - { - "token": "shared-copilot-token", - "expires_at": 4102444800, - "endpoints": {"api": "https://api.githubcopilot.com"}, - } - ) - ) + (tmp_path / "access-token").write_text("shared-copilot-token") monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) monkeypatch.setattr(litellm, "cache", Cache()) - with patch( - "litellm.main.openai_chat_completions.completion", - return_value=_chat_completion_response(), - ) as mock_completion: + with ( + github_http(_user_responder()), + patch( + "litellm.main.openai_chat_completions.completion", + return_value=_chat_completion_response(), + ) as mock_completion, + ): call = { "model": "github_copilot/gpt-4o", "messages": [{"role": "user", "content": "hi"}], @@ -679,8 +759,8 @@ def _anthropic_messages_payload(): def _copilot_surface_responder(model_payload): def respond(request: httpx.Request) -> httpx.Response: - if str(request.url) == EXCHANGE_URL: - return httpx.Response(200, json=_exchange_payload(), request=request) + if str(request.url) == COPILOT_USER_URL: + return httpx.Response(200, json=_user_payload(), request=request) if request.url.host.endswith("githubcopilot.com"): return httpx.Response(200, json=model_payload, request=request) pytest.fail(f"unexpected request to {request.url}") @@ -696,19 +776,19 @@ def _no_authenticator(): ) -def _assert_exchange_then_model_call(requests, github_token, model_host="https://api.githubcopilot.com"): - exchange = [r for r in requests if str(r.url) == EXCHANGE_URL] +def _assert_user_lookup_then_model_call(requests, github_token, model_host="https://api.githubcopilot.com"): + lookup = [r for r in requests if str(r.url) == COPILOT_USER_URL] upstream = [r for r in requests if r.url.host.endswith("githubcopilot.com")] - assert exchange, "expected the GitHub token exchange request" + assert lookup, "expected the GitHub user lookup request" assert upstream, "expected the model request to the Copilot host" - assert exchange[0].headers["Authorization"] == f"token {github_token}" - assert upstream[0].headers["Authorization"].startswith("Bearer ") + assert lookup[0].headers["Authorization"] == f"token {github_token}" + assert upstream[0].headers["Authorization"] == f"Bearer {github_token}" assert upstream[0].url.host == httpx.URL(model_host).host def test_e2e_completion_per_user_credential(per_user_credential): with ( - github_http(_exchange_responder()) as requests, + github_http(_user_responder()) as requests, _no_authenticator(), patch( "litellm.main.openai_chat_completions.completion", @@ -721,13 +801,13 @@ def test_e2e_completion_per_user_credential(per_user_credential): litellm_credential_name="copilot-cred", secret_fields=_secret_fields("user-a", "gho_a"), ) - assert len(requests) == 1 and str(requests[0].url) == EXCHANGE_URL + assert len(requests) == 1 and str(requests[0].url) == COPILOT_USER_URL assert requests[0].headers["Authorization"] == "token gho_a" call_kwargs = mock_completion.call_args.kwargs - assert call_kwargs["api_key"] == "copilot-token" + assert call_kwargs["api_key"] == "gho_a" assert call_kwargs["api_base"] == "https://api.githubcopilot.com" extra_headers = call_kwargs["optional_params"]["extra_headers"] - assert extra_headers["Authorization"] == "Bearer copilot-token" + assert extra_headers["Authorization"] == "Bearer gho_a" assert response.usage.total_tokens > 0 @@ -735,7 +815,7 @@ def test_e2e_completion_per_user_session_token_wins_over_caller_authorization(pe """The session token reaches the wire even when the caller passes their own Authorization in extra_headers (chat dispatches through _complete_custom_openai).""" with ( - github_http(_exchange_responder()), + github_http(_user_responder()), _no_authenticator(), patch( "litellm.main.openai_chat_completions.completion", @@ -750,17 +830,17 @@ def test_e2e_completion_per_user_session_token_wins_over_caller_authorization(pe extra_headers={"Authorization": "Bearer caller-token", "x-custom": "keep"}, ) extra_headers = mock_completion.call_args.kwargs["optional_params"]["extra_headers"] - assert extra_headers["Authorization"] == "Bearer copilot-token" + assert extra_headers["Authorization"] == "Bearer gho_a" assert extra_headers["x-custom"] == "keep" _EVIL_AUTH_HEADERS = {"authorization": "Bearer caller-evil", "x-keep": "1"} -def _assert_session_token_on_the_wire(requests): +def _assert_session_token_on_the_wire(requests, github_token="gho_a"): upstream = [r for r in requests if r.url.host.endswith("githubcopilot.com")] assert upstream, "expected the model request to the Copilot host" - assert upstream[0].headers["authorization"] == "Bearer copilot-token" + assert upstream[0].headers["authorization"] == f"Bearer {github_token}" def test_e2e_completion_per_user_caller_authorization_cannot_override_session(per_user_credential): @@ -769,7 +849,7 @@ def test_e2e_completion_per_user_caller_authorization_cannot_override_session(pe handed to openai_chat_completions.completion.""" caller_headers: dict[str, str] = dict(_EVIL_AUTH_HEADERS) with ( - github_http(_exchange_responder()), + github_http(_user_responder()), _no_authenticator(), patch( "litellm.main.openai_chat_completions.completion", @@ -784,7 +864,7 @@ def test_e2e_completion_per_user_caller_authorization_cannot_override_session(pe extra_headers=caller_headers, ) extra_headers = mock_completion.call_args.kwargs["optional_params"]["extra_headers"] - assert extra_headers["Authorization"] == "Bearer copilot-token" + assert extra_headers["Authorization"] == "Bearer gho_a" assert "authorization" not in extra_headers assert extra_headers["x-keep"] == "1" assert caller_headers == dict(_EVIL_AUTH_HEADERS), "the caller's dict must not be mutated" @@ -805,9 +885,7 @@ def test_e2e_embedding_per_user_caller_authorization_cannot_override_session(per _assert_session_token_on_the_wire(requests) -def test_e2e_custom_client_id_sends_the_github_token_to_the_accounts_copilot_host( - per_user_credential, custom_client_id -): +def test_e2e_per_user_sends_the_github_token_to_the_accounts_copilot_host(per_user_credential, copilot_client_id): with ( github_http(_account_responder(model_payload=_embedding_response_payload())) as requests, _no_authenticator(), @@ -869,7 +947,7 @@ def test_e2e_embedding_per_user_credential(per_user_credential): litellm_credential_name="copilot-cred", secret_fields=_secret_fields("user-a", "gho_a"), ) - _assert_exchange_then_model_call(requests, "gho_a") + _assert_user_lookup_then_model_call(requests, "gho_a") assert requests[-1].url.path == "/embeddings" assert response.usage.prompt_tokens == 4 @@ -885,7 +963,7 @@ def test_e2e_responses_per_user_credential(per_user_credential): litellm_credential_name="copilot-cred", secret_fields=_secret_fields("user-a", "gho_a"), ) - _assert_exchange_then_model_call(requests, "gho_a") + _assert_user_lookup_then_model_call(requests, "gho_a") assert requests[-1].url.path == "/responses" assert response.usage.total_tokens == 6 @@ -903,7 +981,7 @@ async def test_e2e_anthropic_messages_per_user_credential(per_user_credential): litellm_credential_name="copilot-cred", secret_fields=_secret_fields("user-a", "gho_a"), ) - _assert_exchange_then_model_call(requests, "gho_a") + _assert_user_lookup_then_model_call(requests, "gho_a") assert requests[-1].url.path == "/v1/messages" assert response["usage"]["input_tokens"] == 4 @@ -954,26 +1032,19 @@ async def test_e2e_anthropic_messages_not_connected_raises_connect_401(per_user_ def _shared_login(tmp_path, monkeypatch): - import json - - (tmp_path / "api-key.json").write_text( - json.dumps( - { - "token": "shared-copilot-token", - "expires_at": 4102444800, - "endpoints": {"api": "https://api.githubcopilot.com"}, - } - ) - ) + (tmp_path / "access-token").write_text("shared-copilot-token") monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) def test_e2e_completion_shared_mode_still_uses_authenticator(tmp_path, monkeypatch): _shared_login(tmp_path, monkeypatch) - with patch( - "litellm.main.openai_chat_completions.completion", - return_value=_chat_completion_response(), - ) as mock_completion: + with ( + github_http(_user_responder()), + patch( + "litellm.main.openai_chat_completions.completion", + return_value=_chat_completion_response(), + ) as mock_completion, + ): response = litellm.completion( model="github_copilot/gpt-4o", messages=[{"role": "user", "content": "hi"}], @@ -1007,45 +1078,37 @@ def test_e2e_responses_shared_mode_still_uses_authenticator(tmp_path, monkeypatc @pytest.mark.asyncio -async def test_e2e_anthropic_messages_shared_mode_still_uses_authenticator(tmp_path, monkeypatch): - # The shared Authenticator's access-token refresh refuses to run inside an event - # loop, so on this async surface the mounted-login path can't reach it; patch the - # shared-login edge to emulate a mounted api-key.json. +async def test_e2e_anthropic_messages_shared_mode_still_uses_authenticator(): from litellm.llms.github_copilot.authenticator import Authenticator with ( github_http(_copilot_surface_responder(_anthropic_messages_payload())) as requests, - patch.object( # test-quality-ok: patching the shared-login edge is the fixture under test - Authenticator, "get_api_key", return_value="shared-copilot-token" - ) as mock_api_key, - patch.object( # test-quality-ok: same edge; api_key read is mocked so base must be too - Authenticator, "get_api_base", return_value="https://api.githubcopilot.com" - ), + patch.object(Authenticator, "get_api_key", return_value="shared-copilot-token"), + patch.object(Authenticator, "get_api_base", return_value="https://api.githubcopilot.com"), ): response = await litellm.anthropic.messages.acreate( model="github_copilot/claude-haiku-4.5", max_tokens=16, messages=[{"role": "user", "content": "hi"}], ) - mock_api_key.assert_called() upstream = [r for r in requests if r.url.host.endswith("githubcopilot.com")] assert upstream and upstream[0].headers["Authorization"] == "Bearer shared-copilot-token" assert response["usage"]["input_tokens"] == 4 -def test_sync_exchange_single_flight_one_http_call(): +def test_sync_resolve_single_flight_one_http_call(): import threading def slow_respond(request: httpx.Request) -> httpx.Response: time.sleep(0.05) - return httpx.Response(200, json=_exchange_payload(), request=request) + return httpx.Response(200, json=_user_payload(), request=request) sessions: list[GithubCopilotUserSession] = [] with github_http(slow_respond) as requests: threads = [ threading.Thread( target=lambda: sessions.append( - exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) + resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS) ) ) for _ in range(6) @@ -1055,13 +1118,13 @@ def test_sync_exchange_single_flight_one_http_call(): for t in threads: t.join() assert len(requests) == 1 - assert len(sessions) == 6 and all(s.token == "copilot-token" for s in sessions) + assert len(sessions) == 6 and all(s.token == "gh-token" for s in sessions) def test_two_users_same_github_token_do_not_share_session(): - with github_http(_exchange_responder()): - s1: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS) - s2: Final = exchange_github_token("user-b", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS) + with github_http(_user_responder()): + s1: Final = resolve_copilot_user_session("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS) + s2: Final = resolve_copilot_user_session("user-b", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS) user_a_key: Final = _session_cache_key("user-a", "copilot-cred", _GITHUB_ENDPOINTS.host) user_b_key: Final = _session_cache_key("user-b", "copilot-cred", _GITHUB_ENDPOINTS.host) other_credential_key: Final = _session_cache_key("user-a", "another-credential", _GITHUB_ENDPOINTS.host) @@ -1069,7 +1132,7 @@ def test_two_users_same_github_token_do_not_share_session(): assert user_a_key != user_b_key assert user_a_key != other_credential_key assert user_a_key != other_host_key - assert s1.token == s2.token == "copilot-token" + assert s1.token == s2.token == "gh-same" def test_same_user_and_token_do_not_share_sessions_across_hosts(): @@ -1078,11 +1141,11 @@ def test_same_user_and_token_do_not_share_sessions_across_hosts(): ghe_endpoints: Final = github_host_endpoints("octo.ghe.com") def respond(request: httpx.Request) -> httpx.Response: - return httpx.Response(200, json=_exchange_payload(), request=request) + return httpx.Response(200, json=_user_payload(), request=request) with github_http(respond) as requests: - github_session: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS) - ghe_session: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", ghe_endpoints) + github_session: Final = resolve_copilot_user_session("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS) + ghe_session: Final = resolve_copilot_user_session("user-a", "gh-same", "copilot-cred", ghe_endpoints) assert len(requests) == 2 assert github_session.api_base == "https://api.githubcopilot.com" assert ghe_session.api_base == "https://copilot-api.octo.ghe.com" @@ -1141,33 +1204,37 @@ def test_provider_info_shared_github_copilot_still_uses_authenticator(): @pytest.mark.asyncio -async def test_aexchange_cancelled_waiter_does_not_cancel_the_shared_exchange(per_user_credential): - """Two waiters share one in-flight exchange; cancelling one must not cancel the +async def test_aresolve_cancelled_waiter_does_not_cancel_the_shared_session(per_user_credential): + """Two waiters share one in-flight resolution; cancelling one must not cancel the underlying future the other is waiting on.""" gate = asyncio.Event() class _GatedClient: async def get(self, url, headers=None): await gate.wait() - return httpx.Response(200, json=_exchange_payload(), request=httpx.Request("GET", url)) + return httpx.Response(200, json=_user_payload(), request=httpx.Request("GET", url)) with patch( # test-quality-ok: the gate holds the HTTP edge open so two waiters overlap "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=_GatedClient() ): - first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS)) + first = asyncio.ensure_future( + aresolve_copilot_user_session("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS) + ) await asyncio.sleep(0) - second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS)) + second = asyncio.ensure_future( + aresolve_copilot_user_session("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS) + ) await asyncio.sleep(0) second.cancel() gate.set() session = await first - assert session.token == "copilot-token" and session.api_base == "https://api.githubcopilot.com" + assert session.token == "gh-shared" and session.api_base == "https://api.githubcopilot.com" with pytest.raises(asyncio.CancelledError): await second @pytest.mark.asyncio -async def test_aexchange_does_not_share_an_in_flight_session_after_token_changes(per_user_credential): +async def test_aresolve_does_not_share_an_in_flight_session_after_token_changes(per_user_credential): gate = asyncio.Event() auth_headers: asyncio.Queue[str] = asyncio.Queue() @@ -1175,26 +1242,29 @@ async def test_aexchange_does_not_share_an_in_flight_session_after_token_changes async def get(self, url, headers=None): await auth_headers.put(headers["authorization"]) await gate.wait() - token = "copilot-token-a" if headers["authorization"] == "token gh-a" else "copilot-token-b" return httpx.Response( 200, - json=_exchange_payload(token=token), + json=_user_payload(), request=httpx.Request("GET", url), ) with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=_GatedClient()): - first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-a", "copilot-cred", _GITHUB_ENDPOINTS)) + first = asyncio.ensure_future( + aresolve_copilot_user_session("user-a", "gh-a", "copilot-cred", _GITHUB_ENDPOINTS) + ) assert await asyncio.wait_for(auth_headers.get(), timeout=1) == "token gh-a" - second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-b", "copilot-cred", _GITHUB_ENDPOINTS)) + second = asyncio.ensure_future( + aresolve_copilot_user_session("user-a", "gh-b", "copilot-cred", _GITHUB_ENDPOINTS) + ) assert await asyncio.wait_for(auth_headers.get(), timeout=1) == "token gh-b" gate.set() sessions = await asyncio.gather(first, second) - assert [session.token for session in sessions] == ["copilot-token-a", "copilot-token-b"] + assert [session.token for session in sessions] == ["gh-a", "gh-b"] @pytest.mark.asyncio -async def test_aexchange_failure_with_a_single_caller_emits_no_unretrieved_warning(per_user_credential): +async def test_aresolve_failure_with_a_single_caller_emits_no_unretrieved_warning(per_user_credential): """The leader removes the failed future before anyone awaits it; the stored exception must be consumed so the loop does not log 'never retrieved'.""" import litellm.llms.github_copilot.per_user_auth as pua @@ -1205,10 +1275,10 @@ async def test_aexchange_failure_with_a_single_caller_emits_no_unretrieved_warni loop.set_exception_handler(lambda _loop, context: captured.append((_loop, context))) try: with ( - github_http(_exchange_responder(status=401)), + github_http(_user_responder(status=401)), pytest.raises(CallerCredentialAuthenticationError), ): - await aexchange_github_token("user-a", "gh-fail", "copilot-cred", _GITHUB_ENDPOINTS) + await aresolve_copilot_user_session("user-a", "gh-fail", "copilot-cred", _GITHUB_ENDPOINTS) _IN_FLIGHT_KEYS = list(pua._IN_FLIGHT) import gc @@ -1240,7 +1310,7 @@ async def test_per_user_streaming_call_writes_nothing_to_the_response_cache(per_ "model": "github_copilot/gpt-4o", "messages": [{"role": "user", "content": "hi"}], GITHUB_COPILOT_USER_SESSION_KWARG_KEY: GithubCopilotUserSession( - token="copilot-token", api_base="https://api.githubcopilot.com" + token="gh-token", api_base="https://api.githubcopilot.com" ), } await handler.async_get_cache( @@ -1302,24 +1372,30 @@ def _per_user_router(): ) +def _route_with_credential_name(router: litellm.Router, credential_name: str) -> ModelResponse: + return router.completion( + model="copilot", + messages=[{"role": "user", "content": "hi"}], + litellm_credential_name=credential_name, + mock_response="ok", + ) + + def test_per_user_deployment_rejects_a_caller_credential_override(per_user_credential): """Caller kwargs win the litellm_params merge downstream, so the deployment layer is the last place that still sees both names: a different name on a per-user deployment must be refused there even if discovery missed the hop.""" router: Final = _per_user_router() - deployment: Final = router.get_model_list()[0] - kwargs: Final = {"litellm_credential_name": "shared-cred", "metadata": {}} - with pytest.raises(litellm.BadRequestError): - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + with pytest.raises(BadRequestError, match="litellm_credential_name cannot be overridden"): + _route_with_credential_name(router, "shared-cred") def test_per_user_deployment_allows_the_configured_credential_name(per_user_credential): + """The configured name passes the deployment layer and reaches the per-user login, which asks the + caller to connect GitHub Copilot since no user token is stored.""" router: Final = _per_user_router() - deployment: Final = router.get_model_list()[0] - kwargs: Final = {"litellm_credential_name": "copilot-cred", "metadata": {}} - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - assert kwargs["litellm_credential_name"] == "copilot-cred" - assert "model_info" in kwargs + with pytest.raises(CallerCredentialAuthenticationError, match="copilot-cred"): + _route_with_credential_name(router, "copilot-cred") def test_shared_deployment_keeps_caller_credential_override(): @@ -1330,8 +1406,5 @@ def test_shared_deployment_keeps_caller_credential_override(): patch.object(Authenticator, "get_api_key", return_value="shared-token"), ): router: Final = _per_user_router() - deployment: Final = router.get_model_list()[0] - kwargs: Final = {"litellm_credential_name": "other-cred", "metadata": {}} - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - assert kwargs["litellm_credential_name"] == "other-cred" - assert "model_info" in kwargs + response: Final = _route_with_credential_name(router, "other-cred") + assert response.choices[0].message.content == "ok" diff --git a/tests/unit/llms/openai/live/test_handler.py b/tests/unit/llms/openai/live/test_handler.py index bb7cda99e79..63df7a0a1cc 100644 --- a/tests/unit/llms/openai/live/test_handler.py +++ b/tests/unit/llms/openai/live/test_handler.py @@ -1254,7 +1254,7 @@ def _upstream_failing_a_client_send() -> _FakeUpstream: ids=["session-closed", "upstream-1011-after-a-frame", "upstream-1006-after-a-frame", "upstream-send-failed"], ) async def test_live_session_counts_a_router_success_only_when_it_ended_cleanly( - upstream_factory: Callable[[], _FakeUpstream], counted_successes: int + upstream_factory: Callable[[], _FakeUpstream], counted_successes: int, monkeypatch: pytest.MonkeyPatch ) -> None: router: Final = Router( model_list=[ @@ -1276,10 +1276,10 @@ async def test_live_session_counts_a_router_success_only_when_it_ended_cleanly( api_key="live-test-key", ) + monkeypatch.setattr(litellm, "_arealtime", _arealtime) await asyncio.wait_for( - router._ageneric_api_call_with_fallbacks( + router._arealtime( # the routed realtime endpoint the proxy's websocket route calls model="live-voice", - original_function=_arealtime, litellm_logging_obj=Logging( model="gpt-live-1", messages=[], diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index f61cfba9a10..53122246de4 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -1,6 +1,8 @@ import asyncio +import contextlib import gzip import importlib +import inspect import json import os from collections.abc import Awaitable, Callable, Mapping, Sequence @@ -227,6 +229,52 @@ async def test_acompletion_returns_prompt_cache_details_over_injected_transport( assert prompt_tokens_details.cached_tokens == 7 +_SDK_READS_OPENAI_ADMIN_KEY: Final = "admin_api_key" in inspect.signature(AsyncOpenAI.__init__).parameters + + +@pytest.mark.skipif(not _SDK_READS_OPENAI_ADMIN_KEY, reason="openai<3 has no OPENAI_ADMIN_KEY fallback to leak") +@pytest.mark.asyncio +async def test_acompletion_never_sends_the_openai_admin_key_for_a_deployment_without_a_key( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setenv("OPENAI_ADMIN_KEY", "sk-admin-only") + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + outbound: Final = asyncio.Queue() + + def respond(request: httpx.Request) -> httpx.Response: + outbound.put_nowait(request.headers.get("authorization")) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-no-admin-key", + "object": "chat.completion", + "created": 0, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + monkeypatch.setattr(litellm, "aclient_session", http_client) + with contextlib.suppress(litellm.AuthenticationError, litellm.InternalServerError): + await asyncio.wait_for( + litellm.acompletion( + model="openai/gpt-5.4-mini", + api_key="", + messages=[{"role": "user", "content": "who signs this request"}], + num_retries=0, + max_retries=0, + ), + timeout=10, + ) + + sent: Final = tuple(outbound.get_nowait() for _ in range(outbound.qsize())) + assert all(authorization is None for authorization in sent), sent + + class _SearchingChatGate(CustomLogger): def __init__(self) -> None: diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 98e204dd44e..842500a7b27 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -2161,7 +2161,7 @@ class TestVertexClaudeBatch: monkeypatch.setattr(litellm, "drop_params", False) entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5) - with pytest.raises(VertexAIError, match=r"line 1 .*top_p=0.5.*drop_params") as exc_info: + with pytest.raises(VertexAIError, match=r"line 1 .*top_p.*drop_params") as exc_info: _uploaded_rows(config, {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9"}, [entry]) assert exc_info.value.status_code == 400 diff --git a/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py index ea2fe42a01f..7e5b3a7abd7 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -80,7 +80,7 @@ def _make_mock_prisma(): return record async def _update(where, data): - record = store[where["name"]] + record = store[where["name"]] if "name" in where else next(r for r in store.values() if r.id == where["id"]) for k, v in data.items(): setattr(record, k, v) return record @@ -91,9 +91,20 @@ def _make_mock_prisma(): mock_table.update = AsyncMock(side_effect=_update) mock_client.db.litellm_claudecodeplugintable = mock_table mock_client.writer_db = mock_client.db + mock_client.db.query_raw = AsyncMock(return_value=[]) + mock_client.db.tx.return_value.__aenter__.return_value = mock_client.db + mock_client.db.litellm_claudecodepluginversion.count = AsyncMock(return_value=0) + mock_client.db.litellm_claudecodepluginversion.create = AsyncMock() + mock_client.db.litellm_claudecodepluginupload.find_unique = AsyncMock(return_value=None) + mock_client.db.litellm_claudecodepluginupload.delete_many = AsyncMock(return_value=0) + mock_client.db.litellm_claudecodepluginupload.upsert = AsyncMock() return mock_client +def _uploads() -> PluginUploadRepository: + return PluginUploadRepository(litellm.proxy.proxy_server.prisma_client) + + _USER = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key=MASTER_KEY, @@ -131,7 +142,7 @@ async def test_register_plugin_git_subdir_success(): """git-subdir with both url and path fields registers successfully.""" request = RegisterPluginRequest(name="my-monorepo-plugin", source=_GIT_SUBDIR_SOURCE) - response = await register_plugin(request=request, user_api_key_dict=_USER) + response = await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER) assert response.status == "success" assert response.action == "created" @@ -149,7 +160,7 @@ async def _read_stored_manifest(name: str) -> dict: async def test_register_plugin_duplicate_name_conflicts(): """A second POST with an existing name returns 409 and leaves the stored plugin untouched.""" name = "my-monorepo-plugin" - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), user_api_key_dict=_USER, ) @@ -167,7 +178,7 @@ async def test_register_plugin_duplicate_name_conflicts(): version="2.0.0", ) with pytest.raises(HTTPException) as exc_info: - await register_plugin(request=conflicting, user_api_key_dict=_USER) + await register_plugin(uploads=_uploads(), request=conflicting, user_api_key_dict=_USER) assert exc_info.value.status_code == 409 assert "already exists" in exc_info.value.detail["error"] @@ -182,13 +193,13 @@ async def test_register_plugin_duplicate_name_conflicts(): async def test_update_plugin_replaces_existing_source(): """PUT updates an existing plugin: action=updated and the stored source is replaced.""" name = "my-monorepo-plugin" - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), user_api_key_dict=_USER, ) new_source = {"source": "github", "repo": "org/replacement"} - response = await update_plugin( + response = await update_plugin(uploads=_uploads(), plugin_name=name, request=UpdatePluginRequest(source=new_source, version="2.0.0", description="updated"), user_api_key_dict=_USER, @@ -208,13 +219,13 @@ async def test_update_plugin_replaces_existing_source(): async def test_metadata_edits_preserve_uploaded_source_and_reject_a_changed_http_destination() -> None: table: Final = litellm.proxy.proxy_server.prisma_client.writer_db.litellm_claudecodeplugintable await table.create(data={"name": _UPLOAD_NAME, "manifest_json": json.dumps({"source": _UPLOAD_SOURCE})}) - request: Final = UpdatePluginRequest(source=_UPLOAD_SOURCE, description="Updated description") - response: Final = await update_plugin(plugin_name=_UPLOAD_NAME, request=request, user_api_key_dict=_USER) - expected: Final = {"name": _UPLOAD_NAME, "source": _UPLOAD_SOURCE, "description": "Updated description"} + request: Final = UpdatePluginRequest(source=_UPLOAD_SOURCE, description="Updated description", version="1.0.0") + response: Final = await update_plugin(uploads=_uploads(), plugin_name=_UPLOAD_NAME, request=request, user_api_key_dict=_USER) + expected: Final = {"name": _UPLOAD_NAME, "source": _UPLOAD_SOURCE, "description": "Updated description", "version": "1.0.0"} assert response.plugin.description == "Updated description" assert await _read_stored_manifest(_UPLOAD_NAME) == expected with pytest.raises(HTTPException) as error: - await update_plugin( + await update_plugin(uploads=_uploads(), plugin_name=_UPLOAD_NAME, user_api_key_dict=_USER, request=UpdatePluginRequest(source={**_UPLOAD_SOURCE, "url": "http://other.example/skill.zip"}), ) @@ -227,7 +238,7 @@ async def test_metadata_edits_preserve_uploaded_source_and_reject_a_changed_http async def test_update_plugin_not_found(): """PUT on a name that does not exist raises HTTP 404.""" with pytest.raises(HTTPException) as exc_info: - await update_plugin( + await update_plugin(uploads=_uploads(), plugin_name="does-not-exist", request=UpdatePluginRequest(source=_GIT_SUBDIR_SOURCE), user_api_key_dict=_USER, @@ -246,7 +257,7 @@ async def test_register_plugin_create_race_maps_unique_violation_to_409(): table.create = AsyncMock(side_effect=UniqueViolationError({}, message="duplicate name")) with pytest.raises(HTTPException) as exc_info: - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="racy-plugin", source=_GIT_SUBDIR_SOURCE), user_api_key_dict=_USER, ) @@ -262,7 +273,7 @@ async def test_update_plugin_db_error_maps_to_structured_500(): from prisma.errors import PrismaError name = "my-monorepo-plugin" - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), user_api_key_dict=_USER, ) @@ -271,7 +282,7 @@ async def test_update_plugin_db_error_maps_to_structured_500(): table.update = AsyncMock(side_effect=PrismaError("connection lost")) with pytest.raises(HTTPException) as exc_info: - await update_plugin( + await update_plugin(uploads=_uploads(), plugin_name=name, request=UpdatePluginRequest(source={"source": "github", "repo": "org/replacement"}), user_api_key_dict=_USER, @@ -286,16 +297,17 @@ async def test_update_plugin_deleted_mid_update_returns_404(): """A concurrent delete between the find_unique pre-check and the update makes prisma's update return None; that must surface the same 404 as a plain miss, not an AttributeError.""" name = "my-monorepo-plugin" - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), user_api_key_dict=_USER, ) table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable - table.update = AsyncMock(return_value=None) + existing: Final = await table.find_unique(where={"name": name}) + table.find_unique = AsyncMock(side_effect=[existing, None]) with pytest.raises(HTTPException) as exc_info: - await update_plugin( + await update_plugin(uploads=_uploads(), plugin_name=name, request=UpdatePluginRequest(source={"source": "github", "repo": "org/replacement"}), user_api_key_dict=_USER, @@ -307,7 +319,7 @@ async def test_update_plugin_deleted_mid_update_returns_404(): @pytest.mark.asyncio async def test_get_marketplace_skips_plugin_with_null_manifest(): - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="good-plugin", source=_GIT_SUBDIR_SOURCE, version="1.0.0"), user_api_key_dict=_USER, ) @@ -333,7 +345,7 @@ def _granted_user(skills: list[str]) -> UserAPIKeyAuth: async def _register_public_and_private_plugins() -> None: for name, enabled in (("public-skill", True), ("private-skill", False)): - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), user_api_key_dict=_USER, ) @@ -401,7 +413,7 @@ async def test_get_marketplace_key_query_param_adds_granted_disabled_plugins(mon @pytest.mark.asyncio @pytest.mark.parametrize("preference", ["available", "auto_install", "required"]) async def test_get_marketplace_emits_installation_preference(preference): - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest( name="s3-skill", source=_ARCHIVE_SOURCE, @@ -434,7 +446,7 @@ def test_plugin_request_rejects_unknown_installation_preference(build_request): @pytest.mark.asyncio async def test_get_marketplace_omits_installation_preference_when_unset(): - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="manual-install-plugin", source=_GIT_SUBDIR_SOURCE), user_api_key_dict=_USER, ) @@ -447,12 +459,12 @@ async def test_get_marketplace_omits_installation_preference_when_unset(): @pytest.mark.asyncio async def test_update_plugin_propagates_installation_preference_to_get_and_list(): name = "updated-install-plugin" - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_ARCHIVE_SOURCE), user_api_key_dict=_USER, ) - response = await update_plugin( + response = await update_plugin(uploads=_uploads(), plugin_name=name, request=UpdatePluginRequest( source=_ARCHIVE_SOURCE, @@ -472,12 +484,12 @@ async def test_update_plugin_propagates_installation_preference_to_get_and_list( @pytest.mark.asyncio async def test_update_plugin_without_installation_preference_clears_it(): name = "cleared-install-plugin" - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_ARCHIVE_SOURCE, installation_preference="auto_install"), user_api_key_dict=_USER, ) - await update_plugin( + await update_plugin(uploads=_uploads(), plugin_name=name, request=UpdatePluginRequest(source=_ARCHIVE_SOURCE), user_api_key_dict=_USER, @@ -518,7 +530,7 @@ async def test_register_accepts_the_camel_case_installation_preference_key(): {"name": "s3-skill", "source": _ARCHIVE_SOURCE, "installationPreference": "auto_install"} ) - response = await register_plugin(request=request, user_api_key_dict=_USER) + response = await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER) marketplace = json.loads((await get_marketplace(request=MagicMock())).body) assert response.plugin.installation_preference == "auto_install" @@ -536,13 +548,13 @@ _UNPINNED_ARCHIVE_SOURCE = {"source": "archive", "url": "https://artifacts.examp async def _assert_preference_refused(source, preference): with pytest.raises(HTTPException) as register_error: - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="no-op-plugin", source=source, installation_preference=preference), user_api_key_dict=_USER, ) - await register_plugin(request=RegisterPluginRequest(name="no-op-plugin", source=source), user_api_key_dict=_USER) + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="no-op-plugin", source=source), user_api_key_dict=_USER) with pytest.raises(HTTPException) as update_error: - await update_plugin( + await update_plugin(uploads=_uploads(), plugin_name="no-op-plugin", request=UpdatePluginRequest(source=source, installation_preference=preference), user_api_key_dict=_USER, @@ -568,7 +580,7 @@ async def test_self_installing_preference_needs_the_archive_sha256(preference): @pytest.mark.asyncio async def test_available_preference_is_accepted_on_an_unpinned_archive(): - response = await register_plugin( + response = await register_plugin(uploads=_uploads(), request=RegisterPluginRequest( name="listed-plugin", source=_UNPINNED_ARCHIVE_SOURCE, installation_preference="available" ), @@ -626,7 +638,7 @@ async def test_register_plugin_git_subdir_missing_url(): ) with pytest.raises(HTTPException) as exc_info: - await register_plugin(request=request, user_api_key_dict=_USER) + await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER) assert exc_info.value.status_code == 400 assert "url" in exc_info.value.detail["error"] @@ -641,7 +653,7 @@ async def test_register_plugin_git_subdir_empty_url(): ) with pytest.raises(HTTPException) as exc_info: - await register_plugin(request=request, user_api_key_dict=_USER) + await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER) assert exc_info.value.status_code == 400 assert "url" in exc_info.value.detail["error"] @@ -656,7 +668,7 @@ async def test_register_plugin_git_subdir_missing_path(): ) with pytest.raises(HTTPException) as exc_info: - await register_plugin(request=request, user_api_key_dict=_USER) + await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER) assert exc_info.value.status_code == 400 assert "path" in exc_info.value.detail["error"] @@ -675,7 +687,7 @@ async def test_register_plugin_git_subdir_empty_path(): ) with pytest.raises(HTTPException) as exc_info: - await register_plugin(request=request, user_api_key_dict=_USER) + await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER) assert exc_info.value.status_code == 400 assert "path" in exc_info.value.detail["error"] @@ -703,7 +715,7 @@ async def test_register_plugin_git_subdir_path_traversal(): ) with pytest.raises(HTTPException) as exc_info: - await register_plugin(request=request, user_api_key_dict=_USER) + await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER) assert exc_info.value.status_code == 400 assert "relative" in exc_info.value.detail["error"] @@ -718,7 +730,7 @@ async def test_register_plugin_unknown_source_type(): ) with pytest.raises(HTTPException) as exc_info: - await register_plugin(request=request, user_api_key_dict=_USER) + await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER) assert exc_info.value.status_code == 400 assert "git-subdir" in exc_info.value.detail["error"] @@ -727,7 +739,7 @@ async def test_register_plugin_unknown_source_type(): @pytest.mark.asyncio async def test_archive_source_registers_and_is_served_verbatim_in_marketplace(): - response = await register_plugin( + response = await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="s3-skill", source=_ARCHIVE_SOURCE), user_api_key_dict=_USER, ) @@ -743,7 +755,7 @@ async def test_archive_source_registers_and_is_served_verbatim_in_marketplace(): async def test_archive_source_without_sha256_is_accepted(): source = {"source": "archive", "url": "https://artifacts.example.com/plugin.zip"} - response = await register_plugin( + response = await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="unpinned-skill", source=source), user_api_key_dict=_USER, ) @@ -754,12 +766,12 @@ async def test_archive_source_without_sha256_is_accepted(): @pytest.mark.asyncio async def test_update_plugin_to_archive_source(): name = "my-monorepo-plugin" - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE), user_api_key_dict=_USER, ) - response = await update_plugin( + response = await update_plugin(uploads=_uploads(), plugin_name=name, request=UpdatePluginRequest(source=_ARCHIVE_SOURCE), user_api_key_dict=_USER, @@ -787,7 +799,7 @@ async def test_update_plugin_to_archive_source(): ) async def test_register_plugin_archive_rejects_malformed_source(source, expected_fragment): with pytest.raises(HTTPException) as exc_info: - await register_plugin(request=RegisterPluginRequest(name="bad-plugin", source=source), user_api_key_dict=_USER) + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="bad-plugin", source=source), user_api_key_dict=_USER) assert exc_info.value.status_code == 400 assert expected_fragment in exc_info.value.detail["error"] @@ -799,7 +811,7 @@ async def test_register_plugin_rejects_non_admin(): request = RegisterPluginRequest(name="attacker-plugin", source=_GIT_SUBDIR_SOURCE) with pytest.raises(HTTPException) as exc_info: - await register_plugin(request=request, user_api_key_dict=_NON_ADMIN_USER) + await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_NON_ADMIN_USER) assert exc_info.value.status_code == 403 @@ -811,14 +823,14 @@ async def test_register_plugin_rejects_non_admin(): async def test_update_plugin_rejects_non_admin_overwrite(): """A non-admin key cannot overwrite an existing plugin's source.""" name = "trusted-plugin" - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), user_api_key_dict=_USER, ) malicious_source = {"source": "github", "repo": "attacker/malicious-repo"} with pytest.raises(HTTPException) as exc_info: - await update_plugin( + await update_plugin(uploads=_uploads(), plugin_name=name, request=UpdatePluginRequest(source=malicious_source), user_api_key_dict=_NON_ADMIN_USER, @@ -834,7 +846,7 @@ async def test_update_plugin_rejects_non_admin_overwrite(): async def test_enable_disable_delete_plugin_reject_non_admin(): """Non-admin keys cannot enable, disable, or delete catalog entries.""" name = "trusted-plugin-2" - await register_plugin( + await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), user_api_key_dict=_USER, ) @@ -843,6 +855,7 @@ async def test_enable_disable_delete_plugin_reject_non_admin(): enable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), disable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), delete_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + claude_code_marketplace.get_plugin_downloads(plugin_name=name, uploads=None, user_api_key_dict=_NON_ADMIN_USER), ): with pytest.raises(HTTPException) as exc_info: await coro @@ -869,7 +882,7 @@ async def test_register_plugin(): user_id="test-user", ) - response: Final = await register_plugin( + response: Final = await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=user_api_key_dict, ) @@ -904,7 +917,7 @@ async def test_get_marketplace(): user_id="test-user", ) - await register_plugin( + await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=user_api_key_dict, ) @@ -946,7 +959,7 @@ async def test_register_plugin_git_subdir(): user_id="test-user", ) - response: Final = await register_plugin( + response: Final = await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=user_api_key_dict, ) @@ -1211,14 +1224,20 @@ _UPLOAD_SOURCE: Final = {"source": "archive", "url": f"http://test{_UPLOAD_PATH} async def upload_storage(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_UploadClient]: prisma: Final = litellm.proxy.proxy_server.prisma_client prisma.writer_db = MagicMock() + prisma.writer_db.execute_raw = AsyncMock(return_value=1) plugins: Final = prisma.writer_db.litellm_claudecodeplugintable archives: Final = prisma.writer_db.litellm_claudecodepluginupload + prisma.writer_db.litellm_claudecodeplugindownload.create = AsyncMock(return_value=None) plugins.create = AsyncMock(return_value=prisma_models.LiteLLM_ClaudeCodePluginTable( - id="uploaded-id", name=_UPLOAD_NAME, enabled=True, + id="uploaded-id", name=_UPLOAD_NAME, enabled=True, anonymous_download_count=0, )) plugins.find_unique = AsyncMock(return_value=prisma_models.LiteLLM_ClaudeCodePluginTable( - id="uploaded-id", name=_UPLOAD_NAME, enabled=True, manifest_json=json.dumps({"source": _UPLOAD_SOURCE}), + id="uploaded-id", name=_UPLOAD_NAME, enabled=True, anonymous_download_count=0, manifest_json=json.dumps({"source": _UPLOAD_SOURCE}), )) + prisma.writer_db.query_raw = AsyncMock(return_value=[]) + prisma.writer_db.tx.return_value.__aenter__.return_value = prisma.writer_db + prisma.writer_db.litellm_claudecodepluginversion.create = AsyncMock() + archives.create = AsyncMock() archives.find_first = AsyncMock(return_value=prisma_models.LiteLLM_ClaudeCodePluginUpload( plugin_name=_UPLOAD_NAME, content=Base64.encode(b"stored archive"), )) @@ -1252,7 +1271,8 @@ async def _post_skill_upload( async def test_skill_upload_creates_metadata_and_bytes_atomically_and_rejects_duplicates( upload_storage: _UploadClient, ) -> None: - client, plugins, _ = upload_storage + client, plugins, archives = upload_storage + plugins.find_unique = AsyncMock(return_value=None) buffer: Final = io.BytesIO() with zipfile.ZipFile(buffer, "w") as archive: archive.writestr("SKILL.md", "Read the personal instructions") @@ -1261,7 +1281,7 @@ async def test_skill_upload_creates_metadata_and_bytes_atomically_and_rejects_du assert response.status_code == 200, response.text plugins.create.assert_awaited_once() stored: Final = plugins.create.await_args.kwargs["data"] - content: Final = stored["uploaded_archive"]["create"]["content"].decode() + content: Final = archives.create.await_args.kwargs["data"]["content"].decode() digest: Final = hashlib.sha256(content).hexdigest() source: Final = { "source": "archive", @@ -1272,7 +1292,7 @@ async def test_skill_upload_creates_metadata_and_bytes_atomically_and_rejects_du assert stored == { "name": _UPLOAD_NAME, "description": "My notes", "version": "2.0.0", "manifest_json": stored["manifest_json"], "created_by": litellm.proxy.proxy_server.litellm_proxy_admin_name, - "uploaded_archive": {"create": {"content": Base64.encode(content)}}, + "files_json": "{}", "updated_at": stored["updated_at"], } assert response.json() == {"status": "success", "action": "created", "plugin": { "id": "uploaded-id", "name": _UPLOAD_NAME, "version": "2.0.0", "description": "My notes", @@ -1318,7 +1338,7 @@ async def test_archive_reads_only_authorized_current_bytes( client, plugins, archives = upload_storage manifest: Final = json.dumps({"source": source}) plugins.find_unique.return_value = None if source is None else prisma_models.LiteLLM_ClaudeCodePluginTable( - id="uploaded-id", name=_UPLOAD_NAME, enabled=enabled, manifest_json=manifest, + id="uploaded-id", name=_UPLOAD_NAME, enabled=enabled, anonymous_download_count=0, manifest_json=manifest, ) if not blob: archives.find_first.return_value = None @@ -1327,9 +1347,16 @@ async def test_archive_reads_only_authorized_current_bytes( assert (response.content == b"stored archive") is (status == 200) plugins.find_unique.assert_awaited_once_with(where={"name": _UPLOAD_NAME}) assert archives.find_first.await_args_list == ([call(where={ - "plugin_name": _UPLOAD_NAME, "plugin": {"is": {"manifest_json": manifest, "enabled": enabled}}, + "plugin_name": _UPLOAD_NAME, "plugin": {"is": { + "id": "uploaded-id", "manifest_json": manifest, "enabled": enabled, + }}, })] if reads else []) litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable.find_unique.assert_not_awaited() + events: Final = litellm.proxy.proxy_server.prisma_client.writer_db.litellm_claudecodeplugindownload + assert events.create.await_count == int(status == 200 and key is not None) + assert litellm.proxy.proxy_server.prisma_client.writer_db.execute_raw.await_count == int( + status == 200 and key is None + ) @pytest.mark.asyncio @@ -1359,3 +1386,59 @@ async def test_archive_requires_header_credentials_for_private_content( assert response.status_code == status, response.text assert (response.content == b"stored archive") is (status == 200) assert archives.find_first.await_count == int(status == 200) + events: Final = litellm.proxy.proxy_server.prisma_client.writer_db.litellm_claudecodeplugindownload + assert events.create.await_count == int(status == 200) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key, query, exists, status", [ + (_LIMITED_KEY, {}, True, 403), + (None, {}, True, 401), + (MASTER_KEY, {"cursor": "9223372036854775808"}, True, 422), + (MASTER_KEY, {"cursor": "0"}, True, 422), + (MASTER_KEY, {"limit": 101}, True, 422), + (MASTER_KEY, {}, False, 404), +]) +async def test_download_history_rejects_unauthorized_invalid_and_unknown_requests( + upload_storage: _UploadClient, key: str | None, query: Mapping[str, str | int], exists: bool, status: int, +) -> None: + client, plugins, _ = upload_storage + if not exists: + plugins.find_unique.return_value = None + response: Final = await client.get( + f"/claude-code/plugins/{_UPLOAD_NAME}/downloads", + headers={"Authorization": f"Bearer {key}"} if key else {}, params=query, + ) + assert response.status_code == status, response.text + events: Final = litellm.proxy.proxy_server.prisma_client.writer_db.litellm_claudecodeplugindownload + events.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_published_updates_auto_bump_and_reject_stale_editors() -> None: + from datetime import datetime, timezone + + source: Final = {"source": "github", "repo": "example/skill"} + created: Final = await register_plugin( + uploads=_uploads(), request=RegisterPluginRequest(name="versioned", source=source, version=None), + user_api_key_dict=_USER, + ) + updated: Final = await update_plugin( + uploads=_uploads(), plugin_name="versioned", request=UpdatePluginRequest(source=source), + user_api_key_dict=_USER, + ) + assert (created.plugin.version, updated.plugin.version) == ("1.0.0", "1.0.1") + before: Final = await _read_stored_manifest("versioned") + with pytest.raises(HTTPException) as stale: + await update_plugin( + uploads=_uploads(), plugin_name="versioned", user_api_key_dict=_USER, + request=UpdatePluginRequest(source=source, expected_updated_at=datetime(2000, 1, 1, tzinfo=timezone.utc)), + ) + assert stale.value.status_code == 409 + with pytest.raises(HTTPException) as error: + await update_plugin( + uploads=_uploads(), plugin_name="versioned", user_api_key_dict=_USER, + request=UpdatePluginRequest(source=source, version="1.0.0", description="rejected edit"), + ) + assert error.value.status_code == 400 + assert await _read_stored_manifest("versioned") == before diff --git a/tests/unit/proxy/auth/test_fallback_model_access.py b/tests/unit/proxy/auth/test_fallback_model_access.py index 5c6a98d9e42..bcda2c05966 100644 --- a/tests/unit/proxy/auth/test_fallback_model_access.py +++ b/tests/unit/proxy/auth/test_fallback_model_access.py @@ -55,10 +55,9 @@ def _broken_access_group_lookup(*args: object, **kwargs: object) -> NoReturn: @pytest.mark.asyncio -@pytest.mark.usefixtures("python_router_core") # the Rust core keeps this lookup on its host registry async def test_is_model_authorized_for_token_fails_closed_when_the_lookup_breaks(monkeypatch: pytest.MonkeyPatch): router = _router() - monkeypatch.setattr(router, "_core_get_model_access_groups", _broken_access_group_lookup) + monkeypatch.setattr(type(router.registry()), "access_groups", _broken_access_group_lookup) assert ( await is_model_authorized_for_token( diff --git a/tests/unit/proxy/auth/test_model_access_group_budgets.py b/tests/unit/proxy/auth/test_model_access_group_budgets.py index d1b381907fa..6e2e9576fc8 100644 --- a/tests/unit/proxy/auth/test_model_access_group_budgets.py +++ b/tests/unit/proxy/auth/test_model_access_group_budgets.py @@ -222,11 +222,15 @@ async def test_group_without_a_budget_is_not_attributed(): @pytest.mark.asyncio -@pytest.mark.usefixtures("python_router_core") # the Rust core keeps this lookup on its host registry async def test_empty_registry_skips_the_access_group_matching_entirely(monkeypatch: pytest.MonkeyPatch): router = Router(model_list=MODEL_LIST) - access_group_lookups = MagicMock(wraps=router._core_get_model_access_groups) - monkeypatch.setattr(router, "_core_get_model_access_groups", access_group_lookups) + registry_type = type(router.registry()) + access_group_lookups = MagicMock(wraps=registry_type.access_groups) + monkeypatch.setattr( + registry_type, + "access_groups", + lambda registry, *args, **kwargs: access_group_lookups(registry, *args, **kwargs), + ) assert await _matched(key_models=["tier-a"], cache=await _cache(budgeted_groups=()), llm_router=router) == () assert access_group_lookups.call_count == 0 @@ -259,12 +263,11 @@ def _broken_access_group_lookup(*args: object, **kwargs: object) -> NoReturn: @pytest.mark.asyncio -@pytest.mark.usefixtures("python_router_core") # the Rust core keeps this lookup on its host registry async def test_stamp_does_not_break_auth_when_the_access_group_lookup_fails(monkeypatch: pytest.MonkeyPatch): cache = await _cache() valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a"], team_id=TEAM_ID, user_id=USER_ID) router = Router(model_list=MODEL_LIST) - monkeypatch.setattr(router, "_core_get_model_access_groups", _broken_access_group_lookup) + monkeypatch.setattr(type(router.registry()), "access_groups", _broken_access_group_lookup) await stamp_matched_model_access_groups( model="gpt-4o", diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index efb4451f2fc..c1011310af3 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -363,18 +363,24 @@ def test_proxy_admin_viewer_allowed_management_reads(allowed_read_route): ) -def test_virtual_key_allowed_routes_with_litellm_routes_member_name_allowed(): +@pytest.mark.parametrize("group, route", ( + ("openai_routes", "/chat/completions"), + ("skill_hub_routes", "/skill-hub/skills/release-notes/downloads"), + ("management_routes", "/skill-hub/skills/release-notes/downloads"), + ("anthropic_routes", "/claude-code/plugins/release-notes/downloads"), +)) +def test_virtual_key_allowed_routes_with_litellm_routes_member_name_allowed(group: str, route: str) -> None: """Test that virtual key is allowed to call routes when allowed_routes contains LiteLLMRoutes member name""" # Create a UserAPIKeyAuth with allowed_routes containing a LiteLLMRoutes member name valid_token = UserAPIKeyAuth( user_id="test_user", - allowed_routes=["openai_routes"], # This is a member name in LiteLLMRoutes enum + allowed_routes=[group], ) # Test that a route from the openai_routes group is allowed result = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/chat/completions", # This is in LiteLLMRoutes.openai_routes.value + route=route, valid_token=valid_token, ) @@ -5102,3 +5108,25 @@ def test_every_dashboard_role_can_send_telemetry_ui_events( ) except HTTPException as exc: pytest.fail(f"{user_role.value} was denied {method} {route}: {exc.detail}") + + +@pytest.mark.parametrize("route", [ + "/skill-hub/skills/{name}/versions", + "/skill-hub/skills/{name}/versions/{version_id}", + "/skill-hub/skills/{name}/versions/{version_id}/archive", +]) +def test_member_can_reach_skill_history_visibility_checks(route: str) -> None: + request: Final = MagicMock(spec=Request) + request.method = "GET" + request.query_params = {} + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=LiteLLM_UserTable(user_id="reader", user_role=LitellmUserRoles.INTERNAL_USER.value), + _user_role=LitellmUserRoles.INTERNAL_USER, + route=route, + request=request, + valid_token=UserAPIKeyAuth(user_id="reader", user_role=LitellmUserRoles.INTERNAL_USER), + request_data={}, + ) + except HTTPException as exc: + pytest.fail(f"History visibility could not run: {exc.detail}") diff --git a/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py b/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py index 36277de047d..b9a0d583475 100644 --- a/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py +++ b/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py @@ -153,7 +153,7 @@ class TestUnmappedModelBudgetEnforcement: Regression test for the stale-cache bug surfaced in PR review: upgrading an explicitly free deployment to paid via ``upsert_deployment`` (same deployment count, same router instance) must invalidate the - cached ``_is_model_cost_zero=True`` answer so budget checks resume + cached zero-cost answer so budget checks resume immediately — not after the next proxy restart. """ from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -178,7 +178,6 @@ class TestUnmappedModelBudgetEnforcement: ) # Warm the cache as zero-cost. assert is_model_cost_zero(model="ramping-model", llm_router=router) is True - assert router._zero_cost_cache.get("ramping-model") is True # In-place pricing update: same deployment count, same router id, # same model name. The pre-fix cache key was @@ -200,9 +199,7 @@ class TestUnmappedModelBudgetEnforcement: ) ) - # Cache must have been cleared by ``_invalidate_model_group_info_cache``. - assert router._zero_cost_cache == {} - # Subsequent call sees the new pricing and enforces budget. + # The next call sees the new pricing and enforces budget. assert is_model_cost_zero(model="ramping-model", llm_router=router) is False def test_strategy_router_alias_with_zero_pricing_enforces_budget(self): diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 92114b00398..eee5d96511b 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3324,9 +3324,13 @@ class TestMultiPodBatchCostClaim: @pytest.mark.parametrize("failure_phase", ("before_write", "after_write", "external")) @pytest.mark.parametrize("collector_enabled", (False, True)) async def test_builtin_accounting_retries_only_before_a_possible_charge( - self, failure_phase: str, collector_enabled: bool + self, failure_phase: str, collector_enabled: bool, monkeypatch: pytest.MonkeyPatch ) -> None: import base64 + + from litellm.proxy import proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False) from datetime import datetime import httpx diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index cdd18fd4eef..27fffc7e56e 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -2069,7 +2069,10 @@ from litellm.proxy.credential_endpoints import user_provider_credentials as upc _DEVICE_CODE_URL = "https://github.com/login/device/code" _ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token" -_COPILOT_TOKEN_URL = "https://api.github.com/copilot_internal/v2/token" +_COPILOT_USER_URL = "https://api.github.com/copilot_internal/user" +_COPILOT_API_BASE = "https://api.githubcopilot.com" +_COPILOT_MODELS_URL = f"{_COPILOT_API_BASE}/models" +_COPILOT_USER_RESPONSE = {"endpoints": {"api": _COPILOT_API_BASE}} _GITHUB_USER_URL = "https://api.github.com/user" @@ -2245,7 +2248,8 @@ def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch): { _DEVICE_CODE_URL: _DEVICE_FLOW_START, _ACCESS_TOKEN_URL: {"access_token": "gho_1"}, - _COPILOT_TOKEN_URL: {"token": "copilot-tok", "expires_at": 4102444800}, + _COPILOT_USER_URL: _COPILOT_USER_RESPONSE, + _COPILOT_MODELS_URL: {}, _GITHUB_USER_URL: {"login": "octo"}, } ): @@ -2276,19 +2280,18 @@ def test_poll_host_change_race_requires_reconnect_on_the_next_request(monkeypatc table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")]) - def change_host_after_token_exchange(request: httpx.Request) -> dict[str, str]: + def change_host_after_access_token(request: httpx.Request) -> dict[str, str]: monkeypatch.setattr(litellm, "credential_list", [_per_user_credential()]) return {"access_token": "gho_1"} with _github_http( { "https://octo.ghe.com/login/device/code": _DEVICE_FLOW_START, - "https://octo.ghe.com/login/oauth/access_token": change_host_after_token_exchange, - "https://api.octo.ghe.com/copilot_internal/v2/token": { - "token": "copilot-tok", - "expires_at": 4102444800, - "endpoints": {"api": "https://copilot-api.octo.ghe.com"}, + "https://octo.ghe.com/login/oauth/access_token": change_host_after_access_token, + "https://api.octo.ghe.com/copilot_internal/user": { + "endpoints": {"api": "https://copilot-api.octo.ghe.com"} }, + "https://copilot-api.octo.ghe.com/models": {}, "https://api.octo.ghe.com/user": {"login": "octo"}, } ): @@ -2318,7 +2321,7 @@ def test_poll_host_change_race_requires_reconnect_on_the_next_request(monkeypatc }, } module_level_client: Final = MagicMock() - module_level_client.get.side_effect = AssertionError("unexpected token exchange") + module_level_client.get.side_effect = AssertionError("unexpected GitHub user lookup") with _github_http({}), patch.object(litellm, "module_level_client", module_level_client): with pytest.raises(litellm.CallerCredentialAuthenticationError) as exc: attach_github_copilot_user_session(kwargs) @@ -2326,17 +2329,16 @@ def test_poll_host_change_race_requires_reconnect_on_the_next_request(monkeypatc assert stored_connection.github_host == "octo.ghe.com" -def test_ghe_user_connection_uses_tenant_urls_for_device_flow_exchange_and_login(monkeypatch): +def test_ghe_user_connection_uses_tenant_urls_for_device_flow_and_user_lookup(monkeypatch): table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")]) with _github_http( { "https://octo.ghe.com/login/device/code": _DEVICE_FLOW_START, "https://octo.ghe.com/login/oauth/access_token": {"access_token": "gho_1"}, - "https://api.octo.ghe.com/copilot_internal/v2/token": { - "token": "copilot-tok", - "expires_at": 4102444800, - "endpoints": {"api": "https://copilot-api.octo.ghe.com"}, + "https://api.octo.ghe.com/copilot_internal/user": { + "endpoints": {"api": "https://copilot-api.octo.ghe.com"} }, + "https://copilot-api.octo.ghe.com/models": {}, "https://api.octo.ghe.com/user": {"login": "octo"}, } ): @@ -2471,7 +2473,9 @@ def test_user_connection_poll_reports_no_copilot_seat(monkeypatch, seat_status): return httpx.Response(200, json=_DEVICE_FLOW_START, request=request) if url == _ACCESS_TOKEN_URL: return httpx.Response(200, json={"access_token": "gho_seat"}, request=request) - if url == _COPILOT_TOKEN_URL: + if url == _COPILOT_USER_URL: + return httpx.Response(200, json=_COPILOT_USER_RESPONSE, request=request) + if url == _COPILOT_MODELS_URL: return httpx.Response(seat_status, json={}, request=request) pytest.fail(f"unexpected request to {url}") @@ -2499,7 +2503,7 @@ def test_user_connection_poll_rate_limit_returns_429(monkeypatch): return httpx.Response(200, json=_DEVICE_FLOW_START, request=request) if url == _ACCESS_TOKEN_URL: return httpx.Response(200, json={"access_token": "gho_limited"}, request=request) - if url == _COPILOT_TOKEN_URL: + if url == _COPILOT_USER_URL: return httpx.Response(429, json={}, request=request) pytest.fail(f"unexpected request to {url}") @@ -2875,7 +2879,8 @@ def test_user_connection_poll_valid_handle_works_with_a_fresh_cache(monkeypatch) { _DEVICE_CODE_URL: _DEVICE_FLOW_START, _ACCESS_TOKEN_URL: {"access_token": "gho_1"}, - _COPILOT_TOKEN_URL: {"token": "copilot-tok", "expires_at": 4102444800}, + _COPILOT_USER_URL: _COPILOT_USER_RESPONSE, + _COPILOT_MODELS_URL: {}, _GITHUB_USER_URL: {"login": "octo"}, } ): @@ -2959,7 +2964,8 @@ def test_user_connection_poll_reports_503_when_the_cache_overwrite_cannot_land(m { _DEVICE_CODE_URL: _DEVICE_FLOW_START, _ACCESS_TOKEN_URL: {"access_token": "gho_1"}, - _COPILOT_TOKEN_URL: {"token": "copilot-tok", "expires_at": 4102444800}, + _COPILOT_USER_URL: _COPILOT_USER_RESPONSE, + _COPILOT_MODELS_URL: {}, _GITHUB_USER_URL: {"login": "octo"}, } ): diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index d90c7f6d21d..e2debe05c03 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -44,6 +44,15 @@ from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpush from tests.unit.proxy.db.fake_prisma_engine import engine_call + +@pytest.fixture(autouse=True) +def legacy_spend_writers(monkeypatch: pytest.MonkeyPatch) -> None: + """These tests cover the legacy writers, still reachable as the rollback path with use_spend_logs_ledger off.""" + from litellm.proxy import proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False) + + @pytest.mark.asyncio async def test_daily_spend_tracking_with_disabled_spend_logs(): """ @@ -3714,11 +3723,12 @@ async def test_failed_per_entity_increment_from_redis_restores_only_what_may_sti proxy_logging_obj=proxy_logging_obj, ) - mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once() - restored = mock_redis_update_buffer.restore_transactions_to_redis.call_args.kwargs["db_spend_update_transactions"] - assert restored["user_list_transactions"] is None - assert restored["team_list_transactions"] == {"team-1": 1.5} - assert restored["key_list_transactions"] == ({"key-1": 1.5} if safe_to_resend else None) + restores = mock_redis_update_buffer.restore_transactions_to_redis.await_args_list + assert [ + {name: rows for name, rows in call.kwargs["db_spend_update_transactions"].items() if rows} + for call in restores + ] == ([{"key_list_transactions": {"key-1": 1.5}}] if safe_to_resend else []) + mock_batcher.litellm_teamtable.update_many.assert_called_once() mock_batcher.litellm_usertable.update_many.assert_called_once() @@ -3892,13 +3902,13 @@ async def test_commit_spend_updates_to_db_reports_each_completed_table(): ) assert on_table_committed.call_args_list == [ - call("user_list_transactions"), - call("end_user_list_transactions"), - call("key_list_transactions"), - call("team_list_transactions"), - call("team_member_list_transactions"), call("org_list_transactions"), call("org_member_list_transactions"), + call("team_list_transactions"), + call("user_list_transactions"), + call("end_user_list_transactions"), + call("team_member_list_transactions"), + call("key_list_transactions"), call("project_list_transactions"), call("tag_list_transactions"), call("model_access_group_list_transactions"), @@ -5808,7 +5818,9 @@ async def test_completed_spend_survives_error_log_suppression( prisma.get_request_status.return_value = status monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "disable_spend_logs", False) - monkeypatch.setattr(proxy_server, "general_settings", {"disable_error_logs": disable_error_logs}) + monkeypatch.setattr( + proxy_server, "general_settings", {"disable_error_logs": disable_error_logs, "use_spend_logs_ledger": False} + ) now: Final = datetime(2026, 10, 20, 12, tzinfo=timezone.utc) result: Final = ModelResponse( model="gpt-4o-mini", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) diff --git a/tests/unit/proxy/discovery_endpoints/test_agent_skills_endpoints.py b/tests/unit/proxy/discovery_endpoints/test_agent_skills_endpoints.py index 89c31cbdf1d..16eeedab0be 100644 --- a/tests/unit/proxy/discovery_endpoints/test_agent_skills_endpoints.py +++ b/tests/unit/proxy/discovery_endpoints/test_agent_skills_endpoints.py @@ -97,7 +97,7 @@ def hub_record( digest: Final = hashlib.sha256(wrapper).hexdigest() return ( prisma_models.LiteLLM_ClaudeCodePluginTable( - id=f"hub-{name}", name=name, enabled=True, manifest_json=json.dumps({ + id=f"hub-{name}", name=name, enabled=True, anonymous_download_count=0, manifest_json=json.dumps({ "source": { "source": "archive", "sha256": digest, "url": f"http://testserver/claude-code/plugins/{name}/archive/{digest}.zip", @@ -112,6 +112,7 @@ def hub_record( def hub_repository(*names: str) -> PluginUploadRepository: records: Final = tuple(hub_record(name) for name in names) prisma: Final = MagicMock() + prisma.writer_db.execute_raw = AsyncMock(return_value=1) plugins: Final = prisma.writer_db.litellm_claudecodeplugintable plugins.find_many = AsyncMock(return_value=[record for record, _ in records]) diff --git a/tests/unit/proxy/discovery_endpoints/test_skill_hub_endpoints.py b/tests/unit/proxy/discovery_endpoints/test_skill_hub_endpoints.py index 143d26a3ccb..8127fc80bab 100644 --- a/tests/unit/proxy/discovery_endpoints/test_skill_hub_endpoints.py +++ b/tests/unit/proxy/discovery_endpoints/test_skill_hub_endpoints.py @@ -17,13 +17,13 @@ from unittest.mock import AsyncMock, MagicMock import httpx import pytest from prisma import models as prisma_models -from prisma.errors import UniqueViolationError +from prisma.errors import ForeignKeyViolationError, UniqueViolationError import litellm from litellm.caching.caching import DualCache from litellm.proxy import proxy_server from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth -from litellm.proxy.deps.repositories import get_plugin_upload_repository +from litellm.proxy.deps.repositories import get_plugin_upload_repository, get_skill_submission_repository from litellm.proxy.discovery_endpoints.agent_skills_archive import ( MAX_ARCHIVE_ENTRIES, MAX_ARCHIVE_UNPACKED_BYTES, @@ -32,7 +32,9 @@ from litellm.proxy.discovery_endpoints.skill_hub_endpoints import import_http_cl from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache from litellm.repositories.plugin_upload_repository import PluginUploadRepository +from litellm.repositories.skill_submission_repository import SkillSubmissionRepository from tests._master_key import MASTER_KEY +from tests.unit.proxy.skill_submission_rows import SubmissionRows NAME: Final = "portable-notes" MANIFEST: Final = ( @@ -48,6 +50,9 @@ HUB: Final = "/skill-hub/skills" INDEX: Final = "/.well-known/agent-skills/index.json" GRANTED_KEY: Final = "sk-hub-granted" PLAIN_KEY: Final = "sk-hub-plain" +MEMBER_KEY: Final = "sk-hub-member" +OTHER_KEY: Final = "sk-hub-other" +VIEWER_KEY: Final = "sk-hub-viewer" def _zip(files: Mapping[str, bytes], prefix: str = "") -> bytes: @@ -70,22 +75,29 @@ class HubClient: client: httpx.AsyncClient plugins: MagicMock archives: MagicMock + downloads: MagicMock + anonymous_counter: AsyncMock upstream: MagicMock clock: MagicMock + database: MagicMock @pytest.fixture async def hub(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[HubClient]: prisma: Final = MagicMock() prisma.writer_db = prisma.db + prisma.db.execute_raw = AsyncMock(return_value=1) prisma.get_data = AsyncMock(return_value=None) prisma.db.litellm_skillstable.find_many = AsyncMock(return_value=[]) prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) plugins: Final = prisma.db.litellm_claudecodeplugintable archives: Final = prisma.db.litellm_claudecodepluginupload - plugins.create = AsyncMock(return_value=prisma_models.LiteLLM_ClaudeCodePluginTable( - id="hub-id", name=NAME, enabled=True, - )) + async def create_row(data: Mapping[str, object]) -> prisma_models.LiteLLM_ClaudeCodePluginTable: + return prisma_models.LiteLLM_ClaudeCodePluginTable.model_validate({"id": "hub-id", "enabled": True, "anonymous_download_count": 0, **data}) + + plugins.create = AsyncMock(side_effect=create_row) + downloads: Final = prisma.db.litellm_claudecodeplugindownload + downloads.create = AsyncMock() plugins.find_unique = AsyncMock(return_value=None) plugins.find_many = AsyncMock(return_value=[]) @@ -98,6 +110,17 @@ async def hub(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[HubClient]: plugins.update = AsyncMock(side_effect=update_row) plugins.delete = AsyncMock() archives.find_first = AsyncMock(return_value=None) + archives.find_unique = AsyncMock(return_value=None) + archives.create = AsyncMock() + archives.upsert = AsyncMock() + archives.delete_many = AsyncMock() + prisma.db.tx.return_value.__aenter__.return_value = prisma.db + submissions: Final = SubmissionRows() + prisma.db.litellm_skillsubmission = submissions + prisma.db.query_raw = AsyncMock(return_value=[]) + prisma.db.litellm_claudecodepluginversion.create = AsyncMock() + prisma.db.litellm_claudecodepluginversion.count = AsyncMock(return_value=0) + prisma.db.litellm_claudecodepluginversion.find_first = AsyncMock(return_value=None) upstream: Final = MagicMock(return_value=httpx.Response(404)) clock: Final = MagicMock(return_value=datetime.now()) limiter: Final = PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(DualCache()), time_provider=clock) @@ -118,22 +141,34 @@ async def hub(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[HubClient]: object_permission=LiteLLM_ObjectPermissionTable(object_permission_id=key, skills=skills), ), ) + for key, role in ( + (MEMBER_KEY, proxy_server.LitellmUserRoles.INTERNAL_USER), + (OTHER_KEY, proxy_server.LitellmUserRoles.INTERNAL_USER), + (VIEWER_KEY, proxy_server.LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY), + ): + proxy_server.user_api_key_cache.set_cache( + key=proxy_server.hash_token(key), + value=UserAPIKeyAuth(token=proxy_server.hash_token(key), user_role=role, user_id=key), + ) async with ( httpx.AsyncClient(transport=httpx.MockTransport(upstream), trust_env=False) as source_client, httpx.AsyncClient(transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test") as client, ): monkeypatch.setattr(proxy_server.app, "dependency_overrides", { get_plugin_upload_repository: lambda: PluginUploadRepository(prisma), + get_skill_submission_repository: lambda: SkillSubmissionRepository(prisma), import_http_client: lambda: source_client, }) - yield HubClient(client, plugins, archives, upstream, clock) + yield HubClient(client, plugins, archives, downloads, prisma.db.execute_raw, upstream, clock, prisma.db) -async def _upload(hub: HubClient, content: bytes, key: str | None = MASTER_KEY) -> httpx.Response: +async def _upload( + hub: HubClient, content: bytes, key: str | None = MASTER_KEY, version: str | None = "2.0.0", +) -> httpx.Response: return await hub.client.post( f"{HUB}/upload", headers=_headers(key), - data={"version": "2.0.0"}, + data={"version": version} if version is not None else {}, files={"file": ("notes.zip", content, "application/zip")}, ) @@ -141,14 +176,15 @@ async def _upload(hub: HubClient, content: bytes, key: str | None = MASTER_KEY) def _stored_upload(hub: HubClient, enabled: bool = True) -> prisma_models.LiteLLM_ClaudeCodePluginTable: data: Final = hub.plugins.create.await_args.kwargs["data"] record: Final = prisma_models.LiteLLM_ClaudeCodePluginTable( - id="hub-id", name=data["name"], version=data["version"], description=data["description"], - manifest_json=data["manifest_json"], enabled=enabled, + id="hub-id", name=data["name"], version=data["version"], description=data["description"], anonymous_download_count=0, + manifest_json=data["manifest_json"], enabled=enabled, updated_at=data["updated_at"], ) hub.plugins.find_unique.return_value = record hub.plugins.find_many.return_value = [record] hub.archives.find_first.return_value = prisma_models.LiteLLM_ClaudeCodePluginUpload( - plugin_name=record.name, content=data["uploaded_archive"]["create"]["content"], + plugin_name=record.name, content=hub.archives.create.await_args.kwargs["data"]["content"], ) + hub.archives.find_unique.return_value = hub.archives.find_first.return_value return record @@ -160,6 +196,27 @@ def _assert_portable(response: httpx.Response) -> None: assert stat.S_IMODE(archive.getinfo("scripts/check.sh").external_attr >> 16) == 0o755 +@pytest.mark.parametrize("key, expected", ((MASTER_KEY, 7), (PLAIN_KEY, None))) +async def test_catalog_counts_are_admin_only_and_batched_for_the_page( + hub: HubClient, key: str, expected: int | None, +) -> None: + assert (await _upload(hub, _zip(FILES))).status_code == 200 + record: Final = _stored_upload(hub) + hub.plugins.find_many.return_value = [record, record.model_copy(update={"id": "other", "name": "zzz-other"})] + hub.database.query_raw.reset_mock() + hub.database.query_raw.return_value = [{"id": record.id, "download_count": 7}] + response: Final = await hub.client.get(f"{HUB}?limit=1", headers=_headers(key)) + assert response.status_code == 200, response.text + assert [(item["id"], item["download_count"]) for item in response.json()["skills"]] == [(record.id, expected)] + if expected is None: + hub.database.query_raw.assert_not_awaited() + else: + assert hub.database.query_raw.await_count == 1 + assert hub.database.query_raw.await_args.args[1] == (record.id,) + assert (await hub.client.get(f"{HUB}?offset=2", headers=_headers(key))).json()["skills"] == [] + assert hub.database.query_raw.await_count == 1 + + @pytest.mark.parametrize("operation", ("upload", "import")) @pytest.mark.parametrize("public_base, root_path", ( ("http://test", ""), ("https://public.example", ""), ("https://public.example", "/gateway"), @@ -176,6 +233,7 @@ async def test_publication_round_trips_both_formats_without_inflating_on_the_req ) ) assert response.status_code == 200, response.text + assert response.json()["version"] == ("2.0.0" if operation == "upload" else "1.0.0") record: Final = _stored_upload(hub) manifest: Final = json.loads(record.manifest_json) monkeypatch.setattr(proxy_server.app, "root_path", root_path) @@ -190,6 +248,8 @@ async def test_publication_round_trips_both_formats_without_inflating_on_the_req assert details.json()["content_url"] == f"{serving_base}{HUB}/{NAME}/content" catalog: Final = await hub.client.get("/claude-code/plugins", headers=_headers(PLAIN_KEY)) assert catalog.status_code == 200 and "has_agent_skill_metadata" not in catalog.text + hub.downloads.create.assert_not_awaited() + hub.anonymous_counter.assert_not_awaited() request_thread: Final = threading.get_ident() decompress: Final = zlib.decompressobj @@ -209,10 +269,13 @@ async def test_publication_round_trips_both_formats_without_inflating_on_the_req assert (await hub.client.get(entry["url"])).content == portable.content assert (await hub.client.get(f"{HUB}/{NAME}/archive?digest=sha256:stale")).status_code == 404 legacy_path: Final = httpx.URL(manifest["source"]["url"]).path - legacy: Final = await hub.client.get(legacy_path) + legacy: Final = await hub.client.get(legacy_path, headers=_headers(GRANTED_KEY)) assert legacy.status_code == 200 assert hashlib.sha256(legacy.content).hexdigest() == manifest["source"]["sha256"] assert (await hub.client.get(legacy_path.replace(manifest["source"]["sha256"], "stale"))).status_code == 404 + assert hub.anonymous_counter.await_count == 2 + assert hub.anonymous_counter.await_args.args[1:] == (record.id,) + hub.downloads.create.assert_awaited_once() monkeypatch.setattr(zlib, "decompressobj", decompress) _assert_portable(portable) with zipfile.ZipFile(io.BytesIO(legacy.content)) as archive: @@ -285,6 +348,8 @@ async def test_download_policy_requires_publication_or_a_valid_grant( for suffix in ("archive", "content"): response: Final = await hub.client.get(f"{HUB}/{NAME}/{suffix}", headers=_headers(key)) assert response.status_code == status, response.text + assert hub.downloads.create.await_count == int(status == 200 and key is not None) + assert hub.anonymous_counter.await_count == int(status == 200 and key is None) if status != 200: hub.archives.find_first.assert_not_awaited() index: Final = await hub.client.get(INDEX) @@ -293,6 +358,21 @@ async def test_download_policy_requires_publication_or_a_valid_grant( assert [entry["name"] for entry in index.json()["skills"]] == ([NAME] if enabled and not auth_required else []) +@pytest.mark.parametrize("key", (None, GRANTED_KEY)) +async def test_deleted_skill_is_not_served_when_download_cannot_be_recorded(hub: HubClient, key: str | None) -> None: + assert (await _upload(hub, _zip(FILES))).status_code == 200 + record: Final = _stored_upload(hub) + legacy_url: Final = json.loads(record.manifest_json)["source"]["url"] + hub.anonymous_counter.return_value = 0 + hub.downloads.create.side_effect = ForeignKeyViolationError(data={}) + for path in (f"{HUB}/{NAME}/archive", legacy_url): + response: Final = await hub.client.get(path, headers=_headers(key)) + assert response.status_code == 404, response.text + assert "Skill archive not found" in response.text + assert hub.anonymous_counter.await_count == (2 if key is None else 0) + assert hub.downloads.create.await_count == (0 if key is None else 2) + + async def test_search_paginates_visible_matches(hub: HubClient) -> None: assert (await _upload(hub, _zip(FILES))).status_code == 200 record: Final = _stored_upload(hub) @@ -387,24 +467,36 @@ async def test_stored_content_survives_duplicate_upload_and_is_hidden_after_sour assert (await hub.client.get(INDEX)).json()["skills"] == [] -@pytest.mark.parametrize("user_id, admitted", (("publisher", 10), (None, 20))) +@pytest.mark.parametrize("user_id, role, admitted", ( + ("publisher", proxy_server.LitellmUserRoles.PROXY_ADMIN, 10), + (None, proxy_server.LitellmUserRoles.PROXY_ADMIN, 20), + ("publisher", proxy_server.LitellmUserRoles.INTERNAL_USER, 10), +), ids=("admin-user", "admin-key", "member-user")) async def test_publication_limit_is_shared_by_user_and_route_and_resets( - hub: HubClient, monkeypatch: pytest.MonkeyPatch, user_id: str | None, admitted: int, + hub: HubClient, monkeypatch: pytest.MonkeyPatch, user_id: str | None, role: str, admitted: int, ) -> None: for key in (PLAIN_KEY, GRANTED_KEY): proxy_server.user_api_key_cache.set_cache( key=proxy_server.hash_token(key), value=UserAPIKeyAuth( token=proxy_server.hash_token(key), user_id=user_id, - user_role=proxy_server.LitellmUserRoles.PROXY_ADMIN, + user_role=role, ), ) - responses: Final = await asyncio.gather(*( - _upload(hub, b"invalid", PLAIN_KEY) if index % 2 else hub.client.post( - f"{HUB}/import", headers=_headers(GRANTED_KEY), json={"url": "https://example.com/skill"}, - ) for index in range(20) - )) - assert sum(response.status_code in (400, 422) for response in responses) == admitted + async def attempt(index: int) -> httpx.Response: + key: Final = PLAIN_KEY if index % 2 else GRANTED_KEY + if index % 3 == 0: + return await _upload(hub, b"invalid", key) + if index % 3 == 1: + return await hub.client.post( + f"{HUB}/import", headers=_headers(key), json={"url": "https://example.com/skill"}, + ) + return await hub.client.put( + f"{HUB}/{NAME}/upload", headers=_headers(key), files={"file": ("invalid.zip", b"invalid")}, + ) + + responses: Final = await asyncio.gather(*(attempt(index) for index in range(20))) + assert sum(response.status_code in (400, 404, 422) for response in responses) == admitted assert sum(response.status_code == 429 for response in responses) == 20 - admitted assert all(response.headers["retry-after"] == "60" for response in responses if response.status_code == 429) hub.plugins.create.assert_not_awaited() @@ -418,3 +510,299 @@ async def test_publication_limit_is_shared_by_user_and_route_and_resets( assert (await _upload(hub, b"invalid", GRANTED_KEY)).status_code == 422 monkeypatch.setattr(proxy_server.proxy_logging_obj, "proxy_hook_mapping", {}) assert (await _upload(hub, b"invalid")).status_code == 503 + + +@pytest.mark.asyncio +async def test_neutral_metadata_publication_pairs_version_and_preserved_skill_snapshot( + hub: HubClient, monkeypatch: pytest.MonkeyPatch, +) -> None: + created: Final = await _upload(hub, _zip(FILES), version=None) + assert created.status_code == 200, created.text + assert created.json()["version"] == "1.0.0" + original: Final = _stored_upload(hub) + hub.plugins.find_unique.return_value = original.model_copy(update={ + "manifest_json": json.dumps({**json.loads(original.manifest_json), "author": "legacy author"}), + }) + assert (await hub.client.get(f"{HUB}/{NAME}", headers=_headers(MASTER_KEY))).status_code == 200 + monkeypatch.delenv("PROXY_BASE_URL") + hub.client.base_url = httpx.URL("https://public.example") + first: Final = hub.database.litellm_claudecodepluginversion.create.await_args.kwargs["data"] + hub.database.litellm_claudecodepluginversion.count.return_value = 1 + + async def replace_row( + where: Mapping[str, str], data: Mapping[str, object], + ) -> prisma_models.LiteLLM_ClaudeCodePluginTable: + current: Final = original.model_copy(update=dict(data)) + hub.plugins.find_unique.return_value = current + return current + + hub.plugins.update.side_effect = replace_row + updated: Final = await hub.client.put( + f"{HUB}/{NAME}", headers=_headers(MASTER_KEY), + json={"description": "Registry description", "domain": "notes", "namespace": "team", "keywords": ["notes"]}, + ) + assert updated.status_code == 200, updated.text + assert updated.json()["version"] == "1.0.1" + assert updated.json()["download_url"] == f"https://public.example{HUB}/{NAME}/archive" + latest: Final = hub.database.litellm_claudecodepluginversion.create.await_args.kwargs["data"] + manifest: Final = json.loads(latest["manifest_json"]) + assert latest["version"] == manifest["version"] == updated.json()["version"] + assert latest["manifest_json"] == hub.plugins.find_unique.return_value.manifest_json + assert manifest["agent_skill"] == json.loads(original.manifest_json)["agent_skill"] + assert (manifest["domain"], manifest["namespace"], manifest["keywords"]) == ("notes", "team", ["notes"]) + assert first["version"] == "1.0.0" + assert latest["content"].decode() == first["content"].decode() + _assert_portable(await hub.client.get(updated.json()["download_url"])) + + +@pytest.mark.parametrize("key,status", ((MASTER_KEY, 200), (GRANTED_KEY, 200), (PLAIN_KEY, 404))) +@pytest.mark.asyncio +async def test_neutral_history_list_detail_and_portable_download_share_current_visibility( + hub: HubClient, key: str, status: int, +) -> None: + assert (await _upload(hub, _zip(FILES), version=None)).status_code == 200 + original: Final = _stored_upload(hub, enabled=False) + data: Final = hub.database.litellm_claudecodepluginversion.create.await_args.kwargs["data"] + version: Final = prisma_models.LiteLLM_ClaudeCodePluginVersion.model_validate({"id": "retained", **data}) + hub.plugins.find_unique.return_value = original.model_copy(update={"version": "3.0.0"}) + hub.database.litellm_claudecodepluginversion.find_first.return_value = version + hub.database.query_raw.return_value = [{ + "id": version.id, "version": version.version, "description": version.description, + "manifest": json.loads(version.manifest_json), "published_at": version.published_at, + "published_by": version.published_by, "archive_available": version.archive_available, + }] + listed: Final = await hub.client.get(f"{HUB}/{NAME}/versions", headers=_headers(key)) + detail: Final = await hub.client.get(f"{HUB}/{NAME}/versions/retained", headers=_headers(key)) + downloaded: Final = await hub.client.get(f"{HUB}/{NAME}/versions/retained/archive", headers=_headers(key)) + assert listed.status_code == detail.status_code == downloaded.status_code == status + assert hub.downloads.create.await_count == int(status == 200) + if status == 200: + event: Final = hub.downloads.create.await_args.kwargs["data"] + assert event["version"] == version.version != "3.0.0" + assert event["digest"] == hashlib.sha256(downloaded.content).hexdigest() + assert listed.json() == {"versions": [detail.json()], "next_offset": None} + assert detail.json()["manifest"]["agent_skill"] == json.loads(original.manifest_json)["agent_skill"] + assert detail.json()["download_url"] == f"http://test{HUB}/{NAME}/versions/retained/archive" + _assert_portable(downloaded) + assert downloaded.headers["etag"] == f'"sha256:{hashlib.sha256(downloaded.content).hexdigest()}"' + hub.database.query_raw.return_value = [] + assert (await hub.client.get(f"{HUB}/{NAME}/versions/missing", headers=_headers(key))).status_code == 404 + hub.database.litellm_claudecodepluginversion.find_first.return_value = None + assert (await hub.client.get( + f"{HUB}/{NAME}/versions/missing/archive", headers=_headers(key), + )).status_code == 404 + + +@pytest.mark.asyncio +async def test_neutral_history_labels_external_sources_and_bounds_pagination(hub: HubClient) -> None: + assert (await _upload(hub, _zip(FILES))).status_code == 200 + _stored_upload(hub) + row: Final = { + "id": "reference", "version": "3.0.0", "description": None, + "manifest": {"source": {"source": "github", "repo": "example/skill"}}, + "published_at": None, "published_by": None, "archive_available": False, + } + hub.database.query_raw.return_value = [row, {**row, "id": "older"}] + listed: Final = await hub.client.get(f"{HUB}/{NAME}/versions?limit=1", headers=_headers(MASTER_KEY)) + assert listed.status_code == 200, listed.text + assert listed.json()["next_offset"] == 1 + assert len(listed.json()["versions"]) == 1 + assert listed.json()["versions"][0]["content_kind"] == "source_reference" + assert listed.json()["versions"][0]["download_url"] is None + assert (await hub.client.get( + f"{HUB}/{NAME}/versions?limit=101", headers=_headers(MASTER_KEY), + )).status_code == 422 + + +@pytest.fixture +async def pending_hub(hub: HubClient) -> HubClient: + hub.database.query_raw.side_effect = hub.database.litellm_skillsubmission.query_raw + return hub + + +@pytest.mark.parametrize("operation,approve", (("upload", True), ("import", True), ("upload", False)), + ids=("approve-upload", "approve-import", "reject-upload")) +async def test_member_archive_stays_private_until_admin_review( + pending_hub: HubClient, operation: str, approve: bool, +) -> None: + hub: Final = pending_hub + hub.upstream.return_value = httpx.Response(200, content=_zip(FILES, "repository/")) + pending: Final = ( + await _upload(hub, _zip(FILES), MEMBER_KEY, version=None) if operation == "upload" + else await hub.client.post( + f"{HUB}/import", headers=_headers(MEMBER_KEY), + json={"url": f"https://skills.sh/synthetic-author/notes-library/{NAME}"}, + ) + ) + assert pending.status_code == 200, pending.text + receipt: Final = pending.json() + assert receipt["status"] == "pending_review" and receipt["proposal"]["version"] is None + assert receipt["proposal"]["agent_skill"]["name"] == NAME + hub.plugins.create.assert_not_awaited() + hub.database.litellm_claudecodepluginversion.create.assert_not_awaited() + path: Final = f"/skill-hub/submissions/{receipt['submission_id']}" + assert receipt["proposal"]["source"]["url"] == f"http://test{path}/archive" + own: Final = await hub.client.get(path, headers=_headers(MEMBER_KEY)) + assert own.json()["current_skill"] is None + assert own.json()["submission"]["proposal"] == receipt["proposal"] + for suffix in ("", "/archive"): + assert (await hub.client.get(path + suffix, headers=_headers(OTHER_KEY))).status_code == 404 + _assert_portable(await hub.client.get(path + "/archive", headers=_headers(MEMBER_KEY))) + assert (await hub.client.get(path, headers=_headers(VIEWER_KEY))).status_code == 200 + denied: Final = await hub.client.put(path + "/approve", headers=_headers(VIEWER_KEY)) + assert denied.status_code == 403, denied.text + assert (await hub.client.get(f"{HUB}/{NAME}/archive")).status_code == 404 + if not approve: + rejected: Final = await hub.client.put( + path + "/reject", headers=_headers(MASTER_KEY), json={"reason": "Needs revision"}, + ) + assert rejected.status_code == 200, rejected.text + assert rejected.json()["skill"] is None and rejected.json()["action"] == "rejected" + assert rejected.json()["submission"]["rejection_reason"] == "Needs revision" + assert rejected.json()["submission"]["proposal"] == receipt["proposal"] + hub.plugins.create.assert_not_awaited() + return + approved: Final = await hub.client.put(path + "/approve", headers=_headers(MASTER_KEY)) + assert approved.status_code == 200, approved.text + assert approved.json()["action"] == "approved" + assert approved.json()["skill"]["version"] == "1.0.0" + assert approved.json()["submission"]["proposal"] == receipt["proposal"] + _stored_upload(hub) + _assert_portable(await hub.client.get(approved.json()["skill"]["download_url"])) + assert (await hub.client.put(path + "/reject", headers=_headers(MASTER_KEY))).status_code == 409 + listed: Final = await hub.client.get("/skill-hub/submissions?status=approved", headers=_headers(MEMBER_KEY)) + assert (listed.json()["count"], listed.json()["has_more"]) == (1, False) + + +@pytest.mark.parametrize("replacement", (False, True), ids=("metadata", "archive")) +async def test_member_updates_preserve_manifest_and_original_editor_base(pending_hub: HubClient, replacement: bool) -> None: + hub: Final = pending_hub + assert (await _upload(hub, _zip(FILES))).status_code == 200 + original: Final = _stored_upload(hub) + manifest: Final = { + **json.loads(original.manifest_json), "domain": "Tools", "namespace": "research", + "keywords": ["notes"], "extension": {"retained": True}, "author": "legacy author", + } + hub.plugins.find_unique.return_value = original.model_copy(update={"manifest_json": json.dumps(manifest)}) + before: Final = hub.database.litellm_claudecodepluginversion.create.await_count + updated_files: Final = {**FILES, "references/guide.md": b"Updated supporting instructions.\n"} + expected: Final = original.updated_at.isoformat() + response: Final = ( + await hub.client.put( + f"{HUB}/{NAME}/upload", headers=_headers(MEMBER_KEY), + data={"expected_updated_at": expected}, files={"file": ("new.zip", _zip(updated_files), "application/zip")}, + ) if replacement else await hub.client.put( + f"{HUB}/{NAME}", headers=_headers(MEMBER_KEY), + json={"description": "Reviewed description", "expected_updated_at": expected, + "domain": "Tools", "namespace": "research", "keywords": ["notes"]}, + ) + ) + assert response.status_code == 200, response.text + path: Final = f"/skill-hub/submissions/{response.json()['submission_id']}" + detail: Final = (await hub.client.get(path, headers=_headers(MEMBER_KEY))).json() + assert detail["submission"]["base_updated_at"] == original.updated_at.isoformat().replace("+00:00", "Z") + assert detail["submission"]["proposal"]["domain"] == "Tools" + assert detail["submission"]["proposal"]["keywords"] == ["notes"] + assert detail["current_skill"]["namespace"] == "research" + assert hub.database.litellm_claudecodepluginversion.create.await_count == before + hub.plugins.update.assert_not_awaited() + approved: Final = await hub.client.put(path + "/approve", headers=_headers(MASTER_KEY)) + assert approved.status_code == 200, approved.text + stored: Final = json.loads(hub.plugins.find_unique.return_value.manifest_json) + assert (stored["extension"], stored["author"]) == (manifest["extension"], manifest["author"]) + assert stored["agent_skill"]["digest"] == response.json()["proposal"]["agent_skill"]["digest"] + assert approved.json()["skill"]["version"] == "2.0.1" + assert approved.json()["submission"]["proposal"]["version"] is None + wrapper: Final = hub.archives.upsert.await_args.kwargs["data"]["update"]["content"].decode() + with zipfile.ZipFile(io.BytesIO(wrapper)) as archive: + assert archive.read(f"skills/{NAME}/references/guide.md") == ( + updated_files["references/guide.md"] if replacement else FILES["references/guide.md"] + ) + if not replacement: + assert stored["agent_skill"] == manifest["agent_skill"] + + +@pytest.mark.parametrize("failure,status", (("stale", 409), ("wrong-name", 422), ("viewer", 403)), + ids=("stale", "wrong-name", "viewer")) +async def test_replacement_rejections_do_not_stage_or_change_files(pending_hub: HubClient, failure: str, status: int) -> None: + hub: Final = pending_hub + assert (await _upload(hub, _zip(FILES))).status_code == 200 + _stored_upload(hub) + response: Final = await hub.client.put( + f"{HUB}/{NAME}/upload", headers=_headers(VIEWER_KEY if failure == "viewer" else MEMBER_KEY), + data={"expected_updated_at": "2020-01-01T00:00:00Z"} if failure == "stale" else {}, + files={"file": ("new.zip", _zip({"SKILL.md": MANIFEST.replace(NAME.encode(), b"different-name")}), "application/zip")}, + ) + assert response.status_code == status, response.text + assert hub.database.litellm_skillsubmission.rows == () + hub.plugins.update.assert_not_awaited() + hub.archives.upsert.assert_not_awaited() + + +@pytest.mark.parametrize("key", (MEMBER_KEY, OTHER_KEY, VIEWER_KEY)) +@pytest.mark.parametrize("method", ("GET", "POST")) +async def test_skill_scan_requires_full_admin(hub: HubClient, key: str, method: str) -> None: + response = await hub.client.request(method, "/skill-hub/submissions/absent/scan", headers=_headers(key)) + assert response.status_code == 403, response.text + + +async def test_unconfigured_scanning_preserves_manual_approval(pending_hub: HubClient) -> None: + hub = pending_hub + created = await _upload(hub, _zip(FILES), MEMBER_KEY, version=None) + path = f"/skill-hub/submissions/{created.json()['submission_id']}" + state = await hub.client.get(path + "/scan", headers=_headers(MASTER_KEY)) + assert state.status_code == 200, state.text + assert state.json() == {"configured": False, "scan": None} + unavailable = await hub.client.post(path + "/scan", headers=_headers(MASTER_KEY)) + assert unavailable.status_code == 503, unavailable.text + approved = await hub.client.put(path + "/approve", headers=_headers(MASTER_KEY)) + assert approved.status_code == 200, approved.text + assert approved.json()["submission"]["status"] == "approved" + + +@pytest.mark.parametrize("scanner_status", (200, 503), ids=("report", "error")) +async def test_configured_scan_persists_result_without_changing_manual_approval( + pending_hub: HubClient, monkeypatch: pytest.MonkeyPatch, scanner_status: int, +) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy.discovery_endpoints import skill_scanning + + hub: Final = pending_hub + report: Final = { + "scanner_version": "test-version", "complete": True, "scan_mode": "static-only", + "risk_score": 0, "findings": [], "limitations": [], + } + + def respond(request: httpx.Request) -> httpx.Response: + assert request.headers["authorization"] == "Bearer scanner-only-token" + return httpx.Response(scanner_status, json=report) + + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + + def scanner_client(provider: str, params: Mapping[str, object]) -> AsyncHTTPHandler: + return handler + + monkeypatch.setattr(skill_scanning, "get_async_httpx_client", scanner_client) + monkeypatch.setattr(proxy_server, "general_settings", {"skill_scanner": { + "endpoint": "https://scanner.example/scan", "api_key_env": "SCANNER_TEST_KEY", + }}) + monkeypatch.setenv("SCANNER_TEST_KEY", "scanner-only-token") + created: Final = await _upload(hub, _zip(FILES), MEMBER_KEY, version=None) + path: Final = f"/skill-hub/submissions/{created.json()['submission_id']}" + before: Final = (await hub.client.get(path, headers=_headers(MASTER_KEY))).json() + async with handler.client: + started: Final = await hub.client.post(path + "/scan", headers=_headers(MASTER_KEY)) + assert started.status_code == 200, started.text + completed: Final = (await hub.client.get(path + "/scan", headers=_headers(MASTER_KEY))).json()["scan"] + assert completed["id"] == started.json()["scan"]["id"] + expected: Final = ("completed", report) if scanner_status == 200 else ("failed", None) + assert (completed["status"], completed["report"]) == expected + assert completed["error"] is None if scanner_status == 200 else "HTTP 503" in completed["error"] + monkeypatch.setattr(proxy_server, "general_settings", {"skill_scanner": {"endpoint": "invalid"}}) + retained: Final = await hub.client.get(path + "/scan", headers=_headers(MASTER_KEY)) + assert retained.status_code == 200 and retained.json() == {"configured": False, "scan": completed} + unavailable: Final = await hub.client.post(path + "/scan", headers=_headers(MASTER_KEY)) + assert unavailable.status_code == 503 and "configuration is invalid" in unavailable.text + assert (await hub.client.get(path, headers=_headers(MASTER_KEY))).json() == before + approved: Final = await hub.client.put(path + "/approve", headers=_headers(MASTER_KEY)) + assert approved.status_code == 200 and approved.json()["submission"]["status"] == "approved" diff --git a/tests/unit/proxy/discovery_endpoints/test_skill_scanning.py b/tests/unit/proxy/discovery_endpoints/test_skill_scanning.py new file mode 100644 index 00000000000..934620e5de8 --- /dev/null +++ b/tests/unit/proxy/discovery_endpoints/test_skill_scanning.py @@ -0,0 +1,155 @@ +import hashlib +import io +import zipfile +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from fastapi.testclient import TestClient +from pydantic import ValidationError + +from litellm.models.skill_submission import SkillScanInput +from litellm.proxy.discovery_endpoints.skill_scanning import scan_report +from litellm.types.proxy.skill_scanning import SkillScan, SkillScannerConfig, SkillScanReport + +NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +CONFIG: Final = SkillScannerConfig(endpoint="https://scanner.example/scan", api_key_env="SCANNER_TEST_KEY") +REPORT: Final = SkillScanReport( + scanner_version="test-version", complete=True, scan_mode="static-only", risk_score=0, findings=(), +) +TARGET: Final = SkillScanInput( + SkillScan(id="scan-1", status="running", requested_by="admin", started_at=NOW, expires_at=NOW + timedelta(seconds=10)), + {"source": "archive", "url": "https://gateway.example/private.zip"}, + b"exact-private-archive", +) + + +@pytest.mark.asyncio +async def test_scan_posts_stored_bytes_with_only_scanner_credentials() -> None: + def respond(request: httpx.Request) -> httpx.Response: + assert request.url == CONFIG.endpoint + assert request.headers["authorization"] == "Bearer scanner-only-token" + assert b"exact-private-archive" in request.content + assert b'name="use_llm"\r\n\r\nfalse' in request.content + assert f'name="expires_at"\r\n\r\n{TARGET.scan.expires_at.isoformat()}'.encode() in request.content + return httpx.Response(200, json=REPORT.model_dump()) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + result: Final = await scan_report(client, CONFIG, "scanner-only-token", TARGET) + assert result == REPORT + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ("redirect", "unavailable", "malformed", "oversized", "network")) +async def test_scanner_failures_never_become_clean_reports(failure: str) -> None: + def respond(request: httpx.Request) -> httpx.Response: + assert request.url.host == "scanner.example" + if failure == "redirect": + return httpx.Response(302, headers={"Location": "https://untrusted.example/scan"}) + if failure == "unavailable": + return httpx.Response(503, text="provider-internal-secret") + if failure == "malformed": + return httpx.Response(200, json={"complete": True}) + if failure == "oversized": + return httpx.Response(200, content=b" " * (1024 * 1024 + 1)) + raise httpx.ConnectError("provider-internal-secret", request=request) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond), follow_redirects=True) as client: + result: Final = await scan_report(client, CONFIG, "scanner-only-token", TARGET) + assert isinstance(result, str) + assert "provider-internal-secret" not in result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("existing_limitations", (0, 100), ids=("empty", "full")) +async def test_missing_requested_semantic_pass_remains_incomplete(existing_limitations: int) -> None: + report: Final = REPORT.model_copy(update={"limitations": ("Existing limitation",) * existing_limitations}) + + def respond(request: httpx.Request) -> httpx.Response: + assert b'name="use_llm"\r\n\r\ntrue' in request.content + return httpx.Response(200, json=report.model_dump()) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + result: Final = await scan_report(client, CONFIG.model_copy(update={"use_llm": True}), "token", TARGET) + assert isinstance(result, SkillScanReport) + assert result == report.model_copy(update={ + "complete": False, "limitations": (*report.limitations[:99], "Requested semantic analysis did not run."), + }) + assert SkillScanReport.model_validate(result.model_dump()) == result + + +def test_interrupted_scan_becomes_retryable_without_changing_completed_evidence() -> None: + assert TARGET.scan.at(NOW) == TARGET.scan + assert TARGET.scan.at(NOW + timedelta(seconds=10)).status == "failed" + complete: Final = TARGET.scan.model_copy(update={"status": "completed", "report": REPORT}) + assert complete.at(NOW + timedelta(days=1)) == complete + + +@pytest.mark.parametrize("endpoint", ( + "http://scanner.example/scan", "https://user:secret@scanner.example/scan", "https://scanner.example/scan?token=secret", +)) +def test_scanner_config_rejects_unsafe_credential_destinations(endpoint: str) -> None: + with pytest.raises(ValidationError): + SkillScannerConfig(endpoint=endpoint, api_key_env="SCANNER_TEST_KEY") + + +def test_adapter_authentication_and_exact_remote_source(monkeypatch: pytest.MonkeyPatch) -> None: + from cookbook.skill_scanner import server + + monkeypatch.setenv("SKILL_SCANNER_API_KEY", "adapter-secret") + clock: Final = MagicMock() + clock.now.return_value = NOW + timedelta(seconds=1) + monkeypatch.setattr(server, "datetime", clock) + launch: Final = AsyncMock(side_effect=AssertionError("Invalid deadline launched a worker")) + monkeypatch.setattr(server.asyncio, "create_subprocess_exec", launch) + with TestClient(server.app) as client: + assert client.post("/scan", data={"source": "{}"}).status_code == 401 + assert client.post("/scan", headers={"Authorization": "Bearer wrong"}, data={"source": "{}"}).status_code == 401 + response = client.post("/scan", headers={"Authorization": "Bearer adapter-secret"}, data={ + "source": '{"source":"git","url":"https://example.org/repo.git","ref":"release"}', + }) + assert response.status_code == 400 + for deadline, status in ((NOW.isoformat(), 504), (NOW.replace(tzinfo=None).isoformat(), 422)): + rejected: Final = client.post("/scan", headers={"Authorization": "Bearer adapter-secret"}, data={ + "source": '{"source":"archive","url":"https://example.org/skill.zip"}', "expires_at": deadline, + }) + assert rejected.status_code == status + launch.assert_not_called() + assert server.source_target({"source": "archive", "url": "https://example.org/skill.zip"}) == "https://example.org/skill.zip" + assert server.source_target({"source": "github", "repo": "example/skills", "ref": "v2", "path": "skills/review"}) == "https://github.com/example/skills/tree/v2/skills/review" + + +def test_adapter_accepts_expanded_stored_skill_bytes(monkeypatch: pytest.MonkeyPatch) -> None: + from cookbook.skill_scanner import server + from litellm.proxy.anthropic_endpoints.claude_code_endpoints.skill_upload_archive import ( + build_uploaded_skill_archive, + ) + + uploaded: Final = io.BytesIO() + with zipfile.ZipFile(uploaded, "w", zipfile.ZIP_DEFLATED) as archive: + archive.writestr("SKILL.md", "---\nname: large-skill\ndescription: Synthetic archive\n---\n") + for name in ("first.bin", "second.bin"): + archive.writestr(name, b"x" * (6 * 1024 * 1024)) + stored: Final = build_uploaded_skill_archive(uploaded.getvalue(), "large-skill") + assert stored is not None + stored_size: Final = len(stored) + assert stored_size > 10 * 1024 * 1024 + expected: Final = (stored_size, hashlib.sha256(stored).hexdigest()) + + async def run_worker( + target: str, output: Path, use_llm: bool, expires_at: datetime | None = None, + ) -> dict[str, object]: + received: Final = Path(target).read_bytes() + actual: Final = (len(received), hashlib.sha256(received).hexdigest()) + assert actual == expected + return {"complete": True} + + monkeypatch.setattr(server, "run_worker", run_worker) + monkeypatch.setenv("SKILL_SCANNER_API_KEY", "adapter-secret") + with TestClient(server.app) as client: + response: Final = client.post("/scan", headers={"Authorization": "Bearer adapter-secret"}, + data={"source": "{}"}, files={"file": ("skill.zip", stored, "application/zip")}) + assert response.status_code == 200 and response.json() == {"complete": True} diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 84f4d7a4173..de3ced0c583 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3291,3 +3291,21 @@ async def test_pass_through_guardrail_text_is_the_client_body_without_the_inject texts: Final = [call["inputs"]["texts"] for call in guardrail.apply_calls] assert [[json.loads(text) for text in call_texts] for call_texts in texts] == [[body]], texts assert data["litellm_metadata"]["user_api_key_user_email"] == "someone@example.com", data["litellm_metadata"] + + +@pytest.mark.asyncio +async def test_moderation_hook_lets_responses_compact_through_without_a_translation() -> None: + handler: Final = UnifiedLLMGuardrails() + guardrail: Final = RecordingGuardrail() + data: Final = { + "guardrail_to_apply": guardrail, + "model": "gpt-5.4-mini", + "input": [{"role": "user", "content": "Remember the number 42."}], + } + + result: Final = await handler.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), call_type="acompact_responses" + ) + + assert result is data + assert guardrail.apply_calls == [] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py index bccc531793f..712750f702e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/test_zscaler_ai_guard.py @@ -23,6 +23,8 @@ from fastapi import HTTPException import litellm from litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import ZscalerAIGuard +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.types.utils import GenericGuardrailAPIInputs from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -141,6 +143,216 @@ async def test_apply_guardrail_text_concatenation(mock_api_call): assert call_args.kwargs["content"] == "Hello world" +_CONVERSATION: Final = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is two plus two?"}, + {"role": "assistant", "content": "Four."}, +] + + +async def _execute_policy_bodies( + inputs: GenericGuardrailAPIInputs, + input_type: Literal["request", "response"], + scan_only_tool_results: bool = False, +) -> tuple[object, ...]: + received: Final = asyncio.Queue[object]() + + async def execute_policy(request: web.Request) -> web.Response: + await received.put(await request.json()) + return web.json_response({"statusCode": 200, "action": "ALLOW"}) + + app: Final = web.Application() + app.router.add_post("/v1/detection/execute-policy", execute_policy) + handler: Final = InMemoryGuardrailHandler() + async with TestServer(app) as server: + guardrail_id: Final = handler.initialize_guardrail( + guardrail={ + "guardrail_name": "zscaler-last-user-turn", + "litellm_params": { + "guardrail": "zscaler_ai_guard", + "mode": "pre_call", + "api_key": "test_key", + "api_base": str(server.make_url("/v1/detection/execute-policy")), + "policy_id": 1, + "scan_only_tool_results": scan_only_tool_results, + }, + }, + )["guardrail_id"] + try: + guardrail: Final = handler.guardrail_id_to_custom_guardrail[guardrail_id] + assert guardrail is not None + returned: Final = await guardrail.apply_guardrail(inputs, {}, input_type) + finally: + handler.delete_in_memory_guardrail(guardrail_id) + + assert returned is inputs + return tuple(received.get_nowait() for _ in range(received.qsize())) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inputs", "expected_content"), + [ + ( + { + "texts": ["every", "message", "flattened"], + "structured_messages": [*_CONVERSATION, {"role": "user", "content": "And three plus three?"}], + }, + "And three plus three?", + ), + ( + { + "texts": ["every", "message", "flattened"], + "structured_messages": [ + *_CONVERSATION, + { + "role": "user", + "content": [{"type": "text", "text": "And three"}, {"type": "text", "text": "plus three?"}], + }, + ], + }, + "And three plus three?", + ), + ( + { + "texts": ["every", "message", "flattened"], + "structured_messages": [ + *_CONVERSATION, + {"role": "user", "content": "And three plus three?"}, + {"role": "user", "content": "Answer in words."}, + ], + }, + "And three plus three? Answer in words.", + ), + ( + { + "texts": ["Summarize", "this report"], + "structured_messages": [ + { + "role": "user", + "parts": [{"kind": "text", "text": "Summarize"}, {"kind": "text", "text": "this report"}], + "messageId": "msg-1", + } + ], + }, + "Summarize this report", + ), + ({"texts": ["first input", "second input"]}, "first input second input"), + ( + { + "texts": ["every", "message"], + "structured_messages": [*_CONVERSATION, {"role": "assistant", "content": "Prefilled answer:"}], + }, + "What is two plus two?", + ), + ( + { + "texts": ["every", "message"], + "structured_messages": [ + *_CONVERSATION, + {"role": "tool", "tool_call_id": "call_1", "content": "tool output"}, + ], + }, + "What is two plus two?", + ), + ( + { + "texts": ["every", "message"], + "structured_messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "assistant", "content": "Prefilled answer:"}, + ], + }, + None, + ), + ( + { + "texts": ["You are a helpful assistant.", " "], + "structured_messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": " "}, + ], + }, + None, + ), + ( + { + "texts": ["every", "message", "\n"], + "structured_messages": [ + *_CONVERSATION, + { + "role": "user", + "content": [ + {"type": "text", "text": "\n"}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + ], + }, + ], + }, + None, + ), + ], + ids=[ + "string-content", + "content-parts", + "consecutive-user-messages", + "a2a-message-parts", + "texts-only-route", + "ends-on-assistant", + "ends-on-tool", + "no-user-message", + "whitespace-only-latest-turn", + "whitespace-text-beside-image", + ], +) +async def test_request_scan_sends_only_the_last_user_turn( + inputs: GenericGuardrailAPIInputs, expected_content: str | None +) -> None: + bodies: Final = await _execute_policy_bodies(inputs, "request") + + assert bodies == ( + () if expected_content is None else ({"direction": "IN", "content": expected_content, "policyId": 1},) + ) + + +@pytest.mark.asyncio +async def test_response_scan_sends_every_text() -> None: + inputs: Final = GenericGuardrailAPIInputs( + texts=["First choice.", "Second choice."], + structured_messages=[*_CONVERSATION, {"role": "user", "content": "And three plus three?"}], + ) + + bodies: Final = await _execute_policy_bodies(inputs, "response") + + assert bodies == ({"direction": "OUT", "content": "First choice. Second choice.", "policyId": 1},) + + +@pytest.mark.asyncio +async def test_scan_only_tool_results_sends_the_tool_results() -> None: + inputs: Final = GenericGuardrailAPIInputs( + texts=["Sunny, 22C"], + structured_messages=[ + {"role": "user", "content": "What is the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Sunny, 22C"}, + ], + ) + + bodies: Final = await _execute_policy_bodies(inputs, "request", scan_only_tool_results=True) + + assert bodies == ({"direction": "IN", "content": "Sunny, 22C", "policyId": 1},) + + @pytest.mark.asyncio @patch( "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call", diff --git a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py index fe4d3c9bc4b..1fe2cb2ca68 100644 --- a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py +++ b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.anthropic.prompt_cache_prediction import NativePredictionTarget, TokenCounter from litellm.proxy.hooks.autorouter_baseline_cache import AutoRouterBaselineCache, CapturedBaselineObservation from litellm.router import Router +from litellm.router_rust.steps.prerouting import _record_routing_decision # pyright: ignore[reportPrivateUsage] # production trusted stamp owner from litellm.types.router import RetryPolicy from litellm.types.utils import CallTypes, StandardLoggingRoutingDecision @@ -86,7 +87,7 @@ def _kwargs(logging_obj: Logging, trusted: bool = True, *, explicit_logging: boo context: Final = _OBJECTS.validate_json( '{"max_tokens":16,"litellm_metadata":{"user_api_key_hash":"test-caller-hash"}}' ) - Router._record_routing_decision( # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # production trusted stamp owner + _record_routing_decision( context, StandardLoggingRoutingDecision( router_model_name="test-router", diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index c6366a40504..892c891f20c 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1939,6 +1939,7 @@ async def test_callback_handles_every_status_the_interactions_api_can_return(): assert released == { "completed", "requires_action", + "requires_continuation", "failed", "cancelled", "incomplete", @@ -2851,9 +2852,12 @@ async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_requ live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch ): """The window row's one-time seed from the spend log stops at the request start, which keeps the failure row - it already logged from being counted twice.""" + it already logged from being counted twice. Only the legacy writers queue this row, so the ledger is off.""" + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import proxy_logging_obj + monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False) + window_queue: Final = proxy_logging_obj.db_spend_update_writer.window_spend_update_queue await window_queue.flush_and_get_aggregated_window_spend_transactions() key_cache: Final = DualCache() diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index e487c00f25d..d658e59b76c 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -23,7 +23,11 @@ from litellm.proxy.management_endpoints.daily_activity_routes import ( router, ) from litellm.router import Router -from litellm.types.proxy.management_endpoints.common_daily_activity import KeySpendMetrics, SpendMetrics +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + KeySpendMetrics, + SpendAnalyticsPaginatedResponse, + SpendMetrics, +) from litellm.types.repositories.daily_activity import ( AggregatedRows, DailyActivityScope, @@ -600,6 +604,7 @@ def daily_activity_client() -> Iterator[tuple[TestClient, _FakeRepository]]: app.dependency_overrides[get_daily_activity_repository] = lambda: repository app.dependency_overrides[get_daily_activity_prisma_client] = lambda: prisma_client + app.dependency_overrides[get_daily_activity_llm_router] = lambda: None app.dependency_overrides[user_api_key_auth] = resolve_auth with TestClient(app) as client: yield client, repository @@ -910,6 +915,59 @@ def test_key_page_route_rejects_invalid_bounds( repository.key_page.assert_not_awaited() +@pytest.mark.parametrize( + ("headers", "expected_users"), + ( + ({}, {"2025-01-01": {"user-a": 12.5, "user-b": 100.0}, "2025-01-02": {"user-a": 2.0}}), + ( + {"x-user-role": LitellmUserRoles.INTERNAL_USER.value, "x-user-id": "user-a"}, + {"2025-01-01": {"user-a": 12.5}, "2025-01-02": {"user-a": 2.0}}, + ), + ), +) +def test_user_aggregate_returns_scoped_spend_per_user_with_labels( + daily_activity_client: tuple[TestClient, _FakeRepository], + headers: dict[str, str], + expected_users: dict[str, dict[str, float]], +) -> None: + client, _ = daily_activity_client + response: Final = client.get("/user/daily/activity/aggregated", params=_DATE_PARAMS, headers=headers) + + assert response.status_code == 200, response.text + body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(response.content) + assert { + day.date.isoformat(): {user_id: entity.metrics.spend for user_id, entity in day.breakdown.entities.items()} + for day in body.results + } == expected_users, response.text + labels: Final = { + "user-a": {"user_email": "user@example.test", "user_alias": "Alias A"}, + "user-b": {"user_email": "other-user@example.test", "user_alias": None}, + } + for day in body.results: + assert {user_id: entity.metadata for user_id, entity in day.breakdown.entities.items()} == { + user_id: labels[user_id] for user_id in expected_users[day.date.isoformat()] + }, response.text + assert sum(sum(users.values()) for users in expected_users.values()) == body.metadata.total_spend + + +def test_user_aggregate_preserves_spend_when_user_metadata_is_missing( + daily_activity_client: tuple[TestClient, _FakeRepository], +) -> None: + client, repository = daily_activity_client + repository.user_metadata.return_value = {} + repository.user_metadata.side_effect = None + response: Final = client.get("/user/daily/activity/aggregated", params={**_DATE_PARAMS, "user_id": "user-a"}) + + assert response.status_code == 200, response.text + body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(response.content) + assert { + day.date.isoformat(): { + user_id: (entity.metrics.spend, entity.metadata) for user_id, entity in day.breakdown.entities.items() + } + for day in body.results + } == {"2025-01-01": {"user-a": (12.5, {})}, "2025-01-02": {"user-a": (2.0, {})}}, response.text + + def test_user_page_route_ranks_users_by_spend_and_attaches_metadata( daily_activity_client: tuple[TestClient, _FakeRepository], ) -> None: diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index d8aa0239e2d..cee0c90f80c 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3,7 +3,7 @@ import hashlib import json import logging from collections.abc import AsyncIterator, Mapping, Sequence -from contextlib import asynccontextmanager +from contextlib import AbstractAsyncContextManager, asynccontextmanager from datetime import datetime, timezone from types import SimpleNamespace from typing import TYPE_CHECKING, Final, cast @@ -130,6 +130,11 @@ async def test_ui_view_users_proxy_admin_no_org_filter(mocker): UserWhereCondition = InsensitiveContains | Sequence[Mapping[str, InsensitiveContains]] +@asynccontextmanager +async def _transaction_over(db: object) -> AsyncIterator[SimpleNamespace]: + yield SimpleNamespace(query_raw=AsyncMock(return_value=[]), litellm_usertable=getattr(db, "litellm_usertable")) + + def _matches_user_where(row: LiteLLM_UserTableFiltered, where: Mapping[str, UserWhereCondition]) -> bool: def matches(field: str, condition: UserWhereCondition) -> bool: if not isinstance(condition, Mapping): @@ -2419,6 +2424,7 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke prisma_client: Final = mocker.MagicMock() prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) + prisma_client.tx = lambda: _transaction_over(prisma_client.db) mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) @@ -2465,6 +2471,7 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) + prisma_client.tx = lambda: _transaction_over(prisma_client.db) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency @@ -2575,6 +2582,9 @@ async def test_user_update_reaches_cached_user_on_every_worker( self.update_data_payload = data return {"user_id": user_id, "data": self.updated_user} + def tx(self) -> AbstractAsyncContextManager[SimpleNamespace]: + return _transaction_over(self.db) + saved_user: Final = LiteLLM_UserTable( user_id="user-spruce", tpm_limit=100000, diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 44c43713438..32c9535b4c8 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -2288,8 +2288,8 @@ class TestTeamModelSiblingRouting: async def test_router_finds_all_sibling_team_deployments(self): """ When two team deployments share team_public_model_name="gpt-4.1-mini", - the router's _common_checks_available_deployment must return BOTH as - healthy_deployments (not collapse to one). + routing that public name for the team must reach BOTH deployments + (not collapse to one). """ import litellm @@ -2338,14 +2338,16 @@ class TestTeamModelSiblingRouting: result = router.map_team_model(public_name, team_id) assert result == public_name - # _common_checks_available_deployment should return both deployments - model, healthy = router._common_checks_available_deployment( - model=public_name, - request_kwargs={"metadata": {"user_api_key_team_id": team_id}}, - ) - assert isinstance(healthy, list) - assert len(healthy) == 2 - api_bases = {d["litellm_params"]["api_base"] for d in healthy} + # routing the public name for the team must reach both siblings + api_bases = set() + for _ in range(30): + response = await router.acompletion( + model=public_name, + messages=[{"role": "user", "content": "hi"}], + mock_response="ok", + metadata={"user_api_key_team_id": team_id}, + ) + api_bases.add(response._hidden_params["api_base"]) assert api_bases == { "https://eastus.openai.azure.com", "https://westus.openai.azure.com", @@ -8567,12 +8569,19 @@ class TestAccessGroupModelSync: return [{"deployment_count": deployment_count}] if sql.startswith('UPDATE "LiteLLM_AccessGroupTable"'): return [{"access_group_id": "ag-1"}] + if sql.startswith(("SELECT 1 AS locked FROM", 'SELECT team_id FROM "LiteLLM_TeamTable"')): + return [] assert sql.startswith("WITH ") return TestAccessGroupModelSync._ALLOWLIST_ROWS + @contextlib.asynccontextmanager + async def transaction(): + yield mock_prisma.db + mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=query_raw) + mock_prisma.db.tx = MagicMock(side_effect=transaction) mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=row) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=row) @@ -11420,3 +11429,46 @@ async def test_add_new_model_reports_a_database_outage_during_the_insert_as_503( ) assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection") + + +@pytest.mark.asyncio +async def test_rename_rewrites_only_in_a_transaction_holding_every_locked_teams_advisory_lock(): + """A team that starts listing the old name between the advisory snapshot and the row lock must not be + rewritten without its advisory lock, or a roster change on that team can still deadlock with the rename.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model + + snapshots = iter(([{"team_id": "team-1"}], [{"team_id": "team-1"}, {"team_id": "team-2"}])) + transactions: list[list[tuple[str, tuple[str, ...]]]] = [] + + async def query_raw(sql: str, *params: str) -> list[dict[str, object]]: + if transactions: + transactions[-1].append((sql, params)) + if sql.startswith("SELECT COUNT(*)"): + return [{"deployment_count": 0}] + if sql.startswith('SELECT team_id FROM "LiteLLM_TeamTable"') and not sql.endswith("FOR UPDATE"): + return next(snapshots) + if sql.startswith('SELECT team_id FROM "LiteLLM_TeamTable"'): + return [{"team_id": "team-1"}, {"team_id": "team-2"}] + return [] + + @contextlib.asynccontextmanager + async def transaction(): + transactions.append([]) + yield db + + db = SimpleNamespace(query_raw=query_raw, tx=transaction) + with patch("litellm.proxy.management_helpers.model_allowlist_rename_sync.evict_and_broadcast", AsyncMock()): + await sync_model_allowlists_for_renamed_model( + SimpleNamespace(db=db), + model_id="deployment", + old_name="old-model", + new_name="new-model", + llm_router=None, + user_api_key_cache=UserApiKeyCache(), + ) + + rewrites = [calls for calls in transactions if any(sql.startswith("WITH ") for sql, _ in calls)] + assert len(rewrites) == 1 + advisory = [params[0] for sql, params in rewrites[0] if "pg_advisory_xact_lock" in sql] + assert advisory == ["team-1", "team-2"] diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 5e06756f3a9..8d2c6c0054c 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1,5 +1,6 @@ import asyncio import json +from contextlib import asynccontextmanager from types import MappingProxyType, SimpleNamespace from typing import Final, Mapping, Optional, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -1560,6 +1561,12 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke prisma_client.db.litellm_jwtkeymapping = jwt_table prisma_client.db.litellm_organizationtable.delete = AsyncMock(return_value=MagicMock()) + @asynccontextmanager + async def transaction(): + yield prisma_client.db + + prisma_client.tx = transaction + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) diff --git a/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py b/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py index cbdd029fc5d..43d8e56274f 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py @@ -11,6 +11,8 @@ from unittest.mock import patch as patch_ctx import pytest from fastapi import HTTPException +import litellm + from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.llms.gemini.cost_calculator import cost_per_web_search_request @@ -1047,9 +1049,9 @@ class TestPtuDeploymentsAreNotBilledPerToken: def test_the_cost_map_contributes_no_price_to_a_priced_ptu_deployment(self, backend): """The acceptance criterion, read off the entry the router registers for the deployment. - Zeroing only the per-token pair leaves the cache-tier fields unset, which is exactly what - Router._inherit_builtin_cache_pricing back-fills from the public cost map, so a cached - prompt would still be billed at the public rate.""" + Zeroing only the per-token pair leaves the cache-tier fields unset, which the router + back-fills from the public cost map, so a cached prompt would still be billed at the + public rate.""" priced = _ptu_priced_deployment( Deployment( model_name="ptu-deployment", @@ -1062,7 +1064,8 @@ class TestPtuDeploymentsAreNotBilledPerToken: ), ) ) - registered = Router._deployment_model_cost_payload(priced) + Router(model_list=[priced.to_json(exclude_none=True)]) + registered = litellm.model_cost["dep-ptu"] charged = { k: v for k, v in registered.items() diff --git a/tests/unit/proxy/middleware/test_request_body_decompression_middleware.py b/tests/unit/proxy/middleware/test_request_body_decompression_middleware.py index 48981288d00..dbe01804151 100644 --- a/tests/unit/proxy/middleware/test_request_body_decompression_middleware.py +++ b/tests/unit/proxy/middleware/test_request_body_decompression_middleware.py @@ -1,6 +1,7 @@ import asyncio import base64 import gzip +import hashlib import json import math import random @@ -59,6 +60,8 @@ def _client(max_decompressed_body_size_mb: float = CAP_MB, root_path: str = "", routes=[ Route("/v1/systemone", _echo_request, methods=["POST"]), Route("/claude-code/plugins/upload", _echo_request, methods=["POST"]), + Route("/skill-hub/skills/upload", _echo_request, methods=["POST"]), + Route("/skill-hub/skills/{name}/upload", _echo_request, methods=["PUT"]), Route("/plugin-proxy/{path:path}", _echo_request, methods=["POST"]), Route("/claude_code_gateway/v1/{signal}", _echo_request, methods=["POST"]), ] @@ -248,16 +251,21 @@ def test_operator_request_size_limit_caps_the_decoded_body_when_it_is_lower( @pytest.mark.parametrize(("coding", "compress"), [ ("identity", bytes), ("gzip", gzip.compress), ("zstd", zstd.compress), ]) +@pytest.mark.parametrize(("method", "path"), [ + ("POST", "/claude-code/plugins/upload"), + ("POST", "/skill-hub/skills/upload"), + ("PUT", "/skill-hub/skills/example/upload"), +]) def test_proxy_bounds_skill_uploads_before_authentication_or_body_parsing( - coding: str, compress: Callable[[bytes], bytes], + coding: str, compress: Callable[[bytes], bytes], method: str, path: str, ) -> None: from litellm.proxy.proxy_server import app body: Final = random.Random(0).randbytes(6 * MIB) + bytes(6 * MIB) encoded: Final = compress(body) - response: Final = TestClient(app).post( - "/claude-code/plugins/upload", + response: Final = TestClient(app).request( + method, path, content=encoded, headers={"content-encoding": coding, "content-type": "multipart/form-data; boundary=upload"}, ) @@ -269,6 +277,32 @@ def test_proxy_bounds_skill_uploads_before_authentication_or_body_parsing( ) +@pytest.mark.parametrize("root_path", ("", "/gateway")) +@pytest.mark.parametrize(("coding", "compress"), [("gzip", gzip.compress), ("zstd", zstd.compress)]) +def test_replacement_upload_route_limit_caps_decoded_bytes( + root_path: str, coding: str, compress: Callable[[bytes], bytes], +) -> None: + limit: Final = 1 / 32 + client: Final = _client( + root_path=root_path, + route_max_request_size_mb={("PUT", "/skill-hub/skills/{name}/upload"): limit}, + ) + body: Final = _low_entropy_text(int(limit * MIB)) + path: Final = f"{root_path}/skill-hub/skills/example/upload" + headers: Final = {"content-encoding": coding, "content-type": "application/zip"} + + within: Final = client.put(path, content=compress(body), headers=headers) + over: Final = client.put(path, content=compress(body + b"x"), headers=headers) + + assert within.status_code == 200 + decoded: Final = bytes.fromhex(within.json()["body_hex"]) + actual: Final = (len(decoded), hashlib.sha256(decoded).hexdigest()) + expected: Final = (len(body), hashlib.sha256(body).hexdigest()) + assert actual == expected + assert over.status_code == 413 + assert over.json() == _decoded_too_large(limit) + + @pytest.mark.parametrize(("signature", "content_encoding"), [(GZIP_MAGIC, "gzip"), (ZSTD_FRAME_MAGIC, "zstd")]) def test_compressed_body_over_the_cap_is_rejected_before_decoding(signature: bytes, content_encoding: str): over_cap_on_the_wire: Final = signature + bytes(int(CAP_MB * MIB)) diff --git a/tests/unit/proxy/middleware/test_request_size_limit_middleware.py b/tests/unit/proxy/middleware/test_request_size_limit_middleware.py index 7bb73b77000..6d8d8363220 100644 --- a/tests/unit/proxy/middleware/test_request_size_limit_middleware.py +++ b/tests/unit/proxy/middleware/test_request_size_limit_middleware.py @@ -12,15 +12,20 @@ from litellm.proxy.middleware.request_size_limit_middleware import ( ) -@pytest.mark.parametrize(("path", "root_path", "operator_mb", "route_mb"), [ - ("/chat/completions", "", 1, None), - ("/claude-code/plugins/upload", "", None, 1), - ("/gateway/claude-code/plugins/upload", "/gateway", None, 1), - ("/claude-code/plugins/upload", "", 2, 1), - ("/claude-code/plugins/upload", "", 1, 2), +@pytest.mark.parametrize(("method", "path", "root_path", "operator_mb", "route_mb"), [ + ("POST", "/chat/completions", "", 1, None), + ("POST", "/claude-code/plugins/upload", "", None, 1), + ("POST", "/skill-hub/skills/upload", "", None, 1), + ("POST", "/gateway/claude-code/plugins/upload", "/gateway", None, 1), + ("POST", "/claude-code/plugins/upload", "", 2, 1), + ("POST", "/claude-code/plugins/upload", "", 1, 2), + ("PUT", "/skill-hub/skills/example/upload", "", None, 1), + ("PUT", "/gateway/skill-hub/skills/example/upload", "/gateway", None, 1), + ("PUT", "/skill-hub/skills/example/upload", "/gateway", None, 1), + ("PUT", "/skill-hub/skills/example/upload", "", 1, 2), ]) def test_request_size_limit_middleware_rejects_content_length_before_body_read( - path: str, root_path: str, operator_mb: int | None, route_mb: int | None, + method: str, path: str, root_path: str, operator_mb: int | None, route_mb: int | None, ) -> None: downstream_called = False @@ -35,13 +40,17 @@ def test_request_size_limit_middleware_rejects_content_length_before_body_read( app, get_max_request_size_mb=lambda: operator_mb, is_request_size_limit_enabled=lambda: operator_mb is not None, - route_max_request_size_mb={("POST", "/claude-code/plugins/upload"): route_mb} if route_mb else None, + route_max_request_size_mb={ + ("POST", "/claude-code/plugins/upload"): route_mb, + ("POST", "/skill-hub/skills/upload"): route_mb, + ("PUT", "/skill-hub/skills/{name}/upload"): route_mb, + } if route_mb else None, ), root_path=root_path, ) - response = client.post( - path, + response = client.request( + method, path, content=b"x" * (1024 * 1024 + 1), headers={"content-type": "application/json"}, ) @@ -53,14 +62,19 @@ def test_request_size_limit_middleware_rejects_content_length_before_body_read( @pytest.mark.parametrize(("enabled", "operator_mb"), [(True, 0), (False, 1)]) -@pytest.mark.parametrize(("method", "path"), [ - ("POST", "/chat/completions"), - ("PUT", "/claude-code/plugins/upload"), - ("POST", "/claude-code/plugins/upload-other"), - ("POST", "/claude-code/plugins"), +@pytest.mark.parametrize(("method", "path", "root_path"), [ + ("POST", "/chat/completions", ""), + ("PUT", "/claude-code/plugins/upload", ""), + ("POST", "/claude-code/plugins/upload-other", ""), + ("POST", "/claude-code/plugins", ""), + ("POST", "/skill-hub/skills/example/upload", ""), + ("PUT", "/skill-hub/skills/example/upload-other", ""), + ("PUT", "/skill-hub/skills/example/nested/upload", ""), + ("PUT", "/skill-hub/skills//upload", ""), + ("PUT", "/gateway-other/skill-hub/skills/example/upload", "/gateway"), ]) def test_request_size_limit_middleware_zero_or_disabled_global_limit_preserves_other_routes( - enabled: bool, operator_mb: int, method: str, path: str, + enabled: bool, operator_mb: int, method: str, path: str, root_path: str, ) -> None: downstream_called = False @@ -75,8 +89,12 @@ def test_request_size_limit_middleware_zero_or_disabled_global_limit_preserves_o app, get_max_request_size_mb=lambda: operator_mb, is_request_size_limit_enabled=lambda: enabled, - route_max_request_size_mb={("POST", "/claude-code/plugins/upload"): 1}, - ) + route_max_request_size_mb={ + ("POST", "/claude-code/plugins/upload"): 1, + ("PUT", "/skill-hub/skills/{name}/upload"): 1, + }, + ), + root_path=root_path, ) response = client.request( @@ -91,7 +109,15 @@ def test_request_size_limit_middleware_zero_or_disabled_global_limit_preserves_o @pytest.mark.asyncio -async def test_request_size_limit_middleware_rejects_streamed_body_without_content_length(): +@pytest.mark.parametrize(("method", "path", "root_path", "operator_mb"), [ + ("POST", "/chat/completions", "", 1), + ("POST", "/skill-hub/skills/upload", "", None), + ("PUT", "/skill-hub/skills/example/upload", "", None), + ("PUT", "/gateway/skill-hub/skills/example/upload", "/gateway", None), +]) +async def test_request_size_limit_middleware_rejects_streamed_body_without_content_length( + method: str, path: str, root_path: str, operator_mb: int | None, +) -> None: received_body_bytes = 0 async def app(scope, receive, send): @@ -109,8 +135,11 @@ async def test_request_size_limit_middleware_rejects_streamed_body_without_conte middleware = RequestSizeLimitMiddleware( app, - get_max_request_size_mb=lambda: 1, - is_request_size_limit_enabled=lambda: True, + get_max_request_size_mb=lambda: operator_mb, + is_request_size_limit_enabled=lambda: operator_mb is not None, + route_max_request_size_mb={ + ("POST", "/skill-hub/skills/upload"): 1, ("PUT", "/skill-hub/skills/{name}/upload"): 1, + }, ) sent_messages: list[Message] = [] receive_messages: list[Message] = [ @@ -135,8 +164,9 @@ async def test_request_size_limit_middleware_rejects_streamed_body_without_conte await middleware( { "type": "http", - "method": "POST", - "path": "/chat/completions", + "method": method, + "path": path, + "root_path": root_path, "headers": [(b"content-type", b"application/json")], }, receive, @@ -180,9 +210,14 @@ def test_an_oversized_messages_request_gets_the_anthropic_request_too_large_erro @pytest.mark.asyncio -@pytest.mark.parametrize("path", ("/claude-code/plugins/upload", "/chat/completions")) +@pytest.mark.parametrize(("method", "path", "route"), [ + ("POST", "/claude-code/plugins/upload", "/claude-code/plugins/upload"), + ("POST", "/skill-hub/skills/upload", "/skill-hub/skills/upload"), + ("POST", "/chat/completions", "/chat/completions"), + ("PUT", "/skill-hub/skills/example/upload", "/skill-hub/skills/{name}/upload"), +]) async def test_chunked_upload_overflow_stops_fastapi_and_auth_multipart_parsing( - path: str, + method: str, path: str, route: str, ) -> None: from litellm.proxy.proxy_server import app @@ -190,7 +225,7 @@ async def test_chunked_upload_overflow_stops_fastapi_and_auth_multipart_parsing( app, get_max_request_size_mb=lambda: None, is_request_size_limit_enabled=lambda: False, - route_max_request_size_mb={("POST", path): 1 / 1024}, + route_max_request_size_mb={(method, route): 1 / 1024}, ) async def chunks() -> AsyncIterator[bytes]: @@ -199,8 +234,8 @@ async def test_chunked_upload_overflow_stops_fastapi_and_auth_multipart_parsing( pytest.fail("The server consumed more of the upload after its byte limit") async with httpx.AsyncClient(transport=httpx.ASGITransport(app=limited), base_url="http://test") as client: - response: Final = await client.post( - path, + response: Final = await client.request( + method, path, content=chunks(), headers={"content-type": "multipart/form-data; boundary=skill"}, ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py index e654c6d444b..60b263948b6 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py @@ -79,10 +79,7 @@ def test_initialize_deployment_for_pass_through_success(reusable_credentials): ) # Test the initialization - router._initialize_deployment_for_pass_through( - deployment=deployment, - custom_llm_provider="vertex_ai", - ) + router.add_deployment(deployment) # Verify the credentials were properly set from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -114,10 +111,7 @@ def test_initialize_deployment_for_pass_through_missing_params(): ValueError, match="vertex_project, and vertex_location must be set in litellm_params for pass-through endpoints", ): - router._initialize_deployment_for_pass_through( - deployment=deployment, - custom_llm_provider="vertex_ai", - ) + router.add_deployment(deployment) def test_initialize_deployment_when_pass_through_disabled(): @@ -133,10 +127,7 @@ def test_initialize_deployment_when_pass_through_disabled(): ) # This should exit without error, even with missing vertex parameters - router._initialize_deployment_for_pass_through( - deployment=deployment, - custom_llm_provider="vertex_ai", - ) + router.add_deployment(deployment) # If we reach this point, the test passes as the method exited without raising any errors assert True diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index c17a9ba9371..dd3ffad4727 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -1797,25 +1797,40 @@ def test_get_agent_fields_appends_inherited_provider_credential_fields(): @pytest.mark.parametrize("preference", ["auto_install", None]) -def test_public_skill_hub_echoes_registered_installation_preference(preference): +def test_public_skill_hub_echoes_registered_installation_preference(preference: str | None) -> None: + from collections.abc import Mapping + + from prisma.models import LiteLLM_ClaudeCodePluginTable + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import register_plugin + from litellm.repositories.plugin_upload_repository import PluginUploadRepository from litellm.types.proxy.claude_code_endpoints import RegisterPluginRequest + def create(data: Mapping[str, object]) -> LiteLLM_ClaudeCodePluginTable: + defaults: Final = {"id": "plugin-id", "enabled": True, "anonymous_download_count": 0} + return LiteLLM_ClaudeCodePluginTable.model_validate({**defaults, **data}) + + def find_many(where: Mapping[str, object]) -> list[LiteLLM_ClaudeCodePluginTable]: + return [create(table.create.call_args.kwargs["data"])] + table: Final = MagicMock() table.find_unique = AsyncMock(return_value=None) - table.create = AsyncMock(side_effect=lambda data: SimpleNamespace(id="plugin-id", **data)) - table.find_many = AsyncMock( - side_effect=lambda where: [SimpleNamespace(id="plugin-id", **table.create.call_args.kwargs["data"])] - ) + table.create = AsyncMock(side_effect=create) + table.find_many = AsyncMock(side_effect=find_many) prisma: Final = MagicMock() prisma.db.litellm_claudecodeplugintable = table + prisma.writer_db = prisma.db + prisma.db.tx.return_value.__aenter__.return_value = prisma.db + prisma.db.query_raw = AsyncMock(return_value=[]) + prisma.db.litellm_claudecodepluginversion.create = AsyncMock() app: Final = FastAPI() app.include_router(router) with patch("litellm.proxy.proxy_server.prisma_client", prisma): asyncio.run( register_plugin( + uploads=PluginUploadRepository(prisma), request=RegisterPluginRequest( name="security-review", source={ diff --git a/tests/unit/proxy/skill_submission_rows.py b/tests/unit/proxy/skill_submission_rows.py new file mode 100644 index 00000000000..90fc54d0969 --- /dev/null +++ b/tests/unit/proxy/skill_submission_rows.py @@ -0,0 +1,76 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Final + +from prisma import models as prisma_models + + +class SubmissionRows: + def __init__(self) -> None: + self.rows: tuple[prisma_models.LiteLLM_SkillSubmission, ...] = () + + async def create(self, data: Mapping[str, object]) -> prisma_models.LiteLLM_SkillSubmission: + row: Final = prisma_models.LiteLLM_SkillSubmission.model_validate( + { + "id": f"submission-{len(self.rows)}", + "status": "pending_review", + "created_at": datetime(2026, 1, 1, tzinfo=timezone.utc), + **data, + } + ) + self.rows = (*self.rows, row) + return row + + async def find_unique(self, where: Mapping[str, object]) -> prisma_models.LiteLLM_SkillSubmission | None: + return next((row for row in self.rows if row.id == where["id"]), None) + + async def find_first(self, where: Mapping[str, object]) -> prisma_models.LiteLLM_SkillSubmission | None: + row: Final = await self.find_unique(where) + return None if row is None or where.get("submitted_by", row.submitted_by) != row.submitted_by else row + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> prisma_models.LiteLLM_SkillSubmission | None: + current: Final = await self.find_unique(where) + if current is None: + return None + updated: Final = prisma_models.LiteLLM_SkillSubmission.model_validate({**current.model_dump(), **data}) + self.rows = tuple(updated if row.id == current.id else row for row in self.rows) + return updated + + async def query_raw(self, query: str, *args: object) -> tuple[Mapping[str, object], ...]: + if query.startswith("SELECT scan_json"): + return tuple({"scan_json": row.scan_json} for row in self.rows if row.id == args[0]) + if 'FROM "LiteLLM_SkillSubmission"' not in query or "FOR UPDATE" in query: + return () + if "COUNT(*) AS count" in query: + pending: Final = tuple( + row for row in self.rows if row.submitted_by == args[0] and row.status == "pending_review" + ) + return ( + { + "count": len(pending), + "storage_bytes": sum( + len(row.proposal_json.encode("utf-8")) + + (0 if row.archive_content is None else len(row.archive_content.decode())) + for row in pending + ), + }, + ) + records: Final = ( + tuple(row for row in self.rows if row.id == args[0] and (args[1] is None or row.submitted_by == args[1])) + if "WHERE id = $1" in query + else tuple(row for row in self.rows if args[0] is None or row.submitted_by == args[0]) + ) + if "WHERE id = $1" in query: + return tuple({**row.model_dump(), "has_archive": row.archive_content is not None} for row in records) + limit, offset = args[2:] + assert isinstance(limit, int) and isinstance(offset, int) + filtered: Final = tuple(row for row in records if args[1] is None or row.status == args[1]) + ordered: Final = sorted(filtered, key=lambda row: (row.created_at, row.id), reverse=True) + return tuple( + {**row.model_dump(), "has_archive": row.archive_content is not None} + for row in ordered[offset : offset + limit] + ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index e119acd5ffc..c3e916d581d 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3819,7 +3819,10 @@ class TestSpendLogsPayload: await _wait_for_mock_call(mock_client) kwargs = mock_client.call_args.kwargs - payload: SpendLogsPayload = kwargs["payload"] + payload: SpendLogsPayload = { + key: value for key, value in kwargs["payload"].items() if not key.startswith("accounting_") + } + assert kwargs["payload"]["accounting_payload"]["spend"] == kwargs["payload"]["spend"] expected_payload = SpendLogsPayload( **{ "request_id": "chatcmpl-34df56d5-4807-45c1-bb99-61e52586b802", diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 6ce8d400daf..f49c04319aa 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -3057,6 +3057,7 @@ async def test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_c @pytest.mark.asyncio async def test_streaming_cancel_while_holding_back_provider_output_keeps_reservation( spend_counter_state, + monkeypatch: pytest.MonkeyPatch, ): counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) @@ -3086,9 +3087,17 @@ async def test_streaming_cancel_while_holding_back_provider_output_keeps_reserva } ] ) - response = await router._aanthropic_messages_streaming_iterator( - response=AnthropicMessagesStreamingResponse(completion_stream=held_back, hidden_params={"additional_headers": {}}), - initial_kwargs={"model": "claude-haiku-4-5"}, + provider_response = AnthropicMessagesStreamingResponse( + completion_stream=held_back, hidden_params={"additional_headers": {}} + ) + monkeypatch.setattr( # test-quality-ok: the provider call the router streams from + litellm, "anthropic_messages", AsyncMock(return_value=provider_response) + ) + response = await router.aanthropic_messages( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "hi"}], + max_tokens=16, + stream=True, ) async def ping_then_cancel(user_api_key_dict, response, request_data): diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 8143636a435..b68c616457b 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9027,7 +9027,7 @@ class TestStreamingClientDisconnectBilling: from litellm.llms.anthropic.messages.adapters.transformation import ( AnthropicAdapter, ) - from litellm.router import FallbackAwareAnthropicMessagesStream + from litellm.router_rust.invoke.streaming import FallbackAwareAnthropicMessagesStream, StreamAttempt async def _sse_frames() -> AsyncGenerator[bytes, None]: yield b"event: message_start\n\n" @@ -9046,7 +9046,7 @@ class TestStreamingClientDisconnectBilling: ) assert isinstance(source_iterator, AnthropicSSEStream) streamed: Final = prepare_response_for_header_attachment( - FallbackAwareAnthropicMessagesStream(_sse_frames(), source_iterator) + FallbackAwareAnthropicMessagesStream(_sse_frames(), source_iterator, MagicMock(spec=StreamAttempt)) ) billed: Final = await _bill_partial_streamed_spend_on_disconnect( @@ -9586,7 +9586,7 @@ def _apply_stream_usage_tracking( from litellm.llms.anthropic.messages.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, ) - from litellm.router import FallbackAwareAnthropicMessagesStream + from litellm.router_rust.invoke.streaming import FallbackAwareAnthropicMessagesStream, StreamAttempt from litellm.types.llms.openai import ResponsesAPIResponse async def _sse_frames() -> AsyncGenerator[bytes, None]: @@ -9595,7 +9595,9 @@ def _apply_stream_usage_tracking( recorder: Final = _RecordingSuccessLogger() stream: Final = await self._start_partial_responses_stream(recorder) wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=stream, model="gpt-4o-mini") - streamed: Final = FallbackAwareAnthropicMessagesStream(_sse_frames(), wrapper.async_anthropic_sse_wrapper()) + streamed: Final = FallbackAwareAnthropicMessagesStream( + _sse_frames(), wrapper.async_anthropic_sse_wrapper(), MagicMock(spec=StreamAttempt) + ) billed: Final = await self._bill_and_wait(recorder, streamed, stream.logging_obj) assert billed is True diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..42e828f267e 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -9213,7 +9213,9 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_on_resp @pytest.mark.asyncio -async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_strip(): +async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_strip( + monkeypatch: pytest.MonkeyPatch, +): """Regression for the #38586 break: a client that planted a reserved key in litellm_metadata made the router hand downstream a scrubbed copy, so the proxy's post_call write-backs (guardrail telemetry, applied guardrails) landed in a dict the @@ -9247,13 +9249,13 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_ ] ) forwarded_buckets = [] - original_acompletion = router._acompletion + original_acompletion = litellm.acompletion async def _spy(*args, **spy_kwargs): forwarded_buckets.append(spy_kwargs["metadata"]) return await original_acompletion(*args, **spy_kwargs) - router._acompletion = _spy + monkeypatch.setattr(litellm, "acompletion", _spy) await router.acompletion(**request_data) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 62d8db6f810..8536cbaa399 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -5635,6 +5635,13 @@ def _routing_groups_router(): ) +def _routing_groups_of(router: litellm.Router) -> dict[str, tuple[list[str], str]]: + return { + group["group_name"]: (group["models"], group["routing_strategy"]) + for group in router.settings().as_dict()["routing_groups"] or [] + } + + @pytest.mark.asyncio async def test_invalid_db_routing_groups_do_not_abort_other_router_settings(): from unittest.mock import AsyncMock, MagicMock @@ -5656,8 +5663,7 @@ async def test_invalid_db_routing_groups_do_not_abort_other_router_settings(): await ProxyConfig()._add_router_settings_from_db_config(llm_router=router, prisma_client=mock_prisma_client) assert router.num_retries == 7 - assert router._model_to_group == {"m1": "g1"} - assert router._get_routing_context("m1", None)[0] == "latency-based-routing" + assert _routing_groups_of(router) == {"g1": (["m1"], "latency-based-routing")} @pytest.mark.asyncio @@ -5678,8 +5684,7 @@ async def test_valid_db_routing_groups_still_replace_router_groups(): await ProxyConfig()._add_router_settings_from_db_config(llm_router=router, prisma_client=mock_prisma_client) assert router.num_retries == 7 - assert router._model_to_group == {"m2": "g2"} - assert router._get_routing_context("m2", None)[0] == "least-busy" + assert _routing_groups_of(router) == {"g2": (["m2"], "least-busy")} def _content_policy_router(**router_kwargs: object) -> litellm.Router: diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 71f816ea183..cb903fb134a 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1522,6 +1522,37 @@ def test_raise_if_required_body_param_missing_rejects_missing_param(route_type, assert exc_info.value.message == f"{route}: Missing required parameter: '{param}'." +@pytest.mark.parametrize( + "data", + [ + {"model": "gemini-4-argon", "input": "prove the prime number theorem"}, + {"model": "gemini-4-argon", "previous_interaction_id": "v1_abc", "continuation_token": "tok-1"}, + {"model": "gemini-4-argon", "input": "prove it again", "continuation_token": "tok-1"}, + ], +) +def test_create_interaction_accepts_input_or_a_continuation_token(data: dict[str, object]) -> None: + from litellm.proxy.route_llm_request import raise_if_required_body_param_missing + + assert raise_if_required_body_param_missing(route_type="acreate_interaction", data=data, llm_router=None) is None + + +def test_create_interaction_without_input_or_continuation_token_names_input() -> None: + from litellm.proxy.route_llm_request import ( + ProxyMissingRequiredParamError, + raise_if_required_body_param_missing, + ) + + with pytest.raises(ProxyMissingRequiredParamError) as exc_info: + raise_if_required_body_param_missing( + route_type="acreate_interaction", + data={"model": "gemini-4-argon", "previous_interaction_id": "v1_abc"}, + llm_router=None, + ) + + assert exc_info.value.param == "input" + assert exc_info.value.message == "/interactions: Missing required parameter: 'input'." + + @pytest.mark.parametrize( "data, param", [ diff --git a/tests/unit/proxy/test_unit_test_proxy_hooks.py b/tests/unit/proxy/test_unit_test_proxy_hooks.py index 6ce7154f2d1..7990ec12037 100644 --- a/tests/unit/proxy/test_unit_test_proxy_hooks.py +++ b/tests/unit/proxy/test_unit_test_proxy_hooks.py @@ -11,7 +11,7 @@ import litellm @pytest.mark.asyncio async def test_disable_spend_logs(): """ - Test that the spend logs are not written to the database when disable_spend_logs is True + With disable_spend_logs the request is queued only for its charge, flagged so no SpendLogs row is written """ # Mock the necessary components @@ -41,5 +41,4 @@ async def test_disable_spend_logs(): org_id="org_id", kwargs={}, ) - # Verify no spend logs were added - assert len(mock_prisma_client.spend_log_transactions) == 0 + assert [row["accounting_log_disabled"] for row in mock_prisma_client.spend_log_transactions] == [True] diff --git a/tests/unit/proxy/utils/prisma_and_spend/conftest.py b/tests/unit/proxy/utils/prisma_and_spend/conftest.py index afc282a14d4..888f664b638 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/unit/proxy/utils/prisma_and_spend/conftest.py @@ -142,7 +142,9 @@ def mock_prisma_client() -> MagicMock: client.db.stop_token_refresh_task = AsyncMock() client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}]) client.db.execute_raw = AsyncMock() - client.db.tx = MagicMock() + client.db.tx = MagicMock( + return_value=MagicMock(__aenter__=AsyncMock(return_value=client.db), __aexit__=AsyncMock(return_value=False)) + ) client.db.batch_ = MagicMock() for table_name in _PRISMA_TABLES: setattr(client.db, table_name, _make_table_mock()) diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 926566605a5..73d3414abb0 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -60,45 +60,31 @@ def test_direct_vector_store_embedding_executor_rejects_invalid_value(): def test_router_vector_store_search_injects_executor_and_request_metadata(): router = litellm.Router(model_list=[]) - original = MagicMock(return_value="searched") - wrapped = router.factory_function(original, call_type="vector_store_search") - assert ( - wrapped( - vector_store_id="store", - query="query", - custom_llm_provider="valkey", - litellm_metadata={"user_api_key_team_id": "team-a"}, + with patch( # test-quality-ok: isolates the SDK search the router dispatches to + "litellm.vector_stores.main.search", return_value="searched" + ) as search: + assert ( + router.vector_store_search( + vector_store_id="store", + query="query", + custom_llm_provider="valkey", + litellm_metadata={"user_api_key_team_id": "team-a"}, + ) + == "searched" ) - == "searched" - ) - call_kwargs = original.call_args.kwargs + call_kwargs = search.call_args.kwargs assert call_kwargs["custom_llm_provider"] == "valkey" executor = call_kwargs["_direct_vector_store_embedding_executor"] assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) assert executor.metadata == {"user_api_key_team_id": "team-a"} - assert litellm.Router._vector_store_request_metadata({"metadata": {"user_api_key_team_id": "team-b"}}) == { - "user_api_key_team_id": "team-b" - } - assert litellm.Router._vector_store_request_metadata({}) == {} - with patch.object( # test-quality-ok: fallback dispatch is the boundary this wrapper delegates to - router, "_generic_api_call_with_fallbacks", return_value="routed" - ) as fallback: - assert wrapped(model="vector-alias", vector_store_id="store", query="query") == "routed" - assert fallback.call_args.kwargs["model"] == "vector-alias" - assert fallback.call_args.kwargs["original_function"] is original - - create_original = MagicMock(return_value="created") - wrapped_create = router.factory_function(create_original, call_type="vector_store_create") - assert wrapped_create(name="store") == "created" - create_original.assert_called_once_with(name="store") - with patch.object( # test-quality-ok: fallback dispatch is the boundary this wrapper delegates to - router, "_generic_api_call_with_fallbacks", return_value="created-through-router" - ) as fallback: - assert wrapped_create(model="vector-alias", name="store") == "created-through-router" - fallback.assert_called_once_with(original_function=create_original, model="vector-alias", name="store") + with patch( # test-quality-ok: isolates the SDK create the router dispatches to + "litellm.vector_stores.main.create", return_value="created" + ) as create: + assert router.vector_store_create(name="store") == "created" + create.assert_called_once_with(name="store") @pytest.mark.asyncio @@ -171,15 +157,16 @@ async def test_router_avector_store_search_passes_correct_args(): # Create a router router = litellm.Router(model_list=[]) - # Mock the router's _init_vector_store_api_endpoints method to avoid real API calls - with patch.object(router, "_init_vector_store_api_endpoints") as mock_init: - mock_init.return_value = { - "object": "vector_store.search_results.page", - "search_query": "test query", - "data": [], - } - - # Call router's avector_store_search + with patch( # test-quality-ok: isolates the SDK search the router dispatches to + "litellm.vector_stores.main.asearch", + new=AsyncMock( + return_value={ + "object": "vector_store.search_results.page", + "search_query": "test query", + "data": [], + } + ), + ) as asearch: await router.avector_store_search( vector_store_id="test_store_id", query="test query", @@ -187,17 +174,14 @@ async def test_router_avector_store_search_passes_correct_args(): metadata={"user_api_key_team_id": "team-a"}, ) - # Verify the internal method was called with correct args - mock_init.assert_called_once() - call_args = mock_init.call_args - - # Check that the original function is passed correctly - assert call_args[1]["vector_store_id"] == "test_store_id" - assert call_args[1]["query"] == "test query" - assert call_args[1]["custom_llm_provider"] == "bedrock" - executor = call_args[1]["_direct_vector_store_embedding_executor"] - assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) - assert executor.metadata["user_api_key_team_id"] == "team-a" + asearch.assert_awaited_once() + call_kwargs = asearch.call_args.kwargs + assert call_kwargs["vector_store_id"] == "test_store_id" + assert call_kwargs["query"] == "test query" + assert call_kwargs["custom_llm_provider"] == "bedrock" + executor = call_kwargs["_direct_vector_store_embedding_executor"] + assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) + assert executor.metadata["user_api_key_team_id"] == "team-a" @pytest.mark.asyncio diff --git a/tests/unit/repositories/test_plugin_publication.py b/tests/unit/repositories/test_plugin_publication.py new file mode 100644 index 00000000000..0a784553d09 --- /dev/null +++ b/tests/unit/repositories/test_plugin_publication.py @@ -0,0 +1,149 @@ +import hashlib +import json +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final, cast +from unittest.mock import AsyncMock, MagicMock + +import pytest +from prisma import Base64 +from prisma.models import ( + LiteLLM_ClaudeCodePluginTable, + LiteLLM_ClaudeCodePluginUpload, +) + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PublishedPlugin, +) +from litellm.repositories.plugin_publication import ( + PluginPublicationTransaction, + publish_plugin, + resolve_plugin_version, +) + + +@pytest.mark.parametrize("requested,current,updating,expected", [ + (None, None, False, "1.0.0"), + (None, "1.0.0", True, "1.0.1"), + (None, "1.2.3-rc.1+build", True, "1.2.4"), + ("2.0.0", "1.9.9", True, "2.0.0"), + ("1.0.0-rc.10", "1.0.0-rc.9", True, "1.0.0-rc.10"), + ("1.0.0", "1.0.0-rc.1", True, "1.0.0"), + ("1.0.0", None, True, "1.0.0"), + ("2.0.0", "legacy", True, "2.0.0"), +]) +def test_published_version_resolution( + requested: str | None, current: str | None, updating: bool, expected: str, +) -> None: + assert resolve_plugin_version(requested, current, updating=updating) == expected + + +@pytest.mark.parametrize("requested,current,kind", [ + ("1.0.0", "1.0.0", "invalid_version"), + ("0.9.9", "1.0.0", "invalid_version"), + ("1.0.0+new", "1.0.0+old", "invalid_version"), + ("1.0.0-rc.1", "1.0.0", "invalid_version"), + ("1.0.0-rc.9", "1.0.0-rc.10", "invalid_version"), + ("1.2", "1.0.0", "invalid_version"), + (None, None, "conflict"), + (None, "legacy", "conflict"), +]) +def test_invalid_or_non_increasing_versions_cannot_publish( + requested: str | None, current: str | None, kind: str, +) -> None: + result: Final = resolve_plugin_version(requested, current, updating=True) + assert isinstance(result, PluginPublicationError) + assert result.kind == kind + + +@pytest.fixture +def publication_tx() -> MagicMock: + content: Final = b"original uploaded bytes" + tx: Final = MagicMock() + tx.query_raw = AsyncMock(return_value=()) + tx.litellm_claudecodeplugintable.find_unique = AsyncMock(return_value=LiteLLM_ClaudeCodePluginTable( + id="skill-id", name="skill", version="1.0.0", enabled=False, anonymous_download_count=0, + updated_at=datetime(2025, 1, 1, tzinfo=timezone.utc), + manifest_json=json.dumps({"source": {"source": "archive", "sha256": hashlib.sha256(content).hexdigest()}}), + )) + tx.litellm_claudecodepluginupload.find_unique = AsyncMock(return_value=LiteLLM_ClaudeCodePluginUpload( + plugin_name="skill", content=Base64.encode(content), + )) + tx.litellm_claudecodepluginversion.count = AsyncMock(return_value=0) + for table in (tx.litellm_claudecodeplugintable, tx.litellm_claudecodepluginupload, tx.litellm_claudecodepluginversion): + for operation in ("create", "update", "upsert", "delete_many"): + setattr(table, operation, AsyncMock()) + return tx + + +@pytest.mark.parametrize("failure,kind", [ + ("duplicate", "conflict"), ("missing", "not_found"), ("invalid", "invalid_version"), + ("digest", "conflict"), ("identity", "conflict"), ("timestamp", "conflict"), +]) +@pytest.mark.asyncio +async def test_rejected_publication_never_writes_catalog_archive_or_history( + publication_tx: MagicMock, failure: str, kind: str, +) -> None: + current: Final = LiteLLM_ClaudeCodePluginTable.model_validate( + publication_tx.litellm_claudecodeplugintable.find_unique.return_value, + ) + publication_tx.litellm_claudecodeplugintable.find_unique.return_value = None if failure == "missing" else current + expected: Final = None if failure == "duplicate" else PluginPublicationBase( + "old-id" if failure == "identity" else current.id, + current.updated_at - timedelta(seconds=1) if failure == "timestamp" else current.updated_at, + ) + result: Final = await publish_plugin( + cast(PluginPublicationTransaction, publication_tx), + PluginPublication("skill", "0.0.1" if failure == "invalid" else None, "rejected", + json.loads(current.manifest_json), "admin", b"mismatched" if failure == "digest" else None), + expected, + ) + assert isinstance(result, PluginPublicationError) and result.kind == kind + for table in (publication_tx.litellm_claudecodeplugintable, publication_tx.litellm_claudecodepluginupload, + publication_tx.litellm_claudecodepluginversion): + for operation in ("create", "update", "upsert", "delete_many"): + getattr(table, operation).assert_not_awaited() + + +@pytest.mark.parametrize("valid_archive", (True, False)) +@pytest.mark.asyncio +async def test_first_legacy_update_preserves_its_original_snapshot( + publication_tx: MagicMock, valid_archive: bool, +) -> None: + current: Final = LiteLLM_ClaudeCodePluginTable.model_validate( + publication_tx.litellm_claudecodeplugintable.find_unique.return_value, + ).model_copy(update={"version": None}) + publication_tx.litellm_claudecodeplugintable.find_unique.return_value = current + content: Final = b"original uploaded bytes" if valid_archive else b"corrupt legacy bytes" + publication_tx.litellm_claudecodepluginupload.find_unique.return_value = LiteLLM_ClaudeCodePluginUpload( + plugin_name="skill", content=Base64.encode(content), + ) + + async def update(where: Mapping[str, str], data: Mapping[str, object]) -> LiteLLM_ClaudeCodePluginTable: + return current.model_copy(update=dict(data)) + + publication_tx.litellm_claudecodeplugintable.update.side_effect = update + source: Final = { + "source": "archive", "url": "https://external.example.test/skill.zip", + "sha256": hashlib.sha256(b"original uploaded bytes").hexdigest(), + } + result: Final = await publish_plugin( + cast(PluginPublicationTransaction, publication_tx), + PluginPublication("skill", "1.0.0", "replacement", {"source": source}, "admin"), + PluginPublicationBase(current.id, current.updated_at), + ) + assert isinstance(result, PublishedPlugin) and result.version == "1.0.0" and result.enabled is False + original, replacement = publication_tx.litellm_claudecodepluginversion.create.call_args_list + assert original.kwargs["data"] == { + "plugin_id": current.id, "version": None, "description": current.description, + "manifest_json": current.manifest_json, "published_at": current.updated_at, "published_by": None, + "content": Base64.encode(content) if valid_archive else None, "archive_available": valid_archive, + } + assert replacement.kwargs["data"] == { + "plugin_id": current.id, "version": result.version, "description": "replacement", + "manifest_json": json.dumps({"source": source, "version": result.version}), + "published_at": result.updated_at, "published_by": "admin", "content": None, "archive_available": False, + } diff --git a/tests/unit/repositories/test_plugin_upload_repository.py b/tests/unit/repositories/test_plugin_upload_repository.py index dc35973a062..8d11afc2518 100644 --- a/tests/unit/repositories/test_plugin_upload_repository.py +++ b/tests/unit/repositories/test_plugin_upload_repository.py @@ -7,17 +7,28 @@ import pytest from fastapi import Request from prisma import models as prisma_models +from litellm.models.plugin_upload import PluginArchiveMetadata, PluginDownloader, PluginDownloadResult from litellm.proxy import proxy_server from litellm.proxy.discovery_endpoints.skill_hub_endpoints import public_hub_entries from litellm.repositories.plugin_upload_repository import PluginUploadRepository +@pytest.mark.asyncio +async def test_download_without_stored_identity_returns_missing_skill() -> None: + repository: Final = PluginUploadRepository(MagicMock()) + result: Final = await repository.record_download( + PluginArchiveMetadata("missing", True, {}, None), "digest", + PluginDownloader(True, None, None, None, None, None), + ) + assert result is PluginDownloadResult.SKILL_NOT_FOUND + + def plugin( name: str, manifest_json: str | None = None, *, host: str = "testserver", ) -> prisma_models.LiteLLM_ClaudeCodePluginTable: digest: Final = "a" * 64 return prisma_models.LiteLLM_ClaudeCodePluginTable( - id=name, name=name, enabled=True, + id=name, name=name, enabled=True, anonymous_download_count=0, manifest_json=manifest_json if manifest_json is not None else json.dumps({ "source": { "source": "archive", "url": f"http://{host}/claude-code/plugins/{name}/archive/{digest}.zip", diff --git a/tests/unit/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index 082a29b7001..049e5275db6 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -168,7 +168,8 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch( return batches[-1] async with budget_cascade_unit_of_work(_new_batch) as uow: - writes = [getattr(uow, field.name) for field in fields(uow)] + uow.row_locks.queue_row_locks(["budget-1"]) + writes = [getattr(uow, field.name) for field in fields(uow) if field.name != "row_locks"] for write in writes: if isinstance(write, SpendResetWrites): write.queue_spend_zero(where={"budget_id": "budget-1"}) @@ -177,8 +178,10 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch( assert len(batches) == 1, "the cascade must open exactly one transaction" batch = batches[0] - assert len(batch.calls) == len(writes), "a dependent bound to a batch of its own would not land here" - assert len({call[0] for call in batch.calls}) == len(writes), "two dependents share one table" + table_writes = [call for call in batch.calls if call[0] != "execute_raw"] + assert len(table_writes) == len(writes), "a dependent bound to a batch of its own would not land here" + assert len({call[0] for call in table_writes}) == len(writes), "two dependents share one table" + assert [call[0] for call in batch.calls].count("execute_raw") > 0, "the row locks share the cascade's batch" assert batch.commit_count == 1 diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 012926c6cb3..2427e7eacb1 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -4259,75 +4259,6 @@ class TestCompletedResponseLatchedOnStreamEnd: assert iterator.completed_response.type == "response.completed" -class TestFallbackWrapperStopAsyncIterationFallback: - """Regression: FallbackResponsesStreamWrapper.__anext__ only sniffed - terminal events off forwarded chunks. When the inner generator raises - StopAsyncIteration without a sniffable chunk (the bridge path ends this - way), the wrapper re-raised without checking source_iterator for a - latched completed_response, leaving its own completed_response None.""" - - def test_falls_back_to_source_completed_response_on_stop(self): - """When the inner async generator raises StopAsyncIteration and the - wrapper never sniffed a terminal chunk, it must copy - source_iterator.completed_response so the proxy ownership hook - still sees the terminal event.""" - import asyncio - from types import SimpleNamespace - - from litellm.router import Router - - source = SimpleNamespace( - response=None, - model="deepseek/deepseek-chat", - logging_obj=None, - responses_api_provider_config=None, - start_time=None, - litellm_metadata=None, - custom_llm_provider="deepseek", - request_data={}, - call_type="aresponses", - _hidden_params={}, - completed_response=SimpleNamespace( - type="response.completed", - response=SimpleNamespace(id="resp_src", output=[], container=None), - ), - ) - - async def empty_gen(): - return - yield # pragma: no cover - - async def _drive(): - router = Router( - model_list=[ - { - "model_name": "deepseek/deepseek-chat", - "litellm_params": {"model": "deepseek/deepseek-chat", "api_key": "sk-test"}, - } - ] - ) - wrapper = await router._aresponses_streaming_iterator( - response=source, # type: ignore[arg-type] - initial_kwargs={}, - ) - wrapper._async_generator = empty_gen() - out = [] - try: - async for chunk in wrapper: - out.append(chunk) - except StopAsyncIteration: - pass - return wrapper, out - - wrapper, _ = asyncio.run(_drive()) - assert wrapper.completed_response is not None, ( - "FallbackResponsesStreamWrapper.completed_response is None after " - "StopAsyncIteration even though source_iterator had one — the " - "proxy ownership hook will log a spurious warning" - ) - assert wrapper.completed_response.type == "response.completed" - - class TestEnsureOutputItemContentPartAdded: """Test that _ensure_output_item_for_chunk emits content_part.added after output_item.added for message items.""" @@ -6564,7 +6495,6 @@ def test_an_application_inference_profile_gets_no_item_marks_under_drop_params() assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 0 - def test_a_bedrock_deployment_keeps_item_marks_by_its_own_flag_not_a_siblings_on_the_same_model( local_model_cost_map: None, ): @@ -6589,6 +6519,7 @@ def test_a_bedrock_deployment_keeps_item_marks_by_its_own_flag_not_a_siblings_on bridged_system({"id": "opted-in-id", "supports_prompt_cache_breakpoint": True}), ) == ({"role": "system", "content": "House rules"}, marked_system) + def test_the_shared_input_transform_keeps_every_item_mark_for_a_request_it_sends(): messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=cast(ResponseInputParam, list(_marked_turns(5))), # cast-ok: cache_control is not an input TypedDict key @@ -6901,7 +6832,11 @@ def test_a_new_message_after_a_paused_search_turn_replays_only_that_turn_without [ {"role": "user", "content": "Plan a Tokyo trip."}, _reasoning_item("rs_tokyo", tokyo_thinking), - {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Tokyo is sunny."}]}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Tokyo is sunny."}], + }, _hosted_search("ws_srvtoolu_01Tokyo", "completed"), {"role": "user", "content": "Now Osaka."}, _reasoning_item("rs_osaka", _ANSWERED_THINKING), @@ -7264,7 +7199,11 @@ def test_a_tool_output_with_an_unhashable_call_id_is_a_400_naming_its_input_item ) -> None: history: Final = cast( ResponseInputParam, - [_WEATHER_ASK, _weather_call(call_id="call_weather"), {"type": output_type, "call_id": call_id, "output": "sunny"}], + [ + _WEATHER_ASK, + _weather_call(call_id="call_weather"), + {"type": output_type, "call_id": call_id, "output": "sunny"}, + ], ) with pytest.raises(litellm.BadRequestError) as raised: @@ -7368,7 +7307,11 @@ def test_a_non_string_role_still_converts_for_a_target_that_judges_it_itself() - id="dict_function_call_id", ), pytest.param( - [_WEATHER_ASK, _weather_call(call_id=7), {"type": "function_call_output", "call_id": "7", "output": "sunny"}], + [ + _WEATHER_ASK, + _weather_call(call_id=7), + {"type": "function_call_output", "call_id": "7", "output": "sunny"}, + ], "input[1].call_id", "int", id="int_function_call_id", diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 037004ab999..a39bc09cee6 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -2722,7 +2722,7 @@ async def test_stream_auto_execute_bills_a_follow_up_cut_after_streamed_text_wit @pytest.mark.asyncio -async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row( +async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_one_failure_row( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -2743,10 +2743,22 @@ async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_r if not ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked): ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) await logging_obj.async_failure_handler(blocked, "blocked by output-policy") - await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + await _settle(lambda: len(recorder.failed) > 0) - _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) - assert [row["response_cost"] for row in recorder.failed] == [0.0] + assert recorder.logged == () + assert len(recorder.failed) == 1 + row: Final = recorder.failed[0] + charged_once: Final = _round_cost(_ROUND_ONE_BODY) + _round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN + assert row["status"] == "failure" + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + client: Final = _final_response(events) + assert "message" in _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (71 + 149, 19 + 5) + assert client.usage.cost == pytest.approx(charged_once) @pytest.mark.asyncio diff --git a/tests/unit/responses/test_responses_router_cooldown.py b/tests/unit/responses/test_responses_router_cooldown.py index 078890c2395..a1cd5c8517b 100644 --- a/tests/unit/responses/test_responses_router_cooldown.py +++ b/tests/unit/responses/test_responses_router_cooldown.py @@ -11,14 +11,12 @@ from unittest.mock import AsyncMock, patch import httpx import pytest - import litellm -from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments @pytest.mark.asyncio async def test_responses_api_rate_limit_marks_deployment_for_cooldown(): - failing_deployment_id = "deployment-rate-limited" + deployment_ids = {"mock-api-key-1": "deployment-a", "mock-api-key-2": "deployment-b"} router = litellm.Router( model_list=[ @@ -28,7 +26,7 @@ async def test_responses_api_rate_limit_marks_deployment_for_cooldown(): "model": "openai/gpt-5.1-codex", "api_key": "mock-api-key-1", }, - "model_info": {"id": failing_deployment_id}, + "model_info": {"id": "deployment-a"}, }, { "model_name": "openai.gpt-5.1-codex", @@ -36,7 +34,7 @@ async def test_responses_api_rate_limit_marks_deployment_for_cooldown(): "model": "openai/gpt-5.1-codex", "api_key": "mock-api-key-2", }, - "model_info": {"id": "deployment-healthy"}, + "model_info": {"id": "deployment-b"}, }, ], num_retries=0, @@ -53,22 +51,16 @@ async def test_responses_api_rate_limit_marks_deployment_for_cooldown(): ), ) - def pin_to_failing_deployment(seq): - for d in seq: - if d["model_info"]["id"] == failing_deployment_id: - return d - return seq[0] + called_api_keys = [] - with ( - patch( - "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", - new_callable=AsyncMock, - side_effect=rate_limit_error, - ), - patch( - "litellm.router_strategy.simple_shuffle.random.choice", - side_effect=pin_to_failing_deployment, - ), + async def fail_and_record(*args, **kwargs): + called_api_keys.append(kwargs["litellm_params"].api_key) + raise rate_limit_error + + with patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new_callable=AsyncMock, + side_effect=fail_and_record, ): with pytest.raises(litellm.RateLimitError): await router.aresponses( @@ -76,10 +68,15 @@ async def test_responses_api_rate_limit_marks_deployment_for_cooldown(): input="hi", ) - cooldown_ids = await async_get_cooldown_deployments( - ctx=router._core_ctx, parent_otel_span=None - ) - assert failing_deployment_id in cooldown_ids, ( - f"Responses API failure callback did not register cooldown for " - f"{failing_deployment_id!r}; cooldown set was {cooldown_ids}" + # num_retries=0: exactly one deployment took the request and failed. + (called_api_key,) = called_api_keys + failed_id = deployment_ids[called_api_key] + cooldown_ids = [ + model_id + for model_id, _ in await router.cooldown_cache.async_get_active_cooldowns( + model_ids=list(deployment_ids.values()), parent_otel_span=None + ) + ] + assert cooldown_ids == [failed_id], ( + f"Responses API failure callback did not register cooldown for {failed_id!r}; cooldown set was {cooldown_ids}" ) diff --git a/tests/unit/router_rust/test_walk.py b/tests/unit/router_rust/test_walk.py new file mode 100644 index 00000000000..2f2febea3eb --- /dev/null +++ b/tests/unit/router_rust/test_walk.py @@ -0,0 +1,49 @@ +import pytest + +from litellm.router_rust.invoke.walk import WalkMetadata + + +def _no_base_model(_model_id: str) -> str | None: + return None + + +@pytest.mark.parametrize("field", ["metadata", "litellm_metadata"]) +def test_hop_kwargs_keep_the_requests_key_labels_under_the_walk_stamps(field: str) -> None: + """The fallback loggers (prometheus' fallback counters) read the caller's key and team off + the hop's metadata, so a hop keeps the request's bucket and lays its own stamps over it.""" + request = { + "model": "primary", + field: { + "user_api_key_hash": "hashed-key", + "user_api_key_alias": "key-alias", + "user_api_key_team_id": "team-id", + "user_api_key_team_alias": "team-alias", + "attempted_fallbacks": "stale", + }, + } + walk = WalkMetadata(request, field, _no_base_model) + walk.hop({"parent": 0, "scope": 1, "from": "primary", "to": "fallback", "depth": 1, "max_fallbacks": 5}) + + kwargs = walk.hop_kwargs({"scope": 1, "to": "fallback"}) + + assert kwargs["model"] == "fallback" + metadata = kwargs[field] + assert isinstance(metadata, dict) + assert metadata["user_api_key_hash"] == "hashed-key" + assert metadata["user_api_key_alias"] == "key-alias" + assert metadata["user_api_key_team_id"] == "team-id" + assert metadata["user_api_key_team_alias"] == "team-alias" + assert metadata["model_group"] == "fallback" + assert metadata["original_model_group"] == "primary" + assert metadata["attempted_fallbacks"] == 1 + assert request[field]["attempted_fallbacks"] == "stale" + assert "model_group" not in request[field] + + +def test_hop_kwargs_without_request_metadata_hold_the_walk_stamps() -> None: + walk = WalkMetadata({"model": "primary"}, "metadata", _no_base_model) + walk.hop({"parent": 0, "scope": 1, "from": "primary", "to": "fallback", "depth": 1, "max_fallbacks": 5}) + + metadata = walk.hop_kwargs({"scope": 1, "to": "fallback"})["metadata"] + + assert metadata == {"original_model_group": "primary", "model_group": "fallback", "attempted_fallbacks": 1} diff --git a/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py b/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py index 6fe39599b0f..b7b33f4789b 100644 --- a/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py +++ b/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py @@ -26,63 +26,6 @@ def _adaptive(r, name): return r.adaptive_routers[name][0].strategy -def _params(**overrides): - base = {"model": "auto_router/adaptive_router"} - base.update(overrides) - return LiteLLM_Params(**base) - - -# ---- Fix 1 & 2: opt-in prefix routing ----------------------------------- - - -def test_auto_router_check_excludes_adaptive_router_prefix(): - r = Router(model_list=[]) - assert ( - r._is_auto_router_deployment( - litellm_params=_params(model="auto_router/adaptive_router") - ) - is False - ) - - -def test_auto_router_check_excludes_complexity_router_prefix(): - r = Router(model_list=[]) - assert ( - r._is_auto_router_deployment( - litellm_params=_params(model="auto_router/complexity_router") - ) - is False - ) - - -def test_auto_router_check_still_matches_plain_auto_router_prefix(): - r = Router(model_list=[]) - assert ( - r._is_auto_router_deployment( - litellm_params=_params(model="auto_router/my-semantic-router") - ) - is True - ) - - -def test_adaptive_router_check_recognizes_prefix(): - r = Router(model_list=[]) - assert ( - r._is_adaptive_router_deployment( - litellm_params=_params(model="auto_router/adaptive_router") - ) - is True - ) - - -def test_adaptive_router_check_rejects_other_prefixes(): - r = Router(model_list=[]) - assert ( - r._is_adaptive_router_deployment(litellm_params=_params(model="openai/gpt-4o")) - is False - ) - - # ---- Fix 3: cost field path -------------------------------------------- @@ -481,44 +424,6 @@ def test_init_adaptive_router_rejects_duplicate_model_name(): r.init_adaptive_router_deployment(deployment=deployment) -def test_finalize_adaptive_router_if_configured_initializes_and_is_idempotent(): - """`_finalize_adaptive_router_if_configured` walks the model_list, builds an - AdaptiveRouter for each adaptive deployment, and is a safe no-op on - re-entry (models already in self.adaptive_routers are skipped).""" - r = Router( - model_list=[ - { - "model_name": "fast", - "litellm_params": {"model": "openai/gpt-4o-mini"}, - "model_info": {"input_cost_per_token": 0.00000015}, - }, - { - "model_name": "smart", - "litellm_params": {"model": "openai/gpt-4o"}, - "model_info": {"input_cost_per_token": 0.0000025}, - }, - { - "model_name": "my-router", - "litellm_params": { - "model": "auto_router/adaptive_router", - "adaptive_router_config": { - "available_models": ["fast", "smart"], - }, - }, - }, - ] - ) - - # Router __init__ already called _finalize_adaptive_router_if_configured. - assert "my-router" in r.adaptive_routers - original = _adaptive(r, "my-router") - - # Calling again must be idempotent: the existing AdaptiveRouter instance - # is preserved, not rebuilt. - r._finalize_adaptive_router_if_configured() - assert _adaptive(r, "my-router") is original - - def test_finalize_prunes_stale_adaptive_router_hooks_from_callbacks(): """Replacing the Router (hot-reload path) must not leave stale AdaptiveRouterPostCallHook instances in `litellm.callbacks` — otherwise @@ -572,16 +477,3 @@ def test_finalize_prunes_stale_adaptive_router_hooks_from_callbacks(): litellm.callbacks.append(cb) -def test_finalize_adaptive_router_if_configured_noop_when_none_configured(): - """With no adaptive deployments in model_list, the finalizer leaves - `adaptive_routers` empty.""" - r = Router( - model_list=[ - { - "model_name": "fast", - "litellm_params": {"model": "openai/gpt-4o-mini"}, - } - ] - ) - r._finalize_adaptive_router_if_configured() - assert r.adaptive_routers == {} diff --git a/tests/unit/router_strategy/test_baseline_request.py b/tests/unit/router_strategy/test_baseline_request.py index eb785a41d63..c0ca2464a9c 100644 --- a/tests/unit/router_strategy/test_baseline_request.py +++ b/tests/unit/router_strategy/test_baseline_request.py @@ -37,7 +37,7 @@ def test_chat_projection_applies_extra_body_after_top_level_parameters() -> None def test_baseline_projection_keeps_caller_tools_and_request_parameter_precedence() -> None: - from litellm.router import Router + from litellm.router_rust.invoke.shaping import merge_tools_from_deployment deployment: Final = { "tools": [{"type": "function", "function": {"name": "configured"}}], @@ -50,7 +50,7 @@ def test_baseline_projection_keeps_caller_tools_and_request_parameter_precedence "max_tokens": 128, } actual_request: Final = dict(caller) - Router._merge_tools_from_deployment({"litellm_params": deployment}, actual_request) + merge_tools_from_deployment(deployment, actual_request, {}) snapshot: Final = capture_baseline_parameters(caller) assert snapshot is not None assert baseline_request({"tools": [{"name": "routed-only"}], "max_tokens": 4}, snapshot, deployment) == { diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index d0e2b381253..ba829512d8b 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -10,9 +10,8 @@ import logging import math import sys import time -from collections.abc import AsyncIterator, Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from copy import deepcopy -from functools import partial from types import MappingProxyType from typing import Dict, Final, List, Literal from unittest.mock import DEFAULT, AsyncMock, MagicMock, patch @@ -35,7 +34,6 @@ from litellm.constants import ( SESSION_ID_GENERATED_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.router import as_output_cap, stamped_client_ceiling from litellm.router_api import RoutingContext from litellm.router_strategy.auto_router.auto_router_model_naming import ( @@ -95,7 +93,6 @@ from litellm.types.router import ( LiteLLM_Params, PreRoutingHookResponse, RouterErrors, - TaggedPreRoutingStrategy, ) from litellm.types.utils import StandardLoggingRoutingDecision @@ -1684,36 +1681,6 @@ class TestEdgeCases: class TestRouterComplexityDeploymentMethods: """Tests for Router._is_complexity_router_deployment and Router.init_complexity_router_deployment.""" - def test_is_complexity_router_deployment_true(self): - """_is_complexity_router_deployment returns True for complexity router models.""" - router = Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini"}, - } - ] - ) - from litellm.types.router import LiteLLM_Params - - params = LiteLLM_Params(model="auto_router/complexity_router/my-router") - assert router._is_complexity_router_deployment(params) is True - - def test_is_complexity_router_deployment_false(self): - """_is_complexity_router_deployment returns False for regular models.""" - router = Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini"}, - } - ] - ) - from litellm.types.router import LiteLLM_Params - - params = LiteLLM_Params(model="openai/gpt-4o-mini") - assert router._is_complexity_router_deployment(params) is False - def test_init_complexity_router_deployment(self): """init_complexity_router_deployment registers a ComplexityRouter.""" router = Router( @@ -2506,106 +2473,6 @@ class TestComplexityRouterTagBasedRouting: assert us is not None and us.model == "gpt-us" -class TestPreRoutingStrategyRegistry: - """Directly exercise the tag-scoped registry/selection helpers behind #33655.""" - - def _router(self) -> Router: - return Router(model_list=[{"model_name": "x", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) - - @staticmethod - def _deployment(tags: list) -> Deployment: - return Deployment( - model_name="smart", - litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini", tags=tags), - ) - - def test_deployment_tags_normalizes_to_tuple(self): - router = self._router() - assert router._deployment_tags(self._deployment(["cn", "row"])) == ("cn", "row") - untagged = Deployment(model_name="smart", litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini")) - assert router._deployment_tags(untagged) == () - - def test_register_scopes_by_tags_and_rejects_exact_duplicate(self): - router = self._router() - registry: dict = {} - router._register_pre_routing_strategy( - registry=registry, deployment=self._deployment(["cn"]), strategy="CN", strategy_label="Test" - ) - router._register_pre_routing_strategy( - registry=registry, deployment=self._deployment(["us"]), strategy="US", strategy_label="Test" - ) - assert [entry.tags for entry in registry["smart"]] == [("cn",), ("us",)] - assert router._has_registered_strategy(registry, "smart", ("cn",)) is True - assert router._has_registered_strategy(registry, "smart", ("row",)) is False - with pytest.raises(ValueError, match="already exists"): - router._register_pre_routing_strategy( - registry=registry, deployment=self._deployment(["cn"]), strategy="CN2", strategy_label="Test" - ) - - def test_select_prefers_request_tag_then_default_then_first(self): - router = self._router() - cn, us, fallback = object(), object(), object() - router.complexity_routers = { - "smart": [ - TaggedPreRoutingStrategy(tags=("cn",), strategy=cn), - TaggedPreRoutingStrategy(tags=("us",), strategy=us), - ] - } - assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["us"]}}).strategy is us - assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["cn"]}}).strategy is cn - assert router._select_pre_routing_strategy("missing", {"metadata": {"tags": ["cn"]}}) is None - - router.complexity_routers = { - "smart": [ - TaggedPreRoutingStrategy(tags=("cn",), strategy=cn), - TaggedPreRoutingStrategy(tags=("default",), strategy=fallback), - ] - } - assert router._select_pre_routing_strategy("smart", {}).strategy is fallback - router.complexity_routers = { - "smart": [ - TaggedPreRoutingStrategy(tags=("cn",), strategy=cn), - TaggedPreRoutingStrategy(tags=("us",), strategy=us), - ] - } - assert router._select_pre_routing_strategy("smart", {}).strategy is cn - - @staticmethod - def _router_with_plain_smart_deployment(enable_tag_filtering: bool) -> Router: - return Router( - model_list=[{"model_name": "smart", "litellm_params": {"model": "openai/gpt-4o-mini"}}], - enable_tag_filtering=enable_tag_filtering, - ) - - def test_select_falls_through_to_plain_deployments_when_no_tag_matches_under_tag_filtering(self): - router = self._router_with_plain_smart_deployment(enable_tag_filtering=True) - cn, us = object(), object() - - router.complexity_routers = {"smart": [TaggedPreRoutingStrategy(tags=("cn",), strategy=cn)]} - assert router._select_pre_routing_strategy("smart", {}) is None - assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["cn"]}}).strategy is cn - - router.complexity_routers = { - "smart": [ - TaggedPreRoutingStrategy(tags=("cn",), strategy=cn), - TaggedPreRoutingStrategy(tags=("us",), strategy=us), - ] - } - assert router._select_pre_routing_strategy("smart", {}) is None - assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["row"]}}) is None - assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["us"]}}).strategy is us - - router.complexity_routers["router-only"] = [TaggedPreRoutingStrategy(tags=("cn",), strategy=cn)] - assert router._select_pre_routing_strategy("router-only", {}).strategy is cn - - def test_select_keeps_capturing_when_tag_filtering_is_disabled(self): - router = self._router_with_plain_smart_deployment(enable_tag_filtering=False) - cn = object() - - router.complexity_routers = {"smart": [TaggedPreRoutingStrategy(tags=("cn",), strategy=cn)]} - assert router._select_pre_routing_strategy("smart", {}).strategy is cn - - class TestAsyncPreRoutingHookMultiFormat: """Test async_pre_routing_hook with multiple input formats.""" @@ -3637,11 +3504,10 @@ def _native_classifier_router( deployment_model: str = "openai/gpt-6-astra", failure: Exception | None = None, native_router: Router | None = None, - http_handler: AsyncHTTPHandler | None = None, ) -> tuple[ComplexityRouter, MagicMock]: dependency: Final = MagicMock( aresponses=( - native_router.factory_function(partial(litellm.aresponses, client=http_handler), call_type="aresponses") + native_router.aresponses if native_router is not None else AsyncMock(return_value=_native_classifier_response(output), side_effect=failure) ), @@ -3685,15 +3551,14 @@ def _native_classifier_router( @pytest.fixture -async def native_classifier_http() -> AsyncIterator[tuple[AsyncHTTPHandler, MagicMock]]: +def native_classifier_http(monkeypatch: pytest.MonkeyPatch) -> Iterator[MagicMock]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) respond: Final = MagicMock( return_value=httpx.Response(200, json=_native_classifier_response('{"tier":"REASONING"}').model_dump()) ) - async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: - handler: Final = AsyncHTTPHandler() - await handler.client.aclose() - handler.client = client - yield handler, respond + with respx.mock(assert_all_mocked=True, assert_all_called=False) as upstream: + upstream.post(host="classifier.test").mock(side_effect=lambda request: respond(request)) + yield respond class TestEncryptedTaskClassifier: @@ -3871,9 +3736,9 @@ class TestEncryptedTaskClassifier: ["anthropic/test-classifier", "openai/chat_completions/gpt-6-astra", "xai/test-classifier"], ) async def test_incompatible_classifier_does_not_flatten_encryption( - self, deployment_model: str, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock] + self, deployment_model: str, native_classifier_http: MagicMock ): - handler, respond = native_classifier_http + respond = native_classifier_http native: Final = Router( model_list=[ { @@ -3887,7 +3752,7 @@ class TestEncryptedTaskClassifier: ], num_retries=0, ) - router, _ = _native_classifier_router(native_router=native, http_handler=handler) + router, _ = _native_classifier_router(native_router=native) result: Final = await router.async_pre_routing_hook( model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]} @@ -3900,9 +3765,9 @@ class TestEncryptedTaskClassifier: @pytest.mark.asyncio @pytest.mark.parametrize("blocked", [True, False]) async def test_native_classifier_validates_selected_deployment( - self, blocked: bool, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock] + self, blocked: bool, native_classifier_http: MagicMock ): - handler, respond = native_classifier_http + respond = native_classifier_http native: Final = Router( model_list=[ { @@ -3923,7 +3788,7 @@ class TestEncryptedTaskClassifier: ], num_retries=0, ) - router, _ = _native_classifier_router(native_router=native, http_handler=handler) + router, _ = _native_classifier_router(native_router=native) task: Final = _encrypted_agent_task() result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs={"input": [task]}) @@ -3991,7 +3856,6 @@ class TestLLMClassifier: assert outcome.score is not None @pytest.mark.asyncio - @pytest.mark.parametrize("redact", (False, True)) @pytest.mark.parametrize( "override,threshold,tier,model", ( @@ -4007,14 +3871,11 @@ class TestLLMClassifier: async def test_heuristic_v2_routes_directly_to_predicted_builtin_tier( self, mock_router_instance: MagicMock, - redact: bool, - monkeypatch: pytest.MonkeyPatch, override: Mapping[str, float | None], threshold: float, tier: str, model: str, ) -> None: - monkeypatch.setattr(litellm, "turn_off_message_logging", redact) artifact: Final = _heuristic_v2_artifact() router: Final = ComplexityRouter( model_name="tier-router", @@ -4049,11 +3910,7 @@ class TestLLMClassifier: "tier-probability:complex=0.892157", "tier-probability:reasoning=0.980392", ] - redacted: Final = Router._redact_prompt_text_if_needed( - request_kwargs={}, routing_decision=response.routing_decision - ) - assert ("signals" in redacted) is not redact - assert redacted["heuristic_v2_forecast"] == { + assert response.routing_decision["heuristic_v2_forecast"] == { "probabilities": { "SIMPLE": 11 / 102, "MEDIUM": 21 / 102, @@ -5065,20 +4922,6 @@ class TestRouterPreRoutingAliasOverrides: assert "thinking" not in request_kwargs assert "output_config" not in request_kwargs - def test_drop_client_effort_carriers_helper_edge_shapes(self): - no_pin: Dict = {"thinking": {"type": "adaptive"}} - Router._drop_client_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1}) - assert no_pin == {"thinking": {"type": "adaptive"}} - - non_dict_carriers: Dict = {"output_config": "max", "reasoning": 3} - Router._drop_client_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"}) - assert non_dict_carriers == {"output_config": "max", "reasoning": 3} - - effort_only: Dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}} - Router._pop_effort_from_nested_carrier(effort_only, "output_config") - Router._pop_effort_from_nested_carrier(effort_only, "reasoning") - assert effort_only == {} - @pytest.mark.asyncio async def test_client_effort_carriers_survive_when_gate_drops_the_tier_pin(self): """The tier-param gate removes a pin the routed target cannot take, and a @@ -5496,15 +5339,6 @@ class TestRouterPreRoutingSharedAliasName: assert cn_result is not None and cn_result.model == "gpt-cn" assert "drop_params" not in cn_kwargs - def test_forwardable_alias_marker_params_reads_the_marker_entry_only(self): - router = Router(model_list=[self._plain_entry(), self._marker_entry(), self._tier_entry()]) - - forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=(), request_kwargs={})) - - assert forwarded["drop_params"] is True - assert "api_key" not in forwarded and "api_base" not in forwarded - assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=(), request_kwargs={}) == () - @staticmethod def _region_marker_entry() -> dict: return { @@ -5690,27 +5524,6 @@ class TestRouterPreRoutingSharedAliasName: assert not any(key.startswith("_alias_marker") for key in sent) - def test_forwarded_alias_marker_keys_the_deployment_sets(self): - deployment = {"litellm_params": {"model": "bedrock/x", "aws_region_name": "us-east-1", "timeout": None}} - - assert Router._forwarded_alias_marker_keys_the_deployment_sets( - deployment=deployment, forwarded_keys=("aws_region_name", "timeout", "drop_params") - ) == ("aws_region_name",) - assert Router._forwarded_alias_marker_keys_the_deployment_sets(deployment=deployment, forwarded_keys=()) == () - assert Router._forwarded_alias_marker_keys_the_deployment_sets(deployment=deployment, forwarded_keys=None) == () - assert Router._forwarded_alias_marker_keys_the_deployment_sets(deployment={}, forwarded_keys=("x",)) == () - - def test_deployment_sets_litellm_param(self): - params = {"aws_region_name": "us-east-1", "timeout": None, "use_litellm_proxy": False, "custom_flag": False} - - assert Router._deployment_sets_litellm_param(params, "aws_region_name") is True - assert Router._deployment_sets_litellm_param(params, "timeout") is False - assert Router._deployment_sets_litellm_param(params, "missing") is False - assert Router._deployment_sets_litellm_param(params, "use_litellm_proxy") is False - assert Router._deployment_sets_litellm_param({"use_litellm_proxy": True}, "use_litellm_proxy") is True - assert Router._deployment_sets_litellm_param(params, "custom_flag") is True - - class TestAdaptiveSoftFloors: def test_adaptive_defaults_use_cost_weighted_cold_policy(self): config = ComplexityRouterConfig( @@ -9857,51 +9670,6 @@ class TestRoutingDecisionIsPerAttempt: assert bucket["routing_decision"]["router_model_name"] == "smart-router" -class TestRecordRoutingDecision: - """Direct coverage of the single recording point, whose contract is write-or-clear: - the request's metadata must describe the current attempt and nothing else.""" - - DECISION = {"router_model_name": "smart-router", "router_type": "complexity", "routed_model": "gpt-4o-mini"} - - def test_none_clears_a_previous_decision_from_both_buckets(self): - request_kwargs: Dict = { - "metadata": {"routing_decision": self.DECISION, "keep": 1}, - "litellm_metadata": {"routing_decision": self.DECISION}, - } - Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) - assert "routing_decision" not in request_kwargs["metadata"] - assert "routing_decision" not in request_kwargs["litellm_metadata"] - assert request_kwargs["metadata"]["keep"] == 1 - - def test_none_creates_no_bucket_on_a_request_that_had_none(self): - request_kwargs: Dict = {} - Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) - assert request_kwargs == {} - - def test_clearing_the_decision_takes_the_savings_facts_with_it(self) -> None: - """A fallback to a plain model group re-enters the hook with the same - `request_kwargs`. The baseline and the conversation shape ride inside the - decision rather than beside it, so one clear cannot leave either behind and - attribute an auto-router saving to a deployment that never routed.""" - from litellm.types.router import BaselineRouteStamp - - decision: Final = { - "router_model_name": "smart-router", - "router_type": "complexity", - "routed_model": "gpt-4o-mini", - "savings_baseline_model": "anthropic/claude-opus-5", - "savings_baseline_deployment_id": "opus-deployment", - "conversation_continuing": False, - } - request_kwargs: Final[dict[str, dict[str, object]]] = {"litellm_metadata": {}} - Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=decision) - stamp: Final = request_kwargs["litellm_metadata"]["_autorouter_baseline_route"] - assert isinstance(stamp, BaselineRouteStamp) - assert stamp.baseline_deployment_id == "opus-deployment" - Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) - assert request_kwargs["litellm_metadata"] == {} - - class TestEscalationIsRecordedConsistently: """An escalation keyword records two separate facts on every path: that the caller asked, and whether the tier actually moved. Dropping the ask when there is nowhere @@ -10062,44 +9830,6 @@ class TestRedactedLoggingDropsPromptText: assert decision["routed_model"] == "gpt-4o" assert decision["escalated"] is False - def test_only_verbatim_prompt_fields_are_classified_as_prompt_text(self, monkeypatch): - """The field classification is the whole contract, so pin it directly: anything - that quotes the prompt goes, anything derived from it stays.""" - monkeypatch.setattr(litellm, "turn_off_message_logging", True) - full = { - "router_model_name": "smart-router", - "router_type": "complexity", - "routed_model": "gpt-4o", - "cause": "literal_keyword_match", - "tier": "REASONING", - "score": 0.8, - "tier_boundaries": {"simple_medium": 0.15, "medium_complex": 0.35, "complex_reasoning": 0.6}, - "classifier_model": "claude-haiku", - "classifier_crux": "deploy the requested service to k8s", - "classifier_primary_rule": "SUP-2", - "classifier_capability_boundary": "supported", - "classifier_p_solve": 0.8, - "classifier_calibrated_p_solve": 0.65, - "classifier_calibration_version": "fitted-v1", - "classifier_threshold": 0.5, - "escalated": True, - "tier_litellm_params": {"reasoning_effort": "xhigh"}, - "signals": ["code (python)"], - "matched_keyword": "deploy to k8s", - "escalation_keyword": "LITELLM ESCALATE", - } - kept = Router._redact_prompt_text_if_needed(request_kwargs={}, routing_decision=full) - assert set(full) - set(kept) == { - "signals", - "matched_keyword", - "escalation_keyword", - "classifier_crux", - } - assert kept["classifier_p_solve"] == 0.8 - assert kept["classifier_calibrated_p_solve"] == 0.65 - assert kept["classifier_calibration_version"] == "fitted-v1" - assert kept["tier_litellm_params"] == {"reasoning_effort": "xhigh"} - @pytest.mark.asyncio async def test_redaction_via_request_header_is_honored(self): request_kwargs: Dict = {"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}} @@ -15761,14 +15491,6 @@ class TestHealthFallbackDispatch: else: assert decision["tier"] == "REASONING" assert decision["tier_label"] == "Advanced" - redacted: Final = Router._redact_prompt_text_if_needed( - request_kwargs={"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}}, - routing_decision=decision, - ) - assert "classifier_crux" not in redacted and "signals" not in redacted - assert {field: value for field, value in redacted.items() if field.startswith("classifier_")} == { - field: value for field, value in expected.items() if field != "classifier_crux" - } @pytest.mark.asyncio @pytest.mark.parametrize("peer", (True, False), ids=("peer_failover", "default_fallback")) @@ -16718,13 +16440,15 @@ class TestTierHealthFailover: ("live-big", "openai/gpt-4o-mini", 200000), ) router_instance.enable_pre_call_checks = True - dead_ids = {d["model_info"]["id"] for d in router_instance.model_list if d["model_name"] == "dead-big"} - - async def active_cooldowns(model_ids, parent_otel_span): - return [(i, {"exception_received": "boom"}) for i in model_ids if i in dead_ids] - - router_instance.cooldown_cache.async_get_active_cooldowns = active_cooldowns router_instance.cache = DualCache() + for deployment in router_instance.model_list: + if deployment["model_name"] == "dead-big": + router_instance.cooldown_cache.add_deployment_to_cooldown( + model_id=deployment["model_info"]["id"], + original_exception=RuntimeError("boom"), + exception_status=500, + cooldown_time=600, + ) router = ComplexityRouter( model_name="health-window-router", ctx=RoutingContext(router_instance), @@ -17623,24 +17347,6 @@ class TestMaxTokensFromTierModel: assert sent["max_tokens"] == 64000 - @pytest.mark.parametrize( - "tier_params, responses_call, expected", - [ - ({"max_tokens": 1, "temperature": 0.2}, False, {"max_tokens": 1, "temperature": 0.2}), - ({"max_tokens": 1}, True, {"max_output_tokens": 1}), - ({"max_completion_tokens": 2}, False, {"max_tokens": 2}), - ({"max_completion_tokens": 2}, True, {"max_output_tokens": 2}), - ({"max_output_tokens": 3}, False, {"max_tokens": 3}), - ({"max_output_tokens": 3}, True, {"max_output_tokens": 3}), - ({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 1}), - ({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, True, {"max_output_tokens": 3}), - ({"max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 2}), - ({"reasoning_effort": "low"}, True, {"reasoning_effort": "low"}), - ], - ) - def test_every_tier_alias_collapses_onto_the_surface_key(self, tier_params, responses_call, expected): - assert dict(Router._tier_ceiling_under_the_surface_name(tier_params, responses_call=responses_call)) == expected - @pytest.mark.asyncio async def test_the_default_fallback_exit_carries_the_ceiling(self): routed: dict = {"max_tokens": 8192} @@ -17813,47 +17519,6 @@ class TestMaxTokensFromTierModel: assert (None if stamped is None else dict(stamped)) == expected - def test_restoring_the_callers_ceiling_reads_the_stamp_and_replaces_every_carrier(self): - stamped: dict = {"max_output_tokens": 500, "metadata": {"_client_output_ceiling": {"max_tokens": 8192}}} - Router._restore_client_output_ceiling(stamped) - assert {k: v for k, v in stamped.items() if k != "metadata"} == {"max_tokens": 8192} - - coerced: dict = { - "max_tokens": 64000, - "metadata": {"_client_output_ceiling": {"max_tokens": "8192", "max_completion_tokens": 100.0}}, - } - Router._restore_client_output_ceiling(coerced) - assert {k: v for k, v in coerced.items() if k != "metadata"} == { - "max_tokens": 8192, - "max_completion_tokens": 100, - } - - unstamped: dict = {"max_tokens": 64000, "metadata": {}} - Router._restore_client_output_ceiling(unstamped) - assert unstamped["max_tokens"] == 64000 - - @pytest.mark.asyncio - async def test_pinning_stamps_the_callers_carriers_once(self): - router = self._router() - request_kwargs: dict = {"max_completion_tokens": 8192} - router._stamp_client_output_ceiling(request_kwargs) - assert stamped_client_ceiling(request_kwargs) == {"max_completion_tokens": 8192} - - first = router._pin_tier_params_onto_request( - model="big", tier_litellm_params={"max_tokens": 64000}, request_kwargs=request_kwargs, responses_call=False - ) - second = router._pin_tier_params_onto_request( - model="big", tier_litellm_params={"max_tokens": 32000}, request_kwargs=request_kwargs, responses_call=False - ) - none = router._pin_tier_params_onto_request( - model="big", tier_litellm_params=None, request_kwargs=request_kwargs, responses_call=False - ) - - assert (first, second, none) == (True, True, False) - assert request_kwargs["max_tokens"] == 32000 - assert request_kwargs["metadata"]["_client_output_ceiling"] == {"max_completion_tokens": 8192} - - class _OutputCeilingRecorder(CustomLogger): def __init__(self) -> None: super().__init__() diff --git a/tests/unit/router_strategy/test_llm_v2.py b/tests/unit/router_strategy/test_llm_v2.py index 993da777c66..4e566c0ae3b 100644 --- a/tests/unit/router_strategy/test_llm_v2.py +++ b/tests/unit/router_strategy/test_llm_v2.py @@ -389,9 +389,7 @@ async def test_caller_constraints_respect_claude_code_prompt_policy(user_agent: @pytest.mark.asyncio @pytest.mark.parametrize("calibrated", (False, True)) -async def test_routing_metadata_preserves_exact_forecasts_and_redaction( - calibrated: bool, monkeypatch: pytest.MonkeyPatch -) -> None: +async def test_routing_metadata_preserves_exact_forecasts(calibrated: bool) -> None: base: Final = _config().llm_v2_config assert base is not None calibration: Final = LLMV2Calibration( @@ -410,23 +408,18 @@ async def test_routing_metadata_preserves_exact_forecasts_and_redaction( assert result.model == ("capable" if calibrated else "efficient") decision: Final = result.routing_decision assert decision is not None - monkeypatch.setattr(litellm, "turn_off_message_logging", True) - redacted: Final = Router._redact_prompt_text_if_needed(request_kwargs={}, routing_decision=decision) - assert redacted is not None - assert "signals" not in redacted - for record in (decision, redacted): - assert record["classifier_efficient_p_solve"] == 0.900000123 - assert record["classifier_capable_p_solve"] == 0.920000321 - assert record["classifier_max_quality_gap"] == 0.05 - assert record["classifier_prompt_version"] == LLM_V2_PROMPT_VERSION - if calibrated: - assert record["classifier_calibration_version"] == "test-pair-v1" - assert record["classifier_calibrated_efficient_p_solve"] == calibration.efficient.calibrate(0.900000123) - assert record["classifier_calibrated_capable_p_solve"] == calibration.capable.calibrate(0.920000321) - else: - assert "classifier_calibration_version" not in record - assert "classifier_calibrated_efficient_p_solve" not in record - assert "classifier_calibrated_capable_p_solve" not in record + assert decision["classifier_efficient_p_solve"] == 0.900000123 + assert decision["classifier_capable_p_solve"] == 0.920000321 + assert decision["classifier_max_quality_gap"] == 0.05 + assert decision["classifier_prompt_version"] == LLM_V2_PROMPT_VERSION + if calibrated: + assert decision["classifier_calibration_version"] == "test-pair-v1" + assert decision["classifier_calibrated_efficient_p_solve"] == calibration.efficient.calibrate(0.900000123) + assert decision["classifier_calibrated_capable_p_solve"] == calibration.capable.calibrate(0.920000321) + else: + assert "classifier_calibration_version" not in decision + assert "classifier_calibrated_efficient_p_solve" not in decision + assert "classifier_calibrated_capable_p_solve" not in decision @pytest.mark.asyncio diff --git a/tests/unit/router_strategy/test_quality_router.py b/tests/unit/router_strategy/test_quality_router.py index 249e5254412..3878527298c 100644 --- a/tests/unit/router_strategy/test_quality_router.py +++ b/tests/unit/router_strategy/test_quality_router.py @@ -1024,39 +1024,7 @@ class TestSetResponseHeadersLiftsDecision: class TestRouterQualityDeploymentMethods: - """Tests for Router._is_quality_router_deployment and Router.init_quality_router_deployment.""" - - def test_is_quality_router_deployment_true(self): - """_is_quality_router_deployment returns True for quality router models.""" - from litellm.router import Router - from litellm.types.router import LiteLLM_Params - - router = Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini"}, - } - ] - ) - params = LiteLLM_Params(model="auto_router/quality_router/my-router") - assert router._is_quality_router_deployment(params) is True - - def test_is_quality_router_deployment_false(self): - """_is_quality_router_deployment returns False for regular models.""" - from litellm.router import Router - from litellm.types.router import LiteLLM_Params - - router = Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini"}, - } - ] - ) - params = LiteLLM_Params(model="openai/gpt-4o-mini") - assert router._is_quality_router_deployment(params) is False + """Tests for Router.init_quality_router_deployment.""" def test_init_quality_router_deployment(self): """init_quality_router_deployment registers a QualityRouter.""" diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 951182051c4..8a0d5b9a1d6 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -7,9 +7,8 @@ the implicit `"default"` group driven by the router's top-level import asyncio import datetime -import time import uuid -from collections.abc import Callable, Mapping +from collections.abc import Mapping from typing import Final, Literal from unittest.mock import patch @@ -22,7 +21,7 @@ from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.router_api import RoutingContext from litellm.router_strategy.least_busy import picks_reserve -from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy +from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup from litellm.utils import Rules, function_setup @@ -66,63 +65,6 @@ def _build_router(routing_strategy="simple-shuffle", routing_groups=None): ) -def test_no_groups_uses_top_level_strategy_for_all_models(): - router = _build_router(routing_strategy="latency-based-routing") - assert router._get_routing_context("filtered-model")[0] == "latency-based-routing" - assert router._get_routing_context("other-model")[0] == "latency-based-routing" - - -def test_explicit_group_overrides_top_level(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ], - ) - strategy, selector = router._get_routing_context("filtered-model") - assert strategy == "latency-based-routing" - assert selector is not None - # other-model isn't in any explicit group, so it lands in the default fallback - strategy_other, _ = router._get_routing_context("other-model") - assert strategy_other == "simple-shuffle" - - -def test_default_group_is_simple_shuffle_when_top_level_strategy_unset(): - router = _build_router() - strategy, selector = router._get_routing_context("other-model") - assert strategy == "simple-shuffle" - # simple-shuffle does not use a selector - assert selector is None - - -def test_two_groups_same_strategy_have_independent_selectors(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "group_a", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - "routing_strategy_args": {"ttl": 60}, - }, - { - "group_name": "group_b", - "models": ["other-model"], - "routing_strategy": "latency-based-routing", - "routing_strategy_args": {"ttl": 600}, - }, - ], - ) - a_selector = router._group_selectors["group_a"]["latency-based-routing"] - b_selector = router._group_selectors["group_b"]["latency-based-routing"] - assert a_selector is not None and b_selector is not None - assert id(a_selector) != id(b_selector) - - def test_overlapping_models_across_groups_raises(): with pytest.raises(ValueError, match="appears in"): _build_router( @@ -167,23 +109,6 @@ def test_invalid_strategy_in_group_raises(): ) -def test_unknown_model_in_group_warns_but_does_not_raise(caplog): - import logging - - with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): - router = _build_router( - routing_groups=[ - { - "group_name": "g1", - "models": ["model-not-in-list"], - "routing_strategy": "latency-based-routing", - } - ], - ) - assert "model-not-in-list" in caplog.text - assert router._model_to_group.get("model-not-in-list") == "g1" - - def test_duplicate_group_name_raises(): with pytest.raises(ValueError, match="duplicate"): _build_router( @@ -202,250 +127,6 @@ def test_duplicate_group_name_raises(): ) -def test_routing_group_object_input_accepted(): - router = _build_router( - routing_groups=[ - RoutingGroup( - group_name="g1", - models=["filtered-model"], - routing_strategy="latency-based-routing", - ) - ], - ) - assert router._model_to_group["filtered-model"] == "g1" - - -@pytest.mark.asyncio -async def test_async_dispatch_uses_group_strategy_for_grouped_model(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ], - ) - - group_selector = router._group_selectors["fast"]["latency-based-routing"] - - with ( - patch.object( - group_selector, - "async_get_available_deployments", - wraps=group_selector.async_get_available_deployments, - ) as latency_spy, - patch( - "litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle - ) as shuffle_spy, - ): - await router.async_get_available_deployment( - model="filtered-model", request_kwargs={} - ) - - assert ( - latency_spy.called - ), "group's latency selector should run for grouped model" - assert not shuffle_spy.called - - -@pytest.mark.asyncio -async def test_async_dispatch_falls_back_to_default_for_ungrouped_models(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ], - ) - - group_selector = router._group_selectors["fast"]["latency-based-routing"] - - with ( - patch.object( - group_selector, - "async_get_available_deployments", - wraps=group_selector.async_get_available_deployments, - ) as latency_spy, - patch( - "litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle - ) as shuffle_spy, - ): - await router.async_get_available_deployment( - model="other-model", request_kwargs={} - ) - - assert shuffle_spy.called, "default group's simple-shuffle should run" - assert not latency_spy.called - - -def test_update_settings_round_trip_routing_groups(): - router = _build_router() - assert router._model_to_group == {} - - router.update_settings( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ] - ) - assert router._model_to_group == {"filtered-model": "fast"} - assert router._get_routing_context("filtered-model")[0] == "latency-based-routing" - - settings = router.get_settings() - assert any(g["group_name"] == "fast" for g in settings["routing_groups"]) - - -def test_default_group_accepts_routing_strategy_enum_for_top_level_strategy(): - """ - The Router constructor accepts both string and RoutingStrategy enum forms. - The default group must resolve a real selector for either input shape, or - every ungrouped model raises NoDeploymentAvailable. - """ - router = _build_router(routing_strategy=RoutingStrategy.LATENCY_BASED) - strategy, selector = router._get_routing_context("other-model") - assert strategy == "latency-based-routing" - assert selector is not None - assert router.lowestlatency_logger is selector - - -@pytest.mark.asyncio -async def test_async_dispatch_uses_default_selector_when_constructed_with_enum(): - router = _build_router(routing_strategy=RoutingStrategy.LATENCY_BASED) - default_selector = router.lowestlatency_logger - assert default_selector is not None - - # filtered-model has 2 deployments, so the single-deployment short-circuit - # in async_get_available_deployment doesn't kick in and we actually hit - # the latency selector. - with ( - patch.object( - default_selector, - "async_get_available_deployments", - wraps=default_selector.async_get_available_deployments, - ) as latency_spy, - patch( - "litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle - ) as shuffle_spy, - ): - await router.async_get_available_deployment( - model="filtered-model", request_kwargs={} - ) - - assert ( - latency_spy.called - ), "default group's latency selector must run when constructed with the enum" - assert not shuffle_spy.called - - -def test_update_settings_does_not_leak_strategy_callbacks(monkeypatch): - """ - Repeated `update_settings` calls must not accumulate stale selectors in - `litellm.callbacks` / `litellm.input_callback`. Each rebuild owns the - previous generation and is responsible for unregistering it by identity - so the global lists don't grow without bound. - """ - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - - router = _build_router(routing_strategy="latency-based-routing") - initial_default_selector = router.lowestlatency_logger - assert initial_default_selector is not None - assert any( - c is initial_default_selector for c in litellm.callbacks - ), "fresh router should register its default selector" - - # Flip strategies back and forth — each transition triggers - # routing_strategy_init and must replace, not accumulate. - for _ in range(3): - router.update_settings(routing_strategy="usage-based-routing") - router.update_settings(routing_strategy="latency-based-routing") - - assert all( - c is not initial_default_selector for c in litellm.callbacks - ), "old default selector instance leaked into litellm.callbacks" - - # The set of router-owned selectors should be bounded — at most one per - # strategy class. Without the cleanup fix this grows by one per toggle. - router_owned_classes = ( - "LeastBusyLoggingHandler", - "LowestTPMLoggingHandler", - "LowestTPMLoggingHandler_v2", - "LowestLatencyLoggingHandler", - "LowestCostLoggingHandler", - ) - router_owned = [ - c for c in litellm.callbacks if type(c).__name__ in router_owned_classes - ] - assert len(router_owned) <= len( - router_owned_classes - ), f"router selectors leaked: {[type(c).__name__ for c in router_owned]}" - - # Group selectors: capture the v1 instance, swap, confirm v1 is gone by id. - router.update_settings( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ] - ) - group_selector_v1 = router._group_selectors["fast"]["latency-based-routing"] - assert group_selector_v1 is not None - - router.update_settings( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - "routing_strategy_args": {"ttl": 600}, - } - ] - ) - assert all( - c is not group_selector_v1 for c in litellm.callbacks - ), "old group selector instance leaked into litellm.callbacks" - group_selector_v2 = router._group_selectors["fast"]["latency-based-routing"] - assert group_selector_v2 is not group_selector_v1 - - -def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeypatch): - """ - Setting routing_groups to an empty list (or omitting all groups) must - unregister every previously-owned group selector. - """ - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - - router = _build_router( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "least-busy", - } - ] - ) - group_selector = router._group_selectors["fast"]["least-busy"] - assert group_selector in litellm.callbacks - assert group_selector in litellm.input_callback - - router.update_settings(routing_groups=[]) - - assert all(c is not group_selector for c in litellm.callbacks) - assert all(c is not group_selector for c in litellm.input_callback) - assert router._group_selectors == {} - - def test_two_least_busy_groups_count_a_request_once(monkeypatch): """ Least-busy counts a request up from the pre-call hooks on `litellm.input_callback` and @@ -558,326 +239,10 @@ def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): # --------------------------------------------------------------------------- -def test_normalize_strategy_handles_string_enum_and_none(): - assert Router._normalize_strategy(None) is None - assert Router._normalize_strategy("simple-shuffle") == "simple-shuffle" - assert ( - Router._normalize_strategy(RoutingStrategy.LATENCY_BASED) - == RoutingStrategy.LATENCY_BASED.value - ) - - -def test_validate_routing_strategy_accepts_valid_and_rejects_invalid(): - router = _build_router() - # Valid: string, enum, None - router._validate_routing_strategy("simple-shuffle") - router._validate_routing_strategy(RoutingStrategy.LATENCY_BASED) - router._validate_routing_strategy(None) - with pytest.raises(ValueError, match="Invalid routing_strategy"): - router._validate_routing_strategy("not-a-real-strategy") - - -def test_build_strategy_selector_returns_none_for_simple_shuffle(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router() - assert ( - router._build_strategy_selector( - strategy="simple-shuffle", - routing_strategy_args={}, - register_callbacks=False, - ) - is None - ) - - -def test_build_strategy_selector_constructs_for_known_strategies(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router() - selector = router._build_strategy_selector( - strategy="latency-based-routing", - routing_strategy_args={"ttl": 30}, - register_callbacks=False, - ) - assert selector is not None - - -def test_unregister_router_selectors_removes_by_identity(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router() - selector = router._build_strategy_selector( - strategy="least-busy", - routing_strategy_args={}, - register_callbacks=True, - ) - assert selector in litellm.callbacks - assert selector in litellm.input_callback - router._unregister_router_selectors([selector]) - assert all(c is not selector for c in litellm.callbacks) - assert all(c is not selector for c in litellm.input_callback) - - -def test_init_routing_groups_with_none_clears_state(): - router = _build_router( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "least-busy", - } - ] - ) - assert router.routing_groups - router._init_routing_groups(None) - assert router.routing_groups == {} - assert router._model_to_group == {} - assert router._group_selectors == {} - - -@pytest.mark.asyncio -async def test_select_deployment_async_returns_none_without_selector(): - router = _build_router() - result = await router._select_deployment_async( - strategy="simple-shuffle", - selector=None, - model="filtered-model", - healthy_deployments=[], - messages=None, - input=None, - request_kwargs=None, - reserve=False, - ) - assert result is None - - -def test_select_deployment_sync_returns_none_without_selector(): - router = _build_router() - result = router._select_deployment_sync( - strategy="simple-shuffle", - selector=None, - model="filtered-model", - healthy_deployments=[], - messages=None, - input=None, - request_kwargs=None, - reserve=False, - ) - assert result is None - - -def test_select_deployment_sync_does_not_dispatch_cost_based_routing(): - """ - `LowestCostLoggingHandler` only implements `async_get_available_deployments`, - so the sync dispatch must fall through to None for `cost-based-routing` - instead of raising AttributeError. - """ - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "cheap", - "models": ["filtered-model"], - "routing_strategy": "cost-based-routing", - } - ], - ) - cost_selector = router._group_selectors["cheap"]["cost-based-routing"] - assert not hasattr(cost_selector, "get_available_deployments"), ( - "test premise broken: LowestCostLoggingHandler unexpectedly grew a " - "sync method — revisit the sync dispatch arm." - ) - - result = router._select_deployment_sync( - strategy="cost-based-routing", - selector=cost_selector, - model="filtered-model", - healthy_deployments=[], - messages=None, - input=None, - request_kwargs=None, - reserve=False, - ) - assert result is None - - -@pytest.mark.asyncio -async def test_async_dispatch_falls_back_to_sync_for_usage_based_routing_v1(): - """ - `LowestTPMLoggingHandler` (v1) only implements the sync - `get_available_deployments`. The async dispatch must call the sync method - on the v1 selector instead of awaiting a non-existent async one. - """ - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "v1", - "models": ["filtered-model"], - "routing_strategy": "usage-based-routing", - } - ], - ) - v1_selector = router._group_selectors["v1"]["usage-based-routing"] - assert not hasattr(v1_selector, "async_get_available_deployments"), ( - "test premise broken: LowestTPMLoggingHandler v1 unexpectedly grew an " - "async method — revisit the async dispatch arm." - ) - - with patch.object( - v1_selector, - "get_available_deployments", - wraps=v1_selector.get_available_deployments, - ) as v1_spy: - await router._select_deployment_async( - strategy="usage-based-routing", - selector=v1_selector, - model="filtered-model", - healthy_deployments=[ - { - "model_name": "filtered-model", - "litellm_params": {"model": "openai/gpt-4o"}, - "model_info": {"id": "deploy-1"}, - } - ], - messages=[{"role": "user", "content": "hi"}], - input=None, - request_kwargs={}, - reserve=False, - ) - - assert v1_spy.called, "async dispatch must route v1 strategy through sync method" - - -def test_request_routing_strategy_override_beats_top_level(): - router = _build_router(routing_strategy="least-busy") - strategy, selector = router._get_routing_context( - "other-model", {"routing_strategy": "simple-shuffle"} - ) - assert strategy == "simple-shuffle" - assert selector is None - - -def test_request_routing_strategy_override_beats_explicit_group(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ], - ) - strategy, _ = router._get_routing_context( - "filtered-model", {"routing_strategy": "least-busy"} - ) - assert strategy == "least-busy" - - -def test_request_routing_strategy_override_builds_and_caches_selector(): - router = _build_router(routing_strategy="simple-shuffle") - strategy, selector = router._get_routing_context( - "other-model", {"routing_strategy": "latency-based-routing"} - ) - assert strategy == "latency-based-routing" - assert selector is not None - _, selector_again = router._get_routing_context( - "other-model", {"routing_strategy": "latency-based-routing"} - ) - assert selector_again is selector - - -def test_request_routing_strategy_override_matching_global_reuses_default_selector(): - router = _build_router(routing_strategy="least-busy") - _, selector = router._get_routing_context( - "other-model", {"routing_strategy": "least-busy"} - ) - assert selector is router.leastbusy_logger - assert router._override_selectors == {} - - -def test_invalid_request_routing_strategy_override_falls_back(): - router = _build_router(routing_strategy="least-busy") - strategy, selector = router._get_routing_context( - "other-model", {"routing_strategy": "not-a-real-strategy"} - ) - assert strategy == "least-busy" - assert selector is router.leastbusy_logger - - -def test_no_override_key_keeps_existing_behavior(): - router = _build_router(routing_strategy="least-busy") - strategy, _ = router._get_routing_context("other-model", {"messages": []}) - assert strategy == "least-busy" - strategy_none_kwargs, _ = router._get_routing_context("other-model", None) - assert strategy_none_kwargs == "least-busy" - - -def test_request_routing_strategy_override_helper_validates_directly(): - router = _build_router(routing_strategy="least-busy") - assert router._get_request_routing_strategy_override({"routing_strategy": "simple-shuffle"}) == "simple-shuffle" - assert router._get_request_routing_strategy_override({"routing_strategy": RoutingStrategy.LEAST_BUSY}) == "least-busy" - assert router._get_request_routing_strategy_override({"routing_strategy": "lar1"}) is None - assert router._get_request_routing_strategy_override({"routing_strategy": {"bad": "type"}}) is None - assert router._get_request_routing_strategy_override({}) is None - assert router._get_request_routing_strategy_override(None) is None - - -def test_override_strategy_selector_helper_builds_per_strategy(): - router = _build_router(routing_strategy="least-busy") - latency_selector = router._get_override_strategy_selector("latency-based-routing") - assert latency_selector is not None - assert router._get_override_strategy_selector("latency-based-routing") is latency_selector - assert router._get_override_strategy_selector("least-busy") is router.leastbusy_logger - assert router._get_override_strategy_selector("simple-shuffle") is None - - -def test_strategy_reinit_unregisters_override_selectors(): - router = _build_router(routing_strategy="least-busy") - override_selector = router._get_override_strategy_selector("latency-based-routing") - assert override_selector is not None - assert not any(cb is override_selector for cb in litellm.callbacks) - - router.update_settings(routing_strategy="latency-based-routing") - - assert router._override_selectors == {} - assert not any(cb is override_selector for cb in litellm.callbacks) - assert router._get_override_strategy_selector("latency-based-routing") is router.lowestlatency_logger - - def _single_latency_group(): return [{"group_name": "g1", "models": ["filtered-model"], "routing_strategy": "latency-based-routing"}] -def _assert_still_routes_with_original_group(router, selector): - assert list(router.routing_groups) == ["g1"] - assert router._model_to_group == {"filtered-model": "g1"} - assert router._group_selectors["g1"]["latency-based-routing"] is selector - assert router._get_routing_context("filtered-model", None) == ("latency-based-routing", selector) - assert sum(1 for cb in litellm.callbacks if cb is selector) == 1 - - -def test_failed_routing_groups_update_keeps_previous_groups(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_groups=_single_latency_group()) - selector = router._group_selectors["g1"]["latency-based-routing"] - - with pytest.raises(ValueError, match="appears in"): - router.update_settings( - routing_groups=[ - *_single_latency_group(), - {"group_name": "g2", "models": ["filtered-model"], "routing_strategy": "least-busy"}, - ], - ) - - _assert_still_routes_with_original_group(router, selector) - assert sum(1 for cb in litellm.callbacks if type(cb) is not type(selector)) == 0 - assert litellm.input_callback == [] - - def test_failed_routing_groups_update_does_not_poison_later_strategy_changes(monkeypatch): monkeypatch.setattr(litellm, "callbacks", []) monkeypatch.setattr(litellm, "input_callback", []) @@ -918,106 +283,6 @@ def test_overlap_error_names_every_conflicting_model(): assert "'other-model' appears in 'g1' and 'g2'" in message -def test_invalid_group_strategy_keeps_previous_groups(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_groups=_single_latency_group()) - selector = router._group_selectors["g1"]["latency-based-routing"] - - with pytest.raises(ValueError, match="Invalid routing_strategy"): - router.update_settings( - routing_groups=[ - {"group_name": "g2", "models": ["other-model"], "routing_strategy": "not-a-real-strategy"}, - ], - ) - - _assert_still_routes_with_original_group(router, selector) - - -def test_unbuildable_group_selector_keeps_previous_groups(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_groups=_single_latency_group()) - selector = router._group_selectors["g1"]["latency-based-routing"] - - with pytest.raises(ValidationError, match="ttl"): - router.update_settings( - routing_groups=[ - {"group_name": "g0", "models": ["other-model"], "routing_strategy": "least-busy"}, - *_single_latency_group(), - { - "group_name": "g2", - "models": ["other-model-2"], - "routing_strategy": "latency-based-routing", - "routing_strategy_args": {"ttl": "not-a-number"}, - }, - ], - ) - - _assert_still_routes_with_original_group(router, selector) - assert litellm.callbacks == [selector] - assert litellm.input_callback == [] - - -def test_register_router_selector_wires_only_the_hooks_the_strategy_needs(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router() - least_busy = router._build_strategy_selector( - strategy="least-busy", routing_strategy_args={}, register_callbacks=False - ) - latency = router._build_strategy_selector( - strategy="latency-based-routing", routing_strategy_args={}, register_callbacks=False - ) - assert least_busy is not None and latency is not None - assert litellm.callbacks == [] and litellm.input_callback == [] - - router._register_router_selector(least_busy) - router._register_router_selector(latency) - - assert [cb for cb in litellm.callbacks if cb is least_busy or cb is latency] == [least_busy, latency] - assert litellm.input_callback == [least_busy] - - -def test_replace_routing_groups_swaps_state_and_callbacks_in_one_step(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_groups=_single_latency_group()) - old_selector = router._group_selectors["g1"]["latency-based-routing"] - new_selector = router._build_strategy_selector( - strategy="least-busy", routing_strategy_args={}, register_callbacks=False - ) - assert new_selector is not None - - router._replace_routing_groups( - ( - (RoutingGroup(group_name="g2", models=["other-model"], routing_strategy="least-busy"), new_selector), - (RoutingGroup(group_name="g3", models=["other-model-2"], routing_strategy="simple-shuffle"), None), - ) - ) - - assert list(router.routing_groups) == ["g2", "g3"] - assert router._model_to_group == {"other-model": "g2", "other-model-2": "g3"} - assert router._group_selectors == {"g2": {"least-busy": new_selector}, "g3": {}} - assert router._get_routing_context("other-model", None) == ("least-busy", new_selector) - assert router._get_routing_context("filtered-model", None)[0] == router.routing_strategy - assert all(cb is not old_selector for cb in litellm.callbacks) - assert sum(1 for cb in litellm.callbacks if cb is new_selector) == 1 - assert litellm.input_callback == [new_selector] - - -def test_override_selectors_are_not_registered_process_wide(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_strategy="simple-shuffle") - v2_selector = router._get_override_strategy_selector("usage-based-routing-v2") - least_busy_selector = router._get_override_strategy_selector("least-busy") - assert v2_selector is not None and least_busy_selector is not None - - assert litellm.callbacks == [] - assert litellm.input_callback == [] - - def _rpm_limited_model_list(): return [ { @@ -1033,81 +298,6 @@ def _rpm_limited_model_list(): ] -async def _mock_completion(router, **override): - return await router.acompletion( - model="other-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok", **override - ) - - -@pytest.mark.asyncio -async def test_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(): - router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) - - await _mock_completion(router, routing_strategy="usage-based-routing-v2") - override_selector = router._override_selectors["usage-based-routing-v2"] - assert not any(cb is override_selector for cb in litellm.callbacks) - - with patch.object( - override_selector, "async_pre_call_check", wraps=override_selector.async_pre_call_check - ) as pre_call_spy: - for _ in range(2): - plain = await _mock_completion(router) - assert plain.choices[0].message.content == "ok" - assert not pre_call_spy.called - - with pytest.raises(litellm.RateLimitError): - await _mock_completion(router, routing_strategy="usage-based-routing-v2") - - -def test_sync_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(monkeypatch): - monkeypatch.setattr( - "litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", - lambda: datetime.datetime(2026, 1, 1, 12, 0, 30, tzinfo=datetime.timezone.utc), - ) - router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) - messages = [{"role": "user", "content": "hi"}] - - router.completion( - model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2" - ) - override_selector = router._override_selectors["usage-based-routing-v2"] - assert not any(cb is override_selector for cb in litellm.callbacks) - - for _ in range(2): - plain = router.completion(model="other-model", messages=messages, mock_response="ok") - assert plain.choices[0].message.content == "ok" - - with pytest.raises(ValueError, match="No deployments available"): - router.completion( - model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2" - ) - - -@pytest.mark.asyncio -async def test_override_selector_pre_call_check_only_runs_for_override_selectors(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - deployment = _rpm_limited_model_list()[0] - - override_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle") - override_selector = override_router._get_override_strategy_selector("usage-based-routing-v2") - await override_router._async_override_selector_pre_call_check( - "usage-based-routing-v2", override_selector, deployment, None - ) - with pytest.raises(litellm.RateLimitError): - override_router._override_selector_pre_call_check("usage-based-routing-v2", override_selector, deployment) - - default_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="usage-based-routing-v2") - for _ in range(2): - await default_router._async_override_selector_pre_call_check( - "usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment, None - ) - default_router._override_selector_pre_call_check( - "usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment - ) - await default_router._async_override_selector_pre_call_check(None, None, deployment, None) - default_router._override_selector_pre_call_check(None, None, deployment) - - @pytest.mark.asyncio async def test_usage_based_v2_override_enforces_rpm_when_a_specific_deployment_is_requested(): router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) @@ -1191,149 +381,6 @@ def _proxy_shaped_request(**data: object) -> dict[str, object]: return {**kwargs, "litellm_logging_obj": logging_obj} -async def _async_override_pick(router: Router, strategy: str) -> str: - deployment = await router.async_get_available_deployment( - "grp", request_kwargs=_proxy_shaped_request(model="grp", routing_strategy=strategy) - ) - return deployment["model_info"]["id"] - - -def _sync_override_pick(router: Router, strategy: str) -> str: - deployment = router.get_available_deployment( - "grp", request_kwargs=_proxy_shaped_request(model="grp", routing_strategy=strategy) - ) - return deployment["model_info"]["id"] - - -def _in_flight(router: Router, deployment_id: str) -> int | None: - return router.cache.get_cache(f"grp_request_count:{deployment_id}") - - -async def _async_wait_until(predicate: Callable[[], bool]) -> None: - for _ in range(100): - if predicate(): - return - await asyncio.sleep(0.02) - raise AssertionError("lifecycle callback never reached the override selector") - - -def _sync_wait_until(predicate: Callable[[], bool]) -> None: - for _ in range(100): - if predicate(): - return - time.sleep(0.02) - raise AssertionError("lifecycle callback never reached the override selector") - - -def _selector_is_not_global(selector: CustomLogger) -> bool: - global_lists = ( - litellm.callbacks, - litellm.input_callback, - litellm.success_callback, - litellm.failure_callback, - litellm._async_success_callback, - litellm._async_failure_callback, - ) - return not any(cb is selector for cbs in global_lists for cb in cbs) - - -@pytest.mark.asyncio -async def test_least_busy_override_sees_the_overriding_request_in_flight(): - router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle", num_retries=0) - - stream = await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True)) - busy = stream._hidden_params["model_id"] - idle = "d2" if busy == "d1" else "d1" - assert [await _async_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] - - async for _ in stream: - pass - await _async_wait_until(lambda: _in_flight(router, busy) == 0) - assert await _async_override_pick(router, "least-busy") == "d1" - assert _selector_is_not_global(router._override_selectors["least-busy"]) - - -def test_sync_least_busy_override_sees_the_overriding_request_in_flight(): - router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle", num_retries=0) - - stream = router.completion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True)) - busy = stream._hidden_params["model_id"] - idle = "d2" if busy == "d1" else "d1" - assert [_sync_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] - - for _ in stream: - pass - _sync_wait_until(lambda: _in_flight(router, busy) == 0) - assert _sync_override_pick(router, "least-busy") == "d1" - assert _selector_is_not_global(router._override_selectors["least-busy"]) - - -@pytest.mark.asyncio -async def test_least_busy_override_releases_the_slot_when_the_overriding_request_fails(): - router = Router( - model_list=_two_deployment_model_list(mock_response="litellm.InternalServerError"), - routing_strategy="simple-shuffle", - num_retries=0, - ) - - with pytest.raises(litellm.InternalServerError): - await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy")) - - await _async_wait_until(lambda: _in_flight(router, "d1") == 0) - assert await _async_override_pick(router, "least-busy") == "d1" - assert _selector_is_not_global(router._override_selectors["least-busy"]) - - -@pytest.mark.asyncio -async def test_latency_based_override_learns_from_the_overriding_requests(): - router = Router( - model_list=_two_deployment_model_list(mock_delay=0.05), routing_strategy="simple-shuffle", num_retries=0 - ) - - def samples(deployment_id: str) -> list[float]: - recorded = (router.cache.get_cache("grp_map") or {}).get(deployment_id, {}).get("latency_normalized", []) - return [latency for latency in recorded if latency > 0] - - async def overriding_call() -> str: - sampled_before = {"d1": len(samples("d1")), "d2": len(samples("d2"))} - response = await router.acompletion( - **_proxy_shaped_request(model="grp", routing_strategy="latency-based-routing") - ) - deployment_id = response._hidden_params["model_id"] - await _async_wait_until(lambda: len(samples(deployment_id)) > sampled_before[deployment_id]) - return deployment_id - - served = [await overriding_call() for _ in range(6)] - - assert "d1" in served - assert served[2:] == ["d2"] * 4 - assert _selector_is_not_global(router._override_selectors["latency-based-routing"]) - - -def test_override_selector_is_bound_only_to_the_request_that_asked_for_it(): - router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle") - overriding = _proxy_shaped_request(model="grp", routing_strategy="least-busy") - plain = _proxy_shaped_request(model="grp") - - router.get_available_deployment("grp", request_kwargs=overriding) - router.get_available_deployment("grp", request_kwargs=overriding) - router.get_available_deployment("grp", request_kwargs=plain) - - selector = router._override_selectors["least-busy"] - bound = overriding["litellm_logging_obj"] - for callbacks in ( - bound.dynamic_input_callbacks, - bound.dynamic_success_callbacks, - bound.dynamic_async_success_callbacks, - bound.dynamic_failure_callbacks, - bound.dynamic_async_failure_callbacks, - ): - assert callbacks == [selector] - unbound = plain["litellm_logging_obj"] - assert unbound.dynamic_input_callbacks is None and unbound.dynamic_success_callbacks is None - assert unbound.dynamic_failure_callbacks is None and unbound.dynamic_async_failure_callbacks is None - - def test_override_matching_the_router_strategy_is_not_bound_twice(): router = Router(model_list=_two_deployment_model_list(), routing_strategy="least-busy") request = _proxy_shaped_request(model="grp", routing_strategy="least-busy") @@ -1343,52 +390,20 @@ def test_override_matching_the_router_strategy_is_not_bound_twice(): assert request["litellm_logging_obj"].dynamic_input_callbacks is None -@pytest.mark.asyncio -async def test_override_matching_a_routing_group_strategy_records_each_request_once(): - router = Router( - model_list=_two_deployment_model_list(), - routing_strategy="simple-shuffle", - routing_groups=[RoutingGroup(group_name="lat", models=["grp"], routing_strategy="latency-based-routing")], - num_retries=0, - ) - request = _proxy_shaped_request(model="grp", routing_strategy="latency-based-routing") - assert router._globally_registered_strategies() == {"simple-shuffle", "latency-based-routing"} - - response = await router.acompletion(**request) - deployment_id = response._hidden_params["model_id"] - await _async_wait_until(lambda: (router.cache.get_cache("grp_map") or {}).get(deployment_id) is not None) - - assert len(router.cache.get_cache("grp_map")[deployment_id]["latency_normalized"]) == 1 - assert request["litellm_logging_obj"].dynamic_success_callbacks is None - - -def test_bind_override_selector_to_request_binds_once_and_ignores_requests_without_logging(): - router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle") - selector = router._get_override_strategy_selector("least-busy") - request = _proxy_shaped_request(model="grp", routing_strategy="least-busy") - request["litellm_logging_obj"].dynamic_success_callbacks = ["langfuse"] - - router._bind_override_selector_to_request("least-busy", selector, request) - router._bind_override_selector_to_request("least-busy", selector, request) - router._bind_override_selector_to_request("least-busy", selector, None) - router._bind_override_selector_to_request("least-busy", selector, {"model": "grp"}) - - logging_obj = request["litellm_logging_obj"] - assert logging_obj.dynamic_success_callbacks == ["langfuse", selector] - assert logging_obj.dynamic_input_callbacks == [selector] - assert logging_obj.dynamic_async_failure_callbacks == [selector] - assert _selector_is_not_global(selector) - - def _quality_group(strategy="latency-based-routing"): return [{"group_name": "quality", "models": ["filtered-model", "other-model"], "routing_strategy": strategy}] -def test_group_name_is_callable_and_unions_member_deployments(): +async def _candidate_ids(router: Router, model: str) -> list[str]: + deployments: Final = await router.async_get_healthy_deployments(model=model, request_kwargs={}) + assert isinstance(deployments, list) + return sorted(d["model_info"]["id"] for d in deployments) + + +@pytest.mark.asyncio +async def test_group_name_is_callable_and_unions_member_deployments(): router = _build_router(routing_groups=_quality_group()) - model, deployments = router._common_checks_available_deployment(model="quality") - assert model == "quality" - assert sorted(d["model_info"]["id"] for d in deployments) == ["deploy-1", "deploy-2", "deploy-3"] + assert await _candidate_ids(router, "quality") == ["deploy-1", "deploy-2", "deploy-3"] def test_group_name_appears_in_model_names_and_model_list(): @@ -1399,34 +414,8 @@ def test_group_name_appears_in_model_names_and_model_list(): assert sorted(r["model_info"]["id"] for r in rows) == ["deploy-1", "deploy-2", "deploy-3"] -def test_get_routing_context_for_group_name_uses_group_strategy(): - router = _build_router(routing_groups=_quality_group()) - strategy, selector = router._get_routing_context("quality") - assert strategy == "latency-based-routing" - assert selector is router._group_selectors["quality"]["latency-based-routing"] - - @pytest.mark.asyncio -async def test_group_call_dispatches_via_group_selector(): - router = _build_router(routing_groups=_quality_group()) - group_selector = router._group_selectors["quality"]["latency-based-routing"] - - with ( - patch.object( - group_selector, - "async_get_available_deployments", - wraps=group_selector.async_get_available_deployments, - ) as latency_spy, - patch("litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle) as shuffle_spy, - ): - deployment = await router.async_get_available_deployment(model="quality", request_kwargs={}) - - assert latency_spy.called - assert not shuffle_spy.called - assert deployment["model_name"] in {"filtered-model", "other-model"} - - -def test_group_name_colliding_with_model_name_is_shadowed_with_warning(caplog): +async def test_group_name_colliding_with_model_name_is_shadowed_with_warning(caplog): import logging with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): @@ -1437,13 +426,11 @@ def test_group_name_colliding_with_model_name_is_shadowed_with_warning(caplog): ) assert any("shadowed" in record.getMessage() for record in caplog.records) assert router.get_routing_group("filtered-model") is None - assert router._get_routing_context("other-model")[0] == "latency-based-routing" - - model, deployments = router._common_checks_available_deployment(model="filtered-model") - assert sorted(d["model_info"]["id"] for d in deployments) == ["deploy-1", "deploy-2"] + assert await _candidate_ids(router, "filtered-model") == ["deploy-1", "deploy-2"] -def test_group_name_colliding_with_model_group_alias_is_shadowed_with_warning(caplog): +@pytest.mark.asyncio +async def test_group_name_colliding_with_model_group_alias_is_shadowed_with_warning(caplog): import logging with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): @@ -1454,13 +441,11 @@ def test_group_name_colliding_with_model_group_alias_is_shadowed_with_warning(ca ) assert any("shadowed" in record.getMessage() for record in caplog.records) assert router.get_routing_group("quality") is None - - model, deployments = router._common_checks_available_deployment(model="quality") - assert model == "filtered-model" - assert sorted(d["model_info"]["id"] for d in deployments) == ["deploy-1", "deploy-2"] + assert await _candidate_ids(router, "quality") == ["deploy-1", "deploy-2"] -def test_real_model_added_later_shadows_group(): +@pytest.mark.asyncio +async def test_real_model_added_later_shadows_group(): router = _build_router(routing_groups=_quality_group()) assert router.get_routing_group("quality") is not None @@ -1474,34 +459,32 @@ def test_real_model_added_later_shadows_group(): ) ) assert router.get_routing_group("quality") is None - model, deployments = router._common_checks_available_deployment(model="quality") - assert [d["model_info"]["id"] for d in deployments] == ["deploy-shadow"] + assert await _candidate_ids(router, "quality") == ["deploy-shadow"] router.delete_deployment(id="deploy-shadow") assert "quality" not in router.model_names assert router.get_routing_group("quality") is not None - _, restored = router._common_checks_available_deployment(model="quality") - assert sorted(d["model_info"]["id"] for d in restored) == ["deploy-1", "deploy-2", "deploy-3"] + assert await _candidate_ids(router, "quality") == ["deploy-1", "deploy-2", "deploy-3"] -def test_group_with_no_member_deployments_raises_no_healthy(): +@pytest.mark.asyncio +async def test_group_with_no_member_deployments_raises_no_healthy(): router = Router( model_list=_model_list(), routing_groups=[{"group_name": "empty-group", "models": ["ghost-model"], "routing_strategy": "simple-shuffle"}], ) with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment(model="empty-group") + await router.async_get_healthy_deployments(model="empty-group", request_kwargs={}) -def test_alias_pointing_at_group_composes(): +@pytest.mark.asyncio +async def test_alias_pointing_at_group_composes(): router = Router( model_list=_model_list(), model_group_alias={"quality-alias": "quality"}, routing_groups=_quality_group(), ) - model, deployments = router._common_checks_available_deployment(model="quality-alias") - assert model == "quality" - assert sorted(d["model_info"]["id"] for d in deployments) == ["deploy-1", "deploy-2", "deploy-3"] + assert await _candidate_ids(router, "quality-alias") == ["deploy-1", "deploy-2", "deploy-3"] def test_model_group_info_reports_group(): @@ -1525,11 +508,10 @@ def test_routing_group_has_alternatives(): assert solo_router.routing_group_has_alternatives("solo-group") is False -def test_member_direct_call_unchanged_by_callable_groups(): +@pytest.mark.asyncio +async def test_member_direct_call_unchanged_by_callable_groups(): router = _build_router(routing_groups=_quality_group()) - model, deployments = router._common_checks_available_deployment(model="other-model") - assert model == "other-model" - assert [d["model_info"]["id"] for d in deployments] == ["deploy-3"] + assert await _candidate_ids(router, "other-model") == ["deploy-3"] def test_update_settings_group_change_invalidates_model_group_info(): @@ -1587,29 +569,19 @@ def test_group_rows_cache_invalidated_on_model_list_change(): assert sum(1 for row in router.get_model_list() if row["model_name"] == "quality") == 4 -def _pin_choice_to(deployment_id): - def _pick(seq): - for candidate in seq: - if candidate["model_info"]["id"] == deployment_id: - return candidate - return seq[0] - - return _pick - - async def _call_and_get_cooldowns(router, model): - from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments - - with ( - patch("litellm.router_strategy.simple_shuffle.random.choice", side_effect=_pin_choice_to("deploy-3")), - pytest.raises(litellm.RateLimitError), - ): + """Cool down the other members first, so the call has to land on deploy-3.""" + for model_id in ("deploy-1", "deploy-2"): + router.cooldown_cache.add_deployment_to_cooldown( + model_id=model_id, original_exception=RuntimeError("unavailable"), exception_status=500, cooldown_time=60 + ) + with pytest.raises(litellm.RateLimitError): await router.acompletion( model=model, messages=[{"role": "user", "content": "hi"}], mock_response="litellm.RateLimitError", ) - return await async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) + return [model_id for model_id, _ in await router.cooldown_cache.async_get_active_cooldowns(["deploy-3"], None)] @pytest.mark.asyncio @@ -1668,15 +640,6 @@ def test_group_rows_do_not_inherit_member_access_groups(): assert "access_groups" in router.get_model_list(model_name="gated-member")[0]["model_info"] -def test_group_rebuild_invalidates_access_groups_cache(): - router = _build_router(routing_groups=_quality_group()) - router.get_model_access_groups() - assert router._access_groups_cache is not None - - router.update_settings(routing_groups=[]) - assert router._access_groups_cache is None - - def test_get_model_list_from_routing_groups_materializes_rows(): router = _build_router(routing_groups=_quality_group()) rows = router.get_model_list_from_routing_groups() @@ -1688,28 +651,6 @@ def test_get_model_list_from_routing_groups_materializes_rows(): assert router.get_model_list_from_routing_groups(model_name="filtered-model") == () -def test_get_routing_group_deployments_unions_members(): - router = _build_router(routing_groups=_quality_group()) - union = router._get_routing_group_deployments("quality") - assert sorted(d["model_info"]["id"] for d in union) == ["deploy-1", "deploy-2", "deploy-3"] - assert router._get_routing_group_deployments("filtered-model") is None - - -def test_materialize_routing_group_rows_labels_members_with_group_name(): - router = _build_router(routing_groups=_quality_group()) - group = router.get_routing_group("quality") - rows = router._materialize_routing_group_rows((group,)) - assert {row["model_name"] for row in rows} == {"quality"} - assert len(rows) == 3 - - -def test_as_routing_group_row_strips_access_groups(): - source = {"model_name": "member", "model_info": {"id": "d1", "access_groups": ["restricted"]}} - row = Router._as_routing_group_row(source) - assert row["model_info"] == {"id": "d1"} - assert source["model_info"]["access_groups"] == ["restricted"] - - @pytest.mark.asyncio async def test_group_call_429_cools_down_member_across_retries(): router = Router( @@ -1755,30 +696,6 @@ def _priority_deployments( ] -def test_priority_group_affinity_scope_follows_group_aliases_and_settings_reload() -> None: - router: Final = Router( - model_list=_priority_deployments(), - routing_groups=[ - _priority_group(), - RoutingGroup(group_name="legacy", models=["filtered-model"], routing_strategy="simple-shuffle"), - ], - model_group_alias={"priority-alias": "priority-group"}, - ) - requested_models: Final = ( - "priority-group", "priority-alias", "filtered-model", "other-model", "legacy", "missing" - ) - assert tuple(router._is_priority_routing_group(model) for model in requested_models) == ( - True, True, False, False, False, False - ) - - router.update_settings(routing_groups=[{ - "group_name": "priority-group", - "models": ["filtered-model", "other-model"], - "routing_strategy": "simple-shuffle", - }]) - assert tuple(router._is_priority_routing_group(model) for model in requested_models) == (False,) * 6 - - @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("requested_model", ["priority-group", "priority-alias"]) @pytest.mark.asyncio @@ -1837,8 +754,6 @@ async def test_opposite_priority_groups_preserve_deployments_and_legacy_member_p model="reverse-group", messages=[{"role": "user", "content": "hi"}] ) assert (forward.choices[0].message.content, reverse.choices[0].message.content) == ("primary", "backup") - assert router._get_routing_context("filtered-model")[0] == "least-busy" - assert router._get_routing_context("other-model")[0] == "latency-based-routing" assert router.get_model_list(model_name="filtered-model") == before assert all(deployment["litellm_params"]["order"] == 10 for deployment in before) @@ -2002,67 +917,6 @@ async def test_auto_router_selected_priority_group_fails_over_inside_the_group() assert response.choices[0].message.content == "backup" -@pytest.mark.parametrize("endpoint_request", [{"input": "hi"}, {"messages": [{"role": "user", "content": "hi"}]}]) -@pytest.mark.asyncio -async def test_priority_generic_fallback_keeps_routing_controls_out_of_provider_kwargs( - endpoint_request: dict[str, object], -) -> None: - async def provider(model: str, **provider_kwargs: object) -> str: - assert "_target_order" not in provider_kwargs - assert "model_priorities" not in provider_kwargs - model_info: Final = provider_kwargs["model_info"] - assert isinstance(model_info, dict) - if model_info["id"] != "deploy-3": - raise litellm.RateLimitError(message="primary refused", model=model, llm_provider="openai") - return "backup" - - router: Final = Router( - model_list=_priority_deployments(), - routing_groups=[_priority_group()], - num_retries=0, - disable_cooldowns=True, - ) - response: Final = await router._ageneric_api_call_with_fallbacks( - model="priority-group", original_function=provider, **endpoint_request - ) - assert response == "backup" - - -@pytest.mark.parametrize("provider_model_id", [None, "resolved-downstream-id"]) -@pytest.mark.asyncio -async def test_priority_generic_fallback_dict_preserves_served_model_id(provider_model_id: str | None) -> None: - async def provider(model: str, **provider_kwargs: object) -> dict[str, object]: - model_info: Final = provider_kwargs["model_info"] - assert isinstance(model_info, dict) - if model_info["id"] == "deploy-1": - raise litellm.NotFoundError(message="primary missing", model=model, llm_provider="openai") - assert model_info["id"] == "deploy-3" - return { - "content": "backup", - **({"_hidden_params": {"model_id": provider_model_id}} if provider_model_id is not None else {}), - } - - router: Final = Router( - model_list=_priority_deployments()[::2], - routing_groups=[_priority_group()], - num_retries=0, - disable_cooldowns=True, - ) - outer_metadata: Final[dict[str, object]] = {} - response: Final = await router._ageneric_api_call_with_fallbacks( - model="priority-group", - original_function=provider, - messages=[{"role": "user", "content": "hi"}], - litellm_metadata=outer_metadata, - ) - primary_info: Final = outer_metadata["model_info"] - assert isinstance(primary_info, dict) - assert primary_info["id"] == "deploy-1" - assert response["content"] == "backup" - assert response["_hidden_params"]["model_id"] == (provider_model_id or "deploy-3") - assert response["_hidden_params"]["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 - - @pytest.mark.asyncio async def test_priority_group_ignores_cached_backup_affinity() -> None: from litellm.router_utils.prompt_caching_cache import PromptCachingCache diff --git a/tests/unit/router_strategy/test_router_routing_plugins.py b/tests/unit/router_strategy/test_router_routing_plugins.py index 0a54addce5e..b04c5edd3c2 100644 --- a/tests/unit/router_strategy/test_router_routing_plugins.py +++ b/tests/unit/router_strategy/test_router_routing_plugins.py @@ -259,56 +259,6 @@ async def test_router_without_plugins_is_unaffected(): assert response.choices[0].message.content == "hi" -@pytest.mark.asyncio -async def test_run_routing_plugins_narrows_candidates_and_records_signals(): - """Unit-level check of _run_routing_plugins in isolation, independent of acompletion.""" - router = Router( - model_list=_smart_router_model_list(), - plugins=[LanguageDetector(), DomainClassifier(), TenantPolicy(), BudgetPolicy()], - ) - request_kwargs = {"metadata": {"tenant": "acme-corp"}} - - context = await router._run_routing_plugins( - model="smart-router", - request_kwargs=request_kwargs, - messages=[{"role": "user", "content": "hi"}], - ) - - assert context.candidate_models == ["openai/gpt-4o-mini", "anthropic/claude-haiku-4-5"] - assert context.signals["domain-classifier"]["domain"] == "coding" - assert request_kwargs["metadata"]["_routing_plugin_candidate_models"] == context.candidate_models - - -def test_filter_by_routing_plugin_candidates_narrows_and_raises_when_empty(): - """Unit-level check of _filter_by_routing_plugin_candidates in isolation.""" - router = Router(model_list=_smart_router_model_list(), plugins=[TenantPolicy()]) - healthy_deployments = router.model_list - - narrowed = router._filter_by_routing_plugin_candidates( - healthy_deployments=healthy_deployments, - request_kwargs={"metadata": {"_routing_plugin_candidate_models": ["openai/gpt-4o-mini"]}}, - ) - assert [d["litellm_params"]["model"] for d in narrowed] == ["openai/gpt-4o-mini"] - - with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): - router._filter_by_routing_plugin_candidates( - healthy_deployments=healthy_deployments, - request_kwargs={"metadata": {"_routing_plugin_candidate_models": ["nonexistent/model"]}}, - ) - - -def test_json_default_stable_id_is_stable_across_instances(): - """generate_model_id's json.dumps `default=` fallback must not embed an object's - memory address (e.g. plain str() on an object with no custom __repr__ falls back - to object.__repr__'s ``) -- that would make the - deployment id churn on every process restart for any deployment whose - litellm_params contain a live plugin instance.""" - router = Router(model_list=_smart_router_model_list()) - - assert router._json_default_stable_id(LanguageDetector()) == router._json_default_stable_id(LanguageDetector()) - assert router._json_default_stable_id(LanguageDetector()) != router._json_default_stable_id(TenantPolicy()) - - def testgenerate_model_id_is_stable_when_litellm_params_contain_a_plugin_instance(): """End-to-end: a deployment id built from litellm_params containing a routing plugin instance (e.g. complexity_router_config.plugins) must be identical across diff --git a/tests/unit/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py index 0c19cfc72f6..0d3876c64a3 100644 --- a/tests/unit/router_strategy/test_router_tag_routing.py +++ b/tests/unit/router_strategy/test_router_tag_routing.py @@ -4,7 +4,6 @@ import logging from typing import Final, Literal -from unittest.mock import patch import pytest @@ -1846,23 +1845,22 @@ async def test_tag_fail_open_survives_deployment_narrowing( ) -> None: router: Final = _fail_open_router(route=route) requested_model: Final = "chat" if route == "model" else "chat-test" if route == "wildcard" else "public-chat" - with patch("time.time", return_value=1_700_000_000.0): - if narrowing == "cooldown": - router.cooldown_cache.add_deployment_to_cooldown( - model_id="opted-in", - original_exception=RuntimeError("unavailable"), - exception_status=500, - cooldown_time=60, - ) - candidates: Final = await router.async_get_healthy_deployments( - model=requested_model, - request_kwargs={ - "litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"}, - **( - {"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {} - ), - }, + if narrowing == "cooldown": + router.cooldown_cache.add_deployment_to_cooldown( + model_id="opted-in", + original_exception=RuntimeError("unavailable"), + exception_status=500, + cooldown_time=60, ) + candidates: Final = await router.async_get_healthy_deployments( + model=requested_model, + request_kwargs={ + "litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"}, + **( + {"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {} + ), + }, + ) assert isinstance(candidates, (list, tuple)) assert tuple(deployment["model_info"]["id"] for deployment in candidates) == ("remaining",) @@ -2700,6 +2698,13 @@ def _quality_high_cost_low_router(): ) +def _cool_down(router: litellm.Router, *model_ids: str) -> None: + for model_id in model_ids: + router.cooldown_cache.add_deployment_to_cooldown( + model_id=model_id, original_exception=RuntimeError("unavailable"), exception_status=500, cooldown_time=60 + ) + + @pytest.mark.asyncio() async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default(): # Every deployment in the group is tagged "default" (a legitimate cross-cutting @@ -2707,21 +2712,16 @@ async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default # the quality:high deployments unhealthy, a request asking for quality:high # must still hard-fail, not silently get served by a cost:low deployment just # because it happens to also carry "default". - from unittest.mock import AsyncMock, patch - router = _quality_high_cost_low_router() + _cool_down(router, "quality-high-1", "quality-high-2") - with patch( - "litellm.router_python_core.async_get_cooldown_deployments", - new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]), - ): - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: - await router.acompletion( - model="gpt-4", - messages=[{"role": "user", "content": "hi"}], - metadata={"tags": ["quality:high"]}, - mock_response="hi", - ) + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["quality:high"]}, + mock_response="hi", + ) from litellm.types.router import RouterErrors @@ -2734,18 +2734,14 @@ async def test_plain_tag_exhaustion_with_universal_default_tag_falls_open_when_a for deployment in router.model_list: deployment["model_info"]["allow_fail_open"] = True - from unittest.mock import AsyncMock, patch + _cool_down(router, "quality-high-1", "quality-high-2") - with patch( - "litellm.router_python_core.async_get_cooldown_deployments", - new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]), - ): - response = await router.acompletion( - model="gpt-4", - messages=[{"role": "user", "content": "hi"}], - metadata={"tags": ["quality:high"]}, - mock_response="hi", - ) + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["quality:high"]}, + mock_response="hi", + ) assert response._hidden_params["model_id"] in ("cost-low-1", "cost-low-2") @@ -3095,24 +3091,6 @@ async def test_router_selecting_tag_is_consumed_on_litellm_metadata_shaped_reque assert deployment["model_info"]["id"] == "tier-gemini-flash" -def test_consumed_request_tags_stamp_names_the_routed_group_and_spent_tags_only_on_a_tag_match(): - from litellm.types.router import ConsumedRequestTagsStamp, PreRoutingHookResponse - - router = _tagged_marker_router() - strategy = router.auto_routers["gpt4o"][0] - rewrite = PreRoutingHookResponse(model="gemini-flash", messages=None) - - consumed = router._consumed_request_tags_stamp( - selected_strategy=strategy, pre_routing_hook_response=rewrite, request_tags=["route"] - ) - unmatched = router._consumed_request_tags_stamp( - selected_strategy=strategy, pre_routing_hook_response=rewrite, request_tags=["other"] - ) - - assert consumed == ConsumedRequestTagsStamp(model_group="gemini-flash", tags=("route",)) - assert unmatched is None - - @pytest.mark.asyncio() async def test_tagged_request_direct_to_plain_group_still_rejected(): # Sent straight to the tier, no router selection consumed the tag, so strict diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index fa75f980cc0..191f8463cfa 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -440,41 +440,18 @@ async def test_async_user_parameter_does_not_trigger_deployment_affinity(): model_group = "azure-sdk-test" user_id = "sdk-user-123" - choice_calls = {"count": 0} - - def deterministic_choice(seq): - choice_calls["count"] += 1 - if choice_calls["count"] == 1: - return seq[0] - return seq[1] if len(seq) > 1 else seq[0] - - with ( - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post, - patch( - "litellm.router_strategy.simple_shuffle.random.choice", - side_effect=deterministic_choice, - ), - ): + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: mock_post.return_value = MockResponse(mock_response_data, 200) - # First call with 'user' parameter (end-user id) - first_response = await router.aresponses( - model=model_group, - input="Hi", - user=user_id, - ) - first_model_id = first_response._hidden_params["model_id"] - - # Second call with same 'user' parameter should NOT be pinned by affinity - second_response = await router.aresponses( - model=model_group, - input="Follow-up", - user=user_id, - ) - assert second_response._hidden_params["model_id"] != first_model_id + # Repeated calls with the same 'user' parameter must keep spreading across the group + model_ids = set() + for _ in range(20): + response = await router.aresponses(model=model_group, input="Hi", user=user_id) + model_ids.add(response._hidden_params["model_id"]) + assert len(model_ids) == 2 @pytest.mark.asyncio @@ -839,61 +816,29 @@ async def test_model_group_affinity_config_only_applies_to_configured_group(): ) user_api_key_hash = "test-per-group-key" - choice_calls = {"count": 0} - def deterministic_choice(seq): - choice_calls["count"] += 1 - if choice_calls["count"] == 1: - return seq[0] - return seq[1] if len(seq) > 1 else seq[0] + async def served_model_ids(model: str) -> set[str]: + model_ids = set() + for _ in range(20): + response = await router.aresponses( + model=model, + input="Hello", + truncation="auto", + litellm_metadata={"user_api_key_hash": user_api_key_hash}, + ) + model_ids.add(response._hidden_params["model_id"]) + return model_ids - with ( - patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post, - patch( - "litellm.router_strategy.simple_shuffle.random.choice", - side_effect=deterministic_choice, - ), - ): + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: mock_post.return_value = MockResponse(mock_response_data, 200) - # gpt-4: affinity should work — second request pinned to same deployment - first = await router.aresponses( - model="gpt-4", - input="Hello", - truncation="auto", - litellm_metadata={"user_api_key_hash": user_api_key_hash}, - ) - first_model_id = first._hidden_params["model_id"] - - second = await router.aresponses( - model="gpt-4", - input="Follow-up", - truncation="auto", - litellm_metadata={"user_api_key_hash": user_api_key_hash}, - ) - assert second._hidden_params["model_id"] == first_model_id - - # claude-3: no affinity configured — should NOT be pinned - choice_calls["count"] = 0 - first_claude = await router.aresponses( - model="claude-3", - input="Hello", - truncation="auto", - litellm_metadata={"user_api_key_hash": user_api_key_hash}, - ) - first_claude_id = first_claude._hidden_params["model_id"] - - second_claude = await router.aresponses( - model="claude-3", - input="Follow-up", - truncation="auto", - litellm_metadata={"user_api_key_hash": user_api_key_hash}, - ) - # With deterministic choice and len>1, second call picks seq[1] - assert second_claude._hidden_params["model_id"] != first_claude_id + # gpt-4: affinity pins every request from the key to one deployment + assert len(await served_model_ids("gpt-4")) == 1 + # claude-3: no affinity configured, so requests keep spreading across the group + assert len(await served_model_ids("claude-3")) == 2 @pytest.mark.asyncio diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index a4d279e32e6..76a574bae24 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -849,31 +849,19 @@ async def test_affinity_falls_back_to_same_encryption_boundary_on_model_group_sw num_retries=0, ) - def first_call_picks_account_a(seq): - for d in seq: - if d["model_info"]["id"] == "gpt-5.3-codex-account-a": - return d - return seq[0] - - with ( - patch( - "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", - new_callable=AsyncMock, - return_value=first_resp, - ), - patch( - "litellm.router_strategy.simple_shuffle.random.choice", - side_effect=first_call_picks_account_a, - ), + with patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new_callable=AsyncMock, + return_value=first_resp, ): r1 = await router.aresponses(model="gpt-5.3-codex", input="hi") - assert r1._hidden_params["model_id"] == "gpt-5.3-codex-account-a" + origin_account = r1._hidden_params["model_id"].removeprefix("gpt-5.3-codex-") + assert origin_account in ("account-a", "account-b") encoded_id = _extract_encoded_item_id(r1) assert encoded_id.startswith("encitem_") - # simple_shuffle.random.choice NOT patched: prove affinity narrows the - # candidate pool to a single deployment regardless of which one shuffle picks. + # Affinity narrows the candidate pool to the origin's account regardless of which one shuffle picks. with patch( "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", new_callable=AsyncMock, @@ -890,7 +878,7 @@ async def test_affinity_falls_back_to_same_encryption_boundary_on_model_group_sw ], ) - assert r2._hidden_params["model_id"] == "gpt-5.4-account-a" + assert r2._hidden_params["model_id"] == f"gpt-5.4-{origin_account}" @pytest.mark.asyncio @@ -977,25 +965,15 @@ async def test_affinity_falls_back_to_same_boundary_on_alias_switch(): num_retries=0, ) - def pick_account_a(seq): - for d in seq: - if d["model_info"]["id"] == "gpt-5.3-codex-account-a": - return d - return seq[0] - - with ( - patch( - "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", - new_callable=AsyncMock, - return_value=first_resp, - ), - patch( - "litellm.router_strategy.simple_shuffle.random.choice", - side_effect=pick_account_a, - ), + with patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new_callable=AsyncMock, + return_value=first_resp, ): r1 = await router.aresponses(model="gpt-5.3-codex", input="hi") + origin_account = r1._hidden_params["model_id"].removeprefix("gpt-5.3-codex-") + assert origin_account in ("account-a", "account-b") encoded_id = _extract_encoded_item_id(r1) assert encoded_id.startswith("encitem_") @@ -1015,7 +993,7 @@ async def test_affinity_falls_back_to_same_boundary_on_alias_switch(): ], ) - assert r2._hidden_params["model_id"] == "gpt-5.2-codex-account-a" + assert r2._hidden_params["model_id"] == f"gpt-5.2-codex-{origin_account}" def test_boundary_fallback_no_router_ref_returns_empty(): diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 9487e4d6321..5169642bc43 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1480,3 +1480,58 @@ async def test_a_pin_written_without_tools_is_read_when_tools_is_an_empty_list() ) assert filtered == [deployments[0]] + + +class _LoggedMessagesProbe(CustomLogger): + def __init__(self) -> None: + self.logged_messages: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_messages.append((kwargs["standard_logging_object"]["messages"], kwargs["messages"])) + + +@pytest.mark.asyncio +async def test_pin_is_written_with_message_logging_off_while_loggers_still_see_redacted_messages( + monkeypatch, local_model_cost_map +): + """ + `turn_off_message_logging` redacted the shared success-event payload before the check read it, so the + pin was keyed off `redacted-by-litellm`, never written, and a cacheable conversation bounced across the + group. The pin must land whatever the flag says, while a real logger alongside still sees only the + redacted messages. + """ + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + probe = _LoggedMessagesProbe() + monkeypatch.setattr(litellm, "callbacks", [probe]) + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-fake"}, + "model_info": {"id": model_id}, + } + for model_id in ("dep-1", "dep-2") + ], + optional_pre_call_checks=["prompt_caching"], + ) + messages = _messages(word_count=5000) + + first = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok") + served_by = first._hidden_params["model_id"] + + affinity_key = PromptCachingCache.get_prompt_caching_cache_key(messages, None) + pin = await _eventually(lambda: router.prompt_caching_pin_cache.get_cache(key=affinity_key)) + assert pin == {"model_id": served_by}, pin + + subsequent = [ + (await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok"))._hidden_params[ + "model_id" + ] + for _ in range(4) + ] + assert subsequent == [served_by] * 4 + + await _eventually(lambda: len(probe.logged_messages) == 5) + redacted = [{"role": "user", "content": "redacted-by-litellm"}] + assert probe.logged_messages == [(redacted, redacted)] * 5, probe.logged_messages diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 6b4ea74894a..eb4e72cb9e6 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -146,26 +146,6 @@ async def test_healthy_legacy_string_pin_is_preserved_with_repinning_enabled() - router.discard() -@pytest.mark.parametrize("route_bucket", ["metadata", "litellm_metadata"]) -async def test_repin_session_after_success_reads_split_metadata(route_bucket: str) -> None: - router: Final = _repinning_router() - try: - await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) - _cool_down_original(router) - session_bucket: Final = "litellm_metadata" if route_bucket == "metadata" else "metadata" - await router._repin_session_after_success( - { - "model_info": {"id": "replacement"}, - route_bucket: {"deployment_model_name": "repin-group"}, - session_bucket: {"session_id": "repin-session", "user_api_key_hash": "caller-1"}, - }, - {"status": "completed"}, - ) - assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "replacement"} - finally: - router.discard() - - @pytest.mark.parametrize("enabled", [False, True]) async def test_session_selection_observes_another_workers_repin_only_when_enabled(enabled: bool) -> None: server: Final = FakeServer() @@ -320,49 +300,6 @@ async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: M router.discard() -class _ReadCountingCache(InMemoryCache): - def __init__(self) -> None: - super().__init__() - self.read_keys: tuple[str, ...] = () - - def get_cache(self, key: str, **kwargs: object) -> object: - self.read_keys += (key,) - return super().get_cache(key, **kwargs) - - -@pytest.mark.parametrize("case", ["disabled", "no-session", "generated-session", "group-disabled", "active"]) -async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str) -> None: - router: Final = _repinning_router(case != "disabled") - reads: Final = _ReadCountingCache() - try: - router.cooldown_cache.cooldown_store.in_memory_cache = reads - callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) - if case == "group-disabled": - callback.model_group_affinity_config = {"repin-group": []} - await router._repin_session_after_success( - { - "model_info": {"id": "replacement"}, - "metadata": { - "deployment_model_name": "repin-group", - **({"session_id": "repin-session"} if case != "no-session" else {}), - SESSION_ID_GENERATED_METADATA_KEY: case == "generated-session", - }, - }, - {"status": "completed"}, - ) - expected: Final = ( - { - router.cooldown_cache.get_cooldown_cache_key(model_id) - for model_id in ("original", "replacement", "other") - } - if case == "active" - else set() - ) - assert set(reads.read_keys) == expected - finally: - router.discard() - - class MockResponse: def __init__(self, json_data, status_code): self._json_data = json_data @@ -690,15 +627,19 @@ async def test_auto_router_session_affinity_writes_scoped_pin_and_follows_it(): ) async def test_seeded_session_pin_is_invisible_outside_its_scope(model, key_hash): """The pin binds (auto-routed request, api key, session): a direct call to the - group and a different key reusing the session id must both ignore it.""" - router = _smart_router() - try: - await router.cache.async_set_cache( - key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"} - ) - assert await _one_turn(router, model, "scoped-session", key_hash) == "deployment-1" - finally: - _cleanup_router_callbacks(router) + group and a different key reusing the session id must both ignore it: across fresh + routers each seeded with the pin, the request still reaches the other deployment.""" + served = set() + for _ in range(20): + router = _smart_router() + try: + await router.cache.async_set_cache( + key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"} + ) + served.add(await _one_turn(router, model, "scoped-session", key_hash)) + finally: + _cleanup_router_callbacks(router) + assert "deployment-1" in served @pytest.mark.asyncio diff --git a/tests/unit/router_utils/test_client_initalization_utils.py b/tests/unit/router_utils/test_client_initalization_utils.py index 6f9a7b730ac..30791f47623 100644 --- a/tests/unit/router_utils/test_client_initalization_utils.py +++ b/tests/unit/router_utils/test_client_initalization_utils.py @@ -81,14 +81,17 @@ def test_slot_is_released_when_the_held_call_raises(): assert limit.in_flight == 1 -def _router_limit(router: Router, model_name: str) -> MaxParallelRequestsLimit: - deployment: Final = router.get_deployment_by_model_group_name(model_group_name=model_name) - assert deployment is not None - client: Final = router._get_client( - deployment=deployment.model_dump(), kwargs={}, client_type="max_parallel_requests" - ) - assert isinstance(client, MaxParallelRequestsLimit) - return client +async def _concurrent_outcomes(router: Router, calls: int) -> list[str]: + async def attempt() -> str: + try: + await router.acompletion( + model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="ok", mock_delay=0.3 + ) + except litellm.RateLimitError as e: + return e.message + return "ok" + + return await asyncio.gather(*[attempt() for _ in range(calls)]) @pytest.mark.parametrize( @@ -103,23 +106,25 @@ def _router_limit(router: Router, model_name: str) -> MaxParallelRequestsLimit: @pytest.mark.asyncio async def test_router_deployment_rejects_past_its_derived_cap(litellm_params: dict[str, int], expected_cap: int): router: Final = Router( - model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", **litellm_params}}] + model_list=[ + {"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "k", **litellm_params}} + ], + num_retries=0, ) - limit: Final = _router_limit(router, "gpt-5.6") - assert limit.max_parallel_requests == expected_cap - release: Final = asyncio.Event() - holders: Final = [asyncio.create_task(_hold(limit, release)) for _ in range(expected_cap)] - await asyncio.sleep(0) - assert limit.in_flight == expected_cap - assert f"max_parallel_requests={expected_cap}" in _expect_rejection(limit).message - release.set() - assert await asyncio.wait_for(asyncio.gather(*holders), timeout=2) == ["ok"] * expected_cap + + outcomes: Final = await _concurrent_outcomes(router, expected_cap + 1) + + assert outcomes.count("ok") == expected_cap + rejections: Final = [outcome for outcome in outcomes if outcome != "ok"] + assert len(rejections) == 1 + assert f"max_parallel_requests={expected_cap}" in rejections[0] -def test_router_without_any_concurrency_setting_has_no_limit(): - router: Final = Router(model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6"}}]) - deployment: Final = router.get_deployment_by_model_group_name(model_group_name="gpt-5.6") - assert deployment is not None - assert ( - router._get_client(deployment=deployment.model_dump(), kwargs={}, client_type="max_parallel_requests") is None +@pytest.mark.asyncio +async def test_router_without_any_concurrency_setting_has_no_limit(): + router: Final = Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "k"}}], + num_retries=0, ) + + assert await _concurrent_outcomes(router, 20) == ["ok"] * 20 diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 71c37b5b1d5..6a35c8a9602 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -36,6 +36,7 @@ from litellm.router_utils.cooldown_handlers import ( async_get_cooldown_deployments, cast_exception_status_to_int, get_cooldown_deployments, + is_response_lookup_not_found, mark_advisor_orchestration_failure, set_cooldown_deployments, should_cooldown_based_on_allowed_fails_policy, @@ -1277,80 +1278,6 @@ class TestResolveAllowedFailsFromPolicyFallsThrough: assert result == 3, "must fall through to BadRequestErrorAllowedFails when the more specific field is unset" -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -class TestDeploymentCallbackOnFailureCooldownTimePrecedence: - def test_model_info_cooldown_time_used_in_primary_sync_path(self): - """ - Router.deployment_callback_on_failure (the primary sync failure-callback path, - as opposed to the fallback path covered by TestFallbackDeploymentCooldown) must - also honor a model_info.cooldown_time, not just litellm_params.cooldown_time. - """ - router = _make_router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, - "model_info": {"id": "primary", "cooldown_time": 15.0}, - }, - ], - ) - - exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") - kwargs = { - "exception": exc, - "litellm_params": { - "model_info": {"id": "primary", "cooldown_time": 15.0}, - }, - } - - with patch("litellm.router_python_core.set_cooldown_deployments") as mock_set_cooldown: - router.deployment_callback_on_failure( - kwargs=kwargs, - completion_response=None, - start_time=0, - end_time=1, - ) - - mock_set_cooldown.assert_called_once() - call_kwargs = mock_set_cooldown.call_args[1] - assert call_kwargs["time_to_cooldown"] == 15.0, ( - "model_info.cooldown_time must be honored in the primary sync failure-callback path" - ) - - def test_litellm_params_cooldown_time_still_honored_as_fallback(self): - """cooldown_time has pre-existing litellm_params support on this primary - path; it must keep working when model_info doesn't set it.""" - router = _make_router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4", "cooldown_time": 20.0}, - "model_info": {"id": "primary"}, - }, - ], - ) - - exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") - kwargs = { - "exception": exc, - "litellm_params": { - "model_info": {"id": "primary"}, - "cooldown_time": 20.0, - }, - } - - with patch("litellm.router_python_core.set_cooldown_deployments") as mock_set_cooldown: - router.deployment_callback_on_failure( - kwargs=kwargs, - completion_response=None, - start_time=0, - end_time=1, - ) - - call_kwargs = mock_set_cooldown.call_args[1] - assert call_kwargs["time_to_cooldown"] == 20.0, "litellm_params.cooldown_time must still be honored" - - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCallerScopedOAuthAuthFailureCooldown: @staticmethod @@ -2465,22 +2392,19 @@ async def test_cooldown_time_zero_uses_zero_not_default(): num_retries=0, ) - with patch.object(router.cooldown_cache, "add_deployment_to_cooldown") as mock_add_cooldown: - try: - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hey, how's it going?"}], - mock_response="litellm.RateLimitError", - ) - except litellm.RateLimitError: - pass - - mock_add_cooldown.assert_not_called() + try: + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_response="litellm.RateLimitError", + ) + except litellm.RateLimitError: + pass cooldown_list = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) assert len(cooldown_list) == 0 - healthy_deployments, _ = await router._async_get_healthy_deployments(model="gpt-3.5-turbo", parent_otel_span=None) + healthy_deployments = await router.select("gpt-3.5-turbo", {}, kind="healthy") assert len(healthy_deployments) == 1 @@ -2559,20 +2483,17 @@ def test_single_deployment_no_cooldowns(num_deployments: int): router = Router(model_list=model_list, num_retries=0) - with patch.object(router.cooldown_cache, "add_deployment_to_cooldown", new=MagicMock()) as mock_client: - try: - router.completion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hey, how's it going?"}], - mock_response="litellm.RateLimitError", - ) - except litellm.RateLimitError: - pass + try: + router.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_response="litellm.RateLimitError", + ) + except litellm.RateLimitError: + pass - if num_deployments == 1: - mock_client.assert_not_called() - else: - mock_client.assert_called_once() + cooled = get_cooldown_deployments(RoutingContext(router), parent_otel_span=None) + assert len(cooled) == (0 if num_deployments == 1 else 1) @pytest.mark.asyncio @@ -3608,33 +3529,6 @@ def test_naming_tag_routed_deployments_leaves_the_request_metadata_untouched() - assert metadata == {"tags": ["tenant-a"]} -@pytest.mark.parametrize( - ("has_request_input", "expected"), ((True, (TENANT_A, SHARED)), (False, EVERY_TENANT)), ids=("chat", "no-input") -) -def test_region_routable_deployments_follow_the_pre_call_region_check( - has_request_input: bool, expected: tuple[CooldownDeploymentInfo, ...] -) -> None: - router: Final = _tenant_router((), **REGION_CHECKS_ON) - kept: Final = router._region_routable_deployments( - router.get_model_list(model_name="tenants") or [], {"allowed_model_region": "eu"}, has_request_input - ) - assert router._deployment_ids(kept) == frozenset(deployment["model_id"] for deployment in expected) - - -def test_routable_cooled_deployment_ids_skip_healthy_deployments_and_peel_tag_fallbacks() -> None: - router: Final = _tenant_router((TENANT_A, SHARED), **TAGS_ON) - routable: Final = asyncio.run( - router._async_routable_cooled_deployment_ids( - model="tenants", - request_deployments=router.get_model_list(model_name="tenants") or [], - cooldown_deployments=(TENANT_A["model_id"], SHARED["model_id"]), - request_kwargs={"metadata": {}}, - has_request_input=True, - ) - ) - assert routable == frozenset((TENANT_A["model_id"], SHARED["model_id"])) - - REACHABILITY_REQUESTS: Final = ( {"metadata": {"tags": ["tenant-a"]}}, {"metadata": {}}, @@ -3672,3 +3566,25 @@ def test_the_refusal_names_each_cooled_deployment_some_cooldown_state_routes_the with pytest.raises(RouterRateLimitError) as refusal: _select(_tenant_router(EVERY_TENANT, **BOTH_ON), "async", "tenants", copy.deepcopy(request_kwargs)) assert {deployment["model_id"] for deployment in refusal.value.cooldown_deployments or ()} == reachable + + +@pytest.mark.parametrize( + ("status_code", "call_type", "expected"), + [ + (404, "aget_responses", True), + (404, "adelete_responses", True), + (404, "acancel_responses", True), + (404, "alist_input_items", True), + (404, "get_responses", True), + (404, "aresponses", False), + (404, "acompact_responses", False), + (404, "acompletion", False), + (500, "aget_responses", False), + (429, "alist_input_items", False), + ], +) +def test_is_response_lookup_not_found(status_code: int, call_type: str, expected: bool) -> None: + error: Final = litellm.NotFoundError(message="gone", llm_provider="openai", model="gpt-4o-mini") + error.status_code = status_code + + assert is_response_lookup_not_found(error, {"call_type": call_type}) is expected diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index a386542d61b..6263292fb40 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -10,13 +10,11 @@ from unittest.mock import MagicMock, patch import httpx import pytest -import respx from pydantic import TypeAdapter import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger -from litellm.llms.anthropic.messages.native.fake_stream_iterator import FakeAnthropicMessagesStreamIterator from litellm.litellm_core_utils import get_llm_provider_logic from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -39,7 +37,6 @@ from litellm.router_utils.fallback_event_handlers import ( is_bad_request_without_fallback, log_failure_fallback_event, log_success_fallback_event, - mid_stream_fallback_hop_kwargs, mid_stream_fallback_snapshot_kwargs, mid_stream_retry_kwargs, record_pre_routing_selection, @@ -49,9 +46,6 @@ from litellm.router_utils.fallback_event_handlers import ( ) from litellm.types.router import FallbackSkipReason, SkippedFallbackTarget from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome -from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE -from typing import Dict -import os class StreamingWrapper: @@ -245,65 +239,10 @@ async def test_run_async_fallback_preserves_superseded_usage_when_hop_fails_befo assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) -class _RateLimitedHop(Exception): - def __init__(self, deployment_id: str = "fallback-dep-1") -> None: - super().__init__("fallback deployment rate limited") - self.status_code = 429 - self.failed_deployment_id = deployment_id - - _EARLIER_FAILURE_MARKER: Final = datetime(2026, 1, 1) -_FIRST_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 2) _FINAL_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 3) -def _two_failure_fallback_router() -> Router: - router: Final = Router( - model_list=[ - { - "model_name": "fallback-model", - "litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"}, - "model_info": {"id": f"fallback-dep-{index}"}, - } - for index in (1, 2) - ], - allowed_fails=0, - cooldown_time=60, - num_retries=0, - ) - - async def fail_both_hops(*args: object, **kwargs: object) -> NoReturn: - logging_obj: Final[object] = kwargs["litellm_logging_obj"] - assert isinstance(logging_obj, Logging) - assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False - - logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FIRST_HOP_FAILURE_MARKER - try: - raise _RateLimitedHop("fallback-dep-1") - except _RateLimitedHop: - logging_obj.mark_logging_complete("async_failure") - - logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER - raise _RateLimitedHop("fallback-dep-2") - - router.async_function_with_fallbacks = fail_both_hops - return router - - -def _two_deployment_router() -> Router: - return Router( - model_list=[ - { - "model_name": "fallback-model", - "litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"}, - "model_info": {"id": f"fallback-dep-{index}"}, - } - for index in (1, 2) - ], - num_retries=0, - ) - - def _entra_fallback_router() -> Router: return Router( model_list=[ @@ -361,67 +300,6 @@ async def test_fallback_hop_cools_down_only_a_deployment_entra_rejected_its_own_ assert await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) == cooled_down -@pytest.mark.asyncio -@pytest.mark.parametrize("failure_already_logged", [False, True]) -async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_logged: bool) -> None: - """The manual cooldown stands in for the failure callbacks the logging object blocks once a failure has - been logged. A hop whose own failure was the first one logged already ran them, so only a hop that failed - with the flag already set gets the manual cooldown.""" - router: Final = _two_deployment_router() - logging_obj: Final = _request_logging_obj() - if failure_already_logged: - logging_obj.model_call_details["current_attempt_api_call_start_time"] = _EARLIER_FAILURE_MARKER - logging_obj.mark_logging_complete("async_failure") - - async def rate_limited_hop(*args: object, **kwargs: object) -> NoReturn: - hop_logging_obj: Final = kwargs["litellm_logging_obj"] - assert isinstance(hop_logging_obj, Logging) - hop_logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER - if hop_logging_obj.model_call_details.get("has_logged_async_failure", False) is not True: - hop_logging_obj.mark_logging_complete("async_failure") - raise _RateLimitedHop() - - with pytest.raises(_RateLimitedHop): - await run_async_fallback( - litellm_routing_ctx=RoutingContext(router), - fallback_model_group=["fallback-model"], - original_model_group="primary-model", - original_exception=RuntimeError("stream dropped before content"), - max_fallbacks=3, - fallback_depth=0, - litellm_logging_obj=logging_obj, - original_function=rate_limited_hop, - metadata={}, - ) - - cooled_down: Final = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) - assert cooled_down == (["fallback-dep-1"] if failure_already_logged else []) - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("python_router_core") # the Rust core walks fallbacks in Rust, so it has no such member -async def test_run_async_fallback_cools_down_the_final_deployment_after_retry_failure_logging() -> None: - router: Final = _two_failure_fallback_router() - logging_obj: Final = _request_logging_obj() - assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False - - with pytest.raises(_RateLimitedHop): - await run_async_fallback( - litellm_routing_ctx=RoutingContext(router), - fallback_model_group=["fallback-model"], - original_model_group="primary-model", - original_exception=RuntimeError("stream dropped before content"), - max_fallbacks=3, - fallback_depth=0, - litellm_logging_obj=logging_obj, - metadata={}, - ) - - assert logging_obj.model_call_details.get("has_logged_async_failure", False) is True - cooled_down: Final = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) - assert cooled_down == ["fallback-dep-2"] - - @pytest.mark.asyncio async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_call(): """A nested fallback (multi-hop) must keep collecting errors, so the opt-in @@ -1362,7 +1240,7 @@ class TestTriggerCooldownForFailedDeployment: async def _stalled_stream_error() -> Exception: - from litellm.router import _aread_first_chunk + from litellm.router_rust.invoke.streaming import _aread_first_chunk # noqa: PLC2701 async def _stall(**kwargs: object) -> NoReturn: raise httpx.ReadTimeout("Timeout on reading data from socket") @@ -2244,41 +2122,6 @@ async def test_streaming_refusal_without_fallback_row_passes_frames_through(): assert len(fake.calls) == 1 -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -@pytest.mark.asyncio -async def test_refused_stream_built_from_a_finished_agentic_turn_is_delivered_without_a_fallback(): - """Websearch, MCP and advisor loops run their tools before streaming the final reply, so a - fallback would run every tool call again for a reply the client then gets from another model.""" - fake: Final = FakeAnthropicUpstream() - router: Final = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) - - with fake.install(): - stream: Final = await router._aanthropic_messages_streaming_iterator( - response=FakeAnthropicMessagesStreamIterator(REFUSAL_RESPONSE), - initial_kwargs=mid_stream_fallback_hop_kwargs( - model="fable-tier", - original_generic_function=litellm.anthropic_messages, - controls=None, - kwargs={ - "model": "fable-tier", - "max_tokens": 16, - "stream": True, - "messages": [{"role": "user", "content": "hi"}], - }, - ), - ) - events: Final = [ - json.loads(line[len(b"data: ") :]) - for line in (await _collect(stream)).splitlines() - if line.startswith(b"data: ") - ] - - assert fake.calls == [] - message_delta: Final = next(event for event in events if event["type"] == "message_delta") - assert message_delta["delta"]["stop_reason"] == "refusal" - assert message_delta["delta"]["stop_details"] == REFUSAL_RESPONSE["stop_details"] - - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata(): @@ -2424,28 +2267,6 @@ async def test_configured_content_policy_list_stays_authoritative_over_generic_r assert len(fake.calls) == 1 -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -def test_refusal_fallback_available_arms_on_generic_rows_only_without_content_policy(): - router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"tier-group": ["opus-target"]}]) - stamped = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}} - - assert router._refusal_fallback_available("router-group", stamped) is True - assert router._refusal_fallback_available("router-group", {}) is False - assert router._refusal_fallback_available("router-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False - - -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -def test_chat_content_filter_gate_unchanged_by_generic_rows(): - """The generic-row arming is scoped to /v1/messages safeguard refusals; the chat surface's - content_filter gate keeps its long-standing content-policy-only semantics.""" - from litellm.types.utils import Choices, ModelResponse - - router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}]) - response = ModelResponse(choices=[Choices(finish_reason="content_filter")]) - - assert router._should_raise_content_policy_error(model="fable-tier", response=response, kwargs={}) is False - - @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) @@ -2493,74 +2314,6 @@ async def test_disable_fallbacks_beats_a_content_policy_row_too(): assert len(fake.calls) == 1 -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -def test_refusal_gate_keys_on_pre_routing_tier_stamp(): - router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}]) - - def anthropic_messages(**kwargs: Any) -> None: - return None - - refusal_kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}} - assert ( - router._should_raise_anthropic_refusal_error( - model="router-group", - original_generic_function=anthropic_messages, - response=dict(REFUSAL_RESPONSE), - kwargs=refusal_kwargs, - ) - is True - ) - assert ( - router._should_raise_anthropic_refusal_error( - model="router-group", - original_generic_function=anthropic_messages, - response=dict(REFUSAL_RESPONSE), - kwargs={}, - ) - is False - ) - - -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -def test_has_content_policy_fallback_default_fallbacks_arm(): - router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}]) - - assert router._has_content_policy_fallback("any-group", {}) is True - assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False - - -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested(): - router = _router(content_policy_fallbacks=None) - fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}] - - assert router._get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, lookup_groups=("tier1", "smart-router") - ) == ["backup-a"] - assert router._get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, lookup_groups=("tier9", "smart-router") - ) == ["backup-b"] - assert router._get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=()) is None - - -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -def test_refusal_gate_ignores_other_generic_call_types(): - router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) - - def aresponses(**kwargs: Any) -> None: - return None - - assert ( - router._should_raise_anthropic_refusal_error( - model="fable-tier", - original_generic_function=aresponses, - response=dict(REFUSAL_RESPONSE), - kwargs={}, - ) - is False - ) - - class FallbackEventLogger(CustomLogger): def __init__(self) -> None: super().__init__() @@ -2600,61 +2353,6 @@ def _create_fallback_event_router() -> Router: ) -@pytest.mark.parametrize( - "function_name", - ["_acompletion", "_atext_completion", "_aembedding"], -) -@pytest.mark.asyncio -async def test_run_async_fallback(function_name): - """ - Basic test - given a list of fallback models, run the original function with the fallback models - """ - router = create_test_router() - original_function = getattr(router, function_name) - - litellm.set_verbose = True - fallback_model_group = ["gpt-4"] - original_model_group = "gpt-3.5-turbo" - original_exception = litellm.exceptions.InternalServerError( - message="Simulated error", - llm_provider="openai", - model="gpt-3.5-turbo", - ) - - request_kwargs = { - "mock_response": "hello this is a test for run_async_fallback", - "metadata": {"previous_models": ["gpt-3.5-turbo"]}, - } - - if function_name == "_aembedding": - request_kwargs["input"] = "hello this is a test for run_async_fallback" - elif function_name == "_atext_completion": - request_kwargs["prompt"] = "hello this is a test for run_async_fallback" - elif function_name == "_acompletion": - request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}] - - result = await run_async_fallback( - litellm_routing_ctx=RoutingContext(router), - original_function=original_function, - num_retries=1, - fallback_model_group=fallback_model_group, - original_model_group=original_model_group, - original_exception=original_exception, - max_fallbacks=5, - fallback_depth=0, - **request_kwargs, - ) - - assert result is not None - - if function_name == "_acompletion": - assert isinstance(result, litellm.ModelResponse) - elif function_name == "_atext_completion": - assert isinstance(result, litellm.TextCompletionResponse) - elif function_name == "_aembedding": - assert isinstance(result, litellm.EmbeddingResponse) - - @pytest.mark.asyncio async def test_log_success_fallback_event(): """ @@ -2714,71 +2412,6 @@ async def test_log_failure_fallback_event(): ) -@pytest.mark.asyncio -@pytest.mark.parametrize("function_name", ["_acompletion", "_atext_completion"]) -async def test_failed_fallbacks_raise_most_recent_exception(function_name): - """ - Tests that if all fallbacks fail, the most recent occuring exception is raised - - meaning the exception from the last fallback model is raised - """ - router = create_test_router() - original_function = getattr(router, function_name) - - fallback_model_group = ["gpt-4"] - original_model_group = "gpt-3.5-turbo" - original_exception = litellm.exceptions.InternalServerError( - message="Simulated error", - llm_provider="openai", - model="gpt-3.5-turbo", - ) - - request_kwargs: Dict[str, Any] = {"metadata": {"previous_models": ["gpt-3.5-turbo"]}} - - if function_name == "_aembedding": - request_kwargs["input"] = "hello this is a test for run_async_fallback" - elif function_name == "_atext_completion": - request_kwargs["prompt"] = "hello this is a test for run_async_fallback" - elif function_name == "_acompletion": - request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}] - - with pytest.raises(litellm.exceptions.RateLimitError): - await run_async_fallback( - litellm_routing_ctx=RoutingContext(router), - original_function=original_function, - num_retries=1, - fallback_model_group=fallback_model_group, - original_model_group=original_model_group, - original_exception=original_exception, - mock_response="litellm.RateLimitError", - max_fallbacks=5, - fallback_depth=0, - **request_kwargs, - ) - - -def create_test_router(): - return Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), - }, - }, - { - "model_name": "gpt-4", - "litellm_params": { - "model": "gpt-4", - "api_key": os.getenv("OPENAI_API_KEY"), - }, - }, - ], - fallbacks=[{"gpt-3.5-turbo": ["gpt-4"]}], - ) - - class CustomTestLogger(CustomLogger): def __init__(self): super().__init__() @@ -2800,120 +2433,6 @@ class CustomTestLogger(CustomLogger): self.failure_fallback_events.append((original_model_group, kwargs, original_exception)) -def create_test_router_2(): - return Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), - }, - }, - { - "model_name": "gpt-4", - "litellm_params": { - "model": "gpt-4", - "api_key": "very-fake-key", - }, - }, - { - "model_name": "fake-openai-endpoint-2", - "litellm_params": { - "model": "openai/fake-openai-endpoint-2", - "api_key": "working-key-since-this-is-fake-endpoint", - "api_base": FAKE_OPENAI_API_BASE, - }, - }, - ], - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("function_name", ["_acompletion", "_atext_completion"]) -async def test_multiple_fallbacks(function_name, respx_mock: respx.MockRouter, monkeypatch): - """ - Tests that if multiple fallbacks passed: - - fallback 1 = bad configured deployment / failing endpoint - - fallback 2 = working deployment / working endpoint - - Assert that: - - a success response is received from the working endpoint (fallback 2) - """ - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - invalid_key = httpx.Response( - 401, - json={ - "error": { - "message": "Incorrect API key provided", - "type": "invalid_request_error", - "code": "invalid_api_key", - } - }, - ) - respx_mock.post("https://api.openai.com/v1/chat/completions").mock(return_value=invalid_key) - respx_mock.post("https://api.openai.com/v1/completions").mock(return_value=invalid_key) - respx_mock.post(f"{FAKE_OPENAI_API_BASE}/chat/completions").mock( - return_value=httpx.Response( - 200, - json={ - "id": "chatcmpl-fake", - "object": "chat.completion", - "created": 1700000000, - "model": "fake-openai-endpoint-2", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, - }, - ) - ) - respx_mock.post(f"{FAKE_OPENAI_API_BASE}/completions").mock( - return_value=httpx.Response( - 200, - json={ - "id": "cmpl-fake", - "object": "text_completion", - "created": 1700000000, - "model": "fake-openai-endpoint-2", - "choices": [{"index": 0, "text": "hi", "logprobs": None, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, - }, - ) - ) - router_2 = create_test_router_2() - original_function = getattr(router_2, function_name) - - fallback_model_group = ["gpt-4", "fake-openai-endpoint-2"] - original_model_group = "gpt-3.5-turbo" - original_exception = Exception("Simulated error") - - request_kwargs: dict[str, object] = {"metadata": {"previous_models": ["gpt-3.5-turbo"]}} - - if function_name == "_aembedding": - request_kwargs["input"] = "hello this is a test for run_async_fallback" - elif function_name == "_atext_completion": - request_kwargs["prompt"] = "hello this is a test for run_async_fallback" - elif function_name == "_acompletion": - request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}] - - result = await run_async_fallback( - litellm_routing_ctx=RoutingContext(router_2), - original_function=original_function, - num_retries=1, - fallback_model_group=fallback_model_group, - original_model_group=original_model_group, - original_exception=original_exception, - max_fallbacks=5, - fallback_depth=0, - **request_kwargs, - ) - - print(result) - - print(result._hidden_params) - - assert result._hidden_params["api_base"] == FAKE_OPENAI_API_BASE - - def _provider_error(status_code: int) -> httpx.Response: return httpx.Response(status_code=status_code, request=httpx.Request(method="POST", url="https://provider.test")) diff --git a/tests/unit/router_utils/test_health_check_allowed_fails_integration.py b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py index 74b3a8825b3..a064c7346b4 100644 --- a/tests/unit/router_utils/test_health_check_allowed_fails_integration.py +++ b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py @@ -6,7 +6,8 @@ should increment the same counters used by allowed_fails_policy, using the actual exception type from the health check error. """ -from unittest.mock import ANY, patch +import time +from unittest.mock import patch import pytest @@ -14,9 +15,21 @@ import litellm from litellm.proxy.health_check import run_with_timeout from litellm.router import Router from litellm.router_api import RoutingContext +from litellm.router_utils.cooldown_handlers import get_cooldown_deployments +from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + get_deployment_failures_for_current_minute, +) from litellm.types.router import AllowedFailsPolicy +def _cooled(router: Router) -> list[str]: + return list(get_cooldown_deployments(RoutingContext(router), parent_otel_span=None)) + + +def _failures(router: Router, model_id: str) -> int: + return get_deployment_failures_for_current_minute(RoutingContext(router), model_id) + + def _make_model(model_id: str, model_name: str = "gpt-4") -> dict: return { "model_name": model_name, @@ -277,43 +290,29 @@ class TestHealthCheckCooldownIntegration: class TestWriteHealthStateIntegration: - """Test _write_health_state_to_router_cache integrates with cooldown pipeline.""" + """Test _write_health_state_to_router_cache feeds the router's cooldown pipeline.""" - def test_unhealthy_endpoint_triggers_set_cooldown(self): - """_write_health_state_to_router_cache should call set_cooldown_deployments for unhealthy endpoints.""" + def test_unhealthy_endpoint_cools_down_the_deployment(self): + """An unhealthy endpoint's exception counts against the policy and cools the deployment down.""" import litellm.proxy.proxy_server as proxy_module from litellm.proxy.proxy_server import _write_health_state_to_router_cache router = Router( - model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], - allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5), + model_list=[_make_model("deploy-1"), _make_model("deploy-2")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0), enable_health_check_routing=True, ) timeout_exc = litellm.Timeout(message="Health check timeout", model="", llm_provider="") - unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "timeout"}, - ] - healthy_endpoints = [ - {"model_id": "deploy-2"}, - ] - with patch.object(proxy_module, "llm_router", router): - with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown: - _write_health_state_to_router_cache( - healthy_endpoints=healthy_endpoints, - unhealthy_endpoints=unhealthy_endpoints, - exceptions_by_model_id={"deploy-1": timeout_exc}, - ) - mock_cooldown.assert_called_once_with( - ANY, - original_exception=timeout_exc, - exception_status=408, # Timeout has status_code 408 - deployment="deploy-1", - time_to_cooldown=router.settings().cooldown_time, - ) - assert mock_cooldown.call_args.args[0].belongs_to(router) + _write_health_state_to_router_cache( + healthy_endpoints=[{"model_id": "deploy-2"}], + unhealthy_endpoints=[{"model_id": "deploy-1", "error": "timeout"}], + exceptions_by_model_id={"deploy-1": timeout_exc}, + ) + + assert _cooled(router) == ["deploy-1"] def test_unhealthy_endpoint_without_exception_skips_cooldown(self): """Unhealthy endpoints without an exception key should not trigger cooldown.""" @@ -321,270 +320,72 @@ class TestWriteHealthStateIntegration: from litellm.proxy.proxy_server import _write_health_state_to_router_cache router = Router( - model_list=[_make_model("deploy-1")], - allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5), + model_list=[_make_model("deploy-1"), _make_model("deploy-2")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0), enable_health_check_routing=True, ) - unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "unknown failure"}, - ] - with patch.object(proxy_module, "llm_router", router): - with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown: - _write_health_state_to_router_cache( - healthy_endpoints=[], - unhealthy_endpoints=unhealthy_endpoints, - # no exceptions_by_model_id → cooldown should not fire - ) - mock_cooldown.assert_not_called() + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=[{"model_id": "deploy-1", "error": "unknown failure"}], + ) + + assert _failures(router, "deploy-1") == 0 + assert _cooled(router) == [] def test_unhealthy_endpoint_increments_failure_counter(self): - """Unhealthy endpoints should call increment_deployment_failures_for_current_minute.""" + """An unhealthy endpoint counts one failure for its deployment.""" import litellm.proxy.proxy_server as proxy_module from litellm.proxy.proxy_server import _write_health_state_to_router_cache router = Router( - model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + model_list=[_make_model("deploy-1"), _make_model("deploy-2")], allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=10), enable_health_check_routing=True, ) rate_exc = litellm.RateLimitError(message="Rate limited", model="gpt-4", llm_provider="openai") - unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "rate limited"}, - ] - with patch.object(proxy_module, "llm_router", router): - with patch( - "litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute" - ) as mock_increment: - with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments"): - _write_health_state_to_router_cache( - healthy_endpoints=[], - unhealthy_endpoints=unhealthy_endpoints, - exceptions_by_model_id={"deploy-1": rate_exc}, - ) - mock_increment.assert_called_once_with(ANY, "deploy-1") - assert mock_increment.call_args.args[0].belongs_to(router) - - -class TestHealthCheckFilterBypassWithPolicy: - """ - When allowed_fails_policy is set, the binary health check filter should be - bypassed so cooldown is the sole routing exclusion mechanism. - """ - - def test_filter_bypassed_when_policy_set(self): - """Binary health check filter is a no-op when allowed_fails_policy is configured.""" - import time - - from litellm.caching.caching import DualCache - from litellm.router_utils.health_state_cache import DeploymentHealthCache - - router = Router( - model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], - allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=3), - enable_health_check_routing=True, - ) - - # Mark deploy-1 as unhealthy in the health state cache - cache = DualCache() - health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) - health_cache.set_deployment_health_states( - { - "deploy-1": { - "is_healthy": False, - "timestamp": time.time(), - "reason": "test", - }, - } - ) - router.health_state_cache = health_cache - - deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] - - # Filter should pass all through because policy is set - result = router._filter_health_check_unhealthy_deployments(deployments) - assert len(result) == 2, "Binary filter should be bypassed when allowed_fails_policy is set" - - def test_filter_active_when_no_policy(self): - """Binary health check filter still works when no allowed_fails_policy is configured.""" - import time - - from litellm.caching.caching import DualCache - from litellm.router_utils.health_state_cache import DeploymentHealthCache - - router = Router( - model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], - enable_health_check_routing=True, - ) - - cache = DualCache() - health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) - health_cache.set_deployment_health_states( - { - "deploy-1": { - "is_healthy": False, - "timestamp": time.time(), - "reason": "test", - }, - } - ) - router.health_state_cache = health_cache - - deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] - - result = router._filter_health_check_unhealthy_deployments(deployments) - assert len(result) == 1 - assert result[0]["model_info"]["id"] == "deploy-2" - - @pytest.mark.asyncio - async def test_async_filter_bypassed_when_policy_set(self): - """Async version also bypasses when allowed_fails_policy is set.""" - import time - - from litellm.caching.caching import DualCache - from litellm.router_utils.health_state_cache import DeploymentHealthCache - - router = Router( - model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], - allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=2), - enable_health_check_routing=True, - ) - - cache = DualCache() - health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) - health_cache.set_deployment_health_states( - { - "deploy-1": { - "is_healthy": False, - "timestamp": time.time(), - "reason": "test", - }, - } - ) - router.health_state_cache = health_cache - - deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] - - result = await router._async_filter_health_check_unhealthy_deployments(deployments) - assert len(result) == 2 - - def _make_scoped_router_with_unhealthy(self, policy) -> Router: - import time - - from litellm.caching.caching import DualCache - from litellm.router_utils.health_state_cache import DeploymentHealthCache - - router = Router( - model_list=[ - _make_model("bad-listed"), - _make_model("ok-listed"), - _make_model("bad-unlisted", "gpt-5"), - ], - allowed_fails_policy=policy, - enable_health_check_routing=True, - background_health_check_model_groups=["gpt-4"], - ) - cache = DualCache() - health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) - health_cache.set_deployment_health_states( - { - model_id: { - "is_healthy": False, - "timestamp": time.time(), - "reason": "test", - } - for model_id in ("bad-listed", "bad-unlisted") - } - ) - router.health_state_cache = health_cache - return router - - def test_filter_with_policy_still_applies_to_listed_groups(self): - """A model-group allowlist keeps the filter active for listed groups even with a policy set.""" - router = self._make_scoped_router_with_unhealthy(AllowedFailsPolicy(AuthenticationErrorAllowedFails=3)) - deployments = [ - _make_model("bad-listed"), - _make_model("ok-listed"), - _make_model("bad-unlisted", "gpt-5"), - ] - - result = router._filter_health_check_unhealthy_deployments(deployments) - assert [d["model_info"]["id"] for d in result] == ["ok-listed", "bad-unlisted"] - - @pytest.mark.asyncio - async def test_async_filter_with_policy_still_applies_to_listed_groups(self): - """Async version: listed groups stay filtered with a policy set, unlisted stay untouched.""" - router = self._make_scoped_router_with_unhealthy(AllowedFailsPolicy(TimeoutErrorAllowedFails=2)) - deployments = [ - _make_model("bad-listed"), - _make_model("ok-listed"), - _make_model("bad-unlisted", "gpt-5"), - ] - - result = await router._async_filter_health_check_unhealthy_deployments(deployments) - assert [d["model_info"]["id"] for d in result] == ["ok-listed", "bad-unlisted"] - - -class TestAllDeploymentsInCooldownSafetyNet: - """ - When enable_health_check_routing=True and ALL deployments enter cooldown, - the async routing path should bypass the cooldown filter and return all - deployments rather than blocking all traffic. - """ - - def test_raw_cooldown_filter_returns_empty_when_all_cooled(self): - """The raw _filter_cooldown_deployments has no safety net -- it returns empty.""" - router = Router( - model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], - enable_health_check_routing=True, - ) - deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] - result = router._filter_cooldown_deployments( - healthy_deployments=deployments, - cooldown_deployments=["deploy-1", "deploy-2"], - ) - assert result == [] # raw filter has no safety net - - @pytest.mark.asyncio - async def test_async_routing_path_bypasses_all_cooldown(self): - """In the async routing path, all-in-cooldown with enable_health_check_routing - returns the full list instead of empty (safety net).""" - from unittest.mock import AsyncMock - - from litellm.router_utils.cooldown_handlers import ( - _async_get_cooldown_deployments, - ) - - router = Router( - model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], - allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=0), - enable_health_check_routing=True, - ) - - deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] - - # Simulate all deployments in cooldown - with patch( - "litellm.router.async_get_cooldown_deployments", - new=AsyncMock(return_value=["deploy-1", "deploy-2"]), - ): - # The safety net in async_get_available_deployment should restore - # all deployments when the cooldown filter empties the list - _pre = deployments.copy() - filtered = router._filter_cooldown_deployments( - healthy_deployments=deployments, - cooldown_deployments=["deploy-1", "deploy-2"], + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=[{"model_id": "deploy-1", "error": "rate limited"}], + exceptions_by_model_id={"deploy-1": rate_exc}, ) - # If filtered is empty and enable_health_check_routing is True, - # the routing path restores _pre_cooldown_deployments - if not filtered and router.enable_health_check_routing: - filtered = _pre - assert len(filtered) == 2, "Safety net should return all deployments when all are in cooldown" + assert _failures(router, "deploy-1") == 1 + assert _cooled(router) == [] + + +@pytest.mark.parametrize( + ("allowed_fails_policy", "model_groups", "routable"), + [ + (None, None, ["deploy-2"]), + (AllowedFailsPolicy(AuthenticationErrorAllowedFails=3), None, ["deploy-1", "deploy-2"]), + (AllowedFailsPolicy(AuthenticationErrorAllowedFails=3), ["gpt-4"], ["deploy-2"]), + ], + ids=["no-policy-filters", "policy-bypasses-filter", "policy-keeps-filter-for-listed-groups"], +) +@pytest.mark.asyncio +async def test_health_check_filter_is_bypassed_by_an_allowed_fails_policy( + allowed_fails_policy: AllowedFailsPolicy | None, model_groups: list[str] | None, routable: list[str] +) -> None: + """With allowed_fails_policy set, cooldown is the sole routing exclusion, except for groups the + health check is scoped to.""" + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2")], + allowed_fails_policy=allowed_fails_policy, + enable_health_check_routing=True, + background_health_check_model_groups=model_groups, + ) + router.health().health_checks.set_deployment_health_states( + {"deploy-1": {"is_healthy": False, "timestamp": time.time(), "reason": "test"}} + ) + + healthy = await router.select("gpt-4", {}, kind="healthy") + + assert sorted(deployment["model_info"]["id"] for deployment in healthy) == routable class TestHealthCheckIgnoreTransientErrors: @@ -614,17 +415,14 @@ class TestHealthCheckIgnoreTransientErrors: ] with patch.object(proxy_module, "llm_router", router): - with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown: - with patch( - "litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute" - ) as mock_increment: - _write_health_state_to_router_cache( - healthy_endpoints=[], - unhealthy_endpoints=unhealthy_endpoints, - exceptions_by_model_id={"deploy-1": rate_exc}, - ) - mock_cooldown.assert_not_called() - mock_increment.assert_not_called() + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + + assert _failures(router, "deploy-1") == 0 + assert _cooled(router) == [] def test_408_skipped_when_flag_enabled(self): """408 from health check does not trigger cooldown when flag is set.""" @@ -645,13 +443,13 @@ class TestHealthCheckIgnoreTransientErrors: ] with patch.object(proxy_module, "llm_router", router): - with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown: - _write_health_state_to_router_cache( - healthy_endpoints=[], - unhealthy_endpoints=unhealthy_endpoints, - exceptions_by_model_id={"deploy-1": timeout_exc}, - ) - mock_cooldown.assert_not_called() + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": timeout_exc}, + ) + + assert _cooled(router) == [] def test_401_still_triggers_cooldown_when_flag_enabled(self): """Auth errors (401) still trigger cooldown even when flag is set.""" @@ -672,13 +470,13 @@ class TestHealthCheckIgnoreTransientErrors: ] with patch.object(proxy_module, "llm_router", router): - with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown: - _write_health_state_to_router_cache( - healthy_endpoints=[], - unhealthy_endpoints=unhealthy_endpoints, - exceptions_by_model_id={"deploy-1": auth_exc}, - ) - mock_cooldown.assert_called_once() + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": auth_exc}, + ) + + assert _cooled(router) == ["deploy-1"] def test_429_not_written_to_health_state_cache_when_flag_enabled(self): """429 endpoint is excluded from health state cache when flag is set, @@ -729,13 +527,13 @@ class TestHealthCheckIgnoreTransientErrors: ] with patch.object(proxy_module, "llm_router", router): - with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown: - _write_health_state_to_router_cache( - healthy_endpoints=[], - unhealthy_endpoints=unhealthy_endpoints, - exceptions_by_model_id={"deploy-1": rate_exc}, - ) - mock_cooldown.assert_called_once() + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + + assert _cooled(router) == ["deploy-1"] class TestSharedCacheTransientErrorFilter: diff --git a/tests/unit/router_utils/test_kubernetes_pod_discovery.py b/tests/unit/router_utils/test_kubernetes_pod_discovery.py index 0b772b09f6f..88b840e5c2a 100644 --- a/tests/unit/router_utils/test_kubernetes_pod_discovery.py +++ b/tests/unit/router_utils/test_kubernetes_pod_discovery.py @@ -3,6 +3,7 @@ import copy import hashlib import json import logging +import os import re import socket import threading @@ -17,7 +18,6 @@ import respx from openai import AsyncOpenAI, OpenAI import litellm -import litellm.router_python_core from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, KUBERNETES_POD_ROUTING_KEY, @@ -277,12 +277,11 @@ def _empty_proxy_environment() -> Mapping[str, str]: return {} -def _patch_router_discovery( - monkeypatch: pytest.MonkeyPatch, - *discoveries: KubernetesPodDiscovery, -) -> None: - discovery_iter: Final = iter(discoveries) - monkeypatch.setattr(litellm.router_python_core, "KubernetesPodDiscovery", lambda: next(discovery_iter)) +def _clear_proxy_environment(monkeypatch: pytest.MonkeyPatch) -> None: + """Clear the proxy environment the router's own pod discovery reads, so it substitutes pod IPs.""" + for name in tuple(os.environ): + if name.lower().endswith("_proxy"): + monkeypatch.delenv(name) def _stub_sync_dns(monkeypatch: pytest.MonkeyPatch, *ips: str) -> None: @@ -1290,7 +1289,6 @@ def test_concurrent_sync_refresh_uses_cached_ips_or_service_url_while_in_flight( async def test_router_sends_pod_hosts_without_forwarding_discovery_flag( monkeypatch: pytest.MonkeyPatch, ) -> None: - clock: Final = _clock((0.0, 0.0, 0.0, 0.0, 0.0, 5.0, 5.0)) loop: Final = asyncio.get_running_loop() async_lookups: Final = count() sync_lookups: Final = count() @@ -1326,9 +1324,7 @@ async def test_router_sends_pod_hosts_without_forwarding_discovery_flag( r"http://(?:10\.0\.0\.1|10\.0\.0\.2|vllm-headless\.ns\.svc\.cluster\.local):8000/v1/chat/completions" ) ).mock(return_value=httpx.Response(200, json=_CHAT_RESPONSE)) - discovery_enabled: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment) - control: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery_enabled, control) + _clear_proxy_environment(monkeypatch) discovery_enabled_router: Final = Router(model_list=[_deployment()]) control_router: Final = Router(model_list=[_deployment(kubernetes_pod_discovery=None)]) cached_async_client: Final = AsyncOpenAI(api_key="fake", base_url=_SERVICE_URL) @@ -1364,8 +1360,8 @@ async def test_router_sends_pod_hosts_without_forwarding_discovery_flag( "vllm-headless.ns.svc.cluster.local", ) assert all("kubernetes_pod_discovery" not in body for body in bodies) - assert next(async_lookups) == 1 - assert next(sync_lookups) == 1 + # The router's discovery resolves the service once; the sync call reuses the cached pod set. + assert (next(async_lookups), next(sync_lookups)) == (1, 0) @pytest.mark.asyncio @@ -1405,9 +1401,7 @@ async def test_router_session_requests_stick_to_one_pod_for_sync_and_async( route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( return_value=httpx.Response(200, json=_CHAT_RESPONSE) ) - async_discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) - sync_discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, async_discovery, sync_discovery) + _clear_proxy_environment(monkeypatch) async_router: Final = Router(model_list=[_deployment()]) sync_router: Final = Router(model_list=[_deployment()]) async_client: Final = _cache_async_client(async_router) @@ -1451,8 +1445,7 @@ async def test_router_session_mapping_repeats_across_pods( route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( return_value=httpx.Response(200, json=_CHAT_RESPONSE) ) - discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery) + _clear_proxy_environment(monkeypatch) router: Final = Router(model_list=[_deployment()]) client: Final = _cache_async_client(router) for session_id in session_ids: @@ -1493,8 +1486,7 @@ async def test_generated_session_id_uses_round_robin_pod_selection( route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( return_value=httpx.Response(200, json=_CHAT_RESPONSE) ) - discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery) + _clear_proxy_environment(monkeypatch) router: Final = Router(model_list=[_deployment()]) client: Final = _cache_async_client(router) for _ in range(6): @@ -1533,8 +1525,7 @@ async def test_empty_session_id_uses_round_robin_pod_selection( route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( return_value=httpx.Response(200, json=_CHAT_RESPONSE) ) - discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery) + _clear_proxy_environment(monkeypatch) router: Final = Router(model_list=[_deployment()]) client: Final = _cache_async_client(router) for _ in range(6): @@ -1567,8 +1558,7 @@ async def test_session_requests_do_not_advance_round_robin_cursor( route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( return_value=httpx.Response(200, json=_CHAT_RESPONSE) ) - discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery) + _clear_proxy_environment(monkeypatch) router: Final = Router(model_list=[_deployment()]) client: Final = _cache_async_client(router) await _router_acompletion(router) @@ -1587,14 +1577,13 @@ async def test_session_requests_do_not_advance_round_robin_cursor( ) -@pytest.mark.asyncio -async def test_session_pod_membership_change_remaps_only_affected_sessions( +def test_session_pod_membership_change_remaps_only_affected_sessions( monkeypatch: pytest.MonkeyPatch, ) -> None: - loop: Final = asyncio.get_running_loop() lookups: Final = count() + now: Final = [0.0] - async def getaddrinfo( + def getaddrinfo( host: str | None, port: str | int | None, family: int = 0, @@ -1606,44 +1595,34 @@ async def test_session_pod_membership_change_remaps_only_affected_sessions( ips: Final = ("10.0.0.1", "10.0.0.2", "10.0.0.3") if lookup_number == 0 else ("10.0.0.2", "10.0.0.3") return _records(*ips, port=int(port or 0)) - monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo) - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(socket, "getaddrinfo", getaddrinfo) + discovery: Final = KubernetesPodDiscovery( + refresh_interval_seconds=10, + clock=lambda: now[0], + proxy_environment=_empty_proxy_environment, + ) + + def host_for(session_id: str) -> str | None: + return httpx.URL( + _api_base(discovery.resolve_deployment(_deployment(), {"metadata": {"session_id": session_id}})) + ).host session_ids: Final = tuple(f"session-{index}" for index in range(60)) - with respx.mock(assert_all_called=True) as respx_mock: - route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( - return_value=httpx.Response(200, json=_CHAT_RESPONSE) - ) - discovery: Final = KubernetesPodDiscovery( - refresh_interval_seconds=10, - clock=_clock((0.0,) * 62 + (10.0,) * 10), - proxy_environment=_empty_proxy_environment, - ) - _patch_router_discovery(monkeypatch, discovery) - router: Final = Router(model_list=[_deployment()]) - client: Final = _cache_async_client(router) - for session_id in session_ids: - await _router_acompletion(router, metadata={"session_id": session_id}) - initial_hosts: Final = tuple(call.request.url.host for call in route.calls) - removed_session_ids: Final = tuple( - session_id for session_id, host in zip(session_ids, initial_hosts) if host == "10.0.0.1" - ) - surviving_session_ids: Final = tuple( - session_id for session_id, host in zip(session_ids, initial_hosts) if host == "10.0.0.2" - ) - assert removed_session_ids - assert len(surviving_session_ids) >= 2 - await _router_acompletion(router) - await _router_acompletion(router, metadata={"session_id": surviving_session_ids[0]}) - await _router_acompletion(router, metadata={"session_id": removed_session_ids[0]}) - await _router_acompletion(router, metadata={"session_id": surviving_session_ids[1]}) - hosts: Final = tuple(call.request.url.host for call in route.calls) - await client.close() + initial_hosts: Final = {session_id: host_for(session_id) for session_id in session_ids} + removed_session_ids: Final = tuple( + session_id for session_id in session_ids if initial_hosts[session_id] == "10.0.0.1" + ) + surviving_session_ids: Final = tuple( + session_id for session_id in session_ids if initial_hosts[session_id] in ("10.0.0.2", "10.0.0.3") + ) + assert removed_session_ids + assert surviving_session_ids - assert hosts[60] == "10.0.0.1" - assert hosts[61] == "10.0.0.2" - assert hosts[62] in ("10.0.0.2", "10.0.0.3") - assert hosts[63] == "10.0.0.2" + now[0] = 10.0 + remapped_hosts: Final = {session_id: host_for(session_id) for session_id in session_ids} + + assert all(remapped_hosts[session_id] == initial_hosts[session_id] for session_id in surviving_session_ids) + assert all(remapped_hosts[session_id] in ("10.0.0.2", "10.0.0.3") for session_id in removed_session_ids) assert next(lookups) == 2 @@ -1684,15 +1663,7 @@ async def test_custom_routing_strategy_keeps_sync_and_async_sessions_on_one_pod( route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( return_value=httpx.Response(200, json=_CHAT_RESPONSE) ) - async_discovery: Final = KubernetesPodDiscovery( - clock=lambda: 0.0, - proxy_environment=_empty_proxy_environment, - ) - sync_discovery: Final = KubernetesPodDiscovery( - clock=lambda: 0.0, - proxy_environment=_empty_proxy_environment, - ) - _patch_router_discovery(monkeypatch, async_discovery, sync_discovery) + _clear_proxy_environment(monkeypatch) async_router: Final = Router(model_list=[_deployment()]) sync_router: Final = Router(model_list=[_deployment()]) async_router.set_custom_routing_strategy(_ModelListRoutingStrategy(async_router)) @@ -1735,8 +1706,7 @@ async def test_top_level_litellm_session_id_keeps_requests_on_one_pod( route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( return_value=httpx.Response(200, json=_CHAT_RESPONSE) ) - discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery) + _clear_proxy_environment(monkeypatch) router: Final = Router(model_list=[_deployment()]) client: Final = _cache_async_client(router) for _ in range(6): @@ -1849,8 +1819,7 @@ async def test_generated_metadata_ignores_top_level_litellm_session_id( route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock( return_value=httpx.Response(200, json=_CHAT_RESPONSE) ) - discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery) + _clear_proxy_environment(monkeypatch) router: Final = Router(model_list=[_deployment()]) client: Final = _cache_async_client(router) for _ in range(6): @@ -1872,7 +1841,6 @@ async def test_generated_metadata_ignores_top_level_litellm_session_id( async def test_custom_routing_strategy_resolves_pod_hosts_for_sync_and_async_requests( monkeypatch: pytest.MonkeyPatch, ) -> None: - clock: Final = _clock((0.0,) * 40) loop: Final = asyncio.get_running_loop() async_lookups: Final = count() sync_lookups: Final = count() @@ -1908,15 +1876,7 @@ async def test_custom_routing_strategy_resolves_pod_hosts_for_sync_and_async_req r"http://(?:10\.0\.0\.1|10\.0\.0\.2|vllm-headless\.ns\.svc\.cluster\.local):8000/v1/chat/completions" ) ).mock(return_value=httpx.Response(200, json=_CHAT_RESPONSE)) - async_discovery: Final = KubernetesPodDiscovery( - clock=clock, - proxy_environment=_empty_proxy_environment, - ) - sync_discovery: Final = KubernetesPodDiscovery( - clock=clock, - proxy_environment=_empty_proxy_environment, - ) - _patch_router_discovery(monkeypatch, async_discovery, sync_discovery) + _clear_proxy_environment(monkeypatch) async_router: Final = Router(model_list=[_deployment()]) sync_router: Final = Router(model_list=[_deployment()]) async_router.set_custom_routing_strategy(_ModelListRoutingStrategy(async_router)) @@ -1932,23 +1892,8 @@ async def test_custom_routing_strategy_resolves_pod_hosts_for_sync_and_async_req assert next(sync_lookups) == 1 -@pytest.mark.asyncio -async def test_unresolved_router_selectors_keep_service_hostname() -> None: - router: Final = Router(model_list=[_deployment()]) - - sync_selected: Final = router._get_available_deployment_unresolved(model="gpu-model", request_kwargs={}) - async_selected: Final = await router._async_get_available_deployment_unresolved( - model="gpu-model", - request_kwargs={}, - ) - - assert _api_base(sync_selected) == _SERVICE_URL - assert _api_base(async_selected) == _SERVICE_URL - - @pytest.mark.asyncio async def test_router_async_retry_uses_next_discovered_pod(monkeypatch: pytest.MonkeyPatch) -> None: - clock: Final = _clock((0.0,) * 8) loop: Final = asyncio.get_running_loop() attempts: Final = count() @@ -1974,8 +1919,7 @@ async def test_router_async_retry_uses_next_discovered_pod(monkeypatch: pytest.M route: Final = respx_mock.post( url__regex=re.compile(r"http://(?:10\.0\.0\.1|10\.0\.0\.2):8000/v1/chat/completions") ).mock(side_effect=response_for) - discovery: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery) + _clear_proxy_environment(monkeypatch) router: Final = Router(model_list=[_deployment()], num_retries=1) response: Final = await router.acompletion( @@ -1989,7 +1933,6 @@ async def test_router_async_retry_uses_next_discovered_pod(monkeypatch: pytest.M def test_router_sync_retry_uses_next_discovered_pod(monkeypatch: pytest.MonkeyPatch) -> None: - clock: Final = _clock((0.0,) * 8) attempts: Final = count() def getaddrinfo( @@ -2014,8 +1957,7 @@ def test_router_sync_retry_uses_next_discovered_pod(monkeypatch: pytest.MonkeyPa route: Final = respx_mock.post( url__regex=re.compile(r"http://(?:10\.0\.0\.1|10\.0\.0\.2):8000/v1/chat/completions") ).mock(side_effect=response_for) - discovery: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery) + _clear_proxy_environment(monkeypatch) router: Final = Router(model_list=[_deployment()], num_retries=1) response: Final = router.completion( @@ -2088,8 +2030,7 @@ async def test_router_session_retry_uses_next_rendezvous_pod(monkeypatch: pytest session_id: Final = "retry-session" with respx.mock(assert_all_called=True) as respx_mock: route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(side_effect=response_for) - discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment) - _patch_router_discovery(monkeypatch, discovery) + _clear_proxy_environment(monkeypatch) router: Final = Router(model_list=[_deployment()], num_retries=1) response: Final = await router.acompletion( diff --git a/tests/unit/router_utils/test_router_health_check_routing.py b/tests/unit/router_utils/test_router_health_check_routing.py index b1498ea3687..79ca88120e5 100644 --- a/tests/unit/router_utils/test_router_health_check_routing.py +++ b/tests/unit/router_utils/test_router_health_check_routing.py @@ -6,236 +6,70 @@ import time import pytest -from litellm.caching.caching import DualCache -from litellm.router_utils.health_state_cache import DeploymentHealthCache +from litellm.router import Router def _make_deployment(model_id: str, model_name: str = "gpt-4") -> dict: """Helper to create a deployment dict for testing.""" return { "model_name": model_name, - "litellm_params": {"model": model_name, "api_key": "fake"}, + "litellm_params": {"model": "openai/gpt-4", "api_key": "fake"}, "model_info": {"id": model_id}, } -def _make_health_cache( - unhealthy_ids: set = None, staleness_threshold: float = 60.0 -) -> DeploymentHealthCache: - """Create a health cache pre-populated with unhealthy deployment IDs.""" - cache = DualCache() - health_cache = DeploymentHealthCache( - cache=cache, staleness_threshold=staleness_threshold +def _health_routed_router(enable: bool, unhealthy_ids: set[str], model_groups: list[str] | None = None) -> Router: + router = Router( + model_list=[ + _make_deployment("deploy-1", model_name="prod"), + _make_deployment("deploy-2", model_name="prod"), + _make_deployment("deploy-3", model_name="prod"), + _make_deployment("other-1", model_name="other"), + _make_deployment("other-2", model_name="other"), + ], + enable_health_check_routing=enable, + background_health_check_model_groups=model_groups, ) - if unhealthy_ids: - now = time.time() - states = {} - for uid in unhealthy_ids: - states[uid] = { - "is_healthy": False, - "timestamp": now, - "reason": "test_unhealthy", - } - health_cache.set_deployment_health_states(states) - return health_cache + now = time.time() + router.health().health_checks.set_deployment_health_states( + {uid: {"is_healthy": False, "timestamp": now, "reason": "test_unhealthy"} for uid in unhealthy_ids} + ) + return router -class TestFilterHealthCheckUnhealthyDeployments: - """Test the sync filter method.""" +@pytest.mark.parametrize( + ("enable", "unhealthy_ids", "model_groups", "routable"), + [ + (True, {"deploy-2"}, None, {"prod": ["deploy-1", "deploy-3"], "other": ["other-1", "other-2"]}), + (False, {"deploy-2"}, None, {"prod": ["deploy-1", "deploy-2", "deploy-3"], "other": ["other-1", "other-2"]}), + (True, set(), None, {"prod": ["deploy-1", "deploy-2", "deploy-3"], "other": ["other-1", "other-2"]}), + ( + True, + {"deploy-1", "deploy-2", "deploy-3"}, + None, + {"prod": ["deploy-1", "deploy-2", "deploy-3"], "other": ["other-1", "other-2"]}, + ), + (True, {"deploy-2", "other-1"}, ["prod"], {"prod": ["deploy-1", "deploy-3"], "other": ["other-1", "other-2"]}), + (True, {"deploy-2", "other-1"}, None, {"prod": ["deploy-1", "deploy-3"], "other": ["other-2"]}), + ], + ids=[ + "removes-unhealthy", + "noop-when-disabled", + "all-pass-when-cache-empty", + "safety-net-when-all-unhealthy", + "scoped-to-listed-model-groups", + "unscoped-when-model-groups-unset", + ], +) +@pytest.mark.asyncio +async def test_health_check_routing_filters_unhealthy_deployments( + enable: bool, unhealthy_ids: set[str], model_groups: list[str] | None, routable: dict[str, list[str]] +) -> None: + router = _health_routed_router(enable, unhealthy_ids, model_groups) - def _make_router_like( - self, - enable: bool, - health_cache: DeploymentHealthCache, - model_groups: frozenset[str] | None = None, - ): - """Create a minimal object that behaves like Router for filter testing.""" - - class FakeRouter: - def __init__(self): - self._core_enable_health_check_routing = enable - self._core_health_state_cache = health_cache - self._core_allowed_fails_policy = None - self._core_background_health_check_model_groups = model_groups - - # Import the actual method and bind it - from litellm.router import Router - - fake = FakeRouter() - # Use the unbound method - fake._filter_health_check_unhealthy_deployments = ( - Router._filter_health_check_unhealthy_deployments.__get__(fake, FakeRouter) - ) - return fake - - def test_filter_removes_unhealthy_deployments(self): - """Unhealthy deployments should be removed from candidates.""" - health_cache = _make_health_cache(unhealthy_ids={"deploy-2"}) - router = self._make_router_like(enable=True, health_cache=health_cache) - - deployments = [ - _make_deployment("deploy-1"), - _make_deployment("deploy-2"), - _make_deployment("deploy-3"), - ] - result = router._filter_health_check_unhealthy_deployments(deployments) - assert len(result) == 2 - assert all(d["model_info"]["id"] != "deploy-2" for d in result) - - def test_filter_noop_when_disabled(self): - """When enable_health_check_routing=False, filter should be a no-op.""" - health_cache = _make_health_cache(unhealthy_ids={"deploy-1"}) - router = self._make_router_like(enable=False, health_cache=health_cache) - - deployments = [ - _make_deployment("deploy-1"), - _make_deployment("deploy-2"), - ] - result = router._filter_health_check_unhealthy_deployments(deployments) - assert len(result) == 2 # no filtering - - def test_filter_returns_all_when_all_unhealthy(self): - """Safety net: if ALL deployments are unhealthy, return all (don't cause outage).""" - health_cache = _make_health_cache( - unhealthy_ids={"deploy-1", "deploy-2", "deploy-3"} - ) - router = self._make_router_like(enable=True, health_cache=health_cache) - - deployments = [ - _make_deployment("deploy-1"), - _make_deployment("deploy-2"), - _make_deployment("deploy-3"), - ] - result = router._filter_health_check_unhealthy_deployments(deployments) - assert len(result) == 3 # all returned, safety net - - def test_filter_returns_all_when_cache_empty(self): - """When cache is empty, all deployments should pass through.""" - health_cache = _make_health_cache() # empty - router = self._make_router_like(enable=True, health_cache=health_cache) - - deployments = [ - _make_deployment("deploy-1"), - _make_deployment("deploy-2"), - ] - result = router._filter_health_check_unhealthy_deployments(deployments) - assert len(result) == 2 - - def test_filter_scoped_to_listed_model_groups(self): - """With an allowlist, only deployments in listed groups are filtered on health.""" - health_cache = _make_health_cache(unhealthy_ids={"bad-listed", "bad-unlisted"}) - router = self._make_router_like( - enable=True, health_cache=health_cache, model_groups=frozenset({"prod"}) - ) - - deployments = [ - _make_deployment("bad-listed", model_name="prod"), - _make_deployment("ok-listed", model_name="prod"), - _make_deployment("bad-unlisted", model_name="other"), - _make_deployment("ok-unlisted", model_name="other"), - ] - result = router._filter_health_check_unhealthy_deployments(deployments) - assert [d["model_info"]["id"] for d in result] == [ - "ok-listed", - "bad-unlisted", - "ok-unlisted", - ] - - def test_filter_unscoped_when_model_groups_unset(self): - """Without an allowlist, unhealthy deployments in every group are filtered.""" - health_cache = _make_health_cache(unhealthy_ids={"bad-listed", "bad-unlisted"}) - router = self._make_router_like(enable=True, health_cache=health_cache) - - deployments = [ - _make_deployment("bad-listed", model_name="prod"), - _make_deployment("ok-listed", model_name="prod"), - _make_deployment("bad-unlisted", model_name="other"), - _make_deployment("ok-unlisted", model_name="other"), - ] - result = router._filter_health_check_unhealthy_deployments(deployments) - assert [d["model_info"]["id"] for d in result] == ["ok-listed", "ok-unlisted"] - - -class TestAsyncFilterHealthCheckUnhealthyDeployments: - """Test the async filter method.""" - - def _make_router_like( - self, - enable: bool, - health_cache: DeploymentHealthCache, - model_groups: frozenset[str] | None = None, - ): - from litellm.router import Router - - class FakeRouter: - def __init__(self): - self._core_enable_health_check_routing = enable - self._core_health_state_cache = health_cache - self._core_allowed_fails_policy = None - self._core_background_health_check_model_groups = model_groups - - fake = FakeRouter() - fake._async_filter_health_check_unhealthy_deployments = ( - Router._async_filter_health_check_unhealthy_deployments.__get__( - fake, FakeRouter - ) - ) - return fake - - @pytest.mark.asyncio - async def test_async_filter_removes_unhealthy(self): - """Async version: unhealthy deployments removed.""" - health_cache = _make_health_cache(unhealthy_ids={"deploy-2"}) - router = self._make_router_like(enable=True, health_cache=health_cache) - - deployments = [ - _make_deployment("deploy-1"), - _make_deployment("deploy-2"), - _make_deployment("deploy-3"), - ] - result = await router._async_filter_health_check_unhealthy_deployments( - healthy_deployments=deployments - ) - assert len(result) == 2 - assert all(d["model_info"]["id"] != "deploy-2" for d in result) - - @pytest.mark.asyncio - async def test_async_filter_safety_net(self): - """Async version: safety net when all unhealthy.""" - health_cache = _make_health_cache(unhealthy_ids={"deploy-1", "deploy-2"}) - router = self._make_router_like(enable=True, health_cache=health_cache) - - deployments = [ - _make_deployment("deploy-1"), - _make_deployment("deploy-2"), - ] - result = await router._async_filter_health_check_unhealthy_deployments( - healthy_deployments=deployments - ) - assert len(result) == 2 # safety net - - @pytest.mark.asyncio - async def test_async_filter_scoped_to_listed_model_groups(self): - """Async version: only deployments in listed groups are filtered on health.""" - health_cache = _make_health_cache(unhealthy_ids={"bad-listed", "bad-unlisted"}) - router = self._make_router_like( - enable=True, health_cache=health_cache, model_groups=frozenset({"prod"}) - ) - - deployments = [ - _make_deployment("bad-listed", model_name="prod"), - _make_deployment("ok-listed", model_name="prod"), - _make_deployment("bad-unlisted", model_name="other"), - _make_deployment("ok-unlisted", model_name="other"), - ] - result = await router._async_filter_health_check_unhealthy_deployments( - healthy_deployments=deployments - ) - assert [d["model_info"]["id"] for d in result] == [ - "ok-listed", - "bad-unlisted", - "ok-unlisted", - ] + for model_group, expected in routable.items(): + healthy = await router.select(model_group, {}, kind="healthy") + assert sorted(d["model_info"]["id"] for d in healthy) == expected, model_group class TestBuildDeploymentHealthStates: diff --git a/tests/unit/router_utils/test_router_interactions_endpoints.py b/tests/unit/router_utils/test_router_interactions_endpoints.py index 91bea170458..bd83990fd1c 100644 --- a/tests/unit/router_utils/test_router_interactions_endpoints.py +++ b/tests/unit/router_utils/test_router_interactions_endpoints.py @@ -2,10 +2,11 @@ Tests for Router interactions API endpoint initialization functions. """ -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock import pytest +import litellm from litellm import Router @@ -39,191 +40,6 @@ class TestInitializeInteractionsEndpoints: assert callable(router.aget_interaction) assert callable(router.get_interaction) - def test_initialize_interactions_endpoints_can_be_called_directly(self): - """Test that _initialize_interactions_endpoints can be called directly to reinitialize endpoints.""" - router = Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4"}, - } - ] - ) - - # Call _initialize_interactions_endpoints directly - router._initialize_interactions_endpoints() - - # Verify the interaction methods still exist after re-initialization - assert hasattr(router, "acreate_interaction") - assert hasattr(router, "create_interaction") - assert callable(router.acreate_interaction) - - -class TestInitInteractionsApiEndpoints: - """Test cases for _init_interactions_api_endpoints method""" - - @pytest.mark.asyncio - async def test_init_interactions_api_endpoints_passes_custom_llm_provider(self): - """Test that _init_interactions_api_endpoints passes custom_llm_provider to the original function.""" - router = Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4"}, - } - ] - ) - - mock_function = AsyncMock(return_value={"result": "success"}) - - result = await router._init_interactions_api_endpoints( - original_function=mock_function, - custom_llm_provider="gemini", - interaction_id="test-id", - ) - - mock_function.assert_called_once_with( - custom_llm_provider="gemini", - interaction_id="test-id", - ) - assert result == {"result": "success"} - - @pytest.mark.asyncio - async def test_init_interactions_api_endpoints_defaults_to_gemini(self): - """Test that _init_interactions_api_endpoints defaults to gemini when no custom_llm_provider is specified.""" - router = Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4"}, - } - ] - ) - - mock_function = AsyncMock(return_value={"result": "success"}) - - result = await router._init_interactions_api_endpoints( - original_function=mock_function, - interaction_id="test-id", - ) - - mock_function.assert_called_once_with( - custom_llm_provider="gemini", - interaction_id="test-id", - ) - assert result == {"result": "success"} - - @pytest.mark.asyncio - async def test_init_interactions_api_endpoints_does_not_override_existing_provider( - self, - ): - """Test that _init_interactions_api_endpoints does not override custom_llm_provider if already in kwargs.""" - router = Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4"}, - } - ] - ) - - mock_function = AsyncMock(return_value={"result": "success"}) - - # Pass custom_llm_provider in kwargs directly (not as separate param) - result = await router._init_interactions_api_endpoints( - original_function=mock_function, - custom_llm_provider="vertex_ai", - ) - - # Should use the provided custom_llm_provider - mock_function.assert_called_once_with( - custom_llm_provider="vertex_ai", - ) - assert result == {"result": "success"} - - @pytest.mark.asyncio - async def test_init_interactions_api_endpoints_clears_model_when_equals_agent( - self, - ): - """Managed agent interactions must not pass agent name as model to the SDK.""" - router = Router(model_list=[]) - - mock_function = AsyncMock(return_value={"result": "success"}) - - await router._init_interactions_api_endpoints( - original_function=mock_function, - agent="mqy-custom-slides-agent", - model="mqy-custom-slides-agent", - input="hello", - ) - - mock_function.assert_called_once_with( - custom_llm_provider="gemini", - agent="mqy-custom-slides-agent", - model=None, - input="hello", - ) - - -class TestRouterCreateInteractionRouting: - """acreate_interaction routing: agent-only vs model + fallbacks.""" - - @pytest.mark.asyncio - async def test_acreate_interaction_agent_only_uses_init_interactions(self): - """Agent-only create must not use model-group fallback lookup.""" - router = Router( - model_list=[ - { - "model_name": "openai/*", - "litellm_params": {"model": "gpt-4"}, - } - ] - ) - - with ( - patch.object( - router, - "_init_interactions_api_endpoints", - new_callable=AsyncMock, - return_value={"id": "int-1"}, - ) as mock_init, - patch.object( - router, - "_ageneric_api_call_with_fallbacks", - new_callable=AsyncMock, - ) as mock_generic, - ): - result = await router.acreate_interaction( - agent="mqy-custom-slides-agent", - input="hello", - custom_llm_provider="gemini", - ) - - mock_init.assert_called_once() - mock_generic.assert_not_called() - assert result == {"id": "int-1"} - - @pytest.mark.asyncio - async def test_init_interactions_model_uses_generic_fallbacks(self): - """Model-based create uses _ageneric_api_call_with_fallbacks inside _init_interactions.""" - router = Router(model_list=[]) - - with patch.object( - router, - "_ageneric_api_call_with_fallbacks", - new_callable=AsyncMock, - return_value={"id": "int-1"}, - ) as mock_generic: - result = await router._init_interactions_api_endpoints( - original_function=AsyncMock(), - model="gemini-2.5-flash", - input="hello", - custom_llm_provider="gemini", - ) - - mock_generic.assert_called_once() - assert result == {"id": "int-1"} - class TestInitializeManagedAgentsEndpoints: """Tests for _initialize_managed_agents_endpoints.""" @@ -248,51 +64,60 @@ class TestInitializeManagedAgentsEndpoints: assert hasattr(router, method_name), f"missing {method_name}" assert callable(getattr(router, method_name)), f"{method_name} not callable" - def test_initialize_managed_agents_endpoints_can_be_called_directly(self): - router = Router(model_list=[]) - router._initialize_managed_agents_endpoints() - assert callable(router.acreate_agent) - assert callable(router.alist_agents) + +@pytest.mark.parametrize( + ("request_kwargs", "provider"), + [({}, "gemini"), ({"custom_llm_provider": "vertex_ai"}, "vertex_ai")], + ids=["defaults-to-gemini", "keeps-the-caller-provider"], +) +@pytest.mark.asyncio +async def test_interaction_endpoints_default_the_provider_to_gemini( + monkeypatch: pytest.MonkeyPatch, request_kwargs: dict[str, str], provider: str +) -> None: + provider_call = AsyncMock(return_value={"result": "success"}) + monkeypatch.setattr(litellm.interactions, "aget", provider_call) + router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}]) + + result = await router.aget_interaction(interaction_id="test-id", **request_kwargs) + + assert result == {"result": "success"} + assert provider_call.call_args.kwargs["custom_llm_provider"] == provider + assert provider_call.call_args.kwargs["interaction_id"] == "test-id" -class TestInitManagedAgentsApiEndpoints: - """Tests for _init_managed_agents_api_endpoints.""" +@pytest.mark.asyncio +async def test_agent_only_interaction_does_not_send_the_agent_name_as_the_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Managed agent interactions must not pass agent name as model to the SDK.""" + provider_call = AsyncMock(return_value={"id": "int-1"}) + monkeypatch.setattr(litellm.interactions, "acreate", provider_call) + router = Router(model_list=[{"model_name": "openai/*", "litellm_params": {"model": "gpt-4"}}]) - @pytest.mark.asyncio - async def test_init_managed_agents_api_endpoints_defaults_to_gemini(self): - router = Router(model_list=[]) - mock_fn = AsyncMock(return_value={"agents": []}) + result = await router.acreate_interaction( + agent="mqy-custom-slides-agent", model="mqy-custom-slides-agent", input="hello" + ) - await router._init_managed_agents_api_endpoints( - original_function=mock_fn, - ) + assert result == {"id": "int-1"} + assert provider_call.call_args.kwargs["agent"] == "mqy-custom-slides-agent" + assert provider_call.call_args.kwargs.get("model") is None + assert provider_call.call_args.kwargs["custom_llm_provider"] == "gemini" - call_kwargs = mock_fn.call_args.kwargs - assert call_kwargs["custom_llm_provider"] == "gemini" - @pytest.mark.asyncio - async def test_init_managed_agents_api_endpoints_passes_custom_provider(self): - router = Router(model_list=[]) - mock_fn = AsyncMock(return_value={"agents": []}) +@pytest.mark.parametrize( + ("request_kwargs", "provider"), + [({}, "gemini"), ({"custom_llm_provider": "vertex_ai"}, "vertex_ai")], + ids=["defaults-to-gemini", "keeps-the-caller-provider"], +) +@pytest.mark.asyncio +async def test_managed_agent_endpoints_default_the_provider_to_gemini( + monkeypatch: pytest.MonkeyPatch, request_kwargs: dict[str, str], provider: str +) -> None: + provider_call = AsyncMock(return_value={"agents": []}) + monkeypatch.setattr(litellm.interactions.agents, "alist", provider_call) + router = Router(model_list=[]) - await router._init_managed_agents_api_endpoints( - original_function=mock_fn, - custom_llm_provider="vertex_ai", - ) + result = await router.alist_agents(**request_kwargs) - call_kwargs = mock_fn.call_args.kwargs - assert call_kwargs["custom_llm_provider"] == "vertex_ai" - - @pytest.mark.asyncio - async def test_init_managed_agents_api_endpoints_does_not_override_existing_provider( - self, - ): - router = Router(model_list=[]) - mock_fn = AsyncMock(return_value={"agents": []}) - - await router._init_managed_agents_api_endpoints( - original_function=mock_fn, - custom_llm_provider="vertex_ai", - ) - - mock_fn.assert_called_once_with(custom_llm_provider="vertex_ai") + assert result == {"agents": []} + assert provider_call.call_args.kwargs["custom_llm_provider"] == provider diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 6ec72868b6e..0de10a64f2e 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -54,16 +54,12 @@ class TestFilterTeamBasedModels: {"model_info": {"id": "deployment-2"}}, ] - def test_filter_team_based_models_none_request_kwargs( - self, sample_deployments_with_teams - ): + def test_filter_team_based_models_none_request_kwargs(self, sample_deployments_with_teams): """Test that when request_kwargs is None, all deployments are returned unchanged""" result = filter_team_based_models(sample_deployments_with_teams, None) assert result == sample_deployments_with_teams - def test_filter_team_based_models_empty_request_kwargs( - self, sample_deployments_with_teams - ): + def test_filter_team_based_models_empty_request_kwargs(self, sample_deployments_with_teams): """Test with empty request_kwargs""" result = filter_team_based_models(sample_deployments_with_teams, {}) # Should include all deployments since no team_id in request @@ -76,9 +72,7 @@ class TestFilterTeamBasedModels: # Should include only non-team based deployments assert len(result) == 1 - def test_filter_team_based_models_team_match_metadata( - self, sample_deployments_with_teams - ): + def test_filter_team_based_models_team_match_metadata(self, sample_deployments_with_teams): """Test filtering when team_id is in metadata""" request_kwargs = {"metadata": {"user_api_key_team_id": "team-a"}} result = filter_team_based_models(sample_deployments_with_teams, request_kwargs) @@ -93,9 +87,7 @@ class TestFilterTeamBasedModels: result_ids = [d.get("model_info", {}).get("id") for d in result] assert sorted(result_ids) == sorted(expected_ids) - def test_filter_team_based_models_team_match_litellm_metadata( - self, sample_deployments_with_teams - ): + def test_filter_team_based_models_team_match_litellm_metadata(self, sample_deployments_with_teams): """Test filtering when team_id is in litellm_metadata""" request_kwargs = {"litellm_metadata": {"user_api_key_team_id": "team-b"}} result = filter_team_based_models(sample_deployments_with_teams, request_kwargs) @@ -110,9 +102,7 @@ class TestFilterTeamBasedModels: result_ids = [d.get("model_info", {}).get("id") for d in result] assert sorted(result_ids) == sorted(expected_ids) - def test_filter_team_based_models_priority_metadata_over_litellm( - self, sample_deployments_with_teams - ): + def test_filter_team_based_models_priority_metadata_over_litellm(self, sample_deployments_with_teams): """Test that metadata.user_api_key_team_id takes priority over litellm_metadata.user_api_key_team_id""" request_kwargs = { "metadata": { @@ -127,9 +117,7 @@ class TestFilterTeamBasedModels: result_ids = [d.get("model_info", {}).get("id") for d in result] assert sorted(result_ids) == sorted(expected_ids) - def test_filter_team_based_models_no_matching_team( - self, sample_deployments_with_teams - ): + def test_filter_team_based_models_no_matching_team(self, sample_deployments_with_teams): """Test when request team doesn't match any deployment teams""" request_kwargs = {"metadata": {"user_api_key_team_id": "team-nonexistent"}} result = filter_team_based_models(sample_deployments_with_teams, request_kwargs) @@ -139,9 +127,7 @@ class TestFilterTeamBasedModels: result_ids = [d.get("model_info", {}).get("id") for d in result] assert result_ids == expected_ids - def test_filter_team_based_models_no_team_restrictions( - self, sample_deployments_no_teams - ): + def test_filter_team_based_models_no_team_restrictions(self, sample_deployments_no_teams): """Test with deployments that have no team restrictions""" request_kwargs = {"metadata": {"user_api_key_team_id": "any-team"}} result = filter_team_based_models(sample_deployments_no_teams, request_kwargs) @@ -188,9 +174,7 @@ class TestFilterTeamBasedModels: def test_filter_team_based_models_empty_deployments(self): """Test with empty deployments list""" - result = filter_team_based_models( - [], {"metadata": {"user_api_key_team_id": "team-a"}} - ) + result = filter_team_based_models([], {"metadata": {"user_api_key_team_id": "team-a"}}) assert result == [] def test_filter_team_based_models_none_team_id_in_deployment(self): @@ -347,9 +331,7 @@ class TestFilterWebSearchDeployments: deployments = [ {"model_info": {"id": "d1"}}, # No supports_web_search - defaults to True {"model_info": {"id": "d2"}}, # No supports_web_search - defaults to True - { - "model_info": {"id": "d3", "supports_web_search": False} - }, # Explicit False + {"model_info": {"id": "d3", "supports_web_search": False}}, # Explicit False ] request_kwargs = {"tools": [{"type": "web_search"}]} result = filter_web_search_deployments(deployments, request_kwargs) @@ -375,26 +357,8 @@ class TestFilterWebSearchDeployments: assert result == deployment -def test_invalidate_model_group_info_cache(): - """Test that _invalidate_model_group_info_cache clears the LRU cache.""" - router = Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ] - ) - # Populate the cache - router._cached_get_model_group_info("gpt-4") - assert router._cached_get_model_group_info.cache_info().currsize > 0 - - # Invalidate and verify cache is cleared - router._invalidate_model_group_info_cache() - assert router._cached_get_model_group_info.cache_info().currsize == 0 - - -def test_filter_deployments_by_model_access_groups_access_group_only_key(): +@pytest.mark.asyncio +async def test_filter_deployments_by_model_access_groups_access_group_only_key(): """ Access-group-only keys should only route to deployments in allowed groups, even when multiple deployments share the same public model name. @@ -403,13 +367,13 @@ def test_filter_deployments_by_model_access_groups_access_group_only_key(): model_list=[ { "model_name": "gpt-5", - "litellm_params": {"model": "openai/gpt-5.1", "api_key": "key-1"}, - "model_info": {"access_groups": ["AG1"]}, + "litellm_params": {"model": "openai/gpt-5.1", "api_key": "key-1", "mock_response": "ag1"}, + "model_info": {"id": "ag1", "access_groups": ["AG1"]}, }, { "model_name": "gpt-5", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "key-2"}, - "model_info": {"access_groups": ["AG2"]}, + "litellm_params": {"model": "openai/gpt-4o", "api_key": "key-2", "mock_response": "ag2"}, + "model_info": {"id": "ag2", "access_groups": ["AG2"]}, }, ] ) @@ -420,21 +384,20 @@ def test_filter_deployments_by_model_access_groups_access_group_only_key(): models=["AG2"], team_models=["AG2"], ) + metadata = {"user_api_key_team_id": "team-2", "user_api_key_auth": scoped_key} - filtered = router._filter_deployments_by_model_access_groups( - model="gpt-5", - healthy_deployments=router.get_all_deployments(model_name="gpt-5"), - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team-2", - "user_api_key_auth": scoped_key, - } - }, - request_team_id="team-2", - ) + healthy = await router.select("gpt-5", {"metadata": metadata}, kind="healthy") + served = { + ( + await router.acompletion( + model="gpt-5", messages=[{"role": "user", "content": "hi"}], metadata=dict(metadata) + ) + )._hidden_params["model_id"] + for _ in range(5) + } - assert len(filtered) == 1 - assert filtered[0].get("model_info", {}).get("access_groups") == ["AG2"] + assert [deployment["model_info"]["id"] for deployment in healthy] == ["ag2"] + assert served == {"ag2"} class TestAddModelFileIdMappings: @@ -734,12 +697,7 @@ class TestWarnOnProviderCredentialMismatch: ) def test_silent_when_no_provider_scoped_credentials_are_set(self): - assert ( - warn_on_provider_credential_mismatch( - model_name="gpt-5.5", litellm_params={"model": "gpt-5.5"} - ) - is None - ) + assert warn_on_provider_credential_mismatch(model_name="gpt-5.5", litellm_params={"model": "gpt-5.5"}) is None def test_vertex_params_name_vertex_not_bedrock(self): """The hint must follow the params that were actually set, otherwise it @@ -789,8 +747,7 @@ class TestWarnOnProviderCredentialMismatch: mismatch_warnings = [r for r in caplog.records if "resolves to provider" in r.getMessage()] assert len(mismatch_warnings) == 1, ( - "exactly the prefix-less deployment should warn; " - f"got {[r.getMessage() for r in mismatch_warnings]}" + f"exactly the prefix-less deployment should warn; got {[r.getMessage() for r in mismatch_warnings]}" ) assert "aws_region_name" in mismatch_warnings[0].getMessage() @@ -867,7 +824,14 @@ class TestWarnOnProviderCredentialMismatch: ({"model": "no-provider-knows-this-model"}, None), ({"api_base": "https://my-resource.openai.azure.com"}, None), ], - ids=["declared_wins", "prefix_beats_host_flip", "prefix_beats_cohere_chat_flip", "unprefixed_inferred", "unknown", "no_model"], + ids=[ + "declared_wins", + "prefix_beats_host_flip", + "prefix_beats_cohere_chat_flip", + "unprefixed_inferred", + "unknown", + "no_model", + ], ) def test_provider_for_generic_call(litellm_params, expected, monkeypatch): monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") @@ -903,7 +867,9 @@ def test_router_warns_only_for_a_capability_flag_set_under_litellm_params(caplog def test_router_serves_deployments_without_model_capability_flags(): capability_flags: Final = frozenset(ProviderSpecificModelInfo.__annotations__) | frozenset( - key for key in chain.from_iterable(GetModelCostMap.loaded_model_cost_map().values()) if key.startswith("supports_") + key + for key in chain.from_iterable(GetModelCostMap.loaded_model_cost_map().values()) + if key.startswith("supports_") ) declared_params: Final = {"supports_system_message": False, "supports_system_prompt": True} router: Final = Router( diff --git a/tests/unit/test_conftest_isolation.py b/tests/unit/test_conftest_isolation.py index 3fdc2d65da1..e9a153a3f8d 100644 --- a/tests/unit/test_conftest_isolation.py +++ b/tests/unit/test_conftest_isolation.py @@ -9,11 +9,10 @@ import pytest import litellm from litellm import Router -from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module from litellm.caching.caching import DualCache -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from tests.unit import conftest as unit_harness CANARY_MODEL = "conftest-isolation-canary-model" @@ -117,12 +116,12 @@ def test_live_router_membership_is_scoped_to_this_test(): } ] ) - assert _CanaryRouterHolder.router._core in litellm_router_module._live_routers + assert _CanaryRouterHolder.router._core._hosted in unit_harness.live_router_registries def test_live_router_membership_was_rolled_back(): assert _CanaryRouterHolder.router is not None - assert _CanaryRouterHolder.router._core not in litellm_router_module._live_routers + assert _CanaryRouterHolder.router._core._hosted not in unit_harness.live_router_registries def test_aws_cache_reset_sees_new_handlers_and_replaced_caches(monkeypatch: pytest.MonkeyPatch) -> None: @@ -160,7 +159,7 @@ def test_aws_cache_reset_sees_a_replaced_module(monkeypatch: pytest.MonkeyPatch) @pytest.fixture(scope="module", autouse=True) def registered_collection_models() -> Iterator[None]: original_registrations: Final = dict(litellm_utils_module._runtime_registered_model_cost) - original_routers: Final = frozenset(litellm_router_module._live_routers) + original_routers: Final = frozenset(unit_harness.live_router_registries) litellm.register_model( {"collection-isolation-ledger": {"litellm_provider": "openai", "input_cost_per_token": 0.001}} ) @@ -175,8 +174,8 @@ def registered_collection_models() -> Iterator[None]: yield litellm_utils_module._runtime_registered_model_cost.clear() litellm_utils_module._runtime_registered_model_cost.update(original_registrations) - litellm_router_module._live_routers.discard(collection_router._core) - litellm_router_module._live_routers.update(original_routers) + unit_harness.live_router_registries.discard(collection_router._core._hosted) + unit_harness.live_router_registries.update(original_routers) def test_collection_model_metadata_is_not_replayed_into_a_test_catalog() -> None: diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index ceb9ef38887..a446d980bcb 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -57,6 +57,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset( "litellm/llms/databricks/common_utils.py", "litellm/llms/gigachat/authenticator.py", "litellm/llms/github_copilot/per_user_auth.py", + "litellm/llms/github_copilot/authenticator.py", "litellm/llms/litellm_proxy/skills/handler.py", "litellm/llms/openai/common_utils.py", "litellm/llms/openai_like/model_info.py", diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 99c6f6c402c..034b76507cf 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -2533,9 +2533,7 @@ def non_openai_responses_bridge( json.dumps({"access_token": "token", "account_id": "acct", "expires_at": 4102444800}) ) monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) - (tmp_path / "api-key.json").write_text( - json.dumps({"token": "token", "expires_at": 4102444800, "endpoints": {"api": "https://api.githubcopilot.com"}}) - ) + (tmp_path / "access-token").write_text("gho_test") litellm.register_model( { **{ @@ -2579,7 +2577,11 @@ def non_openai_responses_bridge( {"api_base": _FOUNDRY_API_BASE, "api_key": "fake-key"}, ), "chatgpt": ("chatgpt/gpt-6.1-bridge-test", "https://chatgpt.com/backend-api/codex/responses", {}), - "github_copilot": ("github_copilot/gpt-6.1-bridge-test", "https://api.githubcopilot.com/responses", {}), + "github_copilot": ( + "github_copilot/gpt-6.1-bridge-test", + "https://api.githubcopilot.com/responses", + {"api_base": "https://api.githubcopilot.com"}, + ), "databricks": ( "databricks/responses/databricks-gpt-bridge-test", "https://dbx.example/serving-endpoints/responses", diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index a3168ad6d36..d257866b654 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -1949,7 +1949,6 @@ ALIAS_CASES: Final = ( "is_model_access_group_for_wildcard_route", False, ), - ("litellm.router", "Router", "_replay_model_cost_registrations", "replay_model_cost_registrations", False), ("litellm.router", "Router", "_routing_groups", "routing_groups", True), ("litellm.router", "Router", "_update_redis_cache", "update_redis_cache", False), ( diff --git a/tests/unit/test_redact_string_in_error_paths.py b/tests/unit/test_redact_string_in_error_paths.py index eaf2d43a25f..8727df23961 100644 --- a/tests/unit/test_redact_string_in_error_paths.py +++ b/tests/unit/test_redact_string_in_error_paths.py @@ -39,9 +39,7 @@ class TestRedactStringFunction: text = "This is a normal error message with no secrets" assert redact_string(text) == text - @pytest.mark.skipif( - not _ENABLE_SECRET_REDACTION, reason="redaction disabled via env var" - ) + @pytest.mark.skipif(not _ENABLE_SECRET_REDACTION, reason="redaction disabled via env var") def test_redaction_enabled_by_default(self): text = "Bearer sk-9876567890abcdefghij" result = redact_string(text) @@ -77,9 +75,7 @@ class TestOpenAIRealtimeRedaction: from litellm.llms.openai.realtime.handler import OpenAIRealtime handler = OpenAIRealtime() - secret_error = RuntimeError( - "Connection failed for api_key=sk-9876567890abcdefghij" - ) + secret_error = RuntimeError("Connection failed for api_key=sk-9876567890abcdefghij") kwargs = self._call_kwargs() mock_ws = kwargs["websocket"] @@ -96,9 +92,7 @@ class TestBedrockRealtimeRedaction: """Test that redact_string produces safe close reasons for Bedrock-style errors.""" def test_internal_error_message_redacted(self): - secret_error = RuntimeError( - "Failed with aws_secret_access_key=AKIAIOSFODNN7EXAMPLE123456" - ) + secret_error = RuntimeError("Failed with aws_secret_access_key=AKIAIOSFODNN7EXAMPLE123456") reason = redact_string(f"Internal error: {str(secret_error)}") assert "AKIAIOSFODNN7EXAMPLE123456" not in reason @@ -112,17 +106,13 @@ class TestLLMHTTPHandlerRealtimeRedaction: def test_internal_server_error_pattern(self): error_msg = "Connection failed for api_key=sk-secret-key-12345678" - assert "sk-secret-key-12345678" not in redact_string( - f"Internal server error: {error_msg}" - ) + assert "sk-secret-key-12345678" not in redact_string(f"Internal server error: {error_msg}") class TestProxyStreamingDataGeneratorRedaction: def test_redact_traceback_format_exc(self): try: - raise RuntimeError( - "Failed connecting to api_key=sk-9876567890abcdefghij at https://api.example.com" - ) + raise RuntimeError("Failed connecting to api_key=sk-9876567890abcdefghij at https://api.example.com") except RuntimeError: raw_tb = traceback.format_exc() @@ -134,8 +124,7 @@ class TestProxyStreamingDataGeneratorRedaction: class TestRouterFallbackFailureTracebackRedaction: - """Test the fallback-failure logs in router.py's - async_function_with_fallbacks_common_utils. Both call sites must redact the + """Test the router's fallback-failure logs. Both call sites must redact the traceback at the call site with redact_string() rather than hand a live exception to exc_info=True. SecretRedactionFilter rewrites record.exc_text, but record.exc_info stays an exception object no filter can rewrite, so any @@ -144,57 +133,44 @@ class TestRouterFallbackFailureTracebackRedaction: @pytest.mark.asyncio async def test_fallback_failure_does_not_leak_secret_via_exc_info(self, caplog): - """The helper is driven from inside an `except` block because that is the only - way production reaches it, and the entry-point debug log takes its traceback - from the active exception. With no exception in flight sys.exc_info() is empty, - so an exc_info=True regression there would degrade to (None, None, None) and - this test would pass against it. - """ + """Both the primary group and its fallback fail with the secret in the error, so the + router logs the fallback failure while that exception is in flight.""" import litellm + secret = "sk-testsecretvalue1234567890abcdef" router = litellm.Router( model_list=[ { "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key"}, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "fake-key", + "mock_response": Exception(f"primary deployment failed api_key={secret}"), + }, }, { "model_name": "claude-3-haiku", - "litellm_params": {"model": "anthropic/claude-3-haiku-20240307", "api_key": "fake-key"}, + "litellm_params": { + "model": "anthropic/claude-3-haiku-20240307", + "api_key": "fake-key", + "mock_response": Exception(f"boom api_key={secret}"), + }, }, ], + fallbacks=[{"gpt-3.5-turbo": ["claude-3-haiku"]}], ) - secret = "sk-testsecretvalue1234567890abcdef" - - with patch( - "litellm.router.run_async_fallback", - new=AsyncMock(side_effect=RuntimeError(f"boom api_key={secret}")), - ): - try: - raise ValueError(f"primary deployment failed api_key={secret}") - except ValueError as original_exception: - with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): - with pytest.raises(ValueError, match='primary deployment failed api_key=sk-testsecretvalu'): - await router.async_function_with_fallbacks_common_utils( - e=original_exception, - disable_fallbacks=False, - fallbacks=[{"gpt-3.5-turbo": ["claude-3-haiku"]}], - context_window_fallbacks=None, - content_policy_fallbacks=None, - model_group="gpt-3.5-turbo", - args=(), - kwargs={"model": "gpt-3.5-turbo"}, - ) + with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): + with pytest.raises(litellm.InternalServerError): + await router.acompletion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "hi"}]) debug_records = [r for r in caplog.records if r.levelno == logging.DEBUG] assert debug_records, "expected the entry-point debug log, which carries the active traceback" error_records = [r for r in caplog.records if r.levelno == logging.ERROR] - assert error_records, "expected an error log for the fallback failure" - assert any( - "Cooldown Deployments" in r.getMessage() for r in error_records - ), "expected the fallback-failure log, not an unrelated error" + assert any("Error occurred while trying to do fallbacks" in r.getMessage() for r in error_records), ( + "expected the fallback-failure log, not an unrelated error" + ) for record in caplog.records: assert secret not in record.getMessage() @@ -255,9 +231,7 @@ class TestGeminiIngestionHeaders: mock_client = AsyncMock() mock_response = MagicMock() mock_response.status_code = 200 - mock_response.headers = { - "x-goog-upload-url": "https://upload.example.com/upload123" - } + mock_response.headers = {"x-goog-upload-url": "https://upload.example.com/upload123"} mock_client.post.return_value = mock_response with patch( diff --git a/tests/unit/test_responses_streaming_container_ownership.py b/tests/unit/test_responses_streaming_container_ownership.py index 65889052bce..737a6eaafcd 100644 --- a/tests/unit/test_responses_streaming_container_ownership.py +++ b/tests/unit/test_responses_streaming_container_ownership.py @@ -1,9 +1,8 @@ """ Regression for #30210. -When streaming /v1/responses goes through the proxy + Router, the -streaming iterator is wrapped by ``Router._aresponses_streaming_iterator`` -which returns ``FallbackResponsesStreamWrapper``. That wrapper set +When streaming /v1/responses goes through the proxy + Router, the router +returns a ``FallbackResponsesStreamWrapper``. That wrapper set ``self.completed_response = None`` in __init__ and never updated it, so the proxy's container-ownership hook (which reads ``getattr(stream_response, "completed_response", None)`` via @@ -12,91 +11,81 @@ saw None on every streaming call and silently recorded nothing — follow-up ``GET /v1/containers//files`` then 403'd for the very key that created the container. -Tests below construct the wrapper from a fake async generator that -yields one terminal ``response.completed`` chunk and assert the -wrapper now carries that chunk on ``completed_response`` so the -proxy hook can walk it. +Tests below stream a mocked upstream through ``Router.aresponses`` and +assert the returned stream carries the terminal chunk on +``completed_response`` so the proxy hook can walk it. """ import asyncio +import json from types import SimpleNamespace from unittest.mock import MagicMock, patch -import pytest +import httpx +import respx + +import litellm + +_DELTA = { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "hello", + "logprobs": [], +} -def _make_wrapper_class(): - """Pull ``FallbackResponsesStreamWrapper`` out by running - ``Router._aresponses_streaming_iterator`` long enough to construct - the class then return it. Mirrors how the wrapper is actually - instantiated in production.""" - from litellm.router import Router - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) +def _terminal_event(event_type: str, status: str, sequence_number: int = 2, created_at: int = 1761600000): + return { + "type": event_type, + "sequence_number": sequence_number, + "response": { + "id": "resp_test", + "object": "response", + "created_at": created_at, + "status": status, + "model": "gpt-5.5", + "output": [], + "container": {"id": "cntr_test", "type": "code_interpreter"}, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + } - # Minimal source iterator stub with every attribute the wrapper - # copies in __init__ (see router.py:2552-2583). - source = SimpleNamespace( - response=None, - model="openai/gpt-5.5", - logging_obj=None, - responses_api_provider_config=None, - start_time=None, - litellm_metadata=None, - custom_llm_provider="openai", - request_data={}, - call_type="aresponses", - _hidden_params={}, - ) - # The class is defined inside _aresponses_streaming_iterator; capture - # it by patching FallbackResponsesStreamWrapper into a sentinel on - # construction. - captured = {} - - real_router_module = __import__("litellm.router", fromlist=["Router"]) +def _stream_through_router(*events): + """Stream ``events`` from a mocked upstream through ``Router.aresponses``; + return the router's stream object after draining it, and the chunks.""" + body = b"".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n".encode() for e in events) async def _drive(): - async def empty_gen(): - if False: - yield # pragma: no cover - return - - router = Router( + router = litellm.Router( model_list=[ { - "model_name": "openai/gpt-5.5", - "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}, + "model_name": "gpt-5.5", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "api_base": "https://api.openai.com/v1", + }, } ] ) - wrapped = await router._aresponses_streaming_iterator( - response=source, # type: ignore[arg-type] - initial_kwargs={}, - ) - captured["wrapper_cls"] = type(wrapped) - captured["instance"] = wrapped + with respx.mock() as mock_router: + mock_router.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body) + ) + stream = await router.aresponses(model="gpt-5.5", input="hi", stream=True) + chunks = [chunk async for chunk in stream] + return stream, chunks - asyncio.run(_drive()) - return captured["wrapper_cls"], captured["instance"] - - -def _terminal_chunk(event_type: str): - """A SimpleNamespace shaped like the openai responses-api terminal - event chunks the wrapper inspects (.type attribute).""" - return SimpleNamespace( - type=event_type, - response=SimpleNamespace( - id="resp_test", - output=[], - container={"id": "cntr_test", "type": "code_interpreter"}, - ), - ) - - -def _non_terminal_chunk(event_type: str = "response.output_text.delta"): - return SimpleNamespace(type=event_type, delta="hello") + with patch.object(litellm, "disable_aiohttp_transport", True): + litellm.in_memory_llm_clients_cache.flush_cache() + try: + return asyncio.run(_drive()) + finally: + litellm.in_memory_llm_clients_cache.flush_cache() class TestStreamWrapperCapturesTerminalEvent: @@ -104,76 +93,29 @@ class TestStreamWrapperCapturesTerminalEvent: """The #30210 bug: a forwarded ``response.completed`` chunk used to leave ``completed_response`` at None on the wrapper. Verify it now carries the chunk.""" - wrapper_cls, _ = _make_wrapper_class() - - async def gen(): - yield _non_terminal_chunk() - yield _terminal_chunk("response.completed") - - wrapper = wrapper_cls(gen()) - # Drain the wrapper. - out = asyncio.run(_drain(wrapper)) - assert len(out) == 2 - assert wrapper.completed_response is not None, ( + stream, chunks = _stream_through_router(_DELTA, _terminal_event("response.completed", "completed")) + assert len(chunks) == 2 + assert stream.completed_response is not None, ( "FallbackResponsesStreamWrapper.completed_response is still None " "after a response.completed chunk passed through — the proxy " "container-ownership hook will 403 follow-up file lookups (#30210)" ) - assert wrapper.completed_response.type == "response.completed" + assert stream.completed_response.type == "response.completed" def test_terminal_incomplete_event_is_recorded(self): - wrapper_cls, _ = _make_wrapper_class() - - async def gen(): - yield _terminal_chunk("response.incomplete") - - wrapper = wrapper_cls(gen()) - asyncio.run(_drain(wrapper)) - assert wrapper.completed_response is not None - assert wrapper.completed_response.type == "response.incomplete" - - def test_terminal_failed_event_is_recorded(self): - wrapper_cls, _ = _make_wrapper_class() - - async def gen(): - yield _terminal_chunk("response.failed") - - wrapper = wrapper_cls(gen()) - asyncio.run(_drain(wrapper)) - assert wrapper.completed_response is not None - assert wrapper.completed_response.type == "response.failed" - - def test_non_terminal_chunks_do_not_set_completed_response(self): - wrapper_cls, _ = _make_wrapper_class() - - async def gen(): - yield _non_terminal_chunk("response.output_text.delta") - yield _non_terminal_chunk("response.code_interpreter.in_progress") - - wrapper = wrapper_cls(gen()) - asyncio.run(_drain(wrapper)) - assert ( - wrapper.completed_response is None - ), "non-terminal chunks must not set completed_response" + stream, _ = _stream_through_router(_terminal_event("response.incomplete", "incomplete", sequence_number=1)) + assert stream.completed_response is not None + assert stream.completed_response.type == "response.incomplete" def test_first_terminal_event_wins(self): """Real streams only emit one terminal event, but defend against future producers emitting more: keep the first one (the inner source iterator behaves the same way).""" - wrapper_cls, _ = _make_wrapper_class() - - first = _terminal_chunk("response.completed") - first.response.id = "resp_first" - second = _terminal_chunk("response.completed") - second.response.id = "resp_second" - - async def gen(): - yield first - yield second - - wrapper = wrapper_cls(gen()) - asyncio.run(_drain(wrapper)) - assert wrapper.completed_response.response.id == "resp_first" + stream, _ = _stream_through_router( + _terminal_event("response.completed", "completed", sequence_number=1, created_at=1), + _terminal_event("response.completed", "completed", sequence_number=2, created_at=2), + ) + assert stream.completed_response.response.created_at == 1 class TestProxyOwnershipHookReadsCompletedResponse: @@ -186,19 +128,10 @@ class TestProxyOwnershipHookReadsCompletedResponse: ProxyBaseLLMRequestProcessing, ) - wrapper_cls, _ = _make_wrapper_class() + stream, _ = _stream_through_router(_terminal_event("response.completed", "completed", sequence_number=1)) - async def gen(): - yield _terminal_chunk("response.completed") - - wrapper = wrapper_cls(gen()) - asyncio.run(_drain(wrapper)) - - extracted = ProxyBaseLLMRequestProcessing.extract_completed_responses_response( - wrapper - ) + extracted = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream) assert extracted is not None - assert extracted.id == "resp_test" assert extracted.container["id"] == "cntr_test" @@ -209,24 +142,18 @@ class TestSilentSkipNowLogged: def test_warning_logged_when_completed_response_missing(self, caplog): import logging + from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) - # Wrap a generator that produces NO terminal event so the - # wrapper stays at completed_response=None — same shape as the + # A stream that never saw a terminal event: the same shape as the # pre-fix bug. - wrapper_cls, _ = _make_wrapper_class() - - async def gen(): - yield _non_terminal_chunk() - - wrapper = wrapper_cls(gen()) + stream = SimpleNamespace(completed_response=None) async def driver(): async def inner_gen(): - async for c in wrapper: - yield c + yield SimpleNamespace(type="response.output_text.delta", delta="hello") # Patch _record_container_owners_from_responses_if_needed to # a noop async so the warning branch is exercised in @@ -237,7 +164,7 @@ class TestSilentSkipNowLogged: new=MagicMock(), ): wrapped = ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership( - original_stream_response=wrapper, + original_stream_response=stream, wrapped_generator=inner_gen(), user_api_key_dict=MagicMock(), ) @@ -248,14 +175,5 @@ class TestSilentSkipNowLogged: asyncio.run(driver()) assert any( - "Container ownership recording skipped on streaming /v1/responses" - in r.message - for r in caplog.records + "Container ownership recording skipped on streaming /v1/responses" in r.message for r in caplog.records ), "silent-skip warning never fired despite completed_response=None" - - -async def _drain(it): - out = [] - async for chunk in it: - out.append(chunk) - return out diff --git a/tests/unit/test_router/core/conftest.py b/tests/unit/test_router/core/conftest.py index 3477cfc0c3d..7475071c382 100644 --- a/tests/unit/test_router/core/conftest.py +++ b/tests/unit/test_router/core/conftest.py @@ -1,11 +1,11 @@ -from collections.abc import Callable, Generator +from collections.abc import Callable from pathlib import Path from typing import Final import pytest from litellm.router_core import RouterCore -from tests.unit.test_router.core.scripted_host import CORE_BUILDERS, CoreDeclined, ScriptedHost +from tests.unit.test_router.core.scripted_host import CORE_BUILDERS, ScriptedHost @pytest.fixture @@ -16,15 +16,12 @@ def host() -> ScriptedHost: @pytest.fixture(params=sorted(CORE_BUILDERS)) def make_core(request: pytest.FixtureRequest, host: ScriptedHost) -> Callable[..., RouterCore]: """``make_core(**router_arguments)``: the core under test, on the ``host`` fixture. Every test that - uses it runs once per core in ``CORE_BUILDERS`` (``-k python`` or ``-k rust`` picks one).""" + uses it runs once per core in ``CORE_BUILDERS``.""" name: Final[str] = str(request.param) # pyright: ignore[reportAny] # the fixture parameter builder: Final = CORE_BUILDERS[name] def build(**arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments - try: - return builder(host, **arguments) - except CoreDeclined as declined: - pytest.skip(f"{name} core declines this configuration: {declined}") + return builder(host, **arguments) return build @@ -47,23 +44,3 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: name: Final = "::".join((path.relative_to(here).as_posix(), *item.nodeid.split("::")[1:])) if name in _RUST_PARITY_PENDING: item.add_marker(pytest.mark.xfail(reason="Rust core parity pending", strict=True)) - - -@pytest.hookimpl(hookwrapper=True) -def pytest_runtest_call(item: pytest.Item) -> Generator[None, object, None]: - """A test that reaches a member or request the Rust core does not serve yet is skipped with - that member named, so the run separates "not served yet" from a wrong answer.""" - outcome: Final = yield - excinfo: Final[object] = getattr(outcome, "excinfo", None) - if not isinstance(excinfo, tuple): - return - info: Final[tuple[object, ...]] = excinfo # pyright: ignore[reportUnknownVariableType] # pluggy's (type, value, traceback) - if len(info) < 2: - return - from litellm.router_rust.core import RustCoreUnsupported - - error: Final[object] = info[1] - cause: Final[object] = getattr(error, "__cause__", None) - force_exception: Final[object] = getattr(outcome, "force_exception", None) - if (isinstance(error, RustCoreUnsupported) or isinstance(cause, RustCoreUnsupported)) and callable(force_exception): - force_exception(pytest.skip.Exception(f"not served by the Rust core yet: {error}", _use_item_location=True)) diff --git a/tests/unit/test_router/core/scripted_host.py b/tests/unit/test_router/core/scripted_host.py index f673ab98610..33164bd4408 100644 --- a/tests/unit/test_router/core/scripted_host.py +++ b/tests/unit/test_router/core/scripted_host.py @@ -22,7 +22,6 @@ from litellm.litellm_core_utils.agentic_followup_failure import mark_agentic_fol from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager from litellm.router_config import RouterDynamicConfig, build_router_config from litellm.router_core import RouterCore -from litellm.router_python_core import PythonCore from litellm.router_python_host import CallbackList, Endpoint, PythonHost from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -482,43 +481,19 @@ def deployment( } -def python_core(host: ScriptedHost, **arguments: object) -> PythonCore: # kwargs-ok: the router's constructor arguments - """A ``PythonCore`` built the way ``Router`` builds one, on ``host``.""" - config: Final = build_router_config(arguments) - return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the same arguments Router binds - - -class CoreDeclined(Exception): - """The core does not serve this configuration (the router would stay on the Python core).""" - - CoreBuilder = Callable[..., RouterCore] """``builder(host, **router_arguments)``: a router core on ``host``.""" def rust_core(host: ScriptedHost, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments - """A ``RustCore`` on ``host``, built the way ``make_core`` builds one when the Rust router is on.""" + """A ``RustCore`` on ``host``, built the way ``make_core`` builds one.""" from litellm.router_rust.core import RustCore - from litellm.rust_bridge._native import RustBridgeDeclined - config: Final = build_router_config(arguments) # an invalid config raises here, as for the Python core - try: - return RustCore.build(config, host) - except RustBridgeDeclined as declined: # make_core keeps such a router on the Python core - raise CoreDeclined(str(declined)) from declined + return RustCore.build(build_router_config(arguments), host) -CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core, "rust": rust_core} - - -def records_outcomes_through_callbacks(core: RouterCore) -> bool: - """Whether ``core`` learns attempt outcomes from callbacks it registers on its host. The Python - core does; the Rust core records every outcome from the attempt itself and registers none.""" - return isinstance(core, PythonCore) - - -"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a -configuration raises ``CoreDeclined``, and the test is skipped for that core.""" +CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"rust": rust_core} +"""Every core the suite runs on.""" __all__ = [ @@ -526,7 +501,6 @@ __all__ = [ "INFERENCE_OPS", "Call", "CoreBuilder", - "CoreDeclined", "HostCallbacks", "ScriptedHost", "auth_error", @@ -537,7 +511,6 @@ __all__ = [ "deployment", "dynamic_config", "embedding_response", - "python_core", "rate_limit", "rust_core", "server_error", diff --git a/tests/unit/test_router/core/test_construction_precall_characterization.py b/tests/unit/test_router/core/test_construction_precall_characterization.py index a56cd78ccac..f58d9a35110 100644 --- a/tests/unit/test_router/core/test_construction_precall_characterization.py +++ b/tests/unit/test_router/core/test_construction_precall_characterization.py @@ -7,21 +7,28 @@ creates, alerting, optional pre-call check registration, ``observe``, ``register Each test asserts what the Python core does today; oddities are named in the test's docstring. """ +import inspect import logging from collections.abc import Callable, Generator, Mapping from typing import Final import pytest +import redis +import redis.asyncio +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeConnection +from redis.asyncio import BlockingConnectionPool import litellm from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCache +from litellm.router import Router from litellm.router_api import ChangeSet, router_init_args from litellm.router_core import RouterCore -from litellm.router_python_core import PythonCore from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck +from litellm.scheduler import FlowItem from litellm.types.router import ( AlertingConfig, AllowedFailsPolicy, @@ -34,7 +41,7 @@ from tests.unit.test_router.core.construction_precall_helpers import ( collect_other_routers, host_callback_names, ) -from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, records_outcomes_through_callbacks +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rust_core MakeCore = Callable[..., RouterCore] MESSAGES: Final = [{"role": "user", "content": "hi"}] @@ -619,6 +626,41 @@ def test_applying_a_redis_cache_attaches_it_to_the_routing_state_and_the_pins(ma assert core.registry().plugins("prompt_caching_pins").redis_cache is redis_cache +@pytest.fixture +def shared_redis(monkeypatch: pytest.MonkeyPatch) -> FakeServer: + server: Final = FakeServer() + + class Pool(BlockingConnectionPool): + def __init__(self, **kwargs: object) -> None: # kwargs-ok: forwarded to the redis pool + super().__init__(connection_class=FakeConnection, server=server, **kwargs) + + monkeypatch.setattr(redis, "Redis", FakeRedis) + monkeypatch.setattr(redis.asyncio, "BlockingConnectionPool", Pool) + return server + + +async def test_two_workers_on_one_redis_share_the_priority_queue(make_core: MakeCore, shared_redis: FakeServer) -> None: + first: Final = _one(make_core, redis_host="scheduler-shared", redis_port=6399) + second: Final = _one(make_core, redis_host="scheduler-shared", redis_port=6399) + await first.scheduler.add_request(FlowItem(priority=0, request_id="urgent", model_name="g")) # pyright: ignore[reportAttributeAccessIssue] + waiter: Final = FlowItem(priority=5, request_id="waiter", model_name="g") + await second.scheduler.add_request(waiter) # pyright: ignore[reportAttributeAccessIssue] + + assert not await second.scheduler.poll(waiter, []) # pyright: ignore[reportAttributeAccessIssue] + await first.scheduler.remove_request("urgent", "g") # pyright: ignore[reportAttributeAccessIssue] + assert await second.scheduler.poll(waiter, []) # pyright: ignore[reportAttributeAccessIssue] + + +async def test_without_redis_each_worker_keeps_its_own_queue(make_core: MakeCore) -> None: + first: Final = _one(make_core) + second: Final = _one(make_core) + await first.scheduler.add_request(FlowItem(priority=0, request_id="urgent", model_name="g")) # pyright: ignore[reportAttributeAccessIssue] + waiter: Final = FlowItem(priority=5, request_id="waiter", model_name="g") + await second.scheduler.add_request(waiter) # pyright: ignore[reportAttributeAccessIssue] + + assert await second.scheduler.poll(waiter, []) # pyright: ignore[reportAttributeAccessIssue] + + def test_cache_responses_installs_a_local_response_cache_on_the_host(make_core: MakeCore, host: ScriptedHost) -> None: core: Final = _one(make_core, cache_responses=True) assert isinstance(host.cache, litellm.Cache) @@ -694,24 +736,6 @@ def test_alerting_config_registers_slack_alerting_and_its_slow_response_callback ) -def test_the_python_core_registers_its_sync_outcome_callback_before_the_slow_response_callback( - host: ScriptedHost, -) -> None: - """Python core only: it learns sync outcomes from a success callback on its host. The Rust core - records outcomes from the attempt itself and registers no outcome callbacks.""" - from tests.unit.test_router.core.scripted_host import python_core - - python_core( - host, - model_list=[deployment("g", "a")], - alerting_config=AlertingConfig(webhook_url="https://example.invalid/hook", alerting_threshold=10), - ) - assert host_callback_names(host)["success_callback"] == [ - "sync_deployment_callback_on_success", - "response_taking_too_long_callback", - ] - - def test_observe_reports_no_slack_alerting_without_an_alerting_config(make_core: MakeCore) -> None: assert _one(make_core).observe().slack_alerting_logger is None @@ -884,11 +908,7 @@ def test_runtime_removal_keeps_the_host_callback_while_another_live_router_uses_ host) still has one of that type.""" collect_other_routers() other_host: Final = CallbackHost() - from tests.unit.test_router.core.scripted_host import python_core - - other: Final = python_core( - other_host, model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"] - ) + other: Final = rust_core(other_host, model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"]) core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) core.apply(ChangeSet(settings={"optional_pre_call_checks": []})) assert core.optional_callbacks == [] # pyright: ignore[reportAttributeAccessIssue] @@ -948,30 +968,13 @@ def test_retiring_one_router_keeps_the_callbacks_of_another_on_the_same_host( assert host.lists["input_callback"] == [second.leastbusy_logger] # pyright: ignore[reportAttributeAccessIssue] -def test_retiring_one_python_core_keeps_the_outcome_callback_of_another_on_the_same_host( - host: ScriptedHost, -) -> None: - """Python core only: each one registers an async outcome callback bound to itself; the Rust core - records outcomes from the attempt itself and registers none.""" - from tests.unit.test_router.core.scripted_host import python_core - - collect_other_routers() - first: Final = python_core(host, model_list=[deployment("g", "a")], routing_strategy="least-busy") - second: Final = python_core(host, model_list=[deployment("g", "a")], routing_strategy="least-busy") - first.retire() - assert host_callback_names(host)["_async_success_callback"] == ["deployment_callback_on_success"] - assert all(callback.__self__ is second for callback in host.lists["_async_success_callback"]) # pyright: ignore[reportAttributeAccessIssue] - - def test_retire_hands_a_vacated_check_to_another_live_router_on_a_different_host( make_core: MakeCore, host: ScriptedHost ) -> None: """Oddity: retire re-registers the same-keyed check of any other live router in the process into this host's list, even one built on a different host.""" collect_other_routers() - from tests.unit.test_router.core.scripted_host import python_core - - other: Final = python_core( + other: Final = rust_core( CallbackHost(), model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"] ) core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) @@ -997,14 +1000,6 @@ class RecordingLogging: self.calls.append(("add", type(callback).__name__)) -def _bound_outcome_callbacks(core: RouterCore) -> list[tuple[str, object]]: # mutable-ok: a fresh list - return ( - [("functions", ["async_failure", "async_success", "failure"])] - if records_outcomes_through_callbacks(core) - else [] - ) - - def test_observe_request_scoped_removes_every_callback_the_router_registered_globally( make_core: MakeCore, host: ScriptedHost ) -> None: @@ -1018,13 +1013,12 @@ def test_observe_request_scoped_removes_every_callback_the_router_registered_glo def test_observe_with_a_logging_object_binds_the_router_outcome_callbacks( make_core: MakeCore, host: ScriptedHost ) -> None: - """The Python core binds its three outcome callbacks; the Rust core records every attempt's - outcome itself and binds none, as it registers none on the host.""" + """The core records every attempt's outcome itself and binds none, as it registers none on the host.""" core: Final = _one(make_core) before: Final = host_callback_names(host) logging_obj: Final = RecordingLogging() core.observe(logging_obj) # pyright: ignore[reportArgumentType] - assert logging_obj.calls == _bound_outcome_callbacks(core) + assert logging_obj.calls == [] assert host_callback_names(host) == before @@ -1035,7 +1029,6 @@ def test_observe_request_scoped_then_binding_hands_selectors_and_policies_to_the logging_obj: Final = RecordingLogging() core.observe(logging_obj) # pyright: ignore[reportArgumentType] assert logging_obj.calls == [ - *_bound_outcome_callbacks(core), ("replace", "LeastBusyLoggingHandler"), ("add", "PromptCachingDeploymentCheck"), ] @@ -1048,7 +1041,7 @@ def test_binding_without_request_scoping_hands_the_request_no_selector_or_policy core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"]) logging_obj: Final = RecordingLogging() core.observe(logging_obj) # pyright: ignore[reportCallIssue] # observe takes a logging object on every core - assert logging_obj.calls == _bound_outcome_callbacks(core) + assert logging_obj.calls == [] def test_binding_skips_a_policy_the_host_already_runs(make_core: MakeCore, host: ScriptedHost) -> None: @@ -1060,7 +1053,7 @@ def test_binding_skips_a_policy_the_host_already_runs(make_core: MakeCore, host: host.lists["callbacks"].append(PromptCachingDeploymentCheck(cache=DualCache())) logging_obj: Final = RecordingLogging() core.observe(logging_obj) # pyright: ignore[reportCallIssue] # observe takes a logging object on every core - assert logging_obj.calls == [*_bound_outcome_callbacks(core), ("replace", "LeastBusyLoggingHandler")] + assert logging_obj.calls == [("replace", "LeastBusyLoggingHandler")] # --- register_plugin --------------------------------------------------------------------- @@ -1224,7 +1217,7 @@ def test_registry_settings_and_health_views_are_built_once(make_core: MakeCore) def test_get_valid_args_lists_every_constructor_argument_but_self_and_the_host() -> None: - names: Final = PythonCore.get_valid_args() + names: Final = router_init_args() assert names[0] == "model_list" assert "self" not in names assert "_host" not in names @@ -1232,12 +1225,12 @@ def test_get_valid_args_lists_every_constructor_argument_but_self_and_the_host() assert len(names) == len(set(names)) -def test_get_valid_args_matches_the_router_init_args() -> None: - assert tuple(PythonCore.get_valid_args()) == tuple(router_init_args()) +def test_get_valid_args_matches_the_router_signature() -> None: + assert tuple(inspect.signature(Router).parameters) == tuple(router_init_args()) def test_get_valid_args_on_a_built_router(make_core: MakeCore) -> None: - assert _one(make_core).get_valid_args() == PythonCore.get_valid_args() # pyright: ignore[reportAttributeAccessIssue] + assert _one(make_core).get_valid_args() == list(router_init_args()) # pyright: ignore[reportAttributeAccessIssue] # --- print_deployment ---------------------------------------------------------------------- diff --git a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py index 96d4d1879ce..9cdb3f55d17 100644 --- a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py +++ b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py @@ -20,7 +20,7 @@ import litellm from litellm.router_core import RouterCore from litellm.types.router import RouterRateLimitError from tests.unit.test_router.core.construction_precall_helpers import CallbackHost, restoring_model_cost, windowed -from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, python_core +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rust_core MakeCore = Callable[..., RouterCore] PICKS: Final = 60 # three-way shuffles miss a deployment in 20 picks about once in a thousand runs @@ -326,7 +326,7 @@ async def test_a_deployment_context_window_is_written_into_the_process_cost_map( model_list=[deployment("g", "declares-five", model=UNMAPPED_MODEL, model_info={"max_input_tokens": 5})] ) first.retire() - other: Final = python_core( + other: Final = rust_core( CallbackHost(), model_list=[deployment("g", "declares-nothing", model=UNMAPPED_MODEL)], enable_pre_call_checks=True, @@ -343,7 +343,7 @@ async def test_a_deployment_context_window_stays_off_a_mapped_models_process_ent model_list=[deployment("g", "declares-five", model=NO_JSON_MODE_MODEL, model_info={"max_input_tokens": 5})] ) first.retire() - other: Final = python_core( + other: Final = rust_core( CallbackHost(), model_list=[deployment("g", "declares-nothing", model=NO_JSON_MODE_MODEL)], enable_pre_call_checks=True, diff --git a/tests/unit/test_router/core/test_endpoints.py b/tests/unit/test_router/core/test_endpoints.py index ce56ae403ba..d129dc9f408 100644 --- a/tests/unit/test_router/core/test_endpoints.py +++ b/tests/unit/test_router/core/test_endpoints.py @@ -22,7 +22,6 @@ from tests.unit.test_router.core.scripted_host import ( completion_response, deployment, embedding_response, - records_outcomes_through_callbacks, ) MESSAGES: Final = [{"role": "user", "content": "hi"}] @@ -279,21 +278,18 @@ async def test_request_level_api_key_replaces_the_deployment_key_on_the_host_cal ) -async def test_construction_registers_one_success_and_one_failure_callback_per_mode( - make_core: MakeCore, host: ScriptedHost -): +async def test_construction_registers_no_outcome_callbacks(make_core: MakeCore, host: ScriptedHost): + """The core records every attempt's outcome itself, so it registers no callbacks on its host.""" core: Final = _one_group(make_core) - registered: Final = {name: len(callbacks) for name, callbacks in host.lists.items()} - one: Final = 1 if records_outcomes_through_callbacks(core) else 0 - assert registered == { + assert core.registry() is not None + assert {name: len(callbacks) for name, callbacks in host.lists.items()} == { "callbacks": 0, - "success_callback": one, - "_async_success_callback": one, - "failure_callback": one, - "_async_failure_callback": one, + "success_callback": 0, + "_async_success_callback": 0, + "failure_callback": 0, + "_async_failure_callback": 0, "input_callback": 0, } - assert all(callable(callbacks[0]) for callbacks in host.lists.values() if callbacks) async def test_retire_removes_the_callbacks_the_core_registered(make_core: MakeCore, host: ScriptedHost): diff --git a/tests/unit/test_router/core/test_legacy_members.py b/tests/unit/test_router/core/test_legacy_members.py index 1158a44e7fd..213fd5b84ba 100644 --- a/tests/unit/test_router/core/test_legacy_members.py +++ b/tests/unit/test_router/core/test_legacy_members.py @@ -1,11 +1,10 @@ """ -The old public Router members outside the entry points, on each core: registry lookups, the +The old public Router members outside the entry points: registry lookups, the routing-state lifecycle, the routing callbacks' checks, one provider call's accounting and the pre-routing strategy builders (``RustLegacyMembers``, litellm/router_rust/surface/legacy_members.py). """ import asyncio -import inspect from collections.abc import Mapping from typing import Final @@ -13,18 +12,10 @@ import pytest import litellm from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY -from litellm.containers.endpoint_factory import ( - _generated_endpoints, # pyright: ignore[reportPrivateUsage] # the names the routers register -) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.token_counter import served_openai_models from litellm.router_api import RoutingContext from litellm.router_core import RouterCore -from litellm.router_python_core import PythonCore -from litellm.router_rust.surface.legacy_members import ( - _ASSISTANTS_ENDPOINTS, # pyright: ignore[reportPrivateUsage] # pinned against the Python core below - _ROUTER_ENDPOINTS, # pyright: ignore[reportPrivateUsage] # see above -) from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook from litellm.router_strategy.auto_router.auto_router_model_naming import HEURISTIC_V2_CAPABILITY from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy @@ -46,7 +37,6 @@ from tests.unit.test_router.core.scripted_host import ( ScriptedHost, completion_response, deployment, - python_core, rate_limit, rust_core, server_error, @@ -179,12 +169,6 @@ def test_served_openai_models_reads_the_group_and_follows_a_model_list_change(ma ) # pyright: ignore[reportAttributeAccessIssue] # see above -def test_the_tier_params_never_dropped_constant_matches(make_core: MakeCore) -> None: - core: Final = make_core(model_list=list(GROUP)) - - assert core.TIER_PARAMS_NEVER_DROPPED == PythonCore.TIER_PARAMS_NEVER_DROPPED # pyright: ignore[reportAttributeAccessIssue] # a class constant - - def test_auto_router_capability_violation_resolves_the_limit_on_every_call(make_core: MakeCore) -> None: limits: Final[list[int | None]] = [None] # mutable-ok: what the resolver answers @@ -326,19 +310,6 @@ def test_initialize_router_endpoints_replaces_only_what_it_binds(make_core: Make assert core.acreate_file is assigned # pyright: ignore[reportAttributeAccessIssue] # see above -def test_the_endpoint_names_are_the_ones_the_python_core_binds() -> None: - bound: Final = vars(python_core(ScriptedHost(), model_list=[])) - listed: Final = {*_ROUTER_ENDPOINTS, *_ASSISTANTS_ENDPOINTS, *_generated_endpoints} - - assert listed <= set(bound) - wrappers: Final = { - name - for name, value in bound.items() - if inspect.isfunction(value) and value.__qualname__.startswith("PythonCore.factory_function.") - } - assert wrappers <= listed - - # --- pre-routing strategies ---------------------------------------------------------------------- diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index 7c2658d3d0d..f09ad5973e2 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -1476,7 +1476,6 @@ class TestRouterIOTokenIntegration: assert await _io_counters_since(router, first_minute) == (0, 0) - @pytest.mark.asyncio @pytest.mark.parametrize( ("call", "original_function", "mock_error"), @@ -1549,41 +1548,6 @@ class TestContextSlotRetention: set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is None - @pytest.mark.asyncio - async def test_router_does_not_pin_kwargs_without_io_limits(self): - router = Router( - model_list=[ - { - "model_name": "plain", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, - } - ] - ) - set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata") - kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - deployment = router.get_deployment_by_model_group_name("plain") - assert deployment is not None - router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) - assert get_io_token_rate_limit_request_kwargs() is None - - @pytest.mark.asyncio - async def test_router_pins_kwargs_for_io_limited_deployment(self): - router = Router( - model_list=[ - { - "model_name": "limited", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata") - kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} - deployment = router.get_deployment_by_model_group_name("limited") - assert deployment is not None - router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) - assert get_io_token_rate_limit_request_kwargs() is kwargs - @pytest.mark.asyncio async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_token_count() -> None: diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a21777fbee6..dbc47fd6962 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -31,8 +31,7 @@ from pydantic import BaseModel, TypeAdapter import litellm from litellm import APIConnectionError, Router from litellm._logging import verbose_logger, verbose_router_logger -from litellm.caching import RedisCache, RedisClusterCache -from litellm.caching.caching import DualCache +from litellm.caching import RedisClusterCache from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError, ModifyResponseException @@ -49,7 +48,7 @@ from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import ( from litellm.llms.bedrock.common_utils import BedrockError from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, ProxyException, UserAPIKeyAuth -from litellm.router import ( +from litellm.router_rust.invoke.streaming import ( MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS, FallbackAwareAnthropicMessagesStream, @@ -61,9 +60,8 @@ from litellm.router import ( _anthropic_stream_should_decline_fallback, _is_retriable_anthropic_status, _responses_stream_holds_event, - _without_line_breaks, - _search_tool_support, -) +) # noqa: PLC2701 +from litellm.router_rust.surface.registry import _search_tool_support # noqa: PLC2701 from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request from litellm.router_strategy import simple_shuffle from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter @@ -98,7 +96,7 @@ from litellm.types.router import ( RoutingContext, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.llms.base_llm.vector_store.transformation import( +from litellm.llms.base_llm.vector_store.transformation import ( LiteLLMVectorStoreEmbeddingExecutor, RouterVectorStoreEmbeddingExecutor, ) @@ -248,146 +246,6 @@ async def test_chat_on_a_mixed_group_still_reaches_the_leg_without_decisions_sup assert response._hidden_params["model_id"] == "bedrock-leg" -def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): - # initialize a real Router (env‑vars can be empty) - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_AI_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - } - ], - ) - - # override to known defaults for the test - router.default_litellm_params = { - "foo": "bar", - "metadata": {"baz": 123}, - } - original = copy.deepcopy(router.default_litellm_params) - kwargs: dict = {} - - # invoke the helper - router._update_kwargs_with_default_litellm_params( - kwargs=kwargs, - metadata_variable_name="litellm_metadata", - ) - - # 1) router.defaults must be unchanged - assert router.default_litellm_params == original - - # 2) non‑metadata keys get merged - assert kwargs["foo"] == "bar" - - # 3) metadata lands under "metadata" - assert kwargs["litellm_metadata"] == {"baz": 123} - - -def test_router_model_group_llm_provider_resolves_chatgpt_without_login(tmp_path, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - params: Final = LiteLLM_Params(model="chatgpt/gpt-5.4") - - assert Router._model_group_llm_provider(params) == ("gpt-5.4", "chatgpt") - assert Router._model_group_llm_provider(LiteLLM_Params(model="no-such-provider-model-xyz")) == ("", "") - assert not any(tmp_path.iterdir()) - - -def test_router_with_model_info_and_model_group(): - """ - Test edge case where user specifies model_group in model_info - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - }, - "model_info": { - "tpm": 1000, - "rpm": 1000, - "model_group": "gpt-3.5-turbo", - }, - } - ], - ) - - router._set_model_group_info( - model_group="gpt-3.5-turbo", - user_facing_model_group_name="gpt-3.5-turbo", - ) - - -def test_router_model_group_encrypted_content_affinity_callback_registration(): - from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( - DeploymentAffinityCheck, - ) - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, - ) - - model_group = "openai.gpt-5.1-codex" - model_group_affinity_config = { - model_group: ["encrypted_content_affinity"], - } - original_callbacks = list(litellm.callbacks) - litellm.callbacks = [] - router = None - - try: - router = litellm.Router( - model_list=[ - { - "model_name": model_group, - "litellm_params": { - "model": "openai/gpt-5.1-codex", - "api_key": "mock-api-key", - }, - } - ], - model_group_affinity_config=model_group_affinity_config, - num_retries=0, - ) - callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - assert len(encrypted_content_callbacks) == 1 - assert encrypted_content_callbacks[0].enable_global_affinity is False - assert ( - encrypted_content_callbacks[0].model_group_affinity_config - == model_group_affinity_config - ) - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( - litellm.callbacks.index(deployment_callback) - ) - - router._add_encrypted_content_affinity_check(enable_global_affinity=True) - - callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] - assert len(encrypted_content_callbacks) == 1 - assert encrypted_content_callbacks[0].enable_global_affinity is True - assert encrypted_content_callbacks[0].ctx is router._core_ctx - finally: - if router is not None: - router.discard() - litellm.callbacks = original_callbacks - - @pytest.mark.asyncio async def test_encrypted_content_affinity_model_group_config_is_additive(): from litellm.responses.utils import ResponsesAPIRequestUtils @@ -409,13 +267,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - {model_group: ["encrypted_content_affinity"]} - ) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -456,10 +310,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert ( - "encrypted_content_affinity_enabled" - not in disabled_request_kwargs["litellm_metadata"] - ) + assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -479,9 +330,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"][ - "encrypted_content_affinity_enabled" - ] + assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] @pytest.mark.asyncio @@ -528,18 +377,10 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callback) < ( - litellm.callbacks.index(deployment_callback) - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -550,9 +391,7 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -762,7 +601,9 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards from io import BytesIO from unittest.mock import MagicMock, patch - jsonl_content = b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + jsonl_content = ( + b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + ) router = litellm.Router( model_list=[ { @@ -773,9 +614,7 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards ) with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: - await router.acreate_file( - model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True - ) + await router.acreate_file(model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True) forwarded = mock_acreate_file.call_args.kwargs assert forwarded["passthrough"] is True forwarded["file"].seek(0) @@ -1251,9 +1090,7 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object( - litellm, "aretrieve_batch", return_value=AsyncMock() - ) as mock_aretrieve_batch: + with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -1513,6 +1350,7 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( == expected_successes ) + _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") @@ -1524,8 +1362,7 @@ async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]: moved = sorted( f"{key}={cache_dict[key]}" for key in cache_dict - if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) - and cache_dict[key] + if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) and cache_dict[key] ) if moved: return moved @@ -1606,9 +1443,7 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object( - litellm, "afile_content", return_value=AsyncMock() - ) as mock_afile_content: + with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -1669,7 +1504,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: + with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1753,9 +1588,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team_id and team_public_model_name match" + assert result is True, "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -1763,9 +1596,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is False - ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" + assert result is False, ( + "Should return False when team_id matches but model_name doesn't match team_public_model_name" + ) # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -1781,30 +1614,18 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team deployment has no team_public_model_name to match" + assert result is True, "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id=None - ) - assert ( - result is True - ), "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) + assert result is True, "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id="any-team" - ) - assert ( - result is True - ), "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") + assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment( - model_name="different-model", model=deployment_without_team, team_id=None - ) + result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -1813,9 +1634,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert ( - result is True - ), "Should return True when matching model with exact model_name" + assert result is True, "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -1866,9 +1685,7 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = ( - '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' - ) + mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1942,7 +1759,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: + with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1967,7 +1784,9 @@ def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str: model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="served by onprem handler" ) - monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}]) + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}] + ) monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) return "test-onprem-llm" @@ -2006,327 +1825,6 @@ def test_router_add_deployment_accepts_explicit_custom_provider_from_custom_prov assert router.get_model_list(model_name="onprem")[0]["litellm_params"]["model"] == "my-model" -@pytest.mark.asyncio -async def test_router_ageneric_api_call_with_fallbacks_helper(): - """ - Test the _ageneric_api_call_with_fallbacks_helper method with various scenarios - """ - from unittest.mock import patch - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": "test-key", - "api_base": "https://api.openai.com/v1", - }, - "model_info": { - "tpm": 1000, - "rpm": 1000, - }, - }, - ], - ) - - # Test 1: Successful call - async def mock_generic_function(**kwargs): - return {"result": "success", "model": kwargs.get("model")} - - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.return_value = { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": "test-key", - "api_base": "https://api.openai.com/v1", - }, - } - - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: - result = await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_generic_function, - messages=[{"role": "user", "content": "test"}], - ) - - assert result is not None - assert result["result"] == "success" - mock_get_deployment.assert_called_once() - mock_update_kwargs.assert_called_once() - mock_pre_call_checks.assert_called_once() - - # Test 2: Passthrough on no deployment (success case) - async def mock_passthrough_function(**kwargs): - return {"result": "passthrough", "model": kwargs.get("model")} - - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.side_effect = Exception("No deployment available") - - result = await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_passthrough_function, - passthrough_on_no_deployment=True, - messages=[{"role": "user", "content": "test"}], - ) - - assert result is not None - assert result["result"] == "passthrough" - assert result["model"] == "gpt-3.5-turbo" - - # Test 3: No deployment available and passthrough=False (should raise exception) - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.side_effect = Exception("No deployment available") - - with pytest.raises(Exception, match='No deployment available') as exc_info: - await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_generic_function, - passthrough_on_no_deployment=False, - messages=[{"role": "user", "content": "test"}], - ) - - assert "No deployment available" in str(exc_info.value) - - # Test 4: Test with semaphore (rate limiting) - import asyncio - - async def mock_semaphore_function(**kwargs): - return {"result": "semaphore_success", "model": kwargs.get("model")} - - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.return_value = { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": "test-key", - "api_base": "https://api.openai.com/v1", - }, - } - - mock_semaphore = MaxParallelRequestsLimit( - max_parallel_requests=1, model_id="deployment-1", model_group="gpt-3.5-turbo" - ) - - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=mock_semaphore - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - result = await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_semaphore_function, - messages=[{"role": "user", "content": "test"}], - ) - - assert result is not None - assert result["result"] == "semaphore_success" - mock_get_client.assert_called_once() - mock_pre_call_checks.assert_called_once() - - # Test 5: Test call tracking (success and failure counts) - initial_success_count = router.success_calls.get("gpt-3.5-turbo", 0) - initial_fail_count = router.fail_calls.get("gpt-3.5-turbo", 0) - - async def mock_failing_function(**kwargs): - raise Exception("Mock failure") - - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.return_value = { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": "test-key", - "api_base": "https://api.openai.com/v1", - }, - } - - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with pytest.raises(Exception, match='Mock failure') as exc_info: - await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_failing_function, - messages=[{"role": "user", "content": "test"}], - ) - - assert "Mock failure" in str(exc_info.value) - # Check that fail_calls was incremented - assert router.fail_calls["gpt-3.5-turbo"] == initial_fail_count + 1 - - -@pytest.mark.asyncio -async def test_ageneric_api_call_deployment_model_overrides_alias(): - """ - Regression: when a model alias (e.g. "not-gemini-2.5-flash") maps to a deployment - with model="vertex_ai/gemini-2.5-flash", the underlying litellm function must receive - the deployment model, not the alias. Before the fix, **kwargs overwrote data["model"]. - """ - from unittest.mock import patch - - captured: dict = {} - - async def capture_model(**kwargs): - captured["model"] = kwargs.get("model") - return {"result": "ok"} - - router = litellm.Router( - model_list=[ - { - "model_name": "not-gemini-2.5-flash", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-flash", - "api_key": "fake-key", - }, - } - ] - ) - - def inject_alias_into_kwargs( - deployment: dict, - kwargs: dict, - function_name: str | None = None, - splits_stream_timeout: bool = False, - opens_responses_socket: bool = False, - ) -> None: - # Simulate the alias leaking into kwargs (as happens when - # _ageneric_api_call_with_fallbacks sets kwargs["model"] = alias before - # calling the helper through async_function_with_fallbacks). - kwargs["model"] = "not-gemini-2.5-flash" - - with ( - patch.object(router, "async_get_available_deployment") as mock_dep, - patch.object( - router, - "_update_kwargs_with_deployment", - side_effect=inject_alias_into_kwargs, - ), - patch.object(router, "async_routing_strategy_pre_call_checks"), - patch.object(router, "_get_client", return_value=None), - ): - mock_dep.return_value = { - "model_name": "not-gemini-2.5-flash", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-flash", - "api_key": "fake-key", - }, - } - - await router._ageneric_api_call_with_fallbacks_helper( - model="not-gemini-2.5-flash", - original_generic_function=capture_model, - ) - - assert ( - captured["model"] == "vertex_ai/gemini-2.5-flash" - ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" - - -@pytest.mark.asyncio -async def test_ageneric_api_call_resolves_realtime_session_model(): - """ - Regression for #36742: realtime client secret requests carry the model inside `session` too, and the proxy - fills it with the pre-routing model group name. The underlying litellm function reads session.model first, - so it must see the resolved deployment, while a caller's nested transcription model stays untouched. - """ - routed: Final = AsyncMock(return_value={"result": "ok"}) - - router = litellm.Router( - model_list=[ - { - "model_name": "my-realtime-group", - "litellm_params": { - "model": "openai/gpt-realtime-2.1-mini", - "api_key": "fake-key", - }, - "model_info": {"mode": "realtime"}, - } - ] - ) - - await router._ageneric_api_call_with_fallbacks( - model="my-realtime-group", - original_function=routed, - session={ - "type": "realtime", - "model": "my-realtime-group", - "audio": {"input": {"transcription": {"model": "gpt-4o-transcribe"}}}, - }, - ) - - sent: Final = routed.call_args.kwargs - assert sent["model"] == "openai/gpt-realtime-2.1-mini" - assert sent["session"]["model"] == "openai/gpt-realtime-2.1-mini" - assert sent["session"]["audio"]["input"]["transcription"]["model"] == "gpt-4o-transcribe" - - -@pytest.mark.asyncio -async def test_ageneric_api_call_does_not_add_session_model(): - """ - A session that never carried a model must not gain one from routing: the underlying function then falls back - to the resolved `model` kwarg itself, and the outgoing session body keeps the caller's shape. - """ - routed: Final = AsyncMock(return_value={"result": "ok"}) - - router = litellm.Router( - model_list=[ - { - "model_name": "my-realtime-group", - "litellm_params": { - "model": "openai/gpt-realtime-2.1-mini", - "api_key": "fake-key", - }, - "model_info": {"mode": "realtime"}, - } - ] - ) - - await router._ageneric_api_call_with_fallbacks( - model="my-realtime-group", - original_function=routed, - session={"type": "realtime"}, - ) - - sent: Final = routed.call_args.kwargs - assert sent["model"] == "openai/gpt-realtime-2.1-mini" - assert sent["session"] == {"type": "realtime"} - - -@pytest.mark.parametrize( - "session, expected", - [ - ({"type": "realtime", "model": "my-realtime-group"}, {"session": {"type": "realtime", "model": "resolved"}}), - ({"type": "realtime"}, {}), - (None, {}), - ("not-a-session", {}), - ], -) -def test_with_router_resolved_session_model(session, expected): - from litellm.router import _with_router_resolved_session_model - - assert dict(_with_router_resolved_session_model(session, "resolved")) == expected - - def test_router_get_model_access_groups_team_only_models(): """ Test that Router._core_get_model_access_groups returns the correct response for team-only models @@ -2345,147 +1843,13 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id=None - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id="team_1" - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") assert list(access_groups.keys()) == ["default-models"] -def test_cached_get_model_group_info(): - """ - Test that _cached_get_model_group_info caches results and - invalidates on deployment changes. - """ - from litellm.types.router import Deployment, LiteLLM_Params - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake"}, - "model_info": {"tpm": 1000, "rpm": 100}, - }, - ] - ) - - # First call should compute and cache - result1 = router._cached_get_model_group_info("gpt-4") - assert result1 is not None - assert result1.tpm == 1000 - - # Second call should hit cache (same object) - result2 = router._cached_get_model_group_info("gpt-4") - assert result1 is result2 - - # Add a deployment — cache should be invalidated - router.add_deployment( - Deployment( - model_name="gpt-4", - litellm_params=LiteLLM_Params(model="gpt-4", api_key="fake2"), - model_info={"tpm": 2000, "rpm": 200}, - ) - ) - result3 = router._cached_get_model_group_info("gpt-4") - assert result3 is not result2 - assert result3 is not None - assert result3.tpm == 3000 # 1000 + 2000 - - # Delete a deployment — cache should be invalidated - deployment_id = router.model_list[-1]["model_info"]["id"] - router.delete_deployment(id=deployment_id) - result4 = router._cached_get_model_group_info("gpt-4") - assert result4 is not result3 - assert result4 is not None - assert result4.tpm == 1000 - - # set_model_list — cache should be invalidated - router.set_model_list( - [ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake"}, - "model_info": {"tpm": 5000}, - }, - ] - ) - result5 = router._cached_get_model_group_info("gpt-4") - assert result5 is not result4 - assert result5 is not None - assert result5.tpm == 5000 - - # Verify cache still works after invalidation - result6 = router._cached_get_model_group_info("gpt-4") - assert result5 is result6 - - -def test_model_group_info_cost_from_db_model_info(): - """ - When get_deployment_model_info fails (model_info is None fallback), - input_cost_per_token and output_cost_per_token should be read from db model_info. - """ - from unittest.mock import patch - - router = litellm.Router( - model_list=[ - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-model", - "api_key": "fake", - "api_base": "https://my-custom-endpoint.com", - }, - "model_info": { - "input_cost_per_token": 0.0001, - "output_cost_per_token": 0.0002, - }, - }, - ] - ) - - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): - result = router._cached_get_model_group_info("my-custom-model") - assert result is not None - assert result.input_cost_per_token == 0.0001 - assert result.output_cost_per_token == 0.0002 - - -def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): - """ - When get_deployment_model_info fails and db model_info has no cost fields, - input/output_cost_per_token should be None. - """ - from unittest.mock import patch - - router = litellm.Router( - model_list=[ - { - "model_name": "my-custom-model-no-cost", - "litellm_params": { - "model": "openai/my-custom-model-no-cost", - "api_key": "fake", - "api_base": "https://my-custom-endpoint.com", - }, - "model_info": {}, - }, - ] - ) - - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): - result = router._cached_get_model_group_info("my-custom-model-no-cost") - assert result is not None - assert result.input_cost_per_token is None - assert result.output_cost_per_token is None - - def test_model_group_info_cost_none_for_unpriced_deployment_but_zero_when_declared(): """A deployment with no cost fields anywhere must report None, not the 0 that get_model_info defaults to, so the reported price matches what the zero-cost @@ -2881,153 +2245,6 @@ def test_get_model_list_of_routed_group_falls_back_to_the_wildcard_route_serving assert _served_models(router.get_model_list(model_name="openai/gpt-4o")) == ["ollama/qwen3:0.6b"] -def test_switch_routing_strategy_installs_lar1_then_restores_the_default_selector(): - router = _alias_cost_router() - - router._switch_routing_strategy( - "lar1", - { - "routing_strategy_args": { - "confidence_threshold_low": 0.1, - "confidence_threshold_medium": 0.3, - "confidence_threshold_high": 0.9, - } - }, - ) - assert router.routing_strategy == "lar1" - assert "_core_async_get_available_deployment" in router._core.__dict__ - - router._switch_routing_strategy("usage-based-routing-v2", {}) - assert router.lowesttpm_logger_v2 is not None - assert "_core_async_get_available_deployment" not in router._core.__dict__ - - -@pytest.mark.parametrize( - "value,expected", - [ - ("1e-05", 1e-05), - ("0.00001", 1e-05), - (1e-05, 1e-05), - (5, 5.0), - (None, None), - ("not-a-number", None), - ], -) -def test_cost_value_as_float(value, expected): - from litellm.router import _cost_value_as_float - - assert _cost_value_as_float(value) == expected - - -def test_model_group_info_with_stringified_cost_values(): - """ - YAML 1.2 parsers emit '1e-05' (integer mantissa) as a string, so cost - values in deployment model_info can arrive as str. Aggregating the model - group must not raise TypeError('>' between str and float) and must return - float costs. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-backend-1", - "api_key": "fake", - }, - "model_info": { - "input_cost_per_token": "1e-05", - "output_cost_per_token": "1e-05", - }, - }, - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-backend-2", - "api_key": "fake", - }, - "model_info": { - "input_cost_per_token": "2e-05", - "output_cost_per_token": "2e-05", - }, - }, - ] - ) - - def _model_info_with_str_costs(model_id: str, model_name: str): - for model in router.model_list: - if model["model_info"]["id"] == model_id: - return { - "key": model_name, - "input_cost_per_token": model["model_info"]["input_cost_per_token"], - "output_cost_per_token": model["model_info"]["output_cost_per_token"], - "litellm_provider": "openai", - "mode": "chat", - } - return None - - with patch.object( - router, "_core_get_deployment_model_info", side_effect=_model_info_with_str_costs - ): - result = router._set_model_group_info( - model_group="my-custom-model", - user_facing_model_group_name="my-custom-model", - ) - - assert result is not None - assert result.input_cost_per_token == 2e-05 - assert result.output_cost_per_token == 2e-05 - assert isinstance(result.input_cost_per_token, float) - assert isinstance(result.output_cost_per_token, float) - - -def test_model_group_info_db_fallback_with_stringified_cost_values(): - """ - Fallback path: when get_deployment_model_info returns nothing, costs are - read straight from the deployment's model_info dict, which can hold - stringified floats parsed from YAML. They must be coerced to float. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-backend-1", - "api_key": "fake", - }, - "model_info": { - "input_cost_per_token": "1e-05", - "output_cost_per_token": "3e-05", - }, - }, - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-backend-2", - "api_key": "fake", - }, - "model_info": { - "input_cost_per_token": "2e-05", - "output_cost_per_token": "2e-05", - }, - }, - ] - ) - - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): - result = router._set_model_group_info( - model_group="my-custom-model", - user_facing_model_group_name="my-custom-model", - ) - - assert result is not None - assert result.input_cost_per_token == 2e-05 - assert result.output_cost_per_token == 3e-05 - assert isinstance(result.input_cost_per_token, float) - assert isinstance(result.output_cost_per_token, float) - - def test_get_model_access_groups_caching(): """ Test that get_model_access_groups caches the no-args result @@ -3154,502 +2371,12 @@ def test_get_model_access_groups_cache_invalidation_upsert_deployment(): assert "updated-group" in result2 -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator(): - """Test _acompletion_streaming_iterator for normal streaming and fallback behavior.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - # Helper class for creating async iterators - class AsyncIterator: - def __init__(self, items, error_after=None): - self.items = items - self.index = 0 - self.error_after = error_after - - def __aiter__(self): - return self - - async def __anext__(self): - if self.error_after is not None and self.index >= self.error_after: - raise self.error_after - if self.index >= len(self.items): - raise StopAsyncIteration - item = self.items[self.index] - self.index += 1 - return item - - # Set up router with fallback configuration - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, - }, - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key-2"}, - }, - ], - fallbacks=[{"gpt-4": ["gpt-3.5-turbo"]}], - set_verbose=True, - ) - - # Test data - messages = [{"role": "user", "content": "Hello"}] - initial_kwargs = {"model": "gpt-4", "stream": True, "temperature": 0.7} - - # Test 1: Successful streaming (no errors) - print("\n=== Test 1: Successful streaming ===") - - # Mock successful streaming response - mock_chunks = [ - MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content=" there"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))]), - ] - - mock_response = AsyncIterator(mock_chunks) - - setattr(mock_response, "model", "gpt-4") - setattr(mock_response, "custom_llm_provider", "openai") - setattr(mock_response, "logging_obj", MagicMock()) - - result = await router._acompletion_streaming_iterator( - model_response=mock_response, messages=messages, initial_kwargs=initial_kwargs - ) - - # Collect streamed chunks - collected_chunks = [] - async for chunk in result: - collected_chunks.append(chunk) - - assert len(collected_chunks) == 3 - assert all(chunk in mock_chunks for chunk in collected_chunks) - print("✓ Successfully streamed all chunks") - - # Test 2: MidStreamFallbackError with generated content is re-raised, not silently continued - print("\n=== Test 2: MidStreamFallbackError re-raises when content already generated ===") - - # Error with generated content and is_pre_first_chunk=False (the default): - # the router must re-raise instead of attempting a continuation-prompt fallback, - # because partial content has already been sent to the client. - error = MidStreamFallbackError( - message="Connection lost", - model="gpt-4", - llm_provider="openai", - generated_content="Hello", - ) - - class AsyncIteratorWithError: - def __init__(self, items, error_after_index): - self.items = items - self.index = 0 - self.error_after_index = error_after_index - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index >= len(self.items): - raise StopAsyncIteration - if self.index == self.error_after_index: - raise error - item = self.items[self.index] - self.index += 1 - return item - - mock_error_response = AsyncIteratorWithError(mock_chunks, 1) # Error after first chunk - - setattr(mock_error_response, "model", "gpt-4") - setattr(mock_error_response, "custom_llm_provider", "openai") - setattr(mock_error_response, "logging_obj", MagicMock()) - - result = await router._acompletion_streaming_iterator( - model_response=mock_error_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - # Collect streamed chunks — the first chunk succeeds, then the error re-raises - collected_chunks = [] - async def _drain(): - async for chunk in result: - collected_chunks.append(chunk) - - with pytest.raises(MidStreamFallbackError): - await _drain() - - assert len(collected_chunks) == 1, "one chunk yielded before the error" - print("✓ MidStreamFallbackError re-raised correctly when content was already generated") - - print("\n=== All tests passed! ===") - - -def _make_midstream_source(error, chunks=None, first_chunk=None): - from unittest.mock import MagicMock - - opening_chunk = ( - first_chunk if first_chunk is not None else MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]) - ) - - class _Source: - def __init__(self): - self.index = 0 - self.chunks = chunks if chunks is not None else [] - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index == 0: - self.index += 1 - return opening_chunk - raise error - - async def wait_for_failure_logging(self, timeout): - return None - - return _Source() - - -class _FakeFallbackStream: - def __init__(self, item): - self._item = item - self._done = False - - def __aiter__(self): - return self - - async def __anext__(self): - if self._done: - raise StopAsyncIteration - self._done = True - return self._item - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_continues_after_content_when_eligible(): - """Flag on + plain-text break: the router re-enters the fallback chain with an - assistant-prefill continuation and the mid-stream marker, then streams the - fallback's output instead of re-raising.""" - from unittest.mock import AsyncMock, MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_KWARG, - MID_STREAM_CONTINUATION_MARKER, - ) - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - - error = MidStreamFallbackError( - message="Connection lost", - model="gpt-4", - llm_provider="openai", - generated_content="Hello", - is_pre_first_chunk=False, - emitted_disqualifying_content=False, - ) - source = _make_midstream_source(error) - fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}}]) - initial_kwargs = {"model": "gpt-4", "stream": True} - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), - ) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs - ) - collected = [chunk async for chunk in result] - - mock_fallback.assert_awaited_once() - passed_kwargs = mock_fallback.await_args.kwargs["kwargs"] - assert passed_kwargs[MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER - assert passed_kwargs["messages"][-1] == {"role": "assistant", "content": "Hello", "prefix": True} - assert fallback_chunk in collected - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "continuation_text,joined_output_trips_rule", - [("6789", True), ("6788", False)], - ids=["rule_trips_only_on_the_joined_text", "clean_continuation_streams"], -) -async def test_acompletion_streaming_iterator_continuation_runs_post_call_rules_on_the_joined_text( - monkeypatch, continuation_text, joined_output_trips_rule -): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - monkeypatch.setattr(litellm, "post_call_rules", [lambda output: "123-45-6789" not in output]) - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="123-45-", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": continuation_text}}]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), - ): - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - if joined_output_trips_rule: - with pytest.raises(litellm.APIResponseValidationError): - async for _ in result: - pass - else: - assert fallback_chunk in [chunk async for chunk in result] - - -@pytest.mark.asyncio -async def test_a_second_cut_runs_post_call_rules_only_on_the_text_the_caller_received(monkeypatch): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_KWARG, - MID_STREAM_CONTINUATION_MARKER, - ) - - checked_texts = [] - monkeypatch.setattr(litellm, "post_call_rules", [lambda output: checked_texts.append(output) is None]) - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - - def cut(generated_content): - return MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=generated_content, - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - - messages = [{"role": "user", "content": "Hi"}] - middle_leg = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source( - cut(" there"), - first_chunk=litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " there"}}]), - ), - messages=[*messages, {"role": "assistant", "content": "Hello", "prefix": True}], - initial_kwargs={"model": "gpt-4", "stream": True, MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, - ) - backup_leg = _FakeFallbackStream( - litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " friend"}}]) - ) - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(side_effect=[middle_leg, backup_leg]) - ): - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(cut("Hello")), - messages=messages, - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - _ = [chunk async for chunk in result] - - assert checked_texts == ["Hello there", "Hello there friend"] - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallback_texts", - [("Hel", "lo", " world"), (" world",)], - ids=["fallback_echoes_the_prefill", "fallback_continues_after_the_prefill"], -) -async def test_acompletion_streaming_iterator_continuation_never_repeats_the_partial_output(fallback_texts): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-small-latest", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - - async def fallback_stream(): - for text in fallback_texts: - yield litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) - ): - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - collected = [chunk async for chunk in result] - - continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] - assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world" - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_continuation_stays_in_the_primary_stream(): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - primary_chunk = litellm.ModelResponseStream( - id="chatcmpl-primary", created=111, choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hello"}}] - ) - - async def fallback_stream(): - yield litellm.ModelResponseStream( - id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"role": "assistant", "content": " wor"}}] - ) - yield litellm.ModelResponseStream( - id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"content": "ld"}}] - ) - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) - ): - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error, chunks=[primary_chunk]), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - collected = [chunk async for chunk in result] - - continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] - assert [(chunk.id, chunk.created, chunk.choices[0].delta.role) for chunk in continuation] == [ - ("chatcmpl-primary", 111, None), - ("chatcmpl-primary", 111, None), - ] - assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world" - - -def _continuation_router(): - return litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - - -def _usage_chunk(prompt_tokens, completion_tokens): - from litellm.types.utils import Usage - - return litellm.ModelResponseStream( - model="gpt-4", - choices=[], - usage=Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ), - ) - - -async def _drain_with_fallback(router, error, primary_chunks, messages, fallback_items): - from unittest.mock import AsyncMock, patch - - async def fallback_stream(): - for item in fallback_items: - yield item - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) - ) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error, chunks=primary_chunks), - messages=messages, - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - collected = [chunk async for chunk in result] - streamed = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] - return mock_fallback.await_args.kwargs["kwargs"], streamed - - -@pytest.mark.asyncio -@pytest.mark.parametrize("primary_reports_usage", [True, False], ids=["usage_reported", "usage_not_reported"]) -async def test_acompletion_streaming_iterator_continuation_usage_sums_both_legs(primary_reports_usage): - from litellm.exceptions import MidStreamFallbackError - - messages = [{"role": "user", "content": "Explain transformers in one paragraph."}] - partial = "Transformers use " - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - content_chunk = litellm.ModelResponseStream( - model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": partial}}] - ) - primary_chunks = [content_chunk, _usage_chunk(21, 3)] if primary_reports_usage else [content_chunk] - primary_usage = litellm.stream_chunk_builder(chunks=list(primary_chunks), messages=messages).usage - - _, streamed = await _drain_with_fallback( - _continuation_router(), error, primary_chunks, messages, [_usage_chunk(91, 5)] - ) - - assert primary_usage.prompt_tokens > 0 - assert streamed[-1].usage.prompt_tokens == primary_usage.prompt_tokens + 91 - assert streamed[-1].usage.completion_tokens == primary_usage.completion_tokens + 5 - - class _CutSSE(httpx.AsyncByteStream): def __init__(self, content: str) -> None: self._content: Final = content async def __aiter__(self) -> AsyncIterator[bytes]: - yield f'data: {json.dumps({"id": "cut", "object": "chat.completion.chunk", "created": 1700000000, "model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": self._content}, "finish_reason": None}]})}\n\n'.encode() + yield f"data: {json.dumps({'id': 'cut', 'object': 'chat.completion.chunk', 'created': 1700000000, 'model': 'm', 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': self._content}, 'finish_reason': None}]})}\n\n".encode() raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") @@ -3718,7 +2445,9 @@ async def test_continued_stream_success_log_bills_every_cut_leg( recorder: Final = _SpendRecorder() monkeypatch.setattr(litellm, "callbacks", [recorder]) respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use") + ) ) respx_mock.post(host="api.deepseek.com").mock( return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE(" learned")) @@ -3801,7 +2530,9 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t recorder: Final = _SlowFailureRecorder() monkeypatch.setattr(litellm, "callbacks", [recorder]) respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use") + ) ) respx_mock.post(host="api.mistral.ai").mock( return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE) @@ -3874,7 +2605,11 @@ _AGENTIC_FIRST_CALL_BODY: Final = { "role": "assistant", "content": None, "tool_calls": [ - {"id": "call_abc", "type": "function", "function": {"name": "litellm_code_execution", "arguments": "{}"}} + { + "id": "call_abc", + "type": "function", + "function": {"name": "litellm_code_execution", "arguments": "{}"}, + } ], }, } @@ -4165,7 +2900,10 @@ async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fall 200, headers={"content-type": "text/event-stream"}, text=_responses_sse( - {"type": "response.created", "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress")}, + { + "type": "response.created", + "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress"), + }, { "type": "response.completed", "response": _responses_envelope( @@ -4235,420 +2973,6 @@ async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fall ) -def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherited(): - from litellm.types.utils import Usage - - cut_logging: Final = LiteLLMLogging( - model="deepseek/deepseek-chat", messages=[], stream=True, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="cut", function_id="cut", - ) - cut_logging.model_call_details["response_cost"] = 0.5 - cut_logging.update_environment_variables( - litellm_params={"custom_llm_provider": "deepseek"}, optional_params={}, model="deepseek/deepseek-chat" - ) - cut_logging.record_continued_partial( - usage=Usage(prompt_tokens=40, completion_tokens=7, total_tokens=47), - response_cost=0.001, - response_id="chatcmpl-primary", - ) - fallback_logging: Final = LiteLLMLogging( - model="mistral-large-latest", messages=[], stream=True, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="fallback", function_id="fallback", - ) - - Router._carry_partial_spend_to_fallback( - SimpleNamespace(logging_obj=fallback_logging), - Router._price_cut_partial( - cast(CustomStreamWrapper, SimpleNamespace(logging_obj=cut_logging, model="deepseek/deepseek-chat")), - litellm.ModelResponse( - id="chatcmpl-relay", usage=Usage(prompt_tokens=60, completion_tokens=3, total_tokens=63) - ), - ), - ) - - relay_prompt_cost, relay_completion_cost = litellm.cost_per_token( - model="deepseek/deepseek-chat", prompt_tokens=60, completion_tokens=3 - ) - carried_usage: Final = fallback_logging.model_call_details["continued_partial_usage"] - assert (carried_usage.prompt_tokens, carried_usage.completion_tokens) == (100, 10) - assert fallback_logging.model_call_details["continued_partial_cost"] == pytest.approx( - relay_prompt_cost + relay_completion_cost + 0.001 - ) - assert fallback_logging.model_call_details["continued_response_id"] == "chatcmpl-primary" - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_plain_fallback_keeps_the_fallback_usage(): - from litellm.exceptions import MidStreamFallbackError - - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="anthropic", generated_content="", - is_pre_first_chunk=True, - ) - role_chunk = litellm.ModelResponseStream( - model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": ""}}] - ) - - _, streamed = await _drain_with_fallback( - _continuation_router(), - error, - [role_chunk, _usage_chunk(21, 0)], - [{"role": "user", "content": "Hi"}], - [_usage_chunk(91, 5)], - ) - - assert (streamed[-1].usage.prompt_tokens, streamed[-1].usage.completion_tokens) == (91, 5) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "partial,fallback_texts,expected_answer", - [ - ("Transformers use ", (" self", "-attention"), "Transformers use self-attention"), - ("def f():\n ", ("\n", " return 1"), "def f():\n return 1"), - ("First.\n", ("\n\nSecond",), "First.\n\nSecond"), - ], - ids=["space", "indentation", "paragraph_break"], -) -async def test_acompletion_streaming_iterator_continuation_does_not_double_trailing_whitespace( - partial, fallback_texts, expected_answer -): - from litellm.exceptions import MidStreamFallbackError - - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - fallback_items = [ - litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) for text in fallback_texts - ] - - passed_kwargs, streamed = await _drain_with_fallback( - _continuation_router(), error, [], [{"role": "user", "content": "Hi"}], fallback_items - ) - - assert passed_kwargs["messages"][-1]["content"] == partial - assert partial + "".join(chunk.choices[0].delta.content for chunk in streamed) == expected_answer - - -@pytest.mark.asyncio -@pytest.mark.parametrize("enabled", [True, "true", " True "]) -async def test_update_settings_turns_on_mid_stream_continuation(monkeypatch, enabled): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_KWARG, - MID_STREAM_CONTINUATION_MARKER, - ContinuationPrefillDeploymentCheck, - ) - - monkeypatch.setattr(litellm, "callbacks", []) - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - ) - router.update_settings(enable_mid_stream_fallback_continuation=enabled) - - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), - ) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - collected = [chunk async for chunk in result] - - assert mock_fallback.await_args.kwargs["kwargs"][MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER - assert fallback_chunk in collected - assert any(isinstance(callback, ContinuationPrefillDeploymentCheck) for callback in litellm.callbacks) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "error_kwargs,request_kwargs", - [ - ({"emitted_disqualifying_content": True}, {}), - ({"emitted_disqualifying_content": False}, {"response_format": {"type": "json_object"}}), - ({"emitted_disqualifying_content": False}, {"tool_choice": "required"}), - ({"emitted_disqualifying_content": False}, {"merge_reasoning_content_in_choices": True}), - ({"emitted_disqualifying_content": False}, {"n": 2}), - ], - ids=["tool_or_thinking_emitted", "json_mode", "forced_tool_choice", "merged_reasoning", "multiple_choices"], -) -async def test_acompletion_streaming_iterator_declines_ineligible_after_content(error_kwargs, request_kwargs): - """Flag on but the break is not continuation-safe: the router re-raises and - never enters the fallback chain, so no duplicated/rejected request is sent.""" - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, **error_kwargs, - ) - source = _make_midstream_source(error) - initial_kwargs = {"model": "gpt-4", "stream": True, **request_kwargs} - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs - ) - with pytest.raises(MidStreamFallbackError): - async for _ in result: - pass - - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_raises_content_filter_instead_of_continuing(): - """A content filter that stops a stream mid-answer reaches the client as ContentPolicyViolationError, - never as a continuation that lets a backup finish the filtered answer.""" - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "azure/gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - content_filter = litellm.ContentPolicyViolationError(message="filtered", model="gpt-4", llm_provider="azure") - error = MidStreamFallbackError( - message="filtered", model="gpt-4", llm_provider="azure", generated_content="Step one is to obtain the", - is_pre_first_chunk=False, emitted_disqualifying_content=False, original_exception=content_filter, - ) - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - with pytest.raises(litellm.ContentPolicyViolationError): - async for _ in result: - pass - - mock_fallback.assert_not_awaited() - - -async def _continues_after_content( - initial_kwargs: dict[str, object], - primary_params: Mapping[str, object] = MappingProxyType({}), - default_litellm_params: dict[str, object] | None = None, - cut_model_id: str = "primary-deployment", -) -> bool: - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router: Final = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "k1", **primary_params}, - "model_info": {"id": "primary-deployment"}, - }, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - { - "model_name": "openai-backup", - "litellm_params": {"model": "gpt-4o", "api_key": "k3"}, - "model_info": {"id": "no-prefill-deployment", "supports_assistant_prefill": False}, - }, - { - "model_name": "claude-backup", - "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k4"}, - "model_info": {"id": "prefill-deployment", "supports_assistant_prefill": True}, - }, - { - "model_name": "json-claude-backup", - "litellm_params": { - "model": "anthropic/claude-3-opus-20240229", - "api_key": "k5", - "response_format": {"type": "json_object"}, - }, - "model_info": {"id": "json-prefill-deployment", "supports_assistant_prefill": True}, - }, - ], - fallbacks=[{"gpt-4": ["backup"]}], - default_litellm_params=default_litellm_params, - enable_mid_stream_fallback_continuation=True, - ) - error: Final = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - fallback_chunk: Final = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) - cut_stream: Final = _make_midstream_source(error) - cut_stream._hidden_params = {"model_id": cut_model_id} - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), - ) as mock_fallback: - result: Final = await router._acompletion_streaming_iterator( - model_response=cut_stream, - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True, **initial_kwargs}, - ) - with contextlib.suppress(MidStreamFallbackError): - _ = [chunk async for chunk in result] - return mock_fallback.await_count == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "event_hook,continues", - [("post_call", False), (None, False), ("pre_call", True)], - ids=["post_call_guardrail", "every_hook_guardrail", "input_only_guardrail"], -) -async def test_mid_stream_continuation_declines_while_an_output_guardrail_is_registered( - monkeypatch: pytest.MonkeyPatch, event_hook: str | None, continues: bool -) -> None: - from litellm.types.guardrails import GuardrailEventHooks - - guardrail: Final = CustomGuardrail( - guardrail_name="output-check", event_hook=GuardrailEventHooks(event_hook) if event_hook else None - ) - monkeypatch.setattr(litellm, "callbacks", [guardrail]) - - assert await _continues_after_content({}) is continues - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallbacks,continues", - [ - ([{"model": "backup", "messages": [{"role": "user", "content": "Say hi"}]}], False), - ([{"model": "backup", "stream": False}], False), - ([{"model": "backup", "max_tokens": 300}], False), - ([{"model": "backup"}], True), - ([{"gpt-4": [{"model": "backup", "stream": False}]}], False), - ([{"*": [{"model": "backup", "max_tokens": 300}]}], False), - ([{"gpt-4": [{"model": "backup"}, "backup"]}], True), - ([{"gpt-4": ["backup"]}, {"other-group": [{"model": "backup", "stream": False}]}], True), - ([{}], True), - ([{"gpt-4": 5}], True), - ], - ids=[ - "entry_with_its_own_messages", - "entry_turning_off_streaming", - "entry_with_its_own_max_tokens", - "entry_with_model_only", - "group_target_turning_off_streaming", - "generic_group_target_with_its_own_max_tokens", - "group_targets_with_model_only", - "override_only_in_another_groups_chain", - "empty_entry", - "non_list_chain", - ], -) -async def test_mid_stream_continuation_declines_a_request_fallback_that_overrides_request_params( - fallbacks: list[object], continues: bool -) -> None: - assert await _continues_after_content({"fallbacks": fallbacks}) is continues - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallbacks,continues", - [ - ([{"gpt-4": ["no-prefill-deployment"]}], False), - ([{"gpt-4": ["backup", "no-prefill-deployment"]}], False), - ([{"model": "no-prefill-deployment"}], False), - ([{"gpt-4": ["prefill-deployment"]}], True), - ([{"gpt-4": ["openai-backup"]}], True), - ([{"gpt-4": ["json-prefill-deployment"]}], False), - ], - ids=[ - "deployment_id_without_prefill", - "deployment_id_without_prefill_later_in_chain", - "flat_entry_naming_a_deployment_id_without_prefill", - "deployment_id_with_prefill", - "group_without_prefill_left_to_the_deployment_filter", - "deployment_id_with_its_own_response_format", - ], -) -async def test_mid_stream_continuation_declines_a_fallback_naming_a_deployment_that_cannot_continue( - fallbacks: list[object], continues: bool -) -> None: - assert await _continues_after_content({"fallbacks": fallbacks}) is continues - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "primary_params,default_litellm_params,continues", - [ - ({"response_format": {"type": "json_object"}}, None, False), - ({"n": 2}, None, False), - ({"tool_choice": "required"}, None, False), - ({"merge_reasoning_content_in_choices": True}, None, False), - ({}, {"response_format": {"type": "json_object"}}, False), - ({"response_format": {"type": "text"}}, None, True), - ({"temperature": 0.2}, None, True), - ({"tool_choice": "required"}, {"tool_choice": "auto"}, False), - ({"response_format": {"type": "json_object"}}, {"response_format": {"type": "text"}}, True), - ({"extra_body": {"n": 2}}, None, False), - ], - ids=[ - "deployment_response_format", - "deployment_n", - "deployment_tool_choice", - "deployment_merged_reasoning", - "router_default_response_format", - "deployment_plain_text_format", - "deployment_unrelated_param", - "deployment_tool_choice_over_router_default", - "router_default_format_over_deployment", - "deployment_extra_body_n", - ], -) -async def test_mid_stream_continuation_reads_constraints_configured_on_the_cut_deployment( - primary_params: Mapping[str, object], default_litellm_params: dict[str, object] | None, continues: bool -) -> None: - assert await _continues_after_content({}, primary_params, default_litellm_params) is continues - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "extra_body,continues", - [({"response_format": {"type": "json_object"}}, False), ({"n": 2}, False), ({"top_k": 5}, True)], - ids=["json_format_in_extra_body", "n_in_extra_body", "unrelated_extra_body_key"], -) -async def test_mid_stream_continuation_reads_constraints_the_request_sends_in_extra_body( - extra_body: dict[str, object], continues: bool -) -> None: - assert await _continues_after_content({"extra_body": extra_body}) is continues - - -@pytest.mark.asyncio -async def test_mid_stream_continuation_declines_when_the_cut_deployment_is_not_in_the_router() -> None: - assert await _continues_after_content({}, cut_model_id="per-request-credential-deployment") is False - - @pytest.mark.asyncio @pytest.mark.parametrize( "routing_strategy", @@ -4694,645 +3018,6 @@ async def test_continuation_routes_only_to_deployments_that_can_continue_under_e ) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "flag_kwargs", [{}, {"enable_mid_stream_fallback_continuation": "false"}], ids=["default", "string_false"] -) -async def test_acompletion_streaming_iterator_flag_off_declines_after_content(flag_kwargs): - """Regression guard for the opt-in default: with the flag off, a post-content - break re-raises and never falls back, exactly as before this feature.""" - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - **flag_kwargs, - ) - error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - source = _make_midstream_source(error) - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=source, messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - with pytest.raises(MidStreamFallbackError): - async for _ in result: - pass - - mock_fallback.assert_not_awaited() - - -def test_build_completion_continuation_input_appends_assistant_prefill(): - messages = [{"role": "user", "content": "hi"}] - built = litellm.Router._build_completion_continuation_input(messages, "partial answer") - assert built[:-1] == messages - assert built[-1] == {"role": "assistant", "content": "partial answer", "prefix": True} - - -@pytest.mark.parametrize( - "prefill_content,second_partial,expected_content", - [ - ("part one ", " part two ", "part one part two "), - ("part one ", "part two ", "part one part two "), - ("First.\n\n", "Second", "First.\n\nSecond"), - ([{"type": "text", "text": "part one "}], " part two ", [{"type": "text", "text": "part one part two "}]), - ], - ids=["separator_sent_again", "separator_not_sent_again", "paragraph_break_not_sent_again", "content_blocks"], -) -def test_build_completion_continuation_input_folds_into_existing_prefill( - prefill_content, second_partial, expected_content -): - """A nested break must not leave two trailing assistant turns: the new partial folds into the prior prefill, - which keeps reading as the text the caller already received.""" - messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}] - built = litellm.Router._build_completion_continuation_input(messages, second_partial) - assert built == [messages[0], {"role": "assistant", "content": expected_content, "prefix": True}] - - -@pytest.mark.parametrize( - "prefill_content,expected_content", - [ - ("Hello", "Hello world"), - ([{"type": "text", "text": "Hello"}], [{"type": "text", "text": "Hello world"}]), - ], - ids=["string_prefill", "content_block_prefill"], -) -def test_build_completion_continuation_input_does_not_double_an_echoed_prefill(prefill_content, expected_content): - messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}] - built = litellm.Router._build_completion_continuation_input(messages, "Hello world") - assert built is not None - assert built[-1] == {"role": "assistant", "content": expected_content, "prefix": True} - - -@pytest.mark.parametrize( - "cut_partial,expected_content", - [("Hel", None), ("Hello", "Hello"), ("Hello wor", "Hello wor")], - ids=["cut_mid_echo", "cut_after_a_full_echo", "cut_after_new_text"], -) -def test_build_completion_continuation_input_declines_a_cut_in_the_middle_of_a_prefill_echo( - cut_partial, expected_content -): - messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "Hello", "prefix": True}] - built = litellm.Router._build_completion_continuation_input(messages, cut_partial) - expected = None if expected_content is None else [messages[0], {**messages[1], "content": expected_content}] - assert built == expected - - -@pytest.mark.parametrize( - "assistant_content", ["Earlier answer.", [{"type": "text", "text": "Earlier answer."}]], ids=["string", "blocks"] -) -def test_build_completion_continuation_input_declines_a_trailing_plain_assistant_turn(assistant_content): - """The caller's own assistant turn is history, not a prefill, so the partial is never fused into it.""" - messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": assistant_content}] - assert litellm.Router._build_completion_continuation_input(messages, "The capital of France is") is None - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_declines_when_request_ends_in_plain_assistant_turn(): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Earlier answer."}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - with pytest.raises(MidStreamFallbackError): - async for _ in result: - pass - - mock_fallback.assert_not_awaited() - - -def test_build_completion_continuation_input_keeps_structured_assistant_content(): - """Content blocks on a trailing assistant turn stay blocks: the partial continues the last text - block in place (a block boundary would let a provider drop the space between "Sure, " and "here"), - lands as a new text block only after a non-text block, and a content shape that cannot be extended declines.""" - from litellm.router import _extend_trailing_text_block - - messages = [ - {"role": "user", "content": "hi"}, - {"role": "assistant", "content": [{"type": "text", "text": "Sure, "}], "prefix": True}, - ] - built = litellm.Router._build_completion_continuation_input(messages, "here it is") - assert built is not None - assert [m["role"] for m in built] == ["user", "assistant"] - assert built[-1] == {"role": "assistant", "content": [{"type": "text", "text": "Sure, here it is"}], "prefix": True} - image_block = {"type": "image_url", "image_url": {"url": "https://example.test/a.png"}} - assert _extend_trailing_text_block([image_block], "here it is") == [image_block, {"type": "text", "text": "here it is"}] - assert _extend_trailing_text_block([], "here it is") == [{"type": "text", "text": "here it is"}] - assert ( - litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42, "prefix": True}], "x") - is None - ) - - -def test_mid_stream_continuation_eligible_allows_text_response_format(): - """response_format={"type": "text"} is the unconstrained default and must stay - eligible, unlike json_object / json_schema.""" - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k"}}], - enable_mid_stream_fallback_continuation=True, - ) - e = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "text"}}) is True - assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) is False - - -def test_raise_original_mid_stream_error_surfaces_the_provider_exception(): - from litellm.exceptions import MidStreamFallbackError, RateLimitError - - provider_error = RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4") - wrapped = MidStreamFallbackError( - message="rate limited", model="gpt-4", llm_provider="openai", - original_exception=provider_error, generated_content="Hello", - ) - with pytest.raises(RateLimitError) as raised: - litellm.Router._raise_original_mid_stream_error(wrapped) - assert raised.value is provider_error - assert raised.value.__cause__ is wrapped - - bare = MidStreamFallbackError(message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello") - with pytest.raises(MidStreamFallbackError) as bare_raised: - litellm.Router._raise_original_mid_stream_error(bare) - assert bare_raised.value is bare - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_reraises_original_exception_when_available(): - """Async: when the mid-stream MidStreamFallbackError wraps a real provider - exception (original_exception), the router must re-raise that original - exception instead of the internal wrapper, so the client sees the - specific error type/code (e.g. RateLimitError) rather than a generic - MidStreamFallbackError.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError, RateLimitError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - set_verbose=True, - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - original_exception = RateLimitError( - message="rate limited", - llm_provider="vertex_ai", - model="gpt-4", - ) - error = MidStreamFallbackError( - message="rate limited", - model="gpt-4", - llm_provider="openai", - original_exception=original_exception, - generated_content="Hello", - ) - - mock_chunks = [ - MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content=" there"))]), - ] - - class AsyncIteratorWithError: - def __init__(self, items, error_after_index): - self.items = items - self.index = 0 - self.error_after_index = error_after_index - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index >= len(self.items): - raise StopAsyncIteration - if self.index == self.error_after_index: - raise error - item = self.items[self.index] - self.index += 1 - return item - - mock_error_response = AsyncIteratorWithError(mock_chunks, 1) - setattr(mock_error_response, "model", "gpt-4") - setattr(mock_error_response, "custom_llm_provider", "openai") - setattr(mock_error_response, "logging_obj", MagicMock()) - - result = await router._acompletion_streaming_iterator( - model_response=mock_error_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(RateLimitError) as exc_info: - async for _ in result: - pass - assert exc_info.value is original_exception - assert exc_info.value.type == "throttling_error" - assert exc_info.value.code == "429" - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_edge_cases(): - """Test edge cases for _acompletion_streaming_iterator.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - set_verbose=True, - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - # Test: Empty generated content - empty_error = MidStreamFallbackError( - message="Error", - model="gpt-4", - llm_provider="openai", - generated_content="", # Empty content - ) - - class AsyncIteratorImmediateError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - raise empty_error - - mock_response = AsyncIteratorImmediateError() - - # Mock empty fallback response using AsyncIterator - class EmptyAsyncIterator: - def __aiter__(self): - return self - - async def __anext__(self): - raise StopAsyncIteration - - mock_fallback_response = EmptyAsyncIterator() - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=mock_fallback_response, - ) as mock_fallback_utils: - collected_chunks = [] - iterator = await router._acompletion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - async for chunk in iterator: - collected_chunks.append(chunk) - - # Should still call fallback even with empty content - assert mock_fallback_utils.called - fallback_kwargs = mock_fallback_utils.call_args.kwargs["kwargs"] - modified_messages = fallback_kwargs["messages"] - - # Empty content → pre-first-chunk path uses original messages - # (no continuation prompt added) - assert modified_messages == messages - print("✓ Handles empty generated content correctly") - - print("✓ Edge case tests passed!") - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_preserves_hidden_params(): - """ - Regression test: FallbackStreamWrapper must copy _hidden_params from the - original CustomStreamWrapper so that x-litellm-overhead-duration-ms (and - other hidden params) are present in the proxy response headers for streaming. - """ - from unittest.mock import MagicMock - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - # Simulate a CustomStreamWrapper that already has timing metadata set by - # update_response_metadata (litellm_overhead_time_ms, _response_ms, etc.) - mock_response = MagicMock() - mock_response.model = "gpt-4" - mock_response.custom_llm_provider = "openai" - mock_response.logging_obj = MagicMock() - mock_response._hidden_params = { - "litellm_overhead_time_ms": 12.34, - "_response_ms": 500.0, - "litellm_call_id": "test-call-id", - "api_base": "https://api.openai.com", - "additional_headers": {}, - } - - # Make the mock iterable (yields nothing — we only care about hidden_params) - async def _empty(): - return - yield # make it an async generator - - mock_response.__aiter__ = lambda self: _empty().__aiter__() - - result = await router._acompletion_streaming_iterator( - model_response=mock_response, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - - # The returned FallbackStreamWrapper must carry the original _hidden_params - assert hasattr(result, "_hidden_params"), "result must have _hidden_params" - assert result._hidden_params.get("litellm_overhead_time_ms") == 12.34, ( - "litellm_overhead_time_ms must be preserved — " - "this is what drives x-litellm-overhead-duration-ms in streaming responses" - ) - assert result._hidden_params.get("litellm_call_id") == "test-call-id" - assert result._hidden_params.get("_response_ms") == 500.0 - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_preserves_response_headers(): - """LIT-6767: the returned wrapper must carry the provider's raw response headers. - - Proxy callbacks read ``_response_headers`` off the object the router hands - back. The wrapper used to be built without it, so every streaming chat - completion reported zero raw provider headers while the non-streaming path - reported the full set. - """ - from unittest.mock import MagicMock - - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - async def _empty(): - return - yield # make it an async generator - - provider_headers = { - "x-request-id": "req-provider-123", - "x-ratelimit-remaining-requests": "42", - # a provider must never be able to spoof an internal header - "x-litellm-model-id": "spoofed", - } - source = CustomStreamWrapper( - completion_stream=_empty(), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers=provider_headers, - ) - - result = await router._acompletion_streaming_iterator( - model_response=source, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - - assert result._response_headers == provider_headers - additional_headers = result._hidden_params["additional_headers"] - assert additional_headers["llm_provider-x-request-id"] == "req-provider-123" - assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42" - # internal-header protection survives: the provider value is namespaced, never promoted - assert additional_headers["llm_provider-x-litellm-model-id"] == "spoofed" - assert "x-litellm-model-id" not in additional_headers - - -def test_completion_streaming_iterator_preserves_response_headers(): - """LIT-6767, sync counterpart of the async header-preservation test.""" - from unittest.mock import MagicMock - - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - provider_headers = {"x-request-id": "req-provider-sync", "openai-organization": "org-real"} - source = CustomStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers=provider_headers, - ) - - result = router._completion_streaming_iterator( - model_response=source, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - - assert result._response_headers == provider_headers - assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-provider-sync" - - -def test_adopt_fallback_response_headers_replaces_rather_than_merges(): - """LIT-6767: direct unit for FallbackAwareStreamWrapper.adopt_fallback_response_headers. - - Values from the failed attempt must not survive, so the wrapper replaces both - ``_response_headers`` and ``_hidden_params`` instead of merging them. - """ - from unittest.mock import MagicMock - - from litellm.router import FallbackAwareStreamWrapper, Router - - wrapper = FallbackAwareStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers={"x-request-id": "req-FAILED"}, - ) - wrapper._hidden_params = { - "model_id": "failed-deployment", - "only_on_failed_attempt": "stale", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - } - - fallback = MagicMock() - fallback._response_headers = {"x-request-id": "req-FALLBACK"} - fallback._hidden_params = { - "model_id": "fallback-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, - } - - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) - - assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} - assert wrapper._hidden_params["model_id"] == "fallback-deployment" - assert "only_on_failed_attempt" not in wrapper._hidden_params - assert wrapper._hidden_params is not fallback._hidden_params - # the snapshot CustomStreamWrapper caches at init has to follow, or a chunk built - # from it would still be stamped with the deployment that failed - assert wrapper._base_hidden_params["model_id"] == "fallback-deployment" - # the nested header dict is copied too, so a later mutation on the fallback - # response cannot reach headers the proxy has already published - assert wrapper._hidden_params["additional_headers"] is not fallback._hidden_params["additional_headers"] - fallback._hidden_params["additional_headers"]["llm_provider-x-request-id"] = "req-MUTATED" - assert wrapper._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} - - -def test_adopt_fallback_response_headers_survives_a_collected_wrapper(): - """LIT-6767: adoption still returns the fallback's params once the wrapper is gone.""" - import weakref - from unittest.mock import MagicMock - - from litellm.router import FallbackAwareStreamWrapper, Router - - fallback: Final = MagicMock() - fallback._response_headers = {"x-request-id": "req-FALLBACK"} - fallback._hidden_params = { - "model_id": "fallback-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, - } - - wrapper = FallbackAwareStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - ) - live_ref: Final = weakref.ref(wrapper) - prepared: Final = Router._adopt_fallback_response_headers(live_ref, fallback) - assert prepared == (fallback._hidden_params, fallback._hidden_params["additional_headers"]) - assert wrapper.fallback_headers_adopted is True - assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} - - dead_ref: Final = weakref.ref(wrapper) - del wrapper - assert dead_ref() is None - assert Router._adopt_fallback_response_headers(dead_ref, fallback) == prepared - - -def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_replace(): - """LIT-6767: a fallback that carries no raw provider headers publishes none. - - Keeping the failed attempt's raw headers would hand the client and the callbacks a - provider ``x-request-id`` for a request that deployment never served, which is the - leak this fix exists to close. - """ - from unittest.mock import MagicMock - - from litellm.router import FallbackAwareStreamWrapper, Router - - wrapper = FallbackAwareStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers={"x-request-id": "req-FAILED"}, - ) - wrapper._hidden_params = {"model_id": "failed-deployment", "additional_headers": {}} - - fallback = MagicMock() - fallback._response_headers = None - fallback._hidden_params = {"model_id": "fallback-deployment"} - - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) - - assert wrapper._response_headers is None - assert wrapper._hidden_params["model_id"] == "fallback-deployment" - assert wrapper.fallback_headers_adopted is True - - -def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none(): - """A fallback response carrying no hidden params keeps the identity headers. - - Publishing no ``x-litellm-*`` header at all for a request the fallback served is - worse than keeping what is there, so only the raw provider headers are dropped. - """ - from unittest.mock import MagicMock - - from litellm.router import FallbackAwareStreamWrapper, Router - - wrapper = FallbackAwareStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers={"x-request-id": "req-FAILED"}, - ) - hidden_params_before = wrapper._hidden_params - - fallback = object() - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) - - assert wrapper._response_headers is None - assert wrapper._hidden_params is hidden_params_before - assert wrapper.fallback_headers_adopted is True - - @pytest.mark.asyncio @pytest.mark.parametrize("response_kind", ["object", "dict", "async-generator"]) async def test_set_response_headers_exposes_complexity_decision_on_every_response_shape( @@ -5400,299 +3085,7 @@ async def test_set_response_headers_is_the_only_complexity_header_source_for_pro **additional_headers, ) - assert not { - key for key in proxy_headers if key.startswith("x-litellm-complexity-router-") - } - - -@pytest.mark.asyncio -@pytest.mark.parametrize("close_fails", [False, True]) -async def test_acompletion_streaming_iterator_adopts_fallback_response_headers( - close_fails: bool, -) -> None: - """LIT-6767: after a successful pre-first-chunk fallback, the wrapper must - describe the deployment that served the stream, with no value left over - from the attempt that failed.""" - from unittest.mock import MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - failed_error = MidStreamFallbackError( - message="upstream died before the first chunk", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=True, - ) - - class FailedStream: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - self._response_headers = {"x-request-id": "req-FAILED"} - self._hidden_params = { - "model_id": "failed-deployment", - "api_base": "https://failed.example", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - "only_on_failed_attempt": "stale", - } - - def __aiter__(self): - return self - - async def __anext__(self): - raise failed_error - - class FallbackStream: - def __init__(self): - self.close_attempted: bool = False - self._response_headers = {"x-request-id": "req-FALLBACK"} - self._hidden_params = { - "model_id": "fallback-deployment", - "api_base": "https://fallback.example", - "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, - } - self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) - - def __aiter__(self): - return self - - async def __anext__(self): - try: - return next(self._chunks) - except StopIteration: - raise StopAsyncIteration from None - - async def aclose(self) -> None: - self.close_attempted = True - if close_fails: - raise RuntimeError("fallback transport failed during close") - - fallback_stream = FallbackStream() - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=fallback_stream, - ): - result = await router._acompletion_streaming_iterator( - model_response=FailedStream(), - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - # the failed attempt is what the wrapper is built from - assert result._response_headers == {"x-request-id": "req-FAILED"} - # the very first chunk the fallback produces must already be published under - # the fallback's identity: the proxy commits response headers once that chunk - # is buffered, so adopting any later is adopting too late - first_chunk: Final = await result.__anext__() - assert first_chunk.choices[0].delta.content == "OK" - assert result._response_headers == {"x-request-id": "req-FALLBACK"} - assert result._hidden_params["model_id"] == "fallback-deployment" - async for _ in result: - pass - - assert fallback_stream.close_attempted - assert result._response_headers == {"x-request-id": "req-FALLBACK"} - assert result._hidden_params["model_id"] == "fallback-deployment" - assert result._hidden_params["api_base"] == "https://fallback.example" - assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} - # stale values are removed, not merged over - assert "only_on_failed_attempt" not in result._hidden_params - # and the wrapper holds its own copy, so later fallback mutations cannot leak in - assert result._hidden_params is not fallback_stream._hidden_params - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): - """LIT-6767: a fallback that itself fails over before its first chunk. - - The selected fallback still describes its own failed attempt at selection time, so - the wrapper has to re-read it once a chunk exists or it publishes a deployment that - produced no output. - """ - from unittest.mock import MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - failed_error: Final = MidStreamFallbackError( - message="upstream died before the first chunk", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=True, - ) - - class FailedStream: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - self._response_headers = {"x-request-id": "req-FAILED"} - self._hidden_params = { - "model_id": "failed-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - } - - def __aiter__(self): - return self - - async def __anext__(self): - raise failed_error - - class NestedFallbackStream: - """A fallback that repoints itself at a third deployment as it yields.""" - - def __init__(self): - self._response_headers = {"x-request-id": "req-MIDDLE"} - self._hidden_params = { - "model_id": "middle-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, - } - self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) - - def __aiter__(self): - return self - - async def __anext__(self): - try: - chunk = next(self._chunks) - except StopIteration: - raise StopAsyncIteration from None - self._response_headers = {"x-request-id": "req-SERVED"} - self._hidden_params = { - "model_id": "served-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, - } - return chunk - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=NestedFallbackStream(), - ): - result = await router._acompletion_streaming_iterator( - model_response=FailedStream(), - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - first_chunk: Final = await result.__anext__() - # the proxy commits response headers once this chunk is buffered - assert result._response_headers == {"x-request-id": "req-SERVED"} - assert result._hidden_params["model_id"] == "served-deployment" - assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-SERVED"} - # and the chunk itself carries the same deployment - assert first_chunk._hidden_params["model_id"] == "served-deployment" - async for _ in result: - pass - - assert result._response_headers == {"x-request-id": "req-SERVED"} - assert result._hidden_params["model_id"] == "served-deployment" - - -def test_completion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): - """LIT-6767, sync counterpart of the nested-fallback adoption test.""" - from unittest.mock import MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - failed_error: Final = MidStreamFallbackError( - message="upstream died before the first chunk", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=True, - ) - - class FailedStream: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - self._response_headers = {"x-request-id": "req-FAILED"} - self._hidden_params = { - "model_id": "failed-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - } - - def __iter__(self): - return self - - def __next__(self): - raise failed_error - - class NestedFallbackStream: - """A fallback that repoints itself at a third deployment as it yields.""" - - def __init__(self): - self._response_headers = {"x-request-id": "req-MIDDLE"} - self._hidden_params = { - "model_id": "middle-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, - } - self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) - - def __iter__(self): - return self - - def __next__(self): - chunk = next(self._chunks) - self._response_headers = {"x-request-id": "req-SERVED"} - self._hidden_params = { - "model_id": "served-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, - } - return chunk - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=NestedFallbackStream()), - ): - result = router._completion_streaming_iterator( - model_response=FailedStream(), - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - first_chunk: Final = next(result) - assert result._response_headers == {"x-request-id": "req-SERVED"} - assert result._hidden_params["model_id"] == "served-deployment" - assert first_chunk._hidden_params["model_id"] == "served-deployment" - for _ in result: - pass - - assert result._response_headers == {"x-request-id": "req-SERVED"} - assert result._hidden_params["model_id"] == "served-deployment" + assert not {key for key in proxy_headers if key.startswith("x-litellm-complexity-router-")} def test_completion_streaming_fallback_resumes_chain_without_retrying_primary(): @@ -6020,614 +3413,6 @@ async def test_acompletion_mid_stream_fallback_honors_a_per_request_fallbacks_no assert _groups_called(provider_calls) == ["primary"] -def test_refusal_on_the_last_fallback_hop_is_returned_instead_of_raised(): - """LIT-7400 follow-up: a refusal on the final hop of an exhausted list passes through.""" - from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets - - router = litellm.Router( - model_list=[ - {"model_name": "primary", "litellm_params": {"model": "openai/primary-model", "api_key": "fake-key"}}, - {"model_name": "fb1", "litellm_params": {"model": "openai/fb1-model", "api_key": "fake-key"}}, - {"model_name": "fb2", "litellm_params": {"model": "openai/fb2-model", "api_key": "fake-key"}}, - ], - fallbacks=[{"primary": ["fb1", "fb2"]}], - num_retries=0, - ) - - attempted: Final = AttemptedFallbackTargets() - attempted.record("primary") - attempted.record("fb1") - attempted.record("fb2") - kwargs: Final = { - "attempted_targets": attempted, - "metadata": {"model_group": "fb2", "original_model_group": "primary"}, - } - - assert router._refusal_fallback_available("fb2", kwargs) is False - assert ( - router._refusal_fallback_available( - "fb1", {"metadata": {"model_group": "fb1", "original_model_group": "primary"}} - ) - is True - ) - - -def test_completion_streaming_iterator_adopts_fallback_response_headers(): - """LIT-6767, sync counterpart of the fallback-adoption test.""" - from unittest.mock import MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - failed_error = MidStreamFallbackError( - message="upstream died before the first chunk", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=True, - ) - - class FailedStream: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - self._response_headers = {"x-request-id": "req-FAILED"} - self._hidden_params = { - "model_id": "failed-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - "only_on_failed_attempt": "stale", - } - - def __iter__(self): - return self - - def __next__(self): - raise failed_error - - class FallbackStream: - def __init__(self): - self._response_headers = {"x-request-id": "req-FALLBACK"} - self._hidden_params = { - "model_id": "fallback-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, - } - - def __iter__(self): - return iter([]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=FallbackStream()), - ): - result = router._completion_streaming_iterator( - model_response=FailedStream(), - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - assert result._response_headers == {"x-request-id": "req-FAILED"} - for _ in result: - pass - - assert result._response_headers == {"x-request-id": "req-FALLBACK"} - assert result._hidden_params["model_id"] == "fallback-deployment" - assert "only_on_failed_attempt" not in result._hidden_params - - -def test_completion_streaming_iterator_fallback_on_429(): - """Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback. - - This is the sync counterpart of test_acompletion_streaming_iterator. - Before this fix, __next__ raised RateLimitError directly and the Router - never got a chance to fall back. - """ - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - rate_limit_error = MidStreamFallbackError( - message="Resource exhausted", - model="gpt-4", - llm_provider="vertex_ai", - generated_content="", - is_pre_first_chunk=True, - ) - - class SyncIteratorImmediateError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __iter__(self): - return self - - def __next__(self): - raise rate_limit_error - - mock_response = SyncIteratorImmediateError() - - # Fallback returns a simple non-streaming response (fallback may not stream) - mock_fallback_response = MagicMock() - mock_fallback_response.__iter__ = MagicMock(return_value=iter([])) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=mock_fallback_response), - ) as mock_fallback: - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - collected_chunks = list(result) - - mock_fallback.assert_awaited_once() - call_kwargs = mock_fallback.await_args.kwargs["kwargs"] - # Pre-first-chunk: should use original messages, no continuation prompt - assert call_kwargs.get("messages") == messages - # Verify original_function is _completion (sync) - assert call_kwargs.get("original_function") == router._completion - - -def test_completion_streaming_iterator_preserves_hidden_params(): - """SyncFallbackStreamWrapper must copy _hidden_params from original response.""" - from unittest.mock import MagicMock - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - mock_response = MagicMock() - mock_response.model = "gpt-4" - mock_response.custom_llm_provider = "openai" - mock_response.logging_obj = MagicMock() - mock_response._hidden_params = { - "litellm_overhead_time_ms": 42.0, - "litellm_call_id": "test-sync-call", - } - mock_response.__iter__ = MagicMock(return_value=iter([])) - - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - - assert hasattr(result, "_hidden_params") - assert result._hidden_params.get("litellm_overhead_time_ms") == 42.0 - assert result._hidden_params.get("litellm_call_id") == "test-sync-call" - - -def test_completion_streaming_iterator_reraises_mid_chunk_error(): - """Sync: MidStreamFallbackError with generated_content and is_pre_first_chunk=False - must be re-raised immediately; the router cannot recover after partial content - has already been sent to the client.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - mid_chunk_error = MidStreamFallbackError( - message="Connection reset", - model="gpt-4", - llm_provider="openai", - generated_content="Hello, I am", - is_pre_first_chunk=False, - ) - - class SyncIteratorMidChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __iter__(self): - return self - - def __next__(self): - raise mid_chunk_error - - mock_response = SyncIteratorMidChunkError() - - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(MidStreamFallbackError): - list(result) - - -def test_completion_streaming_iterator_reraises_original_exception_when_available(): - """Sync: when the mid-chunk MidStreamFallbackError wraps a real provider - exception (original_exception), the router must re-raise that original - exception instead of the internal wrapper, so the client sees the - specific error type/code (e.g. RateLimitError) rather than a generic - MidStreamFallbackError.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError, RateLimitError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - original_exception = RateLimitError( - message="rate limited", - llm_provider="vertex_ai", - model="gpt-4", - ) - mid_chunk_error = MidStreamFallbackError( - message="rate limited", - model="gpt-4", - llm_provider="openai", - original_exception=original_exception, - generated_content="Hello, I am", - is_pre_first_chunk=False, - ) - - class SyncIteratorMidChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __iter__(self): - return self - - def __next__(self): - raise mid_chunk_error - - mock_response = SyncIteratorMidChunkError() - - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(RateLimitError) as exc_info: - list(result) - assert exc_info.value is original_exception - assert exc_info.value.type == "throttling_error" - assert exc_info.value.code == "429" - - -def test_completion_streaming_iterator_reraises_mid_chunk_error_with_no_text_content(): - """Sync: a reasoning-only chunk sets is_pre_first_chunk=False without populating - generated_content (which only tracks text deltas). The re-raise guard must still - detect this via the raw chunks on the wrapper, or the router silently retries and - the client receives duplicated/inconsistent output.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - from litellm.types.utils import Delta, StreamingChoices - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - mid_chunk_error = MidStreamFallbackError( - message="Connection reset", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=False, - ) - - reasoning_chunk = litellm.ModelResponseStream( - id="chatcmpl-partial-1", - model="gpt-4", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(reasoning_content="Thinking about the answer", role="assistant"), - ) - ], - ) - - class SyncIteratorNoTextChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [reasoning_chunk] - - def __iter__(self): - return self - - def __next__(self): - raise mid_chunk_error - - mock_response = SyncIteratorNoTextChunkError() - - with patch.object(router, "function_with_fallbacks") as mock_fallback: - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(MidStreamFallbackError): - list(result) - - assert not mock_fallback.called, ( - "fallback must not be attempted once any content, text or non-text, has already streamed" - ) - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_pre_first_chunk_skips_continuation(): - """When MidStreamFallbackError has is_pre_first_chunk=True, use original messages.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Hello"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - pre_first_chunk_error = MidStreamFallbackError( - message="429 Resource exhausted", - model="gpt-4", - llm_provider="vertex_ai", - generated_content="", - is_pre_first_chunk=True, - ) - - class AsyncIteratorPreFirstChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - raise pre_first_chunk_error - - mock_response = AsyncIteratorPreFirstChunkError() - - class EmptyAsyncIterator: - def __aiter__(self): - return self - - async def __anext__(self): - raise StopAsyncIteration - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=EmptyAsyncIterator(), - ) as mock_fallback_utils: - iterator = await router._acompletion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - async for _ in iterator: - pass - - assert mock_fallback_utils.called - fallback_kwargs = mock_fallback_utils.call_args.kwargs["kwargs"] - # Pre-first-chunk: should use original messages, no continuation prompt - assert fallback_kwargs["messages"] == messages - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_reraises_mid_chunk_error_with_no_text_content(): - """Async: a reasoning-only chunk sets is_pre_first_chunk=False without populating - generated_content (which only tracks text deltas). The re-raise guard must still - detect this via the raw chunks on the wrapper, or the router silently retries and - the client receives duplicated/inconsistent output.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - from litellm.types.utils import Delta, StreamingChoices - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - mid_chunk_error = MidStreamFallbackError( - message="Connection reset", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=False, - ) - - reasoning_chunk = litellm.ModelResponseStream( - id="chatcmpl-partial-1", - model="gpt-4", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(reasoning_content="Thinking about the answer", role="assistant"), - ) - ], - ) - - class AsyncIteratorNoTextChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [reasoning_chunk] - - def __aiter__(self): - return self - - async def __anext__(self): - raise mid_chunk_error - - mock_response = AsyncIteratorNoTextChunkError() - - with patch.object(router, "async_function_with_fallbacks_common_utils") as mock_fallback_utils: - iterator = await router._acompletion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(MidStreamFallbackError): - async for _ in iterator: - pass - - assert not mock_fallback_utils.called, ( - "fallback must not be attempted once any content, text or non-text, has already streamed" - ) - - -# --------------------------------------------------------------------------- -# Shared helpers for the _aresponses_streaming_iterator test suite. -# --------------------------------------------------------------------------- -def _make_responses_iterator( - *, - chunks=(), - error=None, - bridge=False, - model="gpt-4", - hidden_params=None, - chat_chunks=None, -): - """Build a minimal mock Responses-API streaming iterator. - - Bypasses BaseResponsesAPIStreamingIterator.__init__ but mirrors every - attribute production code reads. Yields *chunks*, then raises *error* - (or StopAsyncIteration). Set bridge=True to inherit from - LiteLLMCompletionStreamingIterator so the wrapper's bridge-path - isinstance check (used by usage extraction) matches. - """ - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) - - class _Iter(base): - def __init__(self): - self._chunks = list(chunks) - self._idx = 0 - self._hidden_params = hidden_params or {} - self.model = model - self.custom_llm_provider = "anthropic" - self.logging_obj = MagicMock() - self.litellm_metadata = None - self.responses_api_provider_config = None - self.finished = False - self.completed_response = None - self.response = None - self.start_time = None - self.request_data = {} - self.call_type = None - if chat_chunks is not None: - self.collected_chat_completion_chunks = chat_chunks - - def __aiter__(self): - return self - - async def __anext__(self): - if self._idx < len(self._chunks): - self._idx += 1 - return self._chunks[self._idx - 1] - if error is not None: - raise error - raise StopAsyncIteration - - return _Iter() - - class _AsyncList: """Generic async iterator over a list — used as the fallback response.""" @@ -6646,106 +3431,6 @@ class _AsyncList: return item -def _make_router_with_fallback(primary="gpt-4", secondary="gpt-3.5-turbo"): - return litellm.Router( - model_list=[ - { - "model_name": primary, - "litellm_params": {"model": primary, "api_key": "k1"}, - }, - { - "model_name": secondary, - "litellm_params": {"model": secondary, "api_key": "k2"}, - }, - ], - fallbacks=[{primary: [secondary]}], - ) - - -# The Rust core walks fallbacks and runs pre-routing in Rust, past the members these tests patch. -_PYTHON_CORE_ONLY: Final = pytest.mark.usefixtures("python_router_core") - - -def _injected_fallback_router(fallback_response: object) -> Router: - router: Final = Router(model_list=[], fallbacks=[{"primary": ["fallback"]}]) - - async def inject_fallback( - e: Exception, - disable_fallbacks: bool | None, - fallbacks: list | None, - context_window_fallbacks: list | None, - content_policy_fallbacks: list | None, - model_group: str | None, - args: tuple[object, ...], - kwargs: dict[str, object], - include_fallback_errors: bool = False, - ) -> object: - return fallback_response - - router.async_function_with_fallbacks_common_utils = inject_fallback - return router - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_fallback(): - """Catches MidStreamFallbackError, re-enters the fallback chain via - async_function_with_fallbacks_common_utils with the per-attempt helper - and original_generic_function preserved. Mirrors - test_acompletion_streaming_iterator for the aresponses path.""" - from litellm.exceptions import MidStreamFallbackError - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - - router = _make_router_with_fallback( - "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" - ) - src = _make_responses_iterator( - chunks=[MagicMock(type="response.created")], - error=MidStreamFallbackError( - message="anthropic socket timeout", - model="anthropic/claude-sonnet-4-6", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="", - ), - model="anthropic/claude-sonnet-4-6", - hidden_params={"model_id": "src-deployment-1"}, - ) - fallback_chunks = [ - MagicMock(type="response.created"), - MagicMock(type="response.output_text.delta"), - MagicMock(type="response.completed"), - ] - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(fallback_chunks), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "anthropic/claude-sonnet-4-6", - "stream": True, - "input": "Hi", - "original_generic_function": litellm.aresponses, - }, - ) - assert isinstance(wrapped, BaseResponsesAPIStreamingIterator) - assert wrapped._hidden_params.get("model_id") == "src-deployment-1" - collected = [c async for c in wrapped] - - assert collected == fallback_chunks - call_kwargs = mock_fallback_utils.call_args.kwargs - fbk = call_kwargs["kwargs"] - # Bound methods compare equal when they share the same instance + __func__. - assert fbk["original_function"] == router._ageneric_api_call_with_fallbacks_responses_attempt - assert fbk["original_generic_function"] is litellm.aresponses - assert call_kwargs["model_group"] == "anthropic/claude-sonnet-4-6" - assert call_kwargs["disable_fallbacks"] is False - - @pytest.mark.asyncio async def test_aresponses_mid_stream_order_fallback_hop_drops_the_encrypted_reasoning_the_next_provider_cannot_decrypt( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch @@ -6975,433 +3660,6 @@ async def test_aresponses_mid_stream_order_fallback_hop_keeps_the_encrypted_reas assert [json.loads(call.request.read())["input"] for call in openai_route.calls] == [history(), history()] -@pytest.mark.asyncio -async def test_aresponses_streaming_content_policy_error_event_routes_to_content_policy_fallback(): - """Regression: a mid-stream content_policy_violation error event never reached - content_policy_fallbacks. The iterator raised a bare APIError the wrapper does not - catch, and even once wrapped, the MidStreamFallbackError envelope was handed to the - fallback dispatch, whose isinstance branch on ContentPolicyViolationError never matched. - The stream below is the customer's shape: a raw OpenAI error event with code - content_policy_violation, transformed by the real OpenAI config, and the router must - call the content_policy_fallbacks target, not the general fallbacks one.""" - from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - - router = litellm.Router( - model_list=[ - {"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k1"}}, - { - "model_name": "content-fallback", - "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "k2"}, - }, - {"model_name": "general-fallback", "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k3"}}, - ], - fallbacks=[{"primary": ["general-fallback"]}], - content_policy_fallbacks=[{"primary": ["content-fallback"]}], - ) - error_event = { - "type": "error", - "sequence_number": 2, - "error": { - "type": "invalid_request_error", - "code": "content_policy_violation", - "message": "This content was flagged for possible cybersecurity risk. The response was halted mid-stream.", - "param": None, - }, - } - - async def aiter_bytes(): - yield f"data: {json.dumps(error_event)}\n\n".encode() - - raw_response = MagicMock() - raw_response.headers = {} - raw_response.aiter_bytes = aiter_bytes - logging_obj = MagicMock(spec=LiteLLMLogging) - logging_obj.model_call_details = {"litellm_params": {}} - logging_obj.completion_start_time = None - source = ResponsesAPIStreamingIterator( - response=raw_response, - model="gpt-5.4", - responses_api_provider_config=OpenAIResponsesAPIConfig(), - logging_obj=logging_obj, - custom_llm_provider="openai", - ) - fallback_chunks = [MagicMock(type="response.output_text.delta"), MagicMock(type="response.completed")] - fallback_call = AsyncMock(return_value=_AsyncList(fallback_chunks)) - - wrapped = await router._aresponses_streaming_iterator( - response=source, - initial_kwargs={ - "model": "primary", - "stream": True, - "input": "Hi", - "original_generic_function": fallback_call, - }, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == fallback_chunks - fallback_call.assert_awaited_once() - assert fallback_call.await_args.kwargs["model"] == "gemini/gemini-2.5-flash" - - -@pytest.mark.asyncio -@_PYTHON_CORE_ONLY -async def test_aresponses_stream_keeps_numbering_after_a_mid_stream_fallback(): - from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents - - def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: - return OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, - item_id="msg_1", - output_index=0, - content_index=0, - delta=text, - sequence_number=sequence_number, - ) - - from litellm.exceptions import MidStreamFallbackError - - router: Final = _injected_fallback_router(_AsyncList([delta(0, "c"), delta(1, "d")])) - primary: Final = _make_responses_iterator( - chunks=[delta(0, "a"), delta(1, "b")], - error=MidStreamFallbackError( - message="primary died", - model="primary", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="ab", - ), - model="primary", - ) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, - ) - events: Final = [event async for event in wrapped] - - assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 1), ("c", 2), ("d", 3)] - - - -@pytest.mark.asyncio -@_PYTHON_CORE_ONLY -async def test_aresponses_stream_fallback_keeps_higher_numbers_and_passes_unnumbered_events(): - from litellm.exceptions import MidStreamFallbackError - from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents - - def delta(text: str) -> OutputTextDeltaEvent: - return OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="msg_1", output_index=0, content_index=0, delta=text - ) - - def numbered(sequence_number: int, text: str) -> OutputTextDeltaEvent: - return delta(text).model_copy(update={"sequence_number": sequence_number}) - - router: Final = _injected_fallback_router(_AsyncList([numbered(7, "c"), delta("d"), numbered(8, "e")])) - primary: Final = _make_responses_iterator( - chunks=[numbered(0, "a"), numbered(1, "b")], - error=MidStreamFallbackError( - message="primary died", - model="primary", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="ab", - ), - model="primary", - ) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, - ) - events: Final = [event async for event in wrapped] - - assert [(event.delta, event.get("sequence_number")) for event in events] == [ - ("a", 0), - ("b", 1), - ("c", 7), - ("d", None), - ("e", 8), - ] - - -@pytest.mark.asyncio -@_PYTHON_CORE_ONLY -async def test_aresponses_stream_without_fallback_passes_upstream_numbers_unchanged(): - from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents - - def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: - return OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, - item_id="msg_1", - output_index=0, - content_index=0, - delta=text, - sequence_number=sequence_number, - ) - - router: Final = _injected_fallback_router(_AsyncList([])) - primary: Final = _make_responses_iterator( - chunks=[delta(0, "a"), delta(1, "b"), delta(1, "c"), delta(0, "d")], model="primary" - ) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, - ) - events: Final = [event async for event in wrapped] - - assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 1), ("c", 1), ("d", 0)] - - -@pytest.mark.asyncio -@_PYTHON_CORE_ONLY -async def test_aresponses_stream_fallback_continues_above_the_highest_number_the_primary_sent(): - from litellm.exceptions import MidStreamFallbackError - from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents - - def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: - return OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, - item_id="msg_1", - output_index=0, - content_index=0, - delta=text, - sequence_number=sequence_number, - ) - - router: Final = _injected_fallback_router(_AsyncList([delta(0, "d")])) - primary: Final = _make_responses_iterator( - chunks=[delta(0, "a"), delta(3, "b"), delta(1, "c")], - error=MidStreamFallbackError( - message="primary died", - model="primary", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="abc", - ), - model="primary", - ) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, - ) - events: Final = [event async for event in wrapped] - - assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 3), ("c", 1), ("d", 4)] - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_unwraps_content_policy_trigger_for_fallback_dispatch(): - """The fallback dispatch matches on the trigger's own type, so the wrapper must hand it the - ContentPolicyViolationError carried inside MidStreamFallbackError, not the envelope.""" - router = _make_router_with_fallback("openai/gpt-5.4", "openai/gpt-5-mini") - content_policy_error = litellm.ContentPolicyViolationError( - message="flagged mid-stream", llm_provider="openai", model="openai/gpt-5.4" - ) - src = _make_responses_iterator( - chunks=[MagicMock(type="response.created")], - error=MidStreamFallbackError( - message=str(content_policy_error), - model="openai/gpt-5.4", - llm_provider="openai", - original_exception=content_policy_error, - is_pre_first_chunk=True, - ), - model="openai/gpt-5.4", - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_AsyncList([MagicMock(type="response.completed")])), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={"model": "openai/gpt-5.4", "stream": True, "input": "Hi"}, - ) - [chunk async for chunk in wrapped] - - mock_fallback_utils.assert_awaited_once() - assert mock_fallback_utils.await_args.kwargs["e"] is content_policy_error - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallback_headers", - [ - {"x-fallback-only": "yes"}, - { - "x-fallback-only": "yes", - "x-litellm-complexity-router-tier": "SIMPLE", - }, - ], - ids=["plain-fallback", "complexity-tier-fallback"], -) -@_PYTHON_CORE_ONLY -async def test_aresponses_streaming_iterator_replaces_complexity_headers_before_fallback_output( - fallback_headers: dict[str, str], -) -> None: - primary_headers: Final = { - "x-litellm-complexity-router-tier": "REASONING", - "x-litellm-complexity-router-reasoning-effort": "xhigh", - } - source: Final = _make_responses_iterator( - error=MidStreamFallbackError( - message="primary failed before output", - model="gpt-4", - llm_provider="openai", - is_pre_first_chunk=True, - generated_content="", - ), - hidden_params={"additional_headers": primary_headers}, - ) - fallback_output: Final = MagicMock(type="response.output_text.delta") - fallback: Final = _AsyncList([fallback_output]) - fallback._hidden_params = { - "model_id": "fallback-deployment", - "additional_headers": fallback_headers, - } - router: Final = _injected_fallback_router(fallback) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=source, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - assert wrapped._hidden_params["additional_headers"] == primary_headers - first_output: Final = await wrapped.__anext__() - - assert first_output is fallback_output - assert wrapped.fallback_headers_adopted is True - assert wrapped._hidden_params == { - "model_id": "fallback-deployment", - "additional_headers": fallback_headers, - } - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback(): - """Regression: model_group must land under "litellm_metadata" (the key - litellm.aresponses reads), not the default "metadata".""" - from litellm.exceptions import MidStreamFallbackError - - router = _make_router_with_fallback() - src = _make_responses_iterator( - error=MidStreamFallbackError( - message="boom", - model="gpt-4", - llm_provider="anthropic", - is_pre_first_chunk=True, - generated_content="", - ) - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - async for _ in wrapped: - pass - - fbk = mock_fallback_utils.call_args.kwargs["kwargs"] - assert "litellm_metadata" in fbk, "wrong metadata_variable_name" - assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get( - "metadata", {} - ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_pre_first_chunk_skips_continuation(): - """Pre-first-chunk error: original input is preserved unchanged.""" - from litellm.exceptions import MidStreamFallbackError - - router = _make_router_with_fallback() - src = _make_responses_iterator( - error=MidStreamFallbackError( - message="socket timeout before first chunk", - model="gpt-4", - llm_provider="anthropic", - is_pre_first_chunk=True, - generated_content="", - ) - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - async for _ in wrapped: - pass - - fbk = mock_fallback_utils.call_args.kwargs["kwargs"] - assert fbk["input"] == "Hello" # original input, no continuation messages - - -def _make_native_responses_iterator(*, sse_payloads: tuple[dict[str, str], ...], trailing_error: Exception | None): - """A real ResponsesAPIStreamingIterator over canned SSE bytes, so the router test covers the - iterator's own transport-error classification instead of a hand-built MidStreamFallbackError.""" - from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - - async def aiter_bytes(): - for payload in sse_payloads: - yield f"data: {json.dumps(payload)}\n\n".encode() - if trailing_error is not None: - raise trailing_error - - def transform(model, parsed_chunk, logging_obj): - return MagicMock(type=parsed_chunk["type"]) - - response: Final = MagicMock() - response.headers = {} - response.aiter_bytes = aiter_bytes - config: Final = MagicMock(spec=BaseResponsesAPIConfig) - config.transform_streaming_response.side_effect = transform - logging_obj: Final = MagicMock(spec=LiteLLMLogging) - logging_obj.completion_start_time = None - logging_obj.model_call_details = {"litellm_params": {}} - return ResponsesAPIStreamingIterator( - response=response, - model="gpt-4", - responses_api_provider_config=config, - logging_obj=logging_obj, - litellm_metadata={}, - custom_llm_provider="openai", - ) - - -_RESPONSES_LIFECYCLE_PAYLOADS: Final = ({"type": "response.created"}, {"type": "response.in_progress"}) - - async def _events_until_error(stream: AsyncIterable[object]) -> AsyncIterator[object]: try: async for chunk in stream: @@ -7410,144 +3668,6 @@ async def _events_until_error(stream: AsyncIterable[object]) -> AsyncIterator[ob yield error -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_falls_back_on_transport_drop_before_output(): - """A connection lost after response.created but before any output item is re-routed to the - fallback with the original input, the same as a provider error event would be, and the client - sees one response lifecycle: the fallback's, whose id the completed event carries.""" - router: Final = _make_router_with_fallback() - src: Final = _make_native_responses_iterator( - sse_payloads=_RESPONSES_LIFECYCLE_PAYLOADS, - trailing_error=httpx.ReadError("Response payload is not completed"), - ) - fallback_chunks: Final = [ - MagicMock(type="response.created", response=MagicMock(id="resp_fallback")), - MagicMock(type="response.in_progress", response=MagicMock(id="resp_fallback")), - MagicMock(type="response.output_text.delta"), - MagicMock(type="response.completed", response=MagicMock(id="resp_fallback")), - ] - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(fallback_chunks), - ) as mock_fallback_utils: - wrapped: Final = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - collected: Final = [chunk async for chunk in wrapped] - - assert collected == fallback_chunks - assert [chunk.response.id for chunk in collected if chunk.type == "response.created"] == ["resp_fallback"] - assert isinstance(mock_fallback_utils.call_args.kwargs["e"], MidStreamFallbackError) - assert mock_fallback_utils.call_args.kwargs["kwargs"]["input"] == "Hello" - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_surfaces_transport_drop_when_no_fallback_lands(): - transport_error: Final = httpx.ReadError("Response payload is not completed") - router: Final = _make_router_with_fallback() - src: Final = _make_native_responses_iterator( - sse_payloads=_RESPONSES_LIFECYCLE_PAYLOADS, trailing_error=transport_error - ) - - async def reraise_trigger(**kwargs): - raise kwargs["e"] - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(side_effect=reraise_trigger) - ) as mock_fallback_utils: - wrapped: Final = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - outcome: Final = [item async for item in _events_until_error(wrapped)] - - assert [item.type for item in outcome[:-1]] == ["response.created", "response.in_progress"] - assert outcome[-1] is transport_error - assert mock_fallback_utils.await_count == 1 - trigger: Final = mock_fallback_utils.await_args.kwargs["e"] - assert isinstance(trigger, MidStreamFallbackError) - assert trigger.original_exception is transport_error - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_forwards_lifecycle_events_in_order_once_output_starts(): - router: Final = _make_router_with_fallback() - chunks: Final = [ - MagicMock(type="response.created"), - MagicMock(type="response.in_progress"), - MagicMock(type="response.output_text.delta"), - MagicMock(type="response.completed"), - ] - wrapped: Final = await router._aresponses_streaming_iterator( - response=_make_responses_iterator(chunks=chunks), - initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, - ) - - assert [chunk async for chunk in wrapped] == chunks - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_flushes_held_lifecycle_events_when_the_stream_ends_without_output(): - router: Final = _make_router_with_fallback() - chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.in_progress")] - wrapped: Final = await router._aresponses_streaming_iterator( - response=_make_responses_iterator(chunks=chunks), - initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, - ) - - assert [chunk async for chunk in wrapped] == chunks - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_forwards_held_lifecycle_events_before_a_non_fallback_error(): - router: Final = _make_router_with_fallback() - chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.in_progress")] - client_error: Final = litellm.BadRequestError(message="bad input", model="gpt-4", llm_provider="openai") - wrapped: Final = await router._aresponses_streaming_iterator( - response=_make_responses_iterator(chunks=chunks, error=client_error), - initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, - ) - outcome: Final = [item async for item in _events_until_error(wrapped)] - - assert outcome[:-1] == chunks - assert outcome[-1] is client_error - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_commits_held_lifecycle_events_at_the_hold_cap(): - router: Final = _make_router_with_fallback() - chunks: Final = [MagicMock(type="response.in_progress") for _ in range(MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS + 1)] - src: Final = _make_responses_iterator( - chunks=chunks, - error=MidStreamFallbackError( - message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True - ), - ) - fallback_chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.completed")] - - with patch.object(router, "async_function_with_fallbacks_common_utils", return_value=_AsyncList(fallback_chunks)): - wrapped: Final = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, - ) - collected: Final = [chunk async for chunk in wrapped] - - assert collected == [*chunks, *fallback_chunks] - - @pytest.mark.parametrize( ("event_type", "held_event_count", "expected"), [ @@ -7566,103 +3686,6 @@ def test_responses_stream_holds_event_holds_only_pre_output_lifecycle_events_und assert _responses_stream_holds_event(MagicMock(type=event_type), held_event_count) is expected -@pytest.mark.asyncio -async def test_aresponses_fallback_attempt_drops_held_lifecycle_events_when_a_fallback_lands(): - router: Final = _make_router_with_fallback() - trigger: Final = MidStreamFallbackError( - message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True - ) - held: Final = (MagicMock(type="response.created"), MagicMock(type="response.in_progress")) - fallback_chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.completed")] - adopt_headers: Final = MagicMock(return_value=({}, {})) - - with patch.object( - router, "async_function_with_fallbacks_common_utils", return_value=_AsyncList(fallback_chunks) - ) as mock_fallback_utils: - collected: Final = [ - chunk - async for chunk in router._aresponses_fallback_attempt( - trigger, - _make_responses_iterator(), - {"model": "gpt-4", "stream": True, "input": "Hello"}, - adopt_headers, - held, - ) - ] - - assert collected == fallback_chunks - adopt_headers.assert_called_once() - assert mock_fallback_utils.await_args.kwargs["e"] is trigger - - -@pytest.mark.asyncio -async def test_aresponses_fallback_attempt_replays_held_lifecycle_events_when_the_fallback_dies_before_its_first_event(): - """A fallback stream that raises before yielding anything announced no response of its own, so the - primary's held created/in_progress pair is replayed ahead of the error and the client sees the - announcement the failure belongs to, the same as when no fallback was attempted at all.""" - router: Final = _make_router_with_fallback() - trigger: Final = MidStreamFallbackError( - message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True - ) - held: Final = (MagicMock(type="response.created"), MagicMock(type="response.in_progress")) - fallback_error: Final = RuntimeError("fallback closed before its first event") - adopt_headers: Final = MagicMock(return_value=({}, {})) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_make_responses_iterator(error=fallback_error), - ): - outcome: Final = [ - item - async for item in _events_until_error( - router._aresponses_fallback_attempt( - trigger, - _make_responses_iterator(), - {"model": "gpt-4", "stream": True, "input": "Hello"}, - adopt_headers, - held, - ) - ) - ] - - assert outcome == [*held, fallback_error] - - -@pytest.mark.asyncio -async def test_aresponses_fallback_attempt_does_not_replay_held_lifecycle_events_once_the_fallback_announced_itself(): - """Once the fallback has yielded its own created event, a later failure must not replay the - primary's held pair on top of it, or the client would again see two announced response ids.""" - router: Final = _make_router_with_fallback() - trigger: Final = MidStreamFallbackError( - message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True - ) - held: Final = (MagicMock(type="response.created"), MagicMock(type="response.in_progress")) - fallback_created: Final = MagicMock(type="response.created") - fallback_error: Final = RuntimeError("fallback dropped after announcing itself") - adopt_headers: Final = MagicMock(return_value=({}, {})) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_make_responses_iterator(chunks=(fallback_created,), error=fallback_error), - ): - outcome: Final = [ - item - async for item in _events_until_error( - router._aresponses_fallback_attempt( - trigger, - _make_responses_iterator(), - {"model": "gpt-4", "stream": True, "input": "Hello"}, - adopt_headers, - held, - ) - ) - ] - - assert outcome == [fallback_created, fallback_error] - - _BACKGROUND_PRIMARY_URL: Final = "https://api.openai.com/v1/responses" _BACKGROUND_BACKUP_URL: Final = "https://backup.openai.test/v1/responses" @@ -7893,7 +3916,6 @@ async def test_aresponses_background_stream_surfaces_a_later_error_instead_of_a_ assert backup_route.call_count == 0 - _BACKGROUND_TEXT_DELTA: Final = { "type": "response.output_text.delta", "sequence_number": 2, @@ -7964,337 +3986,6 @@ async def test_aresponses_background_stream_content_policy_error_reads_the_same_ assert primary_route.call_count == 2 assert backup_route.call_count == 0 -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): - """Mid-stream error: preserve partial assistant output and request continuation in a user turn.""" - from litellm.exceptions import MidStreamFallbackError - - router = _make_router_with_fallback() - src = _make_responses_iterator( - chunks=[MagicMock(type="response.output_text.delta")], - error=MidStreamFallbackError( - message="socket reset mid-stream", - model="gpt-4", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="The capital of France is", - ), - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "What's the capital of France?", - "original_generic_function": litellm.aresponses, - }, - ) - async for _ in wrapped: - pass - - new_input = mock_fallback_utils.call_args.kwargs["kwargs"]["input"] - assert isinstance(new_input, list) - assert new_input[0]["role"] == "user" - assert new_input[0]["content"][0]["text"] == "What's the capital of France?" - assert new_input[1]["role"] == "assistant" - assert new_input[1]["content"][0]["type"] == "output_text" - assert new_input[1]["content"][0]["text"] == "The capital of France is" - assert new_input[2]["role"] == "user" - assert "do not repeat" in new_input[2]["content"][0]["text"].lower() - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_combines_partial_usage(): - """Partial usage from the bridge path is normalized to ResponseAPIUsage - and summed onto the fallback's response.completed event — no token-name - split, clean ResponseAPIUsage on output.""" - from types import SimpleNamespace - - from litellm.exceptions import MidStreamFallbackError - from litellm.types.llms.openai import ( - ResponseAPIUsage, - ResponseCompletedEvent, - ResponsesAPIResponse, - ResponsesAPIStreamEvents, - ) - - router = _make_router_with_fallback() - src = _make_responses_iterator( - bridge=True, - chat_chunks=[MagicMock()], - chunks=[MagicMock(type="response.output_text.delta")], - error=MidStreamFallbackError( - message="boom", - model="gpt-4", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="hello", - ), - ) - - fallback_response_object = ResponsesAPIResponse( - id="resp_test", created_at=0, model="gpt-4", object="response", output=[] - ) - fallback_response_object.usage = ResponseAPIUsage( - input_tokens=20, output_tokens=15, total_tokens=35 - ) - fallback_event = ResponseCompletedEvent( - type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - response=fallback_response_object, - ) - - with ( - patch( - "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace( - usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) - ), - ), - patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList([fallback_event]), - ), - ): - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "hi", - "original_generic_function": litellm.aresponses, - }, - ) - async for _ in wrapped: - pass - - merged = fallback_response_object.usage - assert isinstance(merged, ResponseAPIUsage) - assert merged.input_tokens == 30 # 10 (translated from prompt_tokens) + 20 - assert merged.output_tokens == 19 # 4 (translated from completion_tokens) + 15 - assert merged.total_tokens == 49 - - -def _midstream_rate_limit_error(): - rate_limit_error = litellm.RateLimitError( - message="vertex_ai_betaException - Resource exhausted.", - model="gemini", - llm_provider="vertex_ai_beta", - ) - midstream_error = MidStreamFallbackError( - message=str(rate_limit_error), - model="gemini", - llm_provider="vertex_ai_beta", - original_exception=rate_limit_error, - is_pre_first_chunk=True, - ) - return rate_limit_error, midstream_error - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks(): - """Regression for #26015: a mid-stream 429 with no fallbacks configured must - surface a clean RateLimitError, not leak the internal MidStreamFallbackError - wrapper to the client, and must terminate instead of hanging.""" - rate_limit_error, midstream_error = _midstream_rate_limit_error() - - router = litellm.Router( - model_list=[ - { - "model_name": "gemini", - "litellm_params": { - "model": "vertex_ai/gemini-2.0-flash", - "api_key": "fake-key", - }, - }, - ], - num_retries=0, - ) - - class _RaisingStream: - def __init__(self): - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - raise midstream_error - - stream = _RaisingStream() - setattr(stream, "model", "gemini") - setattr(stream, "custom_llm_provider", "vertex_ai_beta") - setattr(stream, "logging_obj", MagicMock()) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=midstream_error), - ): - result = await router._acompletion_streaming_iterator( - model_response=stream, - messages=[{"role": "user", "content": "Hello"}], - initial_kwargs={"model": "gemini", "stream": True}, - ) - - async def _consume(): - async for _ in result: - pass - - with pytest.raises(litellm.RateLimitError) as exc_info: - await asyncio.wait_for(_consume(), timeout=10) - - assert not isinstance(exc_info.value, MidStreamFallbackError) - assert exc_info.value.status_code == 429 - assert exc_info.value is rate_limit_error - - -def test_completion_streaming_iterator_surfaces_rate_limit_without_fallbacks(): - """Sync counterpart of - test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks.""" - rate_limit_error, midstream_error = _midstream_rate_limit_error() - - router = litellm.Router( - model_list=[ - { - "model_name": "gemini", - "litellm_params": { - "model": "vertex_ai/gemini-2.0-flash", - "api_key": "fake-key", - }, - }, - ], - num_retries=0, - ) - - class _RaisingSyncStream: - def __init__(self): - self.model = "gemini" - self.custom_llm_provider = "vertex_ai_beta" - self.logging_obj = MagicMock() - self.chunks = [] - - def __iter__(self): - return self - - def __next__(self): - raise midstream_error - - with patch.object( - router, - "function_with_fallbacks", - side_effect=midstream_error, - ): - result = router._completion_streaming_iterator( - model_response=_RaisingSyncStream(), - messages=[{"role": "user", "content": "Hello"}], - initial_kwargs={"model": "gemini", "stream": True}, - ) - - with pytest.raises(litellm.RateLimitError) as exc_info: - list(result) - - assert not isinstance(exc_info.value, MidStreamFallbackError) - assert exc_info.value.status_code == 429 - assert exc_info.value is rate_limit_error - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_surfaces_rate_limit_without_fallbacks(): - """Responses-API counterpart of - test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks.""" - rate_limit_error, midstream_error = _midstream_rate_limit_error() - - router = litellm.Router( - model_list=[ - { - "model_name": "gemini", - "litellm_params": { - "model": "vertex_ai/gemini-2.0-flash", - "api_key": "fake-key", - }, - }, - ], - num_retries=0, - ) - src = _make_responses_iterator(error=midstream_error, model="gemini") - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=midstream_error), - ): - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gemini", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - - async def _consume(): - async for _ in wrapped: - pass - - with pytest.raises(litellm.RateLimitError) as exc_info: - await asyncio.wait_for(_consume(), timeout=10) - - assert not isinstance(exc_info.value, MidStreamFallbackError) - assert exc_info.value.status_code == 429 - assert exc_info.value is rate_limit_error - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_common_utils(): - """Test the async_function_with_fallbacks_common_utils method""" - # Create a basic router for testing - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - }, - } - ], - max_fallbacks=5, - ) - - # Test case 1: disable_fallbacks=True should raise original exception - test_exception = Exception("Test error") - with pytest.raises(Exception, match="Test error"): - await router.async_function_with_fallbacks_common_utils( - e=test_exception, - disable_fallbacks=True, - fallbacks=None, - context_window_fallbacks=None, - content_policy_fallbacks=None, - model_group="gpt-3.5-turbo", - args=(), - kwargs=MagicMock(), - ) - - # Test case 2: original_model_group=None should raise original exception - with pytest.raises(Exception, match="Test error"): - await router.async_function_with_fallbacks_common_utils( - e=test_exception, - disable_fallbacks=False, - fallbacks=None, - context_window_fallbacks=None, - content_policy_fallbacks=None, - model_group="gpt-3.5-turbo", - args=(), - kwargs={}, # No model key - ) - def test_should_include_deployment(): """Test that Router._core_should_include_deployment returns the correct response""" @@ -8340,74 +4031,6 @@ def test_should_include_deployment(): ) -def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch): - """ - tiktoken token counting is the dominant on-loop cost for large prompts. When no - deployment in the group declares max_input_tokens, the count is never consumed, so - _pre_call_checks must not run it at all. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {}) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}}, - ] - result = router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - ) - - assert calls == [] - assert len(result) == 2 - - -def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch): - """ - When a deployment declares max_input_tokens the count must still run, be performed - at most once across the group (memoized), and filter deployments whose limit is - exceeded. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}}, - ] - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - ) - - assert calls == [1] - - _MULTILINGUAL_PROMPT: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] @@ -8465,21 +4088,6 @@ def _embedding_context_router(backend: str, window: int = _EMBEDDING_CONTEXT_WIN ) -async def test_context_window_pre_call_check_counts_an_embedding_image_as_media_on_an_image_backend(): - router: Final = _embedding_context_router("vertex_ai/gemini-embedding-2") - deployments: Final = router.registry().deployments("embedder") - - kept: Final = router._pre_call_checks(model="embedder", healthy_deployments=deployments, input=_EMBEDDING_IMAGE_INPUT) - counted: Final = await router._acount_pre_call_check_tokens( - model="embedder", healthy_deployments=deployments, messages=None, input=_EMBEDDING_IMAGE_INPUT, request_kwargs=None - ) - - assert len(kept) == 1 - assert counted is not None and 0 < counted < _EMBEDDING_CONTEXT_WINDOW - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks(model="embedder", healthy_deployments=deployments, input="word " * 20_000) - - @pytest.mark.parametrize( ("backend", "shape"), [ @@ -8515,24 +4123,12 @@ def test_context_window_pre_call_check_counts_an_embedding_image_as_text_on_a_te with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( - model="embedder", healthy_deployments=router.registry().deployments("embedder"), input=_EMBEDDING_IMAGE_INPUT + model="embedder", + healthy_deployments=router.registry().deployments("embedder"), + input=_EMBEDDING_IMAGE_INPUT, ) -async def test_acount_pre_call_check_tokens_counts_a_group_of_one_openai_model_with_that_models_tokenizer(): - router: Final = _context_limited_router("openai/gpt-6.1-sol") - - count: Final = await router._acount_pre_call_check_tokens( - model="prod", - healthy_deployments=router.get_model_list(model_name="prod"), - messages=_MULTILINGUAL_PROMPT, - input=None, - request_kwargs=None, - ) - - assert count == litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT) - - async def test_a_wildcard_route_checks_the_context_window_with_the_requested_openai_models_tokenizer(): limit: Final = ( litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT) @@ -8559,485 +4155,9 @@ async def test_a_wildcard_route_checks_the_context_window_with_the_requested_ope assert (sync_response.choices[0].message.content, async_response.choices[0].message.content) == ("ok", "ok") -def test_pre_call_checks_uses_precounted_tokens(monkeypatch): - """ - An async caller counts off the event loop and passes the result in. _pre_call_checks - must filter on that count instead of re-counting on the loop. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1 - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - input_token_count=1000, - ) - - assert calls == [] - - -async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(monkeypatch): - """ - The async deployment path must hand _pre_call_checks a count taken in a worker thread, - so a multi-MB prompt never blocks the proxy during deployment selection. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} - ) - - counting_threads = [] - monkeypatch.setattr( - litellm, - "token_counter", - lambda *a, **k: counting_threads.append(threading.current_thread()) or 42, - ) - - counts_passed_in = [] - original_pre_call_checks = router._pre_call_checks - - def spy(**kwargs): - counts_passed_in.append(kwargs.get("input_token_count")) - return original_pre_call_checks(**kwargs) - - monkeypatch.setattr(router, "_pre_call_checks", spy) - - result = await router.async_get_healthy_deployments( - model="m", - request_kwargs={}, - messages=[{"role": "user", "content": "hi"}], - input=None, - specific_deployment=False, - parent_otel_span=None, - ) - - assert len(result) == 1 - assert counts_passed_in == [42] - assert len(counting_threads) == 1 - assert counting_threads[0] is not threading.current_thread() - - -@pytest.mark.parametrize( - "model_info,expected", - [ - ({"max_input_tokens": 100}, True), - ({"max_input_tokens": None}, False), - ({}, False), - ], -) -def test_pre_call_checks_need_token_count(monkeypatch, model_info, expected): - """Only a deployment that declares an integer context window makes a token count worth taking.""" - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: model_info) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - assert router._pre_call_checks_need_token_count("m", deployments) is expected - - -def test_deployment_max_input_tokens_survives_an_unmappable_deployment(monkeypatch): - """ - _pre_call_checks skips a deployment it cannot resolve and carries on. The off-loop - pre-count must do the same, or an unmapped first deployment hides the limit declared by - a later one and the count lands back on the event loop. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - - def flaky_model_info(deployment, received_model_name, id=None): - if deployment["model_info"]["id"] == "unmapped": - raise ValueError("This model isn't mapped yet.") - return {"max_input_tokens": 100} - - monkeypatch.setattr(router, "_core_get_router_model_info", flaky_model_info) - - unmapped = {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "unmapped"}} - mapped = {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "mapped"}} - - assert router._deployment_max_input_tokens("m", unmapped) is None - assert router._deployment_max_input_tokens("m", mapped) == 100 - assert router._pre_call_checks_need_token_count("m", [unmapped, mapped]) is True - - -def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monkeypatch): - """ - When the off-loop count failed there is nothing left to filter on, so _pre_call_checks must - return the deployments unfiltered rather than repeating the count on the event loop. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - result = router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - input_token_count=None, - skip_inline_token_count=True, - ) - - assert calls == [] - assert len(result) == 1 - - -async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypatch): - """ - An off-loop count that raises must not send the same work back onto the event loop through - _pre_call_checks' inline fallback. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - counting_threads = [] - - def exploding_counter(*args, **kwargs): - counting_threads.append(threading.current_thread()) - raise ValueError("Invalid content item type: image") - - monkeypatch.setattr(litellm, "token_counter", exploding_counter) - - result = await router.async_get_healthy_deployments( - model="m", - request_kwargs={}, - messages=[{"role": "user", "content": "hi"}], - input=None, - specific_deployment=False, - parent_otel_span=None, - ) - - assert len(result) == 1 - assert len(counting_threads) == 1 - assert counting_threads[0] is not threading.current_thread() - - -async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypatch): - """ - A multi-MB prompt must not stall the proxy: a competing coroutine has to get - scheduled while the router's context-window count is in flight. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - ran = [] - - async def competitor(): - ran.append("competitor") - - task = asyncio.create_task(competitor()) - count = await router._acount_pre_call_check_tokens( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "A" * 512 * 1024}], - input=None, - request_kwargs=None, - ) - ran.append("count") - await task - - assert count is not None and count > 0 - assert ran == ["competitor", "count"] - - -async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monkeypatch): - """No deployment limits its context window, so there is nothing to count.""" - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {}) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) - - count = await router._acount_pre_call_check_tokens( - model="m", - healthy_deployments=[ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ], - messages=[{"role": "user", "content": "hi"}], - input=None, - request_kwargs=None, - ) - - assert count is None - assert calls == [] - - -def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): - """ - Responses API calls pass `input` (str) instead of `messages`. Context-window - checks must count tokens from `input` and filter deployments over the limit. Uses - the real token_counter so the transform + counting path is a true regression guard. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - input="a very long prompt that exceeds the tiny context window", - ) - - -def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): - """ - Responses API `input` can be a list of input items. It must be normalized to - chat messages and counted so oversized requests are filtered out. Uses the real - token_counter (no mock) so the transform + counting path is a true regression guard. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - input=[ - {"role": "user", "content": "count these tokens against the one token limit please"}, - ], - ) - - -def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): - """ - Responses API `instructions` become a system message the model receives, so their - tokens must be counted too. A request whose `input` alone fits under the limit but - whose `input` + `instructions` exceeds it must be filtered (regression for the - context-window check under-filtering when instructions were ignored). - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - - short_input = "hi" - long_instructions = "you are a helpful assistant. " * 20 - - input_only_tokens = router._count_pre_call_check_tokens(messages=None, input=short_input) - with_instructions_tokens = router._count_pre_call_check_tokens( - messages=None, input=short_input, request_kwargs={"instructions": long_instructions} - ) - assert with_instructions_tokens > input_only_tokens - - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} - ) - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - input=short_input, - request_kwargs={"instructions": long_instructions}, - ) - - _OVERSIZED_TOOL_DESCRIPTION = "look up the answer in the knowledge base. " * 40 -@pytest.mark.parametrize( - "prompt_kwargs, tool", - [ - pytest.param( - {"messages": [{"role": "user", "content": "hi"}]}, - { - "type": "function", - "function": { - "name": "lookup", - "description": _OVERSIZED_TOOL_DESCRIPTION, - "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, - }, - }, - id="chat_completions_tool", - ), - pytest.param( - {"input": "hi"}, - { - "type": "function", - "name": "lookup", - "description": _OVERSIZED_TOOL_DESCRIPTION, - "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, - }, - id="responses_tool", - ), - pytest.param( - {"messages": [{"role": "user", "content": "hi"}]}, - { - "name": "lookup", - "description": _OVERSIZED_TOOL_DESCRIPTION, - "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}}, - }, - id="anthropic_messages_tool", - ), - ], -) -def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwargs, tool): - """ - Tool definitions are sent to the model as prompt tokens but never appear in - `messages` or `input`. A request whose prompt alone fits the context window but - whose prompt plus `tools` exceeds it must be rejected before dispatch, for the - Chat Completions, Responses and Anthropic Messages tool shapes alike. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - - prompt_only_tokens = router._count_pre_call_check_tokens( - messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} - ) - - assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - request_kwargs={"tools": [tool]}, - **prompt_kwargs, - ) - - -@pytest.mark.parametrize( - "system", - [ - pytest.param("You are a meticulous assistant. " * 40, id="system_string"), - pytest.param( - [{"type": "text", "text": "You are a meticulous assistant. " * 40}], - id="system_blocks", - ), - ], -) -def test_pre_call_checks_counts_anthropic_system_tokens(monkeypatch, system): - """ - The Anthropic Messages API carries the system prompt as a top-level `system` field, - not as a message. Its tokens reach the model, so a request whose `messages` fit but - whose `messages` plus `system` exceed the context window must be rejected. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - messages = [{"role": "user", "content": "hi"}] - - messages_only_tokens = router._count_pre_call_check_tokens(messages=messages, input=None) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens}) - - assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, messages=messages)) == 1 - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=messages, - request_kwargs={"system": system}, - ) - - @pytest.mark.asyncio async def test_aanthropic_messages_enforces_context_window_with_system_and_tools(): """ @@ -9106,7 +4226,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): + with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): router._count_pre_call_check_tokens(messages=None, input=None) @@ -9133,37 +4253,6 @@ def test_count_pre_call_check_tokens_skips_the_history_a_compaction_item_replace ) -def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): - """ - When neither messages nor input is provided (e.g. endpoints without prompt text), - token counting is skipped gracefully and all deployments are returned. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - counted: list[dict] = [] - original = router._count_pre_call_check_tokens - monkeypatch.setattr( - router, - "_count_pre_call_check_tokens", - lambda **kwargs: counted.append(kwargs) or original(**kwargs), - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - result = router._pre_call_checks(model="m", healthy_deployments=deployments) - assert len(result) == 1 - assert counted == [] # token counting skipped entirely, so no misleading error is logged - - @pytest.mark.asyncio @pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state") async def test_aresponses_enforces_context_window_pre_call_check(): @@ -9237,9 +4326,7 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object( - litellm, "model_cost", {"test-custom-model": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -9247,15 +4334,11 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call( - model="gpt-3.5-turbo" - ) # base model call + mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -9266,26 +4349,18 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert ( - result["input_cost_per_token"] == 0.001 - ) # From custom model (overrides base 0.0015) - assert ( - result["output_cost_per_token"] == 0.002 - ) # From custom model (same as base) + assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) + assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert ( - result["mode"] == "chat" - ) # From base model (overrides litellm "completion") + assert result["mode"] == "chat" # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert ( - result["key"] == "gpt-3.5-turbo" - ) # From base model (overrides litellm key) + assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -9304,9 +4379,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model-no-base", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -9326,9 +4399,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="non-existent-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -9361,9 +4432,7 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info( - model_id="test-custom-model-invalid", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") # Should handle exception gracefully and still return merged result assert result is not None @@ -9372,12 +4441,8 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object( - litellm, "get_model_info", side_effect=Exception("Not found") - ): - result = router.get_deployment_model_info( - model_id="non-existent", model_name="non-existent" - ) + with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): + result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") # Should return None when no model info is found assert result is None @@ -9400,9 +4465,7 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -9438,15 +4501,11 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info( - model_id="custom-with-base", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") # Should return custom_model_info merged with base model info assert result is not None - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom (overrides base) + assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -9493,18 +4552,14 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object( - litellm, "model_cost", {"custom-model-id": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") assert result is not None @@ -9514,29 +4569,17 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom model (overrides base 0.03) - assert ( - result["max_tokens"] == 8000 - ) # From custom model (overrides base 4096) + assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) + assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert ( - result["output_cost_per_token"] == 0.06 - ) # From base model (not in custom) - assert ( - result["litellm_provider"] == "openai" - ) # From base model (not in custom) - assert ( - result["base_only_field"] == "base_value" - ) # From base model (not in custom) + assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) + assert result["litellm_provider"] == "openai" # From base model (not in custom) + assert result["base_only_field"] == "base_value" # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert ( - result["mode"] == "completion" - ) # From litellm model name info (not in custom or base) + assert result["mode"] == "completion" # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -9548,226 +4591,6 @@ def test_get_deployment_model_info_base_model_merge_priority(): print("✓ Base model merge priority test passed!") -@pytest.mark.parametrize( - "model, litellm_params, endpoint, expected", - [ - ( - "gpt", - {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.services.ai.azure.com", "api_key": "key"}, - "gpt/openai/deployments/gpt-5.4-mini/chat/completions", - "gpt-5.4-mini/openai/deployments/gpt-5.4-mini/chat/completions", - ), - ( - "aws/anthropic/bedrock-claude", - {"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"}, - "/model/aws/anthropic/bedrock-claude/invoke", - "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", - ), - ( - "my-gemini", - {"model": "gemini/gemini-3.1-pro-preview", "api_key": "key"}, - "v1beta/models/my-gemini:streamGenerateContent", - "v1beta/models/gemini-3.1-pro-preview:streamGenerateContent", - ), - ], -) -def test_add_deployment_model_to_endpoint_rewrites_the_model_group_only_as_whole_path_segments( - model, litellm_params, endpoint, expected -): - router = litellm.Router(model_list=[{"model_name": model, "litellm_params": litellm_params}]) - - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs={"endpoint": endpoint}, model=model, model_name=litellm_params["model"] - ) - - assert result["endpoint"] == expected - - -def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): - """ - Test that _add_deployment_model_to_endpoint_for_llm_passthrough_route correctly strips bedrock provider prefix - """ - router = litellm.Router( - model_list=[ - { - "model_name": "special-bedrock-model", - "litellm_params": { - "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - }, - } - ], - ) - - # Test Case 1: Bedrock model with provider prefix - should strip "bedrock/" prefix - kwargs = { - "endpoint": "/model/special-bedrock-model/invoke", - "custom_llm_provider": "bedrock", - } - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, - model="special-bedrock-model", - model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" - ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" - - # Test Case 2: Bedrock invoke-with-response-stream endpoint - kwargs = { - "endpoint": "/model/special-bedrock-model/invoke-with-response-stream", - "custom_llm_provider": "bedrock", - } - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, - model="special-bedrock-model", - model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" - ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" - - # Test Case 3: Bedrock converse endpoint - kwargs = { - "endpoint": "/model/bedrock-model/converse", - "custom_llm_provider": "bedrock", - } - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, - model="bedrock-model", - model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", - ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" - - # Test Case 4: Bedrock provider prefix auto-detected from model_name - kwargs = { - "endpoint": "/model/router-model/invoke", - } - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, - model="router-model", - model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", - ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" - - -def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): - router = litellm.Router( - model_list=[ - { - "model_name": "my-bedrock-model", - "litellm_params": { - "model": "bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0", - }, - } - ], - ) - deployment = router.model_list[0] - kwargs: dict = {} - - with patch( - "litellm.proxy.proxy_server.general_settings", - {"pass_through_request_timeout": 6}, - ): - router._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="_ageneric_api_call_with_fallbacks", - ) - - assert kwargs["timeout"] == 6.0 - - -def _passthrough_timeout(router: litellm.Router, deployment: dict, stream: bool) -> float: - kwargs: Final[dict] = {"stream": stream} - router._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="_ageneric_api_call_with_fallbacks", - ) - return kwargs["timeout"] - - -def test_update_kwargs_with_deployment_passthrough_honors_stream_timeout(): - router = litellm.Router( - model_list=[ - { - "model_name": "anthropic-with-stream-timeout", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-5", - "api_key": "fake-key", - "timeout": 60, - "stream_timeout": 1800, - }, - }, - { - "model_name": "anthropic-router-default", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-5", - "api_key": "fake-key", - "timeout": 60, - }, - }, - ], - timeout=120, - stream_timeout=900, - ) - per_deployment, router_default = router.model_list - - assert _passthrough_timeout(router, per_deployment, stream=True) == 1800.0 - assert _passthrough_timeout(router, router_default, stream=True) == 900.0 - assert _passthrough_timeout(router, per_deployment, stream=False) == 60.0 - assert _passthrough_timeout(router, router_default, stream=False) == 60.0 - - -def test_update_kwargs_with_deployment_passthrough_router_stream_timeout_sources(): - deployment: Final[dict] = { - "model_name": "anthropic-router-default", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "fake-key"}, - } - string_router = litellm.Router(model_list=[deployment], timeout=120, stream_timeout="900") - default_router = litellm.Router( - model_list=[deployment], - timeout=120, - default_litellm_params={"stream_timeout": 700}, - ) - - assert _passthrough_timeout(string_router, string_router.model_list[0], stream=True) == 900.0 - assert _passthrough_timeout(default_router, default_router.model_list[0], stream=True) == 700.0 - assert _passthrough_timeout(default_router, default_router.model_list[0], stream=False) == 120.0 - - -def test_update_kwargs_with_deployment_passthrough_honors_global_request_timeout(monkeypatch: pytest.MonkeyPatch): - """litellm_settings.request_timeout must bound the native responses route when neither the - deployment nor the router carries a timeout, while a deployment timeout keeps winning.""" - monkeypatch.setattr("litellm.request_timeout", 44.0, raising=False) - monkeypatch.setattr("litellm.request_timeout_explicitly_set", True, raising=False) - router: Final = litellm.Router( - model_list=[ - { - "model_name": "responses-global-timeout", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key"}, - }, - { - "model_name": "responses-deployment-timeout", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key", "timeout": 3}, - }, - ], - ) - global_only, per_deployment = router.model_list - - with patch("litellm.proxy.proxy_server.general_settings", {"pass_through_request_timeout": 6}): - assert _passthrough_timeout(router, global_only, stream=True) == 44.0 - assert _passthrough_timeout(router, global_only, stream=False) == 44.0 - assert _passthrough_timeout(router, per_deployment, stream=True) == 3.0 - assert _passthrough_timeout(router, per_deployment, stream=False) == 3.0 - - @pytest.mark.asyncio async def test_router_acompletion_with_unknown_model_and_default_fallback(): """ @@ -9790,14 +4613,10 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [ - {"role": "user", "content": "This call should succeed by falling back."} - ] + messages = [{"role": "user", "content": "This call should succeed by falling back."}] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion( - model="completely-unknown-model", messages=messages - ) + response = await router.acompletion(model="completely-unknown-model", messages=messages) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -9908,100 +4727,6 @@ def test_get_credential_deployment_skips_a_paused_deployment(): assert router.get_credential_deployment(model_id="paused-dep") is None -def test_get_team_public_name_deployment_only_resolves_the_owning_team(): - router = litellm.Router( - model_list=[ - { - "model_name": "mistral/mistral-ocr-latest", - "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "sk-team-a"}, - "model_info": {"id": "team-a-ocr", "team_id": "team-a", "team_public_model_name": "ocr"}, - } - ] - ) - - owning_team = router._get_team_public_name_deployment(model_id="ocr", team_id="team-a") - - assert owning_team is not None and owning_team.model_info.id == "team-a-ocr" - assert router._get_team_public_name_deployment(model_id="ocr", team_id="team-b") is None - assert router._get_team_public_name_deployment(model_id="ocr", team_id=None) is None - assert router.get_credential_deployment(model_id="ocr", team_id="team-a").model_info.id == "team-a-ocr" - assert router.get_credential_deployment(model_id="ocr", team_id="team-b") is None - - -def test_get_wildcard_deployment_usable_by_team_prefers_the_team_pattern(): - router = litellm.Router( - model_list=[ - { - "model_name": "mistral/*", - "litellm_params": {"model": "mistral/*", "api_key": "sk-shared"}, - "model_info": {"id": "shared-wildcard"}, - }, - { - "model_name": "mistral/*", - "litellm_params": {"model": "mistral/*", "api_key": "sk-team-a"}, - "model_info": {"id": "team-a-wildcard", "team_id": "team-a", "team_public_model_name": "mistral/*"}, - }, - ] - ) - ocr = "mistral/mistral-ocr-latest" - - team_match = router._get_wildcard_deployment_usable_by_team(model_id=ocr, team_id="team-a") - other_team_match = router._get_wildcard_deployment_usable_by_team(model_id=ocr, team_id="team-b") - anonymous_match = router._get_wildcard_deployment_usable_by_team(model_id=ocr, team_id=None) - - assert team_match is not None and team_match.model_info.id == "team-a-wildcard" - assert other_team_match is not None and other_team_match.model_info.id == "shared-wildcard" - assert anonymous_match is not None and anonymous_match.model_info.id == "shared-wildcard" - assert router._get_wildcard_deployment_usable_by_team(model_id="openai/gpt-5.6", team_id="team-a") is None - assert router.get_credential_deployment(model_id=ocr, team_id="team-b").model_info.id == "shared-wildcard" - - -@pytest.mark.parametrize( - ("with_catch_all", "requested", "team_id", "expected_ids"), - [ - pytest.param(False, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-when-no-proxy-wide-match"), - pytest.param(False, "rwild-chat", "team-b", (), id="another-teams-pattern-never-serves"), - pytest.param(False, "rwild-chat", None, (), id="no-team-no-team-pattern"), - pytest.param(True, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-beats-the-proxy-wide-pattern"), - pytest.param(True, "rwild-chat", "team-b", ("catch-all",), id="another-team-falls-through-to-proxy-wide"), - pytest.param(True, "ptu-chat", "team-r", ("ptu-wildcard",), id="only-the-most-specific-proxy-wide-pattern"), - ], -) -def test_wildcard_route_deployments_is_the_one_pattern_routing_serves_from( - with_catch_all: bool, requested: str, team_id: str | None, expected_ids: tuple[str, ...] -): - catch_all: Final = ( - ({"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-open"}, "model_info": {"id": "catch-all"}},) - if with_catch_all - else () - ) - router: Final = litellm.Router( - model_list=[ - { - "model_name": "ptu-*", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-ptu"}, - "model_info": {"id": "ptu-wildcard"}, - }, - { - "model_name": "team-r-internal", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-team-r"}, - "model_info": {"id": "team-r-wildcard", "team_id": "team-r", "team_public_model_name": "rwild-*"}, - }, - *catch_all, - ] - ) - - served: Final = router.wildcard_route_deployments(requested, team_id) - early: Final = router._try_early_resolve_deployments_for_model_not_in_names(model=requested, request_team_id=team_id) - early_deployments: Final = early[1] if early is not None and isinstance(early[1], list) else [] - - assert tuple(Deployment.model_validate(d).model_info.id for d in served) == expected_ids - assert tuple(Deployment.model_validate(d).model_info.id for d in early_deployments) == expected_ids - assert router.wildcard_routed_models(requested, team_id) == frozenset( - Deployment.model_validate(d).litellm_params.model for d in served - ) - - @pytest.mark.parametrize( ("team_id", "expected_id"), [ @@ -10055,15 +4780,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-claude-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") assert credentials is not None - assert ( - credentials["aws_bedrock_runtime_endpoint"] - == "https://bedrock-runtime.us-east-1.amazonaws.com" - ) + assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -10090,9 +4810,7 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="vertex-gemini" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -10177,9 +4895,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="azure-gpt-4" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -10216,9 +4932,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -10262,9 +4976,7 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None for key, value in aws_auth_params.items(): @@ -10334,15 +5046,11 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2" - ) + global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -10383,15 +5091,11 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro" - ) + unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -10418,81 +5122,8 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert ( - router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-a" - ) - is None - ) - assert ( - router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") - is None - ) - - -def test_deployment_usable_by_team_helpers(): - """ - Direct coverage of the team-ownership filter: a team-scoped deployment is - usable only by its owning team, shared deployments by anyone, and the - model-group picker returns the first usable deployment or None. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gemini-2.5-pro", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-pro", - "vertex_project": "team-b-project", - }, - "model_info": { - "id": "team-b-vertex", - "team_id": "team-b", - "team_public_model_name": "gemini-2.5-pro", - }, - }, - { - "model_name": "gemini-2.5-pro", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-pro", - "vertex_project": "shared-project", - }, - }, - ], - ) - - team_owned, shared = router.model_list - assert router._deployment_usable_by_team(team_owned, "team-b") is True - assert router._deployment_usable_by_team(team_owned, "team-a") is False - assert router._deployment_usable_by_team(team_owned, None) is False - assert router._deployment_usable_by_team(shared, "team-a") is True - assert router._deployment_usable_by_team(shared, None) is True - - picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-a" - ) - assert picked is not None - assert picked.litellm_params.vertex_project == "shared-project" - - owner_picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-b" - ) - assert owner_picked is not None - assert owner_picked.litellm_params.vertex_project == "team-b-project" - - assert ( - router._get_model_group_deployment_usable_by_team( - model_group_name="unknown-model", team_id="team-a" - ) - is None - ) - - -def test_deployment_usable_by_team_uses_dynamic_model_info_get(): - class ModelInfo: - def get(self, key: str) -> str | None: - return "team-b" if key == "team_id" else None - - assert litellm.Router._core_deployment_usable_by_team({"model_info": ModelInfo()}, "team-a") is False + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): @@ -10524,9 +5155,7 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -10538,21 +5167,11 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is not None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None router.delete_deployment(id="team-wildcard-id") - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): @@ -10625,22 +5244,13 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment( - deployment=Deployment(**_team_wildcard_model(api_key="new-key")) - ) - credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) + credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_get_available_guardrail_single_deployment(): @@ -10720,43 +5330,6 @@ def test_get_available_guardrail_not_found(): router.get_available_guardrail(guardrail_name="non-existent") -@pytest.mark.asyncio -async def test_aguardrail_helper(): - """ - Test _aguardrail_helper selects a guardrail and executes the original function. - """ - guardrail_config = { - "guardrail_name": "content-filter", - "litellm_params": {"guardrail": "custom", "mode": "pre_call"}, - "id": "guardrail-1", - } - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo"}, - } - ], - guardrail_list=[guardrail_config], - ) - - # Mock the original function - async def mock_original_function(**kwargs): - return { - "result": "success", - "selected_guardrail": kwargs.get("selected_guardrail"), - } - - result = await router._aguardrail_helper( - model="content-filter", - original_generic_function=mock_original_function, - ) - - assert result["result"] == "success" - assert result["selected_guardrail"] == guardrail_config - - @pytest.mark.asyncio async def test_aguardrail(): """ @@ -10807,7 +5380,6 @@ async def test_anthropic_messages_call_type_is_cached(): PromptCachingDeploymentCheck, ) from litellm.types.utils import ( - CallTypes, StandardLoggingHiddenParams, StandardLoggingMetadata, StandardLoggingModelInformation, @@ -10828,9 +5400,7 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", custom_llm_provider="anthropic", model_id="model-123", @@ -10910,123 +5480,8 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert ( - cached_result is not None - ), "Model ID should be cached for anthropic_messages call type" - assert ( - cached_result["model_id"] == test_model_id - ), f"Expected {test_model_id}, got {cached_result['model_id']}" - - -def test_update_kwargs_with_deployment_propagates_model_tags(): - """ - Test that deployment-level tags from litellm_params are merged into - kwargs metadata when _update_kwargs_with_deployment is called. - - This ensures model-level tags defined in config.yaml appear in SpendLogs. - See: https://github.com/BerriAI/litellm/issues/XXXX - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "fake-key", - "tags": ["openai-account", "production"], - }, - }, - ], - ) - - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Deployment tags should be propagated to kwargs metadata - assert "tags" in kwargs["metadata"] - assert "openai-account" in kwargs["metadata"]["tags"] - assert "production" in kwargs["metadata"]["tags"] - - -def test_update_kwargs_with_deployment_merges_many_client_tags_in_linear_time(): - router: Final = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key", "tags": ["deployment-tag"]}, - }, - ], - ) - client_tags: Final = [f"client-{index}" for index in range(100_000)] - kwargs: Final = {"metadata": {"tags": client_tags}} - deployment: Final = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") - - started: Final = time.perf_counter() - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert time.perf_counter() - started < 2 - assert kwargs["metadata"]["tags"] == [*client_tags, "deployment-tag"] - - -def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): - """ - Test that when both request-level and deployment-level tags exist, - they are merged without duplicates. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "fake-key", - "tags": ["openai-account", "shared-tag"], - }, - }, - ], - ) - - # Simulate request that already has tags (from request body or key/team level) - kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Both sources should be merged, no duplicates - assert "user-tag" in kwargs["metadata"]["tags"] - assert "openai-account" in kwargs["metadata"]["tags"] - assert "shared-tag" in kwargs["metadata"]["tags"] - assert kwargs["metadata"]["tags"].count("shared-tag") == 1 - - -def test_update_kwargs_with_deployment_no_tags(): - """ - Test that when deployment has no tags, kwargs metadata is not affected. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "fake-key", - }, - }, - ], - ) - - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # No tags key should be added if deployment has no tags - assert "tags" not in kwargs["metadata"] + assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" + assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" @pytest.mark.asyncio @@ -11073,581 +5528,6 @@ async def test_retry_does_not_narrow_tag_filtered_group_to_failed_deployments_ta assert metadata["tags"] == ["free"] -def test_update_kwargs_with_deployment_merges_tools(): - """ - Test that when both deployment litellm_params and request have tools, - they are merged (deployment tools first, then request tools). - - Supports proxy-configured tools (e.g. for o3 deep research) merged with - client-provided tools. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - "tool_choice": "auto", - }, - }, - ], - ) - - kwargs: dict = { - "metadata": {}, - "tools": [ - { - "type": "function", - "function": {"name": "get_weather", "description": "Get weather"}, - }, - ], - } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Tools should be merged: deployment first, then request - assert "tools" in kwargs - assert len(kwargs["tools"]) == 2 - assert kwargs["tools"][0] == {"type": "web_search"} - assert kwargs["tools"][1]["function"]["name"] == "get_weather" - # tool_choice from request (none) - deployment's should be used - assert kwargs["tool_choice"] == "auto" - - -def test_update_kwargs_with_deployment_merge_tools_deployment_only(): - """ - Test that when only deployment has tools, they are applied to kwargs. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - "tool_choice": "required", - }, - }, - ], - ) - - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert kwargs["tools"] == [{"type": "web_search"}] - assert kwargs["tool_choice"] == "required" - - -def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice(): - """ - Test that when request has tool_choice, it overrides deployment's. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - "tool_choice": "auto", - }, - }, - ], - ) - - kwargs: dict = { - "metadata": {}, - "tool_choice": "none", - } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Request tool_choice should be preserved (merged tools still applied) - assert kwargs["tool_choice"] == "none" - - -@pytest.mark.parametrize( - "deployment_params, router_defaults, request_params, expected_tools", - [ - ({"drop_params": True}, {}, {}, [{"type": "web_search"}]), - ({}, {"drop_params": True}, {}, [{"type": "web_search"}]), - ({}, {}, {"additional_drop_params": ["tools"]}, [{"type": "web_search"}]), - ({"drop_params": True}, {}, {"drop_params": False}, 5), - ({}, {}, {}, 5), - ], - ids=[ - "deployment_drop_params", - "router_default_drop_params", - "request_drops_tools", - "request_overrides_deployment_drop", - "no_drop_policy", - ], -) -def test_update_kwargs_with_deployment_keeps_deployment_tools_when_malformed_request_tools_drop( - deployment_params: dict[str, object], - router_defaults: dict[str, object], - request_params: dict[str, object], - expected_tools: object, -) -> None: - router: Final = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - **deployment_params, - }, - }, - ], - default_litellm_params={**router_defaults}, - ) - kwargs: Final[dict[str, object]] = {"metadata": {}, "tools": 5, **request_params} - router._update_kwargs_with_deployment( - deployment=router.get_deployment_by_model_group_name(model_group_name="o3-deep-research"), kwargs=kwargs - ) - - assert kwargs["tools"] == expected_tools - - -_NAMED_REQUEST_TOOL_CHOICE: Final = {"type": "function", "function": {"name": "lookup_order"}} - - -@pytest.mark.parametrize( - "deployment_params, expected_tool_fields", - [ - ( - {"drop_params": True, "tool_choice": "required"}, - {"tools": [{"type": "web_search"}], "tool_choice": "required"}, - ), - ({"drop_params": True}, {"tools": [{"type": "web_search"}]}), - ({}, {"tools": 5, "tool_choice": _NAMED_REQUEST_TOOL_CHOICE, "parallel_tool_calls": False}), - ], - ids=["deployment_tool_choice_applies", "no_deployment_tool_choice", "no_drop_policy"], -) -def test_update_kwargs_with_deployment_drops_request_tool_settings_with_malformed_request_tools( - deployment_params: dict[str, object], expected_tool_fields: dict[str, object] -) -> None: - router: Final = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - **deployment_params, - }, - }, - ], - ) - kwargs: Final[dict[str, object]] = { - "metadata": {}, - "tools": 5, - "tool_choice": _NAMED_REQUEST_TOOL_CHOICE, - "parallel_tool_calls": False, - } - router._update_kwargs_with_deployment( - deployment=router.get_deployment_by_model_group_name(model_group_name="o3-deep-research"), kwargs=kwargs - ) - - tool_fields: Final = ("tools", "tool_choice", "parallel_tool_calls") - assert {name: kwargs[name] for name in tool_fields if name in kwargs} == expected_tool_fields - - -def test_credential_name_injected_as_tag(): - """ - Test that litellm_credential_name from deployment litellm_params - is injected as a tag into metadata during _update_kwargs_with_deployment. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "xai-model", - "litellm_params": { - "model": "xai/grok-4-1-fast", - "litellm_credential_name": "xAI", - }, - } - ], - ) - - kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert "Credential: xAI" in kwargs["metadata"]["tags"] - assert "A.101" in kwargs["metadata"]["tags"] - - -def test_credential_name_not_duplicated_in_tags(): - """ - Test that if the credential tag already exists in the tags list, - it is not duplicated. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "xai-model", - "litellm_params": { - "model": "xai/grok-4-1-fast", - "litellm_credential_name": "xAI", - }, - } - ], - ) - - kwargs: dict = {"metadata": {"tags": ["Credential: xAI", "A.101"]}} - deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert kwargs["metadata"]["tags"].count("Credential: xAI") == 1 - - -def test_credential_name_not_injected_when_absent(): - """ - Test that when no litellm_credential_name is set, tags are unchanged. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-model", - "litellm_params": { - "model": "gpt-4o", - }, - } - ], - ) - - kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-model") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert kwargs["metadata"]["tags"] == ["A.101"] - - -def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): - """For generic_api_call, model_info with pricing must go to litellm_metadata. - - Routes like /messages and /responses use generic_api_call which stores - model_info under litellm_metadata. Regression test for #23185. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "claude-sonnet-4", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-20250514", - "api_key": "fake-key", - }, - "model_info": { - "id": "custom-pricing-id", - "input_cost_per_token": 0.0003, - "output_cost_per_token": 0.0015, - }, - }, - ], - ) - - kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="generic_api_call" - ) - - assert "litellm_metadata" in kwargs - model_info = kwargs["litellm_metadata"]["model_info"] - assert model_info["id"] == "custom-pricing-id" - assert model_info["input_cost_per_token"] == 0.0003 - assert model_info["output_cost_per_token"] == 0.0015 - - -def test_update_kwargs_with_deployment_model_info_in_metadata(): - """For acompletion (function_name=None), model_info goes to metadata. - - /chat/completions uses acompletion which stores model_info under metadata. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "claude-sonnet-4", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-20250514", - "api_key": "fake-key", - }, - "model_info": { - "id": "custom-pricing-id", - "input_cost_per_token": 0.0003, - "output_cost_per_token": 0.0015, - }, - }, - ], - ) - - kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name=None - ) - - assert "metadata" in kwargs - model_info = kwargs["metadata"]["model_info"] - assert model_info["id"] == "custom-pricing-id" - assert model_info["input_cost_per_token"] == 0.0003 - assert model_info["output_cost_per_token"] == 0.0015 - - -def test_update_kwargs_with_deployment_clears_pod_routing_on_non_discovery_fallback() -> None: - from litellm.constants import KUBERNETES_POD_ROUTING_KEY - - routing_record: Final = { - "service_host": "vllm-headless.ns.svc.cluster.local", - "pod_ip": "10.0.0.1", - "pod_count": 3, - "selection": "session_affinity", - } - discovery_deployment: Final = { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, - "model_info": {"id": "discovery-id"}, - KUBERNETES_POD_ROUTING_KEY: routing_record, - } - fallback_deployment: Final = { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, - "model_info": {"id": "fallback-id"}, - } - router: Final = litellm.Router(model_list=[discovery_deployment]) - kwargs: Final = {"metadata": {}} - - router._update_kwargs_with_deployment(deployment=discovery_deployment, kwargs=kwargs) - assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] == routing_record - - router._update_kwargs_with_deployment(deployment=fallback_deployment, kwargs=kwargs) - assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] is None - - -def test_update_kwargs_with_deployment_synchronizes_existing_litellm_metadata_pod_routing() -> None: - from litellm.constants import KUBERNETES_POD_ROUTING_KEY - - routing_record: Final = { - "service_host": "vllm-headless.ns.svc.cluster.local", - "pod_ip": "10.0.0.1", - "pod_count": 3, - "selection": "session_affinity", - } - spoofed_record: Final = { - "service_host": "vllm-headless.ns.svc.cluster.local", - "pod_ip": "10.0.0.9", - "pod_count": 3, - "selection": "round_robin", - } - deployment: Final = { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, - "model_info": {"id": "discovery-id"}, - KUBERNETES_POD_ROUTING_KEY: routing_record, - } - original_litellm_metadata: Final = {KUBERNETES_POD_ROUTING_KEY: spoofed_record} - kwargs: Final = {"metadata": {}, "litellm_metadata": original_litellm_metadata} - router: Final = litellm.Router(model_list=[deployment]) - - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="completion") - - assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] == routing_record - assert kwargs["litellm_metadata"][KUBERNETES_POD_ROUTING_KEY] == routing_record - assert kwargs["litellm_metadata"] is not original_litellm_metadata - assert original_litellm_metadata[KUBERNETES_POD_ROUTING_KEY] == spoofed_record - - -def test_update_kwargs_with_non_discovery_deployment_clears_other_pod_routing_bucket() -> None: - from litellm.constants import KUBERNETES_POD_ROUTING_KEY - - previous_record: Final = { - "service_host": "vllm-headless.ns.svc.cluster.local", - "pod_ip": "10.0.0.1", - "pod_count": 3, - "selection": "session_affinity", - } - deployment: Final = { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, - "model_info": {"id": "fallback-id"}, - } - original_litellm_metadata: Final = {KUBERNETES_POD_ROUTING_KEY: previous_record} - kwargs: Final = {"metadata": {}, "litellm_metadata": original_litellm_metadata} - router: Final = litellm.Router(model_list=[deployment]) - - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="completion") - - assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] is None - assert kwargs["litellm_metadata"][KUBERNETES_POD_ROUTING_KEY] is None - assert original_litellm_metadata[KUBERNETES_POD_ROUTING_KEY] == previous_record - - -def test_combine_fallback_usage(): - """Test that _combine_fallback_usage merges partial and fallback usage.""" - from litellm.router import Router - from litellm.types.utils import Usage - - # Create a stream chunk with usage - chunk = litellm.ModelResponseStream( - id="test", - model="gpt-4o", - choices=[], - usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - ) - - # Call _combine_fallback_usage with no extra usage - Router._combine_fallback_usage(chunk, None) - assert chunk.usage is not None - assert chunk.usage.prompt_tokens == 10 - assert chunk.usage.completion_tokens == 5 - assert chunk.usage.total_tokens == 15 - - Router._combine_fallback_usage(chunk, Usage(prompt_tokens=0, completion_tokens=4, total_tokens=4)) - assert chunk.usage.prompt_tokens == 10 - assert chunk.usage.completion_tokens == 9 - assert chunk.usage.total_tokens == 19 - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure(): - """A mid-stream failure with no successful fallback raises and is logged as - a failure, so the router must never dispatch it as a success. Partial-spend - recovery for the failure row happens in the streaming handler, not here, so - this guards only against reintroducing a success log for a failed stream. - """ - from litellm.exceptions import MidStreamFallbackError - from litellm.types.utils import Delta, StreamingChoices, Usage - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, - }, - ], - set_verbose=True, - ) - - error = MidStreamFallbackError( - message="Connection lost", - model="gpt-4", - llm_provider="openai", - generated_content="The Roman Empire began when", - ) - - def _make_interrupted_model_response(): - partial_chunk = litellm.ModelResponseStream( - id="chatcmpl-partial-1", - created=1742056047, - model="gpt-4", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(content="The Roman Empire began when", role="assistant"), - ) - ], - usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), - ) - - class _RaisingStream: - def __init__(self): - self.index = 0 - self.chunks = [partial_chunk] - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index == 0: - self.index += 1 - return partial_chunk - raise error - - stream = _RaisingStream() - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.model_call_details = {} - setattr(stream, "model", "gpt-4") - setattr(stream, "custom_llm_provider", "openai") - setattr(stream, "logging_obj", logging_obj) - return stream, logging_obj - - messages = [{"role": "user", "content": "Hello"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - # Terminal path: no successful fallback -> the error propagates and the - # router never dispatches a success for the failed stream. - model_response, logging_obj = _make_interrupted_model_response() - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=error), - ): - result = await router._acompletion_streaming_iterator( - model_response=model_response, - messages=messages, - initial_kwargs=dict(initial_kwargs), - ) - collected = [] - async def _drain(): - async for chunk in result: - collected.append(chunk) - - with pytest.raises(MidStreamFallbackError): - await _drain() - - assert len(collected) == 1 - logging_obj.dispatch_success_handlers.assert_not_called() - - # Mid-stream errors with generated content are now re-raised immediately; - # no continuation-prompt fallback is attempted. Success handlers must - # still not be dispatched in this path. - model_response, logging_obj = _make_interrupted_model_response() - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=model_response, - messages=messages, - initial_kwargs=dict(initial_kwargs), - ) - collected = [] - async def _drain(): - async for chunk in result: - collected.append(chunk) - - with pytest.raises(MidStreamFallbackError): - await _drain() - - assert len(collected) == 1, "only the partial chunk before the error" - mock_fallback.assert_not_called() - logging_obj.dispatch_success_handlers.assert_not_called() - - @pytest.mark.asyncio async def test_team_scoped_model_fallback(): """ @@ -11780,7 +5660,12 @@ def _shared_ptu_deployment(model_name: str, deployment_id: str, **litellm_params """A one-PTU Azure deployment whose whole reservation is team-a's share.""" return { "model_name": model_name, - "litellm_params": {"model": "azure/gpt-4.1", "api_key": "fake", "api_base": "https://ptu.example", **litellm_params}, + "litellm_params": { + "model": "azure/gpt-4.1", + "api_key": "fake", + "api_base": "https://ptu.example", + **litellm_params, + }, "model_info": { "id": deployment_id, "base_model": "azure/gpt-4.1", @@ -11890,23 +5775,17 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert ( - len(deployment_ids) == 2 - ), "Each deployment must have a unique ID for cooldown tracking" + assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="other-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") assert len(deployments) == 0 @@ -11951,12 +5830,8 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) - assert ( - len(deployments) == 2 - ), "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -11971,436 +5846,6 @@ async def test_multiregion_team_failover_between_regions(): ] -def test_access_group_scoped_key_filters_deployments_with_same_public_model(): - """ - If a key can access a model only via access group membership, - router candidate deployments for that public model should be constrained - to deployments in the allowed access group. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-5.1", - "api_key": "key1", - "mock_response": "response-via-AG1", - }, - "model_info": {"access_groups": ["AG1"]}, - }, - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "key2", - "mock_response": "response-via-AG2", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - ] - ) - - scoped_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG2"], - team_models=["AG2"], - ) - - _model, deployments = router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": scoped_key, - } - }, - ) - - assert len(deployments) == 1 - assert deployments[0].get("model_info", {}).get("access_groups") == ["AG2"] - - seen = set() - for _ in range(20): - response = router.completion( - model="gpt-5", - messages=[{"role": "user", "content": "hello"}], - metadata={"user_api_key_team_id": "team2", "user_api_key_auth": scoped_key}, - ) - seen.add(response.choices[0].message.content) - - assert seen == {"response-via-AG2"} - - -def test_explicit_model_access_does_not_force_access_group_filtering(): - """ - If a key has explicit model access in addition to access group entries, - do not force access-group-only filtering for deployment selection. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-5.1", - "api_key": "key1", - "mock_response": "response-via-AG1", - }, - "model_info": {"access_groups": ["AG1"]}, - }, - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "key2", - "mock_response": "response-via-AG2", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - ] - ) - - explicit_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG2", "gpt-5"], - team_models=["AG2", "gpt-5"], - ) - - _model, deployments = router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": explicit_key, - } - }, - ) - - deployment_groups = [ - d.get("model_info", {}).get("access_groups") for d in deployments - ] - assert ["AG1"] in deployment_groups - assert ["AG2"] in deployment_groups - - -def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( - monkeypatch: pytest.MonkeyPatch, -): - """ - When access-group filtering removes all candidates, _get_deployment_by_litellm_model - must not run: it does not re-apply access groups and could return blocked deployments - that share the same litellm_params.model as the request model string. - - ``get_model_access_groups`` is patched to expose AG1 for the public model (so the - access-group filter runs with a non-empty allowed set) while every deployment - returned for that name is AG2-only — filtered to empty. Without the guard, the - litellm-model fallback would return both rows because ``litellm_params.model`` matches. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "gpt-5", - "api_key": "key1", - "mock_response": "blocked-dep-1", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - { - "model_name": "gpt-5", - "litellm_params": { - "model": "gpt-5", - "api_key": "key2", - "mock_response": "blocked-dep-2", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - ] - ) - - orig_groups = router.get_model_access_groups - - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): - if model_name == "gpt-5" and model_access_group is None: - return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} - return orig_groups( - model_name=model_name, - model_access_group=model_access_group, - team_id=team_id, - ) - - monkeypatch.setattr(router, "_core_get_model_access_groups", fake_get_model_access_groups) - - scoped_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG1"], - team_models=["AG1"], - ) - - with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": scoped_key, - } - }, - ) - - -def test_access_group_block_does_not_silently_use_default_fallback_model( - monkeypatch: pytest.MonkeyPatch, -): - """ - When access-group filtering empties candidates for model X, the router must not use - ``fallbacks`` default ``*`` routing to model Y: Y may have no ``access_groups``, so - ``_filter_deployments_by_model_access_groups`` would not constrain Y and the caller - would be served despite being blocked from X. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "gpt-5", - "api_key": "key1", - "mock_response": "blocked-dep-1", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - { - "model_name": "gpt-5", - "litellm_params": { - "model": "gpt-5", - "api_key": "key2", - "mock_response": "blocked-dep-2", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - { - "model_name": "gpt-4-fallback", - "litellm_params": { - "model": "gpt-4", - "api_key": "fallback-key", - "mock_response": "should-not-reach", - }, - }, - ], - fallbacks=[{"*": ["gpt-4-fallback"]}], - ) - - orig_groups = router.get_model_access_groups - - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): - if model_name == "gpt-5" and model_access_group is None: - return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} - return orig_groups( - model_name=model_name, - model_access_group=model_access_group, - team_id=team_id, - ) - - monkeypatch.setattr(router, "_core_get_model_access_groups", fake_get_model_access_groups) - - scoped_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG1"], - team_models=["AG1"], - ) - - with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": scoped_key, - } - }, - ) - - -def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallback( - monkeypatch: pytest.MonkeyPatch, -): - """ - When the by-name lookup returns no deployments and the litellm-model fallback - branch finds candidates that access-group filtering then empties, the router - must not fall through to default ``fallbacks`` routing — the default fallback - model may have no ``access_groups`` and would short-circuit the filter, - silently serving a caller blocked by access-group restrictions. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5-alias", - "litellm_params": { - "model": "gpt-5", - "api_key": "key1", - "mock_response": "blocked-dep-1", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - { - "model_name": "gpt-4-fallback", - "litellm_params": { - "model": "gpt-4", - "api_key": "fallback-key", - "mock_response": "should-not-reach", - }, - }, - ], - fallbacks=[{"*": ["gpt-4-fallback"]}], - ) - - orig_groups = router.get_model_access_groups - - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): - if model_name == "gpt-5" and model_access_group is None: - return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} - return orig_groups( - model_name=model_name, - model_access_group=model_access_group, - team_id=team_id, - ) - - monkeypatch.setattr(router, "_core_get_model_access_groups", fake_get_model_access_groups) - - scoped_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG1"], - team_models=["AG1"], - ) - - with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": scoped_key, - } - }, - ) - - -def test_try_early_resolve_deployments_for_model_not_in_names(): - """ - Direct coverage for ``_try_early_resolve_deployments_for_model_not_in_names``: - - - Returns ``None`` when the requested model is already in ``self.model_names`` - (the by-name lookup path will handle it). - - Returns ``None`` when there are no team deployments, no pattern matches, and - no default deployment to fall back to. - - Returns the pattern-router match when the model matches a wildcard route. - - Returns the default deployment with the request model substituted in when one - is configured, without mutating the stored default. - """ - router_in_names = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-5", - "api_key": "key1", - }, - }, - ] - ) - - assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names( - model="gpt-5", request_team_id=None - ) - is None - ) - assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names( - model="some-unknown-model", request_team_id=None - ) - is None - ) - - pattern_router = litellm.Router( - model_list=[ - { - "model_name": "openai/*", - "litellm_params": { - "model": "openai/*", - "api_key": "key-pattern", - }, - }, - ] - ) - - pattern_result = ( - pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None - ) - ) - assert pattern_result is not None - resolved_model, pattern_deployments = pattern_result - assert resolved_model == "openai/gpt-4o-mini" - assert isinstance(pattern_deployments, list) and len(pattern_deployments) == 1 - - default_router = litellm.Router( - model_list=[ - { - "model_name": "named-model", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "key-named", - }, - }, - ] - ) - default_router.default_deployment = { - "model_name": "default", - "litellm_params": { - "model": "openai/will-be-overridden", - "api_key": "key-default", - }, - } - - default_result = ( - default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None - ) - ) - assert default_result is not None - resolved_model, default_deployment = default_result - assert resolved_model == "brand-new-model" - assert isinstance(default_deployment, dict) - assert default_deployment["litellm_params"]["model"] == "brand-new-model" - # The original default_deployment must not be mutated. - assert ( - default_router.default_deployment["litellm_params"]["model"] - == "openai/will-be-overridden" - ) - - def _router_with_two_deployments(blocked_flags): import litellm @@ -12446,10 +5891,7 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - { - uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} - for uid in unhealthy_ids - } + {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} ) @@ -12573,42 +6015,10 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol assert await router.async_get_fully_unhealthy_model_names() == set() -@pytest.mark.asyncio -async def test_async_get_healthy_deployments_skips_blocked_deployment(): - router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) - healthy_ids = [d["model_info"]["id"] for d in healthy] - assert "dep-0" not in healthy_ids - assert "dep-1" in healthy_ids - assert len(all_dep) == 2 - - -def test_get_healthy_deployments_sync_skips_blocked_deployment(): - router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) - healthy_ids = [d["model_info"]["id"] for d in healthy] - assert "dep-0" in healthy_ids - assert "dep-1" not in healthy_ids - assert len(all_dep) == 2 - - -def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): - router = _router_with_two_deployments([True, False]) - filtered = router._filter_blocked_deployments(router.get_model_list() or []) - ids = [d["model_info"]["id"] for d in filtered] - assert ids == ["dep-1"] - - @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments( - model="gpt-4o", request_kwargs={} - ) + deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -12696,9 +6106,7 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through( - model="gpt-4o", request_kwargs={} - ) + deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) assert deployment["model_info"]["id"] == "pt-1" @@ -12707,9 +6115,7 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through( - model="pt-0", request_kwargs={} - ) + router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) def test_get_available_deployment_for_pass_through_names_cooldown_despite_healthy_non_pass_through(): @@ -12751,9 +6157,7 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == [ - "bedrock-iam-pt" - ] + assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] def test_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -12764,12 +6168,7 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert ( - passthrough_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) - == "sk-fake-for-tests" - ) + assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -12784,32 +6183,6 @@ def test_get_deployment_credentials_with_provider_returns_none_for_blocked_deplo assert router.get_deployment_credentials_with_provider(model_id="dep-1") is not None -def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): - """ - Exercises Router._core_is_deployment_blocked so router_code_coverage.py (AST call graph) - marks the helper as covered by router-named tests. - """ - import types - - import litellm - - router = _router_with_two_deployments([True, False]) - blocked_dep = router.get_deployment("dep-0") - unblocked_dep = router.get_deployment("dep-1") - assert blocked_dep is not None and unblocked_dep is not None - assert litellm.Router._core_is_deployment_blocked(blocked_dep) is True - assert litellm.Router._core_is_deployment_blocked(unblocked_dep) is False - - # No model_info on deployment object → treated as not blocked - assert litellm.Router._core_is_deployment_blocked(object()) is False - missing_blocked = types.SimpleNamespace() - assert litellm.Router._core_is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False - assert ( - litellm.Router._core_is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) - is True - ) - - def test_synchronous_text_completion_rejects_a_restricted_model_before_the_provider() -> None: model: Final = "text-completion-openai/unit-text-model" router: Final = Router( @@ -12917,9 +6290,7 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set( - self, explicit_request_timeout - ): + def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -12937,341 +6308,12 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330) - assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - - def test_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330) - # stream=True resolves through _get_stream_timeout; request_timeout must win. - assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - - def test_explicit_stream_timeout_still_wins_over_request_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330, stream_timeout=45) - assert router._get_stream_timeout(kwargs={}, data={}) == 45 - - def test_non_stream_falls_through_to_router_timeout_without_request_timeout(self): - original_value = litellm.request_timeout - original_flag = litellm.request_timeout_explicitly_set - litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS - litellm.request_timeout_explicitly_set = False - try: - router = self._make_router(timeout=330) - assert router._get_non_stream_timeout(kwargs={}, data={}) == 330 - finally: - litellm.request_timeout = original_value - litellm.request_timeout_explicitly_set = original_flag - - def test_per_deployment_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330) - assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - - def test_per_request_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330) - assert ( - router._get_non_stream_timeout( - kwargs={"timeout": 60}, data={"timeout": 120} - ) - == 60 - ) - - def test_passthrough_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): - router = self._make_router(timeout=330) - deployment: Final = router.model_list[0] - assert _passthrough_timeout(router, deployment, stream=False) == 300.0 - assert _passthrough_timeout(router, deployment, stream=True) == 300.0 - - def test_passthrough_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): - router = self._make_router(timeout=330, stream_timeout=45) - assert _passthrough_timeout(router, router.model_list[0], stream=True) == 45.0 - # --------------------------------------------------------------------------- # Deferred-stream eager-fetch tests # --------------------------------------------------------------------------- -def _make_deferred_stream_wrapper(make_call_fn): - """Return a CustomStreamWrapper with completion_stream=None and the given make_call.""" - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - logging_obj = MagicMock() - logging_obj.model_call_details = {"litellm_params": {}} - return CustomStreamWrapper( - completion_stream=None, - model="vertex_ai/gemini-2.0-flash", - logging_obj=logging_obj, - custom_llm_provider="vertex_ai_beta", - make_call=make_call_fn, - ) - - -def _make_router_with_vertex_and_fallback(): - return litellm.Router( - model_list=[ - { - "model_name": "my-gemini", - "litellm_params": { - "model": "vertex_ai/gemini-2.0-flash", - "vertex_project": "test-project", - "vertex_location": "us-central1", - }, - }, - { - "model_name": "my-fallback", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-fake", - }, - }, - ], - fallbacks=[{"my-gemini": ["my-fallback"]}], - num_retries=0, - ) - - -@pytest.mark.asyncio -async def test_acompletion_deferred_stream_error_propagates_through_acompletion(): - """Regression: a deferred-stream CustomStreamWrapper whose make_call raises a 429 - must propagate the exception from within _acompletion's except block so that - fail_calls is incremented (i.e., deployment cooldown fires) and the standard - router fallback chain can handle it. - - Before the fix, the HTTP call happened inside __anext__ (outside the except block), - so fail_calls was never incremented. - """ - import litellm as _litellm - - rate_limit_err = _litellm.RateLimitError( - message="Resource exhausted", - llm_provider="vertex_ai", - model="gemini-2.0-flash", - ) - - async def failing_make_call(**kwargs): - raise rate_limit_err - - router = _make_router_with_vertex_and_fallback() - deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) - - with patch( - "litellm.acompletion", - new_callable=AsyncMock, - return_value=deferred_wrapper, - ): - with pytest.raises(_litellm.RateLimitError): - await router._acompletion( - model="vertex_ai/gemini-2.0-flash", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - model_name = router.model_list[0]["litellm_params"]["model"] - assert router.fail_calls[model_name] == 1, ( - "fail_calls must be incremented when the deferred HTTP call fails; " - "without the eager fetch_stream() fix this stays at 0" - ) - - -@pytest.mark.asyncio -async def test_acompletion_deferred_stream_preserves_original_headers_on_error(): - """Router is used both by the proxy and directly as an SDK. HTTP-framing headers - (Content-Length, Transfer-Encoding, ...) must NOT be stripped at this layer, or - direct SDK callers lose legitimate provider metadata (e.g. content-type, - proxy-authenticate) that only the proxy's own response construction needs to - worry about. Stripping happens in the proxy layer instead - (_handle_llm_api_exception).""" - import litellm as _litellm - - err = _litellm.RateLimitError( - message="Resource exhausted", - llm_provider="vertex_ai", - model="gemini-2.0-flash", - ) - err.headers = { - "content-length": "42", - "transfer-encoding": "chunked", - "content-encoding": "gzip", - "content-type": "application/json", - "x-request-id": "abc-123", - } - - async def failing_make_call(**kwargs): - raise err - - router = _make_router_with_vertex_and_fallback() - deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) - - with patch( - "litellm.acompletion", - new_callable=AsyncMock, - return_value=deferred_wrapper, - ): - with pytest.raises(_litellm.RateLimitError) as exc_info: - await router._acompletion( - model="vertex_ai/gemini-2.0-flash", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - raised = exc_info.value - headers = getattr(raised, "headers", {}) - assert headers.get("content-length") == "42" - assert headers.get("transfer-encoding") == "chunked" - assert headers.get("content-encoding") == "gzip" - assert headers.get("content-type") == "application/json" - assert headers.get("x-request-id") == "abc-123" - - -@pytest.mark.asyncio -async def test_acompletion_deferred_stream_skipped_when_stream_already_set(): - """When completion_stream is already populated (non-deferred provider), the eager - fetch_stream() call must be skipped entirely; no exception should be raised even - if make_call would fail. - """ - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - async def would_fail(**kwargs): - raise RuntimeError("should not be called") - - logging_obj = MagicMock() - logging_obj.model_call_details = {"litellm_params": {}} - - async def noop_aiter(): - return - yield - - noop_stream = noop_aiter() - already_set_wrapper = CustomStreamWrapper( - completion_stream=noop_stream, - model="openai/gpt-4o", - logging_obj=logging_obj, - custom_llm_provider="openai", - make_call=would_fail, - ) - - router = litellm.Router( - model_list=[ - { - "model_name": "my-model", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "sk-fake", - }, - } - ], - ) - - with patch( - "litellm.acompletion", - new_callable=AsyncMock, - return_value=already_set_wrapper, - ): - result = await router._acompletion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - assert result is not None, "should return a streaming wrapper without errors" - assert already_set_wrapper.completion_stream is noop_stream, "completion_stream must not be re-fetched" - await noop_stream.aclose() - - -def test_completion_deferred_stream_error_propagates_through_completion(): - """Regression: the sync router path needs the same eager fetch as the async one. - - A deferred-stream CustomStreamWrapper hands back a wrapper whose HTTP call has - not happened yet, so without fetch_sync_stream() the provider error surfaces on - first iteration, outside _completion's except block. The deployment is then never - marked failed and function_with_fallbacks never sees the error. - """ - import litellm as _litellm - - rate_limit_err = _litellm.RateLimitError( - message="Resource exhausted", - llm_provider="vertex_ai", - model="gemini-2.0-flash", - ) - make_call_invocations = [] - - def failing_make_call(**kwargs): - make_call_invocations.append(kwargs) - raise rate_limit_err - - router = _make_router_with_vertex_and_fallback() - deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) - - with patch("litellm.completion", return_value=deferred_wrapper): - with pytest.raises(_litellm.RateLimitError): - router._completion( - model="vertex_ai/gemini-2.0-flash", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - assert len(make_call_invocations) == 1, ( - "the deferred HTTP call must run inside _completion's try block; " - "without the eager fetch_sync_stream() fix it is deferred to first iteration" - ) - - -def test_completion_deferred_stream_skipped_when_stream_already_set(): - """A non-deferred sync provider already has completion_stream populated, so the - eager fetch must be skipped and make_call left untouched. - """ - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - def would_fail(**kwargs): - raise RuntimeError("should not be called") - - logging_obj = MagicMock() - logging_obj.model_call_details = {"litellm_params": {}} - already_set_stream = iter([]) - - already_set_wrapper = CustomStreamWrapper( - completion_stream=already_set_stream, - model="openai/gpt-4o", - logging_obj=logging_obj, - custom_llm_provider="openai", - make_call=would_fail, - ) - - router = litellm.Router( - model_list=[ - { - "model_name": "my-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, - } - ], - ) - - with patch("litellm.completion", return_value=already_set_wrapper): - result = router._completion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - assert result is not None, "should return a streaming wrapper without errors" - assert already_set_wrapper.completion_stream is already_set_stream, "completion_stream must not be re-fetched" - - class TestAdvisorSubCallCooldown: """Regression for LIT-4565: an advisor orchestration failure must not cool down the selected (healthy) deployment, which would reject unrelated @@ -13302,9 +6344,7 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns( - model_ids=["dep-1"], parent_otel_span=None - ) + active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -13313,12 +6353,7 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(self._auth_error()), None, now, now - ) - is True - ) + assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -13333,12 +6368,7 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(exception), None, now, now - ) - is False - ) + assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False assert "dep-1" not in self._cooled_down_ids(router) @@ -13575,7 +6605,7 @@ class TestCallerTimeoutCooldown: def test_stream_chunks_have_generated_content_detects_text_and_non_text(): - from litellm.router import _stream_chunks_have_generated_content + from litellm.router_rust.invoke.streaming import _stream_chunks_have_generated_content # noqa: PLC2701 from litellm.types.utils import ( ChatCompletionDeltaToolCall, Delta, @@ -13627,13 +6657,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) + images_delta = Delta( + images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] + ) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta( - annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] - ) + annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -13677,12 +6707,8 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert router.get_configured_token_limits( - "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" - ) == (None, None) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -13795,20 +6821,6 @@ def test_get_model_listing_info_carries_configured_limits(): assert (info.max_input_tokens, info.max_output_tokens) == (32000, 8000) -def test_widest_configured_limit_ignores_absent_and_malformed_values(): - model_infos = ( - {"max_input_tokens": 32000}, - {}, - {"max_input_tokens": "not-a-number"}, - {"max_input_tokens": "128000"}, - {"max_output_tokens": 4096}, - ) - - assert litellm.Router._widest_configured_limit(model_infos, "max_input_tokens") == 128000 - assert litellm.Router._widest_configured_limit(model_infos, "max_output_tokens") == 4096 - assert litellm.Router._widest_configured_limit((), "max_input_tokens") is None - - def test_get_model_listing_info_dedupes_interchangeable_deployments(): """The ordinary group is N deployments of one model, so it yields exactly one key.""" router = litellm.Router( @@ -13905,13 +6917,8 @@ def test_get_model_listing_info_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_reads_deployment_model_info(): @@ -13953,13 +6960,8 @@ def test_get_configured_mode_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_treats_malformed_values_as_absent(): @@ -14018,13 +7020,8 @@ def test_get_configured_display_name_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None @pytest.mark.parametrize( @@ -14053,7 +7050,10 @@ def test_get_configured_service_tiers_returns_one_value_per_deployment_in_model_ "litellm_params": {"model": "openai/gpt-6-astra"}, "model_info": {"service_tiers": ["ultrafast"]}, }, - {"model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}}, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}, + }, { "model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://b.example"}, @@ -14316,113 +7316,6 @@ def test_get_configured_display_name_treats_malformed_values_as_absent(): assert router.get_configured_display_name(f"bad-display-{i}") is None -@pytest.mark.asyncio -async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error(): - router = litellm.Router( - model_list=[ - { - "model_name": "owning-model", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-owning", - }, - }, - { - "model_name": "fallback-model", - "litellm_params": { - "model": "azure/gpt-4o-mini", - "api_key": "sk-fallback", - "api_base": "https://fallback.openai.azure.com", - "api_version": "2024-08-01-preview", - }, - }, - ], - fallbacks=[{"owning-model": ["fallback-model"]}], - num_retries=0, - ) - owning_provider_error = litellm.BadRequestError( - message="completion_window must be one of: 24h", - model="openai/gpt-4o-mini", - llm_provider="openai", - ) - mock_create = AsyncMock(side_effect=owning_provider_error) - - with patch.object(router, "_acreate_batch", mock_create): - with pytest.raises(litellm.BadRequestError, match="24h"): - await router.acreate_batch( - model="owning-model", - input_file_id="file-owned-by-openai", - endpoint="/v1/chat/completions", - completion_window="5m", - disable_fallbacks=True, - ) - - mock_create.assert_awaited_once() - assert mock_create.call_args.kwargs["model"] == "owning-model" - - -@pytest.mark.asyncio -async def test_acreate_batch_surfaces_owning_provider_error_without_disable_fallbacks(): - """The router itself has to keep a batch inside the group that owns the input file: - the proxy only sets disable_fallbacks on the managed-files route, so the caller - otherwise gets the fallback provider's error for a file it never received.""" - from litellm.types.utils import LiteLLMBatch - - router = litellm.Router( - model_list=[ - { - "model_name": "owning-model", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-owning", - }, - }, - { - "model_name": "fallback-model", - "litellm_params": { - "model": "azure/gpt-4o-mini", - "api_key": "sk-fallback", - "api_base": "https://fallback.openai.azure.com", - "api_version": "2024-08-01-preview", - }, - }, - ], - fallbacks=[{"owning-model": ["fallback-model"]}], - num_retries=0, - ) - attempted_models = [] - - async def _acreate_batch(model, **kwargs): - attempted_models.append(model) - if model == "owning-model": - raise litellm.APIConnectionError( - message="Connection error - openai is unreachable", - model="openai/gpt-4o-mini", - llm_provider="openai", - ) - return LiteLLMBatch( - id="batch-created-on-the-wrong-provider", - completion_window="24h", - created_at=0, - endpoint="/v1/chat/completions", - input_file_id="file-owned-by-openai", - object="batch", - status="validating", - ) - - with patch.object(router, "_acreate_batch", _acreate_batch): - with pytest.raises(litellm.APIConnectionError, match="openai is unreachable"): - await router.acreate_batch( - model="owning-model", - input_file_id="file-owned-by-openai", - endpoint="/v1/chat/completions", - completion_window="24h", - metadata={"team": "batch-jobs"}, - ) - - assert attempted_models == ["owning-model"] - - @pytest.mark.asyncio async def test_acreate_batch_still_falls_back_within_the_owning_model_group(): """Holding a batch inside the model group that owns its input file must not @@ -14775,7 +7668,12 @@ def test_credentials_for_a_mixed_group_name_keep_resolving_to_its_first_deployme _AZURE_CHAT_SKU_REJECTS_BATCH: Final = httpx.Response( 400, - json={"error": {"code": "invalid_deployment_type", "message": "The current deployment is not supported for batch jobs"}}, + json={ + "error": { + "code": "invalid_deployment_type", + "message": "The current deployment is not supported for batch jobs", + } + }, ) @@ -14796,7 +7694,9 @@ async def test_acreate_batch_on_a_mixed_group_reaches_chat_as_on_main_while_its_ _cool_down(router, *(d["model_info"]["id"] for d in router.model_list if d["model_info"]["mode"] == "batch")) with respx.mock(assert_all_called=False) as respx_mock: - chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock(return_value=_AZURE_CHAT_SKU_REJECTS_BATCH) + chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock( + return_value=_AZURE_CHAT_SKU_REJECTS_BATCH + ) with pytest.raises(openai.BadRequestError): await _create_batches(router, "gpt-mixed", attempts=1) @@ -14822,7 +7722,9 @@ async def test_acreate_batch_with_tags_that_match_only_the_chat_deployment_reach ) with respx.mock(assert_all_called=False) as respx_mock: - chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock(return_value=_AZURE_CHAT_SKU_REJECTS_BATCH) + chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock( + return_value=_AZURE_CHAT_SKU_REJECTS_BATCH + ) batch_route: Final = respx_mock.post(host="batch-sku.openai.azure.com").mock( return_value=httpx.Response(200, json=_AZURE_BATCH_CREATED) ) @@ -14912,13 +7814,16 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, + with ( + patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, + patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ), ): await router.acreate_batch( model="bedrock-batch-model", @@ -14947,13 +7852,13 @@ async def test_avector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_asearch = AsyncMock(return_value=expected_response) # Router.__init__ binds asearch via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.asearch", new=mock_asearch): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.asearch", new=mock_asearch + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -14987,7 +7892,9 @@ async def test_avector_store_create_does_not_inject_router(): } ] ) - with patch("litellm.vector_stores.main.acreate", new=mock_acreate): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.acreate", new=mock_acreate + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface create_response = await router.avector_store_create(model=None, custom_llm_provider="openai") assert create_response is expected_response @@ -15003,13 +7910,13 @@ def test_vector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_search = MagicMock(return_value=expected_response) # Router.__init__ binds search via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.search", new=mock_search + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -15018,9 +7925,7 @@ def test_vector_store_search_injects_router(): } ] ) - search_response = router.vector_store_search( - vector_store_id="v", query="q", custom_llm_provider="s3_vectors" - ) + search_response = router.vector_store_search(vector_store_id="v", query="q", custom_llm_provider="s3_vectors") assert search_response is expected_response mock_search.assert_called_once() @@ -15034,7 +7939,9 @@ def test_vector_store_create_does_not_inject_router(): mock_create = MagicMock(return_value=expected_response) # Router.__init__ binds create via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.create", new=mock_create + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface router = litellm.Router( model_list=[ { @@ -15069,9 +7976,7 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - }, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -15240,22 +8145,6 @@ class TestPreRoutingStrategyRegistryLifecycle: assert len(registered) == 1 assert set(registered[0].strategy.config.available_models) == {"gpt-4o", "gpt-4o-mini"} - def test_delete_repairs_indices_even_when_strategy_release_fails(self): - """Structural removal and strategy release are not equally critical. Once the entry - leaves model_list the index maps must be repaired no matter what, so releasing the - registry slot runs after that repair and cannot abandon the router half-updated.""" - router = self._router_with_complexity_router() - idx = router.model_id_to_deployment_index_map["router-1"] - router.model_list[idx] = {"model_name": "smart-router", "litellm_params": None} - - returned = router.delete_deployment(id="router-1") - - assert returned is not None - assert "router-1" not in router.model_id_to_deployment_index_map - assert all(entry.get("model_info", {}).get("id") != "router-1" for entry in router.model_list) - assert router.get_deployment(model_id="router-1") is None - assert "gpt-4o" in self._model_names(router) - def test_delete_of_adaptive_enabled_complexity_router_frees_both_registries(self): """A complexity router with adaptive set is registered in BOTH complexity_routers and adaptive_routers under the same (model_name, tags). Releasing only the first @@ -15376,9 +8265,7 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params( - {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - ) + **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -15413,44 +8300,6 @@ class TestPreRoutingStrategyRegistryLifecycle: assert "smart-router" in router.adaptive_routers - def test_unregister_pre_routing_strategy_scopes_the_drop_by_tags(self): - """The bool return drives the hook re-sync; a tag mismatch must report False and - leave the registry untouched, and dropping the last entry must free the key.""" - from litellm.types.router import TaggedPreRoutingStrategy - - registry = { - "m": [ - TaggedPreRoutingStrategy(tags=("team-a",), strategy=object()), - TaggedPreRoutingStrategy(tags=(), strategy=object()), - ] - } - - assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-b",)) is False - assert len(registry["m"]) == 2 - - assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-a",)) is True - assert [entry.tags for entry in registry["m"]] == [()] - - assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ()) is True - assert "m" not in registry - - def test_unregister_for_deployment_ignores_non_router_deployments(self): - """Direct twin of the endpoint-level test: a regular deployment that shares a - router's model_name must not evict the router's registry slot.""" - from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo - - router = self._router_with_complexity_router() - - router._unregister_pre_routing_strategy_for_deployment( - deployment=Deployment( - model_name="smart-router", - litellm_params=LiteLLM_Params(model="openai/gpt-4o"), - model_info=ModelInfo(id="plain-1", db_model=True), - ) - ) - - assert "smart-router" in router.complexity_routers - def test_sync_adaptive_router_hooks_keeps_one_hook_per_registered_router(self): """Re-syncing must replace, not accumulate: a duplicated hook double-fires bandit signal recording for every request.""" @@ -15465,33 +8314,6 @@ class TestPreRoutingStrategyRegistryLifecycle: hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) assert len(hooks) == 1 - def test_deployment_participates_in_adaptive_routing_matrix(self): - """The upsert finalize re-run keys off this predicate for both the incoming and - outgoing deployment; a false negative silently strands the adaptive companion.""" - from litellm.types.router import LiteLLM_Params - - router = self._router_with_complexity_router() - - cases = [ - ({"model": "auto_router/adaptive_router", "adaptive_router_config": {}}, True), - (self._hybrid_router_params({"SIMPLE": "gpt-4o-mini"}), True), - (self._complexity_router_params("gpt-4o"), False), - ( - { - "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}, "adaptive": False}, - "complexity_router_default_model": "gpt-4o", - }, - False, - ), - ({"model": "openai/gpt-4o"}, False), - ] - for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing( - litellm_params=LiteLLM_Params(**params) - ) - assert actual is expected, params["model"] - def test_model_info_is_active_for_environment_matrix(monkeypatch): """The model-write endpoints consult this predicate to tell a deliberately @@ -15512,57 +8334,6 @@ def test_model_info_is_active_for_environment_matrix(monkeypatch): model_info_is_active_for_environment(model_info={"supported_environments": ["production"]}) -def test_pre_call_checks_uses_deployment_model_when_model_info_lookup_raises(monkeypatch): - """ - The supported-params check must run against the deployment's own - provider-qualified model. Resolving the per-deployment model only after the - model-info lookup leaves it unset whenever that lookup raises (an - unregistered custom model), so the check falls back to the bare model group - name and the request dies with 'LLM Provider NOT provided'. - """ - monkeypatch.setattr(litellm, "drop_params", False) - - router = litellm.Router( - model_list=[ - { - "model_name": "custom-alias", - "litellm_params": {"model": "hosted_vllm/not-in-the-catalog"}, - } - ], - enable_pre_call_checks=True, - ) - - def _raise_unmapped(**kwargs): - raise ValueError("This model isn't mapped yet") - - monkeypatch.setattr(router, "_core_get_router_model_info", _raise_unmapped) - - seen: list[tuple] = [] - original_get_supported_openai_params = litellm.get_supported_openai_params - - def _record(model, custom_llm_provider=None, **kwargs): - seen.append((model, custom_llm_provider)) - return original_get_supported_openai_params(model=model, custom_llm_provider=custom_llm_provider, **kwargs) - - monkeypatch.setattr(litellm, "get_supported_openai_params", _record) - - deployments = [ - { - "litellm_params": {"model": "hosted_vllm/not-in-the-catalog"}, - "model_info": {"id": "d1"}, - } - ] - result = router._pre_call_checks( - model="custom-alias", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - request_kwargs={"response_format": {"type": "json_object"}}, - ) - - assert len(result) == 1 - assert seen == [("not-in-the-catalog", "hosted_vllm")] - - def test_pre_call_checks_keeps_deployment_when_provider_is_unresolvable(monkeypatch): """ Pre-call checks filter deployments; they must never be the thing that fails @@ -15661,39 +8432,6 @@ class TestUpsertDeploymentRollback: assert router.get_deployment(model_id="fresh-1") is None assert router.model_list == [] - def test_restore_re_adds_popped_deployment(self): - router = litellm.Router( - model_list=[ - { - "model_name": "prod-model", - "litellm_params": {"model": "gpt-4o", "api_key": "sk-old"}, - "model_info": {"id": "prod-1", "db_model": True}, - } - ], - ignore_invalid_deployments=True, - ) - previous = router.get_deployment(model_id="prod-1") - router.delete_deployment(id="prod-1") - assert router.has_model_id("prod-1") is False - - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) - - restored = router.get_deployment(model_id="prod-1") - assert restored is not None - assert restored.litellm_params.model == "gpt-4o" - - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) - assert len(router.model_list) == 1 - - router._restore_deployment_after_failed_upsert( - previous_deployment=None, model_id="prod-1" - ) - assert len(router.model_list) == 1 - class TestUpsertDeploymentRename: """ @@ -15980,31 +8718,6 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() - @pytest.mark.asyncio - async def test_redis_read_failure_does_not_reject_a_subagent_request(self): - from litellm.caching.caching import RedisCache - - router = self._router() - request_kwargs = self._request_kwargs(agent_id="agent-1234") - cache_key = router._claude_code_session_router_cache_key(request_kwargs) - assert cache_key is not None - await router._claude_code_session_router_cache.in_memory_cache.async_set_cache( - cache_key, - "smart-router", - ) - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open")) - router.update_redis_cache(cache=redis_cache) - - response = await router.async_pre_routing_hook( - model="expensive-model", - request_kwargs=request_kwargs, - ) - - assert response is None - assert "model_group" not in request_kwargs["metadata"] - redis_cache.async_get_cache.assert_awaited_once() - @pytest.mark.asyncio async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self): from litellm.caching.caching import RedisCache @@ -16083,20 +8796,6 @@ class TestClaudeCodeSubagentSessionRouterBinding: redis_cache.async_set_cache.assert_not_awaited() redis_cache.async_delete_cache.assert_not_awaited() - @pytest.mark.asyncio - async def test_session_bindings_do_not_evict_router_rate_limit_state(self): - router = self._router() - assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 1 - - for session_index in range(201): - request_kwargs = self._request_kwargs() - request_kwargs["proxy_server_request"]["headers"]["X-Claude-Code-Session-Id"] = ( - f"session-{session_index:04d}" - ) - await router.async_pre_routing_hook(model="smart-router", request_kwargs=request_kwargs) - - assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 2 - @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() @@ -16397,9 +9096,6 @@ class TestTaggedAutoRouterOnSharedModelName: ) assert deployment["litellm_params"]["model"] == "openai/gpt-4o" - def test_deployment_without_litellm_params_mapping_is_not_a_marker(self): - assert litellm.Router._is_strategy_marker_deployment({"model_name": "gpt4o"}) is False - class TestAutoRouterSharedModelNameConnectionParams: """A plain deployment sharing its model_name with an `auto_router/` marker must not have @@ -16451,18 +9147,14 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [ - {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} - ], + "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize( - "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] - ) + @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -16594,9 +9286,7 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs - ) + await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -16870,9 +9560,7 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError( - message=huge_message, llm_provider="openai", model="group-a" - ), + mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -16880,17 +9568,6 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): assert max(len(message) for message in capture.messages) < 5_000 -def test_stamp_or_clear_metadata_key_writes_and_clears_both_buckets(): - request_kwargs = {"metadata": {}} - litellm.Router._stamp_or_clear_metadata_key(request_kwargs=request_kwargs, key="probe", value=7) - assert request_kwargs["metadata"]["probe"] == 7 - - stale_kwargs = {"metadata": {"probe": 7}, "litellm_metadata": {"probe": 7}} - litellm.Router._stamp_or_clear_metadata_key(request_kwargs=stale_kwargs, key="probe", value=None) - assert "probe" not in stale_kwargs["metadata"] - assert "probe" not in stale_kwargs["litellm_metadata"] - - @pytest.mark.parametrize( "complexity_router_config,expect_callback", [ @@ -16929,24 +9606,6 @@ def test_complexity_router_registers_affinity_callback_for_deployment_pin(comple litellm.logging_callback_manager.remove_callback_from_all_lists(cb) -def test_ensure_deployment_affinity_callback_is_idempotent(): - from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( - DeploymentAffinityCheck, - ) - - router = litellm.Router(model_list=[]) - try: - router._ensure_deployment_affinity_callback() - router._ensure_deployment_affinity_callback() - affinity_callbacks = [ - cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) - ] - assert len(affinity_callbacks) == 1 - finally: - for cb in router.optional_callbacks or []: - litellm.logging_callback_manager.remove_callback_from_all_lists(cb) - - def test_get_router_model_info_does_not_wipe_cached_pricing(): """A Deployment's model_info declares the mirrored pricing fields with None defaults; merging it must not write those Nones into the lru_cache'd dict get_model_info() owns, @@ -17082,9 +9741,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion( - model="smart-alias", messages=self._messages(), metadata=metadata - ) + response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -17326,64 +9983,6 @@ class TestTeamPublicNameReachesPreRoutingStrategies: assert response is not None assert response.model == "gemini-flash" - @pytest.mark.asyncio - async def test_strategy_resolution_agrees_with_the_deployment_path_for_every_principal(self): - router = self._router( - {self.INTERNAL_NAME: self._RewriteStrategy()}, - extra_deployments=({"model_name": "shared-name", "litellm_params": {"model": "openai/gpt-4o"}},), - ) - principals = { - "team": self._team_request(), - "other-team": self._team_request(self.OTHER_TEAM), - "no-team": self._team_request(None), - "admin": {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}}, - } - for principal, request_kwargs in principals.items(): - for model in (self.PUBLIC_NAME, "shared-name", "gemini-flash", "missing"): - resolved = [d["model_name"] for d in router.deployments_for_request(model, request_kwargs)] - callable_names = [ - name - for name, deployment in zip(resolved, router.deployments_for_request(model, request_kwargs)) - if not router._is_strategy_marker_deployment(deployment) - ] - if resolved and not callable_names: - with pytest.raises(litellm.BadRequestError, match="strategy router marker"): - router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) - elif not resolved: - with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) - else: - _, deployments = router._common_checks_available_deployment( - model=model, request_kwargs=request_kwargs - ) - assert [d["model_name"] for d in deployments] == callable_names, (principal, model) - - def test_drop_strategy_markers_keeps_plain_deployments_and_rejects_marker_only_sets(self): - router = self._router({}) - marker = router.model_list[0] - plain = {"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o"}} - - assert router._drop_strategy_markers("x", [marker, plain]) == [plain] - assert router._drop_strategy_markers("x", [plain]) == [plain] - assert router._drop_strategy_markers("x", []) == [] - with pytest.raises(litellm.BadRequestError, match="strategy router marker"): - router._drop_strategy_markers("x", [marker]) - - def test_team_deployments_across_teams_unions_one_team_and_rejects_two(self): - other_team_marker = { - **self._team_marker(self.SIBLING_INTERNAL_NAME), - "model_info": {"team_id": self.OTHER_TEAM, "team_public_model_name": self.PUBLIC_NAME}, - } - one_team = self._router({}) - two_teams = self._router({}, markers=(self._team_marker(self.INTERNAL_NAME), other_team_marker)) - - assert [d["model_name"] for d in one_team._team_deployments_across_teams(self.PUBLIC_NAME)] == [ - self.INTERNAL_NAME - ] - assert one_team._team_deployments_across_teams("missing") == [] - with pytest.raises(litellm.BadRequestError, match="multiple teams"): - two_teams._team_deployments_across_teams(self.PUBLIC_NAME) - def test_compression_policy_follows_the_same_resolution_for_every_principal(self): from litellm.proxy.guardrails.auto_router_compression import AutoRouterCompressionPolicy, policy_for_model @@ -17608,7 +10207,6 @@ class TestAutoRouterCompressionDecoupling: @pytest.mark.usefixtures("local_model_cost_map") - @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known @@ -17632,41 +10230,6 @@ class TestAzureBaseModelFallbackLogging: ] ) - def test_map_known_deployment_name_resolves_without_error_log(self): - router = self._router_with_azure_deployment("azure/gpt-4o") - - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: - model_info = router.get_router_model_info( - deployment=None, received_model_name="my-group", id="azure-base-model-test-id" - ) - - assert not any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), f"unexpected error log: {mock_error.call_args_list}" - # the fallback resolution must actually surface the map values - assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] - assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] - - def test_unmappable_deployment_name_still_logs_error(self): - router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: - model_info = router.get_router_model_info( - deployment=None, received_model_name="my-group", id="azure-base-model-test-id" - ) - - assert any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), "expected the error log for an unmappable azure deployment name" - # unmappable names resolve to a zeroed stub — unchanged behavior - assert model_info.get("max_input_tokens") is None - def test_explicit_base_model_still_wins(self): router = litellm.Router( model_list=[ @@ -17690,95 +10253,6 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] -def test_model_group_info_intersects_supported_reasoning_efforts(): - router = litellm.Router( - model_list=[ - { - "model_name": "smart-group", - "litellm_params": {"model": "anthropic/opus-like"}, - "model_info": {"id": "opus-like-deployment"}, - }, - { - "model_name": "smart-group", - "litellm_params": {"model": "openai/mini-like"}, - "model_info": {"id": "mini-like-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - if model_id == "opus-like-deployment": - return { - "key": model_name, - "litellm_provider": "anthropic", - "mode": "chat", - "supports_reasoning": True, - "supports_xhigh_reasoning_effort": True, - "supports_max_reasoning_effort": True, - } - return { - "key": model_name, - "litellm_provider": "openai", - "mode": "chat", - "supports_reasoning": True, - "supports_none_reasoning_effort": False, - "supports_minimal_reasoning_effort": True, - "supports_xhigh_reasoning_effort": False, - } - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="smart-group", - user_facing_model_group_name="smart-group", - ) - - assert result is not None - # opus-like offers all seven levels, mini-like lacks none/xhigh/max; only the common set survives, - # so the group never advertises an effort routing could hand to a deployment that rejects it. - assert result.supported_reasoning_efforts == ("minimal", "low", "medium", "high") - - -def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_off_the_map(): - """The router fills every ModelInfo key, so a deployment absent from the model map arrives with - supports_reasoning None rather than with the key missing. Its synthesized entry carries no mode, - which is what separates it from a mapped non-reasoning model, and nothing being known about it is - no evidence that the unknown deployment accepts levels its mapped sibling supports. The group - therefore reports unknown instead of advertising a value routing might send to either one.""" - router = litellm.Router( - model_list=[ - { - "model_name": "smart-group", - "litellm_params": {"model": "anthropic/opus-like"}, - "model_info": {"id": "opus-like-deployment"}, - }, - { - "model_name": "smart-group", - "litellm_params": {"model": "openai/unmapped-model"}, - "model_info": {"id": "unmapped-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - if model_id == "opus-like-deployment": - return { - "key": model_name, - "litellm_provider": "anthropic", - "mode": "chat", - "supports_reasoning": True, - "supports_max_reasoning_effort": True, - } - return {"key": model_name, "litellm_provider": "openai", "mode": None, "supports_reasoning": None} - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="smart-group", - user_facing_model_group_name="smart-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts is None - @pytest.mark.parametrize( "model,provider,expected", @@ -17799,11 +10273,15 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o def test_model_group_info_fast_mode_uses_exact_provider_catalog( local_model_cost_map: None, model: str, provider: str | None, expected: bool, operator_flag: bool ) -> None: - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": operator_flag}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": operator_flag}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -17815,16 +10293,21 @@ def test_model_group_info_fast_mode_uses_exact_provider_catalog( def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, flag: object ) -> None: - entry: Final = {key: value for key, value in litellm.model_cost["claude-opus-5"].items() - if key != "supports_fast_mode"} + entry: Final = { + key: value for key, value in litellm.model_cost["claude-opus-5"].items() if key != "supports_fast_mode" + } if flag is not None: entry["supports_fast_mode"] = flag monkeypatch.setitem(litellm.model_cost, "claude-opus-5", entry) - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": True}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": True}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -17832,13 +10315,16 @@ def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( assert result.supports_fast_mode is False -@pytest.mark.parametrize("other_model,expected", [ - ("anthropic/claude-opus-4-8", True), - ("anthropic/claude-opus-4-7", False), - ("anthropic/off-map-opus", False), - ("vertex_ai/claude-opus-5", False), - ("bedrock/claude-opus-5", False), -]) +@pytest.mark.parametrize( + "other_model,expected", + [ + ("anthropic/claude-opus-4-8", True), + ("anthropic/claude-opus-4-7", False), + ("anthropic/off-map-opus", False), + ("vertex_ai/claude-opus-5", False), + ("bedrock/claude-opus-5", False), + ], +) @pytest.mark.parametrize("reverse", [True, False]) def test_model_group_info_fast_mode_requires_every_deployment( local_model_cost_map: None, other_model: str, expected: bool, reverse: bool @@ -17860,177 +10346,6 @@ def test_model_group_info_fast_mode_requires_every_deployment( assert result.supports_fast_mode is expected -def test_model_group_info_surfaces_supports_parallel_function_calling(local_model_cost_map): - """``/model_group/info`` folds each deployment's registry flags into the group; a deployment whose - registry entry declares parallel function calling must flip the group to True instead of False.""" - router = litellm.Router( - model_list=[ - { - "model_name": "glm-group", - "litellm_params": {"model": "together_ai/zai-org/GLM-5.3-Flash", "api_key": "fake-key"}, - } - ] - ) - - result = router._set_model_group_info(model_group="glm-group", user_facing_model_group_name="glm-group") - - assert result is not None - assert result.supports_parallel_function_calling is True - - -def test_model_group_info_reasoning_efforts_empty_on_a_mapped_non_reasoning_deployment(): - """A group mixing a reasoning model with one the map knows is not a reasoning model shares no - level, so it advertises none and the picker offers nothing rather than a level routing would - hand to a deployment that rejects it.""" - router = litellm.Router( - model_list=[ - { - "model_name": "mixed-group", - "litellm_params": {"model": "anthropic/opus-like"}, - "model_info": {"id": "opus-like-deployment"}, - }, - { - "model_name": "mixed-group", - "litellm_params": {"model": "openai/plain-chat"}, - "model_info": {"id": "plain-chat-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - if model_id == "opus-like-deployment": - return { - "key": model_name, - "litellm_provider": "anthropic", - "mode": "chat", - "supports_reasoning": True, - "supports_max_reasoning_effort": True, - } - return {"key": model_name, "litellm_provider": "openai", "mode": "chat", "supports_reasoning": None} - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="mixed-group", - user_facing_model_group_name="mixed-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts == () - - -def test_model_group_info_reasoning_efforts_ignore_a_value_declared_in_model_info(): - """The group's levels are computed from its deployments, so a value an operator left in one - deployment's model_info must not seed them. Seeding let the first deployment read narrow the - whole group while the same value on any other deployment was silently ignored.""" - router = litellm.Router( - model_list=[ - { - "model_name": "declared-group", - "litellm_params": {"model": "openai/first-reasoner"}, - "model_info": {"id": "first-deployment"}, - }, - { - "model_name": "declared-group", - "litellm_params": {"model": "openai/second-reasoner"}, - "model_info": {"id": "second-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - info = { - "key": model_name, - "litellm_provider": "openai", - "mode": "chat", - "supports_reasoning": True, - "supports_none_reasoning_effort": True, - } - if model_id == "first-deployment": - info["supported_reasoning_efforts"] = ("high",) - return info - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="declared-group", - user_facing_model_group_name="declared-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") - - -def test_model_group_info_survives_a_junk_typed_operator_effort_value(): - """A deployment's registered model_info reads back with whatever the operator wrote under any - key, so a wrong-typed supported_reasoning_efforts must not fail the group's info. Only the - constructor's trailing override keeps the junk away from ModelGroupInfo validation.""" - router = litellm.Router( - model_list=[ - { - "model_name": "junk-declared-group", - "litellm_params": {"model": "openai/lone-reasoner"}, - "model_info": {"id": "junk-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - return { - "key": model_name, - "litellm_provider": "openai", - "mode": "chat", - "supports_reasoning": True, - "supports_none_reasoning_effort": True, - "supported_reasoning_efforts": "high", - } - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="junk-declared-group", - user_facing_model_group_name="junk-declared-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") - - -def test_model_group_info_reasoning_efforts_are_unknown_for_an_operator_declared_mode(): - """A deployment is registered in the cost map under its own id with whatever model_info the - operator wrote, so a mode they set themselves reads back exactly like one the map supplied. Only - a mode the map supplied marks the deployment as known. An off-map deployment carrying an - operator mode remains unknown and must keep the whole group's level support unknown.""" - from litellm.litellm_core_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts - - mapped_model = "openai/gpt-5.6-sol" - expected = resolve_supported_reasoning_efforts( - litellm.get_model_info(model=mapped_model), - deployment_is_mapped=True, - ) - assert expected - - router = litellm.Router( - model_list=[ - { - "model_name": "smart-group", - "litellm_params": {"model": mapped_model, "api_key": "sk-fake"}, - "model_info": {"id": "mapped-deployment"}, - }, - { - "model_name": "smart-group", - "litellm_params": {"model": "openai/a-model-the-map-never-heard-of", "api_key": "sk-fake"}, - "model_info": {"id": "off-map-deployment", "mode": "chat"}, - }, - ] - ) - - result = router._set_model_group_info( - model_group="smart-group", - user_facing_model_group_name="smart-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts is None - - class TestAddDeploymentApiBaseProviderResolution: def test_bare_model_with_known_api_base_initializes(self): router = litellm.Router( @@ -18092,6 +10407,7 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" + # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -18112,121 +10428,10 @@ class TestAddDeploymentApiBaseProviderResolution: # ===================================================================== -async def _anthropic_messages_empty_generator(): - return - yield # pragma: no cover - makes this an async generator - - -def _anthropic_messages_make_wrapper() -> FallbackAwareAnthropicMessagesStream: - """A minimal wrapper for tests that call _aanthropic_messages_fallback_attempt - directly, bypassing _aanthropic_messages_streaming_iterator.""" - return FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object()) - - -def _anthropic_messages_make_router(**router_kwargs) -> Router: - """A fallback-only router: no same-group retries unless a test asks for them.""" - router_kwargs.setdefault("fallbacks", [{"primary": ["fallback"]}]) - router_kwargs.setdefault("num_retries", 0) - return Router( - model_list=[ - { - "model_name": "primary", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-5", - "api_key": "sk-test", - }, - }, - { - "model_name": "fallback", - "litellm_params": { - "model": "bedrock/anthropic.claude-sonnet-4-5", - }, - }, - ], - **router_kwargs, - ) - - -class _AnthropicMessagesFakeByteStream: - """Minimal AsyncIterator[bytes], carrying _hidden_params like - AnthropicMessagesStreamingResponse does.""" - - def __init__(self, chunks: list) -> None: - self._chunks = list(chunks) - self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}} - self.closed = False - self.closed_as_superseded = False - self.close_contexts: tuple[bool, ...] = () - - def __aiter__(self): - return self - - async def __anext__(self) -> bytes: - if not self._chunks: - raise StopAsyncIteration - return self._chunks.pop(0) - - async def aclose(self) -> None: - self.close_contexts = (*self.close_contexts, closing_superseded_attempt()) - self.closed_as_superseded = self.closed_as_superseded or closing_superseded_attempt() - self.closed = True - - -class _AnthropicMessagesRaisingByteStream: - """Simulates the completion-bridge path: no error SSE chunk is ever - yielded, the underlying CustomStreamWrapper raises MidStreamFallbackError - directly out of the iterator instead (a Vertex AI transport drop).""" - - def __init__(self, chunks: list, error: Exception) -> None: - self._chunks = list(chunks) - self._error = error - self._hidden_params: dict = {} - self.closed = False - - def __aiter__(self): - return self - - async def __anext__(self) -> bytes: - if self._chunks: - return self._chunks.pop(0) - raise self._error - - async def aclose(self) -> None: - self.closed = True - - -class _AnthropicMessagesFallbackByteStream: - def __init__(self, chunks: list, hidden_params: dict | None = None) -> None: - self._chunks = list(chunks) - self._hidden_params = hidden_params if hidden_params is not None else {} - - def __aiter__(self): - return self - - async def __anext__(self) -> bytes: - if not self._chunks: - raise StopAsyncIteration - return self._chunks.pop(0) - - def _anthropic_messages_overloaded_error_chunk() -> bytes: return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' -def _anthropic_messages_invalid_request_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}\n\n' - ) - - -def _anthropic_messages_rate_limit_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "rate_limit_error", "message": "Too many requests"}}\n\n' - ) - - def _anthropic_messages_content_chunk(text: str = "hi") -> bytes: payload = f'{{"type": "content_block_delta", "delta": {{"type": "text_delta", "text": "{text}"}}}}' return f"event: content_block_delta\ndata: {payload}\n\n".encode() @@ -18245,777 +10450,15 @@ def _anthropic_messages_ping_chunk() -> bytes: # -------- _aanthropic_messages_streaming_iterator (passthrough) -------- -@pytest.mark.asyncio -async def test_anthropic_messages_streaming_iterator_passthrough(): - """Without any error chunk, the wrapper forwards every chunk unchanged - and carries the source iterator's _hidden_params through (so response - headers like Bedrock's request-id keep flowing to the client).""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] - ) - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - - collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] - assert wrapped._hidden_params["additional_headers"]["x-amzn-requestid"] == "req-1" - - -@pytest.mark.asyncio -async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_frames_in_order(): - """Regression: lifecycle frames held back to guard against a mid-stream - fallback must still reach the client, in order, once real content - arrives - buffering them for the fallback-safety check must not silently - drop them on the happy path.""" - router = _anthropic_messages_make_router() - message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] - ) - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - - collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] - - -@pytest.mark.asyncio -async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_stream_end(): - """Regression: if the primary stream ends with only lifecycle frames and - no content and no error, the buffered frames must still reach the - client rather than being silently swallowed.""" - router = _anthropic_messages_make_router() - message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - - collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), message_stop] - - with pytest.raises(StopAsyncIteration): - await wrapped.__anext__() - - -@pytest.mark.asyncio -async def test_anthropic_messages_content_coalesced_with_error_in_one_physical_chunk_skips_fallback(): - """Greptile review round: transport-level buffering can coalesce a real - content_block_delta and a following retriable error into ONE physical - read from the source iterator. Since the whole chunk (content and error - together) is forwarded to the client atomically, the client genuinely - receives the content - so no fallback must be attempted, exactly as if - the two events had arrived as separate reads.""" - router = _anthropic_messages_make_router() - coalesced_chunk = _anthropic_messages_content_chunk("partial") + _anthropic_messages_overloaded_error_chunk() - source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [coalesced_chunk] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_ping_behind_buffered_lifecycle_frame_is_forwarded_live(): - """A `ping` behind buffered lifecycle frames still reaches the client - live: it carries no lifecycle, so it cannot create overlapping - lifecycles, and it keeps the connection alive while a fallback-able - stream holds message_start back through a long thinking pass.""" - router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}]) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - yield _anthropic_messages_ping_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [ - _anthropic_messages_message_start_chunk(), - _anthropic_messages_content_chunk("hi"), - ] - - -@pytest.mark.asyncio -async def test_anthropic_messages_split_ping_stays_in_order_behind_buffered_lifecycle_frame(): - """A ping the transport splits across two reads is not a whole frame, so - neither fragment may jump ahead of the buffered message_start: yielding - the head live and flushing the tail behind message_start would splice a - lifecycle frame into the middle of the ping on the wire.""" - router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}]) - ping_head, ping_tail = b'event: ping\ndata: {"ty', b'pe": "ping"}\n\n' - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), ping_head, ping_tail, _anthropic_messages_content_chunk("hi")] - ) - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) - - assert [chunk async for chunk in wrapped] == [ - _anthropic_messages_message_start_chunk(), - ping_head, - ping_tail, - _anthropic_messages_content_chunk("hi"), - ] - - -@pytest.mark.asyncio -async def test_anthropic_messages_no_fallback_message_start_reaches_client_before_content(): - """With no fallback able to take over, the stream is committed from the - first frame: message_start reaches the client live instead of waiting - behind the buffer for content that may be a whole thinking pass away.""" - router = _anthropic_messages_make_router(fallbacks=None) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_disabled_fallbacks_message_start_reaches_client_before_content(): - """A router with fallbacks configured cannot take over a request that - opted out with disable_fallbacks=True, so its lifecycle frames reach - the client live exactly like a no-fallback router's.""" - router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}]) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source(), initial_kwargs={"model": "primary", "disable_fallbacks": True} - ) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_no_fallback_error_frame_reaches_client_verbatim(): - """With no fallback able to take over, a retriable provider error frame - is forwarded verbatim instead of triggering a fallback that does not - exist, and the frames already received stay in order ahead of it.""" - router = _anthropic_messages_make_router(fallbacks=None) - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_default_wildcard_fallback_still_buffers_lifecycle_frames(): - """A "*" default fallback can take over for any group, so lifecycle - frames are still held back until real content commits the primary.""" - router = _anthropic_messages_make_router(fallbacks=[{"*": ["fallback"]}]) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - pending = asyncio.ensure_future(wrapped.__anext__()) - await asyncio.sleep(0.2) - assert not pending.done() - content_released.set() - assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -def _anthropic_messages_two_order_primary_model_list() -> list: - return [ - { - "model_name": "primary", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": 1}, - }, - { - "model_name": "primary", - "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-4-5", "order": 2}, - }, - { - "model_name": "fallback", - "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-4-5"}, - }, - ] - - -@pytest.mark.parametrize( - "router_kwargs,request_kwargs,expected", - [ - pytest.param({"fallbacks": None}, {"model": "primary"}, False, id="no-fallbacks"), - pytest.param({"fallbacks": [{"primary": ["fallback"]}]}, {"model": "primary"}, True, id="group-fallback"), - pytest.param({"fallbacks": [{"other": ["fallback"]}]}, {"model": "primary"}, False, id="unrelated-group"), - pytest.param( - {"fallbacks": [{"*": ["fallback"]}]}, - {"model": "primary", "fallbacks": None}, - False, - id="wildcard-overridden-by-request-none", - ), - pytest.param({"fallbacks": [{"*": ["fallback"]}]}, {"model": "primary"}, True, id="wildcard"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": [{"model": "fallback"}]}, True, id="request-dict-fallback"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback"), - pytest.param( - {"fallbacks": [{"primary": ["fallback"]}]}, - {"model": "primary", "disable_fallbacks": True}, - False, - id="disable-fallbacks", - ), - pytest.param( - {"fallbacks": None, "content_policy_fallbacks": [{"primary": ["fallback"]}]}, - {"model": "primary"}, - True, - id="content-policy-fallback", - ), - pytest.param({"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover"), - ], -) -def test_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): - router = _anthropic_messages_make_router(**router_kwargs) - assert router._stream_can_fall_back("primary", request_kwargs) is expected - - -@pytest.mark.parametrize( - "orders,expected", - [ - pytest.param([1, 2], True, id="distinct-orders-can-fall-back"), - pytest.param([1, 1], False, id="same-order-cannot-fall-back"), - ], -) -def test_stream_can_fall_back_order_levels(orders, expected): - router = Router( - model_list=[ - { - "model_name": "primary", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": order}, - } - for order in orders - ], - fallbacks=None, - ) - assert router._stream_can_fall_back("primary", {"model": "primary"}) is expected - - -@pytest.mark.parametrize( - "request_kwargs,expected", - [ - pytest.param({"model": "primary"}, True, id="no-target-order"), - pytest.param({"model": "primary", "_target_order": 1}, True, id="higher-order-remains"), - pytest.param({"model": "primary", "_target_order": 2}, False, id="top-order-no-order-fallback"), - pytest.param( - {"model": "primary", "_target_order": 2, "fallbacks": [{"primary": ["fallback"]}]}, - True, - id="top-order-external-fallback", - ), - ], -) -def test_stream_can_fall_back_order_target(request_kwargs, expected): - router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None) - assert router._stream_can_fall_back("primary", request_kwargs) is expected - - -def test_stream_order_levels_direct_call(): - router = Router( - model_list=[ - { - "model_name": "primary", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": order}, - } - for order in (2, 1, None) - ], - fallbacks=None, - ) - assert router._stream_order_levels("primary", {"model": "primary"}) == (1, 2) - - -@pytest.mark.asyncio -async def test_anthropic_messages_order_fallback_still_buffers_lifecycle_frames(): - """Two order levels in one group are a real fallback target for the - dispatcher, so lifecycle frames stay buffered until content commits.""" - router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - pending = asyncio.ensure_future(wrapped.__anext__()) - await asyncio.sleep(0.2) - assert not pending.done() - content_released.set() - assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_request_fallbacks_none_forwards_message_start_live(): - """A per-request fallbacks=None override disables the router's wildcard - fallback, so lifecycle frames reach the client live before content.""" - router = _anthropic_messages_make_router(fallbacks=[{"*": ["fallback"]}]) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source(), initial_kwargs={"model": "primary", "fallbacks": None} - ) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): - """A `ping` that no lifecycle frame precedes is how a hold-back turn keeps - its connection alive (AgenticAnthropicStreamingIterator), so it must reach - the client at once rather than wait behind the pre-content buffer.""" - router = _anthropic_messages_make_router() - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_ping_chunk() - await content_released.wait() - yield _anthropic_messages_message_start_chunk() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [ - _anthropic_messages_message_start_chunk(), - _anthropic_messages_content_chunk("hi"), - ] - - -@pytest.mark.asyncio -async def test_anthropic_messages_leading_ping_does_not_disqualify_fallback(): - """A live-forwarded leading `ping` commits nothing: a retriable error after - it still falls back, and the fallback's own lifecycle follows the ping cleanly.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_ping_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - fallback_message_start = b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_2"}}\n\n' - fallback_stream = _AnthropicMessagesFallbackByteStream( - [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ): - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [ - _anthropic_messages_ping_chunk(), - fallback_message_start, - _anthropic_messages_content_chunk("fallback answer"), - ] - - -@pytest.mark.asyncio -async def test_anthropic_messages_hold_back_retrieval_failure_reaches_client_without_fallback(): - """The hold-back iterator's own retrieval-failure frame is the gateway's verdict, not a - provider failure: a configured fallback stays untouched and the client reads the error - right after the live keepalive.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_ping_chunk(), SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] - ) - fallback = AsyncMock( - return_value=_AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - ) - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=fallback): - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - fallback.assert_not_called() - assert collected == [_anthropic_messages_ping_chunk(), SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] - - -@pytest.mark.asyncio -async def test_anthropic_messages_pre_content_buffer_cap_forces_commit(): - """Bugbot regression: a hostile or pathological upstream that never emits - real content or an error must not grow the pre-content lifecycle buffer - without bound - hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS commits - to the primary stream early, exactly as real content arriving would.""" - router = _anthropic_messages_make_router() - lifecycle_chunk = _anthropic_messages_message_start_chunk() - error_chunk = _anthropic_messages_overloaded_error_chunk() - chunks = [lifecycle_chunk] * (MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5) + [error_chunk] - source = _AnthropicMessagesFakeByteStream(chunks) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - collected = [chunk async for chunk in wrapped] - - mock_fallback.assert_not_awaited() - assert collected.count(lifecycle_chunk) == MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5 - assert collected[-1] == error_chunk - - -@pytest.mark.asyncio -async def test_anthropic_messages_ping_coalesced_with_content_in_one_physical_chunk_is_forwarded(): - """Greptile/Bugbot regression: transport-level buffering can coalesce a - `ping` keepalive and a real content_block_delta into ONE physical read. - The pre-content ping-drop must only discard PURE ping frames - dropping - the whole coalesced chunk would silently lose generated content.""" - router = _anthropic_messages_make_router() - coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_content_chunk("hi") - source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) - collected = [chunk async for chunk in wrapped] - - assert collected == [coalesced_chunk] - - -@pytest.mark.asyncio -async def test_anthropic_messages_ping_coalesced_with_retriable_error_still_falls_back(): - """Greptile/Bugbot regression: a physical chunk coalescing a `ping` with a - retriable `event: error` must not be discarded as a keepalive - the error - inside it must still trigger the mid-stream fallback.""" - router = _anthropic_messages_make_router() - coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_overloaded_error_chunk() - source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - collected = [chunk async for chunk in wrapped] - - mock_fallback.assert_awaited_once() - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - - # -------- _aanthropic_messages_fallback_attempt -------- -@pytest.mark.asyncio -async def test_aanthropic_messages_fallback_attempt_yields_fallback_stream(): - """Direct-call regression: the fallback-attempt helper re-enters the - Router's fallback chain and forwards whatever the fallback produces.""" - router = _anthropic_messages_make_router() - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - collected = [ - chunk - async for chunk in router._aanthropic_messages_fallback_attempt( - error, - {"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, - _anthropic_messages_make_wrapper(), - ) - ] - - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - mock_fallback.assert_awaited_once() - assert mock_fallback.await_args.kwargs["e"] is error - - -@pytest.mark.asyncio -async def test_aanthropic_messages_fallback_attempt_raises_original_exception_on_double_failure(): - """Direct-call regression: when the fallback attempt itself fails with a - MidStreamFallbackError wrapping a real provider exception, that real - exception must surface rather than the internal wrapper exception.""" - router = _anthropic_messages_make_router() - error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") - original_exception = litellm.APIError( - status_code=503, message="fallback also overloaded", llm_provider="bedrock", model="fallback" - ) - fallback_failure = MidStreamFallbackError( - message="fallback failed", model="fallback", llm_provider="bedrock", original_exception=original_exception - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=fallback_failure), - ): - with pytest.raises(litellm.APIError) as exc_info: - async for _ in router._aanthropic_messages_fallback_attempt( - error, {"model": "primary"}, _anthropic_messages_make_wrapper() - ): - pass - - assert exc_info.value is original_exception - - -@pytest.mark.asyncio -async def test_aanthropic_messages_fallback_attempt_yields_non_streaming_fallback_response(): - """Bugbot regression: a fallback that resolves to a non-streaming - response (no __aiter__, e.g. an agentic tool-use interception loop) must - be synthesized into a valid SSE byte sequence, not yielded as a raw dict - into a byte stream - the generator is typed AsyncGenerator[bytes, None] - and every item reaching the client must be a real SSE frame.""" - router = _anthropic_messages_make_router() - error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") - non_streaming_response = {"id": "msg_1", "type": "message", "content": [{"type": "text", "text": "hi"}]} - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=non_streaming_response), - ): - collected = [ - item - async for item in router._aanthropic_messages_fallback_attempt( - error, {"model": "primary"}, _anthropic_messages_make_wrapper() - ) - ] - - assert all(isinstance(item, bytes) for item in collected) - event_types = [item.split(b"\n")[0].removeprefix(b"event: ") for item in collected] - assert event_types == [ - b"message_start", - b"content_block_start", - b"content_block_delta", - b"content_block_stop", - b"message_delta", - b"message_stop", - ] - assert b'"text": "hi"' in collected[2] - - -@pytest.mark.asyncio -async def test_aanthropic_messages_fallback_attempt_reraises_plain_exception_on_double_failure(): - """Direct-call regression: when the fallback attempt fails with a plain - exception (not a MidStreamFallbackError), that exception itself must - propagate unchanged.""" - router = _anthropic_messages_make_router() - error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") - fallback_failure = ValueError("no healthy deployments") - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=fallback_failure), - ): - with pytest.raises(ValueError, match="no healthy deployments") as exc_info: - async for _ in router._aanthropic_messages_fallback_attempt( - error, {"model": "primary"}, _anthropic_messages_make_wrapper() - ): - pass - - assert exc_info.value is fallback_failure - - # -------- _aanthropic_messages_with_streaming_fallbacks -------- -@pytest.mark.asyncio -async def test_aanthropic_messages_with_streaming_fallbacks_non_streaming_passthrough(): - """A non-streaming response (plain dict) is returned unchanged, never wrapped.""" - router = _anthropic_messages_make_router() - plain_response = {"id": "msg_1", "type": "message"} - - async def fake_original(**_kwargs): - return plain_response - - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_helper", - new=AsyncMock(return_value=plain_response), - ): - out = await router._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=False, - ) - assert out is plain_response - - -@pytest.mark.asyncio -async def test_aanthropic_messages_with_streaming_fallbacks_wraps_streaming_iterator(): - """A streaming response is wrapped via _aanthropic_messages_streaming_iterator.""" - router = _anthropic_messages_make_router() - streaming_iter = _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk()]) - wrapped_marker = object() - - async def fake_original(**_kwargs): - return streaming_iter - - with ( - patch.object( - router, - "_ageneric_api_call_with_fallbacks_helper", - new=AsyncMock(return_value=streaming_iter), - ), - patch.object( - router, - "_aanthropic_messages_streaming_iterator", - new=AsyncMock(return_value=wrapped_marker), - ) as mock_wrap, - ): - out = await router._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - ) - assert out is wrapped_marker - mock_wrap.assert_awaited_once() - - # -------- mid-stream error handling -------- -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_on_pre_first_chunk_error_event(): - """Regression for #24004: a retriable SSE `event: error` frame - (overloaded_error/internal_server_error) that arrives before any real - content must trigger the router's fallback chain instead of passing - through to the client silently. The frame carries the error a 529 answer maps to, an InternalServerError, - so a failed fallback answers the status every other litellm path gives an overload.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()]) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - mock_fallback.assert_awaited_once() - raised = mock_fallback.await_args.kwargs["e"] - assert isinstance(raised, MidStreamFallbackError) - assert isinstance(raised.original_exception, litellm.InternalServerError) - assert raised.status_code == 500 - assert raised.is_pre_first_chunk is True - assert source.closed is True - - -@pytest.mark.asyncio -async def test_anthropic_messages_mid_stream_error_preserves_real_status_code(): - """Bugbot regression: the MidStreamFallbackError raised for a detected SSE - `event: error` frame must carry the error's REAL parsed status code - (via original_exception), not silently default to 503 for every error - type - a rate_limit_error (429) must surface as 429, not 503.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_rate_limit_error_chunk()]) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, - ) - [chunk async for chunk in wrapped] - - raised = mock_fallback.await_args.kwargs["e"] - assert isinstance(raised, MidStreamFallbackError) - assert raised.status_code == 429 - assert raised.original_exception is not None - assert raised.original_exception.status_code == 429 - assert raised.original_exception.llm_provider == "anthropic" - - -def test_merge_fallback_hidden_params_direct_call(): - """Direct-call regression: merge_fallback_hidden_params combines the - fallback's hidden params/headers with whatever was already present, - with the fallback's values winning on key collisions.""" - wrapper = FallbackAwareAnthropicMessagesStream( - _anthropic_messages_empty_generator(), - _AnthropicMessagesFakeByteStream([]), # carries {"additional_headers": {"x-amzn-requestid": "req-1"}} - ) - wrapper.merge_fallback_hidden_params( - {"model_id": "fallback-deployment"}, - {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"}, - ) - assert wrapper._hidden_params["model_id"] == "fallback-deployment" - assert wrapper._hidden_params["additional_headers"] == { - "x-amzn-requestid": "req-2", - "x-fallback-only": "yes", - } - - def test_anthropic_stream_forwards_ping_live_direct_call(): ping = _anthropic_messages_ping_chunk() content = _anthropic_messages_content_chunk("hi") @@ -19031,21 +10474,6 @@ def test_anthropic_stream_error_is_gateway_verdict_direct_call(): assert _anthropic_stream_error_is_gateway_verdict(_anthropic_messages_ping_chunk()) is False -def test_fallback_aware_stream_reports_withheld_output_of_its_current_source(): - """The proxy's cancel-refund guard reads this flag off the router wrapper, so it - must reflect the stream actually being drained: the primary, then the fallback.""" - - class _HoldingBack: - _hidden_params = {"additional_headers": {}} - has_buffered_provider_output = True - - wrapper = FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), _HoldingBack()) - assert wrapper.has_buffered_provider_output is True - - wrapper.adopt_fallback_source(_AnthropicMessagesFakeByteStream([])) - assert wrapper.has_buffered_provider_output is False - - def test_is_retriable_anthropic_status_direct_call(): assert _is_retriable_anthropic_status(429) is True assert _is_retriable_anthropic_status(503) is True @@ -19089,1200 +10517,6 @@ def test_anthropic_stream_commits_now_direct_call(): ) -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_merges_fallback_hidden_params(): - """Bugbot regression: after a successful mid-stream fallback, the - wrapper's _hidden_params must reflect the FALLBACK deployment's own - provider headers (e.g. a different Bedrock request-id), not stay - frozen on the primary's - raw bytes can't carry per-item _hidden_params - the way a ModelResponseStream/ResponsesAPI event can, so the wrapper - itself is the only place left to expose them.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_overloaded_error_chunk()] - ) # carries x-amzn-requestid: req-1 - fallback_stream = _AnthropicMessagesFallbackByteStream( - [_anthropic_messages_content_chunk("fallback answer")], - hidden_params={"additional_headers": {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"}}, - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ): - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - _ = [chunk async for chunk in wrapped] - - headers = wrapped._hidden_params["additional_headers"] - assert headers["x-amzn-requestid"] == "req-2" - assert headers["x-fallback-only"] == "yes" - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallback_headers", - [ - {"x-fallback-only": "yes"}, - { - "x-fallback-only": "yes", - "x-litellm-complexity-router-tier": "SIMPLE", - }, - ], - ids=["plain-fallback", "complexity-tier-fallback"], -) -@_PYTHON_CORE_ONLY -async def test_anthropic_messages_fallback_replaces_complexity_headers_before_output( - fallback_headers: dict[str, str], -) -> None: - primary_headers: Final = { - "x-litellm-complexity-router-tier": "REASONING", - "x-litellm-complexity-router-reasoning-effort": "xhigh", - } - source: Final = _AnthropicMessagesFallbackByteStream( - [_anthropic_messages_overloaded_error_chunk()], - hidden_params={"additional_headers": primary_headers}, - ) - fallback_output: Final = _anthropic_messages_content_chunk("fallback answer") - fallback: Final = _AnthropicMessagesFallbackByteStream( - [fallback_output], - hidden_params={ - "model_id": "fallback-deployment", - "additional_headers": fallback_headers, - }, - ) - router: Final = _injected_fallback_router(fallback) - - wrapped: Final = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - assert wrapped._hidden_params["additional_headers"] == primary_headers - first_output: Final = await wrapped.__anext__() - - assert first_output == fallback_output - assert wrapped.fallback_headers_adopted is True - assert wrapped._hidden_params == { - "model_id": "fallback-deployment", - "additional_headers": fallback_headers, - } - - -@pytest.mark.asyncio -async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_nested_metadata(): - """Bugbot regression: a shallow .copy() of kwargs still shares the - nested litellm_metadata/metadata dict objects with the primary attempt. - _update_kwargs_with_deployment mutates that dict in place with - deployment-specific fields, which must not leak into the fallback - request's metadata.""" - router = _anthropic_messages_make_router() - primary_metadata = {"model_group": "primary"} - streaming_iter_kwargs = {} - - async def fake_original(**_kwargs): - # Simulate _update_kwargs_with_deployment mutating the primary's - # litellm_metadata in place, as the real helper does. - primary_metadata["deployment"] = "primary-deployment-object" - return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")]) - - with patch.object( - router, - "_aanthropic_messages_streaming_iterator", - new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"), - ): - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_helper", - new=AsyncMock(side_effect=fake_original), - ): - await router._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - litellm_metadata=primary_metadata, - ) - - fallback_kwargs = streaming_iter_kwargs["initial_kwargs"] - assert fallback_kwargs["litellm_metadata"] is not primary_metadata - assert "deployment" not in fallback_kwargs["litellm_metadata"] - - -@pytest.mark.asyncio -async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_metadata_field(): - """Same regression as above for the (separate) `metadata` kwarg some - call sites use instead of `litellm_metadata`.""" - router = _anthropic_messages_make_router() - primary_metadata = {"tag": "primary"} - streaming_iter_kwargs = {} - - async def fake_original(**_kwargs): - primary_metadata["deployment"] = "primary-deployment-object" - return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")]) - - with patch.object( - router, - "_aanthropic_messages_streaming_iterator", - new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"), - ): - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_helper", - new=AsyncMock(side_effect=fake_original), - ): - await router._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - metadata=primary_metadata, - ) - - fallback_kwargs = streaming_iter_kwargs["initial_kwargs"] - assert fallback_kwargs["metadata"] is not primary_metadata - assert "deployment" not in fallback_kwargs["metadata"] - - -@pytest.mark.asyncio -async def test_anthropic_messages_hop_stream_failure_reaches_second_fallback_entry(): - """Regression: fallbacks=[{"primary": ["fb1", "fb2"]}]. The primary fails before - streaming, fb1 is reached through the regular fallback chain and then sends an - error frame mid-stream. Only the primary's stream used to be wrapped, so the outer - wrapper re-tried fb1 with a fresh attempted set and forwarded fb1's error frame to - the client on an HTTP 200; fb2 was unreachable.""" - router = Router( - model_list=[ - {"model_name": "primary", "litellm_params": {"model": "anthropic/primary-model", "api_key": "sk-test"}}, - {"model_name": "fb1", "litellm_params": {"model": "anthropic/fb1-model", "api_key": "sk-test"}}, - {"model_name": "fb2", "litellm_params": {"model": "anthropic/fb2-model", "api_key": "sk-test"}}, - ], - num_retries=0, - fallbacks=[{"primary": ["fb1", "fb2"]}], - ) - calls: list = [] - - async def fake_original(**kwargs): - model = kwargs["model"] - calls.append(model) - if model == "anthropic/primary-model": - raise litellm.InternalServerError(message="primary down", llm_provider="anthropic", model=model) - if model == "anthropic/fb1-model": - return _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - return _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb2")] - ) - - stream = await router._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - messages=[{"role": "user", "content": "hi"}], - max_tokens=10, - ) - body = b"".join([chunk async for chunk in stream]) - - assert calls == ["anthropic/primary-model", "anthropic/fb1-model", "anthropic/fb2-model"] - assert b"from fb2" in body - assert b"overloaded_error" not in body - - -_ANTHROPIC_MESSAGES_RETRY_GROUP: Final = ("anthropic/glm-a", "anthropic/glm-b") - - -def _anthropic_messages_retry_router( - num_retries: int, - deployment_params: Mapping[str, object] | None = None, - fallbacks: list[dict[str, list[str]]] | None = None, - context_window_fallbacks: list[dict[str, list[str]]] | None = None, - retry_policy: RetryPolicy | None = None, -) -> Router: - """Two deployments in the group, so a same-group retry waits for no backoff; no fallbacks unless asked.""" - group_deployments = [ - {"model_name": "glm", "litellm_params": {"model": model, "api_key": "sk-test", **(deployment_params or {})}} - for model in _ANTHROPIC_MESSAGES_RETRY_GROUP - ] - return Router( - model_list=[ - *group_deployments, - {"model_name": "fb", "litellm_params": {"model": "anthropic/fb-model", "api_key": "sk-test"}}, - {"model_name": "cw", "litellm_params": {"model": "anthropic/cw-model", "api_key": "sk-test"}}, - ], - num_retries=num_retries, - fallbacks=fallbacks or [], - context_window_fallbacks=context_window_fallbacks or [], - retry_policy=retry_policy, - ) - - -class _AnthropicMessagesScriptedProvider: - """Stands in for litellm.anthropic_messages: answers each call with the next scripted stream and records - the deployment it was routed to plus the retry counters the router stamped for that attempt.""" - - def __init__(self, *streams) -> None: - self._streams = list(streams) - self.calls: list[tuple[str, object, object]] = [] - - async def __call__(self, **kwargs): - litellm_metadata = kwargs.get("litellm_metadata") or {} - self.calls.append((kwargs["model"], litellm_metadata.get("attempted_retries"), litellm_metadata.get("max_retries"))) - assert self._streams, "provider called more times than scripted" - return self._streams.pop(0)() - - -def _anthropic_messages_transport_drop(original_exception: Exception | None = None) -> MidStreamFallbackError: - """What the completion bridge raises when the upstream closes the connection before any content.""" - return MidStreamFallbackError( - message="Connection closed", - model="glm", - llm_provider="databricks", - original_exception=original_exception - or litellm.APIConnectionError(message="Connection closed", llm_provider="databricks", model="glm"), - is_pre_first_chunk=True, - ) - - -def _anthropic_messages_dropped_before_content(): - return _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], _anthropic_messages_transport_drop()) - - -def _anthropic_messages_bridge_error_chunk() -> bytes: - from litellm.anthropic_interface.exceptions.exception_mapping_utils import anthropic_error_sse_frame - - return anthropic_error_sse_frame(status_code=500, raw_message="Connection closed").encode() - - -def _anthropic_messages_retried_stream(): - return _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - ) - - -async def _anthropic_messages_drain_into(stream, received: list) -> None: - async for chunk in stream: - received.append(chunk) - - -async def _anthropic_messages_stream_through_router(router: Router, provider, **request_kwargs): - return await router._aanthropic_messages_with_streaming_fallbacks( - original_function=provider, - model="glm", - stream=True, - messages=[{"role": "user", "content": "ping"}], - max_tokens=16, - **request_kwargs, - ) - - -_ANTHROPIC_MESSAGES_PRE_CONTENT_DROPS: Final = ( - pytest.param(_anthropic_messages_dropped_before_content, id="bridge-raises-before-content"), - pytest.param( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_bridge_error_chunk()] - ), - id="bridge-error-frame", - ), - pytest.param( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ), - id="provider-overloaded-frame", - ), -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("dropped_stream", _ANTHROPIC_MESSAGES_PRE_CONTENT_DROPS) -async def test_anthropic_messages_stream_dropped_before_content_is_retried_within_the_group(dropped_stream): - """Issue #44238: a /v1/messages stream the provider dropped before any content was answered after a - single upstream attempt, num_retries never applied. The drop is retried within the model group, with - the retry counters continuing the request's count, and the client sees one message lifecycle.""" - router = _anthropic_messages_retry_router(num_retries=2) - provider = _AnthropicMessagesScriptedProvider(dropped_stream, _anthropic_messages_retried_stream) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert all(model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls) - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_same_group_retry_skips_an_embedding_only_deployment(): - """The retry of a /v1/messages stream dropped before content must not land on an enforced embedding-only - deployment of the same group, even when that deployment is the group's preferred order.""" - router: Final = Router( - model_list=[ - {"model_name": "glm", "litellm_params": {"model": "anthropic/glm-chat", "api_key": "sk-test", "order": 2}}, - { - "model_name": "glm", - "litellm_params": {"model": "anthropic/glm-embed", "api_key": "sk-test", "order": 1}, - "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, - }, - ], - num_retries=1, - ) - provider: Final = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, _anthropic_messages_retried_stream - ) - - async def anthropic_messages(**kwargs: object) -> object: - return await provider(**kwargs) - - stream: Final = await _anthropic_messages_stream_through_router(router, anthropic_messages) - body: Final = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [model for model, _, _ in provider.calls] == ["anthropic/glm-chat", "anthropic/glm-chat"] - - -class _AnthropicMessagesRetryStashRecorder: - def __init__(self) -> None: - self.observed: tuple[Usage | None, ...] = () - - def record(self, usage: Usage | None) -> None: - self.observed = (*self.observed, usage) - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_partial_usage() -> None: - """The dropped attempt's usage stays available for terminal billing without reaching the retry's success log.""" - router: Final = _anthropic_messages_retry_router(num_retries=1) - logging_obj: Final = LiteLLMLogging( - model="glm", - messages=[{"role": "user", "content": "ping"}], - stream=True, - call_type="anthropic_messages", - start_time=datetime(2026, 1, 1), - litellm_call_id="retry-own-usage", - function_id="retry-own-usage", - ) - partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) - logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) - stash_recorder: Final = _AnthropicMessagesRetryStashRecorder() - usage_adapter: Final = TypeAdapter(Usage | None) - - def retry_stream() -> _AnthropicMessagesFakeByteStream: - stash_recorder.record( - usage_adapter.validate_python(logging_obj.model_call_details.get("combined_usage_object")) - ) - return _anthropic_messages_retried_stream() - - provider: Final = _AnthropicMessagesScriptedProvider(_anthropic_messages_dropped_before_content, retry_stream) - - stream: Final = await _anthropic_messages_stream_through_router( - router, provider, litellm_logging_obj=logging_obj - ) - body: Final = tuple([chunk async for chunk in stream]) - - assert stash_recorder.observed == (None,) - assert body == (_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")) - assert "response_cost" not in logging_obj.model_call_details - assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) - - -@pytest.mark.asyncio -async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_opens(): - """Both attempts log through the request's one logging object, so the superseded attempt's end-of-stream - logging must be queued before the retry starts producing its own: the router closes the dropped stream - before it asks the group for another.""" - router = _anthropic_messages_retry_router(num_retries=1) - superseded = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - superseded_closed_when_retry_opened: list[bool] = [] - - def retry_stream(): - superseded_closed_when_retry_opened.append(superseded.closed_as_superseded) - return _anthropic_messages_retried_stream() - - provider = _AnthropicMessagesScriptedProvider(lambda: superseded, retry_stream) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert superseded_closed_when_retry_opened == [True] - assert superseded.close_contexts[0] is True - - -@pytest.mark.asyncio -async def test_a_ptu_share_holders_messages_stream_dropped_before_content_is_retried_on_its_shared_group(monkeypatch): - """The same-group retry has to look the group's health up as the calling team; looked up with no - team a shared PTU group looks empty, so the share holder's dropped stream is never retried.""" - monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "true") - router: Final = Router( - model_list=[_shared_ptu_deployment("glm", f"shared-ptu-{i}") for i in (1, 2)], - num_retries=2, - ) - provider: Final = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, _anthropic_messages_retried_stream - ) - - stream: Final = await _anthropic_messages_stream_through_router( - router, provider, litellm_metadata={"user_api_key_team_id": "team-a"} - ) - body: Final = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert len(provider.calls) == 2 - - -@pytest.mark.asyncio -async def test_anthropic_messages_stream_dropped_after_content_keeps_the_error_and_is_not_retried(): - """A drop once content reached the client cannot be retried without a second overlapping message - lifecycle, so it keeps surfacing the provider's error after a single attempt.""" - router = _anthropic_messages_retry_router(num_retries=2) - drop = _anthropic_messages_transport_drop() - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesRaisingByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("par")], drop - ) - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - received = [] - with pytest.raises(litellm.APIConnectionError) as raised: - await _anthropic_messages_drain_into(stream, received) - - assert raised.value is drop.original_exception - assert received == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("par")] - assert len(provider.calls) == 1 - - -@pytest.mark.asyncio -async def test_anthropic_messages_retries_stop_at_num_retries_and_raise_the_last_drop(): - """Every retry's own stream continues the same count, so a group that keeps dropping is tried - exactly 1 + num_retries times before the provider's error reaches the client.""" - router = _anthropic_messages_retry_router(num_retries=2) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_dropped_before_content, - _anthropic_messages_dropped_before_content, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - with pytest.raises(litellm.APIConnectionError): - [chunk async for chunk in stream] - - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2), (2, 2)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_retries_run_out_before_the_fallback_chain_is_consulted(): - """Same-group retries come first; the fallback group is reached only once num_retries is spent.""" - router = _anthropic_messages_retry_router(num_retries=1, fallbacks=[{"glm": ["fb"]}]) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_dropped_before_content, - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - assert [model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls] == [True, True, False] - assert provider.calls[-1][0] == "anthropic/fb-model" - - -def _anthropic_messages_fb_deployment_hidden_params() -> dict: - return {"model_id": "fb-deployment", "additional_headers": {"x-litellm-model-group": "fb"}} - - -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_after_exhausted_retries_attributes_the_response_to_the_fallback_deployment(): - """The retry's stream carries a wrapper of its own, so a fallback it makes before its first byte must reach - the wrapper the proxy reads headers off: the response names the deployment that served it, not the primary.""" - router = _anthropic_messages_retry_router(num_retries=1, fallbacks=[{"glm": ["fb"]}]) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_dropped_before_content, - lambda: _AnthropicMessagesFallbackByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")], - hidden_params=_anthropic_messages_fb_deployment_hidden_params(), - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - assert stream._hidden_params["model_id"] == "fb-deployment" - assert stream._hidden_params["additional_headers"]["x-litellm-model-group"] == "fb" - assert stream._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 - - -def test_anthropic_messages_wrapper_follows_the_attribution_of_a_source_that_fell_back(): - inner = FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object()) - outer = FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), inner) - fallback = _AnthropicMessagesFallbackByteStream([], hidden_params=_anthropic_messages_fb_deployment_hidden_params()) - - outer.follow_source_attribution() - assert "model_id" not in outer._hidden_params - - inner.merge_fallback_hidden_params(*Router._prepare_fallback_hidden_params(fallback)) - inner.adopt_fallback_source(fallback) - outer.follow_source_attribution() - assert outer._hidden_params["model_id"] == "fb-deployment" - assert outer._hidden_params["additional_headers"]["x-litellm-model-group"] == "fb" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("drops", [1, 2]) -async def test_anthropic_messages_mid_stream_retries_are_counted_in_the_response_retry_headers(drops: int): - """A retry made after the stream opened never passes through async_function_with_retries, so the wrapper - stamps the retry headers that path would have and the client reads them along with the first byte.""" - router = _anthropic_messages_retry_router(num_retries=2) - provider = _AnthropicMessagesScriptedProvider( - *([_anthropic_messages_dropped_before_content] * drops), _anthropic_messages_retried_stream - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - first = await stream.__anext__() - headers = stream._hidden_params["additional_headers"] - - assert first == _anthropic_messages_message_start_chunk() - assert (headers["x-litellm-attempted-retries"], headers["x-litellm-max-retries"]) == (drops, 2) - assert [chunk async for chunk in stream] == [_anthropic_messages_content_chunk("pong")] - - -def _anthropic_messages_raise_authentication_error(): - raise litellm.AuthenticationError(message="invalid api key", llm_provider="anthropic", model="glm") - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_raising_a_non_retriable_error_is_handed_to_the_fallback_chain(): - """A retry that fails before its stream opens with an error no retry covers ends the retries and reaches - the fallback group the way a pre-stream failure does, instead of surfacing as the client's error.""" - router = _anthropic_messages_retry_router(num_retries=2, fallbacks=[{"glm": ["fb"]}]) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_raise_authentication_error, - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - assert [model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls] == [True, True, False] - - -def _anthropic_messages_raise_timeout(): - raise litellm.Timeout(message="upstream timed out", model="glm", llm_provider="databricks") - - -def _anthropic_messages_raise_internal_server_error(): - raise litellm.InternalServerError(message="upstream reset", llm_provider="databricks", model="glm") - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_raising_a_timeout_is_retried_like_a_pre_stream_timeout(): - """A 408 raised by a retry attempt before its stream opens is retried the way the pre-stream path retries - a 408, instead of ending the retries on the error-frame gate that only knows 429 and 5xx.""" - router = _anthropic_messages_retry_router(num_retries=3) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_raise_timeout, - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls] == [True, True, True] - - -@pytest.mark.asyncio -async def test_anthropic_messages_deployment_num_retries_also_governs_a_failure_before_the_stream_opens(): - """The deployment's num_retries litellm_param sets the budget for a failure raised before the stream opened - on this route, as it does for a mid-stream drop and for chat completions.""" - router = _anthropic_messages_retry_router(num_retries=0, deployment_params={"num_retries": 2}) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_raise_internal_server_error, - _anthropic_messages_raise_internal_server_error, - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert len(provider.calls) == 3 - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_raising_a_non_retriable_error_reaches_the_client_without_fallbacks(): - router = _anthropic_messages_retry_router(num_retries=2) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_raise_authentication_error, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - with pytest.raises(litellm.AuthenticationError): - [chunk async for chunk in stream] - - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] - - -def _anthropic_messages_raise_context_window_error(): - raise litellm.ContextWindowExceededError(message="prompt too long", llm_provider="anthropic", model="glm") - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_raising_a_context_window_error_takes_the_context_window_fallback(): - """The fallback chain sees the retry's own error type, so a context window overflow on the retried - deployment reaches context_window_fallbacks rather than the regular fallbacks.""" - router = _anthropic_messages_retry_router( - num_retries=2, fallbacks=[{"glm": ["fb"]}], context_window_fallbacks=[{"glm": ["cw"]}] - ) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_raise_context_window_error, - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from cw")] - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from cw")] - assert [model for model, _, _ in provider.calls][-1] == "anthropic/cw-model" - - -def test_anthropic_messages_retry_budget_precedence_direct_call(): - """A retry policy naming the error class outranks the request's num_retries, which outranks the routed - deployment's, which outranks the router's; num_retries=0 on the request turns a policy off too.""" - router = _anthropic_messages_retry_router(num_retries=3, deployment_params={"num_retries": 2}) - deployment_id = router.get_model_list(model_name="glm")[0]["model_info"]["id"] - routed = {"model": "glm", "litellm_metadata": {"model_info": {"id": deployment_id}}} - drop = litellm.APIConnectionError(message="closed", llm_provider="databricks", model="glm") - reset = litellm.InternalServerError(message="reset", llm_provider="databricks", model="glm") - policy_router = _anthropic_messages_retry_router( - num_retries=3, retry_policy=RetryPolicy(InternalServerErrorRetries=4) - ) - - assert router._stream_retry_budget(drop, {"model": "glm"}) == (3, False) - assert router._stream_retry_budget(drop, routed) == (2, False) - assert router._stream_retry_budget(drop, {**routed, "num_retries": 1}) == (1, False) - assert policy_router._stream_retry_budget(reset, {"model": "glm", "num_retries": 1}) == (4, True) - assert policy_router._stream_retry_budget(drop, {"model": "glm", "num_retries": 1}) == (1, False) - assert policy_router._stream_retry_budget(reset, {"model": "glm", "num_retries": 0}) == (0, False) - committed = {**routed, "litellm_metadata": {**routed["litellm_metadata"], "attempted_retries": 1, "max_retries": 5}} - assert router._stream_retry_budget(drop, committed) == (5, False) - assert policy_router._stream_retry_budget(reset, committed) == (5, True) - - -def test_anthropic_messages_stream_can_retry_direct_call(): - router = _anthropic_messages_retry_router(num_retries=1) - policy_router = _anthropic_messages_retry_router( - num_retries=0, retry_policy=RetryPolicy(InternalServerErrorRetries=1) - ) - - assert router._anthropic_messages_stream_can_retry({"model": "glm"}) is True - spent = {"model": "glm", "litellm_metadata": {"attempted_retries": 1}} - assert router._anthropic_messages_stream_can_retry(spent) is False - assert router._anthropic_messages_stream_can_retry({"model": "glm", "num_retries": 0}) is False - assert policy_router._anthropic_messages_stream_can_retry({"model": "glm"}) is True - assert policy_router._anthropic_messages_stream_can_retry(spent) is False - assert policy_router._anthropic_messages_stream_can_retry({"model": "glm", "num_retries": 0}) is False - assert policy_router._anthropic_messages_resolved_retry_policy({"model": "glm"}) is not None - assert policy_router._anthropic_messages_resolved_retry_policy({"model": "glm", "num_retries": 0}) is None - - -def test_retry_policy_ceiling_is_the_largest_budget_any_error_class_is_granted(): - from litellm.router import _retry_policy_ceiling - - assert _retry_policy_ceiling(RetryPolicy(InternalServerErrorRetries=1, RateLimitErrorRetries=3)) == 3 - assert _retry_policy_ceiling(RetryPolicy()) == 0 - - -@pytest.mark.asyncio -async def test_anthropic_messages_last_attempt_under_a_retry_policy_forwards_lifecycle_frames_live(): - """A retry policy bounds the hold the way a plain budget does: once the attempts reach the most retries the - policy grants, the stream is the last one, so its frames reach the client as they arrive and a drop after - them is the provider's error in-band rather than an error raised before any byte.""" - router = _anthropic_messages_retry_router(num_retries=0, retry_policy=RetryPolicy(DefaultRetries=1)) - drop = _anthropic_messages_transport_drop() - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - lambda: _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], drop), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - received: list = [] - with pytest.raises(litellm.APIConnectionError) as raised: - await _anthropic_messages_drain_into(stream, received) - - assert raised.value is drop.original_exception - assert received == [_anthropic_messages_message_start_chunk()] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 0), (1, 1)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_request_num_retries_zero_opts_out_of_the_mid_stream_retry(): - router = _anthropic_messages_retry_router(num_retries=2) - drop = _anthropic_messages_transport_drop() - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], drop) - ) - - stream = await _anthropic_messages_stream_through_router(router, provider, num_retries=0) - with pytest.raises(litellm.APIConnectionError) as raised: - [chunk async for chunk in stream] - - assert raised.value is drop.original_exception - assert len(provider.calls) == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("configured", [1, "1"], ids=["int", "config-string"]) -async def test_anthropic_messages_deployment_num_retries_sets_the_mid_stream_retry_budget(configured): - """A deployment's own num_retries litellm_param outranks the router's, as it does for a failure - raised before the stream opened.""" - router = _anthropic_messages_retry_router(num_retries=0, deployment_params={"num_retries": configured}) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, _anthropic_messages_retried_stream - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 0), (1, 1)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_policy_sets_the_mid_stream_retry_budget_per_error_class(): - router = _anthropic_messages_retry_router(num_retries=0, retry_policy=RetryPolicy(InternalServerErrorRetries=1)) - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesRaisingByteStream( - [_anthropic_messages_message_start_chunk()], - _anthropic_messages_transport_drop( - litellm.InternalServerError(message="upstream reset", llm_provider="databricks", model="glm") - ), - ), - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 0), (1, 1)] - - -def _anthropic_messages_error_frame(error_type: str) -> bytes: - return f"event: error\ndata: {json.dumps({'type': 'error', 'error': {'type': error_type, 'message': error_type}})}\n\n".encode() - - -_ANTHROPIC_MESSAGES_ERROR_FRAME_POLICIES: Final = ( - pytest.param("api_error", RetryPolicy(InternalServerErrorRetries=1), id="api_error-500-internal-server"), - pytest.param("overloaded_error", RetryPolicy(InternalServerErrorRetries=1), id="overloaded-internal-server"), - pytest.param("rate_limit_error", RetryPolicy(RateLimitErrorRetries=1), id="rate-limit-429"), - pytest.param("timeout_error", RetryPolicy(TimeoutErrorRetries=1), id="timeout-504"), -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("error_type,policy", _ANTHROPIC_MESSAGES_ERROR_FRAME_POLICIES) -async def test_anthropic_messages_error_frame_is_retried_under_the_class_the_pre_stream_mapping_gives_it( - error_type, policy -): - """An `event: error` frame before content carried a generic error, so a policy naming only error classes - granted it no retry while the hold still counted the policy: one attempt, then an HTTP error with no - bytes out. The frame now takes the class the pre-stream mapping raises for an answer carrying its body, so - an overloaded frame counts as the InternalServerError a 529 answer is, not a ServiceUnavailableError.""" - router = _anthropic_messages_retry_router(num_retries=0, retry_policy=policy) - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_error_frame(error_type)] - ), - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 0), (1, 1)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_error_frame_of_a_class_granted_no_retry_reaches_the_client_as_sent(): - """A rate limit frame is the RateLimitError a 429 answer is, so a policy granting only InternalServerError - retries leaves it unretried. With no fallback to take over either, the frame reaches the client as the - provider sent it, behind the lifecycle frames held back for a retry that never opened, the way the last - exhausted attempt's frames do; raising it instead turned a provider error frame into an HTTP error only - on the first attempt.""" - router = _anthropic_messages_retry_router(num_retries=0, retry_policy=RetryPolicy(InternalServerErrorRetries=1)) - frame = _anthropic_messages_error_frame("rate_limit_error") - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), frame]) - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), frame] - assert len(provider.calls) == 1 - - -@pytest.mark.asyncio -async def test_anthropic_messages_error_frame_of_a_class_granted_no_retry_still_reaches_a_configured_fallback(): - """The same unretried rate limit frame goes to the fallback group when one is configured, since a - fallback can still take over before any byte reached the client.""" - router = _anthropic_messages_retry_router( - num_retries=0, fallbacks=[{"glm": ["fb"]}], retry_policy=RetryPolicy(InternalServerErrorRetries=1) - ) - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_error_frame("rate_limit_error")] - ), - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - assert [model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls] == [True, False] - assert provider.calls[-1][0] == "anthropic/fb-model" - - -def test_anthropic_messages_recoverable_frame_error_direct_call(): - """Which `event: error` frames are intercepted for a retry or a fallback, and which reach the client as sent.""" - policy_router = _anthropic_messages_retry_router( - num_retries=0, retry_policy=RetryPolicy(InternalServerErrorRetries=1) - ) - fallback_router = _anthropic_messages_retry_router(num_retries=0, fallbacks=[{"glm": ["fb"]}]) - api_error = ("api_error", "reset", 500) - rate_limit = ("rate_limit_error", "slow down", 429) - kwargs = {"model": "glm"} - - recovered = policy_router._anthropic_messages_recoverable_frame_error(api_error, b"", False, "glm", kwargs) - assert isinstance(recovered, litellm.InternalServerError) - assert policy_router._anthropic_messages_recoverable_frame_error(rate_limit, b"", False, "glm", kwargs) is None - assert policy_router._anthropic_messages_recoverable_frame_error(api_error, b"", True, "glm", kwargs) is None - assert policy_router._anthropic_messages_recoverable_frame_error(None, b"", False, "glm", kwargs) is None - assert ( - policy_router._anthropic_messages_recoverable_frame_error( - ("invalid_request_error", "bad", 400), b"", False, "glm", kwargs - ) - is None - ) - spent = {"model": "glm", "litellm_metadata": {"attempted_retries": 1, "max_retries": 1}} - assert policy_router._anthropic_messages_recoverable_frame_error(api_error, b"", False, "glm", spent) is None - assert isinstance( - fallback_router._anthropic_messages_recoverable_frame_error(rate_limit, b"", False, "glm", kwargs), - litellm.RateLimitError, - ) - - -_ANTHROPIC_MESSAGES_MALFORMED_POLICIES: Final = ( - pytest.param({"glm": {"RateLimitErrorRetries": "many"}}, id="string-budget"), - pytest.param({"glm": 5}, id="group-policy-is-an-int"), - pytest.param(5, id="policy-map-is-an-int"), - pytest.param({"glm": [1]}, id="group-policy-is-a-list"), -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("policy", _ANTHROPIC_MESSAGES_MALFORMED_POLICIES) -@pytest.mark.parametrize("where", ["request", "router"]) -async def test_anthropic_messages_malformed_retry_policy_leaves_a_healthy_stream_alone(policy, where): - """The hold decision resolves the group's retry policy before the first byte, so a policy that does not - parse used to fail every stream of that group with a 500 before any attempt. It now governs nothing.""" - router = _anthropic_messages_retry_router(num_retries=0) - request_kwargs = {"model_group_retry_policy": policy} if where == "request" else {} - if where == "router": - router.model_group_retry_policy = policy - provider = _AnthropicMessagesScriptedProvider(_anthropic_messages_retried_stream) - - stream = await _anthropic_messages_stream_through_router(router, provider, **request_kwargs) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert len(provider.calls) == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("policy", _ANTHROPIC_MESSAGES_MALFORMED_POLICIES) -async def test_anthropic_messages_malformed_retry_policy_falls_back_to_the_plain_budget(policy): - router = _anthropic_messages_retry_router(num_retries=1) - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_error_frame("overloaded_error")] - ), - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider, model_group_retry_policy=policy) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 1), (1, 1)] - - -class _AnthropicMessagesAlternatingDeployments(CustomRoutingStrategyBase): - """Routes each attempt to the group's next deployment in turn, so which sibling a retry lands on is known.""" - - def __init__(self, router: Router, model_group: str) -> None: - self._deployments = itertools.cycle(router.get_model_list(model_name=model_group) or ()) - - async def async_get_available_deployment( - self, model, messages=None, input=None, specific_deployment=False, request_kwargs=None - ): - return next(self._deployments) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "first_num_retries,sibling_num_retries,expected_counters", - [ - pytest.param(3, 1, [(0, 0), (1, 3), (2, 3), (3, 3)], id="sibling-grants-fewer"), - pytest.param(1, 3, [(0, 0), (1, 1)], id="sibling-grants-more"), - ], -) -async def test_anthropic_messages_retries_keep_the_budget_the_first_drop_committed_to_across_deployments( - first_num_retries, sibling_num_retries, expected_counters -): - """A retry's stream recomputed its budget from the sibling deployment it landed on, so a group whose - deployments grant different num_retries stopped early or overshot the budget the first drop stamped - into the retry headers; later attempts now keep that budget, as the pre-stream retry loop does.""" - router = Router( - model_list=[ - { - "model_name": "glm", - "litellm_params": {"model": "anthropic/glm-a", "api_key": "sk-test", "num_retries": first_num_retries}, - }, - { - "model_name": "glm", - "litellm_params": {"model": "anthropic/glm-b", "api_key": "sk-test", "num_retries": sibling_num_retries}, - }, - ], - num_retries=0, - fallbacks=None, - ) - router.set_custom_routing_strategy(_AnthropicMessagesAlternatingDeployments(router, "glm")) - provider = _AnthropicMessagesScriptedProvider(*[_anthropic_messages_dropped_before_content] * len(expected_counters)) - - stream = await _anthropic_messages_stream_through_router(router, provider) - with pytest.raises(litellm.APIConnectionError): - [chunk async for chunk in stream] - - assert [model for model, _, _ in provider.calls] == (["anthropic/glm-a", "anthropic/glm-b"] * 2)[: len(expected_counters)] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == expected_counters - - -@pytest.mark.asyncio -async def test_anthropic_messages_lifecycle_frames_wait_for_content_while_a_retry_remains(): - """A retry can only restart cleanly while nothing reached the client, so with retries left a - fallback-less group holds message_start back until the first content frame, as a fallback does.""" - router = _anthropic_messages_retry_router(num_retries=2) - content_released = asyncio.Event() - - async def held_stream(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - provider = _AnthropicMessagesScriptedProvider(held_stream) - stream = await _anthropic_messages_stream_through_router(router, provider) - - pending = asyncio.ensure_future(stream.__anext__()) - await asyncio.sleep(0.2) - assert not pending.done() - content_released.set() - assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk() - assert [chunk async for chunk in stream] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_attempt_strips_the_controls_carrier_and_wraps_every_hop_stream(): - """Each attempt of the chain, not only the primary's, comes back wrapped for mid-stream - failover, and the per-request controls carrier never reaches the provider call.""" - from types import MappingProxyType - - from litellm.router_utils.fallback_event_handlers import ( - MID_STREAM_FALLBACK_CONTROLS_KEY, - MidStreamFallbackControls, - ) - - router = Router( - model_list=[ - {"model_name": "fb1", "litellm_params": {"model": "anthropic/fb1-model", "api_key": "sk-test"}}, - ], - num_retries=0, - ) - hop_stream = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb1")] - ) - seen: dict = {} - - async def fake_original(**kwargs): - seen.update(kwargs) - return hop_stream - - controls = MidStreamFallbackControls(MappingProxyType({"fallbacks": [{"primary": ["fb1", "fb2"]}]})) - stream = await router._ageneric_api_call_with_fallbacks_anthropic_messages_attempt( - model="fb1", - original_generic_function=fake_original, - stream=True, - messages=[{"role": "user", "content": "hi"}], - max_tokens=10, - **{MID_STREAM_FALLBACK_CONTROLS_KEY: controls}, - ) - body = b"".join([chunk async for chunk in stream]) - - assert seen["model"] == "anthropic/fb1-model" - assert MID_STREAM_FALLBACK_CONTROLS_KEY not in seen - assert "fallbacks" not in seen - assert stream is not hop_stream - assert b"from fb1" in body - - -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_triggers_after_lifecycle_only_frame(): - """Regression: Anthropic routinely sends a message_start lifecycle frame - before an overload error even fires. A lifecycle-only frame (no real - content) must not disqualify the fallback attempt, and must not reach - the client either - forwarding it and then appending the fallback's own - message_start would produce two overlapping message lifecycles on one - SSE stream. The primary's buffered lifecycle frame is discarded and the - client sees only the fallback's own, single, clean lifecycle.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - fallback_message_start = b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_2"}}\n\n' - fallback_stream = _AnthropicMessagesFallbackByteStream( - [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] - assert collected.count(_anthropic_messages_message_start_chunk()) == 0, ( - "the primary's message_start must never reach the client" - ) - assert sum(1 for c in collected if c.startswith(b"event: message_start")) == 1, ( - "exactly one message_start must reach the client" - ) - mock_fallback.assert_awaited_once() - raised = mock_fallback.await_args.kwargs["e"] - assert raised.is_pre_first_chunk is True - - -@pytest.mark.asyncio -async def test_anthropic_messages_raised_error_after_real_content_does_not_restart_stream(): - """Regression: a MidStreamFallbackError raised directly by the source - iterator (the completion-bridge path's CustomStreamWrapper, e.g. a - transport drop) must not trigger a fallback once real content already - reached the client - that would append a second, overlapping message - lifecycle onto the same SSE stream. The original exception must - propagate to the caller instead.""" - router = _anthropic_messages_make_router() - content = _anthropic_messages_content_chunk("partial answer") - original_exception = litellm.APIError( - status_code=503, - message="stream reset", - llm_provider="vertex_ai", - model="primary", - ) - raised_error = MidStreamFallbackError( - message="stream reset", - model="primary", - llm_provider="vertex_ai", - original_exception=original_exception, - is_pre_first_chunk=False, - ) - source = _AnthropicMessagesRaisingByteStream([content], raised_error) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(litellm.APIError) as exc_info: - await _consume() - - assert collected == [content] - assert exc_info.value is original_exception - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_also_catches_raised_midstream_error(): - """Regression for the completion-bridge path (deployments with no native - /v1/messages endpoint): its CustomStreamWrapper raises - MidStreamFallbackError directly (e.g. on a Vertex AI transport drop) - instead of yielding an SSE error chunk - the wrapper must catch that too.""" - router = _anthropic_messages_make_router() - raised_error = MidStreamFallbackError( - message="stream reset", - model="primary", - llm_provider="vertex_ai", - is_pre_first_chunk=True, - ) - source = _AnthropicMessagesRaisingByteStream([], raised_error) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - mock_fallback.assert_awaited_once() - assert mock_fallback.await_args.kwargs["e"] is raised_error - - -_MID_STREAM_OPT_OUT_SHAPES: Final = ( - pytest.param({"disable_fallbacks": True}, id="raw-kwarg"), - pytest.param({"metadata": {DISABLE_FALLBACKS_METADATA_KEY: True}}, id="metadata-stamp"), - pytest.param({"litellm_metadata": {DISABLE_FALLBACKS_METADATA_KEY: True}}, id="litellm_metadata-stamp"), -) - - def _mid_stream_opt_out_router() -> Router: return Router( model_list=[ @@ -20328,102 +10562,6 @@ class _MidStreamOptOutChatStream(CustomStreamWrapper): raise self._error -@pytest.mark.asyncio -@pytest.mark.parametrize("opt_out", _MID_STREAM_OPT_OUT_SHAPES) -async def test_acompletion_streaming_iterator_honors_disable_fallbacks(opt_out): - """A chat stream that fails before its first chunk on a request that opted out of fallbacks - surfaces the primary's own error and never tries the fallback deployment.""" - router = _mid_stream_opt_out_router() - primary_error = _mid_stream_opt_out_primary_error() - source = _MidStreamOptOutChatStream(_mid_stream_opt_out_trigger(primary_error)) - - with patch.object(router, "_acompletion", new=AsyncMock(return_value=_AsyncList([]))) as fallback_attempt: - wrapped = await router._acompletion_streaming_iterator( - model_response=source, - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "primary", "stream": True, **copy.deepcopy(opt_out)}, - ) - with pytest.raises(litellm.InternalServerError) as raised: - [chunk async for chunk in wrapped] - - assert raised.value is primary_error - fallback_attempt.assert_not_awaited() - - -@pytest.mark.parametrize("opt_out", _MID_STREAM_OPT_OUT_SHAPES) -def test_completion_streaming_iterator_honors_disable_fallbacks(opt_out): - """Sync counterpart of test_acompletion_streaming_iterator_honors_disable_fallbacks.""" - router = _mid_stream_opt_out_router() - primary_error = _mid_stream_opt_out_primary_error() - source = _MidStreamOptOutChatStream(_mid_stream_opt_out_trigger(primary_error)) - - with patch.object(router, "_completion", new=MagicMock(return_value=iter([]))) as fallback_attempt: - wrapped = router._completion_streaming_iterator( - model_response=source, - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "primary", "stream": True, **copy.deepcopy(opt_out)}, - ) - with pytest.raises(litellm.InternalServerError) as raised: - list(wrapped) - - assert raised.value is primary_error - fallback_attempt.assert_not_called() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("opt_out", _MID_STREAM_OPT_OUT_SHAPES) -async def test_aresponses_streaming_iterator_honors_disable_fallbacks(opt_out): - """Same opt-out contract on the Responses API mid-stream fallback path.""" - router = _mid_stream_opt_out_router() - primary_error = _mid_stream_opt_out_primary_error() - source = _make_responses_iterator(error=_mid_stream_opt_out_trigger(primary_error), model="primary") - - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_responses_attempt", - new=AsyncMock(return_value=_AsyncList([])), - ) as fallback_attempt: - wrapped = await router._aresponses_streaming_iterator( - response=source, - initial_kwargs={ - "model": "primary", - "stream": True, - "input": "Hi", - "original_generic_function": litellm.aresponses, - **copy.deepcopy(opt_out), - }, - ) - with pytest.raises(litellm.InternalServerError) as raised: - [chunk async for chunk in wrapped] - - assert raised.value is primary_error - fallback_attempt.assert_not_awaited() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("opt_out", _MID_STREAM_OPT_OUT_SHAPES) -async def test_anthropic_messages_streaming_iterator_honors_disable_fallbacks(opt_out): - """Same opt-out contract on the Anthropic Messages mid-stream fallback path.""" - router = _mid_stream_opt_out_router() - primary_error = _mid_stream_opt_out_primary_error() - source = _AnthropicMessagesRaisingByteStream([], _mid_stream_opt_out_trigger(primary_error)) - - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_anthropic_messages_attempt", - new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), - ) as fallback_attempt: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary", "stream": True, **copy.deepcopy(opt_out)}, - ) - with pytest.raises(litellm.InternalServerError) as raised: - [chunk async for chunk in wrapped] - - assert raised.value is primary_error - fallback_attempt.assert_not_awaited() - - @pytest.mark.asyncio async def test_acompletion_disable_fallbacks_reaches_the_mid_stream_hop(): """`disable_fallbacks=True` sent to the public entrypoint survives the fallback wrapper's handoff @@ -20462,43 +10600,6 @@ def test_completion_disable_fallbacks_reaches_the_mid_stream_hop(): assert [call.kwargs["metadata"]["model_group"] for call in provider_calls.call_args_list] == ["primary"] -@pytest.mark.asyncio -@pytest.mark.parametrize( - "raised_error", - [ - BedrockError(status_code=503, message='serviceUnavailableException {"message": "Service unavailable"}'), - BedrockError(status_code=500, message='internalServerException {"message": "Internal error"}'), - BedrockError(status_code=429, message='throttlingException {"message": "Too many requests"}'), - httpx.ReadError("connection reset by upstream"), - ], - ids=["503", "500", "429", "transport-drop"], -) -async def test_anthropic_messages_raised_provider_error_before_content_triggers_fallback(raised_error): - """A retriable raise before content falls over exactly like a detected SSE error event.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - mock_fallback.assert_awaited_once() - converted = mock_fallback.await_args.kwargs["e"] - assert isinstance(converted, MidStreamFallbackError) - assert converted.original_exception is raised_error - assert converted.is_pre_first_chunk is True - assert source.closed is True - - class _AnthropicMessagesStringStatusError(Exception): def __init__(self): super().__init__("bad request") @@ -20511,78 +10612,6 @@ class _AnthropicMessagesResponseOnlyStatusError(Exception): self.response = SimpleNamespace(status_code=400) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "raised_error", - [ - BedrockError(status_code=400, message='validationException {"message": "Malformed input"}'), - BedrockError(status_code=424, message='modelStreamErrorException {"message": "Model stream error"}'), - _AnthropicMessagesStringStatusError(), - _AnthropicMessagesResponseOnlyStatusError(), - ], - ids=["400", "424", "str-400", "response-only-400"], -) -async def test_anthropic_messages_raised_non_retriable_provider_error_propagates_unchanged(raised_error): - """A raised client error reaches the caller as the same exception, nothing flushed, no fallback.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(type(raised_error)) as exc_info: - await _consume() - - assert collected == [] - assert exc_info.value is raised_error - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_raised_provider_error_after_content_propagates_unchanged(): - """A raise after content propagates unchanged even when its status is retriable.""" - router = _anthropic_messages_make_router() - content = _anthropic_messages_content_chunk("partial answer") - raised_error = BedrockError( - status_code=503, message='serviceUnavailableException {"message": "Service unavailable"}' - ) - source = _AnthropicMessagesRaisingByteStream([content], raised_error) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(BedrockError) as exc_info: - await _consume() - - assert collected == [content] - assert exc_info.value is raised_error - mock_fallback.assert_not_awaited() - - @pytest.mark.parametrize( "error, expected_status", [ @@ -20597,340 +10626,9 @@ def test_anthropic_stream_raised_error_status_reads_every_status_shape(error, ex assert _stream_raised_error_status(error) == expected_status -@pytest.mark.parametrize( - "error, has_generated_content, converts", - [ - (BedrockError(status_code=503, message="unavailable"), False, True), - (httpx.ReadError("connection reset by upstream"), False, True), - (BedrockError(status_code=400, message="malformed"), False, False), - (BedrockError(status_code=503, message="unavailable"), True, False), - ], - ids=["retriable", "no-status", "client-error", "after-content"], -) -def test_anthropic_stream_fallback_error_for_raised_gates_like_a_detected_error_event( - error, has_generated_content, converts -): - converted = _stream_fallback_error_for_raised(error, "primary", has_generated_content) - if not converts: - assert converted is None - return - assert isinstance(converted, MidStreamFallbackError) - assert converted.original_exception is error - assert converted.is_pre_first_chunk is True - assert converted.llm_provider == "anthropic" - - -@pytest.mark.asyncio -async def test_aanthropic_messages_recover_stream_error_flushes_buffered_frames_before_declining(): - router = _anthropic_messages_make_router() - original = BedrockError(status_code=503, message="unavailable") - declined = MidStreamFallbackError( - message="unavailable", - model="primary", - llm_provider="anthropic", - original_exception=original, - is_pre_first_chunk=False, - ) - buffered = (_anthropic_messages_message_start_chunk(),) - flushed = [] - - async def drain(recovery) -> None: - async for chunk in recovery: - flushed.append(chunk) - - with patch.object(router, "_aanthropic_messages_fallback_attempt") as mock_attempt: - recovery = router._aanthropic_messages_recover_stream_error( - declined, True, buffered, "primary", {"model": "primary"}, _anthropic_messages_make_wrapper(), AsyncMock() - ) - with pytest.raises(BedrockError) as exc_info: - await drain(recovery) - assert flushed == list(buffered) - assert exc_info.value is original - mock_attempt.assert_not_called() - - -@pytest.mark.asyncio -async def test_aanthropic_messages_recover_stream_error_hands_converted_raise_to_fallback_attempt(): - router = _anthropic_messages_make_router() - raised = BedrockError(status_code=503, message="unavailable") - handed_over = [] - - async def fake_attempt(fallback_error, initial_kwargs, wrapper): - handed_over.append(fallback_error) - yield b"fallback" - - with patch.object(router, "_aanthropic_messages_fallback_attempt", new=fake_attempt): - recovery = router._aanthropic_messages_recover_stream_error( - raised, False, (), "primary", {"model": "primary"}, _anthropic_messages_make_wrapper(), AsyncMock() - ) - collected = [chunk async for chunk in recovery] - assert collected == [b"fallback"] - assert len(handed_over) == 1 - assert isinstance(handed_over[0], MidStreamFallbackError) - assert handed_over[0].original_exception is raised - - -@pytest.mark.asyncio -async def test_anthropic_messages_non_retriable_client_error_skips_fallback(): - """A 4xx (non-429) error type (e.g. invalid_request_error) is a client - error a fallback attempt cannot fix, so it must be forwarded to the - client as-is rather than burning a fallback attempt.""" - router = _anthropic_messages_make_router() - error_chunk = _anthropic_messages_invalid_request_error_chunk() - source = _AnthropicMessagesFakeByteStream([error_chunk]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [error_chunk] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_post_first_chunk_error_skips_fallback(): - """Once content has already reached the caller, retrying would start a - second, overlapping Anthropic message lifecycle on the same SSE stream - - the error must be forwarded instead of triggering an invisible retry.""" - router = _anthropic_messages_make_router() - content = _anthropic_messages_content_chunk("partial answer") - error_chunk = _anthropic_messages_overloaded_error_chunk() - source = _AnthropicMessagesFakeByteStream([content, error_chunk]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [content, error_chunk] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_stream_the_client_closes_at_a_forwarded_error_is_not_closed_as_superseded(): - """An error frame forwarded to the client ends that request, and the SDK hangs up on it. No retry took the - request over, so the source must not be closed as a superseded attempt, which would log it as a failure and - drop the usage the provider billed for it.""" - router = _anthropic_messages_make_router() - error_chunk = _anthropic_messages_overloaded_error_chunk() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("partial answer"), error_chunk]) - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) - async for chunk in wrapped: - if chunk == error_chunk: - break - await wrapped.aclose() - - assert source.closed is True - assert source.closed_as_superseded is False - - -@pytest.mark.asyncio -async def test_anthropic_messages_non_retriable_error_flushes_buffered_lifecycle_frames(): - """A non-retriable error arriving while lifecycle frames are still - buffered (no content seen yet) must flush those buffered frames before - forwarding the error, so the client still sees the whole primary - attempt rather than losing the buffered message_start silently.""" - router = _anthropic_messages_make_router() - error_chunk = _anthropic_messages_invalid_request_error_chunk() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), error_chunk]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_message_start_chunk(), error_chunk] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_raised_error_declined_flushes_buffered_lifecycle_frames(): - """When a raised MidStreamFallbackError is declined (source says content - was not pre-first-chunk) while lifecycle frames are still buffered, they - must be flushed to the client before the exception propagates.""" - router = _anthropic_messages_make_router() - raised_error = MidStreamFallbackError( - message="stream reset", - model="primary", - llm_provider="vertex_ai", - is_pre_first_chunk=False, - ) - source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(MidStreamFallbackError) as exc_info: - await _consume() - - assert collected == [_anthropic_messages_message_start_chunk()] - assert exc_info.value is raised_error - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_raised_error_without_original_exception_reraises_itself(): - """When a declined MidStreamFallbackError carries no original_exception, - the bare exception itself must propagate rather than being swallowed.""" - router = _anthropic_messages_make_router() - content = _anthropic_messages_content_chunk("partial answer") - raised_error = MidStreamFallbackError( - message="stream reset", - model="primary", - llm_provider="vertex_ai", - is_pre_first_chunk=False, - ) - source = _AnthropicMessagesRaisingByteStream([content], raised_error) - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(MidStreamFallbackError) as exc_info: - await _consume() - - assert collected == [content] - assert exc_info.value is raised_error - - -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_also_failing_raises_original_exception(): - """If the fallback attempt itself fails with a MidStreamFallbackError - wrapping a real provider exception, the client must see that real - exception, not the internal MidStreamFallbackError.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()]) - original_exception = litellm.APIError( - status_code=503, - message="fallback also overloaded", - llm_provider="bedrock", - model="fallback", - ) - fallback_failure = MidStreamFallbackError( - message="fallback failed", - model="fallback", - llm_provider="bedrock", - original_exception=original_exception, - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=fallback_failure), - ): - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - with pytest.raises(litellm.APIError) as exc_info: - async for _ in wrapped: - pass - - assert exc_info.value is original_exception - - # -------- _dispatch_generic_call_type -------- -@pytest.mark.asyncio -async def test_dispatch_generic_call_type_routes_anthropic_messages_through_streaming_fallbacks(): - router = _anthropic_messages_make_router() - - async def fake_original(**_kwargs): - return {"id": "msg_1"} - - with patch.object( - router, - "_aanthropic_messages_with_streaming_fallbacks", - new=AsyncMock(return_value="anthropic-result"), - ) as mock_anthropic: - out = await router._dispatch_generic_call_type( - call_type="anthropic_messages", - original_function=fake_original, - model="primary", - ) - assert out == "anthropic-result" - mock_anthropic.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_dispatch_generic_call_type_other_call_types_use_generic_fallback(): - router = _anthropic_messages_make_router() - - async def fake_original(**_kwargs): - return {"id": "file_1"} - - with patch.object( - router, - "_ageneric_api_call_with_fallbacks", - new=AsyncMock(return_value="generic-result"), - ) as mock_generic: - out = await router._dispatch_generic_call_type( - call_type="afile_delete", - original_function=fake_original, - model="primary", - ) - assert out == "generic-result" - mock_generic.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_factory_function_anthropic_messages_uses_streaming_fallback_dispatch(): - """anthropic_messages must be wired through the mid-stream-fallback-aware - path rather than the bare generic dispatch every other call type without - special handling uses.""" - router = _anthropic_messages_make_router() - wrapped = router.factory_function(litellm.anthropic_messages, call_type="anthropic_messages") - assert callable(wrapped) - - with patch.object( - router, - "_aanthropic_messages_with_streaming_fallbacks", - new=AsyncMock(return_value="ok"), - ) as mock_anthropic: - result = await wrapped(model="primary") - assert result == "ok" - mock_anthropic.assert_awaited_once() - - @pytest.mark.asyncio async def test_async_function_with_fallbacks_stamps_zero_attempted_fallbacks(): """A request served by the primary model group records attempted_fallbacks=0 and @@ -21062,159 +10760,6 @@ async def test_async_function_with_fallbacks_skips_stamp_on_genuine_reentrant_ho assert metadata["original_model_group"] == "prod-chat" -def _record_router_acompletion_kwargs(router: litellm.Router) -> list: - """Spy on router._acompletion, recording each call's kwargs while delegating through.""" - records = [] - original_acompletion = router._acompletion - - @functools.wraps(original_acompletion) - async def _spy(*args, **spy_kwargs): - records.append(spy_kwargs) - return await original_acompletion(*args, **spy_kwargs) - - router._acompletion = _spy - return records - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_scrubs_spoofed_values_from_sibling_bucket(): - """Spend logs read a truthy litellm_metadata dict in preference to metadata, so spoofed - stamp keys planted in the bucket the route does not own are removed on entry, in place, - before they can flow into the spend log row.""" - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, - } - ] - ) - metadata = {} - litellm_metadata = { - "attempted_fallbacks": 99, - "original_model_group": "spoofed-group", - "client_key": "client_value", - } - downstream_calls = _record_router_acompletion_kwargs(router) - - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hey"}], - metadata=metadata, - litellm_metadata=litellm_metadata, - ) - - assert len(downstream_calls) == 1 - downstream_sibling = downstream_calls[0]["litellm_metadata"] - assert "attempted_fallbacks" not in downstream_sibling - assert "original_model_group" not in downstream_sibling - assert downstream_sibling["client_key"] == "client_value" - assert "attempted_fallbacks" not in litellm_metadata - assert "original_model_group" not in litellm_metadata - assert litellm_metadata["client_key"] == "client_value" - assert metadata["attempted_fallbacks"] == 0 - assert metadata["original_model_group"] == "gpt-3.5-turbo" - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_scrubs_sibling_bucket_in_place(): - """Everything below the router resolves the bucket by key presence, so the scrub edits - the caller's dict object like every other router bucket write. Rebinding kwargs to a - scrubbed copy detaches the proxy's request_data write-backs (guardrail telemetry, retry - accounting) from the object the spend row is built from.""" - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, - } - ] - ) - litellm_metadata = { - "attempted_fallbacks": 7, - "original_model_group": "planted-group", - "client_key": "client_value", - } - caller_snapshot = copy.deepcopy(litellm_metadata) - downstream_calls = _record_router_acompletion_kwargs(router) - - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hey"}], - metadata={}, - litellm_metadata=litellm_metadata, - ) - - assert len(downstream_calls) == 1 - assert downstream_calls[0]["litellm_metadata"] is litellm_metadata - assert "attempted_fallbacks" not in litellm_metadata - assert "original_model_group" not in litellm_metadata - assert litellm_metadata["client_key"] == caller_snapshot["client_key"] - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_stamps_aliased_buckets_on_every_call(): - """One dict object passed as both metadata and litellm_metadata: the first call's own - stamp puts the reserved keys into the shared object, so the second call enters the - scrub with them present. Scrubbing in place keeps the stamp and the bucket on the same - object; a scrubbed copy would leave the spend reader's preferred bucket unstamped.""" - router = litellm.Router( - model_list=[ - { - "model_name": "chat-group", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, - } - ] - ) - shared_metadata = {"team": "alpha"} - downstream_calls = _record_router_acompletion_kwargs(router) - - for _ in range(3): - await router.acompletion( - model="chat-group", - messages=[{"role": "user", "content": "hey"}], - metadata=shared_metadata, - litellm_metadata=shared_metadata, - ) - - assert len(downstream_calls) == 3 - for call_kwargs in downstream_calls: - assert call_kwargs["litellm_metadata"] is shared_metadata - assert call_kwargs["metadata"] is shared_metadata - assert call_kwargs["litellm_metadata"]["attempted_fallbacks"] == 0 - assert call_kwargs["litellm_metadata"]["original_model_group"] == "chat-group" - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_passes_clean_sibling_bucket_through_unchanged(): - """A sibling bucket carrying no reserved stamp keys is forwarded downstream as the - caller's own object with no copy made, matching pre-scrub behavior. Retry accounting - stamped into that bucket downstream predates the scrub and is out of its scope.""" - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, - } - ] - ) - litellm_metadata = {"client_key": "client_value"} - downstream_calls = _record_router_acompletion_kwargs(router) - - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hey"}], - metadata={}, - litellm_metadata=litellm_metadata, - ) - - assert len(downstream_calls) == 1 - assert downstream_calls[0]["litellm_metadata"] is litellm_metadata - assert litellm_metadata["client_key"] == "client_value" - assert "attempted_fallbacks" not in litellm_metadata - assert "original_model_group" not in litellm_metadata - - @pytest.mark.asyncio async def test_run_async_fallback_keeps_caller_metadata_keys_on_the_wire(monkeypatch): """Under enable_preview_features, add_openai_metadata forwards only the first 16 @@ -21287,59 +10832,6 @@ async def test_run_async_fallback_keeps_caller_metadata_keys_on_the_wire(monkeyp assert small_wire["model_group"] == "fallback-group" -@pytest.mark.asyncio -async def test_run_async_fallback_two_hop_chain_reports_entry_group_and_hop_count(): - """A two-hop fallback chain stamps attempted_fallbacks=2 on the final leg and keeps - original_model_group at the group requested on entry: a later hop's stamp appends - after caller keys without overriding the value stamped by an earlier hop.""" - router = litellm.Router( - model_list=[ - { - "model_name": "group-a", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "litellm.InternalServerError"}, - }, - { - "model_name": "group-b", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "litellm.InternalServerError"}, - }, - { - "model_name": "group-c", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "ok"}, - }, - ], - fallbacks=[{"group-a": ["group-b"]}, {"group-b": ["group-c"]}], - num_retries=0, - ) - metadata = {} - leg_records = [] - original_acompletion = router._acompletion - - @functools.wraps(original_acompletion) - async def _spy(*args, **spy_kwargs): - leg_records.append((spy_kwargs.get("model"), copy.deepcopy(spy_kwargs.get("metadata")))) - return await original_acompletion(*args, **spy_kwargs) - - router._acompletion = _spy - - await router.acompletion( - model="group-a", - messages=[{"role": "user", "content": "hey"}], - metadata=metadata, - ) - - assert [model for model, _ in leg_records] == ["group-a", "group-b", "group-c"] - hop_one_metadata = leg_records[1][1] - assert hop_one_metadata["attempted_fallbacks"] == 1 - assert hop_one_metadata["original_model_group"] == "group-a" - assert hop_one_metadata["model_group"] == "group-b" - hop_two_metadata = leg_records[2][1] - assert hop_two_metadata["attempted_fallbacks"] == 2 - assert hop_two_metadata["original_model_group"] == "group-a" - assert hop_two_metadata["model_group"] == "group-c" - assert metadata["attempted_fallbacks"] == 0 - assert metadata["original_model_group"] == "group-a" - - def _permission_denied_error() -> litellm.PermissionDeniedError: return litellm.PermissionDeniedError( message="OpenrouterException - this key has no access to the model", @@ -21419,9 +10911,7 @@ def _router_with_failing_primary(fallback_access_check) -> Router: class _DenyFallbackBudgetCheck: - async def __call__( - self, *, model: str, request_kwargs: Mapping[str, object], ctx: RouterContext - ) -> bool: + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: RouterContext) -> bool: return False @@ -21608,237 +11098,6 @@ def test_resolved_litellm_models_base_model_precedence_is_opt_in() -> None: assert router.resolved_litellm_models("never-configured", use_base_model=True) == () -class TestTierParamsTheTargetAccepts: - """A tier's litellm_params are applied to every request that tier routes, so one the target - cannot take raised UnsupportedParamsError before the request left the proxy, turning the whole - tier into a 400.""" - - @pytest.fixture(autouse=True) - def force_local_model_cost(self, monkeypatch): - from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap - - monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) - - @staticmethod - def _router(model: str) -> litellm.Router: - return litellm.Router( - model_list=[{"model_name": "tiered", "litellm_params": {"model": model, "api_key": "sk-x"}}] - ) - - def test_drops_a_param_no_deployment_declares(self): - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("tiered", {"reasoning_effort": "max"}, {}) - - assert accepted == {} - - def test_keeps_a_param_the_deployment_declares(self): - router = self._router("fireworks_ai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("tiered", {"reasoning_effort": "max"}, {}) - - assert accepted == {"reasoning_effort": "max"} - - @pytest.mark.parametrize( - "control, value", - [ - ("api_base", "https://example.invalid"), - ("api_key", "sk-tier"), - ("base_url", "https://example.invalid"), - ("timeout", 30), - ("default_headers", {"x-tier": "1"}), - ("organization", "org-tier"), - ("deployment_id", "dep-tier"), - ], - ) - def test_keeps_credentials_and_transport_controls(self, control, value): - """These are not chat completion params, so get_optional_params never compares them against - a provider's supported list. Filtering on "is this an OpenAI param" would discard the - configuration the request needs while never touching what the provider would reject.""" - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("tiered", {control: value, "reasoning_effort": "max"}, {}) - - assert accepted == {control: value} - - @pytest.mark.parametrize( - "control, value", - [ - ("additional_drop_params", ["seed"]), - ("drop_params", True), - ("allowed_openai_params", ["seed"]), - ("api_version", "2024-02-01"), - ("metadata", {"tier": "complex"}), - ], - ) - def test_keeps_litellm_controls_the_provider_never_lists(self, control, value): - """No provider lists a litellm control among its supported params, so "no deployment - declares it" means litellm consumes it, not that the target refuses it. Dropping - drop_params or additional_drop_params would silently disable the operator's sanitization.""" - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("tiered", {control: value, "reasoning_effort": "max"}, {}) - - assert accepted == {control: value} - - def test_tier_allowlist_protects_the_param_it_names(self): - """allowed_openai_params is the documented escape hatch for an incomplete supported-params - list, and request-time validation extends the supported list with it, so a param the tier - both sets and allowlists would never 400 and must not be dropped.""" - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"reasoning_effort": "max", "allowed_openai_params": ["reasoning_effort"]}, {} - ) - - assert accepted == {"reasoning_effort": "max", "allowed_openai_params": ["reasoning_effort"]} - - def test_request_allowlist_protects_the_param_it_names(self): - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"reasoning_effort": "max"}, {"allowed_openai_params": ["reasoning_effort"]} - ) - - assert accepted == {"reasoning_effort": "max"} - - def test_allowlist_protects_only_the_params_it_names(self): - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"reasoning_effort": "max", "allowed_openai_params": ["seed"]}, {} - ) - - assert accepted == {"allowed_openai_params": ["seed"]} - - def test_declared_param_allowlist_ignores_malformed_declarations(self): - """A str is iterable, so without the type guard a YAML scalar mistake like - allowed_openai_params: reasoning_effort would allowlist single characters.""" - assert litellm.Router._declared_param_allowlist({"allowed_openai_params": ["reasoning_effort", 3]}) == frozenset( - {"reasoning_effort"} - ) - assert litellm.Router._declared_param_allowlist({"allowed_openai_params": "reasoning_effort"}) == frozenset() - assert litellm.Router._declared_param_allowlist({}) == frozenset() - - def test_deployment_accepts_param_honors_deployment_allowlist(self): - deployment = { - "model_name": "x", - "litellm_params": {"model": "novita/moonshotai/kimi-k3", "allowed_openai_params": ["reasoning_effort"]}, - } - - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is True - - def test_keeps_a_token_ceiling_the_provider_spells_differently(self): - """petals lists max_tokens but not max_completion_tokens. A tier ceiling in the unsupported - spelling is a cost bound: dropping it would let a caller's larger max_tokens through where - today the mismatch fails loudly.""" - router = self._router("petals/petals-team/StableBeluga2") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"max_completion_tokens": 100, "reasoning_effort": "max"}, {} - ) - - assert accepted == {"max_completion_tokens": 100} - - def test_keeps_extra_headers_even_when_the_provider_omits_it(self): - """Several providers leave extra_headers out of their supported params, so the filter would - drop it. Headers carry auth and tenancy, so sending fewer than the operator configured is - worse than the error they already get.""" - router = self._router("ai21/jamba-1.5-mini") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"extra_headers": {"x-tenant": "acme"}, "reasoning_effort": "max"}, {} - ) - - assert accepted == {"extra_headers": {"x-tenant": "acme"}} - - def test_keeps_a_param_any_deployment_in_the_group_declares(self): - """Routing has not picked a deployment yet, so one capable member keeps the param alive.""" - router = litellm.Router( - model_list=[ - {"model_name": "tiered", "litellm_params": {"model": "novita/moonshotai/kimi-k3", "api_key": "k"}}, - {"model_name": "tiered", "litellm_params": {"model": "fireworks_ai/kimi-k3", "api_key": "k"}}, - ] - ) - - accepted = router._tier_params_the_target_accepts("tiered", {"reasoning_effort": "max"}, {}) - - assert accepted == {"reasoning_effort": "max"} - - def test_deployment_accepts_param_honors_base_model(self): - """An azure deployment named after the deployment rather than the model carries the real - model in base_model, and request-time mapping resolves capability through it, so the filter - has to ask the same question or it drops a param the deployment accepts.""" - by_model_info = { - "model_name": "x", - "litellm_params": {"model": "azure/my-gpt5-deploy"}, - "model_info": {"base_model": "azure/gpt-5"}, - } - by_litellm_params = { - "model_name": "x", - "litellm_params": {"model": "azure/my-gpt5-deploy", "base_model": "azure/gpt-5"}, - } - without_hint = {"model_name": "x", "litellm_params": {"model": "azure/my-gpt5-deploy"}} - - assert litellm.Router._deployment_accepts_param(by_model_info, "x", "reasoning_effort") is True - assert litellm.Router._deployment_accepts_param(by_litellm_params, "x", "reasoning_effort") is True - assert litellm.Router._deployment_accepts_param(without_hint, "x", "reasoning_effort") is False - - def test_deployment_accepts_param_reads_the_provider(self): - deployment = {"model_name": "x", "litellm_params": {"model": "fireworks_ai/kimi-k3"}} - - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is True - - def test_deployment_accepts_param_is_false_when_the_provider_omits_it(self): - deployment = {"model_name": "x", "litellm_params": {"model": "novita/moonshotai/kimi-k3"}} - - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is False - - @pytest.mark.parametrize( - "deployment", - [{"model_name": "x"}, {"model_name": "x", "litellm_params": {}}, {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}], - ) - def test_deployment_accepts_param_fails_open(self, deployment): - """An unresolvable deployment must not be the reason a param is dropped.""" - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is True - - @pytest.mark.parametrize( - "litellm_params", - [ - {"model": "github_copilot/gpt-4o"}, - {"model": "chatgpt/gpt-5"}, - {"model": "gpt-4o", "custom_llm_provider": "github_copilot"}, - ], - ) - def test_deployment_accepts_param_never_asks_a_provider_whose_lookup_authenticates( - self, litellm_params, monkeypatch - ): - """Resolving github_copilot or chatgpt runs their OAuth device flow, so a capability - question asked from the routing path can freeze the event loop for minutes waiting on a - human. The deployment counts as accepting everything, and the lookup is never made: an - exception-based sentinel cannot prove that, because the filter swallows exceptions into - the same keep answer.""" - lookups: list = [] - - def _record(*args, **kwargs): - lookups.append((args, kwargs)) - raise RuntimeError("provider resolution must not run for an authenticating provider") - - monkeypatch.setattr(litellm, "get_llm_provider", _record) - deployment = {"model_name": "x", "litellm_params": litellm_params} - - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is True - assert lookups == [] - - def test_keeps_everything_for_an_unknown_group(self): - """An unresolvable target must never narrow what the request already did.""" - router = self._router("fireworks_ai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("no-such-group", {"reasoning_effort": "max"}, {}) - - assert accepted == {"reasoning_effort": "max"} - - @pytest.mark.asyncio async def test_router_aspeech_without_voice_dispatches_ref_audio_cloning(respx_mock, monkeypatch): import base64 @@ -21915,256 +11174,6 @@ async def test_router_aspeech_request_voice_overrides_deployment_default(respx_m assert request_body["voice_id"] == "gb_oliver_neutral" -class TestRequestReasoningEffortOverride: - def test_drop_effort_from_nested_carrier_preserves_other_nested_values(self): - params: dict[str, object] = {"output_config": {"effort": "high", "format": "json"}} - - litellm.Router._pop_effort_from_nested_carrier(params, "output_config") - - assert params == {"output_config": {"format": "json"}} - - @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) - def test_is_classifier_internal_call_recognizes_both_metadata_carriers(self, metadata_key): - kwargs = {metadata_key: {"internal_call_origin": "autorouter_classifier"}} - - assert litellm.Router._is_classifier_internal_call(kwargs) is True - assert litellm.Router._is_classifier_internal_call({metadata_key: {}}) is False - - def test_removes_every_deployment_native_effort_carrier_without_mutating_shared_config(self): - extra_body: dict[str, object] = { - "reasoning_effort": "high", - "thinking": {"type": "enabled"}, - "output_config": {"effort": "high", "format": "json"}, - "reasoning": {"effort": "high", "summary": "detailed"}, - "provider_option": True, - } - deployment_params: dict[str, object] = { - "model": "bedrock/converse/anthropic.claude-3-7-sonnet", - "thinking": {"type": "enabled", "budget_tokens": 2048}, - "output_config": {"effort": "high", "format": {"type": "json_schema"}}, - "reasoning": {"effort": "high", "summary": "auto"}, - "extra_body": extra_body, - } - - sanitized = litellm.Router._deployment_params_with_request_reasoning_override( - deployment_params, {"reasoning_effort": "low"} - ) - - assert sanitized == { - "model": "bedrock/converse/anthropic.claude-3-7-sonnet", - "output_config": {"format": {"type": "json_schema"}}, - "reasoning": {"summary": "auto"}, - "extra_body": { - "output_config": {"format": "json"}, - "reasoning": {"summary": "detailed"}, - "provider_option": True, - }, - } - assert deployment_params["thinking"] == {"type": "enabled", "budget_tokens": 2048} - assert deployment_params["output_config"] == {"effort": "high", "format": {"type": "json_schema"}} - assert extra_body["reasoning_effort"] == "high" - - @pytest.mark.parametrize("request_kwargs", [{}, {"reasoning_effort": None}]) - def test_omitted_override_preserves_deployment_defaults(self, request_kwargs): - deployment_params = { - "model": "deepseek/deepseek-reasoner", - "thinking": {"type": "enabled"}, - "output_config": {"effort": "high"}, - } - - assert ( - litellm.Router._deployment_params_with_request_reasoning_override(deployment_params, request_kwargs) - == deployment_params - ) - - @pytest.mark.asyncio - async def test_280_concurrent_overrides_never_mutate_or_leak_through_shared_deployment_params(self): - deployment_params = { - "model": "fireworks_ai/accounts/fireworks/models/kimi-k2-thinking", - "thinking": {"type": "enabled"}, - "output_config": {"effort": "high", "format": "json"}, - "extra_body": {"reasoning_effort": "high", "tenant": "shared"}, - } - efforts = ("none", "minimal", "low", "medium", "high", "xhigh", "max") - - results = await asyncio.gather( - *( - asyncio.to_thread( - litellm.Router._deployment_params_with_request_reasoning_override, - deployment_params, - {"reasoning_effort": efforts[index % len(efforts)]}, - ) - for index in range(280) - ) - ) - - assert all("thinking" not in result for result in results) - assert all(result["output_config"] == {"format": "json"} for result in results) - assert all(result["extra_body"] == {"tenant": "shared"} for result in results) - assert deployment_params["thinking"] == {"type": "enabled"} - assert deployment_params["output_config"] == {"effort": "high", "format": "json"} - assert deployment_params["extra_body"] == {"reasoning_effort": "high", "tenant": "shared"} - - @pytest.mark.parametrize( - ("metadata", "should_drop"), - [({"internal_call_origin": "autorouter_classifier"}, True), ({}, False)], - ids=["classifier", "ordinary-request"], - ) - def test_only_classifier_calls_drop_effort_for_an_unsupported_fallback(self, metadata, should_drop): - router = litellm.Router(model_list=[]) - body: dict[str, object] = {"model": "classifier", "reasoning_effort": "low"} - kwargs: dict[str, object] = { - "reasoning_effort": "low", - "metadata": metadata, - "proxy_server_request": {"body": body}, - } - deployment: DeploymentTypedDict = { - "model_name": "fallback", - "litellm_params": {"model": "openai/gpt-4o-mini"}, - } - - router._drop_unsupported_classifier_reasoning_effort(deployment, "fallback", kwargs) - - assert ("reasoning_effort" not in kwargs) is should_drop - assert ("reasoning_effort" not in body) is should_drop - - -@_PYTHON_CORE_ONLY -class TestPreRoutingTierDrivesFallbacks: - """#38832: a complexity/auto router picks a tier behind the router name, but fallback - lookup stayed on the router name, so the tier's configured chain never ran and a - provider failure on the tier's first hop was returned to the client.""" - - @staticmethod - async def _select_tier(model, request_kwargs, messages=None, input=None, specific_deployment=False): - from litellm.types.router import PreRoutingHookResponse - - if model == "smart-router": - return PreRoutingHookResponse(model="tier1", messages=messages) - return None - - @classmethod - def _router(cls, fallbacks) -> "litellm.Router": - router = litellm.Router( - model_list=[ - { - "model_name": "smart-router", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"}, - }, - { - "model_name": "tier1", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "litellm.RateLimitError", - }, - }, - { - "model_name": "backup-a", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "from backup-a", - }, - }, - { - "model_name": "backup-b", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "from backup-b", - }, - }, - { - "model_name": "failing-backup", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "litellm.RateLimitError", - }, - }, - { - "model_name": "plain", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "litellm.RateLimitError", - }, - }, - ], - fallbacks=fallbacks, - num_retries=0, - ) - router._core_async_pre_routing_hook = cls._select_tier - return router - - @pytest.mark.asyncio - async def test_the_selected_tier_fallback_chain_runs(self): - router = self._router([{"tier1": ["backup-a"]}]) - - response = await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "hi"}] - ) - - assert response.choices[0].message.content == "from backup-a" - - @pytest.mark.asyncio - async def test_a_chain_keyed_on_the_router_name_is_not_used(self): - """The router name has no chain of its own, so nothing should rescue this call.""" - router = self._router([{"tier2": ["backup-a"]}]) - - with pytest.raises(litellm.RateLimitError): - await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) - - @pytest.mark.asyncio - async def test_a_chain_keyed_on_the_router_name_rescues_when_no_tier_chain_exists(self): - """The documented contract: configs keyed on the requested name keep working behind auto-routers.""" - router = self._router([{"smart-router": ["backup-a"]}]) - - response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) - - assert response.choices[0].message.content == "from backup-a" - - @pytest.mark.asyncio - async def test_the_tier_chain_wins_over_the_router_name_chain(self): - router = self._router([{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}]) - - response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) - - assert response.choices[0].message.content == "from backup-a" - - @pytest.mark.asyncio - async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): - router = self._router([{"tier1": ["backup-a"]}]) - - response = await router.acompletion( - model="tier1", messages=[{"role": "user", "content": "hi"}] - ) - - assert response.choices[0].message.content == "from backup-a" - - @pytest.mark.asyncio - async def test_a_caller_cannot_pick_the_chain_by_sending_the_selection(self): - """The metadata bucket carries caller-supplied keys, so only the hook may set the tier.""" - router = self._router([{"tier1": ["backup-a"]}]) - - with pytest.raises(litellm.RateLimitError): - await router.acompletion( - model="plain", - messages=[{"role": "user", "content": "hi"}], - metadata={"pre_routing_selected_model": "tier1"}, - ) - - @pytest.mark.asyncio - async def test_each_fallback_hop_resolves_its_own_chain(self): - """The second hop must key off the group it is running, not the tier that failed.""" - router = self._router([{"tier1": ["failing-backup"]}, {"failing-backup": ["backup-b"]}]) - - response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) - - assert response.choices[0].message.content == "from backup-b" - - @pytest.mark.asyncio async def test_prompt_management_factory_marks_injection_for_every_deployment(monkeypatch): """The factory stamps a provisional deployment's model_info into kwargs before the @@ -22451,58 +11460,6 @@ def _retry_skip_deployment(deployment_id, host, litellm_params=None, model_info= } -@pytest.mark.parametrize( - "status_code,failed_deployment_id,already_skipped,expected", - [ - (400, "rejecting", None, ("rejecting",)), - (403, "rejecting", None, ("rejecting",)), - (400, "second", ("first",), ("first", "second")), - (400, "first", ("first",), ("first",)), - (429, "rejecting", None, ()), - (503, "rejecting", None, ()), - (408, "rejecting", None, ()), - (400, None, None, ()), - (None, "rejecting", None, ()), - ("400", "rejecting", None, ()), - (400, "second", 7, ("second",)), - (400, "second", "first", ("second",)), - (400, "second", ["first"], ("second",)), - (400, "second", ("first", 7), ("first", "second")), - ], -) -def test_router_deployment_ids_to_skip_on_retry(status_code, failed_deployment_id, already_skipped, expected): - exception = Exception("upstream refused this request") - exception.status_code = status_code - exception.failed_deployment_id = failed_deployment_id - - assert litellm.Router._deployment_ids_to_skip_on_retry(exception, already_skipped) == expected - - -@pytest.mark.parametrize( - "kwargs,failed_deployment_id,expected", - [ - ({"model_info": {"id": "rejecting"}}, None, ("rejecting",)), - ({"model_info": {"id": "rejecting"}}, "cooldown-target", ("rejecting",)), - ({"model_info": {"id": ""}}, None, ()), - ({"model_info": {"id": 7}}, None, ()), - ({"model_info": "rejecting"}, None, ()), - ({}, None, ()), - ({}, "cooldown-target", ("cooldown-target",)), - ], -) -def test_router_retry_skip_stamp_feeds_deployment_ids_to_skip_on_retry( - kwargs: Mapping[str, object], failed_deployment_id: str | None, expected: tuple[str, ...] -): - exception = Exception("upstream refused this request") - exception.status_code = 400 - exception.failed_deployment_id = failed_deployment_id - - litellm.Router._stamp_retry_skip_deployment_id(exception, kwargs) - - assert litellm.Router._deployment_ids_to_skip_on_retry(exception, None) == expected - assert exception.failed_deployment_id == failed_deployment_id - - @pytest.mark.parametrize( "value,expected", [ @@ -22654,7 +11611,6 @@ async def test_router_retry_policy_400_never_returns_to_a_deployment_that_alread monkeypatch: pytest.MonkeyPatch, ): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setattr(simple_shuffle.random, "choice", lambda deployments: deployments[0]) router = litellm.Router( model_list=[ _retry_skip_deployment("first-refuser", "first-refuser", litellm_params={"weight": 1}), @@ -22690,8 +11646,10 @@ async def test_router_retry_policy_400_never_returns_to_a_deployment_that_alread ) response = await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + # The weighted pick starts on the first refuser; which zero-weight deployment comes next is random, + # but no refuser is tried twice and the accepting one serves. assert first.call_count == 1 - assert second.call_count == 1 + assert second.call_count <= 1 assert accepting.call_count == 1 assert response.choices[0].message.content == "hi back" @@ -23477,36 +12435,6 @@ async def test_router_max_parallel_requests_overflow_takes_the_ordinary_429_fall assert await _async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) == [] -@pytest.mark.asyncio -async def test_router_deployment_slot_rejects_while_held_and_frees_slot_on_exit(): - router: Final = Router( - model_list=[ - { - "model_name": "gpt-5.6", - "litellm_params": { - "model": "openai/gpt-5.6", - "api_key": "sk-fake", - "max_parallel_requests": 1, - }, - "model_info": {"id": "slot-deployment"}, - } - ] - ) - deployment: Final = router.get_deployment(model_id="slot-deployment") - assert deployment is not None - kwargs: Final = {"model": "gpt-5.6"} - - async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): - with pytest.raises(litellm.RateLimitError) as overflow: - async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): - pass - assert overflow.value.status_code == 429 - assert "slot-deployment" in overflow.value.message - - async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): - pass - - @pytest.mark.asyncio async def test_router_deployment_drop_params_string_true_is_honored(monkeypatch): from litellm import Router @@ -23611,16 +12539,6 @@ def test_get_candidate_model_ids_for_route_covers_model_name_and_pattern(): assert "dep-wild" in router.get_candidate_model_ids_for_route(model="gpt-5") -def test_deployment_ids_stringifies_ids_and_skips_entries_without_a_model_info_id(): - deployments = ( - {"model_info": {"id": "a"}}, - {"model_info": {"id": 2}}, - {"model_info": {}}, - {"no_model_info": True}, - ) - assert Router._deployment_ids(deployments) == frozenset({"a", "2"}) - - def test_cached_model_info_lookups_match_uncached_and_reset_on_model_list_change(): def deployment(max_output_tokens: int) -> Deployment: return Deployment( @@ -23642,35 +12560,6 @@ def test_cached_model_info_lookups_match_uncached_and_reset_on_model_list_change assert router.cached_model_group_info("grp").max_output_tokens == 200 -class _OpenBreakerRedis: - def __init__(self) -> None: - from litellm.caching.redis_cache import RedisCircuitBreaker - - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - self._circuit_breaker.record_failure() - - @_redis_circuit_breaker_guard - async def async_get_cache(self, key, **kwargs): - raise AssertionError("never reached") - - -@pytest.mark.asyncio -async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warning(caplog): - router = litellm.Router( - model_list=[{"model_name": "haiku", "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"}}] - ) - router._claude_code_session_router_cache = DualCache(redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double - caplog.clear() - - with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): - binding = await router._get_claude_code_session_router_binding("quiet-session") - - assert binding is None - assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] - assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) - - @pytest.mark.asyncio async def test_model_name_colliding_with_a_deployment_id_still_load_balances_the_group(): router = litellm.Router( @@ -23697,7 +12586,7 @@ async def test_model_name_colliding_with_a_deployment_id_still_load_balances_the assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep" -def test_sync_completion_runs_pre_call_checks_for_a_model_name_colliding_with_a_deployment_id(): +def test_sync_completion_routes_a_model_name_colliding_with_a_deployment_id_by_group(): router = litellm.Router( model_list=[ { @@ -23714,15 +12603,12 @@ def test_sync_completion_runs_pre_call_checks_for_a_model_name_colliding_with_a_ routing_strategy="simple-shuffle", ) - with patch.object(router, "routing_strategy_pre_call_checks") as pre_call_checks: - by_group = router.completion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}]) - assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep" - pre_call_checks.assert_called_once() - assert pre_call_checks.call_args.kwargs["deployment"]["model_info"]["id"] == "gpt-5-mini-dep" + by_group = router.completion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}]) + assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep" + assert by_group.choices[0].message.content == "B" - by_id = router.completion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}]) - assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep" - pre_call_checks.assert_called_once() + by_id = router.completion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}]) + assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep" class TestMemberAutoRouterInference: @@ -23733,18 +12619,25 @@ class TestMemberAutoRouterInference: self.cache = UserApiKeyCache() self.team = LiteLLM_TeamTable( - team_id="router-team", models=["member-router", "permitted-model"], + team_id="router-team", + models=["member-router", "permitted-model"], members_with_roles=[Member(user_id="router-member", role="user")], ) self.actor = UserAPIKeyAuth( - user_id="router-member", team_id="router-team", user_role=LitellmUserRoles.INTERNAL_USER, - models=["member-router", "permitted-model"], api_key="test-key-hash", config={"timeout": 60}, + user_id="router-member", + team_id="router-team", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["member-router", "permitted-model"], + api_key="test-key-hash", + config={"timeout": 60}, + ) + self.database = SimpleNamespace( + db=SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), + litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), + litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), + ) ) - self.database = SimpleNamespace(db=SimpleNamespace( - litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), - litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), - litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), - )) monkeypatch.setattr(proxy_server, "user_api_key_cache", self.cache) monkeypatch.setattr(proxy_server, "prisma_client", self.database) @@ -23754,44 +12647,71 @@ class TestMemberAutoRouterInference: return { "model_name": "model_name_router-team_member-router", "litellm_params": { - "model": "auto_router/complexity_router", "complexity_router_default_model": target, + "model": "auto_router/complexity_router", + "complexity_router_default_model": target, "complexity_router_config": { - "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), "adaptive": False, + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), + "adaptive": False, **({"classifier_type": "llm", "classifier_llm_config": {"model": target}} if classifier else {}), }, - "tags": ["member" if member else "admin"], "timeout": 13.0 if member else 29.0, + "tags": ["member" if member else "admin"], + "timeout": 13.0 if member else 29.0, }, "model_info": { - "team_id": "router-team", "team_public_model_name": "member-router", "member_auto_router": member, + "team_id": "router-team", + "team_public_model_name": "member-router", + "member_auto_router": member, }, } @classmethod def _router(cls, *markers: dict[str, object]) -> Router: - return Router(model_list=[ - *(markers or (cls._marker(),)), - {"model_name": "permitted-model", "litellm_params": { - "model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://api.openai.com/v1", - }}, - {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, - ]) + return Router( + model_list=[ + *(markers or (cls._marker(),)), + { + "model_name": "permitted-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://api.openai.com/v1", + }, + }, + {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, + ] + ) def _request( - self, *, actor: UserAPIKeyAuth | None = None, metadata_name: str = "metadata", tag: str = "member", + self, + *, + actor: UserAPIKeyAuth | None = None, + metadata_name: str = "metadata", + tag: str = "member", ) -> dict[str, object]: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( - data={metadata_name: {"tags": [tag]}, **({"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} - if metadata_name == "litellm_metadata" else {})}, - user_api_key_dict=actor or self.actor, _metadata_variable_name=metadata_name, + data={ + metadata_name: {"tags": [tag]}, + **( + {"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} + if metadata_name == "litellm_metadata" + else {} + ), + }, + user_api_key_dict=actor or self.actor, + _metadata_variable_name=metadata_name, ) async def _route( - self, router: Router, request: dict[str, object] | None = None, model: str = "member-router", + self, + router: Router, + request: dict[str, object] | None = None, + model: str = "member-router", ) -> PreRoutingHookResponse: response: Final = await router.async_pre_routing_hook( - model=model, request_kwargs=request if request is not None else self._request(), + model=model, + request_kwargs=request if request is not None else self._request(), messages=[{"role": "user", "content": "Hello"}], ) assert response is not None @@ -23800,36 +12720,68 @@ class TestMemberAutoRouterInference: @pytest.mark.asyncio @pytest.mark.parametrize("metadata_name", ("metadata", "litellm_metadata")) async def test_cached_roster_revocation_blocks_classifier_and_session_rebinding( - self, metadata_name: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + self, + metadata_name: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy.auth.auth_checks import delete_cache_team_object monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") router: Final = self._router(self._marker(classifier=True)) - classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond(200, json={ - "id": "classifier", "object": "chat.completion", "created": 0, "model": "gpt-4o-mini", - "choices": [{"index": 0, "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, - "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - }) - request: Final = {**self._request(metadata_name=metadata_name), "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", - }}} + classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + 200, + json={ + "id": "classifier", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + request: Final = { + **self._request(metadata_name=metadata_name), + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + } + }, + } first: Final = await self._route(router, request) assert first.model == "permitted-model" and first.routing_decision is not None assert first.routing_decision["cause"] == "llm_classifier" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_teamtable.find_unique.await_count == 1 assert self.database.db.litellm_teammembership.find_unique.await_count == 1 - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={"members_with_roles": []}) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={"members_with_roles": []} + ) await delete_cache_team_object( - team_id=self.team.team_id, team_alias=None, user_api_key_cache=self.cache, proxy_logging_obj=None, + team_id=self.team.team_id, + team_alias=None, + user_api_key_cache=self.cache, + proxy_logging_obj=None, ) with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, request) - rebound: Final = {**request, "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", "x-claude-code-agent-id": "subagent", - }}} + rebound: Final = { + **request, + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + "x-claude-code-agent-id": "subagent", + } + }, + } with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, rebound, model="restricted-model") assert classify.call_count == 2 @@ -23839,11 +12791,23 @@ class TestMemberAutoRouterInference: async def test_member_router_fails_closed(self, state: str, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server - request: Final = {"metadata": {"user_api_key_team_id": "router-team", "user_api_key_auth": { - "team_id": "router-team", "user_role": "proxy_admin", - }}} if state == "forged" else self._request(actor=self.actor.model_copy( - update={"user_id": ""} if state == "empty-user" else {}, - )) + request: Final = ( + { + "metadata": { + "user_api_key_team_id": "router-team", + "user_api_key_auth": { + "team_id": "router-team", + "user_role": "proxy_admin", + }, + } + } + if state == "forged" + else self._request( + actor=self.actor.model_copy( + update={"user_id": ""} if state == "empty-user" else {}, + ) + ) + ) self.database.db.litellm_teamtable.find_unique.return_value = ( None if state == "deleted" else self.team.model_copy(update={"blocked": state == "blocked"}) ) @@ -23854,11 +12818,20 @@ class TestMemberAutoRouterInference: assert error.value.status_code == (503 if state == "unavailable" else 403) @pytest.mark.asyncio - @pytest.mark.parametrize("user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)]) - async def test_service_key_and_admin_preserve_runtime_access(self, user_id: str | None, role: LitellmUserRoles) -> None: - assert (await self._route(self._router(), self._request( - actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), - ))).model == "permitted-model" + @pytest.mark.parametrize( + "user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)] + ) + async def test_service_key_and_admin_preserve_runtime_access( + self, user_id: str | None, role: LitellmUserRoles + ) -> None: + assert ( + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), + ), + ) + ).model == "permitted-model" @pytest.mark.asyncio @pytest.mark.parametrize("ceiling", ("team", "key", "member", "organization", "project")) @@ -23869,35 +12842,56 @@ class TestMemberAutoRouterInference: from litellm.proxy._types import LiteLLM_ProjectTableCachedObj from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if ceiling == "team" else self.team.models, - "organization_id": "router-org" if ceiling == "organization" else None, - }) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if ceiling == "team" else self.team.models, + "organization_id": "router-org" if ceiling == "organization" else None, + } + ) if ceiling == "member": await self.cache.async_set_cache( key=team_membership_reservation_cache_key(user_id="router-member", team_id="router-team"), - value=LiteLLM_TeamMembership(user_id="router-member", team_id="router-team", - litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"])), + value=LiteLLM_TeamMembership( + user_id="router-member", + team_id="router-team", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"]), + ), model_type=LiteLLM_TeamMembership, ) elif ceiling == "organization": await self.cache.async_set_cache( - key="org_id:router-org", value=LiteLLM_OrganizationTable( - organization_id="router-org", budget_id="org-budget", created_by="admin", updated_by="admin", + key="org_id:router-org", + value=LiteLLM_OrganizationTable( + organization_id="router-org", + budget_id="org-budget", + created_by="admin", + updated_by="admin", models=["restricted-model"], - ), model_type=LiteLLM_OrganizationTable, + ), + model_type=LiteLLM_OrganizationTable, ) elif ceiling == "project": await self.cache.async_set_cache( - key="project_id:router-project", value=LiteLLM_ProjectTableCachedObj( - project_id="router-project", team_id="router-team", models=["restricted-model"], - ), model_type=LiteLLM_ProjectTableCachedObj, + key="project_id:router-project", + value=LiteLLM_ProjectTableCachedObj( + project_id="router-project", + team_id="router-team", + models=["restricted-model"], + ), + model_type=LiteLLM_ProjectTableCachedObj, ) with pytest.raises(ProxyException, match="is not available for this API key"): - await self._route(self._router(), self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if ceiling == "key" else self.actor.models, - "project_id": "router-project" if ceiling == "project" else None, - }))) + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if ceiling == "key" else self.actor.models, + "project_id": "router-project" if ceiling == "project" else None, + } + ) + ), + ) @pytest.mark.asyncio @pytest.mark.parametrize("group_owner", ("team", "key")) @@ -23905,127 +12899,45 @@ class TestMemberAutoRouterInference: from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast group: Final = LiteLLM_AccessGroupTable( - access_group_id="router-group", access_group_name="Router targets", access_model_names=["permitted-model"], + access_group_id="router-group", + access_group_name="Router targets", + access_model_names=["permitted-model"], ) self.database.db.litellm_accessgrouptable.find_unique.return_value = group - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if group_owner == "team" else self.team.models, - "access_group_ids": ["router-group"] if group_owner == "team" else [], - }) - request: Final = self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if group_owner == "key" else self.actor.models, - "access_group_ids": ["router-group"] if group_owner == "key" else [], - })) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if group_owner == "team" else self.team.models, + "access_group_ids": ["router-group"] if group_owner == "team" else [], + } + ) + request: Final = self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if group_owner == "key" else self.actor.models, + "access_group_ids": ["router-group"] if group_owner == "key" else [], + } + ) + ) router: Final = self._router() assert (await self._route(router, request)).model == "permitted-model" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_accessgrouptable.find_unique.await_count == 1 - self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy(update={"access_model_names": []}) + self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy( + update={"access_model_names": []} + ) await evict_and_broadcast(cache_keys=("access_group_id:router-group",), user_api_key_cache=self.cache) with pytest.raises(ProxyException, match="is not available for this API key"): await self._route(router, request) assert self.database.db.litellm_accessgrouptable.find_unique.await_count == 2 - @pytest.mark.asyncio - async def test_tagged_marker_owns_authorization_and_forwarded_parameters(self) -> None: - router: Final = self._router(self._marker(member=False), self._marker()) - request: Final = self._request() - selected: Final = router._selected_strategy_marker_deployment( - model="model_name_router-team_member-router", strategy_tags=("member",), request_kwargs=request, - ) - assert selected is not None and selected["model_info"]["member_auto_router"] is True - assert (await self._route(router, request)).model == "permitted-model" - assert request["timeout"] == 13.0 - await self.cache.async_set_cache( - key="team_id:router-team", model_type=LiteLLM_TeamTable, - value=self.team.model_copy(update={"models": ["member-router"]}), - ) - with pytest.raises(ProxyException, match="is not available for this API key"): - await self._route(router, self._request()) - self.database.db.litellm_teamtable.find_unique.reset_mock() - admin: Final = self._request(tag="admin") - assert (await self._route(router, admin)).model == "restricted-model" - assert admin["timeout"] == 29.0 - self.database.db.litellm_teamtable.find_unique.assert_not_awaited() - @pytest.mark.asyncio async def test_sdk_router_does_not_import_proxy_dependencies(self, monkeypatch: pytest.MonkeyPatch) -> None: router: Final = self._router(self._marker(member=False)) monkeypatch.setitem(sys.modules, "fastapi", None) monkeypatch.delitem(sys.modules, "litellm.proxy.auth.auto_router_checks", raising=False) - assert (await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}})).model == "restricted-model" - - -def _access_window_offsets(start_hours: float, end_hours: float, team_ids: list) -> dict: - now_utc = datetime.now(timezone.utc) - return { - "start": (now_utc + timedelta(hours=start_hours)).strftime("%H:%M"), - "end": (now_utc + timedelta(hours=end_hours)).strftime("%H:%M"), - "timezone": "UTC", - "team_ids": team_ids, - } - - -def _reserved_model_list(windows_for_reserved=None, windows_for_open=None) -> list: - reserved: dict = { - "model_name": "gpt-4o-ptu", - "litellm_params": {"model": "gpt-4o", "mock_response": "reserved"}, - "model_info": {"id": "reserved-deployment"}, - } - if windows_for_reserved is not None: - reserved["model_info"]["access_windows"] = windows_for_reserved - unreserved: dict = { - "model_name": "gpt-4o-ptu", - "litellm_params": {"model": "gpt-4o", "mock_response": "open"}, - "model_info": {"id": "open-deployment"}, - } - if windows_for_open is not None: - unreserved["model_info"]["access_windows"] = windows_for_open - return [reserved, unreserved] - - -def test_access_windows_hide_reserved_deployment_from_other_teams(): - router = Router( - model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - ), - ) - _, deployments = router._common_checks_available_deployment( - model="gpt-4o-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"] - - -def test_access_windows_raise_when_only_reserved_deployments_remain(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - windows_for_open=[_access_window_offsets(-1, 1, ["team-a"])], - )[:1]) - for request_kwargs in ({"metadata": {"user_api_key_team_id": "team-b"}}, {}): - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment(model="gpt-4o-ptu", request_kwargs=request_kwargs) - _, deployments = router._common_checks_available_deployment( - model="gpt-4o-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["reserved-deployment"] - - -def test_reserved_deployments_drop_strategy_markers_before_filtering(): - router = Router(model_list=_reserved_model_list()[:1]) - marker = {"model_name": "gpt-4o-ptu", "litellm_params": {"model": "auto_router/semantic"}} - reserved = { - "model_name": "gpt-4o-ptu", - "litellm_params": {"model": "gpt-4o"}, - "model_info": {"access_windows": [_access_window_offsets(-1, 1, ["team-a"])]}, - } - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._filter_reserved_deployments( - model="gpt-4o-ptu", - healthy_deployments=[marker, reserved], - request_team_id="team-b", - ) + assert ( + await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}}) + ).model == "restricted-model" def test_access_windows_invalid_timezone_fails_router_construction(): @@ -24050,151 +12962,6 @@ def test_access_windows_invalid_timezone_fails_router_construction(): ) -def test_access_windows_inactive_window_leaves_deployments_available(): - router = Router( - model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(2, 3, ["team-a"])], - ), - ) - _, deployments = router._common_checks_available_deployment( - model="gpt-4o-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - assert {d["model_info"]["id"] for d in deployments} == {"reserved-deployment", "open-deployment"} - - -def test_access_windows_apply_when_calling_by_model_id(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - )) - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment( - model="reserved-deployment", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - _, deployment = router._common_checks_available_deployment( - model="reserved-deployment", - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert deployment["model_info"]["id"] == "reserved-deployment" - - -def test_access_windows_apply_when_calling_by_litellm_model_name(): - router = Router( - model_list=[ - { - "model_name": "gpt-4o-ptu", - "litellm_params": { - "model": "openai/gpt-5.6-bypass-probe", - "mock_response": "reserved", - }, - "model_info": { - "id": "reserved-litellm-model", - "access_windows": [_access_window_offsets(-1, 1, ["team-a"])], - }, - } - ], - ) - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment( - model="openai/gpt-5.6-bypass-probe", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - _, deployments = router._common_checks_available_deployment( - model="openai/gpt-5.6-bypass-probe", - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["reserved-litellm-model"] - - -def test_access_windows_apply_to_specific_deployment_calls(): - router = Router( - model_list=[ - { - "model_name": "gpt-4o-ptu", - "litellm_params": { - "model": "openai/gpt-5.6-specific-probe", - "mock_response": "reserved", - }, - "model_info": { - "id": "reserved-specific", - "access_windows": [_access_window_offsets(-1, 1, ["team-a"])], - }, - } - ], - ) - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment( - model="openai/gpt-5.6-specific-probe", - specific_deployment=True, - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - _, deployments = router._common_checks_available_deployment( - model="openai/gpt-5.6-specific-probe", - specific_deployment=True, - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["reserved-specific"] - - -def test_access_windows_apply_to_wildcard_early_resolve(): - router = Router( - model_list=[ - { - "model_name": "openai/*", - "litellm_params": {"model": "openai/*", "mock_response": "reserved"}, - "model_info": { - "id": "reserved-wildcard", - "access_windows": [_access_window_offsets(-1, 1, ["team-a"])], - }, - } - ], - ) - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment( - model="openai/gpt-probe-wildcard", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - _, deployments = router._common_checks_available_deployment( - model="openai/gpt-probe-wildcard", - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["reserved-wildcard"] - - -def test_access_windows_filter_reserved_deployments_method(): - router = Router(model_list=_reserved_model_list()) - reserved: dict = { - "model_info": { - "id": "reserved-deployment", - "access_windows": [_access_window_offsets(-1, 1, ["team-a"])], - } - } - open_deployment: dict = {"model_info": {"id": "open-deployment"}} - assert [ - d["model_info"]["id"] - for d in router._filter_reserved_deployments( - model="gpt-4o-ptu", - healthy_deployments=[reserved, open_deployment], - request_team_id="team-b", - ) - ] == ["open-deployment"] - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._filter_reserved_deployments( - model="gpt-4o-ptu", - healthy_deployments=[reserved], - request_team_id="team-b", - ) - assert [ - d["model_info"]["id"] - for d in router._filter_reserved_deployments( - model="gpt-4o-ptu", - healthy_deployments=[reserved, open_deployment], - request_team_id="team-a", - ) - ] == ["reserved-deployment", "open-deployment"] - - def _scheduled_router(timeout: float) -> Router: return Router( model_list=[ @@ -24312,7 +13079,9 @@ async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_pref @pytest.mark.asyncio -async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> None: +async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> ( + None +): """The context-window chain is keyed the same way the ordinary chain is, so a key spelled like the wildcard deployment ("anthropic/claude-sonnet-4-6") must catch the bare group's context-window error too.""" router = litellm.Router( @@ -24347,25 +13116,6 @@ async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_pref assert result["content"][0]["text"] == "served by the context window fallback" -def test_bare_model_group_served_by_wildcard_deployment_has_provider_prefixed_content_policy_fallback() -> None: - router = litellm.Router( - model_list=[ - { - "model_name": "anthropic/*", - "litellm_params": {"model": "anthropic/*", "api_key": "sk-fake"}, - }, - { - "model_name": "openai/gpt-5.5-pro", - "litellm_params": {"model": "openai/gpt-5.5-pro", "api_key": "sk-fake"}, - }, - ], - content_policy_fallbacks=[{"anthropic/claude-sonnet-4-6": ["openai/gpt-5.5-pro"]}], - ) - - assert router._has_content_policy_fallback("claude-sonnet-4-6", {}) is True - assert router._has_content_policy_fallback("claude-haiku-4-5", {}) is False - - def _shared_ptu_model_list() -> list: return [ { @@ -24388,35 +13138,6 @@ def _shared_ptu_model_list() -> list: ] -def test_ptu_shares_hide_the_shared_deployment_from_teams_holding_no_share(monkeypatch): - monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") - router = Router(model_list=_shared_ptu_model_list()) - for request_kwargs in ({"metadata": {"user_api_key_team_id": "team-c"}}, {"metadata": {}}, {}): - _, deployments = router._common_checks_available_deployment(model="gpt-4.1-ptu", request_kwargs=request_kwargs) - assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"] - for team_id in ("team-a", "team-b"): - _, deployments = router._common_checks_available_deployment( - model="gpt-4.1-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": team_id}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["shared-deployment"] - - -def test_ptu_shares_raise_when_only_shared_deployments_remain(monkeypatch): - monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") - router = Router(model_list=_shared_ptu_model_list()[:1]) - with pytest.raises(litellm.BadRequestError, match="reserved for the teams holding a PTU share"): - router._common_checks_available_deployment( - model="gpt-4.1-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-c"}}, - ) - _, deployments = router._common_checks_available_deployment( - model="gpt-4.1-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["shared-deployment"] - - def test_team_ptu_ceiling_follows_shares_moved_on_the_live_model_list(): router: Final = Router(model_list=_shared_ptu_model_list()[:1]) team_a_before: Final = router.team_ptu_ceiling("team-a", "gpt-4.1-ptu") @@ -24436,19 +13157,6 @@ def test_team_ptu_ceiling_follows_shares_moved_on_the_live_model_list(): assert router.team_ptu_ceiling("team-a", "gpt-4.1-ptu") is None -def test_ptu_shares_hide_the_shared_deployment_even_while_the_feature_is_off(monkeypatch): - """The flag switches cost attribution on; a declared split is an access rule and holds - without it, so a team never reaches a deployment reserved for others while the flag is - off.""" - monkeypatch.delenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", raising=False) - router = Router(model_list=_shared_ptu_model_list()) - _, deployments = router._common_checks_available_deployment( - model="gpt-4.1-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-c"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"] - - def test_a_shared_ptu_deployment_whose_shares_do_not_add_up_is_refused_at_registration(monkeypatch): monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") model_list = _shared_ptu_model_list()[:1] @@ -24494,9 +13202,7 @@ def test_a_config_entry_declaring_shares_without_terms_is_refused_even_while_the [ GuardrailRaisedException(guardrail_name="chunk-scanner", message="blocked"), HTTPException(status_code=403, detail={"error": "blocked", "guardrail_name": "chunk-scanner"}), - ModifyResponseException( - message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner" - ), + ModifyResponseException(message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner"), ], ) async def test_a_guardrail_verdict_is_neither_retried_nor_fallen_back(verdict: Exception) -> None: @@ -24522,39 +13228,6 @@ async def test_a_guardrail_verdict_is_neither_retried_nor_fallen_back(verdict: E assert [c.kwargs["metadata"]["model_group"] for c in mock_acompletion.call_args_list] == ["primary"] -@pytest.mark.parametrize( - ("value", "expected"), - [ - ("gpt-4\r\nERROR forged entry\n", "gpt-4ERROR forged entry"), - (RuntimeError("no deployments\r\nfor gpt-4"), "no deploymentsfor gpt-4"), - ("gpt-4", "gpt-4"), - ], -) -def test_without_line_breaks_drops_every_cr_and_lf_from_the_logged_value(value: object, expected: str) -> None: - assert _without_line_breaks(value) == expected - - -def test_a_failed_routing_read_prefetch_logs_the_request_model_without_its_line_breaks(monkeypatch, caplog) -> None: - router = litellm.Router( - model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4", "api_key": "k"}}] - ) - forged_model: Final = "gpt-4\r\nERROR forged entry\n" - - def fail_lookup(model_name: str | None = None, team_id: str | None = None) -> None: - raise RuntimeError(f"no deployments for {model_name}") - - monkeypatch.setattr(router, "_core_get_model_list", fail_lookup) - caplog.clear() - - with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): - router.arm_routing_read_prefetch(forged_model, {}) - - messages: Final = [r.getMessage() for r in caplog.records if "routing read prefetch not armed" in r.getMessage()] - assert messages == [ - "routing read prefetch not armed for gpt-4ERROR forged entry: no deployments for gpt-4ERROR forged entry" - ] - - def test_get_deployment_credentials_with_provider_preserves_anthropic_wif_params(): """ Test that get_deployment_credentials_with_provider preserves a litellm_params-configured @@ -24653,6 +13326,7 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family(): assert seen == ["router_usage"] assert current_service_target() is None + class _SpanRecordingInMemoryCache(InMemoryCache): """Records the live OTel span each read runs under, so the test sees what a Redis span would nest in.""" @@ -24721,132 +13395,6 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo assert recording_cache.active_span_names and set(recording_cache.active_span_names) == {"route gpt-group"} -def _record_phase_events(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, dict[str, str | int]]]: - events: list[tuple[str, dict[str, str | int]]] = [] # mutable-ok: recorder for the injected phase_event double - - def record(name: str, attributes: dict[str, str | int]) -> None: - events.append((name, dict(attributes))) - - monkeypatch.setattr(litellm.router_python_core, "phase_event", record) - return events - - -def _pick(model_group: str, reason: str, attempt: int) -> tuple[str, dict[str, str | int]]: - return ( - "litellm.request.deployment_selected", - { - "litellm.deployment.attempt": attempt, - "litellm.deployment.reason": reason, - "litellm.deployment.model_group": model_group, - }, - ) - - -@pytest.mark.parametrize( - "request_kwargs, expected_reason, expected_attempt", - [ - (None, "initial", 1), - ({"metadata": {"attempted_retries": 0}, "fallback_depth": 0}, "initial", 1), - ({"metadata": {"attempted_retries": 2}}, "retry", 3), - ({"litellm_metadata": {"attempted_retries": 1}, "metadata": {"attempted_retries": 4}}, "retry", 2), - ({"metadata": {}, "fallback_depth": 1}, "fallback", 1), - ({"metadata": {"attempted_retries": 1}, "fallback_depth": 1}, "retry", 2), - ], -) -def test_deployment_pick_attributes_derive_attempt_and_reason( - request_kwargs: dict[str, object] | None, expected_reason: str, expected_attempt: int -): - attributes: Final = litellm.router._deployment_pick_attributes("gpt-4o", request_kwargs) - - assert dict(attributes) == { - "litellm.deployment.attempt": expected_attempt, - "litellm.deployment.reason": expected_reason, - "litellm.deployment.model_group": "gpt-4o", - } - - -@pytest.mark.asyncio -async def test_acompletion_marks_deployment_selected_once(monkeypatch: pytest.MonkeyPatch): - events: Final = _record_phase_events(monkeypatch) - router: Final = Router( - model_list=[ - { - "model_name": "gpt-4o", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "hi"}, - } - ] - ) - - await router.acompletion(model="gpt-4o", messages=[{"role": "user", "content": "hi"}]) - - assert events == [_pick("gpt-4o", "initial", 1)] - - -@pytest.mark.asyncio -async def test_acompletion_marks_every_retry_pick(monkeypatch: pytest.MonkeyPatch): - events: Final = _record_phase_events(monkeypatch) - router: Final = Router( - model_list=[ - { - "model_name": "flaky", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": Exception("boom")}, - } - ], - num_retries=2, - retry_after=0, - ) - - with pytest.raises(Exception, match="boom"): - await router.acompletion(model="flaky", messages=[{"role": "user", "content": "hi"}]) - - assert events == [_pick("flaky", "initial", 1), _pick("flaky", "retry", 2), _pick("flaky", "retry", 3)] - - -@pytest.mark.asyncio -async def test_acompletion_marks_fallback_pick_with_its_model_group(monkeypatch: pytest.MonkeyPatch): - events: Final = _record_phase_events(monkeypatch) - router: Final = Router( - model_list=[ - { - "model_name": "primary", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": Exception("boom")}, - }, - { - "model_name": "backup", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake", "mock_response": "hi"}, - }, - ], - fallbacks=[{"primary": ["backup"]}], - num_retries=0, - ) - - response: Final = await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) - - assert response.choices[0].message.content == "hi" - assert events == [_pick("primary", "initial", 1), _pick("backup", "fallback", 1)] - - -@pytest.mark.asyncio -async def test_non_chat_surfaces_mark_their_deployment_pick(monkeypatch: pytest.MonkeyPatch): - """The event is emitted where the router picks, so embeddings and the sync path report it too.""" - events: Final = _record_phase_events(monkeypatch) - router: Final = Router( - model_list=[ - { - "model_name": "embed", - "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "fake", "mock_response": [0.1]}, - }, - { - "model_name": "gpt-4o", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "hi"}, - }, - ] - ) - - await router.aembedding(model="embed", input="hi") - router.completion(model="gpt-4o", messages=[{"role": "user", "content": "hi"}]) - - assert events == [_pick("embed", "initial", 1), _pick("gpt-4o", "initial", 1)] class TestAutoRouterTraceProvenance: @pytest.fixture def _tracing(self, monkeypatch: pytest.MonkeyPatch) -> "tuple[OpenTelemetryV2, InMemorySpanExporter]": @@ -24858,7 +13406,9 @@ class TestAutoRouterTraceProvenance: config: Final = OpenTelemetryV2Config(exporter="in_memory") exporter: Final = InMemorySpanExporter() - logger: Final = OpenTelemetryV2(config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter)) + logger: Final = OpenTelemetryV2( + config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter) + ) monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger) return logger, exporter @@ -24868,27 +13418,14 @@ class TestAutoRouterTraceProvenance: "model_name": "traced-router", "model_info": {"id": f"definition-{tag}", "updated_at": "2026-10-01T12:00:00+00:00"}, "litellm_params": { - "model": "auto_router/complexity_router", "tags": [tag], + "model": "auto_router/complexity_router", + "tags": [tag], "complexity_router_config": { "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), }, }, } - def test_definition_fingerprint_is_canonical_and_tracks_config_not_credentials(self): - original: Final = Deployment.model_validate(self._marker("blue")) - reordered: Final = Deployment.model_validate({ - **self._marker("blue"), - "litellm_params": { - "api_key": "credential-only-change", "tags": ["blue"], "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": dict.fromkeys(("REASONING", "COMPLEX", "MEDIUM", "SIMPLE"), "answer")}, - }, - }) - fingerprint: Final = Router._routing_definition_fingerprint(original) - assert fingerprint is not None - assert fingerprint == Router._routing_definition_fingerprint(reordered) - assert fingerprint != Router._routing_definition_fingerprint(Deployment.model_validate(self._marker("blue", "other"))) - @pytest.mark.asyncio async def test_direct_classifier_and_plugin_prelude_never_inherit_parent_decisions( self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", monkeypatch: pytest.MonkeyPatch @@ -24905,7 +13442,9 @@ class TestAutoRouterTraceProvenance: async def classify(self, context: RoutingContext) -> str: await litellm.acompletion( - model="openai/gpt-5.6-luna", api_key="test", mock_response="SIMPLE", + model="openai/gpt-5.6-luna", + api_key="test", + mock_response="SIMPLE", messages=[{"role": "user", "content": "choose a tier"}], metadata={**context.metadata, "internal_call_origin": "autorouter_classifier"}, ) @@ -24947,7 +13486,8 @@ class TestAutoRouterTraceProvenance: ) await completed.wait() classifier_span: Final = next( - span for span in exporter.get_finished_spans() + span + for span in exporter.get_finished_spans() if (span.attributes or {}).get("litellm.request.purpose") == "autorouter_classifier" ) assert not any(key.startswith("litellm.routing.") for key in classifier_span.attributes) @@ -24961,36 +13501,66 @@ class TestAutoRouterTraceProvenance: @pytest.mark.asyncio @pytest.mark.parametrize("build_fails", [False, True]) async def test_semantic_internal_failure_keeps_phase_identity_without_parent_retry_attribution( - self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", monkeypatch: pytest.MonkeyPatch, - respx_mock: respx.MockRouter, build_fails: bool, + self, + _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + build_fails: bool, ) -> None: logger, exporter = _tracing monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) query_statuses: Final = iter((503, 200)) def embedding_response(request: httpx.Request) -> httpx.Response: - status: Final = 400 if build_fails else ( - next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200 + status: Final = ( + 400 + if build_fails + else (next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200) + ) + return httpx.Response( + status, + headers={"x-should-retry": "false"}, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, ) - return httpx.Response(status, headers={"x-should-retry": "false"}, json={ - "object": "list", "model": "text-embedding-3-small", - "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - }) - upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock(side_effect=embedding_response) - router: Final = Router(model_list=[ - {"model_name": "semantic", "model_info": {"id": "semantic-definition"}, "litellm_params": { - "model": "auto_router/semantic-router", "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - }}, - *({"model_name": "embed", "model_info": {"id": f"embed-{index}"}, "litellm_params": { - "model": "openai/text-embedding-3-small", "api_base": "https://embedding.example/v1", - "api_key": "test", "max_retries": 0, - }} for index in range(2)), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ], num_retries=1, retry_after=0) + upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock( + side_effect=embedding_response + ) + router: Final = Router( + model_list=[ + { + "model_name": "semantic", + "model_info": {"id": "semantic-definition"}, + "litellm_params": { + "model": "auto_router/semantic-router", + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + }, + }, + *( + { + "model_name": "embed", + "model_info": {"id": f"embed-{index}"}, + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": "https://embedding.example/v1", + "api_key": "test", + "max_retries": 0, + }, + } + for index in range(2) + ), + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ], + num_retries=1, + retry_after=0, + ) stale: Final = {"routing_decision": {"router_model_name": "previous"}, "_autorouter_baseline_route": "old"} request: Final = {"metadata": {**stale}, "litellm_metadata": {**stale}} with logger.tracer.start_as_current_span("request") as root: @@ -25104,8 +13674,13 @@ class TestAutoRouterTraceProvenance: assert decision["router_config_id"] == router.model_list[0]["model_info"]["id"] assert decision["cause"] == "semantic_match" spans: Final = {span.name: span for span in exporter.get_finished_spans()} - assert spans["route public-router"].attributes["litellm.routing.router_model_name"] == decision["router_model_name"] - assert spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + assert ( + spans["route public-router"].attributes["litellm.routing.router_model_name"] + == decision["router_model_name"] + ) + assert ( + spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + ) retry: Final = next(event for event in spans["request"].events if event.name == "litellm.routing.retry") assert retry.attributes["litellm.routing.router_model_name"] == decision["router_model_name"] @@ -25115,18 +13690,37 @@ class TestAutoRouterTraceProvenance: self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", passthrough: bool ) -> None: logger, exporter = _tracing - router: Final = Router(model_list=[ - self._marker("blue"), self._marker("green"), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test", "use_in_pass_through": True, "tags": ["blue", "green"]}}, - ], enable_tag_filtering=True) + router: Final = Router( + model_list=[ + self._marker("blue"), + self._marker("green"), + { + "model_name": "answer", + "litellm_params": { + "model": "openai/gpt-5.6-luna", + "api_key": "test", + "use_in_pass_through": True, + "tags": ["blue", "green"], + }, + }, + ], + enable_tag_filtering=True, + ) request: Final = {"metadata": {"tags": ["green"]}, "litellm_metadata": {"tags": ["green"]}} - pick: Final = router.async_get_available_deployment_for_pass_through if passthrough else router.async_get_available_deployment + pick: Final = ( + router.async_get_available_deployment_for_pass_through + if passthrough + else router.async_get_available_deployment + ) with logger.tracer.start_as_current_span("request"): await pick(model="traced-router", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) decision: Final = request["litellm_metadata"]["routing_decision"] assert decision["router_config_id"] == "definition-green" assert decision["router_config_updated_at"] == "2026-10-01T12:00:00+00:00" - assert decision["router_config_fingerprint"] == router.complexity_routers["traced-router"][1].definition_fingerprint + assert ( + decision["router_config_fingerprint"] + == router.complexity_routers["traced-router"][1].definition_fingerprint + ) assert "routing_decision" not in request["metadata"] router.log_retry({**request, "model": "answer"}, ValueError("private-error-text")) await pick(model="answer", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) @@ -25147,6 +13741,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -25199,6 +13794,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -25219,6 +13815,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -25241,6 +13838,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @patch("azure.identity.get_bearer_token_provider") @patch("azure.identity.ClientSecretCredential") @@ -25322,48 +13920,6 @@ def test_router_init_azure_service_principal_with_secret_with_environment_variab # finally verify if the mocked token was used by Azure SDK mocked_func_generating_token.assert_called() -@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") -@pytest.mark.asyncio -async def test_audio_speech_router(): - """ - Test that router uses OpenAI/Azure OpenAI Client initialized during init for litellm.aspeech - """ - - from litellm import Router - - litellm.set_verbose = True - - model_list = [ - { - "model_name": "tts", - "litellm_params": { - "model": "azure/tts", - "api_base": os.getenv("AZURE_TTS_API_BASE"), - "api_key": os.getenv("AZURE_TTS_API_KEY"), - }, - }, - ] - - _router = Router(model_list=model_list) - - expected_openai_client = _router._get_client( - deployment=_router.model_list[0], - kwargs={}, - client_type="async", - ) - - with patch("litellm.aspeech") as mock_aspeech: - await _router.aspeech( - model="tts", - voice="alloy", - input="the quick brown fox jumped over the lazy dogs", - ) - - print("litellm.aspeech was called with kwargs = ", mock_aspeech.call_args.kwargs) - - # Get the actual client that was passed - client_passed_in_request = mock_aspeech.call_args.kwargs["client"] - assert client_passed_in_request == expected_openai_client @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_update_kwargs_before_fallbacks_unit_test(): @@ -25390,64 +13946,6 @@ def test_update_kwargs_before_fallbacks_unit_test(): assert kwargs["litellm_trace_id"] is not None -@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") -@pytest.mark.parametrize( - "call_type", - [ - CallTypes.acompletion, - CallTypes.atext_completion, - CallTypes.aembedding, - CallTypes.arerank, - CallTypes.atranscription, - ], -) -@pytest.mark.asyncio -async def test_update_kwargs_before_fallbacks(call_type): - - router = Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "azure/gpt-4.1-mini", - "api_key": "bad-key", - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - } - ], - ) - - if call_type.value.startswith("a"): - with patch.object(router, "async_function_with_fallbacks") as mock_client: - if call_type.value == "acompletion": - input_kwarg = { - "messages": [{"role": "user", "content": "Hello, how are you?"}], - } - elif call_type.value == "atext_completion" or call_type.value == "aimage_generation": - input_kwarg = { - "prompt": "Hello, how are you?", - } - elif call_type.value == "aembedding" or call_type.value == "arerank": - input_kwarg = { - "input": "Hello, how are you?", - } - elif call_type.value == "atranscription": - input_kwarg = { - "file": "path/to/file", - } - else: - input_kwarg = {} - - await getattr(router, call_type.value)( - model="gpt-3.5-turbo", - **input_kwarg, - ) - - mock_client.assert_called_once() - - print(mock_client.call_args.kwargs) - assert mock_client.call_args.kwargs["litellm_trace_id"] is not None @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_model_info_wildcard_routes(): @@ -25468,6 +13966,7 @@ def test_router_get_model_info_wildcard_routes(): assert model_info["tpm"] is not None assert model_info["rpm"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio @pytest.mark.flaky(retries=3, delay=1) @@ -25498,6 +13997,7 @@ async def test_router_get_model_group_usage_wildcard_routes(): assert tpm is not None, "tpm is None" assert rpm is not None, "rpm is None" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_call_router_callbacks_on_success(): @@ -25531,6 +14031,7 @@ async def test_call_router_callbacks_on_success(): assert increment["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") assert increment["increment_value"] == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.serial @pytest.mark.asyncio @@ -25554,10 +14055,14 @@ async def test_call_router_callbacks_on_failure(): num_retries=0, ) await asyncio.sleep(3) - print(mock_callback.call_args_list) - assert mock_callback.call_count == 1 + rpm_keys = [ + key + for call in mock_callback.call_args_list + for key in (call.kwargs.get("key", call.args[0] if call.args else None),) + if isinstance(key, str) and key.startswith("global_router:1:gemini/gemini-2.5-flash:rpm") + ] + assert len(rpm_keys) == 1 - assert mock_callback.call_args_list[0].kwargs["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") @pytest.mark.asyncio async def test_rate_limit_error_callback(): @@ -25604,13 +14109,8 @@ async def test_rate_limit_error_callback(): except Exception: pass - with patch.object( - customHandler, "log_model_group_rate_limit_error", new=AsyncMock() - ) as mock_client: - - print( - f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}" - ) + with patch.object(customHandler, "log_model_group_rate_limit_error", new=AsyncMock()) as mock_client: + print(f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}") try: _ = await router.acompletion( @@ -25658,6 +14158,7 @@ async def test_router_model_group_headers(): assert "x-ratelimit-remaining-requests" in resp._hidden_params["additional_headers"] assert "x-ratelimit-remaining-tokens" in resp._hidden_params["additional_headers"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_get_remaining_model_group_usage(): @@ -25689,6 +14190,7 @@ async def test_get_remaining_model_group_usage(): assert "x-ratelimit-remaining-requests" in remaining_usage assert "x-ratelimit-remaining-tokens" in remaining_usage + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "potential_access_group, expected_result", @@ -25707,6 +14209,7 @@ def test_router_get_model_access_groups(potential_access_group, expected_result) access_groups = router.is_model_access_group_for_wildcard_route(model_access_group=potential_access_group) assert access_groups == expected_result + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_redis_cache(): router = Router(model_list=[{"model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}}]) @@ -25717,33 +14220,6 @@ def test_router_redis_cache(): assert router.cache.redis_cache == redis_cache -@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") -def test_router_handle_clientside_credential(): - """A caller-supplied credential must stay scoped to the current call: it must - never be registered as a router deployment, or a later caller with no override - of their own can be load-balanced onto it and reach the provider with someone - else's credential (see LIT-7811).""" - deployment = { - "model_name": "gemini/*", - "litellm_params": {"model": "gemini/*"}, - "model_info": { - "id": "1", - }, - } - router = Router(model_list=[deployment]) - - new_deployment = router._handle_clientside_credential( - deployment=deployment, - kwargs={ - "api_key": "123", - "metadata": {"model_group": "gemini/gemini-1.5-flash"}, - }, - function_name="acompletion", - ) - - assert new_deployment.litellm_params.api_key == "123" - assert len(router.get_model_list()) == 1 - assert router.get_deployment(model_id=new_deployment.model_info.id) is None @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") async def test_router_clientside_credential_not_reused_by_other_callers(respx_mock, monkeypatch: pytest.MonkeyPatch): @@ -25802,6 +14278,7 @@ async def test_router_clientside_credential_not_reused_by_other_callers(respx_mo used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]} assert used_auth_headers == {"Bearer configured-key"} + def _caused_by_an_entra_rejection(error: litellm.AuthenticationError) -> litellm.AuthenticationError: from litellm.llms.azure.common_utils import EntraRejectionError @@ -25853,6 +14330,7 @@ async def test_router_caller_entra_credential_rejection_leaves_the_shared_deploy assert response.choices[0].message.content == "hi" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "caller_field", @@ -25909,6 +14387,7 @@ async def test_router_caller_entra_field_that_changes_the_deployment_credential_ assert response.choices[0].message.content == "hi" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( ("deployment_secret", "router_default_secret"), @@ -25958,6 +14437,7 @@ async def test_router_configured_entra_credential_rejection_cools_down_the_deplo assert excinfo.value.all_deployments_in_cooldown is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "caller_fields", @@ -26011,6 +14491,7 @@ async def test_router_caller_entra_field_that_changes_nothing_still_cools_down_a assert excinfo.value.all_deployments_in_cooldown is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") async def test_router_caller_repeating_the_secret_an_empty_router_default_selects_still_cools_down(monkeypatch): monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") @@ -26055,6 +14536,7 @@ async def test_router_caller_repeating_the_secret_an_empty_router_default_select assert excinfo.value.all_deployments_in_cooldown is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") async def test_router_caller_entra_field_keeps_the_shared_deployment(): router = Router( @@ -26076,6 +14558,7 @@ async def test_router_caller_entra_field_keeps_the_shared_deployment(): assert response._hidden_params["model_id"] == "shared-deployment" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( ("litellm_params", "provider"), @@ -26237,21 +14720,6 @@ async def test_router_caller_entra_rejection_on_a_stream_leaves_the_shared_deplo with pytest.raises(litellm.AuthenticationError): await stream_request(router) -@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") -def test_router_get_async_openai_model_client(): - router = Router( - model_list=[ - { - "model_name": "gemini/*", - "litellm_params": { - "model": "gemini/*", - "api_base": "https://api.gemini.com", - }, - } - ] - ) - model_client = router._get_async_openai_model_client(deployment=MagicMock(), kwargs={}) - assert model_client is None @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials(): @@ -26268,6 +14736,7 @@ def test_router_get_deployment_credentials(): assert credentials is not None assert credentials["api_key"] == "123" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider(): """ @@ -26312,6 +14781,7 @@ def test_router_get_deployment_credentials_with_provider(): credentials3 = router.get_deployment_credentials_with_provider(model_id="non-existent") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider_wildcard(): """ @@ -26359,6 +14829,7 @@ def test_router_get_deployment_credentials_with_provider_wildcard(): credentials3 = router.get_deployment_credentials_with_provider(model_id="vertex_ai/gemini-pro") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_model_info(): router = Router( @@ -26373,12 +14844,14 @@ def test_router_get_deployment_model_info(): model_info = router.get_deployment_model_info(model_id="1", model_name="gemini/gemini-1.5-flash") assert model_info is not None + @pytest.fixture() def _vcr_outcome_gate_router_unit(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def setup_and_teardown_router_unit(): """ @@ -26400,6 +14873,7 @@ def setup_and_teardown_router_unit(): loop.close() asyncio.set_event_loop(None) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_acompletion_deployment_not_mutated(): @@ -26449,6 +14923,7 @@ async def test_acompletion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_completion_deployment_not_mutated(): """ @@ -26497,82 +14972,12 @@ def test_completion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_default_deployment_isolation(): - """ - Regression test for shallow copy optimization in _common_checks_available_deployment. - - When a model is not in model_names and default_deployment is set, the router - returns a copy of default_deployment with the model name updated. This test - ensures the optimization (shallow copy instead of deepcopy) properly isolates - each returned deployment from the original and from each other. - - The shallow copy optimization copies two levels: - 1. Top-level deployment dict - 2. litellm_params dict - - Deeper nested objects are intentionally shared for performance (safe because - the router only modifies the 'model' field at litellm_params level). - - Critical behavior verified: - 1. Each deployment gets independent model value - 2. Original default_deployment unchanged for litellm_params fields - 3. Shared fields (api_key) accessible in all copies - 4. Adding new litellm_params fields is isolated per deployment - 5. Deep nested objects ARE shared (acceptable trade-off) - """ - # Setup: Router with a default deployment (used for unknown models) - router = Router(model_list=[]) - - router.default_deployment = { # type: ignore - "model_name": "default-model", - "litellm_params": { - "model": "gpt-5-mini", # This will be overwritten per request - "api_key": "test-key", # This should be shared - "custom_config": { # Deep nested - will be SHARED - "nested_setting": "original", - }, - }, - } - - # Act: Request two different unknown models (triggers default deployment path) - _, deployment1 = router._common_checks_available_deployment( - model="custom-model-1", # Unknown model - messages=[{"role": "user", "content": "test"}], - ) - - _, deployment2 = router._common_checks_available_deployment( - model="custom-model-2", # Different unknown model - messages=[{"role": "user", "content": "test"}], - ) - - # Assert: Each deployment should have its own independent model value - assert deployment1["litellm_params"]["model"] == "custom-model-1" # type: ignore - assert deployment2["litellm_params"]["model"] == "custom-model-2" # type: ignore - - # Assert: Original default_deployment must remain unchanged (not mutated by requests) - assert router.default_deployment["litellm_params"]["model"] == "gpt-5-mini" # type: ignore - - # Assert: Shared fields should still be accessible in all copies - assert deployment1["litellm_params"]["api_key"] == "test-key" # type: ignore - assert deployment2["litellm_params"]["api_key"] == "test-key" # type: ignore - - # Assert: Modifying litellm_params in one deployment doesn't affect others - # This tests the shallow copy properly isolated the litellm_params dict level - deployment1["litellm_params"]["temperature"] = 0.9 # type: ignore - assert "temperature" not in deployment2["litellm_params"] # type: ignore - assert "temperature" not in router.default_deployment["litellm_params"] # type: ignore - - # Assert: Deep nested objects ARE shared (intentional trade-off for 100x perf gain) - # Safe because router only modifies top-level litellm_params fields - deployment1["litellm_params"]["custom_config"]["nested_setting"] = "modified" # type: ignore - assert deployment2["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore - assert router.default_deployment["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore class NoItemsAliasDict(dict): def items(self): raise AssertionError("Unexpected full alias iteration via items()") + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup(): router = Router( @@ -26591,6 +14996,7 @@ def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup model_alias_list = router.get_model_list_from_model_alias(model_name="gpt-5-mini") assert model_alias_list == [] + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name(): router = Router( @@ -26821,62 +15227,10 @@ class TestPreCallChecksOptimization: assert deployments[0].get("model_info", {}).get("id") == "small", "First deployment ID changed!" assert deployments[1].get("model_info", {}).get("id") == "large", "Second deployment ID changed!" + if __name__ == "__main__": pytest.main([__file__, "-v"]) -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_is_prompt_management_model_optimization(): - """ - Test early exit optimization works correctly for all cases. - - Optimization: Check if "/" in model name before calling expensive - get_model_list(). This short-circuits 99% of requests that use - standard model names like "gpt-5.5", "claude-3", etc. - - Tests both negative (early exit) and positive (actual detection) cases. - """ - import litellm - - # Test 1: Standard models without "/" -> early exit returns False - router = Router( - model_list=[ - { - "model_name": "gpt-5.5", - "litellm_params": {"model": "gpt-5.5"}, - }, - { - "model_name": "claude-3", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5-20250929"}, - }, - ] - ) - - assert router._is_prompt_management_model("gpt-5.5") is False - assert router._is_prompt_management_model("claude-3") is False - - # Test 2: Models with "/" but not in model_list -> False after check - assert router._is_prompt_management_model("unknown/model") is False - - # Test 3: Actual prompt management models ARE detected (critical positive case) - original_callbacks = litellm._known_custom_logger_compatible_callbacks.copy() - if "langfuse_prompt" not in litellm._known_custom_logger_compatible_callbacks: - litellm._known_custom_logger_compatible_callbacks.append("langfuse_prompt") - - try: - router_with_prompt = Router( - model_list=[ - { - "model_name": "my-langfuse-prompt/test_id", - "litellm_params": {"model": "langfuse_prompt/actual_prompt_id"}, - }, - ] - ) - - # Critical: Must still detect prompt management models correctly - assert router_with_prompt._is_prompt_management_model("my-langfuse-prompt/test_id") is True - - finally: - litellm._known_custom_logger_compatible_callbacks = original_callbacks @pytest.fixture def router(): @@ -26893,6 +15247,7 @@ def router(): ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch(router): @@ -26916,6 +15271,7 @@ async def test_router_acancel_batch(router): assert response.id == "batch_123" assert response.status == "cancelled" + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch_resolves_credential_name(): @@ -26956,114 +15312,11 @@ async def test_router_acancel_batch_resolves_credential_name(): finally: litellm.credential_list = [] -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_router_acancel_batch_removes_unresolved_credential_name(): - router = Router( - model_list=[ - { - "model_name": "gpt-5.5", - "litellm_params": { - "model": "openai/gpt-5.5", - "litellm_credential_name": "missing-openai-credential", - }, - } - ] - ) - mock_response = MagicMock() - mock_response.id = "batch_123" - mock_response.status = "cancelled" - - with ( - patch.object(router, "_core_get_deployment_credentials_with_provider", return_value=None), - patch.object(litellm, "acancel_batch", new_callable=AsyncMock) as mock_cancel, - ): - mock_cancel.return_value = mock_response - - await router.acancel_batch( - model="gpt-5.5", - batch_id="batch_123", - ) - - call_kwargs = mock_cancel.call_args.kwargs - assert "litellm_credential_name" not in call_kwargs @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingHeaders: """Test that embedding methods properly propagate headers from router configuration.""" - def test_embedding_calls_update_kwargs_before_fallbacks(self): - """ - Test that router.embedding() calls _update_kwargs_before_fallbacks. - - This ensures that metadata is properly set up before the fallback mechanism, - which is necessary for header propagation to work correctly. - """ - model_list = [ - { - "model_name": "text-embedding-3-small", - "litellm_params": { - "model": "text-embedding-3-small", - "api_key": "fake-key", - }, - } - ] - - router = Router(model_list=model_list) - - # Mock the _update_kwargs_before_fallbacks method to verify it's called - with patch.object( - router, - "_update_kwargs_before_fallbacks", - wraps=router._update_kwargs_before_fallbacks, - ) as mock_update: - with patch("litellm.embedding") as mock_litellm_embedding: - mock_litellm_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2, 0.3]}]) - - router.embedding(model="text-embedding-3-small", input=["test input"]) - - # Verify _update_kwargs_before_fallbacks was called - mock_update.assert_called_once() - call_kwargs = mock_update.call_args[1] - assert call_kwargs["model"] == "text-embedding-3-small" - assert "kwargs" in call_kwargs - - @pytest.mark.asyncio - async def test_aembedding_calls_update_kwargs_before_fallbacks(self): - """ - Test that router.aembedding() calls _update_kwargs_before_fallbacks. - - This ensures consistency between sync and async embedding methods. - """ - model_list = [ - { - "model_name": "text-embedding-3-small", - "litellm_params": { - "model": "text-embedding-3-small", - "api_key": "fake-key", - }, - } - ] - - router = Router(model_list=model_list) - - # Mock the _update_kwargs_before_fallbacks method to verify it's called - with patch.object( - router, - "_update_kwargs_before_fallbacks", - wraps=router._update_kwargs_before_fallbacks, - ) as mock_update: - with patch("litellm.aembedding", new_callable=AsyncMock) as mock_litellm_aembedding: - mock_litellm_aembedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2, 0.3]}]) - - await router.aembedding(model="text-embedding-3-small", input=["test input"]) - - # Verify _update_kwargs_before_fallbacks was called - mock_update.assert_called_once() - call_kwargs = mock_update.call_args[1] - assert call_kwargs["model"] == "text-embedding-3-small" - assert "kwargs" in call_kwargs - def test_embedding_propagates_default_litellm_params(self): """ Test that embedding calls properly propagate default_litellm_params including headers. @@ -27302,6 +15555,7 @@ class TestRouterEmbeddingHeaders: assert "litellm_trace_id" in completion_kwargs assert "litellm_trace_id" in embedding_kwargs + if __name__ == "__main__": # Run a simple test test = TestRouterEmbeddingHeaders() @@ -27318,6 +15572,7 @@ OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route: return respx_mock.post(url).mock( return_value=httpx.Response( @@ -27331,11 +15586,13 @@ def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route ) ) + def _sent(route: respx.Route, index: int) -> tuple[str, str, list[str]]: request = route.calls[index].request body = json.loads(request.read()) return request.headers["authorization"], body["model"], body["input"] + def _alias_router() -> Router: return Router( model_list=[ @@ -27349,59 +15606,11 @@ def _alias_router() -> Router: ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingIntegration: """Integration tests for embedding with router configuration.""" - def test_vector_store_request_metadata_prefers_litellm_metadata(self): - assert Router._vector_store_request_metadata( - { - "litellm_metadata": {"user_api_key_team_id": "team-a"}, - "metadata": {"user_api_key_team_id": "team-b"}, - } - ) == {"user_api_key_team_id": "team-a"} - - assert Router._vector_store_request_metadata({"metadata": {"user_api_key_team_id": "team-b"}}) == { - "user_api_key_team_id": "team-b" - } - assert Router._vector_store_request_metadata({}) == {} - - def test_sync_vector_store_wrapper_injects_router_embedding_executor(self): - router = Router(model_list=[]) - original = MagicMock(return_value="searched") - wrapped = router.factory_function(original, call_type="vector_store_search") - - assert ( - wrapped( - vector_store_id="store", - query="query", - custom_llm_provider="valkey", - metadata={"user_api_key_team_id": "team-a"}, - ) - == "searched" - ) - - call_kwargs = original.call_args.kwargs - assert call_kwargs["custom_llm_provider"] == "valkey" - executor = call_kwargs["_direct_vector_store_embedding_executor"] - assert isinstance(executor, RouterVectorStoreEmbeddingExecutor) - assert executor.metadata == {"user_api_key_team_id": "team-a"} - - def test_sync_vector_store_wrapper_preserves_model_routing(self): - router = Router(model_list=[]) - original = MagicMock() - wrapped = router.factory_function(original, call_type="vector_store_search") - - with patch.object(router, "_generic_api_call_with_fallbacks", return_value="routed") as fallback: - assert wrapped(model="vector-alias", vector_store_id="store", query="query") == "routed" - - assert fallback.call_args.kwargs["model"] == "vector-alias" - assert fallback.call_args.kwargs["original_function"] is original - assert isinstance( - fallback.call_args.kwargs["_direct_vector_store_embedding_executor"], - RouterVectorStoreEmbeddingExecutor, - ) - @pytest.mark.asyncio async def test_vector_store_embedding_executors_cover_sdk_and_router_paths( self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch @@ -27835,318 +16044,11 @@ class TestRouterEmbeddingIntegration: assert "headers" in call_kwargs assert call_kwargs["headers"]["X-Custom-Azure-Header"] == "azure-value" + if __name__ == "__main__": # Run tests pytest.main([__file__, "-v"]) -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -class TestRouterIndexManagement: - """Test cases for router index management functions""" - - @pytest.fixture - def router(self): - """Create a router instance for testing""" - return Router(model_list=[]) - - def test_deletion_updates_model_name_indices(self, router): - """Test that deleting a deployment updates model_name_to_deployment_indices correctly""" - router.model_list = [ - {"model_name": "gpt-3.5", "model_info": {"id": "model-1"}}, - {"model_name": "gpt-5.5", "model_info": {"id": "model-2"}}, - {"model_name": "gpt-5.5", "model_info": {"id": "model-3"}}, - {"model_name": "claude", "model_info": {"id": "model-4"}}, - ] - router.model_id_to_deployment_index_map = { - "model-1": 0, - "model-2": 1, - "model-3": 2, - "model-4": 3, - } - router.model_name_to_deployment_indices = { - "gpt-3.5": [0], - "gpt-5.5": [1, 2], - "claude": [3], - } - - # Remove one of the duplicate gpt-5.5 deployments - router._update_deployment_indices_after_removal(model_id="model-2", removal_idx=1) - - # Verify indices are shifted correctly - assert router.model_name_to_deployment_indices["gpt-3.5"] == [0] - assert router.model_name_to_deployment_indices["gpt-5.5"] == [1] # was [1,2], removed 1, shifted 2->1 - assert router.model_name_to_deployment_indices["claude"] == [2] # was [3], shifted to [2] - - # Remove the last gpt-5.5 deployment - router._update_deployment_indices_after_removal(model_id="model-3", removal_idx=1) - - # Verify gpt-5.5 is removed from dict when no deployments remain - assert "gpt-5.5" not in router.model_name_to_deployment_indices - assert router.model_name_to_deployment_indices["gpt-3.5"] == [0] - assert router.model_name_to_deployment_indices["claude"] == [1] - - def test_build_model_id_to_deployment_index_map(self, router): - """Test _build_model_id_to_deployment_index_map function""" - model_list = [ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini"}, - "model_info": {"id": "model-1"}, - }, - { - "model_name": "gpt-5.5", - "litellm_params": {"model": "gpt-5.5"}, - "model_info": {"id": "model-2"}, - }, - ] - - # Test: Build index from model list - router._build_model_id_to_deployment_index_map(model_list) - - # Verify: model_list is populated - assert len(router.model_list) == 2 - # Verify: model_id_to_deployment_index_map is correctly built - assert router.model_id_to_deployment_index_map["model-1"] == 0 - assert router.model_id_to_deployment_index_map["model-2"] == 1 - - def test_add_model_to_list_and_index_map_from_model_info(self, router): - """Test _add_model_to_list_and_index_map extracting model_id from model_info""" - # Setup: Empty router - router.model_list = [] - router.model_id_to_deployment_index_map = {} - - # Test: Add model without explicit model_id - model = {"model": "test-model", "model_info": {"id": "model-info-id"}} - router._add_model_to_list_and_index_map(model=model) - - # Verify: Model added to list - assert len(router.model_list) == 1 - assert router.model_list[0] == model - - # Verify: Index map uses model_info.id - assert router.model_id_to_deployment_index_map["model-info-id"] == 0 - - def test_add_model_to_list_and_index_map_multiple_models(self, router): - """Test _add_model_to_list_and_index_map with multiple models to verify indexing""" - # Setup: Empty router - router.model_list = [] - router.model_id_to_deployment_index_map = {} - - # Test: Add multiple models - model1 = {"model": "model1", "model_info": {"id": "id-1"}} - model2 = {"model": "model2", "model_info": {"id": "id-2"}} - model3 = {"model": "model3", "model_info": {"id": "id-3"}} - - router._add_model_to_list_and_index_map(model=model1, model_id="id-1") - router._add_model_to_list_and_index_map(model=model2, model_id="id-2") - router._add_model_to_list_and_index_map(model=model3, model_id="id-3") - - # Verify: All models added to list - assert len(router.model_list) == 3 - assert router.model_list[0] == model1 - assert router.model_list[1] == model2 - assert router.model_list[2] == model3 - - # Verify: Correct indices in map - assert router.model_id_to_deployment_index_map["id-1"] == 0 - assert router.model_id_to_deployment_index_map["id-2"] == 1 - assert router.model_id_to_deployment_index_map["id-3"] == 2 - - def test_update_team_model_index(self, router): - """Test _update_team_model_index updates team_model_to_deployment_indices.""" - model = { - "model_name": "team-alias", - "model_info": { - "id": "dep-1", - "team_id": "team-abc", - "team_public_model_name": "gpt-5.5", - }, - } - router._update_team_model_index(model, 0) - assert router.team_model_to_deployment_indices[("team-abc", "gpt-5.5")] == [0] - router._update_team_model_index(model, 2) - assert router.team_model_to_deployment_indices[("team-abc", "gpt-5.5")] == [0, 2] - - router._update_team_model_index({"model_name": "x", "model_info": {"id": "dep-2"}}, 5) - assert router.team_model_to_deployment_indices == { - ("team-abc", "gpt-5.5"): [0, 2], - } - - def test_has_model_id(self, router): - """Test has_model_id function for O(1) membership check""" - # Setup: Add models to router - router.model_list = [ - {"model": "test1", "model_info": {"id": "model-1"}}, - {"model": "test2", "model_info": {"id": "model-2"}}, - {"model": "test3", "model_info": {"id": "model-3"}}, - ] - router.model_id_to_deployment_index_map = { - "model-1": 0, - "model-2": 1, - "model-3": 2, - } - - # Test: Check existing model IDs - assert router.has_model_id("model-1") == True - assert router.has_model_id("model-2") == True - assert router.has_model_id("model-3") == True - - # Test: Check non-existing model IDs - assert router.has_model_id("non-existent") == False - assert router.has_model_id("") == False - assert router.has_model_id("model-4") == False - - # Test: Empty router - empty_router = Router(model_list=[]) - assert empty_router.has_model_id("any-id") == False - - def test_build_model_name_index(self, router): - """Test _build_model_name_index function""" - model_list = [ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini"}, - "model_info": {"id": "model-1"}, - }, - { - "model_name": "gpt-5.5", - "litellm_params": {"model": "gpt-5.5"}, - "model_info": {"id": "model-2"}, - }, - { - "model_name": "gpt-5.5", # Duplicate model_name, different deployment - "litellm_params": {"model": "gpt-5.5"}, - "model_info": {"id": "model-3"}, - }, - ] - - # Test: Build index from model list - router._build_model_name_index(model_list) - - # Verify: model_name_to_deployment_indices is correctly built - assert "gpt-5-mini" in router.model_name_to_deployment_indices - assert "gpt-5.5" in router.model_name_to_deployment_indices - - # Verify: gpt-5-mini has single deployment - assert router.model_name_to_deployment_indices["gpt-5-mini"] == [0] - - # Verify: gpt-5.5 has multiple deployments - assert router.model_name_to_deployment_indices["gpt-5.5"] == [1, 2] - - # Test: Rebuild index (should clear and rebuild) - new_model_list = [ - { - "model_name": "claude-3", - "litellm_params": {"model": "claude-3"}, - "model_info": {"id": "model-4"}, - }, - ] - router._build_model_name_index(new_model_list) - - # Verify: Old entries are cleared - assert "gpt-5-mini" not in router.model_name_to_deployment_indices - assert "gpt-5.5" not in router.model_name_to_deployment_indices - - # Verify: New entry is added - assert "claude-3" in router.model_name_to_deployment_indices - assert router.model_name_to_deployment_indices["claude-3"] == [0] - - def test_no_linear_scans_in_router(self): - """ - Static analysis test to ensure Router doesn't use O(n) linear scans. - - Scans router.py for 'in self.model_list' pattern which indicates - inefficient O(n) iteration instead of using index-based O(1) lookups. - - Methods should use: - - model_id_to_deployment_index_map for O(1) model_id lookups - - model_name_to_deployment_indices for O(1) + O(k) model_name lookups - """ - # Methods that are allowed to iterate through self.model_list - ALLOWED_METHODS = { - "_get_deployment_by_litellm_model": "lookup by litellm_params.model, which is not indexed", - "_finalize_adaptive_router_if_configured": 'init-time prefix scan for "auto_router/adaptive_router"; no index for prefix match', - "config_deployments": "filters the whole list on model_info.db_model; admin path only (model add/upsert)", - "auto_router_capability_violation": "counts gated auto-routers across the whole list; admin path only (auto-router init/upsert)", - } - - # Get path to router.py - router_file = os.path.join( - os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))), - "litellm", - "router.py", - ) - - # Read the file - with open(router_file, "r") as f: - content = f.read() - - # Parse with AST - tree = ast.parse(content) - - # Find violations - violations = [] - ignore_methods = set(ALLOWED_METHODS) - - for node in ast.walk(tree): - if isinstance(node, ast.FunctionDef): - method_name = node.name - - # Skip ignored methods - if method_name in ignore_methods: - continue - - # Get source for this method - try: - method_source = ast.get_source_segment(content, node) - if not method_source: - continue - - # Check for the anti-pattern: "in self.model_list" - # This catches: for x in self.model_list, if x in self.model_list, etc. - if "in self.model_list" in method_source: - # Extract the specific line for better error reporting - lines = method_source.split("\n") - pattern_line = None - for line in lines: - if "in self.model_list" in line: - pattern_line = line.strip() - break - - violations.append( - { - "method": method_name, - "line": node.lineno, - "pattern": pattern_line or "in self.model_list", - } - ) - except Exception: - # Skip if we can't get source segment - pass - - # Assert no violations - if violations: - error_msg = "\n".join([f" - {v['method']}() at line {v['line']}: {v['pattern']}" for v in violations]) - - pytest.fail( - f"\n{'=' * 70}\n" - f"Found O(n) linear scan pattern in router.py:\n\n" - f"{error_msg}\n\n" - f"These methods should use index maps instead:\n" - f" - model_id_to_deployment_index_map (for model_id lookups)\n" - f" - model_name_to_deployment_indices (for model_name lookups)\n\n" - f"If a method legitimately needs O(n) iteration, add it to\n" - f"ALLOWED_METHODS in this test method.\n" - f"{'=' * 70}\n" - ) - - def test_model_names_is_set(self): - """Verify that model_names uses a set for O(1) lookups, not a list (O(n))""" - router = Router(model_list=[]) - - assert isinstance(router.model_names, set), ( - f"model_names should be a set for O(1) lookups, but got {type(router.model_names)}" - ) - _PRIORITY_CHAT_COMPLETION: Final = { "id": "chatcmpl-priority", @@ -28599,7 +16501,6 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info() """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -28653,7 +16554,6 @@ def test_router_context_window_check_pre_call_check(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -28707,7 +16607,6 @@ def test_router_context_window_check_pre_call_check_out_group(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -28733,7 +16632,13 @@ def test_router_context_window_check_pre_call_check_out_group(): }, ] - router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0, context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}]) # type: ignore + router = Router( + model_list=model_list, + set_verbose=True, + enable_pre_call_checks=True, + num_retries=0, + context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}], + ) # type: ignore response = router.completion( model="gpt-3.5-turbo-small", @@ -28829,13 +16734,9 @@ def test_router_region_pre_call_check(allowed_model_region): if allowed_model_region is None: assert len(_healthy_deployments) == 2 else: - assert len(_healthy_deployments) == 1, "{} models selected as healthy".format( - len(_healthy_deployments) - ) - assert ( - _healthy_deployments[0]["model_info"]["id"] == "1" - ), "Incorrect model id picked. Got id={}, expected id=1".format( - _healthy_deployments[0]["model_info"]["id"] + assert len(_healthy_deployments) == 1, "{} models selected as healthy".format(len(_healthy_deployments)) + assert _healthy_deployments[0]["model_info"]["id"] == "1", ( + "Incorrect model id picked. Got id={}, expected id=1".format(_healthy_deployments[0]["model_info"]["id"]) ) @@ -29008,104 +16909,12 @@ def test_router_context_window_pre_call_check(model, base_model, llm_provider): else: pytest.fail("Expected to raise an error. Got={}".format(updated_list)) except Exception as e: - if ( - llm_provider == "azure" and base_model is not None - ) or llm_provider == "openai": + if (llm_provider == "azure" and base_model is not None) or llm_provider == "openai": pass except Exception as e: pytest.fail(f"Got unexpected exception on router! - {str(e)}") -@pytest.fixture -def restore_retry_after_header_parser(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - litellm.utils, "_get_retry_after_from_exception_header", litellm.utils._get_retry_after_from_exception_header - ) - - -@pytest.mark.usefixtures("restore_retry_after_header_parser") -def test_router_dynamic_cooldown_correct_retry_after_time(): - """ - User feedback: litellm says "No deployments available for selected model, Try again in 60 seconds" - but Azure says to retry in at most 9s - - ``` - {"message": "litellm.proxy.proxy_server.embeddings(): Exception occured - No deployments available for selected model, Try again in 60 seconds. Passed model=text-embedding-ada-002. pre-call-checks=False, allowed_model_region=n/a, cooldown_list=[('b49cbc9314273db7181fe69b1b19993f04efb88f2c1819947c538bac08097e4c', {'Exception Received': 'litellm.RateLimitError: AzureException RateLimitError - Requests to the Embeddings_Create Operation under Azure OpenAI API version 2023-09-01-preview have exceeded call rate limit of your current OpenAI S0 pricing tier. Please retry after 9 seconds. Please go here: https://aka.ms/oai/quotaincrease if you would like to further increase the default rate limit.', 'Status Code': '429'})]", "level": "ERROR", "timestamp": "2024-08-22T03:25:36.900476"} - ``` - """ - router = Router( - model_list=[ - { - "model_name": "text-embedding-ada-002", - "litellm_params": { - "model": "openai/text-embedding-ada-002", - }, - } - ] - ) - - openai_client = openai.OpenAI(api_key="") - - cooldown_time = 30 - - def _return_exception(*args, **kwargs): - from httpx import Headers, Request, Response - - kwargs = { - "request": Request("POST", "https://www.google.com"), - "message": "Error code: 429 - Rate Limit Error!", - "body": {"detail": "Rate Limit Error!"}, - "code": None, - "param": None, - "type": None, - "response": Response( - status_code=429, - headers=Headers( - { - "date": "Sat, 21 Sep 2024 22:56:53 GMT", - "server": "uvicorn", - "retry-after": f"{cooldown_time}", - "content-length": "30", - "content-type": "application/json", - } - ), - request=Request("POST", "http://0.0.0.0:9000/chat/completions"), - ), - "status_code": 429, - "request_id": None, - } - - exception = Exception() - for k, v in kwargs.items(): - setattr(exception, k, v) - raise exception - - with patch.object( - openai_client, - "post", - side_effect=_return_exception, - ): - new_retry_after_mock_client = MagicMock(return_value=-1) - - litellm.utils._get_retry_after_from_exception_header = ( - new_retry_after_mock_client - ) - - try: - router.embedding( - model="text-embedding-ada-002", - input="Hello world!", - client=openai_client, - ) - except litellm.RateLimitError: - pass - - new_retry_after_mock_client.assert_called() - - response_headers: httpx.Headers = new_retry_after_mock_client.call_args[0][0] - assert int(response_headers["retry-after"]) == cooldown_time - - @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode): @@ -29175,13 +16984,9 @@ async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode): ) if sync_mode: - cooldown_deployments = get_cooldown_deployments( - ctx=router._core_ctx, parent_otel_span=None - ) + cooldown_deployments = get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) else: - cooldown_deployments = await async_get_cooldown_deployments( - ctx=router._core_ctx, parent_otel_span=None - ) + cooldown_deployments = await async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) assert len(cooldown_deployments) > 0 @@ -29282,93 +17087,6 @@ def test_model_group_alias(hidden: bool) -> None: assert len(model_names) == len(_model_list) + 1 -def test_get_team_specific_model() -> None: - """ - Test that _get_team_specific_model returns: - - team_public_model_name when team_id matches - - None when team_id doesn't match - - None when no team_id in model_info - """ - router = Router(model_list=[]) - - deployment = DeploymentTypedDict( - model_name="model-x", - litellm_params={}, - model_info=ModelInfo(team_id="team1", team_public_model_name="public-model-x"), - ) - assert router._get_team_specific_model(deployment, "team1") == "public-model-x" - - assert router._get_team_specific_model(deployment, "team2") is None - - deployment = DeploymentTypedDict( - model_name="model-y", - litellm_params={}, - model_info=ModelInfo(team_public_model_name="public-model-y"), - ) - assert router._get_team_specific_model(deployment, "team1") is None - - deployment = DeploymentTypedDict(model_name="model-z", litellm_params={}, model_info=ModelInfo()) - assert router._get_team_specific_model(deployment, "team1") is None - - -def test_is_team_specific_model() -> None: - """ - Test that _is_team_specific_model returns: - - True when model_info contains team_id - - False when model_info doesn't contain team_id - - False when model_info is None - """ - router = Router(model_list=[]) - - model_info = ModelInfo(team_id="team1", team_public_model_name="public-model-x") - assert router._is_team_specific_model(model_info) is True - - model_info = ModelInfo(team_public_model_name="public-model-y") - assert router._is_team_specific_model(model_info) is False - - model_info = ModelInfo() - assert router._is_team_specific_model(model_info) is False - - assert router._is_team_specific_model(None) is False - - -def test_router_prompt_management_factory(): - router = Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo"}, - }, - { - "model_name": "chatbot_actions", - "litellm_params": { - "model": "langfuse/openai-gpt-3.5-turbo", - "tpm": 1000000, - "prompt_id": "jokes", - }, - }, - { - "model_name": "openai-gpt-3.5-turbo", - "litellm_params": { - "model": "openai/gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), - }, - }, - ] - ) - - assert router._is_prompt_management_model("chatbot_actions") is True - assert router._is_prompt_management_model("openai-gpt-3.5-turbo") is False - - response = router._prompt_management_factory( - model="chatbot_actions", - messages=[{"role": "user", "content": "Hello world!"}], - kwargs={}, - ) - - print(response) - - def test_router_get_model_list_from_model_alias() -> None: router = Router( model_list=[ @@ -29533,39 +17251,6 @@ async def test_enforce_model_rate_limits_returns_429_at_deployment_rpm(): assert rate_limited.value.status_code == 429 -@pytest.mark.parametrize( - "startup_nodes, expected_cache_type", - [ - pytest.param( - [dict(host="node1.localhost", port=6379)], - RedisClusterCache, - id="Expects a RedisClusterCache instance when startup_nodes provided", - ), - pytest.param( - None, - RedisCache, - id="Expects a RedisCache instance when there is no startup nodes", - ), - ], -) -def test_create_correct_redis_cache_instance( - startup_nodes: list[dict] | None, - expected_cache_type: type[RedisClusterCache] | type[RedisCache], -): - cache_config = dict( - host="mockhost", - port=6379, - password="mock-password", - startup_nodes=startup_nodes, - ) - - def _mock_redis_cache_init(*args, **kwargs): ... - - with patch.object(RedisCache, "__init__", _mock_redis_cache_init): - redis_cache = Router._create_redis_cache(cache_config) - assert isinstance(redis_cache, expected_cache_type) - - class CompletionCustomHandler( CustomLogger ): # https://docs.litellm.ai/docs/observability/custom_callback#callback-class @@ -29614,12 +17299,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception as e: print(f"Assertion Error: {traceback.format_exc()}") @@ -29645,9 +17326,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.iscoroutine(kwargs["original_response"]) or inspect.isasyncgen(kwargs["original_response"]) ) @@ -29659,12 +17338,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -29681,23 +17356,18 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponseStream) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -29718,22 +17388,17 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponse) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) ) + assert isinstance(kwargs["api_key"], (str, type(None))) + assert isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) assert isinstance(kwargs["additional_args"], (dict, type(None))) assert isinstance(kwargs["log_event_type"], str) assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool) @@ -29752,23 +17417,18 @@ class CompletionCustomHandler( assert response_obj == None ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or kwargs["original_response"] == None ) assert isinstance(kwargs["additional_args"], (dict, type(None))) @@ -29801,29 +17461,19 @@ class CompletionCustomHandler( ## END TIME assert isinstance(end_time, datetime) ## RESPONSE OBJECT - assert isinstance( - response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse) - ) + assert isinstance(response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse)) ## KWARGS assert isinstance(kwargs["model"], str) # checking we use base_model for azure cost calculation - base_model = litellm.utils.get_base_model_from_metadata( - model_call_details=kwargs - ) + base_model = litellm.utils.get_base_model_from_metadata(model_call_details=kwargs) - if ( - kwargs["model"] == "chatgpt-v-3" - and base_model is not None - and kwargs["stream"] != True - ): + if kwargs["model"] == "chatgpt-v-3" and base_model is not None and kwargs["stream"] != True: # when base_model is set for azure, we should use pricing for the base_model # this checks response_cost == litellm.cost_per_token(model=base_model) assert isinstance(kwargs["response_cost"], float) response_cost = kwargs["response_cost"] - print( - f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}" - ) + print(f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}") prompt_tokens = response_obj.usage.prompt_tokens completion_tokens = response_obj.usage.completion_tokens # ensure the pricing is based on the base_model here @@ -29834,9 +17484,9 @@ class CompletionCustomHandler( ) expected_price = prompt_price + completion_price print(f"expected price: {expected_price}") - assert ( - response_cost == expected_price - ), f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + assert response_cost == expected_price, ( + f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + ) assert isinstance(kwargs["messages"], list) assert isinstance(kwargs["optional_params"], dict) @@ -29847,9 +17497,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -29862,12 +17510,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -29894,9 +17538,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, str, dict)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) or kwargs["original_response"] == None @@ -30052,27 +17694,6 @@ class SpecialEndpointRoutingStrategy(CustomRoutingStrategyBase): pass -def test_reset_custom_routing_strategy(): - """ - Setting a custom routing strategy installs instance-level overrides for - get_available_deployment / async_get_available_deployment. Re-initializing the - routing strategy must clear them so the class implementations are used again. - """ - router = _create_custom_routing_router() - router.set_custom_routing_strategy(SpecialEndpointRoutingStrategy(router)) - - assert "get_available_deployment" in router._core.__dict__ - assert "_core_async_get_available_deployment" in router._core.__dict__ - - router._reset_custom_routing_strategy() - - assert "get_available_deployment" not in router._core.__dict__ - assert "_core_async_get_available_deployment" not in router._core.__dict__ - assert router.async_get_available_deployment.__wrapped__.__func__ is Router.async_get_available_deployment - - router._reset_custom_routing_strategy() - - @pytest.mark.asyncio async def test_select_picks_through_a_custom_routing_strategy() -> None: router: Final = _create_custom_routing_router() @@ -30182,9 +17803,7 @@ def test_multiple_deployments_route_to_a_configured_model( }, ) - route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - side_effect=_deployment_response - ) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_deployment_response) router: Final = Router( model_list=[ { @@ -30333,9 +17952,7 @@ def test_router_retries_transient_server_error( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response(500, json={"error": {"message": "temporary failure"}}), _openai_chat_response("recovered"), @@ -30368,9 +17985,7 @@ def test_router_authentication_error_is_not_swallowed( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response( 401, @@ -30468,9 +18083,9 @@ def test_router_reads_api_key_from_model_list( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock(side_effect=[_openai_chat_response(), _openai_chat_stream_response()]) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=[_openai_chat_response(), _openai_chat_stream_response()] + ) router: Final = Router( model_list=[ { @@ -30493,9 +18108,7 @@ def test_router_reads_api_key_from_model_list( messages=[{"role": "user", "content": "stream key from model list"}], stream=True, ) - streamed_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in stream - ) + streamed_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in stream) assert response.choices[0].message.content == "scripted reply" assert streamed_content == "streamed reply" @@ -30511,9 +18124,9 @@ def test_router_calls_requested_specific_deployment( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock(return_value=_openai_chat_response()) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_openai_chat_response() + ) router: Final = Router( model_list=[ { @@ -30555,9 +18168,7 @@ async def test_router_context_window_error_uses_configured_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response( 400, @@ -30600,9 +18211,7 @@ async def test_router_context_window_error_uses_configured_fallback( assert response.choices[0].message.content == "fallback response" assert route.call_count == 2 - assert tuple( - json.loads(call.request.content)["model"] for call in route.calls - ) == ("gpt-4o-mini", "gpt-4.1-mini") + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("gpt-4o-mini", "gpt-4.1-mini") def test_router_function_calling_sends_tools_and_returns_tool_call( @@ -30610,9 +18219,7 @@ def test_router_function_calling_sends_tools_and_returns_tool_call( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( return_value=httpx.Response( 200, json={ @@ -30691,9 +18298,7 @@ async def test_router_openai_completion_supports_sync_async_and_streaming( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ _openai_chat_response("async reply"), _openai_chat_response("sync reply"), @@ -30728,17 +18333,13 @@ async def test_router_openai_completion_supports_sync_async_and_streaming( stream=True, ) async_stream_chunks: Final = tuple([chunk async for chunk in async_stream]) - async_stream_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in async_stream_chunks - ) + async_stream_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in async_stream_chunks) sync_stream: Final = router.completion( model="openai-model", messages=[{"role": "user", "content": "sync stream"}], stream=True, ) - sync_stream_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in sync_stream - ) + sync_stream_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in sync_stream) assert async_response.choices[0].message.content == "async reply" assert sync_response.choices[0].message.content == "sync reply" @@ -30782,9 +18383,7 @@ async def test_router_moderation_returns_provider_result( num_retries=0, ) - response: Final = await router.amoderation( - model="moderation-model", input="safe scripted text" - ) + response: Final = await router.amoderation(model="moderation-model", input="safe scripted text") assert response.results[0].flagged is False assert json.loads(route.calls[0].request.content)["input"] == "safe scripted text" @@ -30953,9 +18552,7 @@ async def test_batch_completion_multiple_models_returns_each_model_response( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock( - side_effect=_batch_response - ) + route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock(side_effect=_batch_response) router: Final = _batch_router() responses: Final = cast( @@ -30978,9 +18575,7 @@ async def test_batch_completion_multiple_models_and_messages_preserves_each_pair respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock( - side_effect=_batch_response - ) + route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock(side_effect=_batch_response) router: Final = _batch_router() responses: Final = cast( @@ -31061,9 +18656,7 @@ async def test_batch_completion_fastest_response_streams_mocked_chunks() -> None ) chunks: Final = tuple([chunk async for chunk in response]) text: Final = "".join( - chunk.choices[0].delta.content or "" - for chunk in chunks - if chunk.choices and chunk.choices[0].delta is not None + chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices and chunk.choices[0].delta is not None ) assert text == "stream answer" @@ -31219,7 +18812,9 @@ def _first_chunk_cases(): from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="hi"))]) - provider_error = litellm.InternalServerError(message="primary died at stream start", llm_provider="openai", model="m") + provider_error = litellm.InternalServerError( + message="primary died at stream start", llm_provider="openai", model="m" + ) wrapped = MidStreamFallbackError( message=str(provider_error), model="m", @@ -31236,7 +18831,7 @@ def _first_chunk_cases(): @pytest.mark.parametrize("first, expected, raised", _first_chunk_cases()) def test_read_first_chunk_direct_call(first, expected, raised): - from litellm.router import _read_first_chunk + from litellm.router_rust.invoke.streaming import _read_first_chunk # noqa: PLC2701 stream = _FirstPullStream(first) if raised is None: @@ -31251,7 +18846,7 @@ def test_read_first_chunk_direct_call(first, expected, raised): @pytest.mark.asyncio @pytest.mark.parametrize("first, expected, raised", _first_chunk_cases()) async def test_aread_first_chunk_direct_call(first, expected, raised): - from litellm.router import _aread_first_chunk + from litellm.router_rust.invoke.streaming import _aread_first_chunk # noqa: PLC2701 stream = _FirstPullStream(first) if raised is None: @@ -31263,31 +18858,6 @@ async def test_aread_first_chunk_direct_call(first, expected, raised): assert caught.value is raised -def _chunks_with_first_cases(): - return [ - pytest.param(None, ["a", "b"], ["a", "b"], id="nothing-read-stream-as-is"), - pytest.param((None, ["a", "b"]), ["a", "b"], [], id="empty-stream-not-touched"), - pytest.param(("first", ["a", "b"]), ["never-iterated"], ["first", "a", "b"], id="first-chunk-put-back"), - ] - - -@pytest.mark.parametrize("taken, stream, expected", _chunks_with_first_cases()) -def test_chunks_with_first_direct_call(taken, stream, expected): - from litellm.router import _chunks_with_first, _FirstChunk - - first = None if taken is None else _FirstChunk(chunk=taken[0], rest=iter(taken[1])) - assert list(_chunks_with_first(iter(stream), first)) == expected - - -@pytest.mark.asyncio -@pytest.mark.parametrize("taken, stream, expected", _chunks_with_first_cases()) -async def test_achunks_with_first_direct_call(taken, stream, expected): - from litellm.router import _achunks_with_first, _FirstChunk - - first = None if taken is None else _FirstChunk(chunk=taken[0], rest=_AsyncList(taken[1])) - assert [item async for item in _achunks_with_first(_AsyncList(stream), first)] == expected - - class _IterationHandedOffStream(CustomStreamWrapper): """A CustomStreamWrapper that, like the MCP gateway's MCPStreamWrapper, has no completion_stream of its own and hands iteration to another object from __aiter__ / __iter__, a fresh one per call so a second @@ -31505,30 +19075,9 @@ def test_pre_call_checks_keeps_deployment_whose_unsupported_param_is_allowed_or_ assert [d["model_info"]["id"] for d in filtered] == ["only"] -def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supports_the_param(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "drop_params", False) - model_list: Final = [ - { - "model_name": "mixed-json-mode", - "litellm_params": {"model": f"openai/{deployment_model}", "api_key": "my-fake-key"}, - "model_info": {"id": deployment_model, "max_input_tokens": 100}, - } - for deployment_model in ("gpt-3.5-turbo-16k", "gpt-4o") - ] - router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) - without_json_mode: Final = [d for d in router.model_list if d["model_info"]["id"] == "gpt-3.5-turbo-16k"] - - filtered: Final = router._pre_call_checks( - model="mixed-json-mode", - healthy_deployments=without_json_mode, - messages=[{"role": "user", "content": "hi"}], - request_kwargs={"response_format": {"type": "json_object"}}, - ) - - assert filtered == [] - - -def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param(monkeypatch: pytest.MonkeyPatch): +def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) router: Final = Router( model_list=[ @@ -31597,7 +19146,9 @@ def test_pre_call_checks_prefers_a_deployment_that_takes_the_param_over_dropping assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] -def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region(monkeypatch: pytest.MonkeyPatch): +def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) model_list: Final = [ { @@ -31827,106 +19378,6 @@ async def test_router_configured_stream_usage_keeps_a_choice_on_the_usage_chunk( assert all(len(chunk.choices) == 1 for chunk in chunks[:-1]) -@pytest.mark.parametrize( - ("stream_options", "marked_injected"), - [({"include_usage": True}, True), (None, False)], - ids=["include-usage", "no-stream-options"], -) -def test_update_kwargs_with_deployment_reads_stream_options_from_pydantic_litellm_params( - stream_options: dict[str, bool] | None, marked_injected: bool -): - from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD - from litellm.types.router import LiteLLM_Params - - kwargs: Final[dict[str, object]] = {"stream": True, "metadata": {}} - Router()._update_kwargs_with_deployment( - deployment={ - "model_name": "gpt-5-mini", - "litellm_params": LiteLLM_Params(model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options), - "model_info": {"id": "pydantic-params"}, - }, - kwargs=kwargs, - ) - - assert (kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True) is marked_injected - - -@pytest.mark.asyncio -async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_streams_tokens(): - """After a mid-stream hop, a disconnect bills the primary's relayed text plus the fallback's, not the fallback alone.""" - import json - - from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - - def _relayed_iterator(text: str) -> ResponsesAPIStreamingIterator: - logging_obj: Final = MagicMock(spec=LiteLLMLogging) - logging_obj.model_call_details = {"litellm_params": {}} - logging_obj.completion_start_time = None - iterator: Final = ResponsesAPIStreamingIterator( - response=httpx.Response(200), - model="gpt-5.4", - responses_api_provider_config=OpenAIResponsesAPIConfig(), - logging_obj=logging_obj, - request_data={"model": "gpt-5.4", "input": "Hi"}, - custom_llm_provider="openai", - ) - created: Final = { - "id": "resp_1", - "object": "response", - "created_at": 1700000000.0, - "model": "gpt-5.4", - "status": "in_progress", - "output": [], - "parallel_tool_calls": True, - "tool_choice": "auto", - "tools": [], - } - iterator._process_chunk(json.dumps({"type": "response.created", "sequence_number": 0, "response": created})) - iterator._process_chunk( - json.dumps( - { - "type": "response.output_text.delta", - "sequence_number": 1, - "item_id": "msg_1", - "output_index": 0, - "content_index": 0, - "delta": text, - } - ) - ) - return iterator - - router: Final = Router( - model_list=[{"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k"}}] - ) - primary: Final = _relayed_iterator("Once upon a time there was a lighthouse keeper") - fallback: Final = _relayed_iterator("who lit the lamp") - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": AsyncMock()}, - ) - primary_tokens: Final = primary.partial_response_on_disconnect().response.usage.output_tokens - fallback_tokens: Final = fallback.partial_response_on_disconnect().response.usage.output_tokens - - assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens - - wrapped.adopt_fallback_headers(fallback) - - assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens + fallback_tokens - - -_SYNC_SECONDS: Final = 10 -_ROUTER_FILE: Final = Router._update_deployment_indices_after_removal.__code__.co_filename -_INDEX_LOCK_LINES: Final = frozenset( - number - for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1) - if "with self._deployment_index_lock" in line -) - -_Read = TypeVar("_Read") - - def _concurrent_delete_deployment(index: int) -> dict[str, object]: return { "model_name": f"model-{index}", @@ -31943,272 +19394,6 @@ def _concurrent_delete_router() -> Router: return Router(model_list=[_concurrent_delete_deployment(index) for index in range(4)]) -def _read_during_held_write( - read: Callable[[], _Read], - write: Callable[[], object], - event: Literal["call", "c_call"] = "call", - write_after: CodeType | None = None, -) -> _Read: - """Run `read` on its own thread while `write` waits at its first `event` in the re-index, after the list shrank and - before the maps caught up. With `write_after`, the write starts when `read` returns from that function. The write - resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets the half-updated - list and a read that takes the lock waits for the write""" - reindex: Final = Router._update_deployment_indices_after_removal.__code__ - held: Final = threading.Event() - release: Final = threading.Event() - settled: Final = threading.Event() - released_in_time: Final[list[bool]] = [] - outcome: Final[list[_Read]] = [] - failure: Final[list[BaseException]] = [] - - def hold_in_reindex(frame: FrameType, observed: str, _arg: object) -> None: - if observed == event and frame.f_code is reindex and not held.is_set(): - held.set() - released_in_time.append(release.wait(timeout=_SYNC_SECONDS)) - - def run_write() -> None: - previous: Final = sys.getprofile() - sys.setprofile(hold_in_reindex) - try: - write() - finally: - sys.setprofile(previous) - - writer: Final = threading.Thread(target=run_write) - - def start_write() -> None: - writer.start() - assert held.wait(timeout=_SYNC_SECONDS), "the write never reached re-indexing" - - def watch_read(frame: FrameType, observed: str, _arg: object) -> Callable[..., object] | None: - if frame.f_code.co_filename != _ROUTER_FILE: - return None - if observed == "return" and frame.f_code is write_after and writer.ident is None: - start_write() - elif observed == "line" and held.is_set() and frame.f_lineno in _INDEX_LOCK_LINES: - settled.set() - return watch_read - - def run_read() -> None: - previous: Final = sys.gettrace() - sys.settrace(watch_read) - try: - outcome.append(read()) - except BaseException as error: - failure.append(error) - finally: - sys.settrace(previous) - settled.set() - - reader: Final = threading.Thread(target=run_read) - try: - if write_after is None: - start_write() - reader.start() - assert settled.wait(timeout=_SYNC_SECONDS), "the read neither returned nor reached the index lock" - finally: - release.set() - if reader.ident is not None: - reader.join() - if writer.ident is not None: - writer.join() - assert released_in_time == [True], "the write was not held for the whole read" - if failure: - raise failure[0] - return outcome[0] - - -def _deployment_ids(deployments: list[DeploymentTypedDict]) -> tuple[str, ...]: - return tuple(deployment["model_info"]["id"] for deployment in deployments) - - -@pytest.mark.parametrize( - ("read", "expected"), - [ - pytest.param( - lambda router: router.get_deployment("deployment-3").model_info.id, "deployment-3", id="get_deployment" - ), - pytest.param( - lambda router: router.get_model_info("deployment-3")["model_info"]["id"], - "deployment-3", - id="get_model_info", - ), - pytest.param( - lambda router: _deployment_ids(router.get_all_deployments("model-3")), - ("deployment-3",), - id="get_all_deployments", - ), - pytest.param( - lambda router: _deployment_ids(router.get_all_deployments("public-shared", team_id="team-1")), - ("deployment-2", "deployment-3"), - id="get_all_deployments-team", - ), - pytest.param( - lambda router: router.map_team_model("public-shared", "team-1"), "public-shared", id="map_team_model" - ), - pytest.param(lambda router: router.get_model_ids("model-3"), ["deployment-3"], id="get_model_ids-by-name"), - pytest.param( - lambda router: router.get_model_ids(), ["deployment-1", "deployment-2", "deployment-3"], id="get_model_ids" - ), - pytest.param( - lambda router: router.get_deployment_by_model_group_name("model-3").model_info.id, - "deployment-3", - id="get_deployment_by_model_group_name", - ), - pytest.param( - lambda router: router.get_model_listing_info("model-3").cost_map_keys, - ("openai/gpt-3",), - id="get_model_listing_info", - ), - pytest.param( - lambda router: router.team_model_has_alternatives("deployment-3"), True, id="team_model_has_alternatives" - ), - pytest.param( - lambda router: ( - router._get_model_group_deployment_usable_by_team( - model_group_name="model-3", team_id="team-1" - ).model_info.id - ), - "deployment-3", - id="model_group_deployment_usable_by_team", - ), - pytest.param( - lambda router: ( - router._get_team_public_name_deployment(model_id="public-shared", team_id="team-1").model_info.id - ), - "deployment-2", - id="team_public_name_deployment", - ), - pytest.param( - lambda router: _deployment_ids(router._team_deployments_across_teams("public-shared")), - ("deployment-2", "deployment-3"), - id="team_deployments_across_teams", - ), - pytest.param( - lambda router: _deployment_ids(list(router.deployments_named("model-3") or ())), - ("deployment-3",), - id="deployments_named", - ), - pytest.param( - lambda router: should_use_polling_for_request( - background_mode=True, - polling_via_cache_enabled=["openai"], - redis_cache=MagicMock(), - model="model-3", - llm_router=router, - ), - True, - id="response_polling_provider_lookup", - ), - pytest.param( - lambda router: _models_in(router._core_ctx, "model-3"), - (Baseline("openai/gpt-3", "deployment-3"),), - id="savings_baseline_candidates", - ), - ], -) -def test_lookup_during_a_concurrent_delete_returns_the_requested_deployment( - read: Callable[[Router], object], expected: object -) -> None: - router: Final = _concurrent_delete_router() - - observed: Final = _read_during_held_write(lambda: read(router), lambda: router.delete_deployment("deployment-0")) - - assert observed == expected - - -def test_lookup_during_an_upsert_returns_the_requested_deployment() -> None: - router: Final = _concurrent_delete_router() - rotated: Final = Deployment( - **{**_concurrent_delete_deployment(0), "litellm_params": {"model": "openai/gpt-0", "api_key": "rotated"}} - ) - - observed: Final = _read_during_held_write( - lambda: router.get_deployment("deployment-3").model_info.id, lambda: router.upsert_deployment(rotated) - ) - - assert observed == "deployment-3" - - -def test_add_during_a_concurrent_delete_indexes_the_added_deployment() -> None: - router: Final = _concurrent_delete_router() - - _read_during_held_write( - lambda: router.add_deployment(Deployment(**_concurrent_delete_deployment(4))), - lambda: router.delete_deployment("deployment-0"), - ) - - assert [router.get_deployment(f"deployment-{index}").model_info.id for index in range(1, 5)] == [ - f"deployment-{index}" for index in range(1, 5) - ] - - -def test_set_model_list_during_a_concurrent_delete_indexes_the_new_list() -> None: - router: Final = _concurrent_delete_router() - - _read_during_held_write( - lambda: router.set_model_list([_concurrent_delete_deployment(index) for index in range(10, 14)]), - lambda: router.delete_deployment("deployment-0"), - event="c_call", - ) - - assert router.get_model_ids() == [f"deployment-{index}" for index in range(10, 14)] - - -def test_team_model_has_alternatives_when_a_delete_starts_after_its_deployment_lookup() -> None: - router: Final = _concurrent_delete_router() - - has_alternatives: Final = _read_during_held_write( - lambda: router.team_model_has_alternatives("deployment-3"), - lambda: router.delete_deployment("deployment-0"), - write_after=Router._core_get_deployment.__code__, - ) - - assert has_alternatives is True - - -def test_init_adaptive_router_deployment_during_a_concurrent_delete_prices_the_listed_model() -> None: - priced: Final = { - **_concurrent_delete_deployment(3), - "litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06}, - } - router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced]) - adaptive: Final = Deployment( - model_name="adaptive", - litellm_params=LiteLLM_Params( - model="auto_router/adaptive_router", adaptive_router_config={"available_models": ["model-3"]} - ), - model_info=ModelInfo(id="adaptive"), - ) - - _read_during_held_write( - lambda: router.init_adaptive_router_deployment(deployment=adaptive), - lambda: router.delete_deployment("deployment-0"), - ) - - assert router.adaptive_routers["adaptive"][0].strategy.model_to_cost == {"model-3": 3e-06} - - -def test_complexity_router_adaptive_pool_during_a_concurrent_delete_prices_the_listed_model() -> None: - priced: Final = { - **_concurrent_delete_deployment(3), - "litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06}, - } - router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced]) - complexity: Final = ComplexityRouter( - model_name="hybrid", - ctx=router._core_ctx, - complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["model-3"]}}, - ) - - adaptive: Final = _read_during_held_write( - complexity._ensure_adaptive_router, lambda: router.delete_deployment("deployment-0") - ) - - assert adaptive is not None - assert adaptive.model_to_cost == {"model-3": 3e-06} - - def test_delete_of_an_unknown_deployment_returns_none_and_keeps_the_list() -> None: router: Final = _concurrent_delete_router() @@ -32240,7 +19425,10 @@ _IMAGE_RESPONSES_INPUT = [ ], } ] -_VISION_LEG: Final[dict[str, object]] = {"litellm_params": {"model": "openai/gpt-5.6-terra"}, "model_info": {"id": "vision"}} +_VISION_LEG: Final[dict[str, object]] = { + "litellm_params": {"model": "openai/gpt-5.6-terra"}, + "model_info": {"id": "vision"}, +} _DECLARED_TEXT_LEG: Final[dict[str, object]] = { "litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"id": "declared-text", "supports_vision": False}, @@ -32254,17 +19442,6 @@ def _mixed_router(*deployments: dict[str, object], routing_strategy: str = "late ) -def _candidate_ids( - router: litellm.Router, - messages: list[dict[str, object]] | None = None, - request_kwargs: dict[str, object] | None = None, -) -> list[str]: - candidates: Final = router._without_declared_text_only_deployments( - router._core_get_model_list(model_name="mixed"), "mixed", messages, request_kwargs - ) - return [d["model_info"]["id"] for d in candidates] - - def test_router_picks_the_vision_leg_for_image_requests_without_enable_pre_call_checks(): router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) assert router.enable_pre_call_checks is False @@ -32302,81 +19479,6 @@ def test_router_sends_an_image_to_the_text_only_leg_when_the_vision_leg_is_cooli assert router.get_available_deployment(model="mixed", messages=_IMAGE_CHAT)["model_info"]["id"] == "declared-text" -def test_without_declared_text_only_deployments_passes_a_chosen_deployment_through(): - assert litellm.Router._without_declared_text_only_deployments(_DECLARED_TEXT_LEG, "mixed", _IMAGE_CHAT, None) is _DECLARED_TEXT_LEG - - -def test_image_routing_drop_declared_text_only_legs_for_image_requests(): - router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) - assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision"] - - -def test_image_routing_keep_declared_text_only_legs_for_text_requests(): - router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) - assert _candidate_ids(router, messages=[{"role": "user", "content": "Say hi"}]) == [ - "vision", - "declared-text", - ] - - -def test_image_routing_drop_declared_text_only_legs_for_responses_input_image(): - router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) - assert _candidate_ids(router, request_kwargs={"input": _IMAGE_RESPONSES_INPUT}) == ["vision"] - - -def test_image_routing_send_image_to_the_group_when_every_leg_is_declared_text_only(): - other_text_leg = {**_DECLARED_TEXT_LEG, "model_info": {"id": "declared-text-2", "supports_vision": False}} - router: Final = _mixed_router(_DECLARED_TEXT_LEG, other_text_leg) - assert _candidate_ids(router, messages=_IMAGE_CHAT) == [ - "declared-text", - "declared-text-2", - ] - - -def test_image_routing_keep_undeclared_custom_legs_for_image_requests(): - custom_leg = { - "litellm_params": {"model": "hosted_vllm/Qwen2.5-VL-7B-Instruct", "api_base": "http://localhost:8000"}, - "model_info": {"id": "custom"}, - } - router: Final = _mixed_router(_VISION_LEG, custom_leg) - assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision", "custom"] - - -def test_image_routing_resolve_vision_through_the_cost_map_and_base_model(): - catalog_text_leg = {"litellm_params": {"model": "bedrock/deepseek.v3.2"}, "model_info": {"id": "catalog-text"}} - azure_text_leg = { - "litellm_params": {"model": "azure/my-audio-deployment", "api_base": "https://x.openai.azure.com"}, - "model_info": {"id": "azure-text", "base_model": "azure/gpt-audio"}, - } - router: Final = _mixed_router(_VISION_LEG, catalog_text_leg, azure_text_leg) - assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision"] - - -@pytest.mark.parametrize( - ("deployment", "expected"), - [ - (_DECLARED_TEXT_LEG, True), - ({"litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"supports_vision": True}}, False), - ({"litellm_params": {"model": "bedrock/deepseek.v3.2"}, "model_info": {}}, True), - ({"litellm_params": {"model": "azure/my-deploy"}, "model_info": {"base_model": "azure/gpt-audio"}}, True), - ({"litellm_params": {"model": "hosted_vllm/Qwen2.5-VL-7B-Instruct"}, "model_info": {}}, False), - ({"litellm_params": {"model": "no-such-model-anywhere"}, "model_info": {}}, False), - ({"litellm_params": {}, "model_info": {}}, False), - ], - ids=["override-false", "override-true", "cost-map-false", "base-model-false", "undeclared", "unresolvable", "no-params"], -) -def test_deployment_declares_no_vision_reads_the_override_then_the_cost_map(deployment: dict[str, object], expected: bool): - assert litellm.Router._deployment_declares_no_vision(deployment, "mixed") is expected - - -def test_request_carries_image_sees_chat_parts_and_responses_input_items(): - assert litellm.Router._request_carries_image(_IMAGE_CHAT, {}) is True - assert litellm.Router._request_carries_image(None, {"input": _IMAGE_RESPONSES_INPUT}) is True - assert litellm.Router._request_carries_image([{"role": "user", "content": "Say hi"}], {}) is False - assert litellm.Router._request_carries_image(None, {"input": "Say hi"}) is False - assert litellm.Router._request_carries_image(None, {"input": ["safe text", "more safe text"]}) is False - - @pytest.mark.asyncio @pytest.mark.usefixtures("isolate_litellm_state") @pytest.mark.parametrize("surface", ["chat", "messages"]) @@ -32398,19 +19500,31 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( headers={"retry-after": "168307"}, json={"type": "error", "error": {"type": "rate_limit_error", "message": "usage limit"}}, ) - return httpx.Response(200, json={ - "id": "msg-healthy", "type": "message", "role": "assistant", "model": "claude-sonnet-5-5", - "content": [{"type": "text", "text": "healthy account"}], "stop_reason": "end_turn", - "usage": {"input_tokens": 2, "output_tokens": 2}, - }) + return httpx.Response( + 200, + json={ + "id": "msg-healthy", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "healthy account"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 2, "output_tokens": 2}, + }, + ) route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=upstream) router: Final = Router( - model_list=[{ - "model_name": "claude", "litellm_params": { - "model": "anthropic/claude-sonnet-5-5", "api_key": "fake-gateway-key", - }, "model_info": {"id": "shared-anthropic"}, - }], + model_list=[ + { + "model_name": "claude", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "fake-gateway-key", + }, + "model_info": {"id": "shared-anthropic"}, + } + ], num_retries=0, allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=1), ) @@ -32418,22 +19532,36 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( async def call(account: str | None) -> object: credentials: Final = ( - {"provider_specific_header": { - "custom_llm_provider": "anthropic", "extra_headers": { - "Authorization": f"Bearer sk-ant-oat-fake-{account}", - }, - }} if credential_kind == "oauth" else {"api_key": f"fake-{account}"} - ) if account is not None else {} + ( + { + "provider_specific_header": { + "custom_llm_provider": "anthropic", + "extra_headers": { + "Authorization": f"Bearer sk-ant-oat-fake-{account}", + }, + } + } + if credential_kind == "oauth" + else {"api_key": f"fake-{account}"} + ) + if account is not None + else {} + ) if surface == "messages": max_tokens: Final = litellm.get_max_tokens("anthropic/claude-sonnet-5-5") assert max_tokens is not None return await router.aanthropic_messages( - model="claude", messages=[{"role": "user", "content": "hello"}], - client=client, max_tokens=max_tokens, **credentials, + model="claude", + messages=[{"role": "user", "content": "hello"}], + client=client, + max_tokens=max_tokens, + **credentials, ) return await router.acompletion( - model="claude", messages=[{"role": "user", "content": "hello"}], - client=client, **credentials, + model="claude", + messages=[{"role": "user", "content": "hello"}], + client=client, + **credentials, ) try: @@ -32671,7 +19799,9 @@ async def test_a_client_model_id_cannot_move_a_sagemaker_embedding_call_off_the_ served: Final = component_embeddings[str(request["InferenceComponentName"])] return { "Body": io.BytesIO( - json.dumps({"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]}).encode() + json.dumps( + {"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]} + ).encode() ) } @@ -32704,7 +19834,9 @@ async def test_a_header_pin_holds_when_the_deployment_names_sagemaker_through_cu async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clients_model_id_alone( respx_mock: respx.MockRouter, _router_httpx_transport: None ) -> None: - route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond(json=dict(_CHAT_COMPLETION_REPLY)) + route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond( + json=dict(_CHAT_COMPLETION_REPLY) + ) router: Final = Router( model_list=[ { @@ -32719,7 +19851,9 @@ async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clien ] ) - await router.acompletion(model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value") + await router.acompletion( + model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value" + ) assert json.loads(route.calls.last.request.content)["model_id"] == "client-value" @@ -32745,65 +19879,3 @@ def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog assert "Deployment 'embed-from-config' sets enforce_supported_endpoints without" in caplog.text assert "Deployment 'embed-from-api' sets enforce_supported_endpoints without" in caplog.text - - -@pytest.mark.asyncio -@pytest.mark.parametrize("mode", ["sync", "async"]) -@pytest.mark.parametrize("metadata_key", [None, "metadata", "litellm_metadata"]) -@pytest.mark.parametrize( - ("route", "messages", "request_kwargs"), - [ - ("/v1/moderations", None, {"input": ["first", "second"]}), - ("/v1/embeddings", None, {"messages": []}), - ("/v1/messages", ["malformed"], {}), - ("/v1/messages", "malformed", {}), - ("/v1/messages", 5, {}), - ("/v1/messages", None, {"messages": ["malformed"]}), - ("/v1/responses", None, {"input": [{"type": "computer_call_output", "output": {"type": "text", "text": "no image"}}]}), - ("/v1/responses", None, {"input": [{"type": "computer_call_output", "output": "malformed"}]}), - ], - ids=["moderation-list", "embedding-empty-messages", "message-entry", "message-string", "message-number", "kwargs-messages", "computer-output-text", "computer-output-string"], -) -async def test_image_routing_leaves_non_image_and_malformed_requests_to_the_endpoint( - mode: str, metadata_key: str | None, route: str, messages: object, request_kwargs: dict[str, object] -) -> None: - router: Final = _mixed_router( - {**_DECLARED_TEXT_LEG, "litellm_params": {**_DECLARED_TEXT_LEG["litellm_params"], "order": 0}}, - {**_VISION_LEG, "litellm_params": {**_VISION_LEG["litellm_params"], "order": 1}}, - ) - kwargs: Final = { - **request_kwargs, - **({metadata_key: {"user_api_key_request_route": route}} if metadata_key is not None else {}), - } - original: Final = copy.deepcopy(request_kwargs) - deployment: Final = ( - await router._core_async_get_available_deployment(model="mixed", messages=messages, request_kwargs=kwargs) - if mode == "async" - else router.get_available_deployment(model="mixed", messages=messages, request_kwargs=kwargs) - ) - assert deployment["model_info"]["id"] == "declared-text" - assert all(kwargs[key] == value for key, value in original.items()) - - -@pytest.mark.parametrize( - "request_kwargs", - [ - {"messages": _IMAGE_CHAT}, - {"input": _IMAGE_RESPONSES_INPUT}, - {"messages": [{"role": "user", "content": [{"type": "image", "source": {"type": "url", "url": "https://example.invalid/image.png"}}]}]}, - {"messages": [{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_image", "content": [_IMAGE_PART]}]}]}, - {"input": [{"type": "function_call_output", "call_id": "call_image", "output": [{"type": "input_image", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "custom_tool_call_output", "call_id": "call_image", "output": [{"type": "input_image", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "computer_call_output", "call_id": "call_image", "output": {"type": "computer_screenshot", "image_url": "https://example.invalid/image.png"}}]}, - {"input": [{"role": "user", "content": [{"type": "computer_screenshot", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "function_call_output", "call_id": "call_image", "output": [{"type": "computer_screenshot", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "custom_tool_call_output", "call_id": "call_image", "output": [{"type": "computer_screenshot", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "computer_call_output", "call_id": "call_image", "output": {"type": "computer_screenshot", "url": "https://example.invalid/image.png"}}]}, - ], - ids=["chat", "responses", "anthropic", "tool-result", "function-output", "custom-output", "computer-output", "screenshot-content", "screenshot-function-output", "screenshot-custom-output", "screenshot-url"], -) -def test_image_routing_inspects_native_request_content_without_translation(request_kwargs: dict[str, object]) -> None: - router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) - original: Final = copy.deepcopy(request_kwargs) - assert _candidate_ids(router, request_kwargs=request_kwargs) == ["vision"] - assert request_kwargs == original diff --git a/tests/unit/test_router/test_router_endpoints.py b/tests/unit/test_router/test_router_endpoints.py index 0c57fb662ac..d59c3f19c50 100644 --- a/tests/unit/test_router/test_router_endpoints.py +++ b/tests/unit/test_router/test_router_endpoints.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import litellm import pytest +import respx from litellm import Router from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -129,7 +130,7 @@ async def test_aspeech_sets_deployment_metadata(): mock_response = MagicMock() with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech: - response = await router._aspeech( + response = await router.aspeech( model="tts", input="the quick brown fox jumped over the lazy dogs", voice="alloy", @@ -165,9 +166,7 @@ async def test_moderation_endpoint_with_api_base(): ) # Mock the OpenAI client to verify api_base is passed - with patch( - "litellm.main.openai_chat_completions.get_openai_client" - ) as mock_get_client: + with patch("litellm.main.openai_chat_completions.get_openai_client") as mock_get_client: mock_client = AsyncMock() mock_response = MagicMock() mock_response.model_dump.return_value = { @@ -185,16 +184,14 @@ async def test_moderation_endpoint_with_api_base(): mock_client.moderations.create = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_client - response = await router.amoderation( - model="openai/omni-moderation-latest", input="hello this is a test" - ) + response = await router.amoderation(model="openai/omni-moderation-latest", input="hello this is a test") # Verify that get_openai_client was called with the custom api_base mock_get_client.assert_called() call_kwargs = mock_get_client.call_args.kwargs - assert ( - call_kwargs.get("api_base") == custom_api_base - ), f"Expected api_base to be {custom_api_base}, but got {call_kwargs.get('api_base')}" + assert call_kwargs.get("api_base") == custom_api_base, ( + f"Expected api_base to be {custom_api_base}, but got {call_kwargs.get('api_base')}" + ) print(f"✓ Moderation endpoint correctly uses api_base: {custom_api_base}") @@ -217,13 +214,6 @@ async def test_aaaaatext_completion_endpoint(model_list, sync_mode): mock_response="I'm fine, thank you!", ) - response_2 = await router._atext_completion( - model="gpt-5-mini", - prompt="Hello, how are you?", - mock_response="I'm fine, thank you!", - ) - assert response_2.choices[0].text == "I'm fine, thank you!" - assert response.choices[0].text == "I'm fine, thank you!" @@ -252,147 +242,9 @@ async def test_router_with_empty_choices(model_list): assert response is not None -@pytest.mark.parametrize("sync_mode", [True, False]) -def test_generic_api_call_with_fallbacks_basic(sync_mode): - """ - Test both the sync and async versions of generic_api_call_with_fallbacks with a basic successful call - """ - # Create a mock function that will be passed to generic_api_call_with_fallbacks - if sync_mode: - from unittest.mock import Mock - - mock_function = Mock() - mock_function.__name__ = "test_function" - else: - mock_function = AsyncMock() - mock_function.__name__ = "test_function" - - # Create a mock response - mock_response = { - "id": "resp_123456", - "role": "assistant", - "content": "This is a test response", - "model": "test-model", - "usage": {"input_tokens": 10, "output_tokens": 20}, - } - mock_function.return_value = mock_response - - # Create a router with a test model - router = Router( - model_list=[ - { - "model_name": "test-model-alias", - "litellm_params": { - "model": "anthropic/test-model", - "api_key": "fake-api-key", - }, - } - ] - ) - - # Call the appropriate generic_api_call_with_fallbacks method - if sync_mode: - response = router._generic_api_call_with_fallbacks( - model="test-model-alias", - original_function=mock_function, - messages=[{"role": "user", "content": "Hello"}], - max_tokens=100, - ) - else: - response = asyncio.run( - router._ageneric_api_call_with_fallbacks( - model="test-model-alias", - original_function=mock_function, - messages=[{"role": "user", "content": "Hello"}], - max_tokens=100, - ) - ) - - # Verify the mock function was called - mock_function.assert_called_once() - - # Verify the response - assert response == mock_response - - @pytest.mark.asyncio -async def test_aadapter_completion(): - """ - Test the aadapter_completion method which uses async_function_with_fallbacks - """ - - mock_response = { - "id": "adapter_resp_123", - "object": "adapter.completion", - "created": 1677858242, - "model": "test-model-with-adapter", - "choices": [ - { - "text": "This is a test adapter response", - "index": 0, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - } - - with patch.object(Router, "_aadapter_completion", new_callable=AsyncMock) as mock_method: - mock_method.return_value = mock_response - - router = Router( - model_list=[ - { - "model_name": "test-adapter-model", - "litellm_params": { - "model": "anthropic/test-model", - "api_key": "fake-api-key", - }, - } - ] - ) - - router.async_function_with_fallbacks = AsyncMock(return_value=mock_response) - - response = await router.aadapter_completion( - adapter_id="test-adapter-id", - model="test-adapter-model", - prompt="This is a test prompt", - max_tokens=100, - ) - - assert response == mock_response - - router.async_function_with_fallbacks.assert_called_once() - call_kwargs = router.async_function_with_fallbacks.call_args.kwargs - assert call_kwargs["adapter_id"] == "test-adapter-id" - assert call_kwargs["model"] == "test-adapter-model" - assert call_kwargs["prompt"] == "This is a test prompt" - assert call_kwargs["max_tokens"] == 100 - assert call_kwargs["original_function"] == router._aadapter_completion - assert "metadata" in call_kwargs - assert call_kwargs["metadata"]["model_group"] == "test-adapter-model" - - -@pytest.mark.asyncio -async def test__aadapter_completion(): - """ - Test the _aadapter_completion method directly - """ - - mock_response = { - "id": "adapter_resp_123", - "object": "adapter.completion", - "created": 1677858242, - "model": "test-model-with-adapter", - "choices": [ - { - "text": "This is a test adapter response", - "index": 0, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, - } +async def test_aadapter_completion_calls_the_adapter_with_the_deployments_params(): + mock_response = {"id": "adapter_resp_123", "object": "adapter.completion"} with patch("litellm.aadapter_completion", new_callable=AsyncMock) as mock_adapter_completion: mock_adapter_completion.return_value = mock_response @@ -409,43 +261,22 @@ async def test__aadapter_completion(): ] ) - router.async_get_available_deployment = AsyncMock( - return_value={ - "model_name": "test-adapter-model", - "litellm_params": { - "model": "test-model", - "api_key": "fake-api-key", - }, - "model_info": { - "id": "test-unique-id", - }, - } - ) - - router.async_routing_strategy_pre_call_checks = AsyncMock() - - response = await router._aadapter_completion( + response = await router.aadapter_completion( adapter_id="test-adapter-id", model="test-adapter-model", prompt="This is a test prompt", max_tokens=100, ) - assert response == mock_response - - mock_adapter_completion.assert_called_once() - call_kwargs = mock_adapter_completion.call_args.kwargs - assert call_kwargs["adapter_id"] == "test-adapter-id" - assert call_kwargs["model"] == "test-model" - assert call_kwargs["prompt"] == "This is a test prompt" - assert call_kwargs["max_tokens"] == 100 - assert call_kwargs["api_key"] == "fake-api-key" - assert call_kwargs["caching"] == router.cache_responses - - assert router.success_calls["test-model"] == 1 - assert router.total_calls["test-model"] == 1 - - router.async_routing_strategy_pre_call_checks.assert_called_once() + assert response["id"] == "adapter_resp_123" + mock_adapter_completion.assert_called_once() + call_kwargs = mock_adapter_completion.call_args.kwargs + assert call_kwargs["adapter_id"] == "test-adapter-id" + assert call_kwargs["model"] == "anthropic/test-model" + assert call_kwargs["prompt"] == "This is a test prompt" + assert call_kwargs["max_tokens"] == 100 + assert call_kwargs["api_key"] == "fake-api-key" + assert call_kwargs["metadata"]["model_group"] == "test-adapter-model" def test_initialize_router_endpoints(): @@ -482,101 +313,6 @@ def test_initialize_router_endpoints(): assert callable(router.adelete_responses) -@pytest.mark.asyncio -async def test_init_responses_api_endpoints(): - """ - A simpler test for _init_responses_api_endpoints that focuses on the basic functionality - """ - from litellm.responses.utils import ResponsesAPIRequestUtils - - # Create a router with a basic model - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/test-model", - "api_key": "fake-api-key", - }, - } - ] - ) - - # Just mock the _ageneric_api_call_with_fallbacks method - router._ageneric_api_call_with_fallbacks = AsyncMock() - - # Add a mock implementation of _get_model_id_from_response_id to the Router instance - ResponsesAPIRequestUtils.get_model_id_from_response_id = MagicMock( - return_value=None - ) - - # Call without a response_id (no model extraction should happen) - await router._init_responses_api_endpoints( - original_function=AsyncMock(), thread_id="thread_xyz" - ) - - # Verify _ageneric_api_call_with_fallbacks was called but model wasn't changed - first_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs - assert "model" not in first_call_kwargs - assert first_call_kwargs["thread_id"] == "thread_xyz" - - # Reset the mock - router._ageneric_api_call_with_fallbacks.reset_mock() - - # Change the return value for the second call - ResponsesAPIRequestUtils.get_model_id_from_response_id.return_value = ( - "claude-3-sonnet" - ) - - # Call with a response_id - await router._init_responses_api_endpoints( - original_function=AsyncMock(), response_id="resp_claude_123" - ) - - # Verify model was updated in the kwargs - second_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs - assert second_call_kwargs["model"] == "claude-3-sonnet" - assert second_call_kwargs["response_id"] == "resp_claude_123" - - -@pytest.mark.asyncio -async def test_init_vector_store_api_endpoints(): - """ - Test that _init_vector_store_api_endpoints correctly passes custom_llm_provider to kwargs - """ - - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/test-model", - "api_key": "fake-api-key", - }, - } - ] - ) - - mock_original_function = AsyncMock(return_value={"status": "success"}) - - result = await router._init_vector_store_api_endpoints( - original_function=mock_original_function, vector_store_id="test-store" - ) - - mock_original_function.assert_called_once_with(vector_store_id="test-store") - assert result == {"status": "success"} - - mock_original_function.reset_mock() - - await router._init_vector_store_api_endpoints( - original_function=mock_original_function, - custom_llm_provider="openai", - vector_store_id="test-store", - ) - - mock_original_function.assert_called_once_with(vector_store_id="test-store", custom_llm_provider="openai") - - @pytest.mark.asyncio async def test_avector_store_create_sends_only_the_callers_metadata_to_the_provider(): def echo_vector_store(request: httpx.Request) -> httpx.Response: @@ -612,10 +348,8 @@ async def test_avector_store_create_sends_only_the_callers_metadata_to_the_provi def test_apply_default_settings(): """ - Test the apply_default_settings method. - - This test verifies that apply_default_settings correctly initializes - default pre-call checks and doesn't modify existing router state. + apply_default_settings initializes the default pre-call checks and leaves the optional + callbacks the router already has in place. """ router = Router() @@ -637,16 +371,6 @@ def test_apply_default_settings(): assert mock_callback in router_with_callbacks.optional_callbacks - with patch.object(Router, "apply_default_settings") as mock_apply: - Router() - mock_apply.assert_called_once() - - router_test = Router() - with patch.object(router_test, "add_optional_pre_call_checks") as mock_add_checks: - router_test.apply_default_settings() - - mock_add_checks.assert_called_once_with([]) - def test_initialize_core_endpoints(): """ @@ -664,8 +388,6 @@ def test_initialize_core_endpoints(): ] ) - router._initialize_core_endpoints() - core_endpoints = [ "amoderation", "aanthropic_messages", @@ -679,7 +401,6 @@ def test_initialize_core_endpoints(): "acancel_responses", "adelete_responses", "alist_input_items", - "_arealtime", "acreate_fine_tuning_job", "acancel_fine_tuning_job", "alist_fine_tuning_jobs", @@ -710,8 +431,6 @@ def test_initialize_specialized_endpoints(): ] ) - router._initialize_specialized_endpoints() - specialized_endpoints = [ "avector_store_search", "avector_store_create", @@ -770,8 +489,6 @@ def test_initialize_vector_store_endpoints(): ] ) - router._initialize_vector_store_endpoints() - vector_store_endpoints = [ "avector_store_search", "avector_store_create", @@ -800,8 +517,6 @@ def test_initialize_vector_store_file_endpoints(): ] ) - router._initialize_vector_store_file_endpoints() - vector_store_file_endpoints = [ "avector_store_file_create", "vector_store_file_create", @@ -838,8 +553,6 @@ def test_initialize_google_genai_endpoints(): ] ) - router._initialize_google_genai_endpoints() - google_genai_endpoints = [ "agenerate_content", "generate_content", @@ -868,8 +581,6 @@ def test_initialize_ocr_search_endpoints(): ] ) - router._initialize_ocr_search_endpoints() - ocr_search_endpoints = [ "aocr", "ocr", @@ -898,8 +609,6 @@ def test_initialize_video_endpoints(): ] ) - router._initialize_video_endpoints() - video_endpoints = [ "avideo_generation", "video_generation", @@ -934,8 +643,6 @@ def test_initialize_container_endpoints(): ] ) - router._initialize_container_endpoints() - container_endpoints = [ "acreate_container", "create_container", @@ -968,8 +675,6 @@ def test_initialize_skills_endpoints(): ] ) - router._initialize_skills_endpoints() - skills_endpoints = [ "acreate_skill", "alist_skills", @@ -982,35 +687,34 @@ def test_initialize_skills_endpoints(): assert callable(getattr(router, endpoint)) -@pytest.mark.asyncio -async def test_init_containers_api_endpoints(): - """ - Test that _init_containers_api_endpoints calls the original function - directly when there is no managed container ID (no embedded model_id). - """ - router = Router(model_list=[]) +_CONTAINER_FILE: Final = { + "id": "cfile_xyz", + "object": "container.file", + "container_id": "cfile_upstream_abc", + "created_at": 0, + "bytes": 1, + "path": "/a", + "source": "user", +} - mock_response = {"id": "cntr_test", "name": "Test Container"} - mock_original_function = AsyncMock(return_value=mock_response) - result = await router._init_containers_api_endpoints( - original_function=mock_original_function, - custom_llm_provider="openai", - name="Test Container", - ) - - mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container") - assert result == mock_response +async def _container_file_request(router: Router, container_id: str) -> httpx.Request: + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route().mock(return_value=httpx.Response(200, json=_CONTAINER_FILE)) + await router.aretrieve_container_file(container_id=container_id, file_id="cfile_xyz") + assert route.call_count == 1 + return route.calls.last.request @pytest.mark.asyncio -async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallbacks(): +async def test_containers_managed_id_routes_to_its_deployment_with_its_credentials(monkeypatch): """ - Managed ``cntr_`` IDs embed ``model_id``; router should decode and use - ``_ageneric_api_call_with_fallbacks`` so deployment credentials apply. + Managed ``cntr_`` IDs embed ``model_id``; the router decodes it so the deployment's + credentials and api_base apply, and the upstream gets the unwrapped container ID. """ from litellm.responses.utils import ResponsesAPIRequestUtils + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) router = Router( model_list=[ { @@ -1023,31 +727,21 @@ async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallb } ] ) - router._ageneric_api_call_with_fallbacks = AsyncMock() - managed_id = ResponsesAPIRequestUtils.build_container_id( custom_llm_provider="azure", model_id="azure-router-model", container_id="cfile_upstream_abc", ) - await router._init_containers_api_endpoints( - original_function=AsyncMock(), - custom_llm_provider="openai", - container_id=managed_id, - file_id="cfile_xyz", - ) + request = await _container_file_request(router, managed_id) - router._ageneric_api_call_with_fallbacks.assert_called_once() - call_kw = router._ageneric_api_call_with_fallbacks.call_args.kwargs - assert call_kw["model"] == "azure-router-model" - assert call_kw["container_id"] == "cfile_upstream_abc" - assert call_kw["file_id"] == "cfile_xyz" - assert call_kw["custom_llm_provider"] == "azure" + assert request.url.host == "westus.api.cognitive.microsoft.com" + assert request.url.path.endswith("/containers/cfile_upstream_abc/files/cfile_xyz") + assert request.headers["api-key"] == "fake-key" @pytest.mark.asyncio -async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps(): +async def test_containers_managed_id_without_model_id_unwraps(monkeypatch): """ Managed ``cntr_`` IDs may be encoded with an empty ``model_id`` (e.g. when a streaming response had no router metadata). The router must still unwrap the @@ -1056,64 +750,22 @@ async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps """ from litellm.responses.utils import ResponsesAPIRequestUtils + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "sk-global") router = Router(model_list=[]) - mock_original_function = AsyncMock(return_value={"ok": True}) - managed_id = ResponsesAPIRequestUtils.build_container_id( custom_llm_provider="openai", model_id=None, container_id="cfile_upstream_abc", ) - await router._init_containers_api_endpoints( - original_function=mock_original_function, - custom_llm_provider="openai", - container_id=managed_id, - file_id="cfile_xyz", - ) + request = await _container_file_request(router, managed_id) - mock_original_function.assert_called_once() - call_kw = mock_original_function.call_args.kwargs - assert call_kw["container_id"] == "cfile_upstream_abc" - assert call_kw["file_id"] == "cfile_xyz" - assert call_kw["custom_llm_provider"] == "openai" + assert request.url.path == "/v1/containers/cfile_upstream_abc/files/cfile_xyz" @pytest.mark.asyncio -async def test_init_containers_api_endpoints_managed_id_without_model_id_applies_decoded_provider(): - """ - A managed ``cntr_`` ID can encode a non-OpenAI provider (e.g. ``azure``) with - an empty ``model_id`` (streaming events without router ``model_info.id``). - The router must still apply the decoded provider so the request routes to - the correct upstream — not stay on the default ``openai``. - """ - from litellm.responses.utils import ResponsesAPIRequestUtils - - router = Router(model_list=[]) - mock_original_function = AsyncMock(return_value={"ok": True}) - - managed_id = ResponsesAPIRequestUtils.build_container_id( - custom_llm_provider="azure", - model_id=None, - container_id="cfile_upstream_abc", - ) - - await router._init_containers_api_endpoints( - original_function=mock_original_function, - custom_llm_provider="openai", - container_id=managed_id, - file_id="cfile_xyz", - ) - - mock_original_function.assert_called_once() - call_kw = mock_original_function.call_args.kwargs - assert call_kw["container_id"] == "cfile_upstream_abc" - assert call_kw["file_id"] == "cfile_xyz" - assert call_kw["custom_llm_provider"] == "azure" - - -@pytest.mark.asyncio -async def test_init_containers_api_endpoints_create_with_model_uses_deployment_credentials(monkeypatch): +async def test_containers_create_with_model_uses_deployment_credentials(monkeypatch): """ ``POST /v1/containers`` carries no container ID, so a ``model`` in the request body is the only way to pick a deployment. The upstream call must receive that @@ -1121,29 +773,24 @@ async def test_init_containers_api_endpoints_create_with_model_uses_deployment_c ``OPENAI_API_KEY`` (which may be unset on the proxy). """ monkeypatch.delenv("OPENAI_API_KEY", raising=False) - router = Router( - model_list=[ - { - "model_name": "gpt-5.4", - "litellm_params": { - "model": "openai/gpt-5.4", - "api_key": "sk-model-list-key", - "api_base": "https://custom.openai.example/v1", - }, - } - ] - ) - mock_original_function = AsyncMock(return_value={"id": "cntr_test", "name": "Test Container"}) + with patch("litellm.containers.acreate_container", new_callable=AsyncMock) as create: + create.return_value = {"id": "cntr_test", "name": "Test Container"} + router = Router( + model_list=[ + { + "model_name": "gpt-5.4", + "litellm_params": { + "model": "openai/gpt-5.4", + "api_key": "sk-model-list-key", + "api_base": "https://custom.openai.example/v1", + }, + } + ] + ) + await router.acreate_container(name="Test Container", model="gpt-5.4") - await router._init_containers_api_endpoints( - original_function=mock_original_function, - custom_llm_provider="openai", - name="Test Container", - model="gpt-5.4", - ) - - mock_original_function.assert_called_once() - call_kw = mock_original_function.call_args.kwargs + create.assert_called_once() + call_kw = create.call_args.kwargs assert call_kw["api_key"] == "sk-model-list-key" assert call_kw["api_base"] == "https://custom.openai.example/v1" assert call_kw["model"] == "openai/gpt-5.4" @@ -1151,53 +798,40 @@ async def test_init_containers_api_endpoints_create_with_model_uses_deployment_c @pytest.mark.asyncio -async def test_init_containers_api_endpoints_create_without_model_calls_directly(): +async def test_containers_create_without_model_calls_directly(): """ - Without ``model`` (or with ``model=None`` as the proxy forwards it), create/list - must keep calling the handler directly with global provider credentials. + Without ``model``, create must keep calling the handler directly with global provider credentials. """ - router = Router(model_list=[]) - router._ageneric_api_call_with_fallbacks = AsyncMock() - mock_original_function = AsyncMock(return_value={"id": "cntr_test"}) + with patch("litellm.containers.acreate_container", new_callable=AsyncMock) as create: + create.return_value = {"id": "cntr_test"} + router = Router(model_list=[]) + await router.acreate_container(name="Test Container") - await router._init_containers_api_endpoints( - original_function=mock_original_function, - custom_llm_provider="openai", - name="Test Container", - model=None, - ) - - router._ageneric_api_call_with_fallbacks.assert_not_called() - mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None) + create.assert_called_once_with(name="Test Container") @pytest.mark.asyncio -async def test_init_containers_api_endpoints_create_with_unknown_model_passes_through(monkeypatch): +async def test_containers_create_with_unknown_model_passes_through(monkeypatch): """ A ``model`` that names no configured deployment must not turn into a 400. The call falls through to the handler with the caller's model and no injected deployment credentials, matching the behaviour before model-based routing existed. """ monkeypatch.delenv("OPENAI_API_KEY", raising=False) - router = Router( - model_list=[ - { - "model_name": "gpt-5.4", - "litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"}, - } - ] - ) - mock_original_function = AsyncMock(return_value={"id": "cntr_test"}) + with patch("litellm.containers.acreate_container", new_callable=AsyncMock) as create: + create.return_value = {"id": "cntr_test"} + router = Router( + model_list=[ + { + "model_name": "gpt-5.4", + "litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"}, + } + ] + ) + await router.acreate_container(name="Test Container", model="does-not-exist") - await router._init_containers_api_endpoints( - original_function=mock_original_function, - custom_llm_provider="openai", - name="Test Container", - model="does-not-exist", - ) - - mock_original_function.assert_called_once() - call_kw = mock_original_function.call_args.kwargs + create.assert_called_once() + call_kw = create.call_args.kwargs assert call_kw["model"] == "does-not-exist" assert call_kw["name"] == "Test Container" assert "api_key" not in call_kw @@ -1232,30 +866,11 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): try: callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is False - assert ( - encrypted_content_callbacks[0].model_group_affinity_config - == model_group_affinity_config - ) - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( - deployment_callback - ) - - router._add_encrypted_content_affinity_check(enable_global_affinity=True) - - callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] - assert len(encrypted_content_callbacks) == 1 - assert encrypted_content_callbacks[0].enable_global_affinity is True - assert encrypted_content_callbacks[0].ctx is router._core_ctx + assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback) finally: router.discard() diff --git a/tests/unit/test_router/test_router_entry_points.py b/tests/unit/test_router/test_router_entry_points.py index 25dd8fe47cf..c3d5d962123 100644 --- a/tests/unit/test_router/test_router_entry_points.py +++ b/tests/unit/test_router/test_router_entry_points.py @@ -332,15 +332,6 @@ def test_settings_snapshot_is_reused_until_a_setting_is_written(router): warnings.simplefilter("ignore") router.fallbacks = [{"gpt": ["team-gpt"]}] assert router.settings().fallbacks == [{"gpt": ["team-gpt"]}] - router._core_fallbacks = [{"gpt": ["other"]}] # the router's own writes invalidate it too - assert router.settings().fallbacks == [{"gpt": ["other"]}] - router.settings().fallbacks.append( - {"x": ["y"]} - ) # the snapshot's lists are copies: settings change only via apply() - assert router._core_fallbacks == [{"gpt": ["other"]}] - unrelated = router.settings() - router._core_model_list = list(router._core_model_list) - assert router.settings() is unrelated def test_deleting_a_setting_shows_as_unset(router): @@ -358,40 +349,6 @@ def test_settings_snapshot_is_frozen(router): assert set(SETTING_FIELDS) <= set(RouterSettings.__dataclass_fields__) -def test_unset_setting_raises_attribute_error_through_old_name(): - r = _router() - try: - del r._core.__dict__["_core_kubernetes_pod_discovery"] - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - assert getattr(r, "kubernetes_pod_discovery", "missing") == "missing" - assert r.settings().is_set("kubernetes_pod_discovery") is False - finally: - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - r.retire() - - -def test_unset_attribute_error_names_the_public_attribute(router): - # Reading (or deleting) an old attribute the router never set raises what reading it raised before. - router._core.__dict__.pop("_core_slack_alerting_logger", None) - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - with pytest.raises(AttributeError) as read: - _ = router.slack_alerting_logger - with pytest.raises(AttributeError) as deleted: - del router.slack_alerting_logger - del router._core.__dict__["_core_kubernetes_pod_discovery"] - with pytest.raises(AttributeError) as setting: - _ = router.kubernetes_pod_discovery - for caught, name in ((read, "slack_alerting_logger"), (setting, "kubernetes_pod_discovery")): - assert str(caught.value) == f"'Router' object has no attribute '{name}'" - assert caught.value.name == name - assert caught.value.obj is router - # A failed del sets no name/obj in Python either. - assert str(deleted.value) == "'Router' object has no attribute 'slack_alerting_logger'" - - def test_bind_to_request_logging_none_raises_as_before(router): # The old member bound unconditionally, so None fails on the first callback registration. with warnings.catch_warnings(): @@ -456,17 +413,7 @@ def test_router_cannot_be_subclassed(): def test_core_classes_can_still_be_subclassed(): - from litellm.router_python_core import PythonCore - - class PythonSub(PythonCore): - pass - - assert isinstance(object.__new__(PythonSub), Router) - - try: - from litellm.router_rust.core import RustCore - except ImportError: - pytest.skip("the native lib does not load") + from litellm.router_rust.core import RustCore class RustSub(RustCore): pass @@ -490,15 +437,6 @@ def test_alias_attribute_writes_the_same_storage(router): assert router.registry().routing_groups() is replacement -def test_core_members_keep_their_old_names_in_errors(): - from litellm.router import _name_core_members_after_old_names - - _name_core_members_after_old_names() - - assert Router._core_get_deployment_model_info.__qualname__ == "Router.get_deployment_model_info" - assert Router._core_get_deployment.__qualname__ == "Router.get_deployment" - - def test_wrong_arity_errors_still_name_router(router): with warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -510,177 +448,10 @@ def test_wrong_arity_errors_still_name_router(router): router.has_model_id(**None) # pyright: ignore[reportCallIssue, reportArgumentType] with pytest.raises(TypeError, match=r"^Router\.is_deployment_blocked\(\) missing 1 required"): router.is_deployment_blocked() # pyright: ignore[reportCallIssue] - with pytest.raises(TypeError, match=r"^Router\.get_deployment_model_info\(\) missing 2 required"): - router.cached_deployment_model_info() # pyright: ignore[reportCallIssue] - - -_SENTINEL = object() - -# (entry point call, Router core member it forwards to, the arguments the core member must get) -FORWARDS = [ - ( - lambda r: r.registry().deployment_model_info("dep-a", "gpt", cached=True), - "_core_cached_deployment_model_info", - (("dep-a", "gpt"), {}), - ), - ( - lambda r: r.registry().deployment_model_info("dep-a", "gpt"), - "_core_get_deployment_model_info", - ((), {"model_id": "dep-a", "model_name": "gpt"}), - ), - ( - lambda r: r.registry().router_model_info(None, "gpt", id="dep-a"), - "_core_get_router_model_info", - ((), {"deployment": None, "received_model_name": "gpt", "id": "dep-a"}), - ), - (lambda r: r.registry().discovered_model_info("dep-a"), "_core_get_discovered_model_info", (("dep-a",), {})), - ( - lambda r: r.registry().deployments("gpt", source="routing_groups"), - "_core_get_model_list_from_routing_groups", - ((), {"model_name": "gpt"}), - ), - ( - lambda r: r.registry().deployments("gpt", source="routed_group", team_id="t"), - "_core_get_model_list_of_routed_group", - (("gpt", "t"), {}), - ), - ( - lambda r: r.registry().deployments("gpt", source="wildcard", team_id="t"), - "_core_wildcard_route_deployments", - (("gpt", "t"), {}), - ), - ( - lambda r: r.registry().deployments("gpt", source="routable", request={"a": 1}, call_type="aembedding"), - "_core_get_routable_deployments", - (("gpt", {"a": 1}, "aembedding"), {}), - ), - ( - lambda r: r.registry().deployments("gpt", source="for_request"), - "_core_deployments_for_request", - (("gpt", {}), {}), - ), - ( - lambda r: r.registry().includes({"model_name": "gpt"}, model_name="gpt", team_id="t"), - "_core_should_include_deployment", - ((), {"model_name": "gpt", "model": {"model_name": "gpt"}, "team_id": "t"}), - ), - ( - lambda r: r.registry().names("gpt", kind="litellm_models", team_id="t"), - "_core_resolved_litellm_models", - (("gpt",), {"team_id": "t"}), - ), - ( - lambda r: r.registry().names("gpt", kind="litellm_base_models"), - "_core_resolved_litellm_models", - (("gpt",), {"team_id": None, "use_base_model": True}), - ), - ( - lambda r: r.registry().names("gpt", kind="wildcard", team_id="t"), - "_core_wildcard_routed_models", - (("gpt", "t"), {}), - ), - ( - lambda r: r.registry().names("gpt", kind="route_candidates"), - "_core_get_candidate_model_ids_for_route", - (("gpt",), {"team_id": None}), - ), - ( - lambda r: r.registry().resolve("dep-a", kind="from_id", custom_llm_provider="openai"), - "_core_resolve_model_name_from_model_id", - (("dep-a",), {"custom_llm_provider": "openai"}), - ), - ( - lambda r: r.registry().resolve("rg", kind="routing_group_alternatives"), - "_core_routing_group_has_alternatives", - (("rg",), {}), - ), - ( - lambda r: r.registry().resolve("gpt", kind="served", served_model="gpt-4o"), - "_core_resolve_served_model_name", - (("gpt", "gpt-4o"), {}), - ), - ( - lambda r: r.registry().resolve("gpt", kind="upstream", team_id="t"), - "_core_get_routable_upstream_model", - (("gpt",), {"team_id": "t"}), - ), - ( - lambda r: r.registry().resolve("gpt", kind="search_tool_support", team_id="t"), - "_core_get_search_tool_support", - (("gpt",), {"team_id": "t"}), - ), - ( - lambda r: r.registry().resolve("grp", kind="wildcard_access_group"), - "_core_is_model_access_group_for_wildcard_route", - ((), {"model_access_group": "grp"}), - ), - ( - lambda r: r.registry().resolve("dep-a", kind="team_model_alternatives"), - "_core_team_model_has_alternatives", - (("dep-a",), {}), - ), - ( - lambda r: r.registry().group("gpt", include_hidden=True, team_id="t"), - "_core_get_model_group_info", - (("gpt",), {"include_hidden": True, "team_id": "t"}), - ), - (lambda r: r.registry().group("gpt", field="cached_info"), "_core_cached_model_group_info", (("gpt",), {})), - (lambda r: r.registry().group("gpt", field="display_name"), "_core_get_configured_display_name", (("gpt",), {})), - ( - lambda r: r.registry().group("gpt", field="service_tiers", team_id="t"), - "_core_get_configured_service_tiers", - (("gpt",), {"team_id": "t"}), - ), - (lambda r: r.registry().group("gpt", field="listing_info"), "_core_get_model_listing_info", (("gpt",), {})), - (lambda r: r.registry().plugins("budget_limiter"), "_core_get_router_deployment_budget_limiter", ((), {})), - (lambda r: r.settings().allowed_fails_for(ValueError("x")), "_core_get_allowed_fails_from_policy", None), - (lambda r: r.settings().guardrail("g"), "_core_get_available_guardrail", (("g",), {})), - (lambda r: r.health().team_ptu_ceiling("t", "gpt"), "_core_team_ptu_ceiling", (("t", "gpt"), {})), - ( - lambda r: r.health().split_priority("gpt", {"priority": 1}), - "_core_split_request_priority", - (("gpt", {"priority": 1}), {}), - ), -] - - -@pytest.mark.parametrize("call, core, expected", FORWARDS, ids=[f[1] for f in FORWARDS]) -def test_entry_point_forwards_to_its_core_member(router, monkeypatch, call, core, expected): - from unittest.mock import MagicMock - - mock = MagicMock(return_value=_SENTINEL) - monkeypatch.setattr(router, core, mock) - assert call(router) is _SENTINEL - if expected is not None: - assert mock.call_args == expected - else: - mock.assert_called_once() - - -ASYNC_FORWARDS = [ - (lambda r: r.health().fully_unhealthy_model_names(), "_core_async_get_fully_unhealthy_model_names", ((), {})), - (lambda r: r.health().usage("gpt"), "_core_get_model_group_usage", ((), {"model_group": "gpt"})), - ( - lambda r: r.health().usage("gpt", kind="remaining"), - "_core_get_remaining_model_group_usage", - ((), {"model_group": "gpt"}), - ), -] - - -@pytest.mark.parametrize("call, core, expected", ASYNC_FORWARDS, ids=[f[1] for f in ASYNC_FORWARDS]) -def test_async_entry_point_forwards_to_its_core_member(router, monkeypatch, call, core, expected): - from unittest.mock import AsyncMock - - mock = AsyncMock(return_value=_SENTINEL) - monkeypatch.setattr(router, core, mock) - assert asyncio.run(call(router)) is _SENTINEL - assert mock.call_args == expected def test_registry_reads_without_a_core_forward(router): registry = router.registry() - assert registry.zero_cost() is getattr(router, "_zero_cost_cache", None) assert registry.includes(router.registry().deployment("dep-a").model_dump(), team_id=None) is True @@ -727,30 +498,13 @@ def test_old_members_refuse_objects_that_are_not_routers(): _as_router(object()) # pyright: ignore[reportArgumentType] -def test_routing_context_forwards_to_the_router(router, monkeypatch): - from unittest.mock import AsyncMock, MagicMock - +def test_routing_context_forwards_to_the_router(router): from litellm.router_api import RoutingContext ctx = RoutingContext(router) assert ctx.registry().version == router.registry().version assert ctx.settings() is router.settings() assert isinstance(ctx.health(), type(router.health())) - for name in ( - "_pre_call_checks", - "_count_pre_call_check_tokens", - "_update_kwargs_before_fallbacks", - "_core_arm_routing_read_prefetch", - "_core_log_retry", - ): - monkeypatch.setattr(router, name, MagicMock(return_value=_SENTINEL)) - monkeypatch.setattr(router, "_core_async_pre_routing_hook", AsyncMock(return_value=_SENTINEL)) - assert ctx.pre_call_checks("gpt", []) is _SENTINEL - assert ctx.count_pre_call_check_tokens(None, "hi") is _SENTINEL - assert ctx.update_kwargs_before_fallbacks("gpt", {}) is _SENTINEL - assert ctx.arm_read_prefetch("gpt") is _SENTINEL - assert asyncio.run(ctx.pre_routing_hook("gpt", {})) is _SENTINEL - assert ctx.log_retry({}, ValueError("x")) is _SENTINEL def test_attribute_errors_about_other_objects_pass_through_unchanged(router): diff --git a/tests/unit/test_router/test_router_facade.py b/tests/unit/test_router/test_router_facade.py index 0539ebe43b9..07d85d36e47 100644 --- a/tests/unit/test_router/test_router_facade.py +++ b/tests/unit/test_router/test_router_facade.py @@ -1,12 +1,10 @@ """ The ``litellm.Router`` facade's own helpers in litellm/router.py: how it binds the constructor -arguments, which members it leaves to the facade, and the module attributes it forwards to the core. +arguments and which members it leaves to the facade. """ import pytest -import litellm.router as router_module -from litellm import router_python_core from litellm.router import Router, _class_descriptor, router_arguments @@ -41,12 +39,3 @@ def test_class_descriptor_is_true_for_a_data_descriptor_on_the_class(): def test_class_descriptor_is_false_for_a_plain_method_or_a_missing_name(): assert not _class_descriptor(Router, "from_config") assert not _class_descriptor(Router, "no_such_member") - - -def test_module_getattr_forwards_old_helpers_to_the_core_module(): - assert router_module.__getattr__("_stream_pre_content_error") is router_python_core._stream_pre_content_error - - -def test_module_getattr_raises_attribute_error_for_an_unknown_name(): - with pytest.raises(AttributeError): - router_module.__getattr__("no_such_router_helper") diff --git a/tests/unit/test_router/test_router_fallbacks.py b/tests/unit/test_router/test_router_fallbacks.py index b25d5c35892..937f3f40e5b 100644 --- a/tests/unit/test_router/test_router_fallbacks.py +++ b/tests/unit/test_router/test_router_fallbacks.py @@ -1176,7 +1176,7 @@ def test_context_stream_fallback_obeys_all_output_receipts( wrapped: bool, has_output: bool, generated_content: str, disqualifying: bool, recoverable: bool ) -> None: from litellm.exceptions import MidStreamFallbackError - from litellm.router import _context_window_stream_fallback_error + from litellm.router_rust.invoke.streaming import _context_window_stream_fallback_error # noqa: PLC2701 original: Final = litellm.ContextWindowExceededError(message="too long", model="primary", llm_provider="openai") error: Final = ( @@ -1201,7 +1201,7 @@ def test_context_stream_fallback_obeys_all_output_receipts( def test_context_stream_fallback_preserves_other_errors_for_existing_continuation_policy() -> None: from litellm.exceptions import MidStreamFallbackError - from litellm.router import _context_window_stream_fallback_error + from litellm.router_rust.invoke.streaming import _context_window_stream_fallback_error # noqa: PLC2701 error: Final = MidStreamFallbackError( message="provider unavailable", model="primary", llm_provider="openai", generated_content="already visible" diff --git a/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py b/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py new file mode 100644 index 00000000000..8551e6ea884 --- /dev/null +++ b/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py @@ -0,0 +1,515 @@ +""" +Router behavior the router-code differential fuzz run found changed once the Rust core became the +only core: count settings below zero, an allowed_fails_policy write that is no policy, a model +response object as a deployment's mock_response, and endpoints read or patched on the Router class. +""" + +import asyncio +import gc +import json +import socket +import threading +import weakref +from collections.abc import Iterator, Mapping +from datetime import datetime +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final +from unittest import mock + +import pytest + +import litellm +from litellm import Router +from litellm.router import RetryPolicy +from litellm.types.router import Deployment +from litellm.types.utils import ModelResponse + +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +def _deployment(name: str, reply: object) -> dict[str, object]: + return {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "mock_response": reply}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("setting", ["num_retries", "max_fallbacks", "allowed_fails"]) +async def test_a_negative_count_setting_builds_and_serves(setting: str) -> None: + router: Final = Router(model_list=[_deployment("g", "ok")], **{setting: -1}) + + response: Final = await router.acompletion(model="g", messages=MESSAGES) + + assert response.choices[0].message.content == "ok" + + +@pytest.mark.asyncio +async def test_a_negative_max_fallbacks_takes_no_fallback() -> None: + router: Final = Router( + model_list=[_deployment("g", Exception("boom")), _deployment("fb", "fb")], + fallbacks=[{"g": ["fb"]}], + max_fallbacks=-1, + num_retries=0, + ) + + with pytest.raises(litellm.InternalServerError): + await router.acompletion(model="g", messages=MESSAGES) + + +@pytest.mark.asyncio +async def test_a_negative_num_retries_still_falls_back() -> None: + router: Final = Router( + model_list=[_deployment("g", Exception("boom")), _deployment("fb", "fb")], + fallbacks=[{"g": ["fb"]}], + num_retries=-1, + ) + + response: Final = await router.acompletion(model="g", messages=MESSAGES) + + assert response.choices[0].message.content == "fb" + + +@pytest.mark.asyncio +async def test_an_allowed_fails_policy_written_as_a_list_reads_as_no_policy() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")]) + + with pytest.warns(DeprecationWarning, match="allowed_fails_policy"): + router.allowed_fails_policy = [] # pyright: ignore[reportAttributeAccessIssue] # the old attribute write + response: Final = await router.acompletion(model="g", messages=MESSAGES) + + assert response.choices[0].message.content == "ok" + assert router.settings().allowed_fails_policy == [] + + +def test_a_model_response_mock_reads_back_from_the_model_list_as_a_dict() -> None: + reply: Final = ModelResponse(id="chatcmpl-fixed", choices=[{"message": {"role": "assistant", "content": "fb"}}]) + router: Final = Router(model_list=[_deployment("g", reply)]) + + params: Final = router.get_model_list()[0]["litellm_params"] + + assert isinstance(params["mock_response"], dict) + assert params["mock_response"]["id"] == "chatcmpl-fixed" + + +def test_an_endpoint_the_core_serves_by_name_is_readable_on_the_class() -> None: + endpoint: Final = Router.aretrieve_batch + + assert callable(endpoint) + assert endpoint is Router.aretrieve_batch + assert endpoint.__name__ == "aretrieve_batch" + + +@pytest.mark.asyncio +async def test_a_class_level_patch_of_an_endpoint_reaches_the_instance_and_unwinds() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")]) + + with mock.patch.object(Router, "asearch", new=mock.AsyncMock(return_value="patched")) as patched: + assert await router.asearch(query="q", search_tool_name="t") == "patched" + patched.assert_awaited_once() + + assert router.asearch is not patched + + +def test_an_assigned_affinity_ttl_builds_the_affinity_check_added_afterwards() -> None: + from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck + + router: Final = Router(model_list=[_deployment("g", "ok")]) + + router.deployment_affinity_ttl_seconds = 7 + router.add_optional_pre_call_checks(["deployment_affinity"]) + + checks: Final = [c for c in litellm.callbacks if isinstance(c, DeploymentAffinityCheck)] + assert [check.ttl_seconds for check in checks][-1:] == [7] + + +@pytest.mark.asyncio +async def test_a_deleted_timeout_leaves_get_settings_and_still_serves() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0) + + del router.timeout + + assert "timeout" not in router.get_settings() + response: Final = await router.acompletion(model="g", messages=MESSAGES) + assert response.choices[0].message.content == "ok" + + +def test_litellm_params_json_cannot_carry_read_back_as_given_and_hash_like_python() -> None: + from datetime import datetime + + from litellm.router_api import deployment_id + + params: Final = {"model": "openai/gpt-4o", "api_key": "k", "cfg": {"s": {"a"}}, "when": datetime(2026, 1, 1, 12)} + router: Final = Router(model_list=[{"model_name": "g", "litellm_params": dict(params)}]) + + served: Final = router.get_model_list()[0] + + assert served["litellm_params"]["cfg"] == {"s": {"a"}} + assert served["litellm_params"]["when"] == datetime(2026, 1, 1, 12) + assert served["model_info"]["id"] == deployment_id("g", params) + + +def test_a_bad_deployment_before_a_secret_reference_raises_its_own_validation_error() -> None: + from pydantic import ValidationError + + with pytest.raises(ValidationError): + Router( + model_list=[ + {"model_name": "t", "litellm_params": {"api_key": "k"}}, + {"model_name": "u", "litellm_params": {"model": "openai/gpt-4o", "api_key": "os.environ/NOPE_KEY"}}, + ] + ) + + +@pytest.mark.asyncio +async def test_a_failed_vector_store_create_names_no_failed_deployment() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "vs", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": "http://127.0.0.1:9/v1"}, + } + ], + num_retries=0, + ) + + with pytest.raises(Exception) as raised: # noqa: PT011 # the provider's connection error, whatever class maps it + await router.avector_store_create(model="vs", name="x") + + assert getattr(raised.value, "failed_deployment_id", None) is None + + +class _Upstream(BaseHTTPRequestHandler): + """An OpenAI upstream whose path's first segment names how it answers: ``cut`` sends two + Responses stream events then cuts the chunked body, ``filtered`` answers a chat completion the + provider blocked, ``down`` answers 502, and anything else a plain chat completion.""" + + protocol_version = "HTTP/1.1" + + def log_message(self, format: str, *args: object) -> None: # noqa: A002 # the base class's name + del format, args + + def _json(self, status: int, body: Mapping[str, object]) -> None: + data: Final = json.dumps(body).encode() + self.send_response(status) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(data))) + self.end_headers() + _ = self.wfile.write(data) + + def _chunk(self, event: Mapping[str, object]) -> None: + data: Final = f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + _ = self.wfile.write(f"{len(data):x}\r\n".encode() + data + b"\r\n") + self.wfile.flush() + + def _cut(self) -> None: + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("transfer-encoding", "chunked") + self.end_headers() + response: Final = {"id": "r1", "object": "response", "created_at": 1, "status": "in_progress", "output": []} + self._chunk({"type": "response.created", "sequence_number": 0, "response": {**response, "model": "gpt-4o"}}) + self._chunk( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "m1", + "output_index": 0, + "content_index": 0, + "delta": "hello ", + } + ) + self.connection.shutdown(socket.SHUT_RDWR) + self.close_connection = True + + def do_POST(self) -> None: + _ = self.rfile.read(int(self.headers.get("content-length") or 0)) + mode: Final = self.path.split("/")[1] + if mode == "cut": + self._cut() + return + if mode == "down": + self._json(502, {"error": {"message": "provider error", "type": "api_error"}}) + return + finish: Final = "content_filter" if mode == "filtered" else "stop" + message: Final = {"role": "assistant", "content": mode} + choice: Final = {"index": 0, "finish_reason": finish, "message": message} + usage: Final = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + self._json(200, {"id": "c1", "object": "chat.completion", "created": 1, "choices": [choice], "usage": usage}) + + +@pytest.fixture +def upstream() -> Iterator[str]: + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), _Upstream) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield f"http://127.0.0.1:{server.server_address[1]}" + server.shutdown() + thread.join() + + +def _served(name: str, deployment_id: str, model: str, api_base: str) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": model, "api_key": "k", "api_base": api_base}, + "model_info": {"id": deployment_id}, + } + + +def _request_logging(call_type: str, *, stream: bool) -> object: + from litellm.litellm_core_utils.litellm_logging import Logging + + return Logging( + model="g", + messages=MESSAGES, + stream=stream, + call_type=call_type, + start_time=datetime.now(), # noqa: DTZ005 # litellm's own naive start time + litellm_call_id="call-1", + function_id="f", + ) + + +async def _recorded() -> None: + """Let the background recordings of a stream's outcome run.""" + for _ in range(50): + await asyncio.sleep(0.02) + + +@pytest.mark.asyncio +async def test_a_responses_stream_cut_after_it_opened_counts_one_fail_with_a_shared_logging_object( + upstream: str, +) -> None: + router: Final = Router( + model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/cut/v1")], num_retries=0, disable_cooldowns=True + ) + + stream: Final = await router.aresponses( + model="g", input="hi", stream=True, litellm_logging_obj=_request_logging("aresponses", stream=True) + ) + with pytest.raises(Exception): # noqa: B017, PT011 # the transport's read error, whatever class maps it + async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True + pass + await _recorded() + + assert router.health().store.get_cache(key="d1:fails") == 1 + + +@pytest.mark.asyncio +async def test_a_blocked_response_raised_for_a_fallback_counts_no_success_with_a_shared_logging_object( + upstream: str, +) -> None: + router: Final = Router( + model_list=[ + _served("a", "a1", "openai/gpt-4o", f"{upstream}/filtered/v1"), + _served("safe", "s1", "openai/gpt-4o", f"{upstream}/ok/v1"), + ], + content_policy_fallbacks=[{"a": ["safe"]}], + num_retries=0, + disable_cooldowns=True, + ) + + response: Final = await router.acompletion( + model="a", messages=MESSAGES, litellm_logging_obj=_request_logging("acompletion", stream=False) + ) + await _recorded() + + assert response.choices[0].message.content == "ok" + assert router.health().store.get_cache(key="a1:successes") is None + assert router.health().store.get_cache(key="s1:successes") == 1 + + +@pytest.mark.asyncio +async def test_a_text_completion_stream_that_fails_before_its_first_chunk_counts_nothing(upstream: str) -> None: + router: Final = Router( + model_list=[_served("g", "d1", "openai/gpt-3.5-turbo-instruct", f"{upstream}/down/v1")], + num_retries=0, + disable_cooldowns=True, + ) + + stream: Final = await router.atext_completion(model="g", prompt="hi", stream=True) + with pytest.raises(litellm.BadGatewayError): + async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True + pass + await _recorded() + + assert router.health().store.get_cache(key="d1:successes") is None + assert router.health().store.get_cache(key="d1:fails") is None + + +@pytest.mark.asyncio +async def test_a_text_completion_without_a_prompt_raises_type_error_before_routing() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0) + + with pytest.raises(TypeError, match="'prompt'"): + await router.atext_completion(model="g") # pyright: ignore[reportCallIssue] # the missing argument + with pytest.raises(TypeError, match="'prompt'"): + router.text_completion(model="g") # pyright: ignore[reportCallIssue] # the missing argument + + +@pytest.mark.asyncio +async def test_a_default_priority_assigned_after_construction_reads_back_and_queues_calls() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0) + + router.default_priority = 5 + response: Final = await router.acompletion(model="g", messages=MESSAGES) + + assert router.default_priority == 5 + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + + +@pytest.mark.asyncio +async def test_a_prioritized_call_on_an_unknown_group_raises_the_calls_own_error() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0) + + with pytest.raises(litellm.InternalServerError): + await router.schedule_acompletion( + model="nope", messages=MESSAGES, priority=0, mock_testing_fallbacks=True, num_retries=0 + ) + + +def test_a_registry_lookup_of_a_record_with_no_model_name_raises_key_error() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")]) + + with pytest.raises(KeyError): + router.registry().includes({"litellm_params": {}}, model_name="g") # pyright: ignore[reportArgumentType] # the record with no name + + +@pytest.mark.asyncio +async def test_a_refused_pick_by_deployment_id_names_the_deployments_group() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "t2-pub", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k", "mock_response": "p2"}, + "model_info": {"team_id": "T2", "team_public_model_name": "pub", "id": "p2"}, + } + ], + num_retries=0, + ) + + with pytest.raises(Exception, match="model=t2-pub"): # noqa: PT011 # the refusal, whatever class names it + await router.async_get_available_deployment(model="p2", messages=MESSAGES, request_kwargs={}) + + +@pytest.mark.asyncio +async def test_a_request_refused_before_the_provider_counts_against_no_deployment(upstream: str) -> None: + router: Final = Router( + model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/ok/v1")], num_retries=0, disable_cooldowns=True + ) + + with pytest.raises(litellm.UnsupportedParamsError): + await router.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + await _recorded() + + assert router.health().store.get_cache(key="d1:fails") is None + assert not [key for key in router.health().store.in_memory_cache.cache_dict if ":rpm" in key] + + +@pytest.mark.asyncio +async def test_the_async_pick_read_on_the_class_is_a_function_taking_the_router() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")]) + + pick: Final = Router.async_get_available_deployment + + assert pick.__name__ == "async_get_available_deployment" + assert (await pick(router, "g", {}))["model_name"] == "g" + + +def _sharing_ids(*ids: str) -> Router: + return Router( + model_list=[{**_deployment("g", f"from-{index}"), "model_info": {"id": id}} for index, id in enumerate(ids)] + ) + + +@pytest.mark.asyncio +async def test_async_picks_split_over_deployments_sharing_an_id() -> None: + router: Final = _sharing_ids("d1", "d1", "b") + + served: Final = { + (await router.acompletion(model="g", messages=MESSAGES)).choices[0].message.content for _ in range(60) + } + + assert served == {"from-0", "from-1", "from-2"} + + +def test_sync_picks_split_over_deployments_sharing_an_id() -> None: + router: Final = _sharing_ids("d1", "d1") + + served: Final = {router.completion(model="g", messages=MESSAGES).choices[0].message.content for _ in range(40)} + + assert served == {"from-0", "from-1"} + + +@pytest.mark.asyncio +async def test_no_model_is_served_by_the_first_default_fallback() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")], default_fallbacks=["g"]) + + async_served: Final = await router.acompletion(model=None, messages=MESSAGES) + sync_served: Final = router.completion(model=None, messages=MESSAGES) + + assert (async_served.choices[0].message.content, sync_served.choices[0].message.content) == ("ok", "ok") + + +@pytest.mark.asyncio +async def test_no_model_is_served_by_a_catch_all_wildcard() -> None: + router: Final = Router( + model_list=[ + _deployment("g", "ok"), + { + **_deployment("*", "star"), + "litellm_params": {"model": "openai/*", "api_key": "k", "mock_response": "star"}, + }, + ] + ) + + served: Final = await router.acompletion(model=None, messages=MESSAGES) + + assert served.choices[0].message.content == "star" + + +@pytest.mark.asyncio +async def test_no_model_is_retried_by_the_retry_policy_and_raised_without_a_fallback_note() -> None: + router: Final = Router( + model_list=[_deployment("g", "ok")], + default_fallbacks=["h"], + retry_after=0, + retry_policy=RetryPolicy(BadRequestErrorRetries=1), + ) + + with pytest.raises(litellm.BadRequestError) as raised: + await router.acompletion(model=None, messages=MESSAGES) + + assert raised.value.num_retries == 1 + assert "You passed in model=h. There are no healthy deployments for this model" in str(raised.value) + assert "failed with the error above" not in str(raised.value) + + +class _TokenProvider: + def __call__(self) -> str: + return "token" + + +def test_a_deleted_deployment_releases_the_callable_it_was_built_with() -> None: + router: Final = Router(model_list=[_deployment("g", "ok")]) + provider = _TokenProvider() + released: Final = weakref.ref(provider) + router.add_deployment( + Deployment( + model_name="g", + litellm_params={"model": "openai/gpt-4o", "api_key": "k", "azure_ad_token_provider": provider}, + model_info={"id": "rotated"}, + ) + ) + del provider + + router.delete_deployment("rotated") + gc.collect() + + assert released() is None + + +@pytest.mark.asyncio +async def test_a_pick_keeps_model_info_keys_that_are_not_strings() -> None: + router: Final = Router(model_list=[{**_deployment("g", "ok"), "model_info": {"id": "d1", "tiers": {1: "a"}}}]) + + picked: Final = await router.async_get_available_deployment(model="g", messages=MESSAGES, request_kwargs={}) + + assert picked["model_info"]["tiers"] == {1: "a"} diff --git a/tests/unit/test_router/test_router_get_deployments.py b/tests/unit/test_router/test_router_get_deployments.py index 07486762a64..acfcb999c36 100644 --- a/tests/unit/test_router/test_router_get_deployments.py +++ b/tests/unit/test_router/test_router_get_deployments.py @@ -213,66 +213,6 @@ async def test_async_get_available_deployment_for_pass_through(): pytest.fail(f"Error occurred: {e}") -def test_filter_pass_through_deployments(): - """ - Test _filter_pass_through_deployments function - - Tests that it correctly filters deployments with use_in_pass_through=True - """ - try: - litellm.set_verbose = False - model_list = [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), - "use_in_pass_through": True, - }, - }, - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), - "api_version": os.getenv("AZURE_API_VERSION"), - "use_in_pass_through": False, - }, - }, - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "azure/gpt-35-turbo", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), - "api_version": os.getenv("AZURE_API_VERSION"), - "use_in_pass_through": True, - }, - }, - ] - router = Router( - model_list=model_list, - ) - - # Get all healthy deployments - healthy_deployments = router.get_model_list() - - # Filter pass-through deployments - pass_through_deployments = router._filter_pass_through_deployments(healthy_deployments) - - # Should only have 2 deployments with use_in_pass_through=True - assert len(pass_through_deployments) == 2 - - # Verify all returned deployments have use_in_pass_through=True - for deployment in pass_through_deployments: - assert deployment["litellm_params"]["use_in_pass_through"] is True - - router.reset() - except Exception as e: - traceback.print_exc() - pytest.fail(f"Error occurred: {e}") - - async def _weighted_async_deployment_ids(router: Router) -> tuple[str, ...]: random.seed(2025) deployments: Final = tuple( diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 3bbb69413d9..32b38030498 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1,21 +1,18 @@ import asyncio -import functools import gc import json import os -import uuid +import time import weakref -from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from datetime import datetime, timezone -from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from websockets.exceptions import ConnectionClosed -from websockets.frames import Close +from pydantic import BaseModel import litellm from litellm import Router @@ -24,18 +21,11 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.litellm_core_utils.realtime_errors import close_after_upstream_handshake_refusal -from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig -from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook -from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.litellm_params import RoutingStrategyName from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( - CallTypes, ModelResponse, StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -208,29 +198,6 @@ def test_print_deployment_with_redact_enabled(model_list): litellm.redact_user_api_key_info = original_setting -def test_completion(model_list): - """Test if the completion function is working correctly""" - router = Router(model_list=model_list) - response = router._completion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hello, how are you?"}], - mock_response="I'm fine, thank you!", - ) - assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!" - - -@pytest.mark.asyncio -async def test_router_acompletion_util(model_list): - """Test if the underlying '_acompletion' function is working correctly""" - router = Router(model_list=model_list) - response = await router._acompletion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hello, how are you?"}], - mock_response="I'm fine, thank you!", - ) - assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!" - - @pytest.mark.asyncio async def test_router_abatch_completion_one_model_multiple_requests_util(model_list): """Test if the 'abatch_completion_one_model_multiple_requests' function is working correctly""" @@ -261,193 +228,6 @@ async def test_router_schedule_acompletion(model_list): assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!" -@pytest.mark.asyncio -async def test_router_schedule_atext_completion(model_list): - """Test if the 'schedule_atext_completion' function is working correctly""" - from litellm.types.utils import TextCompletionResponse - - router = Router(model_list=model_list) - with patch.object(router, "_atext_completion", AsyncMock()) as mock_atext_completion: - mock_atext_completion.return_value = TextCompletionResponse() - response = await router.atext_completion( - model="gpt-5-mini", - prompt="Hello, how are you?", - priority=1, - ) - mock_atext_completion.assert_awaited_once() - assert "priority" not in mock_atext_completion.call_args.kwargs - - -@pytest.mark.asyncio -async def test_router_schedule_factory(model_list): - """Test if the 'schedule_atext_completion' function is working correctly""" - from litellm.types.utils import TextCompletionResponse - - router = Router(model_list=model_list) - with patch.object(router, "_atext_completion", AsyncMock()) as mock_atext_completion: - mock_atext_completion.return_value = TextCompletionResponse() - response = await router._schedule_factory( - model="gpt-5-mini", - args=( - "gpt-5-mini", - "Hello, how are you?", - ), - priority=1, - kwargs={}, - original_function=router.atext_completion, - ) - mock_atext_completion.assert_awaited_once() - assert "priority" not in mock_atext_completion.call_args.kwargs - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_router_function_with_fallbacks(model_list, sync_mode): - """Test if the router 'async_function_with_fallbacks' + 'function_with_fallbacks' are working correctly""" - router = Router(model_list=model_list) - data = { - "model": "gpt-5-mini", - "messages": [{"role": "user", "content": "Hello, how are you?"}], - "mock_response": "I'm fine, thank you!", - "num_retries": 0, - } - if sync_mode: - response = router.function_with_fallbacks( - original_function=router._completion, - **data, - ) - else: - response = await router.async_function_with_fallbacks( - original_function=router._acompletion, - **data, - ) - assert response.choices[0].message.content == "I'm fine, thank you!" - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_router_function_with_retries(model_list, sync_mode): - """Test if the router 'async_function_with_retries' + 'function_with_retries' are working correctly""" - router = Router(model_list=model_list) - data = { - "model": "gpt-5-mini", - "messages": [{"role": "user", "content": "Hello, how are you?"}], - "mock_response": "I'm fine, thank you!", - "num_retries": 0, - } - response = await router.async_function_with_retries( - original_function=router._acompletion, - **data, - ) - - assert response.choices[0].message.content == "I'm fine, thank you!" - - -@pytest.mark.asyncio -async def test_router_make_call(model_list): - """Test if the router 'make_call' function is working correctly""" - - router = Router(model_list=model_list) - response = await router.make_call( - original_function=router._acompletion, - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hello, how are you?"}], - mock_response="I'm fine, thank you!", - ) - assert response.choices[0].message.content == "I'm fine, thank you!" - - response = await router.make_call( - original_function=router._atext_completion, - model="gpt-5-mini", - prompt="Hello, how are you?", - mock_response="I'm fine, thank you!", - ) - assert response.choices[0].text == "I'm fine, thank you!" - - response = await router.make_call( - original_function=router._aembedding, - model="gpt-5-mini", - input="Hello, how are you?", - mock_response=[0.1, 0.2, 0.3], - ) - assert response.data[0].embedding == [0.1, 0.2, 0.3] - - response = await router.make_call( - original_function=router._aimage_generation, - model="gpt-image-1", - prompt="A cute baby sea otter", - mock_response="https://example.com/image.png", - ) - assert response.data[0].url == "https://example.com/image.png" - - -def test_update_kwargs_with_deployment(model_list): - """Test if the '_update_kwargs_with_deployment' function is working correctly""" - router = Router(model_list=model_list) - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") - router._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - ) - set_fields = ["deployment", "api_base", "model_info"] - assert all(field in kwargs["metadata"] for field in set_fields) - - -def test_update_kwargs_with_default_litellm_params(model_list): - """Test if the '_update_kwargs_with_default_litellm_params' function is working correctly""" - router = Router( - model_list=model_list, - default_litellm_params={"api_key": "test", "metadata": {"key": "value"}}, - ) - kwargs: dict = {"metadata": {"key2": "value2"}} - router._update_kwargs_with_default_litellm_params(kwargs=kwargs) - assert kwargs["api_key"] == "test" - assert kwargs["metadata"]["key"] == "value" - assert kwargs["metadata"]["key2"] == "value2" - - -def test_get_timeout(model_list): - """Test if the '_get_timeout' function is working correctly""" - router = Router(model_list=model_list) - timeout = router._get_timeout(kwargs={}, data={"timeout": 100}) - assert timeout == 100 - - -@pytest.mark.parametrize( - "fallback_kwarg, expected_error", - [ - ("mock_testing_fallbacks", litellm.InternalServerError), - ("mock_testing_context_fallbacks", litellm.ContextWindowExceededError), - ("mock_testing_content_policy_fallbacks", litellm.ContentPolicyViolationError), - ], -) -def test_handle_mock_testing_fallbacks(model_list, fallback_kwarg, expected_error): - """Test if the '_handle_mock_testing_fallbacks' function is working correctly""" - router = Router(model_list=model_list) - data = { - fallback_kwarg: True, - } - - with pytest.raises(expected_error): - router._handle_mock_testing_fallbacks( - kwargs=data, - ) - - -def test_handle_mock_testing_rate_limit_error(model_list): - """Test if the '_handle_mock_testing_rate_limit_error' function is working correctly""" - router = Router(model_list=model_list) - data = { - "mock_testing_rate_limit_error": True, - } - - with pytest.raises(litellm.RateLimitError): - router._handle_mock_testing_rate_limit_error( - kwargs=data, - ) - - @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_deployment_callback_on_success(sync_mode): @@ -467,9 +247,7 @@ async def test_deployment_callback_on_success(sync_mode): ] router = Router(model_list=model_list) # Get the actual deployment ID that was generated - gpt_deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + gpt_deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") deployment_id = gpt_deployment["model_info"]["id"] standard_logging_payload = create_standard_logging_payload() @@ -549,447 +327,6 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( assert await router.cache.async_get_cache(key=tpm_key) == 17 -class _DiscardingLoggingWorker: - def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: - async_coroutine.close() - - -async def _client_never_sends() -> str: - await asyncio.Event().wait() - raise AssertionError("unreachable") - - -def _realtime_client_ws() -> MagicMock: - client_ws: Final = MagicMock() - client_ws.headers = {} - client_ws.receive_text = AsyncMock(side_effect=_client_never_sends) - client_ws.send_text = AsyncMock() - client_ws.close = AsyncMock() - return client_ws - - -def _realtime_session_logging() -> LiteLLMLogging: - return LiteLLMLogging( - model="gpt-realtime-2.1", - messages=[], - stream=False, - call_type="_arealtime", - start_time=datetime.now(), - litellm_call_id=str(uuid.uuid4()), - function_id="realtime-session", - ) - - -async def _relay_upstream( - litellm_logging_obj: LiteLLMLogging, - *upstream_frames: bytes | Exception, - client_ws: MagicMock | None = None, - provider_config: BaseRealtimeConfig | None = None, - upstream_stays_open: bool = False, -) -> None: - backend_ws: Final = MagicMock() - backend_ws.recv = AsyncMock( - side_effect=_frames_then_silence(*upstream_frames) if upstream_stays_open else list(upstream_frames) - ) - backend_ws.send = AsyncMock() - relay: Final = RealTimeStreaming( - client_ws or _realtime_client_ws(), - backend_ws, - litellm_logging_obj, - provider_config=provider_config, - model="gpt-realtime-2.1", - logging_worker=_DiscardingLoggingWorker(), - ) - await asyncio.wait_for(relay.bidirectional_forward(), timeout=2) - - -_REALTIME_SESSION_CREATED: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() - - -async def _clean_realtime_session(litellm_logging_obj: LiteLLMLogging) -> None: - await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(Close(1000, ""), None)) - - -async def _upstream_internal_error_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None: - await _relay_upstream( - litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(Close(1011, "upstream failed"), None) - ) - - -async def _upstream_drop_without_close_frame_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None: - await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(None, None)) - - -async def _upstream_close_before_any_frame(litellm_logging_obj: LiteLLMLogging) -> None: - await _relay_upstream(litellm_logging_obj, ConnectionClosed(Close(1008, "model not found"), None)) - - -def _frames_then_silence(*frames: bytes | str) -> Callable[[], Awaitable[bytes | str]]: - pending: Final = iter(frames) - - async def _next_frame() -> bytes | str: - frame: Final = next(pending, None) - if frame is None: - await asyncio.Event().wait() - raise AssertionError("unreachable") - return frame - - return _next_frame - - -def _client_ws_sending(*client_frames: str) -> MagicMock: - client_ws: Final = _realtime_client_ws() - client_ws.receive_text = AsyncMock(side_effect=_frames_then_silence(*client_frames)) - return client_ws - - -async def _vertex_access_token() -> str: - return "token" - - -async def _client_frame_the_provider_rejects(litellm_logging_obj: LiteLLMLogging) -> None: - deferred_setup_provider: Final = VertexChirpRealtimeConfig( - resolve_access_token=_vertex_access_token, project="proj-1", location="us" - ) - await _relay_upstream( - litellm_logging_obj, - client_ws=_client_ws_sending( - json.dumps( - { - "type": "session.update", - "session": {"type": "transcription", "audio": {"input": {"format": {"type": "audio/pcmu"}}}}, - } - ) - ), - provider_config=deferred_setup_provider, - upstream_stays_open=True, - ) - - -async def _client_hangs_up_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None: - client_ws: Final = _realtime_client_ws() - client_ws.receive_text = AsyncMock(side_effect=ConnectionClosed(Close(1000, ""), None)) - await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, client_ws=client_ws, upstream_stays_open=True) - - -async def _refused_handshake(litellm_logging_obj: LiteLLMLogging) -> None: - await close_after_upstream_handshake_refusal(_realtime_client_ws(), 403) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("realtime_session", "counted_successes"), - [ - (_clean_realtime_session, 1), - (_upstream_close_before_any_frame, 0), - (_refused_handshake, 0), - (_upstream_internal_error_after_a_frame, 0), - (_upstream_drop_without_close_frame_after_a_frame, 0), - (_client_frame_the_provider_rejects, 0), - (_client_hangs_up_after_a_frame, 1), - ], - ids=[ - "clean-session", - "upstream-close-before-any-frame", - "refused-handshake", - "upstream-1011-after-a-frame", - "upstream-1006-after-a-frame", - "client-frame-the-provider-rejects", - "client-hangs-up-after-a-frame", - ], -) -async def test_realtime_session_counts_a_router_success_only_when_it_ended_without_a_failure( - realtime_session: Callable[[LiteLLMLogging], Awaitable[None]], counted_successes: int -) -> None: - router: Final = Router( - model_list=[ - { - "model_name": "realtime-voice", - "litellm_params": {"model": "openai/gpt-realtime-2.1"}, - "model_info": {"id": "realtime-a"}, - } - ], - num_retries=0, - ) - - async def _arealtime(litellm_logging_obj: LiteLLMLogging, **kwargs: object) -> None: - await realtime_session(litellm_logging_obj) - - await router._ageneric_api_call_with_fallbacks( - model="realtime-voice", original_function=_arealtime, litellm_logging_obj=_realtime_session_logging() - ) - - assert get_deployment_successes_for_current_minute(router, "realtime-a") == counted_successes - - -@pytest.mark.asyncio -@pytest.mark.parametrize("realtime_model", ["bedrock/amazon.nova-sonic-v1:0", "openai/gpt-realtime-2.1"]) -@pytest.mark.parametrize(("clean_sessions", "cooled_down"), [(10, False), (3, True)]) -async def test_realtime_cooldown_reads_the_failure_rate_of_the_sessions_a_deployment_served( - realtime_model: str, clean_sessions: int, cooled_down: bool -) -> None: - import time - - provider, _, provider_model = realtime_model.partition("/") - router: Final = Router( - model_list=[ - { - "model_name": "realtime-voice", - "litellm_params": {"model": realtime_model, "order": order}, - "model_info": {"id": deployment_id}, - } - for order, deployment_id in ((1, "realtime-a"), (2, "realtime-b")) - ], - num_retries=0, - ) - drop: Final = litellm.ServiceUnavailableError( - message="realtime provider stream ended before the session closed", - llm_provider=provider, - model=provider_model, - ) - - def realtime_session(failure: Exception | None): - async def _arealtime(litellm_logging_obj: LiteLLMLogging, **kwargs: object) -> None: - if failure is not None: - raise failure - await _clean_realtime_session(litellm_logging_obj) - - return _arealtime - - for _ in range(clean_sessions): - await router._ageneric_api_call_with_fallbacks( - model="realtime-voice", - original_function=realtime_session(None), - litellm_logging_obj=_realtime_session_logging(), - ) - - async def drop_session() -> bool: - with pytest.raises(litellm.ServiceUnavailableError): - await router._ageneric_api_call_with_fallbacks( - model="realtime-voice", - original_function=realtime_session(drop), - litellm_logging_obj=_realtime_session_logging(), - ) - return router.deployment_callback_on_failure( - kwargs={ - "exception": drop, - "litellm_params": {"metadata": {"model_group": "realtime-voice"}, "model_info": {"id": "realtime-a"}}, - }, - completion_response=None, - start_time=time.time(), - end_time=time.time(), - ) - - drop_cooldowns: Final = [await drop_session() for _ in range(5)] - - assert any(drop_cooldowns) is cooled_down - - -async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]: - routed: Final[dict[str, object]] = {} - - async def call(**kwargs: object) -> None: - routed.update(kwargs) - - call.__name__ = function_name - await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call, **request) - return routed - - -async def _finish_35_token_websocket_turn(routed: Mapping[str, object], call_id: str) -> None: - import time - - from litellm.litellm_core_utils.litellm_logging import Logging - - turn: Final = Logging( - model="gpt-5-mini", - messages=[], - stream=False, - call_type=CallTypes.aresponses_websocket.value, - start_time=time.time(), - litellm_call_id=call_id, - function_id=call_id, - ) - turn.update_environment_variables(litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={}) - await turn.dispatch_success_handlers( - [{"type": "response.completed", "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}}], - prefer_async_handlers=True, - ) - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_budget(): - router = Router( - model_list=[ - { - "model_name": "ws-budget", - "litellm_params": { - "model": "openai/gpt-5-mini", - "api_key": "sk-fake", - "max_budget": 0.001, - "budget_duration": "1d", - }, - "model_info": {"id": "ws-budget-dep"}, - } - ] - ) - assert "turn_admission" not in await _kwargs_routed_to(router, "ws-budget", "aresponses") - admit_turn: Final = (await _kwargs_routed_to(router, "ws-budget", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - - await admit_turn() - await router.cache.async_set_cache("deployment_spend:ws-budget-dep:1d", 0.002) - - with pytest.raises(litellm.RateLimitError, match="crossed budget"): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_to_least_busy(): - router: Final = Router( - model_list=[ - { - "model_name": "ws-least-busy", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake"}, - "model_info": {"id": deployment_id}, - } - for deployment_id in ("ws-least-busy-a", "ws-least-busy-b") - ], - routing_strategy="least-busy", - ) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-least-busy", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - for _ in range(3): - await admit_turn() - - picks: Final = [ - (await router._core_async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[ - "model_info" - ]["id"] - for _ in range(2) - ] - - assert sorted(picks) == ["ws-least-busy-a", "ws-least-busy-b"] - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm(): - router = Router( - model_list=[ - { - "model_name": "ws-rpm", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 1}, - "model_info": {"id": "ws-rpm-dep"}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - - with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_is_checked_against_limits_changed_after_the_socket_opened(): - router = Router( - model_list=[ - { - "model_name": "ws-rpm-lowered", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 100}, - "model_info": {"id": "ws-rpm-lowered-dep"}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm-lowered", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - await admit_turn() - - router.upsert_deployment( - Deployment( - model_name="ws-rpm-lowered", - litellm_params=LiteLLM_Params(model="openai/gpt-5-mini", api_key="sk-fake", rpm=1), - model_info=ModelInfo(id="ws-rpm-lowered-dep"), - ) - ) - - with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_counts_toward_deployment_tpm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "callbacks", []) - router = Router( - model_list=[ - { - "model_name": "ws-tpm", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30}, - "model_info": {"id": "ws-tpm-dep"}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks)) - routed: Final = await _kwargs_routed_to(router, "ws-tpm", "_aresponses_websocket") - admit_turn: Final = routed["turn_admission"] - assert callable(admit_turn) - await admit_turn() - - await _finish_35_token_websocket_turn(routed, "ws-tpm-turn") - - with pytest.raises(litellm.RateLimitError, match="TPM limit=30, current usage=35"): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_is_refused_by_usage_based_routing_once_deployment_tpm_is_used( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setattr(litellm, "callbacks", []) - router = Router( - model_list=[ - { - "model_name": "ws-tpm-v2", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30}, - "model_info": {"id": "ws-tpm-v2-dep"}, - } - ], - routing_strategy="usage-based-routing-v2", - ) - monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks)) - routed: Final = await _kwargs_routed_to(router, "ws-tpm-v2", "_aresponses_websocket") - admit_turn: Final = routed["turn_admission"] - assert callable(admit_turn) - await admit_turn() - - await _finish_35_token_websocket_turn(routed, "ws-tpm-v2-turn") - - with pytest.raises(litellm.RateLimitError): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turns_do_not_reserve_deployment_otpm(): - router = Router( - model_list=[ - { - "model_name": "ws-otpm", - "litellm_params": {"model": "openai/ws-otpm-unpriced", "api_key": "sk-fake", "otpm": 5000}, - "model_info": {"id": "ws-otpm-dep"}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-otpm", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - - for _ in range(3): - await admit_turn() - - def _io_limited_router(kind: str, **litellm_params: object) -> Router: return Router( model_list=[ @@ -1012,178 +349,6 @@ def _minutes_counter_key(kind: str, minute: datetime, model: str = "openai/ws-io return f"global_router:ws-io-dep:{model}:{kind}:{minute.strftime('%H-%M')}" -async def _count_minutes_tokens( - router: Router, kind: str, tokens: int, minute: datetime, model: str = "openai/ws-io-unpriced" -) -> None: - await router.cache.async_increment_cache(key=_minutes_counter_key(kind, minute, model), value=tokens, ttl=120) - - -async def _settle_turn( - router: Router, reservation: Mapping[str, object], input_tokens: int, output_tokens: int -) -> None: - from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor - from litellm.litellm_core_utils.get_litellm_params import get_litellm_params - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck - - rate_limits: Final = next( - check for check in router.optional_callbacks or () if isinstance(check, ModelRateLimitingCheck) - ) - await rate_limits.async_log_success_event( - {"litellm_params": get_litellm_params(litellm_metadata=dict(reservation))}, - ResponsesWebSocketTokenUsageProcessor.create_logging_object( - [ - { - "type": "response.completed", - "response": {"usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}}, - } - ] - ), - None, - None, - ) - - -_LONG_OUTPUT_MODEL: Final = "openai/gpt-5.4-mini" - - -def _long_output_router() -> tuple[Router, Mapping[str, int]]: - max_output: Final = litellm.get_model_info(_LONG_OUTPUT_MODEL)["max_output_tokens"] - assert isinstance(max_output, int) - limits: Final = {"itpm": 10_000, "otpm": max_output // 2} - return _io_limited_router("itpm", model=_LONG_OUTPUT_MODEL, otpm=limits["otpm"]), limits - - -@pytest.mark.asyncio -@pytest.mark.parametrize("kind", ["itpm", "otpm"]) -async def test_responses_websocket_opens_without_counting_tokens_until_the_minute_reaches_deployment_io_limit( - kind: str, router_minute_pinned: datetime -): - router, limits = _long_output_router() - - await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) - - assert [ - await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) - for counted in ("itpm", "otpm") - ] == [None, None] - await _count_minutes_tokens(router, kind, limits[kind], router_minute_pinned, _LONG_OUTPUT_MODEL) - with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit={limits[kind]}, current usage="): - await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) - - -@pytest.mark.asyncio -async def test_responses_websocket_first_turn_counts_its_actual_tokens_toward_deployment_itpm_and_otpm( - router_minute_pinned: datetime, -): - router, _ = _long_output_router() - socket_metadata: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()))[ - "litellm_metadata" - ] - assert isinstance(socket_metadata, dict) - - await _settle_turn(router, socket_metadata, input_tokens=20, output_tokens=15) - - assert [ - await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) - for counted in ("itpm", "otpm") - ] == [20, 15] - - -_ADMITTED_LATE_IN_THE_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 50, tzinfo=timezone.utc) -_FINISHED_IN_THE_NEXT_MINUTE: Final = datetime(2026, 1, 1, 12, 1, 15, tzinfo=timezone.utc) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("kind", ["itpm", "otpm"]) -@pytest.mark.parametrize("turn", ["first", "later"]) -async def test_responses_websocket_turn_finishing_in_the_next_minute_counts_its_tokens_toward_that_minute( - kind: str, turn: str, monkeypatch: pytest.MonkeyPatch -): - _pin_router_clock(monkeypatch, _ADMITTED_LATE_IN_THE_MINUTE) - router: Final = _io_limited_router(kind) - socket: Final = await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) - admit_turn: Final = socket["turn_admission"] - assert callable(admit_turn) - reservation: Final = socket["litellm_metadata"] if turn == "first" else await admit_turn() - assert isinstance(reservation, Mapping) - _pin_router_clock(monkeypatch, _FINISHED_IN_THE_NEXT_MINUTE) - - await _settle_turn(router, reservation, input_tokens=10_000, output_tokens=10_000) - - with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage=10000"): - await admit_turn() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("body", [{}, {"websocket": True}], ids=["plain", "websocket_key"]) -async def test_messages_call_reserves_max_tokens_against_deployment_otpm_whatever_websocket_key_its_body_carries( - body: Mapping[str, object], router_minute_pinned: datetime -): - router: Final = _io_limited_router("otpm") - - for _ in range(3): - await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) - - assert await router.cache.async_get_cache(key=_minutes_counter_key("otpm", router_minute_pinned)) == 9_000 - with pytest.raises(litellm.RateLimitError, match="OTPM limit=10000, current usage=12000"): - await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) - - -async def _realtime_opens_until_refused(monkeypatch: pytest.MonkeyPatch, **socket: object) -> tuple[int, str]: - monkeypatch.setattr(litellm, "callbacks", []) - router: Final = _io_limited_router("otpm", otpm=50) - for opened in range(51): - try: - await _kwargs_routed_to(router, "ws-io", "_arealtime", **socket) - except litellm.RateLimitError as refused: - return opened, str(refused) - pytest.fail("realtime opens never reached the deployment's otpm") - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("router_minute_pinned") -async def test_realtime_open_reserves_and_is_refused_at_deployment_otpm_like_an_open_without_a_socket( - monkeypatch: pytest.MonkeyPatch, -): - with_socket: Final = await _realtime_opens_until_refused(monkeypatch, websocket=object()) - - assert "OTPM limit=50, current usage=" in with_socket[1] - assert with_socket == await _realtime_opens_until_refused(monkeypatch) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("kind", ["itpm", "otpm"]) -async def test_responses_websocket_turn_is_refused_once_the_minutes_tokens_reach_deployment_io_limit( - kind: str, router_minute_pinned: datetime -): - router: Final = _io_limited_router(kind) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) - - with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage="): - await admit_turn() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("kind", ["itpm", "otpm"]) -async def test_responses_websocket_turn_refused_for_tokens_leaves_the_daily_request_limit_unused( - kind: str, router_minute_pinned: datetime -): - router: Final = _io_limited_router(kind, rpd=2) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) - with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000"): - await admit_turn() - await _count_minutes_tokens(router, kind, -10_000, router_minute_pinned) - - await admit_turn() - - with pytest.raises(litellm.RateLimitError, match="RPD"): - await admit_turn() - - _PROVIDER_RESPONSE: Final = { "id": "resp_io", "object": "response", @@ -1432,41 +597,6 @@ async def test_deployment_callback_on_failure(model_list): ) -def test_deployment_callback_respects_cooldown_time(model_list): - """Ensure per-model cooldown_time is honored even when exception headers are present.""" - import time - from unittest.mock import patch - - import httpx - - router = Router(model_list=model_list) - - class FakeException(Exception): - def __init__(self): - self.status_code = 429 - self.headers = httpx.Headers({"x-test": "1"}) - - kwargs = { - "exception": FakeException(), - "litellm_params": { - "metadata": {"model_group": "gpt-5-mini"}, - "model_info": {"id": 100}, - "cooldown_time": 0, - }, - } - - with patch("litellm.router_python_core.set_cooldown_deployments") as mock_set: - router.deployment_callback_on_failure( - kwargs=kwargs, - completion_response=None, - start_time=time.time(), - end_time=time.time(), - ) - - mock_set.assert_called_once() - assert mock_set.call_args.kwargs["time_to_cooldown"] == 0 - - @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) -> None: """log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the @@ -1537,57 +667,6 @@ def test_log_retry_reads_base_model_from_the_deployments_litellm_params() -> Non assert (record["model"], record["base_model"]) == ("azure/my-named-deployment", "azure/gpt-4o-mini") -@pytest.mark.usefixtures("router_minute_pinned") -def test_update_usage(model_list): - """Test if the '_update_usage' function is working correctly""" - router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") - deployment_id = deployment["model_info"]["id"] - request_count = router._update_usage(deployment_id=deployment_id, parent_otel_span=None) - assert request_count == 1 - - request_count = router._update_usage(deployment_id=deployment_id, parent_otel_span=None) - - assert request_count == 2 - - -@pytest.mark.parametrize("finish_reason, expected_fallback", [("content_filter", True), ("stop", False)]) -@pytest.mark.parametrize("fallback_type", ["model-specific", "default"]) -def test_should_raise_content_policy_error(model_list, finish_reason, expected_fallback, fallback_type): - """Test if the '_should_raise_content_policy_error' function is working correctly""" - router = Router( - model_list=model_list, - default_fallbacks=["gpt-5.5"] if fallback_type == "default" else None, - ) - - assert ( - router._should_raise_content_policy_error( - model="gpt-5-mini", - response=litellm.ModelResponse( - model="gpt-5-mini", - choices=[ - { - "finish_reason": finish_reason, - "message": {"content": "I'm fine, thank you!"}, - } - ], - usage={"total_tokens": 100}, - ), - kwargs={ - "content_policy_fallbacks": ([{"gpt-5-mini": "gpt-5.5"}] if fallback_type == "model-specific" else None) - }, - ) - is expected_fallback - ) - - -def test_get_healthy_deployments(model_list): - """Test if the '_get_healthy_deployments' function is working correctly""" - router = Router(model_list=model_list) - deployments = router._get_healthy_deployments(model="gpt-5-mini", parent_otel_span=None) - assert len(deployments) > 0 - - @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_routing_strategy_pre_call_checks(model_list, sync_mode): @@ -1600,9 +679,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") litellm_logging_obj = Logging( model="gpt-5-mini", @@ -1617,9 +694,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): router.routing_strategy_pre_call_checks(deployment) else: ## NO EXCEPTION - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) ## WITH EXCEPTION - rate limit error with patch.object( @@ -1634,49 +709,12 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): ), ): with pytest.raises(litellm.RateLimitError): - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) ## WITH EXCEPTION - generic error - with patch.object( - callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error")) - ): + with patch.object(callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error"))): with pytest.raises(Exception, match="Error"): - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) - - -@pytest.mark.parametrize( - "set_supported_environments, supported_environments, is_supported", - [(True, ["staging"], True), (False, None, True), (True, ["development"], False)], -) -def test_create_deployment( - model_list, set_supported_environments, supported_environments, is_supported -): - """Test if the '_create_deployment' function is working correctly""" - router = Router(model_list=model_list) - - if set_supported_environments: - os.environ["LITELLM_ENVIRONMENT"] = "staging" - deployment = router._create_deployment( - deployment_info={}, - _model_name="gpt-5-mini", - _litellm_params={ - "model": "gpt-5-mini", - "api_key": "test", - "custom_llm_provider": "openai", - }, - _model_info={ - "id": 100, - "supported_environments": supported_environments, - }, - ) - if is_supported: - assert deployment is not None - else: - assert deployment is None + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) @pytest.mark.parametrize( @@ -1688,20 +726,14 @@ def test_deployment_is_active_for_environment( ): """Test if the '_deployment_is_active_for_environment' function is working correctly""" router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") if set_supported_environments: os.environ["LITELLM_ENVIRONMENT"] = "staging" deployment["model_info"]["supported_environments"] = supported_environments if is_supported: - assert ( - router.deployment_is_active_for_environment(deployment=deployment) is True - ) + assert router.deployment_is_active_for_environment(deployment=deployment) is True else: - assert ( - router.deployment_is_active_for_environment(deployment=deployment) is False - ) + assert router.deployment_is_active_for_environment(deployment=deployment) is False def test_set_model_list(model_list): @@ -1711,18 +743,6 @@ def test_set_model_list(model_list): assert len(router.model_list) == len(model_list) -def test_add_deployment(model_list): - """Test if the '_add_deployment' function is working correctly""" - router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") - deployment["model_info"]["id"] = "100" - - router.add_deployment(deployment=deployment) - - router._add_deployment(deployment=deployment) - assert len(router.model_list) == len(model_list) + 1 - - def test_upsert_deployment(model_list): """Test if the 'upsert_deployment' function is working correctly""" router = Router(model_list=model_list) @@ -1758,63 +778,6 @@ def test_get_model_group(model_list): assert model_group[0]["model_name"] == "gpt-5-mini" -@pytest.mark.parametrize("user_facing_model_group_name", ["gpt-5-mini", "gpt-5.5"]) -def test_set_model_group_info(model_list, user_facing_model_group_name): - """Test if the 'set_model_group_info' function is working correctly""" - router = Router(model_list=model_list) - resp = router._set_model_group_info( - model_group="gpt-5-mini", - user_facing_model_group_name=user_facing_model_group_name, - ) - assert resp is not None - assert resp.model_group == user_facing_model_group_name - - -def test_set_model_group_info_reports_supports_continuation_when_any_deployment_has_it(): - router: Final = Router( - model_list=[ - { - "model_name": "long-decode", - "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, - }, - { - "model_name": "long-decode", - "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "fake"}, - "model_info": {"supports_continuation": True}, - }, - { - "model_name": "long-decode", - "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"}, - }, - { - "model_name": "plain", - "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, - }, - ] - ) - long_decode: Final = router._set_model_group_info( - model_group="long-decode", user_facing_model_group_name="long-decode" - ) - plain: Final = router._set_model_group_info(model_group="plain", user_facing_model_group_name="plain") - assert long_decode is not None and long_decode.supports_continuation is True - assert plain is not None and plain.supports_continuation is False - - -@pytest.mark.parametrize("value", ["maybe", "", ["x"], "true", 1]) -def test_set_model_group_info_counts_only_a_boolean_true_supports_continuation(value: object) -> None: - router: Final = Router( - model_list=[ - { - "model_name": "long-decode", - "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, - "model_info": {"supports_continuation": value}, - }, - ] - ) - group: Final = router._set_model_group_info(model_group="long-decode", user_facing_model_group_name="long-decode") - assert group is not None and group.supports_continuation is False - - @pytest.mark.asyncio async def test_set_response_headers(model_list): """Test if the 'set_response_headers' function is working correctly""" @@ -1823,42 +786,6 @@ async def test_set_response_headers(model_list): assert resp is None -@pytest.mark.asyncio -async def test_set_response_headers_passes_through_post_increment_counters(model_list): - from pydantic import BaseModel - - class _Usage(BaseModel): - total_tokens: int = 42 - - class _Resp(BaseModel): - usage: _Usage = _Usage() - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 958, - "x-ratelimit-limit-tokens": 1000, - "x-ratelimit-remaining-requests": 99, - "x-ratelimit-limit-requests": 100, - "x-ratelimit-remaining-input-tokens": 1000, - "x-ratelimit-remaining-output-tokens": 500, - } - ) - - resp = _Resp() - resp._hidden_params = {} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") - - headers = resp._hidden_params["additional_headers"] - assert headers["x-ratelimit-remaining-tokens"] == 958 - assert headers["x-ratelimit-remaining-requests"] == 99 - assert headers["x-ratelimit-limit-tokens"] == 1000 - assert headers["x-ratelimit-limit-requests"] == 100 - assert headers["x-ratelimit-remaining-input-tokens"] == 1000 - assert headers["x-ratelimit-remaining-output-tokens"] == 500 - - def _rpm_tpm_router(model_id: str) -> Router: return Router( model_list=[ @@ -1872,7 +799,6 @@ def _rpm_tpm_router(model_id: str) -> Router: def _pin_router_clock(monkeypatch: pytest.MonkeyPatch, pinned: datetime) -> None: - monkeypatch.setattr("litellm.router_python_core.get_utc_datetime", lambda: pinned) monkeypatch.setattr( "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: pinned ) @@ -1880,6 +806,10 @@ def _pin_router_clock(monkeypatch: pytest.MonkeyPatch, pinned: datetime) -> None @pytest.fixture def router_minute_pinned(monkeypatch: pytest.MonkeyPatch) -> datetime: + # The router's own usage counters read the real clock: start well inside a minute so a test never spans two + seconds_left: Final = 60 - datetime.now(timezone.utc).second + if seconds_left < 10: + time.sleep(seconds_left + 0.1) pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) _pin_router_clock(monkeypatch, pinned) return pinned @@ -1958,121 +888,6 @@ async def test_increment_deployment_usage_for_response_skips_session_wrappers(): assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in request_kwargs["litellm_metadata"] -@pytest.mark.asyncio -@pytest.mark.usefixtures("router_minute_pinned") -async def test_increment_deployment_usage_writes_only_positive_deltas_for_limited_deployments(): - router = _rpm_tpm_router("lit-3058-delta") - unlimited = Router( - model_list=[ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"}, - "model_info": {"id": "lit-3058-unlimited"}, - } - ] - ) - - tpm_key = await router._increment_deployment_usage( - deployment_id="lit-3058-delta", - deployment_name="gpt-5-mini", - model_group="gpt-5-mini", - total_tokens=25, - rpm_increment=1, - parent_otel_span=None, - ) - assert tpm_key is not None - assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) - - assert ( - await router._increment_deployment_usage( - deployment_id="lit-3058-delta", - deployment_name="gpt-5-mini", - model_group="gpt-5-mini", - total_tokens=0, - rpm_increment=0, - parent_otel_span=None, - ) - is None - ) - assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) - - assert ( - await unlimited._increment_deployment_usage( - deployment_id="lit-3058-unlimited", - deployment_name="gpt-5-mini", - model_group="gpt-5-mini", - total_tokens=25, - rpm_increment=1, - parent_otel_span=None, - ) - is None - ) - assert await unlimited.get_model_group_usage("gpt-5-mini") == (None, None) - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("router_minute_pinned") -async def test_increment_deployment_usage_starts_tracking_after_limit_is_added(): - router = Router( - model_list=[ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"}, - "model_info": {"id": "lit-3058-upsert"}, - } - ] - ) - increment = functools.partial( - router._increment_deployment_usage, - deployment_id="lit-3058-upsert", - deployment_name="gpt-5-mini", - model_group="gpt-5-mini", - total_tokens=25, - rpm_increment=1, - parent_otel_span=None, - ) - - assert await increment() is None - - router.upsert_deployment( - Deployment( - model_name="gpt-5-mini", - litellm_params=LiteLLM_Params(model="gpt-5-mini", api_key="sk-fake", rpm=10), - model_info={"id": "lit-3058-upsert"}, - ) - ) - - assert await increment() is not None - assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) - - -@pytest.mark.parametrize( - "litellm_params, model_info, deployment_id, tracks", - [ - ({}, {}, "lit-3058-tracks", False), - ({"rpm": 10}, {}, "lit-3058-tracks", True), - ({"tpm": 1000}, {}, "lit-3058-tracks", True), - ({"otpm": 500}, {}, "lit-3058-tracks", True), - ({}, {"itpm": 500}, "lit-3058-tracks", True), - ({"rpm": 10}, {}, "lit-3058-unknown", False), - ], -) -def test_deployment_tracks_usage_for_every_limit_source( - litellm_params: dict, model_info: dict, deployment_id: str, tracks: bool -): - router = Router( - model_list=[ - { - "model_name": "my-unmapped-model", - "litellm_params": {"model": "openai/my-unmapped-model", "api_key": "sk-fake", **litellm_params}, - "model_info": {"id": "lit-3058-tracks", **model_info}, - } - ] - ) - - assert router._deployment_tracks_usage(deployment_id, "my-unmapped-model") is tracks - - def _shared_redis_stub(store: dict) -> MagicMock: from litellm.caching.redis_cache import RedisCache @@ -2158,27 +973,19 @@ async def test_get_model_group_io_token_usage_sums_across_deployments(): minute = get_utc_datetime().strftime("%H-%M") keys_and_values = [ ( - RouterCacheEnum.ITPM.value.format( - id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute - ), + RouterCacheEnum.ITPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), 30, ), ( - RouterCacheEnum.OTPM.value.format( - id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute - ), + RouterCacheEnum.OTPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), 10, ), ( - RouterCacheEnum.ITPM.value.format( - id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute - ), + RouterCacheEnum.ITPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), 70, ), ( - RouterCacheEnum.OTPM.value.format( - id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute - ), + RouterCacheEnum.OTPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), 20, ), ] @@ -2200,66 +1007,57 @@ async def test_get_model_group_io_token_usage_no_deployments_returns_none(): assert current_otpm is None +def _limited_router(**io_limits: int) -> Router: + return Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + "tpm": 1000, + "rpm": 100, + **io_limits, + }, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + @pytest.mark.asyncio -async def test_get_remaining_model_group_usage_merges_io_and_tpm_headers(model_list): +async def test_get_remaining_model_group_usage_merges_io_and_tpm_headers(): """ A model group with both itpm/otpm and tpm/rpm limits must expose the standard remaining-tokens/requests headers alongside the input/output token headers, so clients and prometheus gauges relying on either still get data. """ - from unittest.mock import Mock + router = _limited_router(itpm=900, otpm=450) - from litellm.types.router import ModelGroupInfo + headers = await router.get_remaining_model_group_usage("limited") - router = Router(model_list=model_list) - router._cached_get_model_group_info = Mock( - return_value=ModelGroupInfo( - model_group="gpt-3.5-turbo", - providers=["openai"], - itpm=2000, - otpm=1000, - tpm=5000, - rpm=50, - ) - ) - router.get_model_group_io_token_usage = AsyncMock(return_value=(100, 40)) - router._core_get_model_group_usage = AsyncMock(return_value=(500, 5)) + assert headers["x-ratelimit-remaining-input-tokens"] == 900 + assert headers["x-ratelimit-remaining-output-tokens"] == 450 + assert headers["x-ratelimit-remaining-tokens"] == 1000 + assert headers["x-ratelimit-remaining-requests"] == 100 - headers = await router.get_remaining_model_group_usage("gpt-3.5-turbo") - assert headers["x-ratelimit-remaining-input-tokens"] == 1900 - assert headers["x-ratelimit-remaining-output-tokens"] == 960 - assert headers["x-ratelimit-remaining-tokens"] == 4500 - assert headers["x-ratelimit-remaining-requests"] == 45 +class _HeaderlessResponse(BaseModel): + _hidden_params: dict = {} @pytest.mark.asyncio -async def test_set_response_headers_native_input_token_header_does_not_suppress_router_headers(model_list): +async def test_set_response_headers_native_input_token_header_does_not_suppress_router_headers(): """ A provider that natively returns `x-ratelimit-remaining-input-tokens` must not suppress the router's own remaining-tokens/requests headers for a non-IO model group. """ - from pydantic import BaseModel + router = _limited_router() - class _Usage(BaseModel): - total_tokens: int = 42 - - class _Resp(BaseModel): - usage: _Usage = _Usage() - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 1000, - "x-ratelimit-remaining-requests": 100, - } - ) - - resp = _Resp() + resp = _HeaderlessResponse() resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-input-tokens": 5}} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") + await router.set_response_headers(response=resp, model_group="limited") headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 1000 @@ -2269,29 +1067,12 @@ async def test_set_response_headers_native_input_token_header_does_not_suppress_ @pytest.mark.asyncio -async def test_set_response_headers_native_token_header_does_not_suppress_io_headers(model_list): - from pydantic import BaseModel +async def test_set_response_headers_native_token_header_does_not_suppress_io_headers(): + router = _limited_router(itpm=900, otpm=450) - class _Usage(BaseModel): - total_tokens: int = 42 - - class _Resp(BaseModel): - usage: _Usage = _Usage() - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 1000, - "x-ratelimit-remaining-requests": 100, - "x-ratelimit-remaining-input-tokens": 900, - "x-ratelimit-remaining-output-tokens": 450, - } - ) - - resp = _Resp() + resp = _HeaderlessResponse() resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-tokens": 5}} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") + await router.set_response_headers(response=resp, model_group="limited") headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 5 @@ -2301,27 +1082,16 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea @pytest.mark.asyncio -async def test_set_response_headers_handles_missing_usage(model_list): +async def test_set_response_headers_handles_missing_usage(): """ Streaming chunks and some response shapes may lack a `usage` attribute or populated `total_tokens`. Header composition must not depend on usage and never raise. """ - from pydantic import BaseModel + router = _limited_router() - class _Resp(BaseModel): - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 1000, - "x-ratelimit-remaining-requests": 100, - } - ) - - resp = _Resp() + resp = _HeaderlessResponse() resp._hidden_params = {} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") + await router.set_response_headers(response=resp, model_group="limited") headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 1000 @@ -2329,17 +1099,9 @@ async def test_set_response_headers_handles_missing_usage(model_list): @pytest.mark.asyncio -async def test_set_response_headers_dict_anthropic_messages_response(model_list): +async def test_set_response_headers_dict_anthropic_messages_response(): """Anthropic /v1/messages returns a dict; IO rate-limit headers must attach.""" - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-limit-input-tokens": 25, - "x-ratelimit-remaining-input-tokens": 20, - "x-ratelimit-limit-output-tokens": 100, - "x-ratelimit-remaining-output-tokens": 95, - } - ) + router = _limited_router(itpm=25, otpm=100) resp = { "id": "msg_123", @@ -2348,18 +1110,18 @@ async def test_set_response_headers_dict_anthropic_messages_response(model_list) "content": [{"type": "text", "text": "hi"}], "usage": {"input_tokens": 5, "output_tokens": 1}, } - await router.set_response_headers(response=resp, model_group="io-itpm-strict") + await router.set_response_headers(response=resp, model_group="limited") assert "_hidden_params" in resp headers = resp["_hidden_params"]["additional_headers"] - assert headers["x-litellm-model-group"] == "io-itpm-strict" + assert headers["x-litellm-model-group"] == "limited" assert headers["x-ratelimit-limit-input-tokens"] == 25 - assert headers["x-ratelimit-remaining-input-tokens"] == 20 - assert headers["x-ratelimit-remaining-output-tokens"] == 95 + assert headers["x-ratelimit-remaining-input-tokens"] == 25 + assert headers["x-ratelimit-remaining-output-tokens"] == 100 @pytest.mark.asyncio -async def test_set_response_headers_wraps_bare_async_generator(model_list): +async def test_set_response_headers_wraps_bare_async_generator(): """ Streaming responses that never go through Router.make_call's usual object-based wrappers (e.g. the Anthropic /v1/messages -> Responses API @@ -2371,21 +1133,15 @@ async def test_set_response_headers_wraps_bare_async_generator(model_list): yield {"type": "message_start"} yield {"type": "message_stop"} - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-limit-input-tokens": 25, - "x-ratelimit-remaining-input-tokens": 20, - } - ) + router = _limited_router(itpm=25) - wrapped = await router.set_response_headers(response=_raw_generator(), model_group="io-itpm-strict") + wrapped = await router.set_response_headers(response=_raw_generator(), model_group="limited") assert hasattr(wrapped, "_hidden_params") headers = wrapped._hidden_params["additional_headers"] - assert headers["x-litellm-model-group"] == "io-itpm-strict" + assert headers["x-litellm-model-group"] == "limited" assert headers["x-ratelimit-limit-input-tokens"] == 25 - assert headers["x-ratelimit-remaining-input-tokens"] == 20 + assert headers["x-ratelimit-remaining-input-tokens"] == 25 from collections.abc import AsyncIterator @@ -2417,29 +1173,6 @@ def test_update_settings(model_list): assert router.allowed_fails == 20 -def test_common_checks_available_deployment(model_list): - """Test if the 'common_checks_available_deployment' function is working correctly""" - router = Router(model_list=model_list) - _, available_deployments = router._common_checks_available_deployment( - model="gpt-5-mini", - messages=[{"role": "user", "content": "hi"}], - input="hi", - specific_deployment=False, - ) - - assert len(available_deployments) > 0 - - -def test_filter_cooldown_deployments(model_list): - """Test if the 'filter_cooldown_deployments' function is working correctly""" - router = Router(model_list=model_list) - deployments = router._filter_cooldown_deployments( - healthy_deployments=router.get_all_deployments(model_name="gpt-5-mini"), - cooldown_deployments=[], - ) - assert len(deployments) == len(router.get_all_deployments(model_name="gpt-5-mini")) - - @pytest.mark.parametrize( "exception_type, exception_name, num_retries", [ @@ -2495,23 +1228,6 @@ def test_get_allowed_fails_from_policy(model_list, exception_type, exception_nam assert calc_allowed_fails == allowed_fails -def test_initialize_alerting(model_list): - """Test if the 'initialize_alerting' function is working correctly""" - from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting - from litellm.types.router import AlertingConfig - - router = Router( - model_list=model_list, alerting_config=AlertingConfig(webhook_url="test") - ) - router._initialize_alerting() - - callback_added = False - for callback in litellm.callbacks: - if isinstance(callback, SlackAlerting): - callback_added = True - assert callback_added is True - - def test_flush_cache(model_list): """Test if the 'flush_cache' function is working correctly""" router = Router(model_list=model_list) @@ -2666,9 +1382,6 @@ def test_discard_drops_routing_group_selectors(): {"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "latency-based-routing"} ], ) - _, group_selector = router._get_routing_context("grouped-model") - assert any(callback is group_selector for callback in litellm.callbacks) - router.discard() assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) @@ -2712,27 +1425,6 @@ def test_routing_group_reload_keeps_every_hook_of_an_in_flight_callback_loop(): router.discard() -def _routers_seen_in_success_loop_discarding_first(first: Router, second: Router) -> Iterator[Router]: - for callback in litellm.success_callback: - if not isinstance(callback, MethodType): - continue - if callback.__self__ is first._core: - first.discard() - yield first - if callback.__self__ is second._core: - yield second - - -def test_discard_mid_success_loop_keeps_the_next_router_callback(): - first: Final = Router(model_list=_GROUPED_MODEL_LIST) - second: Final = Router(model_list=_GROUPED_MODEL_LIST) - - seen: Final = tuple(_routers_seen_in_success_loop_discarding_first(first, second)) - - assert seen == (first, second) - second.discard() - - @pytest.mark.asyncio async def test_discard_purges_optional_pre_call_checks_from_event_lists_after_serving_a_request(): router: Final = Router( @@ -2761,23 +1453,6 @@ def _is_registered(callback: object) -> bool: return any(registered is callback for registered in litellm.callbacks) -@pytest.mark.asyncio -async def test_discard_releases_per_request_override_selector(): - def build_override_then_discard() -> weakref.ReferenceType[object]: - router: Final = Router(model_list=_GROUPED_MODEL_LIST) - selector: Final = router._get_override_strategy_selector("usage-based-routing-v2") - assert selector is not None - router.discard() - return weakref.ref(selector) - - selector_ref: Final = build_override_then_discard() - await asyncio.sleep(0) - gc.collect() - - assert selector_ref() is None - - - def test_discard_registers_the_selector_of_a_same_strategy_survivor_built_without_a_model_list(): retired: Final = Router(model_list=[], routing_strategy="latency-based-routing") survivor: Final = Router(routing_strategy="latency-based-routing") @@ -2855,6 +1530,7 @@ async def test_discard_releases_its_router_budget_limiter(): assert budget_limiter_ref() is None + def test_initialize_assistants_endpoint(model_list): """Test if the 'initialize_assistants_endpoint' function is working correctly""" router = Router(model_list=model_list) @@ -2879,13 +1555,6 @@ def test_get_model_from_alias(model_list): assert model == "gpt-5-mini" -def test_get_deployment_by_litellm_model(model_list): - """Test if the 'get_deployment_by_litellm_model' function is working correctly""" - router = Router(model_list=model_list) - deployment = router._get_deployment_by_litellm_model(model="gpt-5-mini") - assert deployment is not None - - def test_get_pattern(model_list): router = Router(model_list=model_list) pattern = router.pattern_router.get_pattern(model="claude-3") @@ -2930,7 +1599,6 @@ def test_deployments_by_pattern(model_list): ], ) def test_pattern_match_deployment_set_model_name(user_request_model, model_name, litellm_model, expected_model): - from re import Match from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter @@ -2964,18 +1632,6 @@ def test_pattern_match_deployment_set_model_name(user_request_model, model_name, assert model["litellm_params"]["model"] == expected_model -@pytest.mark.parametrize( - "has_default_fallbacks, expected_result", - [(True, True), (False, False)], -) -def test_has_default_fallbacks(model_list, has_default_fallbacks, expected_result): - router = Router( - model_list=model_list, - default_fallbacks=(["my-default-fallback-model"] if has_default_fallbacks else None), - ) - assert router._has_default_fallbacks() is expected_result - - def test_add_optional_pre_call_checks(model_list): router = Router(model_list=model_list) @@ -2985,7 +1641,6 @@ def test_add_optional_pre_call_checks(model_list): @pytest.mark.asyncio async def test_async_callback_filter_deployments(model_list): - from litellm.router_strategy.budget_limiter import RouterBudgetLimiting router = Router(model_list=model_list) @@ -3001,20 +1656,6 @@ async def test_async_callback_filter_deployments(model_list): assert len(new_healthy_deployments) == len(healthy_deployments) -def test_cached_get_model_group_info(model_list): - """Test if the '_cached_get_model_group_info' function is working correctly with LRU cache""" - router = Router(model_list=model_list) - - result1 = router._cached_get_model_group_info("gpt-5-mini") - - result2 = router._cached_get_model_group_info("gpt-5-mini") - - assert result1 == result2 - - cache_info = router._cached_get_model_group_info.cache_info() - assert cache_info.hits > 0 - - def test_init_responses_api_endpoints(model_list): """Test if the '_init_responses_api_endpoints' function is working correctly""" from typing import Callable @@ -3072,23 +1713,6 @@ def test_mock_router_testing_params_str_to_bool_conversion( assert kwargs["other_param"] == "should_remain" -def test_is_auto_router_deployment(model_list): - """Test if the '_is_auto_router_deployment' function correctly identifies auto-router deployments""" - router = Router(model_list=model_list) - - litellm_params_auto = LiteLLM_Params(model="auto_router/my-auto-router") - assert router._is_auto_router_deployment(litellm_params_auto) is True - - litellm_params_regular = LiteLLM_Params(model="gpt-5-mini") - assert router._is_auto_router_deployment(litellm_params_regular) is False - - litellm_params_empty = LiteLLM_Params(model="") - assert router._is_auto_router_deployment(litellm_params_empty) is False - - litellm_params_contains = LiteLLM_Params(model="prefix_auto_router/something") - assert router._is_auto_router_deployment(litellm_params_contains) is False - - @patch("litellm.router_strategy.auto_router.auto_router.AutoRouter") def test_init_auto_router_deployment_success(mock_auto_router, model_list): """Test if the 'init_auto_router_deployment' function successfully initializes auto-router when all params provided""" @@ -3244,494 +1868,6 @@ def test_handle_clientside_credential_with_deployment_model_name(model_list): print("✓ _handle_clientside_credential test passed!") -def test_sync_generic_api_call_preserves_requested_model_group_in_logs(): - router = Router( - model_list=[ - { - "model_name": "claude-sonnet-4-6", - "litellm_params": { - "model": "bedrock/global.anthropic.claude-sonnet-4-6", - "aws_access_key_id": "test-access-key", - "aws_secret_access_key": "test-secret-key", - "aws_region_name": "us-west-2", - }, - } - ] - ) - - try: - captured_kwargs = {} - - def mock_original_function(**kwargs): - captured_kwargs.update(kwargs) - return {"status": "ok"} - - response = router._generic_api_call_with_fallbacks( - model="claude-sonnet-4-6", - original_function=mock_original_function, - ) - - assert response == {"status": "ok"} - assert captured_kwargs["model"] == "bedrock/global.anthropic.claude-sonnet-4-6" - assert captured_kwargs["litellm_metadata"]["model_group"] == "claude-sonnet-4-6" - assert captured_kwargs["litellm_metadata"]["deployment"] == "bedrock/global.anthropic.claude-sonnet-4-6" - finally: - router.discard() - - -def test_sync_generic_api_call_uses_request_kwargs_for_deployment_selection(): - router = Router( - model_list=[ - { - "model_name": "regional-model", - "litellm_params": { - "model": "anthropic/us-model", - "api_key": "test-api-key", - "region_name": "us", - }, - }, - { - "model_name": "regional-model", - "litellm_params": { - "model": "anthropic/eu-model", - "api_key": "test-api-key", - "region_name": "eu", - }, - }, - ], - enable_pre_call_checks=True, - ) - - try: - captured_kwargs = {} - - def mock_original_function(**kwargs): - captured_kwargs.update(kwargs) - return {"status": "ok"} - - response = router._generic_api_call_with_fallbacks( - model="regional-model", - original_function=mock_original_function, - messages=[{"role": "user", "content": "Hello from Europe"}], - allowed_model_region="eu", - ) - - assert response == {"status": "ok"} - assert captured_kwargs["model"] == "anthropic/eu-model" - finally: - router.discard() - - -@pytest.mark.parametrize( - "function_name, expected_metadata_key", - [ - ("acompletion", "metadata"), - ("_ageneric_api_call_with_fallbacks", "litellm_metadata"), - ("batch", "litellm_metadata"), - ("completion", "metadata"), - ("acreate_file", "litellm_metadata"), - ("aget_file", "litellm_metadata"), - ], -) -def test_handle_clientside_credential_metadata_loading(model_list, function_name, expected_metadata_key): - """Test that _handle_clientside_credential correctly loads metadata based on function name""" - router = Router(model_list=model_list) - - deployment = { - "model_name": "gpt-4.1", - "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, - "model_info": {"id": "original-id-123"}, - } - - kwargs = { - "api_key": "client_side_key", - "api_base": "https://api.openai.com/v1", - expected_metadata_key: {"model_group": "gpt-4.1", "custom_field": "test_value"}, - } - - result_deployment = router._handle_clientside_credential( - deployment=deployment, kwargs=kwargs, function_name=function_name - ) - - assert isinstance(result_deployment, Deployment) - - assert result_deployment.model_name == "gpt-4.1" - - assert result_deployment.litellm_params.api_key == "client_side_key" - assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1" - - assert result_deployment.model_info.id != "original-id-123" - assert result_deployment.model_info.original_model_id == "original-id-123" - - assert len(router.model_list) == len(model_list) - assert router.get_deployment(model_id=result_deployment.model_info.id) is None - - if function_name == "acompletion": - assert "metadata" in kwargs - assert "litellm_metadata" not in kwargs - elif function_name in [ - "_ageneric_api_call_with_fallbacks", - "batch", - "acreate_file", - "aget_file", - ]: - assert "litellm_metadata" in kwargs - - print(f"✓ Success with function_name '{function_name}' using '{expected_metadata_key}' metadata key") - - -@pytest.mark.parametrize( - "function_name, metadata_key", - [ - ("acompletion", "metadata"), - ("_ageneric_api_call_with_fallbacks", "litellm_metadata"), - ], -) -def test_handle_clientside_credential_metadata_variable_name(model_list, function_name, metadata_key): - """Test that _handle_clientside_credential uses the correct metadata variable name based on function name""" - from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name - - router = Router(model_list=model_list) - - expected_metadata_key = get_router_metadata_variable_name(function_name=function_name) - assert expected_metadata_key == metadata_key - - deployment = { - "model_name": "gpt-4.1", - "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, - "model_info": {"id": "original-id-456"}, - } - - kwargs = { - "api_key": "client_side_key", - "api_base": "https://api.openai.com/v1", - metadata_key: {"model_group": "gpt-4.1", "test_field": "test_value"}, - } - - result_deployment = router._handle_clientside_credential( - deployment=deployment, kwargs=kwargs, function_name=function_name - ) - - assert result_deployment.model_name == "gpt-4.1" - - assert result_deployment.litellm_params.api_key == "client_side_key" - assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1" - - print(f"✓ Success with function_name '{function_name}' correctly using '{metadata_key}' for metadata") - - -def test_handle_clientside_credential_with_responses_function(model_list): - """Test that _handle_clientside_credential works correctly with responses function name""" - router = Router(model_list=model_list) - - deployment = { - "model_name": "gpt-4.1", - "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, - "model_info": {"id": "original-id-responses"}, - } - - kwargs = { - "api_key": "client_side_key", - "api_base": "https://api.openai.com/v1", - "litellm_metadata": { - "model_group": "gpt-4.1", - "responses_field": "responses_value", - }, - } - - result_deployment = router._handle_clientside_credential( - deployment=deployment, - kwargs=kwargs, - function_name="_ageneric_api_call_with_fallbacks", - ) - - assert isinstance(result_deployment, Deployment) - assert result_deployment.model_name == "gpt-4.1" - assert result_deployment.litellm_params.api_key == "client_side_key" - assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1" - assert result_deployment.model_info.id != "original-id-responses" - assert result_deployment.model_info.original_model_id == "original-id-responses" - - assert len(router.model_list) == len(model_list) - assert router.get_deployment(model_id=result_deployment.model_info.id) is None - - print("✓ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata") - - -def test_handle_clientside_credential_still_registers_custom_pricing(model_list): - """A clientside-credential call must still price against the deployment's own - custom rate, even though the call's ephemeral deployment is never added to the - router (see LIT-7811): losing that registration would silently fall back to - public catalog pricing for every clientside-credential call on a deployment - with a custom rate configured.""" - router = Router(model_list=model_list) - deployment = { - "model_name": "gpt-4.1", - "litellm_params": { - "model": "gpt-4.1", - "api_key": "test_key", - "input_cost_per_token": 0.0001234, - "output_cost_per_token": 0.0005678, - }, - "model_info": {"id": "original-id-pricing"}, - } - kwargs = {"api_key": "client_side_key", "metadata": {"model_group": "gpt-4.1"}} - - result_deployment = router._handle_clientside_credential( - deployment=deployment, kwargs=kwargs, function_name="acompletion" - ) - - registered = litellm.model_cost.get(result_deployment.model_info.id) - assert registered is not None - assert registered["input_cost_per_token"] == 0.0001234 - assert registered["output_cost_per_token"] == 0.0005678 - - -def test_register_deployment_pricing_direct_call(): - """Direct-call unit test for the pricing-registration helper `_handle_clientside_credential` - relies on, so it prices a deployment that is deliberately never added to `self.model_list`.""" - deployment = Deployment( - model_name="gpt-4.1", - litellm_params=LiteLLM_Params( - model="gpt-4.1", - api_key="test_key", - input_cost_per_token=0.0009999, - ), - model_info=ModelInfo(id="direct-call-pricing-id"), - ) - - Router._register_deployment_pricing(deployment=deployment) - - assert litellm.model_cost["direct-call-pricing-id"]["input_cost_per_token"] == 0.0009999 - - -def test_get_metadata_variable_name_from_kwargs(model_list): - """ - Test _get_metadata_variable_name_from_kwargs method returns correct metadata variable name based on kwargs content. - """ - router = Router(model_list=model_list) - - kwargs_with_litellm_metadata = { - "litellm_metadata": {"user": "test"}, - "metadata": {"other": "data"}, - } - result = router._get_metadata_variable_name_from_kwargs(kwargs_with_litellm_metadata) - assert result == "litellm_metadata" - - kwargs_with_metadata_only = {"metadata": {"user": "test"}} - result = router._get_metadata_variable_name_from_kwargs(kwargs_with_metadata_only) - assert result == "metadata" - - kwargs_empty = {} - result = router._get_metadata_variable_name_from_kwargs(kwargs_empty) - assert result == "metadata" - - kwargs_other = { - "model": "gpt-5.5", - "messages": [{"role": "user", "content": "hello"}], - } - result = router._get_metadata_variable_name_from_kwargs(kwargs_other) - assert result == "metadata" - - -@pytest.fixture -def search_tools(): - """Fixture for search tools configuration""" - return [ - { - "search_tool_name": "test-search-tool", - "litellm_params": { - "search_provider": "perplexity", - "api_key": "test-api-key", - "api_base": "https://api.perplexity.ai", - "mode": "turbo", - }, - }, - { - "search_tool_name": "test-search-tool", - "litellm_params": { - "search_provider": "perplexity", - "api_key": "test-api-key-2", - "api_base": "https://api.perplexity.ai", - "mode": "turbo", - }, - }, - ] - - -@pytest.mark.asyncio -async def test_asearch_with_fallbacks(search_tools): - """ - Test _asearch_with_fallbacks method of Router. - - Tests that the _asearch_with_fallbacks method correctly: - - Accepts search parameters - - Calls async_function_with_fallbacks with correct configuration - - Returns SearchResponse - """ - from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult - - router = Router(search_tools=search_tools) - - mock_response = SearchResponse( - object="search", - results=[ - SearchResult( - title="Test Result", - url="https://example.com", - snippet="Test snippet content", - ) - ], - ) - - with patch.object(router, "async_function_with_fallbacks", new_callable=AsyncMock) as mock_fallbacks: - mock_fallbacks.return_value = mock_response - - async def mock_asearch(**kwargs): - return mock_response - - response = await router._asearch_with_fallbacks( - original_function=mock_asearch, - search_tool_name="test-search-tool", - query="test query", - max_results=5, - ) - - assert mock_fallbacks.called - - assert isinstance(response, SearchResponse) - assert response.object == "search" - assert len(response.results) == 1 - assert response.results[0].title == "Test Result" - - -@pytest.mark.asyncio -async def test_asearch_with_fallbacks_helper(search_tools): - """ - Test _asearch_with_fallbacks_helper method of Router. - - Tests that the _asearch_with_fallbacks_helper method correctly: - - Selects a search tool from available options - - Calls the original search function with correct provider parameters - - Returns SearchResponse - """ - from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult - - router = Router(search_tools=search_tools) - - mock_response = SearchResponse( - object="search", - results=[ - SearchResult( - title="Helper Test Result", - url="https://example.com/helper", - snippet="Helper test snippet", - ) - ], - ) - - async def mock_original_function(**kwargs): - - assert "search_provider" in kwargs - assert kwargs["search_provider"] == "perplexity" - assert "api_key" in kwargs - assert kwargs["mode"] == "turbo" - assert kwargs["query"] == "helper test query" - return mock_response - - response = await router._asearch_with_fallbacks_helper( - model="test-search-tool", - original_generic_function=mock_original_function, - query="helper test query", - max_results=3, - ) - - assert isinstance(response, SearchResponse) - assert response.object == "search" - assert len(response.results) == 1 - assert response.results[0].title == "Helper Test Result" - assert response.results[0].url == "https://example.com/helper" - - -@pytest.mark.asyncio -async def test_asearch_with_fallbacks_helper_missing_search_tool(): - """ - Test _asearch_with_fallbacks_helper raises error when search tool not found. - - Tests that the helper method raises a ValueError when the requested - search tool name doesn't exist in the router's search_tools configuration. - """ - - router = Router(model_list=[]) - - async def mock_original_function(**kwargs): - return None - - with pytest.raises(ValueError, match="Search tool 'nonexistent-tool' not found"): - await router._asearch_with_fallbacks_helper( - model="nonexistent-tool", - original_generic_function=mock_original_function, - query="test query", - ) - - -@pytest.mark.asyncio -async def test_asearch_with_fallbacks_helper_missing_search_provider(): - """ - Test _asearch_with_fallbacks_helper raises error when search_provider not configured. - - Tests that the helper method raises a ValueError when a search tool - is found but doesn't have search_provider in its litellm_params. - """ - - search_tools_bad = [ - { - "search_tool_name": "bad-tool", - "litellm_params": {"api_key": "test-key"}, - } - ] - - router = Router(search_tools=search_tools_bad) - - async def mock_original_function(**kwargs): - return None - - with pytest.raises(ValueError, match="search_provider not found in litellm_params"): - await router._asearch_with_fallbacks_helper( - model="bad-tool", - original_generic_function=mock_original_function, - query="test query", - ) - - -def test_get_first_default_fallback(): - """Test _get_first_default_fallback method""" - - model_list = [ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini", "api_key": "fake-key"}, - } - ] - - router = Router(model_list=model_list, fallbacks=[{"*": ["gpt-5-mini"]}]) - - result = router._get_first_default_fallback() - assert result == "gpt-5-mini" - - router_no_fallbacks = Router(model_list=model_list) - result = router_no_fallbacks._get_first_default_fallback() - assert result is None - - router_no_default = Router(model_list=model_list, fallbacks=[{"gpt-5.5": ["gpt-5-mini"]}]) - result = router_no_default._get_first_default_fallback() - assert result is None - - router_empty_list = Router(model_list=model_list, fallbacks=[{"*": []}]) - result = router_empty_list._get_first_default_fallback() - assert result is None - - def test_resolve_model_name_from_model_id(): """Test resolve_model_name_from_model_id function with various scenarios""" @@ -3872,6 +2008,7 @@ def test_resolve_model_name_from_model_id_keeps_a_wildcard_name_only_for_the_ids router = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-fake"}}]) assert router.resolve_model_name_from_model_id(model_id, custom_llm_provider) == expected + def test_get_valid_args(): """Test get_valid_args static method returns valid Router.__init__ arguments""" @@ -3922,89 +2059,6 @@ def test_get_router_model_info_with_deployment_object(): assert isinstance(model_info, dict) -def test_deployment_has_budget_limits(): - router = Router(model_list=[]) - - with_budget = Deployment( - model_name="budgeted-model", - litellm_params=LiteLLM_Params( - model="openai/gpt-4o-mini", - max_budget=0.001, - budget_duration="1d", - ), - model_info=ModelInfo(id="budget-deployment-id"), - ) - without_budget = Deployment( - model_name="unbudgeted-model", - litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"), - model_info=ModelInfo(id="no-budget-deployment-id"), - ) - - assert router._deployment_has_budget_limits(deployment=with_budget) is True - assert router._deployment_has_budget_limits(deployment=without_budget) is False - - -def test_sync_deployment_budget_config(monkeypatch): - import asyncio - - monkeypatch.setattr(asyncio, "create_task", lambda coro: None) - - router = Router(model_list=[], optional_pre_call_checks=[]) - deployment = Deployment( - model_name="dynamic-budget-model", - litellm_params=LiteLLM_Params( - model="openai/gpt-4o-mini", - api_key="fake-key", - max_budget=0.000000000001, - budget_duration="1d", - ), - model_info=ModelInfo(id="runtime-budget-deployment"), - ) - - router._sync_deployment_budget_config(deployment=deployment) - - budget_limiter = router.get_router_deployment_budget_limiter() - assert budget_limiter is not None - config = budget_limiter._get_budget_config_for_deployment("runtime-budget-deployment") - assert config is not None - assert config.max_budget == 0.000000000001 - - -def test_sync_deployment_budget_config_clears_removed_limits(monkeypatch): - import asyncio - - monkeypatch.setattr(asyncio, "create_task", lambda coro: None) - - router = Router(model_list=[], optional_pre_call_checks=[]) - model_id = "runtime-budget-deployment" - budgeted = Deployment( - model_name="dynamic-budget-model", - litellm_params=LiteLLM_Params( - model="openai/gpt-4o-mini", - api_key="fake-key", - max_budget=0.000000000001, - budget_duration="1d", - ), - model_info=ModelInfo(id=model_id), - ) - unbudgeted = Deployment( - model_name="dynamic-budget-model", - litellm_params=LiteLLM_Params( - model="openai/gpt-4o-mini", - api_key="fake-key", - ), - model_info=ModelInfo(id=model_id), - ) - - router._sync_deployment_budget_config(deployment=budgeted) - budget_limiter = router.get_router_deployment_budget_limiter() - assert budget_limiter is not None - assert budget_limiter._get_budget_config_for_deployment(model_id) is not None - - router._sync_deployment_budget_config(deployment=unbudgeted) - assert budget_limiter._get_budget_config_for_deployment(model_id) is None - - def test_upsert_deployment_clears_stale_budget_config(monkeypatch): import asyncio diff --git a/tests/unit/test_router/test_router_max_parallel_requests.py b/tests/unit/test_router/test_router_max_parallel_requests.py index fce273cd1ab..1b350f31f23 100644 --- a/tests/unit/test_router/test_router_max_parallel_requests.py +++ b/tests/unit/test_router/test_router_max_parallel_requests.py @@ -4,7 +4,6 @@ from typing import Final import litellm import pytest -from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit from litellm.utils import calculate_max_parallel_requests @@ -36,53 +35,6 @@ def test_scenario( assert calculated == expected -@pytest.mark.parametrize( - ("max_parallel_requests", "tpm", "rpm", "default_max_parallel_requests", "expected"), - [ - pytest.param(10, 300_000, 30, 40, 10, id="explicit-limit"), - pytest.param(None, 300_000, 30, 40, 30, id="rpm"), - pytest.param(None, 300_000, None, 40, 1_800, id="tpm"), - pytest.param(None, 20, None, 40, 1, id="minimum-tpm-limit"), - pytest.param(None, None, None, 40, 40, id="router-default"), - pytest.param(None, None, None, None, None, id="unset"), - ], -) -def test_setting_mpr_limits_per_model( - max_parallel_requests: int | None, - tpm: int | None, - rpm: int | None, - default_max_parallel_requests: int | None, - expected: int | None, -) -> None: - deployment: Final = { - "model_name": "limited", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "test-key", - "max_parallel_requests": max_parallel_requests, - "tpm": tpm, - "rpm": rpm, - }, - "model_info": {"id": "limited-deployment"}, - } - router: Final = litellm.Router( - model_list=[deployment], - default_max_parallel_requests=default_max_parallel_requests, - ) - limit: Final = router._get_client( - deployment=deployment, - kwargs={}, - client_type="max_parallel_requests", - ) - - if expected is None: - assert limit is None - return - - assert isinstance(limit, MaxParallelRequestsLimit) - assert limit.max_parallel_requests == expected - - @pytest.mark.asyncio async def test_max_parallel_requests_rpm_rate_limiting() -> None: router: Final = litellm.Router( diff --git a/tests/unit/test_router/test_router_provider_endpoints.py b/tests/unit/test_router/test_router_provider_endpoints.py index d2331c00804..baf3e94f1f0 100644 --- a/tests/unit/test_router/test_router_provider_endpoints.py +++ b/tests/unit/test_router/test_router_provider_endpoints.py @@ -2,6 +2,7 @@ import asyncio import contextlib import io import json +import time import uuid from collections.abc import Callable, Iterator from datetime import datetime, timezone @@ -17,6 +18,7 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import UserAPIKeyAuth +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.router_callbacks.track_deployment_metrics import ( get_deployment_successes_for_current_minute, ) @@ -51,9 +53,11 @@ def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: @pytest.fixture -def router_minute_pinned(monkeypatch: pytest.MonkeyPatch) -> None: - pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) - monkeypatch.setattr("litellm.router_python_core.get_utc_datetime", lambda: pinned) +def router_minute_pinned() -> None: + """The usage counters these tests read are per minute: start well inside one so a test never spans two.""" + seconds_left: Final = 60 - datetime.now(timezone.utc).second + if seconds_left < 10: + time.sleep(seconds_left + 0.1) class _RouterLoggingCapture(CustomLogger): @@ -145,14 +149,14 @@ async def test_router_transcription_reuses_router_level_client_for_each_deployme model="whisper", file=("speech.wav", io.BytesIO(b"offline audio"), "audio/wav") ) public_client, public_logging = await capture.next_event() - internal_response: Final = await router._atranscription( + second_response: Final = await router.atranscription( model="whisper", file=("speech.wav", io.BytesIO(b"offline audio"), "audio/wav") ) - internal_client, _ = await capture.next_event() + second_client, _ = await capture.next_event() upstream_requests: Final = tuple(call.request for call in route.calls) assert public_client == str(router_level_client) - assert internal_client == str(router_level_client) + assert second_client == str(router_level_client) assert tuple(str(request.url) for request in upstream_requests) == (route_url, route_url) assert all( request.headers.get("content-type", "").startswith("multipart/form-data") @@ -163,7 +167,7 @@ async def test_router_transcription_reuses_router_level_client_for_each_deployme assert isinstance(public_logging, dict) assert public_logging.get("model_group") == "whisper" assert response.text == "hello" - assert internal_response.text == "hello" + assert second_response.text == "hello" @pytest.mark.asyncio @@ -197,7 +201,7 @@ async def test_router_speech_returns_binary_content_and_logs_model_group( @pytest.mark.asyncio -async def test_router_rerank_returns_valid_response_from_public_and_underlying_calls( +async def test_router_rerank_returns_valid_response_on_every_call( respx_mock: respx.MockRouter, ) -> None: router: Final = Router( @@ -223,7 +227,7 @@ async def test_router_rerank_returns_valid_response_from_public_and_underlying_c documents=["hello", "world"], top_n=1, ) - underlying_response: Final = await router._arerank( + second_response: Final = await router.arerank( model="cohere-rerank", query="hello", documents=["hello", "world"], @@ -242,9 +246,9 @@ async def test_router_rerank_returns_valid_response_from_public_and_underlying_c public_validated: Final = RerankResponse.model_validate(public_response) assert public_validated.id == "rerank-1" assert public_validated.results[0]["relevance_score"] == 0.9 - underlying_validated: Final = RerankResponse.model_validate(underlying_response) - assert underlying_validated.id == "rerank-1" - assert underlying_validated.results[0]["relevance_score"] == 0.9 + second_validated: Final = RerankResponse.model_validate(second_response) + assert second_validated.id == "rerank-1" + assert second_validated.results[0]["relevance_score"] == 0.9 @pytest.mark.asyncio @@ -317,9 +321,9 @@ async def test_router_image_generation_returns_valid_image_response( ) response: Final = ( - router._image_generation(model="gpt-image-1", prompt="a cat") + router.image_generation(model="gpt-image-1", prompt="a cat") if sync_mode - else await router._aimage_generation(model="gpt-image-1", prompt="a cat") + else await router.aimage_generation(model="gpt-image-1", prompt="a cat") ) assert route.call_count == 1 @@ -516,24 +520,6 @@ def test_router_routing_strategy_init_accepts_provider_budget_strategy(strategy: assert router.get_settings()["routing_strategy"] == "provider-budget-routing" -def test_router_track_deployment_metrics_updates_observable_usage() -> None: - router: Final = Router( - model_list=[ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"}, - "model_info": {"id": "metrics-deployment"}, - } - ] - ) - deployment: Final = router.model_list[0] - - router._track_deployment_metrics(deployment=deployment, parent_otel_span=None) - router._track_deployment_metrics(deployment=deployment, parent_otel_span=None) - - assert router.cache.get_cache(key="metrics-deployment", local_only=True) == 2 - - class _GatedIncrementCache(DualCache): def __init__(self) -> None: super().__init__(in_memory_cache=InMemoryCache()) @@ -548,6 +534,12 @@ class _GatedIncrementCache(DualCache): self.deployment_success_incremented.set() return result + async def async_increment_cache(self, key: str, value: float, *args: object, **kwargs: object) -> float | None: + result: Final = await super().async_increment_cache(key, value, *args, **kwargs) + if key.endswith(":successes"): + self.deployment_success_incremented.set() + return result + async def async_increment_cache_pipeline( self, increment_list: list[RedisPipelineIncrementOperation], @@ -581,6 +573,12 @@ class _UnavailableIncrementCache(DualCache): self.deployment_success_incremented.set() return result + async def async_increment_cache(self, key: str, value: float, *args: object, **kwargs: object) -> float | None: + result: Final = await super().async_increment_cache(key, value, *args, **kwargs) + if key.endswith(":successes"): + self.deployment_success_incremented.set() + return result + async def async_increment_cache_pipeline( self, increment_list: list[RedisPipelineIncrementOperation], @@ -637,7 +635,6 @@ async def test_router_failed_pre_header_increment_clears_counted_tokens_stamp() await asyncio.wait_for(cache.first_increment_started.wait(), timeout=EVENT_TIMEOUT_SECONDS) await asyncio.wait_for(cache.deployment_success_incremented.wait(), timeout=EVENT_TIMEOUT_SECONDS) - assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30 assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1 assert cache.increment_calls == 1 cache.release_first_increment.set() @@ -662,7 +659,7 @@ ASSISTANT_RESPONSE: Final = { @pytest.mark.asyncio -async def test_router_assistants_endpoint_factory_invokes_provider( +async def test_router_create_assistants_invokes_provider( respx_mock: respx.MockRouter, ) -> None: router: Final = Router(model_list=[]) @@ -670,8 +667,7 @@ async def test_router_assistants_endpoint_factory_invokes_provider( 200, json={**ASSISTANT_RESPONSE, "id": "asst-1"} ) - response: Final = await router._pass_through_assistants_endpoint_factory( - original_function=litellm.acreate_assistants, + response: Final = await router.acreate_assistants( custom_llm_provider="openai", model="gpt-4o-mini", api_key="sk-fake", @@ -684,29 +680,7 @@ async def test_router_assistants_endpoint_factory_invokes_provider( @pytest.mark.asyncio -async def test_router_factory_function_returns_invokable_assistants_wrapper( - respx_mock: respx.MockRouter, -) -> None: - router: Final = Router(model_list=[]) - route: Final = respx_mock.post("https://api.openai.com/v1/assistants").respond( - 200, json={**ASSISTANT_RESPONSE, "id": "asst-2"} - ) - wrapper: Final = router.factory_function(litellm.acreate_assistants, call_type="assistants") - - response: Final = await wrapper( - custom_llm_provider="openai", - model="gpt-4o-mini", - api_key="sk-fake", - name="offline", - ) - - assert route.call_count == 1 - assert json.loads(route.calls[0].request.content) == {"model": "gpt-4o-mini", "name": "offline"} - assert response.id == "asst-2" - - -@pytest.mark.asyncio -async def test_router_moderation_endpoint_factory_invokes_default_model( +async def test_router_moderation_without_a_model_invokes_the_default_model( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -721,9 +695,7 @@ async def test_router_moderation_endpoint_factory_invokes_default_model( }, ) - response: Final = await router._pass_through_moderation_endpoint_factory( - original_function=litellm.amoderation, - custom_llm_provider="openai", + response: Final = await router.amoderation( input="hello", model=None, api_key="sk-fake", @@ -801,7 +773,11 @@ async def test_video_status_on_a_wildcard_polls_the_deployment_that_created_the_ ) -> None: router: Final = Router( model_list=[ - {"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": key}, "model_info": {"id": dep}} + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": key}, + "model_info": {"id": dep}, + } for dep, key in (("wc-a", "key-a"), ("wc-b", "key-b")) ] ) @@ -946,6 +922,43 @@ async def test_video_lookup_404_never_cools_down_the_deployment_that_answered(re assert await router.cooldown_cache.async_get_active_cooldowns(["dep-a", "dep-b"], parent_otel_span=None) == [] +@pytest.mark.parametrize( + ("endpoint", "method", "path"), + [ + ("aget_responses", "GET", ""), + ("adelete_responses", "DELETE", ""), + ("acancel_responses", "POST", "/cancel"), + ("alist_input_items", "GET", "/input_items"), + ], +) +@pytest.mark.asyncio +async def test_response_lookup_404_never_cools_down_the_deployment_that_answered( + respx_mock: respx.MockRouter, endpoint: str, method: str, path: str +) -> None: + respx_mock.route(method=method, url=f"https://api.openai.com/v1/responses/resp_gone{path}").respond( + 404, json={"error": {"code": "NotFound", "message": "Response not found"}} + ) + router: Final = Router( + model_list=[ + { + "model_name": "gpt", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": api_key}, + "model_info": {"id": deployment_id}, + } + for api_key, deployment_id in (("key-a", "dep-a"), ("key-b", "dep-b")) + ], + num_retries=0, + allowed_fails=1, + ) + response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "dep-b", "resp_gone") + + for _ in range(3): + with pytest.raises(litellm.NotFoundError): + await getattr(router, endpoint)(response_id=response_id) + + assert await router.cooldown_cache.async_get_active_cooldowns(["dep-a", "dep-b"], parent_otel_span=None) == [] + + @pytest.mark.asyncio async def test_chat_404_still_cools_down_the_deployment(respx_mock: respx.MockRouter) -> None: respx_mock.post("https://api.openai.com/v1/chat/completions").respond( diff --git a/tests/unit/test_router/test_router_retries.py b/tests/unit/test_router/test_router_retries.py index a5664c1348f..857b5052d6c 100644 --- a/tests/unit/test_router/test_router_retries.py +++ b/tests/unit/test_router/test_router_retries.py @@ -169,9 +169,7 @@ async def test_router_retries_errors(sync_mode: bool, error_type: str, monkeypat async def test_router_retry_policy(error_type): from litellm.router import AllowedFailsPolicy, RetryPolicy - retry_policy = RetryPolicy( - ContentPolicyViolationErrorRetries=3, AuthenticationErrorRetries=0 - ) + retry_policy = RetryPolicy(ContentPolicyViolationErrorRetries=3, AuthenticationErrorRetries=0) allowed_fails_policy = AllowedFailsPolicy( ContentPolicyViolationErrorAllowedFails=1000, @@ -361,9 +359,7 @@ def test_retry_rate_limit_error_with_healthy_deployments(): # Act & Assert try: - response = router.should_retry_this_error( - error=rate_limit_error, healthy_deployments=healthy_deployments - ) + response = router.should_retry_this_error(error=rate_limit_error, healthy_deployments=healthy_deployments) print("response from should_retry_this_error: ", response) except Exception as e: pytest.fail( @@ -394,9 +390,7 @@ def test_do_retry_rate_limit_error_with_no_fallbacks_and_no_healthy_deployments( # Act & Assert try: - response = router.should_retry_this_error( - error=rate_limit_error, healthy_deployments=healthy_deployments - ) + response = router.should_retry_this_error(error=rate_limit_error, healthy_deployments=healthy_deployments) pytest.fail("Should have raised an error") except Exception as e: print("got an exception", e) @@ -438,9 +432,7 @@ def test_raise_context_window_exceeded_error(): healthy_deployments=None, context_window_fallbacks=context_window_fallbacks, ) - pytest.fail( - "Expected to raise context window exceeded error -> trigger fallback" - ) + pytest.fail("Expected to raise context window exceeded error -> trigger fallback") except Exception as e: pass @@ -480,116 +472,11 @@ def test_raise_context_window_exceeded_error_no_retry(): healthy_deployments=None, context_window_fallbacks=context_window_fallbacks, ) - assert ( - response == True - ), "Should not have raised exception since we do not have context window fallbacks" + assert response == True, "Should not have raised exception since we do not have context window fallbacks" except litellm.ContextWindowExceededError: pass -@pytest.mark.parametrize("num_deployments, expected_timeout", [(1, 60), (2, 0.0)]) -def test_timeout_for_rate_limit_error_with_healthy_deployments( - num_deployments, expected_timeout -): - """ - Test 1. Timeout is 0.0 when RateLimit Error and healthy deployments are > 0 - """ - cooldown_time = 60 - rate_limit_error = litellm.RateLimitError( - message="{RouterErrors.no_deployments_available.value}. 12345 Passed model={model_group}. Deployments={deployment_dict}", - llm_provider="", - model="gpt-3.5-turbo", - response=httpx.Response( - status_code=429, - content="", - headers={"retry-after": str(cooldown_time)}, # type: ignore - request=httpx.Request(method="tpm_rpm_limits", url="https://github.com/BerriAI/litellm"), # type: ignore - ), - ) - model_list = [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), - }, - } - ] - if num_deployments == 2: - model_list.append( - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-3.5-turbo"}, - } - ) - - router = litellm.Router(model_list=model_list) - - _timeout = router._time_to_sleep_before_retry( - e=rate_limit_error, - remaining_retries=2, - num_retries=2, - healthy_deployments=[ - { - "model_name": "gpt-4", - "litellm_params": { - "api_key": "my-key", - "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com", - "model": "azure/gpt-4.1-mini", - }, - "model_info": { - "id": "0e30bc8a63fa91ae4415d4234e231b3f9e6dd900cac57d118ce13a720d95e9d6", - "db_model": False, - }, - } - ], - all_deployments=model_list, - ) - - if expected_timeout == 0.0: - assert _timeout == expected_timeout - else: - assert _timeout > 0.0 - - -def test_timeout_for_rate_limit_error_with_no_healthy_deployments(): - """ - Test 2. Timeout is > 0.0 when RateLimit Error and healthy deployments == 0 - """ - healthy_deployments = [] - model_list = [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), - }, - } - ] - - router = litellm.Router(model_list=model_list) - - _timeout = router._time_to_sleep_before_retry( - e=rate_limit_error, - remaining_retries=4, - num_retries=4, - healthy_deployments=healthy_deployments, - all_deployments=model_list, - ) - - print( - "timeout=", - _timeout, - "error is rate_limit_error and there are no healthy deployments", - ) - - assert _timeout > 0.0 - - def test_no_retry_for_not_found_error_404(): healthy_deployments = [] @@ -614,9 +501,7 @@ def test_no_retry_for_not_found_error_404(): llm_provider="azure", ) try: - response = router.should_retry_this_error( - error=error, healthy_deployments=healthy_deployments - ) + response = router.should_retry_this_error(error=error, healthy_deployments=healthy_deployments) pytest.fail( "Should have raised an exception 404 NotFoundError should never be retried, it's typically model_not_found error" ) @@ -652,12 +537,8 @@ def test_no_retry_for_bad_request_error_400(): llm_provider="azure", ) try: - response = router.should_retry_this_error( - error=error, healthy_deployments=healthy_deployments - ) - pytest.fail( - "Should have raised BadRequestError - 400 errors should never be retried" - ) + response = router.should_retry_this_error(error=error, healthy_deployments=healthy_deployments) + pytest.fail("Should have raised BadRequestError - 400 errors should never be retried") except litellm.BadRequestError as e: print("Correctly raised BadRequestError without retry:", e) @@ -693,12 +574,8 @@ def test_no_retry_for_unprocessable_entity_error_422(): ), ) try: - response = router.should_retry_this_error( - error=error, healthy_deployments=healthy_deployments - ) - pytest.fail( - "Should have raised UnprocessableEntityError - 422 errors should never be retried" - ) + response = router.should_retry_this_error(error=error, healthy_deployments=healthy_deployments) + pytest.fail("Should have raised UnprocessableEntityError - 422 errors should never be retried") except litellm.UnprocessableEntityError as e: print("Correctly raised UnprocessableEntityError without retry:", e) @@ -727,9 +604,7 @@ def test_no_retry_when_no_healthy_deployments(): timeout_error, ]: try: - response = router.should_retry_this_error( - error=error, healthy_deployments=healthy_deployments - ) + response = router.should_retry_this_error(error=error, healthy_deployments=healthy_deployments) pytest.fail( "Should have raised an exception, there's no point retrying an error when there are 0 healthy deployments" ) @@ -739,8 +614,6 @@ def test_no_retry_when_no_healthy_deployments(): @pytest.mark.asyncio async def test_router_retries_model_specific_and_global(): - from unittest.mock import MagicMock, patch - litellm.num_retries = 0 router = Router( model_list=[ @@ -748,28 +621,21 @@ async def test_router_retries_model_specific_and_global(): "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), + "api_key": "sk-test", "num_retries": 1, }, } ] ) - with patch.object( - router, "_time_to_sleep_before_retry" - ) as mock_async_function_with_retries: - try: - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello, how are you?"}], - mock_response="litellm.RateLimitError", - ) - except Exception as e: - print("got exception", e) + with pytest.raises(litellm.RateLimitError) as exc_info: + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello, how are you?"}], + mock_response="litellm.RateLimitError", + ) - mock_async_function_with_retries.assert_called_once() - - assert mock_async_function_with_retries.call_args.kwargs["num_retries"] == 1 + assert exc_info.value.num_retries == 1 @pytest.mark.usefixtures("fake_provider_credentials") @@ -817,68 +683,28 @@ async def test_router_retry_num_retries_tracking(): This verifies the fix for the bug where num_retries was incorrectly set to current_attempt (0-indexed) instead of the actual number of retries attempted. """ - from unittest.mock import AsyncMock, patch - router = Router( model_list=[ { "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), - }, + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}, } ], - num_retries=3, # Set at router level to ensure it's used + num_retries=2, ) - # Mock make_call to always raise a RateLimitError - async def mock_make_call(*args, **kwargs): - raise litellm.RateLimitError( - message="Rate limit exceeded", + with pytest.raises(litellm.RateLimitError) as exc_info: + await router.acompletion( model="gpt-3.5-turbo", - llm_provider="openai", + messages=[{"role": "user", "content": "Hello"}], + mock_response="litellm.RateLimitError", ) - with patch.object(router, "make_call", side_effect=mock_make_call): - with patch.object( - router, - "_async_get_healthy_deployments", - return_value=( - [{"model_info": {"id": "test-id"}}], - [{"model_info": {"id": "test-id"}}], - ), - ): - with patch.object( - router, "_time_to_sleep_before_retry", return_value=0.01 - ): # Fast retries for testing - with pytest.raises(litellm.RateLimitError) as exc_info: - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello"}], - ) - e = exc_info.value - assert hasattr( - e, "num_retries" - ), "Exception should have num_retries attribute" - assert hasattr( - e, "max_retries" - ), "Exception should have max_retries attribute" - assert ( - e.num_retries == 3 - ), f"Expected num_retries to be 3, got {e.num_retries}" - assert ( - e.max_retries == 3 - ), f"Expected max_retries to be 3, got {e.max_retries}" - - # Verify the error message includes correct retry information - error_str = str(e) - assert ( - "LiteLLM Retried: 3 times" in error_str - ), f"Error message should indicate 3 retries: {error_str}" - assert ( - "LiteLLM Max Retries: 3" in error_str - ), f"Error message should show max retries: {error_str}" + e = exc_info.value + assert e.num_retries == 2 + assert e.max_retries == 2 + assert "LiteLLM Retried: 2 times" in str(e) + assert "LiteLLM Max Retries: 2" in str(e) @pytest.mark.asyncio @@ -886,51 +712,25 @@ async def test_router_retry_num_retries_single_retry(): """ Test num_retries tracking with a single retry to verify edge case handling. """ - from unittest.mock import patch - router = Router( model_list=[ { "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), - }, + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}, } ], - num_retries=1, # Set at router level - single retry + num_retries=1, ) - # Mock make_call to always raise a Timeout error - async def mock_make_call(*args, **kwargs): - raise litellm.Timeout( - message="Request timed out", + with pytest.raises(litellm.InternalServerError) as exc_info: + await router.acompletion( model="gpt-3.5-turbo", - llm_provider="openai", + messages=[{"role": "user", "content": "Hello"}], + mock_response="litellm.InternalServerError", ) - with patch.object(router, "make_call", side_effect=mock_make_call): - with patch.object( - router, - "_async_get_healthy_deployments", - return_value=( - [{"model_info": {"id": "test-id"}}], - [{"model_info": {"id": "test-id"}}], - ), - ): - with patch.object(router, "_time_to_sleep_before_retry", return_value=0.01): - with pytest.raises(litellm.Timeout) as exc_info: - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello"}], - ) - e = exc_info.value - assert ( - e.num_retries == 1 - ), f"Expected num_retries to be 1, got {e.num_retries}" - assert ( - e.max_retries == 1 - ), f"Expected max_retries to be 1, got {e.max_retries}" + assert exc_info.value.num_retries == 1 + assert exc_info.value.max_retries == 1 class MyCustomHandler(CustomLogger): @@ -940,18 +740,14 @@ class MyCustomHandler(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): print(f"Pre-API Call") - print( - f"previous_models: {kwargs['litellm_params']['metadata'].get('previous_models', None)}" - ) + print(f"previous_models: {kwargs['litellm_params']['metadata'].get('previous_models', None)}") self.previous_models = len( kwargs["litellm_params"]["metadata"].get("previous_models", []) ) # {"previous_models": [{"model": litellm_model_name, "exception_type": AuthenticationError, "exception_string": }]} print(f"self.previous_models: {self.previous_models}") def log_post_api_call(self, kwargs, response_obj, start_time, end_time): - print( - f"Post-API Call - response object: {response_obj}; model: {kwargs['model']}" - ) + print(f"Post-API Call - response object: {response_obj}; model: {kwargs['model']}") def log_stream_event(self, kwargs, response_obj, start_time, end_time): print(f"On Stream") diff --git a/tests/unit/test_router/test_router_timeout.py b/tests/unit/test_router/test_router_timeout.py index 8fc32620f71..3fafb9e0f04 100644 --- a/tests/unit/test_router/test_router_timeout.py +++ b/tests/unit/test_router/test_router_timeout.py @@ -84,73 +84,26 @@ def test_router_timeout_with_retries_anthropic_model(num_retries, expected_call_ assert mock_client.call_count == expected_call_count -@pytest.mark.parametrize( - "stream", - [ - True, - False, - ], -) -def test_unit_test_streaming_timeout(stream): - import os - from dotenv import load_dotenv - import litellm - from litellm.router import Router, RetryPolicy, AllowedFailsPolicy +async def _timeout_sent_to_the_provider_call(router: Router, request_kwargs: Mapping[str, object]) -> object: + sent: Final[list[object]] = [] + real_acompletion: Final = litellm.acompletion - litellm.set_verbose = True + async def spy(*args: object, **kwargs: object) -> object: + # A stream is bounded by the stream timeout litellm.acompletion is handed, any other call by its timeout + stream_timeout = kwargs.get("stream_timeout") + sent.append( + stream_timeout if kwargs.get("stream") is True and stream_timeout is not None else kwargs.get("timeout") + ) + return await real_acompletion(*args, **kwargs) - model_list = [ - { - "model_name": "llama3", - "litellm_params": { - "model": "watsonx/meta-llama/llama-3-1-8b-instruct", - "api_base": os.getenv("WATSONX_URL_US_SOUTH"), - "api_key": os.getenv("WATSONX_API_KEY"), - "project_id": os.getenv("WATSONX_PROJECT_ID_US_SOUTH"), - "timeout": 0.01, - "stream_timeout": 0.0000001, - }, - }, - { - "model_name": "bedrock-anthropic", - "litellm_params": { - "model": "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0", - "timeout": 0.01, - "stream_timeout": 0.0000001, - }, - }, - { - "model_name": "llama3-fallback", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), - }, - }, - ] - - router = Router(model_list=model_list) - - stream_timeout = 0.0000001 - normal_timeout = 0.01 - - args = { - "kwargs": {"stream": stream}, - "data": {"timeout": normal_timeout, "stream_timeout": stream_timeout}, - } - - assert router._get_stream_timeout(**args) == stream_timeout - - assert router._get_non_stream_timeout(**args) == normal_timeout - - stream_timeout_val = router._get_timeout( - kwargs={"stream": stream}, - data={"timeout": normal_timeout, "stream_timeout": stream_timeout}, - ) - - if stream: - assert stream_timeout_val == stream_timeout - else: - assert stream_timeout_val == normal_timeout + with patch.object(litellm, "acompletion", spy): + response = await router.acompletion( + model="openai-gpt-5.1", messages=[{"role": "user", "content": "hi"}], mock_response="ok", **request_kwargs + ) + if request_kwargs.get("stream"): + _ = [chunk async for chunk in response] + assert len(sent) == 1 + return sent[0] @pytest.mark.parametrize( @@ -165,7 +118,8 @@ def test_unit_test_streaming_timeout(stream): ({"service_tier": "flex", "stream": True, "stream_timeout": 30}, 30), ], ) -def test_the_deployment_timeout_for_the_requests_service_tier_wins_over_its_general_timeout( +@pytest.mark.asyncio +async def test_the_deployment_timeout_for_the_requests_service_tier_wins_over_its_general_timeout( request_kwargs: dict[str, object], expected_timeout: float ) -> None: router: Final = Router( @@ -183,12 +137,12 @@ def test_the_deployment_timeout_for_the_requests_service_tier_wins_over_its_gene } ] ) - deployment_params: Final = router.model_list[0]["litellm_params"] - assert router._get_timeout(kwargs=request_kwargs, data=deployment_params) == expected_timeout + assert await _timeout_sent_to_the_provider_call(router, request_kwargs) == expected_timeout -def test_a_stream_keeps_the_deployment_stream_timeout_when_only_the_request_timeout_map_names_its_tier() -> None: +@pytest.mark.asyncio +async def test_a_stream_keeps_the_deployment_stream_timeout_when_only_the_request_timeout_map_names_its_tier() -> None: router: Final = Router( model_list=[ { @@ -203,10 +157,9 @@ def test_a_stream_keeps_the_deployment_stream_timeout_when_only_the_request_time } ] ) - deployment_params: Final = router.model_list[0]["litellm_params"] - assert router._get_timeout(kwargs={"service_tier": "flex", "stream": True}, data=deployment_params) == 30.0 - assert router._get_timeout(kwargs={"service_tier": "flex"}, data=deployment_params) == 900.0 + assert await _timeout_sent_to_the_provider_call(router, {"service_tier": "flex", "stream": True}) == 30.0 + assert await _timeout_sent_to_the_provider_call(router, {"service_tier": "flex"}) == 900.0 async def _drain_responses_stream(router: Router) -> None: @@ -274,7 +227,6 @@ async def test_prompt_managed_stream_keeps_the_target_deployments_own_stream_tim ) - @pytest.mark.asyncio @pytest.mark.parametrize( ("agent_params", "target_params", "request_params"), @@ -350,7 +302,9 @@ _INVOKE_ANSWER: Final = json.dumps( async def _streamed_chat_text(router: Router) -> str: - response: Final = await router.acompletion(model="claude", messages=[{"role": "user", "content": "hi"}], stream=True) + response: Final = await router.acompletion( + model="claude", messages=[{"role": "user", "content": "hi"}], stream=True + ) return "".join([chunk.choices[0].delta.content or "" async for chunk in response]) @@ -422,7 +376,11 @@ class _FailureDetails(CustomLogger): async def _drain_messages_stream(router: Router) -> None: response: Final = await router.aanthropic_messages( - model="native", messages=[{"role": "user", "content": "hi"}], max_tokens=10, stream=True, client_side_timeout=True + model="native", + messages=[{"role": "user", "content": "hi"}], + max_tokens=10, + stream=True, + client_side_timeout=True, ) _ = [chunk async for chunk in response] @@ -477,7 +435,10 @@ async def test_native_live_stream_cut_at_the_stream_timeout_is_not_blamed_on_the started: Final = failure_details.details.get("api_call_start_time") assert isinstance(started, datetime) - assert is_caller_timeout_408(failure_details.details, 408, ended=started + timedelta(seconds=STREAM_TIMEOUT_SECONDS)) + assert is_caller_timeout_408( + failure_details.details, 408, ended=started + timedelta(seconds=STREAM_TIMEOUT_SECONDS) + ) + @pytest.mark.asyncio async def test_bedrock_timeout_preserves_the_provider_timeout_error( @@ -489,9 +450,9 @@ async def test_bedrock_timeout_preserves_the_provider_timeout_error( monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret-key") monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1") model_id: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0" - route: Final = respx_mock.post( - url__regex=r"^https://bedrock-runtime\.us-east-1\.amazonaws\.com/.*" - ).mock(side_effect=httpx.ReadTimeout("upstream timeout")) + route: Final = respx_mock.post(url__regex=r"^https://bedrock-runtime\.us-east-1\.amazonaws\.com/.*").mock( + side_effect=httpx.ReadTimeout("upstream timeout") + ) router: Final = Router( model_list=[ { @@ -556,9 +517,7 @@ def test_stream_timeout_uses_the_configured_fallback( ) chunks: Final = tuple(response) - assert tuple(chunk.choices[0].delta.content for chunk in chunks if chunk.choices[0].delta.content) == ( - "fallback", - ) + assert tuple(chunk.choices[0].delta.content for chunk in chunks if chunk.choices[0].delta.content) == ("fallback",) request_bodies: Final = tuple(json.loads(call.request.content) for call in route.calls) assert route.call_count == 2 assert tuple(body["model"] for body in request_bodies) == ("primary", "fallback") @@ -598,9 +557,9 @@ async def test_anthropic_timeout_raises_for_sync_and_streaming( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.anthropic.com/v1/messages" - ).mock(side_effect=httpx.ReadTimeout("scripted Anthropic timeout")) + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + side_effect=httpx.ReadTimeout("scripted Anthropic timeout") + ) request_kwargs: Final = { "model": "anthropic/claude-haiku-4-5-20251001", "messages": [{"role": "user", "content": "timeout contract"}], @@ -634,9 +593,9 @@ def test_openai_router_timeout_raises_after_httpx_read_timeout( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock(side_effect=httpx.ReadTimeout("scripted OpenAI timeout")) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=httpx.ReadTimeout("scripted OpenAI timeout") + ) router: Final = Router( model_list=[ { @@ -665,9 +624,9 @@ def test_azure_router_timeout_raises_after_httpx_read_timeout( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://azure.example.com/openai/deployments/gpt-4o-mini/chat/completions" - ).mock(side_effect=httpx.ReadTimeout("scripted Azure timeout")) + route: Final = respx_mock.post("https://azure.example.com/openai/deployments/gpt-4o-mini/chat/completions").mock( + side_effect=httpx.ReadTimeout("scripted Azure timeout") + ) router: Final = Router( model_list=[ { diff --git a/tests/unit/test_router_block_helpers.py b/tests/unit/test_router_block_helpers.py index c89d514b166..431ce5c25ee 100644 --- a/tests/unit/test_router_block_helpers.py +++ b/tests/unit/test_router_block_helpers.py @@ -45,13 +45,3 @@ class TestAreAllDeploymentsBlocked: def test_empty_list_returns_false(self): router = _make_router("gpt-4o") assert router.are_all_deployments_blocked([]) is False - - -class TestIsModelFullyBlocked: - def test_all_deployments_blocked_returns_true(self): - router = _make_router("gpt-4o", blocked=True) - assert router._is_model_fully_blocked("gpt-4o") is True - - def test_unblocked_deployment_returns_false(self): - router = _make_router("gpt-4o", blocked=False) - assert router._is_model_fully_blocked("gpt-4o") is False diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 724c07f8c31..c0c8a0655fe 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -29,12 +29,10 @@ import respx import litellm from litellm import Router -from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE from litellm.litellm_core_utils.ptu_pricing import ptu_config_error from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.proxy.utils import create_model_info_response from litellm.types.router import ( Deployment, @@ -117,13 +115,15 @@ async def test_discovered_limits_survive_deployment_growth_and_removal( ) for deployment in deployments[initial_count:]: router.add_deployment(deployment) - await router._arefresh_deployment_model_info(router.model_list[-1], client=handler) + router.cache.in_memory_cache.flush_cache() + await router.arefresh_model_info(client=handler) assert all( router.get_configured_token_limits(deployment.model_name) == (4096, 4096) for deployment in deployments ) for deployment in deployments[-2:]: router.delete_deployment(deployment.model_info.id or "") - await router._arefresh_deployment_model_info(router.model_list[0], client=handler) + router.cache.in_memory_cache.flush_cache() + await router.arefresh_model_info(client=handler) assert all( router.get_configured_token_limits(deployment.model_name) == (4096, 4096) for deployment in deployments[:-2] ) @@ -167,8 +167,9 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch await handler.client.aclose() async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: handler.client = client - await router._arefresh_deployment_model_info(router.model_list[0], client=handler) + await router.arefresh_model_info(client=handler) assert router.get_configured_token_limits("local") == (None, None) + router.cache.in_memory_cache.flush_cache() await router.arefresh_model_info(client=handler) assert router.get_configured_token_limits("local") == (2048, 2048) _invalidate_model_cost_lowercase_map() @@ -204,7 +205,7 @@ async def test_discovery_sends_the_secret_behind_an_oidc_api_key( await handler.client.aclose() async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: handler.client = client - await router._arefresh_deployment_model_info(router.model_list[0], client=handler) + await router.arefresh_model_info(client=handler) assert authorizations == ["Bearer local-secret"] _invalidate_model_cost_lowercase_map() @@ -299,54 +300,6 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa _invalidate_model_cost_lowercase_map() -async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) - clock: Final = Mock(return_value=0.0) - router: Final = Router( - model_list=[ - { - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://expiry.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "expiring-discovery"}, - } - ] - ) - router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS) - responses: Final = iter( - ( - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}), - httpx.Response(503), - ) - ) - handler: Final = AsyncHTTPHandler() - await handler.client.aclose() - async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client: - handler.client = client - await router.arefresh_model_info(client=handler) - clock.return_value = MODEL_INFO_REFRESH_SECONDS - router.cache.in_memory_cache.flush_cache() - await router.arefresh_model_info(client=handler) - clock.return_value = 2 * MODEL_INFO_REFRESH_SECONDS + 1 - router.cache.in_memory_cache.flush_cache() - await router.arefresh_model_info(client=handler) - assert router.get_configured_token_limits("local") == (8192, 8192) - group: Final = router.get_model_group_info("local") - assert group is not None - assert group.max_input_tokens == 8192 - clock.return_value = 3 * MODEL_INFO_REFRESH_SECONDS + 1 - await router.arefresh_model_info(client=handler) - assert router.get_configured_token_limits("local") == (None, None) - expired_group: Final = router.get_model_group_info("local") - assert expired_group is not None - assert expired_group.max_input_tokens is None - _invalidate_model_cost_lowercase_map() - - @pytest.mark.parametrize("provider", ("hosted_vllm", "openai", "openai_like", "text-completion-openai")) async def test_discovered_limits_are_isolated_overridable_and_refreshable( provider: str, monkeypatch: pytest.MonkeyPatch @@ -390,11 +343,6 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable( assert listing is not None assert listing.max_input_tokens == 8192 assert router.get_configured_token_limits("local") == (8192, 8192) - assert router._deployment_max_input_tokens("local", router.model_list[1]) == 4096 - allowed: Final = router._pre_call_checks( - model="local", healthy_deployments=router.model_list, input="prompt", input_token_count=5000 - ) - assert [deployment["model_info"]["id"] for deployment in allowed] == ["one"] assert router.model_list[0]["model_info"].get("max_input_tokens") is None assert litellm.model_cost[f"{provider}/org/local-model"].get("max_input_tokens") is None router.cache.in_memory_cache.flush_cache() @@ -892,32 +840,19 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): litellm.model_cost[backend_model] = backend_entry _invalidate_model_cost_lowercase_map() - router = Router(model_list=[]) - with patch.object(Router, "_add_deployment", lambda self, deployment: deployment): - router._create_deployment( - deployment_info={}, - _model_name="chatgpt/gpt-5.4", - _litellm_params={ - "model": "gpt-5.4", - "custom_llm_provider": "chatgpt", - }, - _model_info={ - "id": "chatgpt-shared-mode-base", - "mode": "responses", - }, - ) - router._create_deployment( - deployment_info={}, - _model_name="chatgpt/gpt-5.4-medium", - _litellm_params={ - "model": "gpt-5.4", - "custom_llm_provider": "chatgpt", - }, - _model_info={ - "id": "chatgpt-shared-mode-alias", - "mode": "chat", - }, - ) + Router( + model_list=[ + { + "model_name": model_name, + "litellm_params": {"model": "gpt-5.4", "custom_llm_provider": "chatgpt"}, + "model_info": {"id": model_id, "mode": mode}, + } + for model_name, model_id, mode in ( + ("chatgpt/gpt-5.4", "chatgpt-shared-mode-base", "responses"), + ("chatgpt/gpt-5.4-medium", "chatgpt-shared-mode-alias", "chat"), + ) + ] + ) assert litellm.model_cost[backend_model]["mode"] == "responses" assert "mode" in litellm.model_cost[backend_model] @@ -940,7 +875,7 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): @pytest.mark.parametrize("rate", (0.0, 0.004)) @pytest.mark.parametrize("explicit_base", (False, True)) @pytest.mark.parametrize("provider", (None, "bedrock")) -@pytest.mark.parametrize("registration_path", ("direct", "startup", "reload")) +@pytest.mark.parametrize("registration_path", ("startup", "reload")) def test_bedrock_deployment_inherits_invoked_profile_around_custom_prices( monkeypatch: pytest.MonkeyPatch, field: str, @@ -986,12 +921,9 @@ def test_bedrock_deployment_inherits_invoked_profile_around_custom_prices( } ), ) - if registration_path == "direct": - Router._register_deployment_pricing(deployment) - else: - router: Final = Router(model_list=[deployment.to_json(exclude_none=True)]) - if registration_path == "reload": - router.replay_model_cost_registrations() + router: Final = Router(model_list=[deployment.to_json(exclude_none=True)]) + if registration_path == "reload": + router.replay_model_cost_registrations() response: Final = ModelResponse( model=model, usage=Usage( @@ -1260,50 +1192,6 @@ def test_partial_custom_pricing_bills_explicit_cache_hits( _restore_model_cost_entries(model_keys) -def test_inherit_builtin_cache_pricing_fills_only_missing_fields(): - """Direct unit test of the helper: missing cache fields are filled from the - backend model's built-in entry, while an explicitly set cache field and the - user's input/output pricing are left untouched. - """ - backend_model = "anthropic/claude-sonnet-4-5-20250929" - builtin_info = litellm.get_model_info(model=backend_model) - builtin_cache_create = builtin_info["cache_creation_input_token_cost"] - builtin_cache_read = builtin_info["cache_read_input_token_cost"] - assert builtin_cache_create is not None and builtin_cache_create > 0 - assert builtin_cache_read is not None and builtin_cache_read > 0 - - explicit_cache_read = builtin_cache_read + 1 - model_info = { - "input_cost_per_token": 0.000003, - "cache_read_input_token_cost": explicit_cache_read, - } - - Router._inherit_builtin_cache_pricing( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider="anthropic", - ) - - assert model_info["input_cost_per_token"] == 0.000003 - assert model_info["cache_read_input_token_cost"] == explicit_cache_read - assert model_info["cache_creation_input_token_cost"] == builtin_cache_create - - -def test_inherit_builtin_cache_pricing_noop_for_unknown_backend(): - """No canonical entry for the backend model means the helper leaves the - passed-in dict unchanged rather than raising. - """ - model_info = {"input_cost_per_token": 0.000003} - - Router._inherit_builtin_cache_pricing( - model_info=model_info, - backend_model="this-backend-model-does-not-exist-x9y8z7", - custom_llm_provider=None, - ) - - assert model_info == {"input_cost_per_token": 0.000003} - - _TIER_BACKEND_MODEL: Final = "tier-priced-backend" _TIER_BACKEND_KEY: Final = f"openai/{_TIER_BACKEND_MODEL}" _CUSTOM_STANDARD_INPUT_RATE: Final = 0.00011 @@ -1349,86 +1237,6 @@ def _register_azure_tier_backend() -> None: _invalidate_model_cost_lowercase_map() -def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> None: - model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) - } - try: - _register_tier_backend() - model_info: Final = { - "id": "custom-priced-tier-deployment", - "input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE, - "output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE, - "output_cost_per_token_ultrafast": 0.00999, - } - - Router._inherit_builtin_service_tier_pricing( - model_info=model_info, - backend_model=_TIER_BACKEND_MODEL, - custom_llm_provider="openai", - ) - - assert model_info == { - "id": "custom-priced-tier-deployment", - "input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE, - "output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE, - "input_cost_per_token_ultrafast": _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"], - "output_cost_per_token_ultrafast": 0.00999, - "input_cost_per_token_priority": _TIER_BACKEND_ENTRY["input_cost_per_token_priority"], - "output_cost_per_token_priority": _TIER_BACKEND_ENTRY["output_cost_per_token_priority"], - "input_cost_per_token_flex": _TIER_BACKEND_ENTRY["input_cost_per_token_flex"], - "output_cost_per_token_flex": _TIER_BACKEND_ENTRY["output_cost_per_token_flex"], - "input_cost_per_token_balanced": _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"], - "output_cost_per_token_balanced": _TIER_BACKEND_ENTRY["output_cost_per_token_balanced"], - "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY["cache_read_input_token_cost_ultrafast"], - "input_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[ - "input_cost_per_token_above_272k_tokens_ultrafast" - ], - "output_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[ - "output_cost_per_token_above_272k_tokens_ultrafast" - ], - } - finally: - _restore_model_cost_entries(model_cost_entries) - litellm.get_model_info.cache_clear() - - -def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend() -> None: - model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) - } - try: - _register_tier_backend() - model_info_without_base_rate: Final = { - "id": "custom-priced-no-base-rate", - "input_cost_per_token_ultrafast": 0.00031, - } - expected_without_base_rate: Final = copy.deepcopy(model_info_without_base_rate) - Router._inherit_builtin_service_tier_pricing( - model_info=model_info_without_base_rate, - backend_model=_TIER_BACKEND_MODEL, - custom_llm_provider="openai", - ) - - model_info_with_unknown_backend: Final = { - "id": "custom-priced-unknown-backend", - "input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE, - "output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE, - } - expected_with_unknown_backend: Final = copy.deepcopy(model_info_with_unknown_backend) - Router._inherit_builtin_service_tier_pricing( - model_info=model_info_with_unknown_backend, - backend_model="tier-priced-backend-unknown", - custom_llm_provider="openai", - ) - - assert model_info_without_base_rate == expected_without_base_rate - assert model_info_with_unknown_backend == expected_with_unknown_backend - finally: - _restore_model_cost_entries(model_cost_entries) - litellm.get_model_info.cache_clear() - - def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() -> None: model_id: Final = "tier-priced-deployment" model_cost_entries: Final = { @@ -1644,8 +1452,8 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay + 100 * _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"] ) - payload: Final = Router._deployment_model_cost_payload( - deployment=Deployment( + router.add_deployment( + Deployment( model_name="azure/tier-priced-alias-from-params", litellm_params=LiteLLM_Params( model="azure/tier-priced-alias", @@ -1657,6 +1465,7 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay model_info=ModelInfo(id=payload_id), ) ) + payload: Final = litellm.model_cost[payload_id] assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["input_cost_per_token_priority"] assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"] @@ -1665,203 +1474,6 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay litellm.get_model_info.cache_clear() -@pytest.mark.parametrize( - ("model_info_base_model", "params_base_model", "model", "expected"), - ( - pytest.param( - "azure/tier-priced-model-info-base", - "azure/tier-priced-params-base", - "azure/tier-priced-deployment-alias", - "azure/tier-priced-model-info-base", - id="model-info-base-model-wins", - ), - pytest.param( - None, - "azure/tier-priced-params-base", - "azure/tier-priced-deployment-alias", - "azure/tier-priced-params-base", - id="params-base-model-fallback", - ), - pytest.param( - None, - None, - "azure/tier-priced-deployment-alias", - "azure/tier-priced-deployment-alias", - id="model-fallback", - ), - pytest.param( - "", - "azure/tier-priced-params-base", - "azure/tier-priced-deployment-alias", - "azure/tier-priced-params-base", - id="empty-model-info-base-model-falls-through", - ), - ), -) -def test_cost_map_backend_model_uses_canonical_model_precedence( - model_info_base_model: str | None, - params_base_model: str | None, - model: str, - expected: str, -) -> None: - deployment: Final = Deployment( - model_name="azure/tier-priced-cost-map-backend", - litellm_params=LiteLLM_Params(model=model, base_model=params_base_model), - model_info=ModelInfo(id="tier-priced-cost-map-backend", base_model=model_info_base_model), - ) - - assert Router._cost_map_backend_model(deployment) == expected - - -def test_inherit_builtin_base_rates_for_off_peak_fills_missing_rates(): - """Direct unit test of the helper: an entry carrying only an - off_peak_pricing block inherits the backend model's built-in base token - rates, so cost lookup via the deployment id can bill standard rates - outside the windows. - """ - backend_model = "gpt-4o-mini" - builtin_info = litellm.get_model_info(model=backend_model, custom_llm_provider="openai") - off_peak_block = { - "hours_utc": "00:00-00:00", - "input_cost_per_token": 5e-07, - "output_cost_per_token": 1e-06, - } - model_info = {"off_peak_pricing": off_peak_block} - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider="openai", - ) - - assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] - assert model_info["output_cost_per_token"] == builtin_info["output_cost_per_token"] - assert model_info["off_peak_pricing"] == off_peak_block - - -def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates(): - """A backend with above-threshold pricing hands the whole rate structure to - the deployment entry, so peak-hour billing of large prompts through that - entry matches the shared backend entry instead of flattening to the base - rate. - """ - backend_model = "gemini/gemini-2.5-pro" - builtin_info = litellm.get_model_info(model=backend_model) - assert builtin_info["input_cost_per_token_above_200k_tokens"] is not None - - model_info = { - "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, - } - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider="gemini", - ) - - assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] - assert ( - model_info["input_cost_per_token_above_200k_tokens"] == builtin_info["input_cost_per_token_above_200k_tokens"] - ) - assert ( - model_info["output_cost_per_token_above_200k_tokens"] == builtin_info["output_cost_per_token_above_200k_tokens"] - ) - - -def test_inherit_builtin_base_rates_for_off_peak_carries_companion_billing_fields(): - """Billing rules that are not literal cost rates, like the web search - billing unit, must ride along, or grounding and regional uplifts would - bill differently through the deployment entry than through the shared - backend entry. - """ - backend_model = "gemini-3-pro-image" - raw_entry = litellm.model_cost[backend_model] - assert raw_entry.get("web_search_billing_unit") is not None - - model_info = { - "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, - } - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=None, - ) - - assert model_info["web_search_billing_unit"] == raw_entry["web_search_billing_unit"] - assert model_info["input_cost_per_token"] == raw_entry["input_cost_per_token"] - - -def test_inherit_builtin_base_rates_for_off_peak_tiered_only_backend_stores_no_zero(): - """A tiered-only backend has no flat token rates; get_model_info synthesizes - zeros for them, and storing those would mark the deployment explicitly - priced free. The tier table itself must carry over as an isolated copy so - mutating the deployment entry never touches the shared cost map. - """ - backend_model = "dashscope/qwen-flash" - raw_tiers = litellm.model_cost[backend_model]["tiered_pricing"] - - model_info = { - "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, - } - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider="dashscope", - ) - - assert model_info.get("input_cost_per_token") != 0 - assert model_info.get("output_cost_per_token") != 0 - assert model_info["tiered_pricing"] == raw_tiers - assert model_info["tiered_pricing"] is not raw_tiers - assert model_info["tiered_pricing"][0] is not raw_tiers[0] - - original_first_tier = copy.deepcopy(raw_tiers[0]) - model_info["tiered_pricing"][0]["input_cost_per_token"] = 123.0 - assert raw_tiers[0] == original_first_tier - - -def test_inherit_builtin_base_rates_for_off_peak_leaves_explicit_rates_alone(): - """An entry that sets its own base rate beside the block already counts as - a full custom pricing entry; the helper must not mix builtin rates into it. - """ - model_info = { - "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, - "input_cost_per_token": 3e-06, - } - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model="gpt-4o-mini", - custom_llm_provider="openai", - ) - - assert model_info["input_cost_per_token"] == 3e-06 - assert "output_cost_per_token" not in model_info - - -def test_inherit_builtin_base_rates_for_off_peak_noop_without_block_or_backend(): - """Nothing happens without an off_peak_pricing block, and an unmapped - backend model leaves the entry unchanged rather than raising. - """ - plain_info = {"id": "dep-1"} - Router._inherit_builtin_base_rates_for_off_peak( - model_info=plain_info, - backend_model="gpt-4o-mini", - custom_llm_provider="openai", - ) - assert plain_info == {"id": "dep-1"} - - off_peak_info = {"off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}} - Router._inherit_builtin_base_rates_for_off_peak( - model_info=off_peak_info, - backend_model="this-backend-model-does-not-exist-x9y8z7", - custom_llm_provider=None, - ) - assert "input_cost_per_token" not in off_peak_info - - def test_custom_pricing_field_denylist_covers_all_builtin_pricing_fields(): """The shared-backend-key stripping in Router relies on CustomPricingLiteLLMParams enumerating every per-deployment pricing field. @@ -2635,28 +2247,6 @@ def test_repointing_a_deployment_drops_its_previous_backend_key(monkeypatch): _invalidate_model_cost_lowercase_map() -@pytest.mark.parametrize( - "model, custom_llm_provider, expected", - [ - ("gpt-4o", None, ("gpt-4o",)), - ("gpt-4o", "openai", ("openai/gpt-4o",)), - ("openai/gpt-4o", None, ("openai/gpt-4o",)), - ("responses/gpt-4o", "openai", ("openai/responses/gpt-4o", "openai/gpt-4o")), - ("responses/gpt-4o", None, ("responses/gpt-4o", "gpt-4o")), - ], -) -def test_backend_cost_map_keys_matches_what_registration_writes(model, custom_llm_provider, expected): - """ - The withdrawal path drops exactly the keys the registration wrote, so the two - have to agree on the provider prefix and on the responses/ alias. The first - key is also the one the registration uses as the shared backend key, so its - position is load-bearing rather than incidental. - """ - keys = Router._backend_cost_map_keys(model=model, custom_llm_provider=custom_llm_provider) - assert keys == expected - assert keys[0] == (model if custom_llm_provider is None else f"{custom_llm_provider}/{model}") - - def test_mantle_chat_deployment_does_not_rewrite_the_shared_responses_row(): """ A Bedrock Mantle deployment that pins ``mode: chat`` routes its own chat @@ -3135,55 +2725,34 @@ def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered() -> None: _invalidate_model_cost_lowercase_map() -def test_replay_model_cost_registrations_survives_a_malformed_deployment(): - """ - The rebuild reads whatever dicts are sitting in model_list, so one entry that - cannot be rebuilt into a Deployment must not stop the rest being restored. - """ - saved_model_cost = litellm.model_cost - try: - router = Router( - model_list=[ - { - "model_name": "healthy", - "litellm_params": {"model": "hosted_vllm/healthy-backend"}, - "model_info": {"id": "healthy-id", "max_input_tokens": 777}, - } - ], - ) - router.model_list.insert(0, {"litellm_params": {}}) - - litellm.model_cost = {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}} - _invalidate_model_cost_lowercase_map() - router.replay_model_cost_registrations() - - assert litellm.model_cost["healthy-id"]["max_input_tokens"] == 777 - finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() - - def test_deployment_model_cost_payload_folds_in_litellm_params_pricing(): """ Custom pricing is configured on litellm_params but has to land in the cost-map entry, and setting it pulls in the built-in cache pricing for the backend model. Both are what make the entry reproducible from a deployment. """ - payload = Router._deployment_model_cost_payload( - deployment=Deployment( - model_name="priced", - litellm_params=LiteLLM_Params( - model="gemini/gemini-2.5-pro", - input_cost_per_token=0.000123, - ), - model_info=ModelInfo(id="payload-id", max_input_tokens=4242), + model_keys = {"payload-id": copy.deepcopy(litellm.model_cost.get("payload-id"))} + try: + Router( + model_list=[ + Deployment( + model_name="priced", + litellm_params=LiteLLM_Params( + model="gemini/gemini-2.5-pro", + input_cost_per_token=0.000123, + ), + model_info=ModelInfo(id="payload-id", max_input_tokens=4242), + ).to_json(exclude_none=True) + ] ) - ) + payload = litellm.model_cost["payload-id"] - assert payload["id"] == "payload-id" - assert payload["max_input_tokens"] == 4242 - assert payload["input_cost_per_token"] == 0.000123 - assert payload["cache_read_input_token_cost"] > 0 + assert payload["id"] == "payload-id" + assert payload["max_input_tokens"] == 4242 + assert payload["input_cost_per_token"] == 0.000123 + assert payload["cache_read_input_token_cost"] > 0 + finally: + _restore_model_cost_entries(model_keys) def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -> None: @@ -3193,18 +2762,21 @@ def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() - } try: _register_tier_backend() - payload: Final = Router._deployment_model_cost_payload( - deployment=Deployment( - model_name="tier-priced-payload", - litellm_params=LiteLLM_Params( - model=_TIER_BACKEND_MODEL, - custom_llm_provider="openai", - input_cost_per_token=_CUSTOM_STANDARD_INPUT_RATE, - output_cost_per_token=_CUSTOM_STANDARD_OUTPUT_RATE, - ), - model_info=ModelInfo(id=model_id), - ) + Router( + model_list=[ + Deployment( + model_name="tier-priced-payload", + litellm_params=LiteLLM_Params( + model=_TIER_BACKEND_MODEL, + custom_llm_provider="openai", + input_cost_per_token=_CUSTOM_STANDARD_INPUT_RATE, + output_cost_per_token=_CUSTOM_STANDARD_OUTPUT_RATE, + ), + model_info=ModelInfo(id=model_id), + ).to_json(exclude_none=True) + ] ) + payload: Final = litellm.model_cost[model_id] assert payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"] assert payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"] @@ -3227,15 +2799,14 @@ def test_register_deployment_in_model_cost_writes_both_key_families(): "hosted_vllm/both-families-backend": copy.deepcopy(litellm.model_cost.get("hosted_vllm/both-families-backend")), } try: - Router._register_deployment_in_model_cost( - model_id="both-families-id", - model_info={ - "id": "both-families-id", - "max_input_tokens": 999, - "litellm_provider": "hosted_vllm", - }, - model="hosted_vllm/both-families-backend", - custom_llm_provider=None, + Router( + model_list=[ + { + "model_name": "both-families", + "litellm_params": {"model": "hosted_vllm/both-families-backend"}, + "model_info": {"id": "both-families-id", "max_input_tokens": 999}, + } + ] ) assert litellm.model_cost["both-families-id"]["max_input_tokens"] == 999 @@ -3345,8 +2916,6 @@ def test_replay_live_router_model_cost_rebuilds_every_live_router(): A process can hold more than one Router, so the rebuild has to fan out across all of them rather than restoring whichever one happens to be reachable. """ - from litellm.router import _replay_live_router_model_cost - saved_model_cost = litellm.model_cost try: first = Router( @@ -3368,9 +2937,7 @@ def test_replay_live_router_model_cost_rebuilds_every_live_router(): ], ) - litellm.model_cost = {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}} - _invalidate_model_cost_lowercase_map() - _replay_live_router_model_cost() + _simulate_price_data_reload({"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}}) assert litellm.model_cost["first-id"]["max_input_tokens"] == 111 assert litellm.model_cost["second-id"]["max_input_tokens"] == 222 @@ -3439,61 +3006,6 @@ def test_strategy_router_alias_pricing_never_enters_model_cost(monkeypatch): _invalidate_model_cost_lowercase_map() -def test_inherit_builtin_tiered_output_rate_fills_the_backend_flat_rate(): - """ - A deployment entry whose custom tiers publish only input rates would bill - completions at 0, so the backend model's flat output rate is copied in at - registration. - """ - model_info = {"tiered_pricing": [{"range": [0, 3000], "input_cost_per_token": 3.25e-07}]} - - Router._inherit_builtin_tiered_output_rate( - model_info=model_info, - backend_model="claude-haiku-4-5", - custom_llm_provider="anthropic", - ) - - backend_rate = litellm.get_model_info(model="claude-haiku-4-5", custom_llm_provider="anthropic")[ - "output_cost_per_token" - ] - assert backend_rate > 0 - assert model_info["output_cost_per_token"] == backend_rate - - -def test_inherit_builtin_tiered_output_rate_never_stores_a_synthesized_zero(): - """ - Regression: get_model_info reports output_cost_per_token 0 for a backend that - only publishes tiered rates (e.g. dashscope/qwen-flash), and storing that zero - would mark the deployment as explicitly priced free. - """ - backend_info = litellm.get_model_info(model="qwen-flash", custom_llm_provider="dashscope") - assert backend_info["output_cost_per_token"] == 0 - - model_info = {"tiered_pricing": [{"range": [0, 3000], "input_cost_per_token": 3.25e-07}]} - Router._inherit_builtin_tiered_output_rate( - model_info=model_info, - backend_model="qwen-flash", - custom_llm_provider="dashscope", - ) - - assert "output_cost_per_token" not in model_info - - -def test_inherit_builtin_tiered_output_rate_leaves_a_user_rate_alone(): - model_info = { - "tiered_pricing": [{"range": [0, 3000], "input_cost_per_token": 3.25e-07}], - "output_cost_per_token": 9e-07, - } - - Router._inherit_builtin_tiered_output_rate( - model_info=model_info, - backend_model="claude-haiku-4-5", - custom_llm_provider="anthropic", - ) - - assert model_info["output_cost_per_token"] == 9e-07 - - # --- a config.yaml PTU deployment must not also bill per token ------------------ _PTU_MODEL_INFO = { diff --git a/tests/unit/test_router_order_fallback.py b/tests/unit/test_router_order_fallback.py index 6dd445e2b5d..3ea176b9790 100644 --- a/tests/unit/test_router_order_fallback.py +++ b/tests/unit/test_router_order_fallback.py @@ -561,34 +561,6 @@ async def test_router_order_fallback_retries_keep_target_order(): assert seen_target_orders.count(2) >= 2 -@pytest.mark.asyncio -async def test_generic_api_call_strips_target_order_from_provider_kwargs(): - captured: Final = {} - - async def _fake_provider(**provider_kwargs): - captured.update(provider_kwargs) - return "ok" - - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key", "order": 2}, - "model_info": {"id": "2"}, - }, - ], - ) - response = await router._ageneric_api_call_with_fallbacks_helper( - model="test-model", - original_generic_function=_fake_provider, - _target_order=2, - messages=[{"role": "user", "content": "hi"}], - ) - assert response == "ok" - assert captured["model"] == "gpt-4o" - assert "_target_order" not in captured - - @pytest.mark.asyncio async def test_text_completion_order_fallback_hop_does_not_send_target_order_upstream(): upstream_bodies: Final[list[dict]] = [] @@ -793,24 +765,6 @@ async def test_responses_order_fallback_hop_keeps_the_encrypted_reasoning_the_sa ] -def test_fallback_hop_reads_the_deployment_that_just_failed_from_the_metadata_bucket_it_writes(): - router: Final = _two_openai_orders_on_one_encryption_boundary_router() - order_2: Final = router.get_deployment(model_id="openai-order-2").model_dump(exclude_none=True) - hop_input: Final = _responses_history_with_order_1_reasoning() - hop_kwargs: Final = { - "model": "gpt-6-astra", - "input": hop_input, - "fallback_depth": 1, - "metadata": {"model_info": {"id": "openai-order-1"}}, - "litellm_metadata": {"previous_models": [{"deployment_id": None}]}, - } - - router._update_kwargs_with_deployment(deployment=order_2, kwargs=hop_kwargs) - - assert hop_input == _responses_history_with_order_1_reasoning() - assert hop_kwargs["metadata"]["model_info"]["id"] == "openai-order-2" - - def test_check_non_standard_fallback_format(): from litellm.router_utils.fallback_event_handlers import ( check_non_standard_fallback_format, diff --git a/tests/unit/test_router_per_deployment_num_retries.py b/tests/unit/test_router_per_deployment_num_retries.py index b08ae25c438..76d30e1c58a 100644 --- a/tests/unit/test_router_per_deployment_num_retries.py +++ b/tests/unit/test_router_per_deployment_num_retries.py @@ -84,105 +84,6 @@ async def test_routed_images_do_not_multiply_sdk_and_router_retries( class TestPerDeploymentNumRetries: """Test that per-deployment num_retries in litellm_params is correctly used.""" - def test_set_deployment_num_retries_on_exception(self): - """ - Test that _set_deployment_num_retries_on_exception sets num_retries - on the exception from the deployment's litellm_params. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "test-key", - "num_retries": 5, # Per-deployment setting - }, - }, - ], - num_retries=1, # Global setting - ) - - deployment = router.model_list[0] - - # Create a mock exception without num_retries - class MockException(Exception): - pass - - exc = MockException("test error") - assert not hasattr(exc, "num_retries") or exc.num_retries is None - - # Call the helper - router._set_deployment_num_retries_on_exception(exc, deployment) - - # Verify num_retries was set from deployment - assert exc.num_retries == 5 - - def test_set_deployment_num_retries_does_not_override_existing(self): - """ - Test that _set_deployment_num_retries_on_exception does NOT override - if exception already has num_retries set. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "test-key", - "num_retries": 5, - }, - }, - ], - num_retries=1, - ) - - deployment = router.model_list[0] - - # Create an exception that already has num_retries - class MockException(Exception): - num_retries = 10 # Already set - - exc = MockException("test error") - - # Call the helper - router._set_deployment_num_retries_on_exception(exc, deployment) - - # Verify num_retries was NOT overridden - assert exc.num_retries == 10 - - def test_deployment_without_num_retries(self): - """ - Test that _set_deployment_num_retries_on_exception does nothing - if deployment has no num_retries set. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "test-key", - # No num_retries set - }, - }, - ], - num_retries=3, - ) - - deployment = router.model_list[0] - - class MockException(Exception): - pass - - exc = MockException("test error") - - # Call the helper - router._set_deployment_num_retries_on_exception(exc, deployment) - - # Verify num_retries was not set (deployment has no num_retries) - assert not hasattr(exc, "num_retries") or exc.num_retries is None - def test_request_level_num_retries_takes_precedence(self): """ Test that request-level num_retries (passed in kwargs) is still respected. @@ -230,39 +131,6 @@ class TestPerDeploymentNumRetries: router._update_kwargs_before_fallbacks(model="test-model", kwargs=kwargs) assert kwargs.get("num_retries") is None - def test_set_deployment_num_retries_with_string_value(self): - """ - Test that _set_deployment_num_retries_on_exception handles string values - from environment variables correctly. - GitHub Issue: #19481 - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "test-key", - "num_retries": "6", # String value (as from env var) - }, - }, - ], - num_retries=0, # Global setting - ) - - deployment = router.model_list[0] - - class MockException(Exception): - pass - - exc = MockException("test error") - - # Call the helper - router._set_deployment_num_retries_on_exception(exc, deployment) - - # Verify num_retries was converted from string to int - assert exc.num_retries == 6 - class TestNumRetriesNoneGuard: """ @@ -330,7 +198,7 @@ class TestNumRetriesNoneGuard: assert response.choices[0].message.content == "ok" @pytest.mark.asyncio - async def test_async_function_with_retries_none_falls_back_to_zero(self): + async def test_acompletion_num_retries_none_with_router_none_falls_back_to_zero(self): """ When both the per-request value AND the router-level setting are None (e.g. after Router.update_settings(num_retries=None), #28126), num_retries must @@ -339,49 +207,15 @@ class TestNumRetriesNoneGuard: router = self._mock_router(num_retries=0) router.num_retries = None # simulate update_settings(num_retries=None) - async def failing_fn(*args, **kwargs): - raise litellm.RateLimitError( - message="boom", model="mock-model", llm_provider="openai" - ) - with patch("asyncio.sleep", return_value=None): with pytest.raises(litellm.RateLimitError): - await router.async_function_with_retries( - original_function=failing_fn, + await router.acompletion( model="mock-model", messages=[{"role": "user", "content": "hi"}], num_retries=None, + mock_response=litellm.RateLimitError(message="boom", model="mock-model", llm_provider="openai"), ) - @pytest.mark.asyncio - async def test_async_function_with_retries_none_falls_back_to_router_default(self): - """ - A None per-request num_retries falls back to the router-level setting, so retries - still happen (original_function is invoked more than once) before the real error - is raised - proving None did not silently disable retries or crash. - """ - router = self._mock_router(num_retries=3) - calls = {"n": 0} - - async def failing_fn(*args, **kwargs): - calls["n"] += 1 - raise litellm.InternalServerError( - message="boom", model="mock-model", llm_provider="openai" - ) - - with patch("asyncio.sleep", return_value=None): - with pytest.raises(litellm.InternalServerError): - await router.async_function_with_retries( - original_function=failing_fn, - model="mock-model", - messages=[{"role": "user", "content": "hi"}], - metadata={}, # populated by acompletion in the real path; log_retry needs it - num_retries=None, - ) - - # 1 initial attempt + at least 1 retry -> proves None fell back to a positive int - assert calls["n"] >= 2 - class TestNoProviderRetryAmplification: """ @@ -435,9 +269,7 @@ class TestNoProviderRetryAmplification: counter = self._install_counting_upstream() with patch("asyncio.sleep", return_value=None): with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="mock", messages=[{"role": "user", "content": "hi"}], **call_kwargs - ) + await router.acompletion(model="mock", messages=[{"role": "user", "content": "hi"}], **call_kwargs) return counter["n"] @pytest.mark.asyncio @@ -449,9 +281,7 @@ class TestNoProviderRetryAmplification: an unfixed build sends 9 (N=2) or 36 (N=5). """ counter = self._install_counting_upstream() - router = self._router( - f"https://amp-{num_retries}.local/v1", {"num_retries": num_retries}, num_retries=1 - ) + router = self._router(f"https://amp-{num_retries}.local/v1", {"num_retries": num_retries}, num_retries=1) with patch("asyncio.sleep", return_value=None): with pytest.raises(litellm.InternalServerError): await router.acompletion(model="mock", messages=[{"role": "user", "content": "hi"}]) @@ -474,9 +304,7 @@ class TestNoProviderRetryAmplification: retries for a routed call: deployment ``num_retries=5`` plus ``max_retries=3`` still sends exactly ``6`` upstream requests. """ - router = self._router( - "https://nest-dep.local/v1", {"num_retries": 5, "max_retries": 3}, num_retries=1 - ) + router = self._router("https://nest-dep.local/v1", {"num_retries": 5, "max_retries": 3}, num_retries=1) assert await self._call_and_count(router) == 6 @pytest.mark.asyncio @@ -660,9 +488,7 @@ class TestRequestNumRetriesBeatsGlobal: ) with patch("asyncio.sleep", return_value=None): with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="mock", messages=[{"role": "user", "content": "hi"}] - ) + await router.acompletion(model="mock", messages=[{"role": "user", "content": "hi"}]) assert counter.attempts == 3 @@ -729,9 +555,7 @@ class TestRequestNumRetriesBeatsDeployment: @pytest.mark.asyncio @pytest.mark.parametrize("request_num_retries, expected_attempts", [(3, 4), (0, 1)]) - async def test_request_num_retries_overrides_deployment( - self, request_num_retries, expected_attempts - ): + async def test_request_num_retries_overrides_deployment(self, request_num_retries, expected_attempts): """ global=1, deployment=2, request=3 -> 4 attempts, and request=0 -> a single attempt. Before the fix the deployment value won, so both cases sent 3 attempts. @@ -749,9 +573,7 @@ class TestRequestNumRetriesBeatsDeployment: The deployment value keeps its place directly above litellm_settings: global=1 and deployment=2 with no request value -> 1 initial attempt + 2 retries. """ - attempts = await self._count_attempts( - global_num_retries=1, deployment_num_retries=2, request_num_retries=None - ) + attempts = await self._count_attempts(global_num_retries=1, deployment_num_retries=2, request_num_retries=None) assert attempts == 3 @pytest.mark.asyncio diff --git a/tests/unit/test_router_retry_backoff_headers.py b/tests/unit/test_router_retry_backoff_headers.py index 03c3af692ce..73d1091b609 100644 --- a/tests/unit/test_router_retry_backoff_headers.py +++ b/tests/unit/test_router_retry_backoff_headers.py @@ -2,20 +2,24 @@ Tests for router retry backoff behavior. """ +import asyncio from unittest.mock import patch import httpx import pytest +import respx import litellm from litellm import Router @pytest.mark.asyncio -async def test_retry_backoff_uses_current_exception_headers(): +async def test_retry_backoff_uses_current_exception_headers(monkeypatch: pytest.MonkeyPatch): """ Ensure retry backoff uses the current retry exception, not the initial one. """ + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() router = Router( model_list=[ { @@ -29,60 +33,36 @@ async def test_retry_backoff_uses_current_exception_headers(): num_retries=2, ) - first_error = litellm.RateLimitError( - message="Rate limited on first attempt", - model="gpt-3.5-turbo", - llm_provider="openai", - ) - first_error.litellm_response_headers = httpx.Headers({"retry-after": "1"}) + retry_after_values = (1, 15, 30) + replies = [ + httpx.Response( + 429, + headers={"retry-after": str(retry_after)}, + json={"error": {"message": f"Rate limited, retry after {retry_after}", "type": "rate_limit"}}, + ) + for retry_after in retry_after_values + ] + real_sleep = asyncio.sleep + slept: list[float] = [] - second_error = litellm.RateLimitError( - message="Rate limited on second attempt", - model="gpt-3.5-turbo", - llm_provider="openai", - ) - second_error.litellm_response_headers = httpx.Headers({"retry-after": "15"}) + async def record_sleep(delay: float, *args: object, **kwargs: object) -> None: + slept.append(delay) + await real_sleep(0) - third_error = litellm.RateLimitError( - message="Rate limited on third attempt", - model="gpt-3.5-turbo", - llm_provider="openai", - ) - third_error.litellm_response_headers = httpx.Headers({"retry-after": "30"}) + with respx.mock() as mock_router: + mock_router.post("https://api.openai.com/v1/chat/completions").mock(side_effect=replies) + with patch("asyncio.sleep", record_sleep): + with pytest.raises(litellm.RateLimitError): + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello"}], + ) + litellm.in_memory_llm_clients_cache.flush_cache() - raised_errors = [first_error, second_error, third_error] - captured_backoff_errors = [] - - async def mock_make_call(*args, **kwargs): - raise raised_errors.pop(0) - - def mock_time_to_sleep_before_retry(*args, **kwargs): - captured_backoff_errors.append(kwargs["e"]) - return 0.01 - - with patch.object(router, "make_call", side_effect=mock_make_call): - with patch.object( - router, - "_async_get_healthy_deployments", - return_value=( - [{"model_info": {"id": "test-id"}}], - [{"model_info": {"id": "test-id"}}], - ), - ): - with patch.object( - router, - "_time_to_sleep_before_retry", - side_effect=mock_time_to_sleep_before_retry, - ): - with pytest.raises(litellm.RateLimitError): - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello"}], - ) - - # Router computes backoff once after the initial failure, then once per failed retry. - # With num_retries=2 and all attempts failing, that's 1 + 2 = 3 invocations. - assert len(captured_backoff_errors) == router.num_retries + 1 - assert captured_backoff_errors[0] is first_error - assert captured_backoff_errors[1] is second_error - assert captured_backoff_errors[2] is third_error + # Router backs off once after the initial failure, then once per failed retry. + # Each backoff honours that attempt's retry-after (plus jitter), so each sleep + # sits between its own header value and the next one. + assert len(slept) == router.num_retries + 1 + assert retry_after_values[0] <= slept[0] < retry_after_values[1] + assert retry_after_values[1] <= slept[1] < retry_after_values[2] + assert retry_after_values[2] <= slept[2] diff --git a/tests/unit/test_router_retry_non_retryable_errors.py b/tests/unit/test_router_retry_non_retryable_errors.py index 98a7db7a079..0ae03d18e63 100644 --- a/tests/unit/test_router_retry_non_retryable_errors.py +++ b/tests/unit/test_router_retry_non_retryable_errors.py @@ -10,60 +10,34 @@ Verifies that: Regression tests for https://github.com/BerriAI/litellm/issues/21343 """ -import asyncio -import datetime -from collections.abc import Awaitable, Callable -from typing import Final -from unittest.mock import AsyncMock, patch +from collections.abc import Iterator +import httpx import pytest +import respx import litellm from litellm import Router -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( - _union_duration_ms, - response_timing_metrics, -) -from litellm.litellm_core_utils.logging_utils import track_llm_api_timing -from litellm.litellm_core_utils.rules import Rules -from litellm.utils import function_setup -def _make_rate_limit_error(message="Rate limited"): - """Create a RateLimitError for testing.""" - return litellm.RateLimitError( - message=message, - llm_provider="bedrock", - model="anthropic.claude-v2", +@pytest.fixture(autouse=True) +def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +def _error(status: int, message: str, code: str | None = None) -> httpx.Response: + return httpx.Response( + status, + headers={"retry-after": "0"}, + json={"error": {"message": message, "type": "invalid_request_error", "code": code}}, ) -def _make_context_window_error(message="prompt is too long: 1205821 tokens > 200000"): - """Create a ContextWindowExceededError for testing.""" - return litellm.ContextWindowExceededError( - message=message, - llm_provider="vertex_ai", - model="claude-3-opus", - ) - - -def _make_bad_request_error(message="Invalid request"): - """Create a BadRequestError for testing.""" - return litellm.BadRequestError( - message=message, - llm_provider="openai", - model="gpt-4", - ) - - -def _make_not_found_error(message="Model not found"): - """Create a NotFoundError for testing.""" - return litellm.NotFoundError( - message=message, - llm_provider="openai", - model="gpt-99", - ) +def _rate_limit(attempt: int) -> httpx.Response: + return _error(429, f"Rate limit attempt {attempt}") def _create_router(num_retries=2): @@ -74,29 +48,21 @@ def _create_router(num_retries=2): "model_name": "test-model", "litellm_params": { "model": "openai/gpt-4", - "api_key": "fake-key-1", + "api_key": f"fake-key-{i}", + "api_base": "https://api.openai.com/v1", }, - }, - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "fake-key-2", - }, - }, + } + for i in (1, 2) ], num_retries=num_retries, ) -def _base_kwargs(): - """Return kwargs required by async_function_with_retries.""" - return { - "model": "test-model", - "messages": [{"role": "user", "content": "test"}], - "original_function": AsyncMock(), - "metadata": {}, - } +async def _call(router: Router, responses: list[httpx.Response]) -> None: + """Run one acompletion against upstream replies served in order.""" + with respx.mock() as mock_router: + mock_router.post("https://api.openai.com/v1/chat/completions").mock(side_effect=responses) + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "test"}]) @pytest.mark.asyncio @@ -110,35 +76,14 @@ async def test_non_retryable_error_in_retry_loop_raises_immediately(): Expected: ContextWindowExceededError is raised, NOT RateLimitError """ router = _create_router(num_retries=2) + replies = [ + _rate_limit(1), + _error(400, "This model's maximum context length is 8192 tokens", "context_length_exceeded"), + _rate_limit(3), + ] - rate_limit_error = _make_rate_limit_error() - context_window_error = _make_context_window_error() - - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise rate_limit_error - else: - raise context_window_error - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): - with pytest.raises(litellm.ContextWindowExceededError): - await router.async_function_with_retries( - num_retries=2, - **_base_kwargs(), - ) + with pytest.raises(litellm.ContextWindowExceededError): + await _call(router, replies) @pytest.mark.asyncio @@ -149,34 +94,8 @@ async def test_bad_request_error_in_retry_loop_raises_immediately(): """ router = _create_router(num_retries=2) - rate_limit_error = _make_rate_limit_error() - bad_request_error = _make_bad_request_error() - - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise rate_limit_error - else: - raise bad_request_error - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): - with pytest.raises(litellm.BadRequestError): - await router.async_function_with_retries( - num_retries=2, - **_base_kwargs(), - ) + with pytest.raises(litellm.BadRequestError): + await _call(router, [_rate_limit(1), _error(400, "Invalid request"), _rate_limit(3)]) @pytest.mark.asyncio @@ -187,30 +106,10 @@ async def test_original_exception_updated_to_latest_error(): """ router = _create_router(num_retries=2) - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - raise _make_rate_limit_error(f"Rate limit attempt {call_count}") - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): - with pytest.raises(litellm.RateLimitError) as exc_info: - await router.async_function_with_retries( - num_retries=2, - **_base_kwargs(), - ) - # Should be the LAST error, not the first - assert "Rate limit attempt 3" in str(exc_info.value) + with pytest.raises(litellm.RateLimitError) as exc_info: + await _call(router, [_rate_limit(i) for i in range(1, 6)]) + # Should be the LAST error, not the first + assert "Rate limit attempt 3" in str(exc_info.value) @pytest.mark.asyncio @@ -220,32 +119,15 @@ async def test_retryable_errors_still_retry_normally(): configured number of times before raising. """ router = _create_router(num_retries=3) + replies = [_rate_limit(i) for i in range(1, 9)] - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - raise _make_rate_limit_error(f"Rate limit attempt {call_count}") - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): + with respx.mock() as mock_router: + route = mock_router.post("https://api.openai.com/v1/chat/completions").mock(side_effect=replies) with pytest.raises(litellm.RateLimitError): - await router.async_function_with_retries( - num_retries=3, - **_base_kwargs(), - ) + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "test"}]) - # Initial call + 3 retries = 4 total calls - assert call_count == 4 + # Initial call + 3 retries = 4 total calls + assert route.call_count == 4 @pytest.mark.asyncio @@ -254,106 +136,12 @@ async def test_not_found_error_in_retry_loop_raises_immediately(): A 404 NotFoundError inside the retry loop should break out immediately. """ router = _create_router(num_retries=2) + replies = [_rate_limit(1), _error(404, "Model not found"), _rate_limit(3)] - rate_limit_error = _make_rate_limit_error() - not_found_error = _make_not_found_error() - - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise rate_limit_error - else: - raise not_found_error - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): + with respx.mock() as mock_router: + route = mock_router.post("https://api.openai.com/v1/chat/completions").mock(side_effect=replies) with pytest.raises(litellm.NotFoundError): - await router.async_function_with_retries( - num_retries=2, - **_base_kwargs(), - ) + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "test"}]) - # Only 2 calls: initial + first retry that hits non-retryable - assert call_count == 2 - - -@pytest.mark.asyncio -async def test_retry_attempts_accumulate_timing_in_shared_request_metadata(): - received_at: Final = datetime.datetime.now() - metadata: dict[str, object] = { - "model_group": "test-model", - "litellm_received_at": received_at, - } - logging_obj_raw, _ = function_setup( - "acompletion", - Rules(), - datetime.datetime.now(), - model="test-model", - messages=[{"role": "user", "content": "test"}], - metadata=metadata, - litellm_call_id="retry-timing-test", - is_async_call=True, - ) - assert isinstance(logging_obj_raw, Logging) - logging_obj: Final[Logging] = logging_obj_raw - attempt_numbers: list[int] = [] - metadata_ids: list[int] = [] - - @track_llm_api_timing() - async def timed_attempt(*, logging_obj: Logging, **kwargs: object) -> str: - del kwargs - attempt_numbers.append(len(attempt_numbers) + 1) - metadata_ids.append(id(logging_obj.model_call_details["litellm_params"]["metadata"])) - await asyncio.sleep(0.01) - if len(attempt_numbers) == 1: - raise _make_rate_limit_error() - return "success" - - async def invoke(original_function: Callable[..., Awaitable[str]], *args: object, **kwargs: object) -> str: - return await original_function(*args, **kwargs) - - router = _create_router(num_retries=1) - with ( - patch.object(router, "make_call", new=AsyncMock(side_effect=invoke)), - patch.object( - router, - "_async_get_healthy_deployments", - new=AsyncMock(return_value=(["d1"], ["d1"])), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - ): - result = await router.async_function_with_retries( - original_function=timed_attempt, - model="test-model", - messages=[{"role": "user", "content": "test"}], - metadata=metadata, - logging_obj=logging_obj, - num_retries=1, - ) - - request_metadata: Final = logging_obj.model_call_details["litellm_params"]["metadata"] - windows: Final = request_metadata["llm_api_timing_windows"] - end_time: Final = datetime.datetime.fromtimestamp(max(window[1] for window in windows)) - timing_metrics: Final = response_timing_metrics(received_at, end_time, logging_obj) - assert result == "success" - assert attempt_numbers == [1, 2] - assert request_metadata is metadata - assert metadata_ids == [id(metadata), id(metadata)] - assert len(windows) == 2 - union_duration_ms: Final = _union_duration_ms(windows, received_at.timestamp(), end_time.timestamp()) - assert union_duration_ms is not None - total_response_time_ms: Final = (end_time.timestamp() - received_at.timestamp()) * 1000 - assert timing_metrics["litellm_overhead_time_ms"] == pytest.approx( - round(total_response_time_ms - union_duration_ms, 4) - ) + # Only 2 calls: initial + first retry that hits non-retryable + assert route.call_count == 2 diff --git a/tests/unit/test_router_retry_policy_update.py b/tests/unit/test_router_retry_policy_update.py index 6ab672d7b04..95a0d495456 100644 --- a/tests/unit/test_router_retry_policy_update.py +++ b/tests/unit/test_router_retry_policy_update.py @@ -88,9 +88,7 @@ def test_update_router_config_rejects_malformed_model_group_retry_policy(): """model_group_retry_policy is Dict[str, RetryPolicy], so each per-group policy is validated the same way.""" with pytest.raises(ValidationError): - UpdateRouterConfig( - model_group_retry_policy={"gpt-4": {"RateLimitErrorRetries": "x"}} - ) + UpdateRouterConfig(model_group_retry_policy={"gpt-4": {"RateLimitErrorRetries": "x"}}) # --------------------------------------------------------------------------- @@ -146,16 +144,6 @@ def test_set_optional_pre_call_checks_reconciles_callback_types(): assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks) -def test_remove_optional_pre_call_check_removes_local_and_global_callbacks(): - router = _build_router() - - router.set_optional_pre_call_checks(["prompt_caching"]) - router._remove_optional_callbacks_of_type(PromptCachingDeploymentCheck) - - assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router.optional_callbacks or [])) - assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks) - - def test_remove_optional_pre_call_check_keeps_global_callback_for_another_router(): router_a = _build_router() router_b = _build_router() @@ -255,7 +243,9 @@ def test_update_settings_keeps_per_group_encrypted_content_affinity_when_global_ router.update_settings(optional_pre_call_checks=[]) affinity_checks: Final = [ - callback for callback in (router.optional_callbacks or []) if isinstance(callback, EncryptedContentAffinityCheck) + callback + for callback in (router.optional_callbacks or []) + if isinstance(callback, EncryptedContentAffinityCheck) ] assert len(affinity_checks) == 1 assert affinity_checks[0].enable_global_affinity is False @@ -470,9 +460,7 @@ async def test_config_update_persists_and_reads_back_retry_policy(monkeypatch): read_back = ( await proxy_server.get_config( - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-9876" - ) + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-9876") ) )["router_settings"]["retry_policy"] assert read_back.BadRequestErrorRetries == 5 diff --git a/tests/unit/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py index fc787fcf74a..c7dd3bfa78c 100644 --- a/tests/unit/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -4,7 +4,6 @@ import json import threading import time from collections.abc import Callable, Generator, Mapping -from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -16,8 +15,6 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.router import Router -from litellm.router import _silent_experiment_kwargs_snapshot -from litellm.router import _silent_experiment_targets class _RecordingLogger(CustomLogger): @@ -197,105 +194,6 @@ def _streaming_model_list(silent_model: object) -> list[dict[str, object]]: ] -class _NonCopyableSpan: - """Mimics an OTel Span which raises on deepcopy, forcing safe_deep_copy - to fall back to the original reference.""" - - def __deepcopy__(self, memo): - raise TypeError("OTel spans cannot be deepcopied") - - -class _FakeUserAPIKeyAuth: - """Mimics UserAPIKeyAuth which contains a parent_otel_span that is not - deepcopy-able. This is what actually causes safe_deep_copy to fail for - the metadata dict in production — safe_deep_copy handles the top-level - litellm_parent_otel_span specially (pops it before copying), but does - NOT handle user_api_key_auth.parent_otel_span inside it.""" - - def __init__(self, key_alias, parent_otel_span): - self.key_alias = key_alias - self.parent_otel_span = parent_otel_span - - def __deepcopy__(self, memo): - raise TypeError("Contains OTel span that cannot be deepcopied") - - -def test_get_silent_experiment_kwargs(): - """ - Test _get_silent_experiment_kwargs returns isolated kwargs with silent experiment metadata. - - Uses a non-copyable user_api_key_auth (mimicking the real proxy scenario) - so that safe_deep_copy falls back to the original metadata reference — - exercising the identity-check fix path. - """ - model_list = [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key"}, - }, - ] - router = Router(model_list=model_list) - mock_span = _NonCopyableSpan() - mock_auth = _FakeUserAPIKeyAuth( - key_alias="HaneefKeyNonTeamProd", - parent_otel_span=mock_span, - ) - kwargs = { - "metadata": { - "foo": "bar", - "litellm_parent_otel_span": mock_span, - "user_api_key_auth": mock_auth, - }, - "litellm_call_id": "call-123", - "stream": True, - "proxy_server_request": {"body": {"model": "test"}}, - } - result = router._get_silent_experiment_kwargs(**kwargs) - assert result["metadata"]["is_silent_experiment"] is True - assert result["metadata"]["foo"] == "bar" - assert "litellm_call_id" not in result - assert result["stream"] is True - # proxy_server_request must be preserved for spend log metadata - assert "proxy_server_request" in result - # CRITICAL: metadata must be a DIFFERENT dict object than the original, - # so that setting model_group / is_silent_experiment on the silent dict - # doesn't corrupt the primary call's metadata. - assert result["metadata"] is not kwargs["metadata"] - # OTel span must be stripped from the silent copy — it's not safe to use - # across event loops (silent experiment runs in a new event loop). - assert "litellm_parent_otel_span" not in result["metadata"] - # Original metadata must NOT be mutated — must carry the real span, - # not safe_deep_copy's temporary "placeholder" string. - assert "is_silent_experiment" not in kwargs["metadata"] - assert kwargs["metadata"]["litellm_parent_otel_span"] is mock_span - assert kwargs["metadata"]["user_api_key_auth"] is mock_auth - # Shallow copy must preserve user_api_key_auth so the silent experiment - # can attribute billing / spend logs to the correct key/team. - assert result["metadata"]["user_api_key_auth"] is mock_auth - - -def test_get_silent_experiment_kwargs_without_stream_stays_non_streaming(): - router = Router(model_list=[{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k"}}]) - result = router._get_silent_experiment_kwargs(metadata={"foo": "bar"}, stream=False) - assert result["stream"] is False - assert "stream" not in router._get_silent_experiment_kwargs(metadata={"foo": "bar"}) - - -@pytest.mark.parametrize( - "silent_model, expected", - [ - ("shadow-a", ("shadow-a",)), - (["shadow-a", "shadow-b"], ("shadow-a", "shadow-b")), - ([], ()), - (None, ()), - (42, ()), - (["shadow-a", 42], ()), - ], -) -def test_silent_experiment_targets(silent_model, expected): - assert _silent_experiment_targets(silent_model) == expected - - @pytest.mark.asyncio async def test_streaming_shadow_is_streamed_and_drained_async(recording_logger): router = Router(model_list=_streaming_model_list("shadow-a")) @@ -594,87 +492,6 @@ def _tagged_primary_model_list() -> list[dict[str, object]]: ] -def test_silent_experiment_kwargs_snapshot_is_isolated_from_later_primary_mutations(): - metadata = {"foo": "bar"} - kwargs: dict[str, object] = {"metadata": metadata, "stream": True} - snapshot = _silent_experiment_kwargs_snapshot(kwargs) - kwargs["messages"] = [{"role": "user", "content": "added by the primary"}] - metadata["tags"] = ["primary-only"] - - assert dict(snapshot) == {"metadata": {"foo": "bar"}, "stream": True} - assert dict(_silent_experiment_kwargs_snapshot({"stream": False, "metadata": None})) == { - "stream": False, - "metadata": None, - } - - -def test_sync_shadow_gets_kwargs_snapshot_taken_before_primary_mutates_them(recording_logger): - deferred: list[Callable[[], None]] = [] - - class _DeferredThread: - def __init__(self, target, args, kwargs, daemon) -> None: - deferred.append(lambda: target(*args, **kwargs)) - - def start(self) -> None: - return None - - router = Router(model_list=_tagged_primary_model_list()) - with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race - "litellm.router_python_core.threading", SimpleNamespace(Thread=_DeferredThread) - ): - response = router.completion( - model="primary-model", - messages=[{"role": "user", "content": "hi"}], - mock_response="pong", - metadata={"foo": "bar"}, - ) - assert response.choices[0].message.content == "pong" - assert len(deferred) == 1 - deferred[0]() - _wait_for_shadow_successes_sync(recording_logger, expected=1) - - shadow_successes = recording_logger.shadow_successes() - assert len(shadow_successes) == 1 - shadow_metadata = shadow_successes[0]["litellm_params"]["metadata"] - assert shadow_metadata["model_group"] == "shadow-b" - assert "primary-only" not in shadow_metadata.get("tags", []) - - -def test_sync_shadow_workers_do_not_share_metadata_with_each_other(recording_logger): - workers: list[tuple[Mapping[str, object], Callable[[], None]]] = [] - - class _DeferredThread: - def __init__(self, target, args, kwargs, daemon) -> None: - workers.append((kwargs, lambda: target(*args, **kwargs))) - - def start(self) -> None: - return None - - router = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"])) - with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race - "litellm.router_python_core.threading", SimpleNamespace(Thread=_DeferredThread) - ): - router.completion( - model="primary-model", - messages=[{"role": "user", "content": "hi"}], - mock_response="pong", - metadata={"foo": "bar"}, - ) - assert len(workers) == 2 - (first_kwargs, run_first), (_, run_second) = workers - first_kwargs["metadata"].pop("foo") - run_second() - run_first() - _wait_for_shadow_successes_sync(recording_logger, expected=2) - - metadata_by_group = { - call["litellm_params"]["metadata"]["model_group"]: call["litellm_params"]["metadata"] - for call in recording_logger.shadow_successes() - } - assert metadata_by_group["shadow-b"]["foo"] == "bar" - assert "foo" not in metadata_by_group["shadow-a"] - - @pytest.mark.asyncio async def test_async_shadow_does_not_inherit_primary_deployment_tags(recording_logger): router = Router(model_list=_tagged_primary_model_list()) @@ -707,25 +524,6 @@ async def test_shadow_of_a_shadow_is_not_launched(recording_logger): assert model_groups == ["shadow-a"] -@pytest.mark.asyncio -async def test_run_silent_experiment_drains_stream_so_callbacks_fire(recording_logger): - router = Router(model_list=_streaming_model_list(None)) - silent_kwargs: Final = { - "stream": True, - "stream_options": {"include_usage": True}, - "mock_response": "pong", - "metadata": {"is_silent_experiment": True, "model_group": "shadow-b"}, - } - await router._run_silent_experiment("shadow-b", [{"role": "user", "content": "hi"}], silent_kwargs) - await _wait_for_shadow_successes(recording_logger, expected=1) - - shadow_successes = recording_logger.shadow_successes() - assert len(shadow_successes) == 1 - assert shadow_successes[0]["stream"] is True - assert shadow_successes[0]["async_complete_streaming_response"] is not None - assert silent_kwargs["stream"] is True - - @pytest.mark.asyncio async def test_router_silent_experiment_acompletion(): """ @@ -784,11 +582,7 @@ async def test_router_silent_experiment_acompletion(): # Find the silent call silent_call = next( - ( - c - for c in call_args_list - if c[1].get("metadata", {}).get("is_silent_experiment") is True - ), + (c for c in call_args_list if c[1].get("metadata", {}).get("is_silent_experiment") is True), None, ) assert silent_call is not None @@ -796,11 +590,7 @@ async def test_router_silent_experiment_acompletion(): # Find the primary call primary_call = next( - ( - c - for c in call_args_list - if not c[1].get("metadata", {}).get("is_silent_experiment") - ), + (c for c in call_args_list if not c[1].get("metadata", {}).get("is_silent_experiment")), None, ) assert primary_call is not None @@ -869,11 +659,7 @@ def test_router_silent_experiment_completion(): # Find the silent call silent_call = next( - ( - c - for c in call_args_list - if c[1].get("metadata", {}).get("is_silent_experiment") is True - ), + (c for c in call_args_list if c[1].get("metadata", {}).get("is_silent_experiment") is True), None, ) assert silent_call is not None @@ -882,35 +668,6 @@ def test_router_silent_experiment_completion(): assert silent_call[1]["metadata"]["model_group"] == "silent-model" -SILENT_EXPERIMENT_RUNNERS: Final = ( - pytest.param(lambda router, **kwargs: router._silent_experiment_completion(**kwargs), id="sync"), - pytest.param(lambda router, **kwargs: asyncio.run(router._silent_experiment_acompletion(**kwargs)), id="async"), -) - - -@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS) -def test_silent_experiment_sends_shadow_request_attributed_to_the_silent_model(run_silent_experiment): - router = Router(model_list=_streaming_model_list(["shadow-a"])) - primary_metadata: Final = {"model_group": "primary-model"} - with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion: - run_silent_experiment( - router, - silent_model="shadow-a", - messages=[{"role": "user", "content": "hi"}], - metadata=primary_metadata, - ) - - acompletion.assert_awaited_once() - shadow_call: Final = acompletion.await_args.kwargs - assert shadow_call["model"] == "shadow-a" - assert shadow_call["messages"] == [{"role": "user", "content": "hi"}] - assert shadow_call["metadata"]["model_group"] == "shadow-a" - assert shadow_call["metadata"]["is_silent_experiment"] is True - assert primary_metadata == {"model_group": "primary-model"} - - - - _EMBEDDING_API_BASE: Final = "https://embeddings.example.test/v1" @@ -972,20 +729,6 @@ async def test_aembedding_with_silent_model_sends_provider_body_without_it( assert response.data[0]["embedding"] == [0.1, 0.2] -@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS) -def test_silent_experiment_does_not_launch_from_a_shadow_request(run_silent_experiment): - router = Router(model_list=_streaming_model_list(["shadow-a"])) - with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion: - run_silent_experiment( - router, - silent_model="shadow-a", - messages=[{"role": "user", "content": "hi"}], - metadata={"is_silent_experiment": True}, - ) - - acompletion.assert_not_awaited() - - def _judged_router(shadow_outcome: str) -> Router: return Router( model_list=[ diff --git a/tests/unit/test_router_streaming_fallback_metadata.py b/tests/unit/test_router_streaming_fallback_metadata.py index 40449d8cc66..09fa9d41c1b 100644 --- a/tests/unit/test_router_streaming_fallback_metadata.py +++ b/tests/unit/test_router_streaming_fallback_metadata.py @@ -1,6 +1,5 @@ import json from typing import Final -from unittest.mock import MagicMock import pytest @@ -8,110 +7,6 @@ import litellm from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.proxy.proxy_server import _should_include_fallback_errors from litellm.router import Router -from litellm.types.llms.base import HiddenParams - - -def test_apply_fallback_hidden_params_copies_from_fallback_response(): - fallback_errors = [ - { - "message": "litellm.RateLimitError: upstream limited request", - "type": "RateLimitError", - "param": None, - "code": "429", - } - ] - chunk = litellm.ModelResponseStream( - id="test", - model="openai/internal-fallback", - choices=[], - ) - chunk._hidden_params = { - "additional_headers": {"x-existing-chunk-header": "keep"}, - "model_id": "chunk-model-id", - } - fallback_response = MagicMock() - fallback_response._hidden_params = { - "additional_headers": { - "x-litellm-attempted-fallbacks": 1, - "x-litellm-model-group": "fallback-model", - "x-litellm-fallback-errors": json.dumps(fallback_errors), - }, - "api_base": "https://fallback.example", - } - - Router._apply_fallback_hidden_params_to_item( - fallback_item=chunk, - prepared_fallback_hidden_params=Router._prepare_fallback_hidden_params( - fallback_response - ), - ) - - assert chunk._hidden_params["api_base"] == "https://fallback.example" - assert chunk._hidden_params["model_id"] == "chunk-model-id" - assert chunk._hidden_params["additional_headers"] == { - "x-existing-chunk-header": "keep", - "x-litellm-attempted-fallbacks": 1, - "x-litellm-model-group": "fallback-model", - "x-litellm-fallback-errors": json.dumps(fallback_errors), - } - - -def test_apply_fallback_hidden_params_updates_plain_duck_chunk(): - class PlainChunk: - def __init__(self) -> None: - self._hidden_params = { - "additional_headers": {"x-existing-chunk-header": "keep"}, - "model_id": "chunk-model-id", - } - - chunk: Final = PlainChunk() - fallback_response = { - "_hidden_params": { - "additional_headers": { - "x-litellm-attempted-fallbacks": 1, - }, - "api_base": "https://fallback.example", - } - } - - Router._apply_fallback_hidden_params_to_item( - fallback_item=chunk, - prepared_fallback_hidden_params=Router._prepare_fallback_hidden_params(fallback_response), - ) - - assert chunk._hidden_params["api_base"] == "https://fallback.example" - assert chunk._hidden_params["model_id"] == "chunk-model-id" - assert chunk._hidden_params["additional_headers"] == { - "x-existing-chunk-header": "keep", - "x-litellm-attempted-fallbacks": 1, - } - - -def test_apply_fallback_hidden_params_normalizes_hidden_params_on_plain_duck_chunk(): - class PlainChunk: - def __init__(self) -> None: - self._hidden_params = HiddenParams( - api_base="https://original.example", - model_id="original-model-id", - additional_headers={"x-existing-chunk-header": "keep"}, - ) - - chunk: Final = PlainChunk() - Router._apply_fallback_hidden_params_to_item( - fallback_item=chunk, - prepared_fallback_hidden_params=( - {"api_base": "https://fallback.example", "model_id": "fallback-model-id"}, - {"x-litellm-attempted-fallbacks": 1}, - ), - ) - - assert isinstance(chunk._hidden_params, dict) - assert chunk._hidden_params["api_base"] == "https://fallback.example" - assert chunk._hidden_params["model_id"] == "fallback-model-id" - assert chunk._hidden_params["additional_headers"] == { - "x-existing-chunk-header": "keep", - "x-litellm-attempted-fallbacks": 1, - } def _two_group_fallback_router() -> Router: @@ -183,43 +78,6 @@ async def test_router_omits_fallback_errors_without_opt_in(): assert "x-litellm-fallback-errors" not in headers -def test_prepare_fallback_hidden_params_no_additional_headers(): - class FakeResponse: - _hidden_params = {"api_base": "http://example.com"} - - hidden_params, headers = Router._prepare_fallback_hidden_params(FakeResponse()) - assert hidden_params == {"api_base": "http://example.com"} - assert headers == {} - - -def test_apply_fallback_hidden_params_to_item_none_item(): - Router._apply_fallback_hidden_params_to_item( - None, ({"api_base": "http://fallback.example"}, {"x-custom": "value"}) - ) - - -def test_apply_fallback_hidden_params_to_item_no_existing_additional_headers(): - chunk: Final = litellm.ModelResponseStream( - id="test", - model="openai/internal-fallback", - choices=[], - ) - chunk.hidden_params["model_id"] = "test-id" - Router._apply_fallback_hidden_params_to_item( - chunk, - ( - {"api_base": "http://fallback.example"}, - {"x-litellm-attempted-fallbacks": 1}, - ), - ) - - assert chunk.hidden_params["api_base"] == "http://fallback.example" - assert chunk.hidden_params["model_id"] == "test-id" - assert chunk.hidden_params["additional_headers"] == { - "x-litellm-attempted-fallbacks": 1 - } - - @pytest.mark.asyncio async def test_set_response_headers_adds_model_group_to_streaming_wrapper(): class StreamingWrapper: diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index 70231b620be..521122291c1 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -9,7 +9,6 @@ cross-group fallback runs. from collections import Counter from typing import Final, Optional -from unittest.mock import AsyncMock, patch import pytest @@ -75,287 +74,6 @@ class TestGetExcludedFilteredDeployments: assert len(result) == 1 -# --------------------------------------------------------------------------- -# Router helpers (router_code_coverage.py requires these names in a *router* test file) -# --------------------------------------------------------------------------- - - -def test_set_failed_deployment_id_on_exception(): - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - exc = Exception("fail") - dep = _make_dep("dep-a") - router._set_failed_deployment_id_on_exception(exc, dep) - assert getattr(exc, "failed_deployment_id", None) == "dep-a" - router._set_failed_deployment_id_on_exception(exc, _make_dep("dep-b")) - assert exc.failed_deployment_id == "dep-a" - - -def test_stamp_failed_deployment_id_with_effective_model_info_prefers_kwargs(): - """kwargs["model_info"] (the dynamic client-side-credential id, when present) must win - over the static deployment's model_info, so a bad-credential tenant's failures are - attributed to their own dynamic deployment id, not the shared static one.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - exc = Exception("fail") - router._stamp_failed_deployment_id_with_effective_model_info( - exc, _make_dep("dep-a"), {"model_info": {"id": "dynamic-dep"}} - ) - assert exc.failed_deployment_id == "dynamic-dep" - - -def test_stamp_failed_deployment_id_with_effective_model_info_falls_back_to_deployment(): - """With no dynamic id in kwargs (the common, non-client-side-credential case), the - static deployment's own model_info.id must still be stamped.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - exc = Exception("fail") - router._stamp_failed_deployment_id_with_effective_model_info(exc, _make_dep("dep-a"), {}) - assert exc.failed_deployment_id == "dep-a" - - -@pytest.mark.asyncio -async def test_ageneric_api_call_with_fallbacks_helper_stamps_failed_deployment_id(): - """_ageneric_api_call_with_fallbacks_helper must stamp failed_deployment_id on a - failure, same as _completion/_acompletion, so callers identifying the failed - deployment (cooldown, weighted failover) work for this call type too instead of - depending on which metadata bucket ("metadata" vs "litellm_metadata") it uses.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - async def _failing_original_function(**kwargs): - raise RuntimeError("boom") - - with pytest.raises(RuntimeError) as exc_info: - await router._ageneric_api_call_with_fallbacks_helper( - model="test-model", - original_generic_function=_failing_original_function, - ) - - assert getattr(exc_info.value, "failed_deployment_id", None) == "dep-a" - - -@pytest.mark.asyncio -async def test_ageneric_api_call_with_fallbacks_helper_stamps_dynamic_id_for_clientside_credentials(): - """A client-side-credential call (tenant-supplied api_key) generates a dynamic - deployment id distinct from the shared static deployment. Stamping the static id - instead would let one tenant's bad credentials cool down the deployment every - other tenant sharing this config relies on.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - async def _failing_original_function(**kwargs): - raise RuntimeError("boom") - - with pytest.raises(RuntimeError) as exc_info: - await router._ageneric_api_call_with_fallbacks_helper( - model="test-model", - original_generic_function=_failing_original_function, - api_key="tenant-supplied-key", - litellm_metadata={"model_group": "test-model"}, - ) - - failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) - assert failed_deployment_id is not None - assert failed_deployment_id != "dep-a" - - -@pytest.mark.asyncio -async def test_acompletion_stamps_dynamic_id_for_clientside_credentials(): - """Same bug as the generic-API-call helper above, but in the regular completion - path: _acompletion's exception handlers must stamp the dynamic client-side-credential - deployment id, not the shared static deployment's id.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - with patch("litellm.acompletion", new_callable=AsyncMock, side_effect=RuntimeError("boom")): - with pytest.raises(RuntimeError) as exc_info: - await router._acompletion( - model="test-model", - messages=[{"role": "user", "content": "Hello"}], - api_key="tenant-supplied-key", - metadata={"model_group": "test-model"}, - ) - - failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) - assert failed_deployment_id is not None - assert failed_deployment_id != "dep-a" - - -@pytest.mark.asyncio -async def test_acompletion_stamps_dynamic_id_for_clientside_credentials_on_timeout(): - """Same bug as the RuntimeError case above, but for the separate `except litellm.Timeout` - branch in `_acompletion`: it has its own call to the stamping helper, so a fix that only - covers the generic `except Exception` branch would leave a caller-supplied timeout - (`litellm.Timeout` is what `x-litellm-timeout` maps to) stamping the shared static id.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - timeout_exc = litellm.Timeout(message="boom", model="test-model", llm_provider="openai") - with patch("litellm.acompletion", new_callable=AsyncMock, side_effect=timeout_exc): - with pytest.raises(litellm.Timeout) as exc_info: - await router._acompletion( - model="test-model", - messages=[{"role": "user", "content": "Hello"}], - api_key="tenant-supplied-key", - metadata={"model_group": "test-model"}, - ) - - failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) - assert failed_deployment_id is not None - assert failed_deployment_id != "dep-a" - - -def test_completion_stamps_dynamic_id_for_clientside_credentials(): - """Sync counterpart: _completion's exception handler must stamp the dynamic - client-side-credential deployment id, not the shared static deployment's id.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - with patch("litellm.completion", side_effect=RuntimeError("boom")): - with pytest.raises(RuntimeError) as exc_info: - router._completion( - model="test-model", - messages=[{"role": "user", "content": "Hello"}], - api_key="tenant-supplied-key", - metadata={"model_group": "test-model"}, - ) - - failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) - assert failed_deployment_id is not None - assert failed_deployment_id != "dep-a" - - -@pytest.mark.asyncio -async def test_maybe_run_weighted_failover_returns_none_without_failed_id(): - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key", "weight": 1}, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key", "weight": 1}, - "model_info": {"id": "B"}, - }, - ], - routing_strategy="simple-shuffle", - enable_weighted_failover=True, - ) - result = await router._maybe_run_weighted_failover( - exception=Exception("fail"), - original_model_group="test-model", - all_deployments=[_make_dep("A"), _make_dep("B")], - args=(), - kwargs={"metadata": {}}, - input_kwargs={}, - ) - assert result is None - - -@pytest.mark.asyncio -async def test_maybe_run_weighted_failover_persists_excluded_ids_to_kwargs(monkeypatch): - """Regression: writing to the metadata dict returned by `setdefault` must - update the dict in `kwargs` itself so the next hop sees prior exclusions. - Previously `setdefault(..., {}) or {}` returned a disconnected dict on the - first hop, dropping `_failover_excluded_ids` writes. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "B"}, - }, - ], - routing_strategy="simple-shuffle", - enable_weighted_failover=True, - ) - - async def _stub_run_async_fallback(*args, **kwargs): - return "ok" - - monkeypatch.setattr("litellm.router_python_core.run_async_fallback", _stub_run_async_fallback) - - exc = Exception("fail") - exc.failed_deployment_id = "A" - kwargs: dict = {"metadata": {}} - await router._maybe_run_weighted_failover( - exception=exc, - original_model_group="test-model", - all_deployments=[_make_dep("A"), _make_dep("B")], - args=(), - kwargs=kwargs, - input_kwargs={}, - ) - # The dict inside kwargs must reflect the write — proves `meta` was the - # same object as kwargs["metadata"] (no disconnected copy). - assert kwargs["metadata"].get("_failover_excluded_ids") == ["A"] - - # --------------------------------------------------------------------------- # Integration tests for weighted-failover end-to-end via Router # --------------------------------------------------------------------------- @@ -758,195 +476,6 @@ async def test_user_config_two_region_failover(): assert successes["eastus2"] == 0 -# --------------------------------------------------------------------------- -# Tests for healthy-deployment-only check in _maybe_run_weighted_failover -# (Issue: weighted failover checked all deployments, not just healthy ones) -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_maybe_run_weighted_failover_skips_when_remaining_all_in_cooldown( - monkeypatch, -): - """When every non-excluded deployment is in cooldown, _maybe_run_weighted_failover - must return None immediately without invoking run_async_fallback. - - Previously the check was against all_deployments (including cooldown ones), so - run_async_fallback would be called unnecessarily and would raise RouterRateLimitError. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "B"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "C"}, - }, - ], - routing_strategy="simple-shuffle", - enable_weighted_failover=True, - ) - - # A just failed; B and C are both in cooldown. - exc = Exception("A down") - exc.failed_deployment_id = "A" - - run_async_fallback_called = False - - async def _should_not_be_called(*args, **kwargs): - nonlocal run_async_fallback_called - run_async_fallback_called = True - return "should not reach here" - - monkeypatch.setattr("litellm.router_python_core.run_async_fallback", _should_not_be_called) - - # Patch cooldown so B and C appear in cooldown. - with patch( - "litellm.router_python_core.async_get_cooldown_deployments", - new=AsyncMock(return_value=["B", "C"]), - ): - result = await router._maybe_run_weighted_failover( - exception=exc, - original_model_group="test-model", - all_deployments=[_make_dep("A"), _make_dep("B"), _make_dep("C")], - args=(), - kwargs={"metadata": {}}, - input_kwargs={}, - ) - - assert result is None, "Should return None when all remaining deployments are in cooldown" - assert not run_async_fallback_called, "run_async_fallback must NOT be called when no healthy deployments remain" - - -@pytest.mark.asyncio -async def test_maybe_run_weighted_failover_proceeds_when_one_healthy_remains( - monkeypatch, -): - """When at least one non-excluded deployment is healthy (not in cooldown), - _maybe_run_weighted_failover should still invoke run_async_fallback normally. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "B"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "C"}, - }, - ], - routing_strategy="simple-shuffle", - enable_weighted_failover=True, - ) - - # A just failed; B is in cooldown; C is healthy. - exc = Exception("A down") - exc.failed_deployment_id = "A" - - run_async_fallback_called = False - - async def _stub_run_async_fallback(*args, **kwargs): - nonlocal run_async_fallback_called - run_async_fallback_called = True - return "ok from C" - - monkeypatch.setattr("litellm.router_python_core.run_async_fallback", _stub_run_async_fallback) - - with patch( - "litellm.router_python_core.async_get_cooldown_deployments", - new=AsyncMock(return_value=["B"]), - ): - result = await router._maybe_run_weighted_failover( - exception=exc, - original_model_group="test-model", - all_deployments=[_make_dep("A"), _make_dep("B"), _make_dep("C")], - args=(), - kwargs={"metadata": {}}, - input_kwargs={}, - ) - - assert result == "ok from C" - assert run_async_fallback_called, "run_async_fallback must be called when a healthy deployment remains" - - -@pytest.mark.asyncio -async def test_failover_falls_through_to_external_fallback_when_remaining_in_cooldown(): - """End-to-end: when the only non-failed deployments are in cooldown, - weighted failover must fall through to the configured cross-group fallback. - - Without the fix the _maybe_run_weighted_failover would invoke run_async_fallback - unnecessarily (because it counted cooldown deployments as "remaining"), get back - RouterRateLimitError, return None, and reach the same fallback path — but only - incidentally. With the fix the early-exit path is taken directly. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "gpt-4o", - "api_key": "bad", - "mock_response": Exception("A down"), - "weight": 1_000_000, # always picked first - }, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": { - "model": "gpt-4o", - "api_key": "bad", - "mock_response": Exception("B down"), - "weight": 1, - }, - "model_info": {"id": "B"}, - }, - { - "model_name": "fallback-model", - "litellm_params": { - "model": "gpt-4o", - "api_key": "good", - "mock_response": "ok from fallback", - }, - "model_info": {"id": "fallback"}, - }, - ], - routing_strategy="simple-shuffle", - num_retries=0, - enable_weighted_failover=True, - fallbacks=[{"test-model": ["fallback-model"]}], - ) - - # Put B in cooldown so weighted failover can't use it after A fails. - with patch( - "litellm.router_python_core.async_get_cooldown_deployments", - new=AsyncMock(return_value=["B"]), - ): - response = await router.acompletion( - model="test-model", - messages=[{"role": "user", "content": "hi"}], - ) - - assert response._hidden_params["model_id"] == "fallback" - - @pytest.mark.asyncio @pytest.mark.parametrize( ("pre_call_check", "metadata"), @@ -1115,14 +644,14 @@ async def test_weighted_failover_skips_the_failed_deployment_for_clientside_cred @pytest.mark.parametrize( - ("endpoint", "attempt", "arguments"), + ("endpoint", "arguments"), [ - ("embedding", "_embedding", {"input": ["hi"]}), - ("image_generation", "_image_generation", {"prompt": "a cat"}), + ("embedding", {"input": ["hi"]}), + ("image_generation", {"prompt": "a cat"}), ], ) def test_a_sync_attempt_stamps_the_failed_deployment( - monkeypatch: pytest.MonkeyPatch, endpoint: str, attempt: str, arguments: dict[str, object] + monkeypatch: pytest.MonkeyPatch, endpoint: str, arguments: dict[str, object] ) -> None: def provider(**kwargs: object) -> object: raise litellm.InternalServerError(message="down", llm_provider="openai", model="m") @@ -1139,7 +668,7 @@ def test_a_sync_attempt_stamps_the_failed_deployment( ) with pytest.raises(litellm.InternalServerError) as raised: - getattr(router, attempt)(model="test-model", **arguments) + getattr(router, endpoint)(model="test-model", **arguments) assert getattr(raised.value, "failed_deployment_id", None) == "A" @@ -1157,6 +686,6 @@ async def test_a_failed_pick_stamps_no_deployment() -> None: ) with pytest.raises(litellm.BadRequestError, match="no healthy deployments") as raised: - await router._aembedding(input=["hi"], model="no-such-model") + await router.aembedding(input=["hi"], model="no-such-model") assert getattr(raised.value, "failed_deployment_id", None) is None diff --git a/tests/unit/test_streaming_connection_cleanup.py b/tests/unit/test_streaming_connection_cleanup.py index 39fcee8d44d..c525e917b4c 100644 --- a/tests/unit/test_streaming_connection_cleanup.py +++ b/tests/unit/test_streaming_connection_cleanup.py @@ -3,13 +3,15 @@ Regression tests for streaming connection pool leak fix. """ import asyncio -from unittest.mock import MagicMock, patch +import json +from unittest.mock import MagicMock import anyio import httpx import pytest +import respx - +import litellm from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.custom_httpx.aiohttp_transport import ( AiohttpResponseStream, @@ -55,9 +57,7 @@ async def test_aiohttp_transport_response_uses_stream_not_content(): return Resp() transport = LiteLLMAiohttpTransport(client=lambda: FakeSession()) # type: ignore - response = await transport.handle_async_request( - httpx.Request("GET", "http://example.com") - ) + response = await transport.handle_async_request(httpx.Request("GET", "http://example.com")) assert isinstance(response.stream, AiohttpResponseStream) @@ -167,224 +167,121 @@ async def test_aclose_completes_under_cancellation(): # ── Router stream_with_fallbacks cleanup tests ────────────────── -@pytest.mark.asyncio -async def test_stream_with_fallbacks_closes_stream_on_generator_close(): - """Closing the FallbackStreamWrapper must aclose() the underlying model_response - via stream_with_fallbacks' finally block.""" +def _sse_chunk(content: str, finish_reason: str | None = None) -> bytes: + payload = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4", + "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(payload)}\n\n".encode() + + +_DONE = b"data: [DONE]\n\n" + + +class _TrackingStream(httpx.AsyncByteStream): + """An upstream SSE body that records whether the router closed it.""" + + def __init__(self, parts: list[bytes], fail: bool = False): + self._parts = parts + self._fail = fail + self.closed = False + + async def __aiter__(self): + for part in self._parts: + yield part + if self._fail: + raise httpx.ReadError("upstream dropped the connection") + + async def aclose(self): + self.closed = True + + +def _full_stream() -> _TrackingStream: + return _TrackingStream([_sse_chunk("c1"), _sse_chunk("c2"), _sse_chunk("c3", "stop"), _DONE]) + + +@pytest.fixture +def stream_router(monkeypatch: pytest.MonkeyPatch): from litellm.router import Router - stream_closed = False - - class FakeStream(CustomStreamWrapper): - def __init__(self): - super().__init__( - completion_stream=None, - model="test-model", - logging_obj=MagicMock(), - custom_llm_provider="openai", - ) - self._items = ["chunk1", "chunk2", "chunk3"] - self._index = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - if self._index >= len(self._items): - raise StopAsyncIteration - item = self._items[self._index] - self._index += 1 - return item - - async def aclose(self): - nonlocal stream_closed - stream_closed = True - - router = Router( + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield Router( model_list=[ { "model_name": "test-model", "litellm_params": { "model": "openai/test", "api_key": "fake", + "api_base": "https://primary.example/v1", }, - } - ] - ) - - fake_stream = FakeStream() - - # Call _acompletion_streaming_iterator directly so we go through - # stream_with_fallbacks and its finally block - result = await router._acompletion_streaming_iterator( - model_response=fake_stream, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "test-model"}, - ) - - # Consume one chunk then close (simulates client disconnect) - async for _ in result: - break - await result.aclose() - - assert ( - stream_closed - ), "model_response stream was not closed by stream_with_fallbacks finally block" - - -@pytest.mark.asyncio -async def test_stream_with_fallbacks_closes_stream_on_normal_completion(): - """stream_with_fallbacks must aclose() model_response even on normal completion.""" - from litellm.router import Router - - stream_closed = False - - class FakeStream(CustomStreamWrapper): - def __init__(self): - super().__init__( - completion_stream=None, - model="test-model", - logging_obj=MagicMock(), - custom_llm_provider="openai", - ) - self._items = ["chunk1"] - self._index = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - if self._index >= len(self._items): - raise StopAsyncIteration - item = self._items[self._index] - self._index += 1 - return item - - async def aclose(self): - nonlocal stream_closed - stream_closed = True - - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/test", - "api_key": "fake", - }, - } - ] - ) - - fake_stream = FakeStream() - - result = await router._acompletion_streaming_iterator( - model_response=fake_stream, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "test-model"}, - ) - - # Exhaust the stream fully - async for _ in result: - pass - await result.aclose() - - assert stream_closed, "model_response stream was not closed after normal completion" - - -@pytest.mark.asyncio -async def test_stream_with_fallbacks_closes_both_on_fallback_disconnect(): - """When a fallback is triggered and the client disconnects during fallback - iteration, both model_response and fallback_response must be closed.""" - from litellm.exceptions import MidStreamFallbackError - from litellm.router import Router - - model_closed = False - fallback_closed = False - - class FakeModelStream(CustomStreamWrapper): - """Stream that raises MidStreamFallbackError immediately to trigger fallback.""" - - def __init__(self): - super().__init__( - completion_stream=None, - model="test-model", - logging_obj=MagicMock(), - custom_llm_provider="openai", - ) - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - raise MidStreamFallbackError( - message="test mid-stream error", - model="test-model", - llm_provider="openai", - generated_content="", - ) - - async def aclose(self): - nonlocal model_closed - model_closed = True - - class FakeFallbackStream: - """Fallback stream that yields chunks.""" - - def __init__(self): - self._items = ["fb1", "fb2", "fb3"] - self._index = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - if self._index >= len(self._items): - raise StopAsyncIteration - item = self._items[self._index] - self._index += 1 - return item - - async def aclose(self): - nonlocal fallback_closed - fallback_closed = True - - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/test", - "api_key": "fake", - }, - } - ] - ) - - fake_model_stream = FakeModelStream() - fake_fallback_stream = FakeFallbackStream() - - # Mock async_function_with_fallbacks_common_utils to return the fallback stream - # instead of actually calling through the full fallback machinery - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=fake_fallback_stream, - ): - result = await router._acompletion_streaming_iterator( - model_response=fake_model_stream, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={ - "model": "test-model", - "fallbacks": ["other-model"], }, - ) + { + "model_name": "other-model", + "litellm_params": { + "model": "openai/test", + "api_key": "fake", + "api_base": "https://fallback.example/v1", + }, + }, + ], + fallbacks=[{"test-model": ["other-model"]}], + ) + litellm.in_memory_llm_clients_cache.flush_cache() - # Consume one fallback chunk then close (simulates client disconnect) + +async def _open_stream(router, primary: _TrackingStream, fallback: _TrackingStream, mock_router): + for host, body in (("primary", primary), ("fallback", fallback)): + mock_router.post(f"https://{host}.example/v1/chat/completions").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body) + ) + return await router.acompletion(model="test-model", messages=[{"role": "user", "content": "hi"}], stream=True) + + +@pytest.mark.asyncio +async def test_stream_with_fallbacks_closes_stream_on_generator_close(stream_router): + """Closing the router's stream after one chunk (a client disconnect) must + close the upstream response.""" + primary = _full_stream() + + with respx.mock(assert_all_called=False) as mock_router: + result = await _open_stream(stream_router, primary, _full_stream(), mock_router) async for _ in result: break await result.aclose() - assert model_closed, "model_response stream was not closed" - assert fallback_closed, "fallback_response stream was not closed" + assert primary.closed, "upstream stream was not closed when the client disconnected" + + +@pytest.mark.asyncio +async def test_stream_with_fallbacks_closes_stream_on_normal_completion(stream_router): + """The upstream response must be closed even on normal completion.""" + primary = _full_stream() + + with respx.mock(assert_all_called=False) as mock_router: + result = await _open_stream(stream_router, primary, _full_stream(), mock_router) + async for _ in result: + pass + await result.aclose() + + assert primary.closed, "upstream stream was not closed after normal completion" + + +@pytest.mark.asyncio +async def test_stream_with_fallbacks_closes_both_on_fallback_disconnect(stream_router): + """When a fallback is triggered and the client disconnects during fallback + iteration, both the primary and the fallback upstream must be closed.""" + primary = _TrackingStream([], fail=True) + fallback = _full_stream() + + with respx.mock() as mock_router: + result = await _open_stream(stream_router, primary, fallback, mock_router) + async for _ in result: + break + await result.aclose() + + assert primary.closed, "primary upstream stream was not closed" + assert fallback.closed, "fallback upstream stream was not closed" diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index e24499db071..c41c4db2c02 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -104,6 +104,7 @@ from litellm.utils import ( _snapshot_exception_for_hook, async_post_call_failure_deployment_hook, async_post_call_success_deployment_hook, + async_pre_call_deployment_hook, calculate_max_parallel_requests, client, get_non_default_completion_params, @@ -8393,6 +8394,7 @@ def test_get_model_info_gemini(monkeypatch): and "veo" not in model and "lyria" not in model and "robotics" not in model + and "omni" not in model and "3.8-flash-tts" not in model and "3.8-flash-lite-tts" not in model ): @@ -11625,3 +11627,25 @@ async def test_success_enqueue_is_tracked_without_starting_callbacks_synchronous assert GLOBAL_LOGGING_WORKER.pending_count > 0, "Tracked enqueue must finish before worker draining starts" await GLOBAL_LOGGING_WORKER.flush() assert logging_obj.model_call_details["has_logged_async_success"] is True + + +class _RecordingPreCallDeploymentHook(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen_call_types: tuple[CallTypes | None, ...] = () + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + self.seen_call_types = (*self.seen_call_types, call_type) + return None + + +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_hands_responses_compact_its_call_type(monkeypatch: pytest.MonkeyPatch) -> None: + hook: Final = _RecordingPreCallDeploymentHook() + monkeypatch.setattr(litellm, "callbacks", [hook]) + + await async_pre_call_deployment_hook({"model": "gpt-5.4-mini"}, "acompact_responses") + + assert hook.seen_call_types == (CallTypes.acompact_responses,) diff --git a/tests/unit/types/test_utils.py b/tests/unit/types/test_utils.py index 32479c62372..bdbcbb698d4 100644 --- a/tests/unit/types/test_utils.py +++ b/tests/unit/types/test_utils.py @@ -1,7 +1,12 @@ import pytest from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_excluded_services_value -from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VALUES, CaptureMessageContent, captures_span_content +from litellm.types.utils import ( + CAPTURE_MESSAGE_CONTENT_VALUES, + CallTypes, + CaptureMessageContent, + captures_span_content, +) @pytest.mark.parametrize("value", ["", " ", "auth", "none,redis", "redis,", ",postgres"]) @@ -45,3 +50,8 @@ def test_capture_message_content_constructor_accepts_any_string() -> None: ) def test_captures_span_content(mode: str | None, captures: bool) -> None: assert captures_span_content(mode) is captures + + +@pytest.mark.parametrize("call_type", ["compact_responses", "acompact_responses"]) +def test_responses_compact_call_types_round_trip_their_public_names(call_type: str) -> None: + assert CallTypes(call_type).value == call_type diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index dd428b286bd..e9bbeb94355 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1037,11 +1037,6 @@ "count": 1 } }, - "src/app/(dashboard)/skills/_components/ClaudeCodePluginsPanel.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/app/(dashboard)/skills/_components/add_plugin_form.tsx": { "local/filename-pascal-case": { "count": 1 @@ -1745,7 +1740,7 @@ "count": 5 }, "no-restricted-syntax": { - "count": 133 + "count": 132 }, "prefer-const": { "count": 26 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts index 652b2c0b480..e44fca35d88 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts @@ -15,11 +15,12 @@ import { export interface FieldConfig { name: string; - label: string; + labelKey: ParseKeys; type: "text" | "textarea" | "url" | "switch" | "list" | "select"; required?: boolean; - tooltip?: string; + tooltipKey?: ParseKeys; placeholder?: string; + placeholderKey?: ParseKeys; defaultValue?: any; rows?: number; validation?: any[]; @@ -50,42 +51,41 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "name", - label: "Display Name", + labelKey: "agents.formConfig.fields.name.label", type: "text", required: true, - placeholder: "e.g., Customer Support Agent", + placeholderKey: "agents.formConfig.fields.name.placeholder", }, { name: "description", - label: "Description", + labelKey: "agents.formConfig.fields.description.label", type: "textarea", required: false, - placeholder: "Describe what this agent does...", + placeholderKey: "agents.formConfig.fields.description.placeholder", rows: 3, }, { name: "url", - label: "URL", + labelKey: "agents.formConfig.fields.url.label", type: "url", required: false, placeholder: "http://localhost:9999/", - tooltip: "Base URL where the agent is hosted (optional)", + tooltipKey: "agents.formConfig.fields.url.tooltip", }, { name: "version", - label: "Version", + labelKey: "agents.formConfig.fields.version.label", type: "text", placeholder: "1.0.0", defaultValue: "1.0.0", }, { name: "protocolVersion", - label: "Protocol Version", + labelKey: "agents.formConfig.fields.protocolVersion.label", type: "select", options: ["1.0", "0.3"], defaultValue: "1.0", - tooltip: - "The A2A protocol version LiteLLM serves to clients for this agent. LiteLLM converts the upstream agent's responses to this version, so clients always see the version you pick here regardless of the original agent's version.", + tooltipKey: "agents.formConfig.fields.protocolVersion.tooltip", helpTextKey: "agents.formConfig.basic.protocolVersionHelp", }, ], @@ -96,7 +96,7 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "skills", - label: "Skills", + labelKey: "agents.formConfig.fields.skills.label", type: "list", defaultValue: [], }, @@ -108,18 +108,18 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "streaming", - label: "Streaming", + labelKey: "agents.formConfig.fields.streaming.label", type: "switch", defaultValue: false, }, { name: "pushNotifications", - label: "Push Notifications", + labelKey: "agents.formConfig.fields.pushNotifications.label", type: "switch", }, { name: "stateTransitionHistory", - label: "State Transition History", + labelKey: "agents.formConfig.fields.stateTransitionHistory.label", type: "switch", }, ], @@ -130,19 +130,19 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "iconUrl", - label: "Icon URL", + labelKey: "agents.formConfig.fields.iconUrl.label", type: "url", placeholder: "https://example.com/icon.png", }, { name: "documentationUrl", - label: "Documentation URL", + labelKey: "agents.formConfig.fields.documentationUrl.label", type: "url", placeholder: "https://docs.example.com", }, { name: "supportsAuthenticatedExtendedCard", - label: "Supports Authenticated Extended Card", + labelKey: "agents.formConfig.fields.supportsAuthenticatedExtendedCard.label", type: "switch", }, ], @@ -153,12 +153,12 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "model", - label: "Model (Optional)", + labelKey: "agents.formConfig.fields.model.label", type: "text", }, { name: "make_public", - label: "Make Public", + labelKey: "agents.formConfig.fields.makePublic.label", type: "switch", }, ], @@ -169,24 +169,24 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "cost_per_query", - label: "Cost Per Query ($)", + labelKey: "agents.formConfig.fields.costPerQuery.label", type: "text", placeholder: "0.0", - tooltip: "Fixed cost per query", + tooltipKey: "agents.formConfig.fields.costPerQuery.tooltip", }, { name: "input_cost_per_token", - label: "Input Cost Per Token ($)", + labelKey: "agents.formConfig.fields.inputCostPerToken.label", type: "text", placeholder: "0.000001", - tooltip: "Cost per input token", + tooltipKey: "agents.formConfig.fields.inputCostPerToken.tooltip", }, { name: "output_cost_per_token", - label: "Output Cost Per Token ($)", + labelKey: "agents.formConfig.fields.outputCostPerToken.label", type: "text", placeholder: "0.000002", - tooltip: "Cost per output token", + tooltipKey: "agents.formConfig.fields.outputCostPerToken.tooltip", }, ], }, @@ -196,10 +196,10 @@ export const AGENT_FORM_CONFIG: { fields: [ { name: "enable_tracing", - label: "Enable Tracing", + labelKey: "agents.formConfig.fields.enableTracing.label", type: "switch", defaultValue: false, - tooltip: "Enable request tracing for this agent", + tooltipKey: "agents.formConfig.fields.enableTracing.tooltip", }, ], }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx index 41e25d14054..2f69336e6c6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx @@ -256,23 +256,24 @@ const AgentFormFields: React.FC = ({ panels, showAgentName {({ value, onChange, ref, ...control }) => { const text = typeof value === "string" ? value : ""; + const placeholder = field.placeholderKey ? t(field.placeholderKey) : field.placeholder; if (field.type === "textarea") { return (