diff --git a/.github/actions/cache-cargo-build/action.yml b/.github/actions/cache-cargo-build/action.yml index c3b8ce22c68..222fad637fb 100644 --- a/.github/actions/cache-cargo-build/action.yml +++ b/.github/actions/cache-cargo-build/action.yml @@ -15,6 +15,12 @@ description: >- cache the same directory for different workloads, and a shared key would let whichever ran first deny the others a save. +inputs: + profile: + description: "Cargo profile the build uses (dev or release)" + required: false + default: "dev" + runs: using: composite steps: @@ -25,6 +31,6 @@ runs: ~/.cargo/registry ~/.cargo/git litellm-rust/target - key: ${{ runner.os }}-maturin-dev-${{ hashFiles('litellm-rust/Cargo.lock') }} + key: ${{ runner.os }}-maturin-${{ inputs.profile }}-${{ hashFiles('litellm-rust/Cargo.lock') }} restore-keys: | - ${{ runner.os }}-maturin-dev- + ${{ runner.os }}-maturin-${{ inputs.profile }}- diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index dbc4ae8f5c2..2386b184e54 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -9,6 +9,7 @@ UNSUPPORTED: Final = re.compile( r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$" r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$" r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$" + r"|^tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e\.py$" ) HARNESS: Final = re.compile( r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$" diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index ea6d2401084..f2b82f86b47 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -134,7 +134,16 @@ def main( uncompressed_wheel_size: Final = sum(member.file_size for member in wheel_members) native_path: Final = wheel.parent / "native" / Path(native_member.filename).name native_path.parent.mkdir(parents=True, exist_ok=True) - native_path.write_bytes(archive.read(native_member)) + native_bytes: Final = archive.read(native_member) + native_path.write_bytes(native_bytes) + duplicated_vocabularies: Final = tuple( + member.filename + for member in wheel_members + if member.filename.startswith("litellm/litellm_core_utils/tokenizers/") + and re.fullmatch(r"[0-9a-f]{40}", PurePosixPath(member.filename).name) + and member.file_size > 0 + and archive.read(member) in native_bytes + ) wheel_metadata_tags_match: Final = ( len(wheel_metadata_tags) == len(expanded_filename_tags) @@ -205,7 +214,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 40_000_000 + native_size_limit: Final = 35_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), @@ -223,6 +232,7 @@ def main( ("Native module loads", native_module_loads), ("Production module omits the panic test hook", panic_test_hook_absent), (f"Native extension does not exceed {native_size_limit / 1_000_000:.0f} MB", native_size_within_limit), + ("Tokenizer vocabularies are not duplicated in the native extension", not duplicated_vocabularies), ("Wheel contents are valid", not unexpected_members), ) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index fd7513a3937..ec7e211faa1 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -13,6 +13,7 @@ on: - ".github/workflows/codspeed.yml" - ".github/actions/setup-uv-with-retries/**" - ".github/actions/cache-cargo-build/**" + - ".github/scripts/uv_sync_with_retries.sh" pull_request: branches: - main @@ -25,6 +26,7 @@ on: - ".github/workflows/codspeed.yml" - ".github/actions/setup-uv-with-retries/**" - ".github/actions/cache-cargo-build/**" + - ".github/scripts/uv_sync_with_retries.sh" # Allow CodSpeed to trigger backtest performance analysis # in order to generate initial data workflow_dispatch: @@ -59,19 +61,27 @@ jobs: - name: Cache the Rust build uses: ./.github/actions/cache-cargo-build + with: + profile: release # Build the wheel and resolve every dependency outside the CodSpeed # runner: the same maturin build took 42 minutes inside `codspeed run` # versus under 3 minutes as a plain step (LIT-6183) - - name: Build environment + - name: Build the release wheel + run: uv build --wheel --out-dir dist + + - name: Install the wheel into the benchmark environment + run: | + UV_PROJECT_ENVIRONMENT="${RUNNER_TEMP}/benchmark-venv" .github/scripts/uv_sync_with_retries.sh --frozen --no-default-groups --group benchmarks --no-install-project --python 3.12 + uv pip install --python "${RUNNER_TEMP}/benchmark-venv/bin/python" --no-deps dist/*.whl + + - name: Collect benchmarks + env: + PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1" + LITELLM_REQUIRE_INSTALLED_WHEEL: "1" run: > - env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 - uv run --frozen --no-default-groups - --with pytest==8.3.5 - --with pytest-codspeed==4.3.0 - --with "mcp>=2.2.0,<3.0" - --with "a2a-sdk>=1.1.0,<2.0" - pytest + "${RUNNER_TEMP}/benchmark-venv/bin/python" -I -m pytest + --import-mode=importlib -p pytest_codspeed.plugin tests/benchmarks/ --codspeed @@ -82,13 +92,9 @@ jobs: with: mode: simulation run: > - env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 - uv run --frozen --no-default-groups - --with pytest==8.3.5 - --with pytest-codspeed==4.3.0 - --with "mcp>=2.2.0,<3.0" - --with "a2a-sdk>=1.1.0,<2.0" - pytest + env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 LITELLM_REQUIRE_INSTALLED_WHEEL=1 + "${RUNNER_TEMP}/benchmark-venv/bin/python" -I -m pytest + --import-mode=importlib -p pytest_codspeed.plugin tests/benchmarks/ --codspeed diff --git a/backend/Dockerfile b/backend/Dockerfile index 622fedcd70d..57e0a43a98d 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -61,6 +61,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra saml \ --python python3.13 +RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/ + RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ prisma generate --schema=./schema.prisma diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index 50aa40ba220..5842cf6f1ac 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -1,8 +1,11 @@ """Guard the cost map on pull requests. -Every pull request gets the file checks: the three cost map files parse, the backup copy matches the root file, -and the JSON schema is in sync and validates the map. Pull requests from the cost map sync bot (branches named -litellm_cost_map_sync_*) additionally may only touch those three files and may only add or update models. +Every pull request whose diff against its merge base touches one of the three cost map files gets the file +checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the +map. A pull request that leaves all three untouched skips them, since merging it keeps the base branch's copies +and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot +(branches named litellm_cost_map_sync_*) always get the file checks and additionally may only touch those three +files and may only add or update models. """ from __future__ import annotations @@ -108,20 +111,37 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str ) +def touches_cost_map(changed_files: Sequence[str]) -> bool: + return any(path in GUARDED_PATHS for path in changed_files) + + +def contract_for(bot: bool, changed_files: Sequence[str]) -> str: + if bot: + return "bot contract enforced" + return "human PR, file checks only" if touches_cost_map(changed_files) else "human PR, cost map untouched" + + def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]: + if not bot and not touches_cost_map(changed_files): + return () head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH) if isinstance(head_map, str): return (head_map,) return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ())) -def _git(*args: str) -> str: +def _git(*args: str) -> str | None: result: Final = subprocess.run(("git", *args), check=False, capture_output=True, text=True) - return result.stdout if result.returncode == 0 else "" + return result.stdout if result.returncode == 0 else None def snapshot(revision: str) -> Snapshot: - return Snapshot(*(_git("show", f"{revision}:{path}") for path in GUARDED_PATHS)) + return Snapshot(*(_git("show", f"{revision}:{path}") or "" for path in GUARDED_PATHS)) + + +def changed_files(base: str, head: str) -> tuple[str, ...] | None: + diff: Final = _git("diff", "--name-only", "--no-renames", base, head) + return None if diff is None else tuple(diff.splitlines()) def main(argv: Sequence[str]) -> int: @@ -131,9 +151,12 @@ def main(argv: Sequence[str]) -> int: parser.add_argument("--head-ref", required=True, help="head branch name of the pull request") args: Final = parser.parse_args(argv) bot: Final = args.head_ref.startswith(BOT_BRANCH_PREFIX) - changed_files: Final = tuple(_git("diff", "--name-only", args.base, args.head).splitlines()) - failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed_files, bot) - contract: Final = "bot contract enforced" if bot else "human PR, file checks only" + changed: Final = changed_files(args.base, args.head) + if changed is None: + print(f"cost map guard failed: git diff {args.base} {args.head} failed, so the changed files are unknown") + return 1 + failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed, bot) + contract: Final = contract_for(bot, changed) if failures: print(f"cost map guard failed ({contract}):") print("\n".join(f"- {failure}" for failure in failures)) diff --git a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json index d8cb122417a..af88708166f 100644 --- a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json +++ b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json @@ -6267,6 +6267,63 @@ ], "title": "Spend update queue sizes (litellm__size)", "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "Requests that carried usage but were logged at $0 on a model whose pricing entry has a non-zero rate, by requested model and reason", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "drawStyle": "line", + "fillOpacity": 10, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": false + }, + "unit": "short" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 430 + }, + "id": 110, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "sum(rate(litellm_zero_cost_requests_total[$__rate_interval])) by (requested_model, reason)", + "legendFormat": "{{requested_model}} / {{reason}}", + "range": true, + "refId": "A" + } + ], + "title": "litellm_zero_cost_requests rate", + "type": "timeseries" } ], "preload": false, diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index f7557983b91..fe58e2dd58c 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -96,16 +96,19 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( "/assemblyai/", "/eu.assemblyai/", "/deepgram/", + "/fal_ai/", "/langfuse/", "/vllm/", "/mistral/", "/typesafe/", + "/openrouter/", "/nvidia_nim/", "/groq/", "/voyage/", "/cursor/", "/milvus/", "/openai_passthrough/", + "/tinyfish/", # Dynamic provider / toolset passthrough (path templates) "/{provider}/", "/toolset/", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql new file mode 100644 index 00000000000..d594b0056df --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index f802a141d41..c55456b2a40 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -73,6 +73,7 @@ model LiteLLM_AgentsTable { static_headers Json? @default("{}") extra_headers String[] @default([]) agent_access_groups String[] @default([]) + access_group_ids String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c02f460f22e..37384cbfa53 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -650,6 +650,49 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "futures-util", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "sync_wrapper", + "tower", + "tower-layer", + "tower-service", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", +] + [[package]] name = "azure_core" version = "1.1.0" @@ -1439,7 +1482,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2241,7 +2284,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.5", "tokio", "tower-service", "tracing", @@ -2729,6 +2772,26 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-qdrant-semantic" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-cache", + "litellm-cache-response", + "qdrant-client", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tokio-stream", + "tonic", + "tonic-prost", + "uuid", +] + [[package]] name = "litellm-cache-redis" version = "0.1.0" @@ -2741,6 +2804,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-redis", + "litellm-cache-response", + "r2d2", + "redis", + "redis-test", + "serde_json", + "sha2 0.10.9", + "tokio", +] + [[package]] name = "litellm-cache-response" version = "0.1.0" @@ -2816,6 +2894,7 @@ dependencies = [ "litellm-host", "litellm-http", "litellm-llms", + "litellm-secrets", "litellm-types", "mime_guess", "moka", @@ -2928,6 +3007,7 @@ dependencies = [ "litellm-framing", "litellm-host", "litellm-http", + "litellm-secrets", "litellm-types", "reqwest 0.12.28", "rstest", @@ -2946,6 +3026,7 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "aws-sdk-secretsmanager", "bytes", "criterion", "futures-util", @@ -2957,7 +3038,9 @@ dependencies = [ "litellm-cache-disk", "litellm-cache-gcs", "litellm-cache-memory", + "litellm-cache-qdrant-semantic", "litellm-cache-redis", + "litellm-cache-redis-semantic", "litellm-cache-response", "litellm-cache-s3", "litellm-cache-valkey-semantic", @@ -2967,11 +3050,16 @@ dependencies = [ "litellm-host-python", "litellm-http", "litellm-llms", + "litellm-secrets", + "litellm-secrets-aws", + "litellm-secrets-types", "litellm-token-counter", "litellm-types", "pyo3", "pyo3-async-runtimes", + "qdrant-client", "redis", + "reqwest 0.12.28", "rstest", "serde", "serde_json", @@ -2979,6 +3067,8 @@ dependencies = [ "sha2 0.10.9", "tokio", "tokio-tungstenite", + "url", + "wiremock", ] [[package]] @@ -3163,6 +3253,7 @@ dependencies = [ name = "litellm-token-counter-huggingface" version = "0.1.0" dependencies = [ + "serde_json", "thiserror 2.0.19", "tokenizers", ] @@ -3171,6 +3262,9 @@ dependencies = [ name = "litellm-token-counter-tiktoken" version = "0.1.0" dependencies = [ + "base64 0.22.1", + "once_cell", + "rustc-hash", "thiserror 2.0.19", "tiktoken-rs", ] @@ -3235,6 +3329,12 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c" +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + [[package]] name = "md-5" version = "0.11.0" @@ -3858,6 +3958,27 @@ dependencies = [ "serde", ] +[[package]] +name = "qdrant-client" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dddc19df129bad7346ebd027288621ab1ac7e52678371f906b9a8622d7aaf87e" +dependencies = [ + "anyhow", + "derive_builder", + "futures", + "parking_lot", + "prost", + "prost-types", + "semver", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tonic", + "tonic-prost", +] + [[package]] name = "quick-error" version = "1.2.3" @@ -3887,7 +4008,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.42", - "socket2 0.5.10", + "socket2 0.6.5", "thiserror 2.0.19", "tokio", "tracing", @@ -3926,9 +4047,9 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.5", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4414,7 +4535,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4485,7 +4606,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5057,10 +5178,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5349,8 +5470,12 @@ version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ + "async-trait", + "axum", "base64 0.22.1", "bytes", + "flate2", + "h2 0.4.15", "http 1.4.2", "http-body 1.1.0", "http-body-util", @@ -5360,6 +5485,7 @@ dependencies = [ "percent-encoding", "pin-project", "rustls-native-certs", + "socket2 0.6.5", "sync_wrapper", "tokio", "tokio-rustls 0.26.4", @@ -5933,7 +6059,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index be86240ac46..813d0713128 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -36,7 +36,9 @@ litellm-cache-redis = { path = "crates/cache-redis" } litellm-cache-s3 = { path = "crates/cache-s3" } litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-disk = { path = "crates/cache-disk" } +litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" } litellm-cache-response = { path = "crates/cache-response" } +litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" } @@ -54,6 +56,8 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] } +qdrant-client = { version = "1.19.0", default-features = false } +uuid = { version = "1", features = ["v4"] } rstest = "0.26.1" rstest_reuse = "0.7.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } @@ -85,7 +89,7 @@ veil = "0.3.0" [profile.release] opt-level = 3 -lto = "thin" +lto = "fat" codegen-units = 1 panic = "unwind" debug = false diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index bbcb0f016c8..cb9195ffeb6 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -621,6 +621,26 @@ mod tests { None } + #[test] + fn secret_names_cover_environment_reads() { + let seen = std::sync::Arc::new(std::sync::Mutex::new( + std::collections::BTreeSet::::new(), + )); + let recorded = seen.clone(); + let env = |name: &str| { + recorded.lock().unwrap().insert(name.to_string()); + None + }; + resolve_aws_region(None, &Map::new(), &env); + aws_auth_config(&Map::new(), &env); + assert!( + seen.lock() + .unwrap() + .iter() + .all(|name| crate::constants::SECRET_NAMES.contains(&name.as_str())) + ); + } + #[test] fn a_region_comes_from_the_call_then_the_model_then_the_environment() { let params = Map::from_iter([("aws_region_name".to_string(), Value::from("eu-west-1"))]); diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index 9e7c6bfab43..26df4f2a350 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -14,6 +14,19 @@ pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT"; pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID"; pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK"; +pub const SECRET_NAMES: &[&str] = &[ + AWS_ACCESS_KEY_ID, + AWS_SECRET_ACCESS_KEY, + AWS_SESSION_TOKEN, + AWS_REGION_NAME, + AWS_REGION, + AWS_SESSION_NAME, + AWS_PROFILE_NAME, + AWS_ROLE_NAME, + AWS_WEB_IDENTITY_TOKEN, + AWS_STS_ENDPOINT, + AWS_EXTERNAL_ID, +]; /// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors /// Python's `_filter_headers_for_aws_signature` allowlist. diff --git a/litellm-rust/crates/auth-azure/src/lib.rs b/litellm-rust/crates/auth-azure/src/lib.rs index 5c7c654b69d..9ed505e4160 100644 --- a/litellm-rust/crates/auth-azure/src/lib.rs +++ b/litellm-rust/crates/auth-azure/src/lib.rs @@ -3,5 +3,5 @@ mod native; mod resolve; mod types; -pub use resolve::AzureAuthService; +pub use resolve::{AzureAuthService, SECRET_NAMES}; pub use types::{AzureAuthInputs, ConfigValue}; diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 4d564b6e68a..9a7afe645db 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -19,6 +19,17 @@ const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST"; const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL"; const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; +pub const SECRET_NAMES: &[&str] = &[ + AZURE_AD_TOKEN_ENV, + AZURE_TENANT_ID_ENV, + AZURE_CLIENT_ID_ENV, + AZURE_CLIENT_SECRET_ENV, + AZURE_SCOPE_ENV, + AZURE_AUTHORITY_HOST_ENV, + AZURE_CREDENTIAL_ENV, + AZURE_FEDERATED_TOKEN_FILE_ENV, +]; + #[derive(Clone, Debug)] pub(crate) enum AzureCredentialPlan { Supplied(Sourced), @@ -440,13 +451,14 @@ fn non_empty_reference(value: &str, kind: &str) -> Result { #[cfg(test)] mod tests { + use std::collections::BTreeSet; use std::future::Future; use std::sync::{Arc, Mutex}; use serde_json::json; use super::{ - AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, oidc_reference, + AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, SECRET_NAMES, oidc_reference, resolve_reference, select_auth_plan, }; use crate::native::ValidatedAzureRequest; @@ -517,6 +529,24 @@ mod tests { assert!(matches!(plan, AzureCredentialPlan::Native(_))); } + #[test] + fn secret_names_cover_environment_reads() { + let seen = std::sync::Arc::new(std::sync::Mutex::new(BTreeSet::::new())); + let recorded = seen.clone(); + let inputs = AzureAuthInputs::default(); + select_auth_plan(&inputs, &|name| { + recorded.lock().unwrap().insert(name.to_string()); + None + }) + .unwrap(); + assert!( + seen.lock() + .unwrap() + .iter() + .all(|name| SECRET_NAMES.contains(&name.as_str())) + ); + } + #[test] fn supplied_token_does_not_require_refresh() { let params = json!({"azure_ad_token": "token"}); diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 534d85acdb0..682f1af5fe1 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -23,6 +23,16 @@ const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; +pub const SECRET_NAMES: &[&str] = &[ + VERTEX_AI_API_KEY_ENV, + VERTEXAI_API_KEY_ENV, + VERTEXAI_CREDENTIALS_ENV, + GOOGLE_APPLICATION_CREDENTIALS_ENV, + VERTEXAI_PROJECT_ENV, + VERTEXAI_LOCATION_ENV, + VERTEX_LOCATION_ENV, +]; + #[derive(Clone, Debug, Default)] pub struct VertexConfig { credentials: Option>, @@ -406,6 +416,7 @@ fn auth_acquisition_error(error: gcp_auth::Error) -> Error { #[cfg(test)] mod tests { + use std::collections::BTreeSet; use std::sync::atomic::{AtomicUsize, Ordering}; use serde_json::json; @@ -476,6 +487,27 @@ mod tests { ); } + #[tokio::test] + async fn secret_names_cover_environment_reads() { + let seen = Arc::new(std::sync::Mutex::new(BTreeSet::::new())); + let recorded = seen.clone(); + let env = |name: &str| { + recorded.lock().unwrap().insert(name.to_string()); + None + }; + let auth = auth(Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))); + auth.validate_environment(Vec::new(), None, &VertexConfig::default(), &env) + .await + .unwrap(); + get_vertex_ai_location(&VertexConfig::default(), &env); + assert!( + seen.lock() + .unwrap() + .iter() + .all(|name| SECRET_NAMES.contains(&name.as_str())) + ); + } + #[test] fn empty_primary_values_fall_back_to_python_aliases() { let config = config(json!({ diff --git a/litellm-rust/crates/auth-types/src/secret.rs b/litellm-rust/crates/auth-types/src/secret.rs index a07fe3eaad9..7e6789deef7 100644 --- a/litellm-rust/crates/auth-types/src/secret.rs +++ b/litellm-rust/crates/auth-types/src/secret.rs @@ -1,4 +1,5 @@ use serde::Deserialize; +use std::hash::{Hash, Hasher}; use veil::Redact; #[derive(Redact, Clone, Deserialize)] @@ -23,6 +24,12 @@ impl PartialEq for SecretValue { impl Eq for SecretValue {} +impl Hash for SecretValue { + fn hash(&self, state: &mut H) { + self.0.hash(state); + } +} + #[cfg(test)] mod tests { use super::SecretValue; diff --git a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml new file mode 100644 index 00000000000..09d6a9637f3 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "litellm-cache-qdrant-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-cache.workspace = true +qdrant-client = { workspace = true, features = ["serde"] } +reqwest.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio.workspace = true +uuid.workspace = true + +[dev-dependencies] +litellm-cache-response.workspace = true +rstest.workspace = true +tonic = "0.14" +tonic-prost = "0.14" +tokio-stream = "0.1" diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs new file mode 100644 index 00000000000..47b898d6f4e --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs @@ -0,0 +1,75 @@ +use std::time::Duration; + +use litellm_cache::Error; +use reqwest::Client; +use serde_json::Value; + +use crate::Embedder; + +pub struct OpenAiEmbedder { + client: Client, + api_base: String, + api_key: String, + model: String, + timeout: Option, +} + +pub struct OpenAiEmbedderConfig { + pub api_base: String, + pub api_key: String, + pub model: String, + pub timeout: Option, +} + +impl OpenAiEmbedder { + pub fn new(client: Client, config: OpenAiEmbedderConfig) -> Self { + Self { + client, + api_base: config.api_base.trim_end_matches('/').to_owned(), + api_key: config.api_key, + model: config.model, + timeout: config.timeout, + } + } +} + +impl Embedder for OpenAiEmbedder { + fn model(&self) -> &str { + &self.model + } + + async fn embed(&self, input: &str) -> Result, Error> { + let request = self + .client + .post(format!("{}/embeddings", self.api_base)) + .bearer_auth(&self.api_key) + .json(&serde_json::json!({ + "model": self.model, + "input": input, + "encoding_format": "float", + })); + let response = if let Some(timeout) = self.timeout { + request.timeout(timeout) + } else { + request + } + .send() + .await + .map_err(|_| Error::Unavailable)? + .error_for_status() + .map_err(|_| Error::Unavailable)?; + let body: Value = response.json().await.map_err(|_| Error::Unavailable)?; + body.get("data") + .and_then(Value::as_array) + .and_then(|data| data.first()) + .and_then(|item| item.get("embedding")) + .and_then(Value::as_array) + .and_then(|embedding| { + embedding + .iter() + .map(|value| value.as_f64().map(|value| value as f32)) + .collect::>>() + }) + .ok_or(Error::Unavailable) + } +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs b/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs new file mode 100644 index 00000000000..0f346a9155b --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs @@ -0,0 +1,7 @@ +mod embedder; +mod prompt; +mod semantic; + +pub use embedder::{OpenAiEmbedder, OpenAiEmbedderConfig}; +pub use prompt::prompt_from_messages; +pub use semantic::{Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization}; diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs b/litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs new file mode 100644 index 00000000000..ef1a2306658 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs @@ -0,0 +1,59 @@ +use serde_json::Value; + +fn search_results_text(search_results: Option<&Value>) -> String { + let Some(Value::Array(results)) = search_results else { + return String::new(); + }; + results + .iter() + .filter_map(Value::as_object) + .flat_map(|result| { + let source = result + .get("source") + .and_then(Value::as_str) + .map(str::to_owned); + let title = result + .get("title") + .and_then(Value::as_str) + .map(str::to_owned); + let content = result + .get("content") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_object) + .filter_map(|block| block.get("text").and_then(Value::as_str).map(str::to_owned)); + let citations = result + .get("citations") + .filter(|value| !value.is_null()) + .map(|value| serde_json::to_string(value).unwrap_or_default()); + source + .into_iter() + .chain(title) + .chain(content) + .chain(citations) + }) + .collect() +} + +pub fn prompt_from_messages(messages: &[Value]) -> String { + messages + .iter() + .filter_map(Value::as_object) + .map(|message| { + let content = match message.get("content") { + Some(Value::String(content)) => content.clone(), + Some(Value::Array(parts)) => parts + .iter() + .filter_map(Value::as_object) + .filter_map(|part| part.get("text").and_then(Value::as_str)) + .collect(), + _ => String::new(), + }; + format!( + "{content}{}", + search_results_text(message.get("search_results")) + ) + }) + .collect() +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs new file mode 100644 index 00000000000..fb165ed5a8e --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs @@ -0,0 +1,262 @@ +use std::future::Future; + +use futures_util::future::try_join_all; +use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext}; +use qdrant_client::{ + Payload, Qdrant, + qdrant::{ + BinaryQuantizationBuilder, CompressionRatio, Condition, CreateCollectionBuilder, + CreateFieldIndexCollectionBuilder, Distance, FieldType, Filter, PointStruct, + ProductQuantizationBuilder, QuantizationSearchParamsBuilder, ScalarQuantizationBuilder, + SearchParamsBuilder, SearchPointsBuilder, UpsertPointsBuilder, VectorParamsBuilder, + }, +}; +use serde_json::{Map, Value, json}; +use uuid::Uuid; + +use crate::prompt_from_messages; + +pub trait Embedder: Send + Sync + 'static { + fn model(&self) -> &str; + fn embed(&self, input: &str) -> impl Future, Error>> + Send; +} + +#[derive(Clone, Debug, PartialEq)] +pub enum Quantization { + Binary, + Scalar, + Product, +} + +pub struct QdrantSemanticConfig { + pub collection_name: String, + pub similarity_threshold: f64, + pub vector_size: u64, + pub quantization: Quantization, +} + +pub struct QdrantSemanticCache { + client: Qdrant, + embedder: E, + codec: C, + config: QdrantSemanticConfig, + runtime: tokio::runtime::Handle, +} + +impl QdrantSemanticCache { + pub async fn connect( + client: Qdrant, + embedder: E, + codec: C, + config: QdrantSemanticConfig, + runtime: tokio::runtime::Handle, + ) -> Result { + let exists = client + .collection_exists(config.collection_name.clone()) + .await + .map_err(|_| Error::Unavailable)?; + if !exists { + client + .create_collection( + CreateCollectionBuilder::new(config.collection_name.clone()) + .vectors_config(VectorParamsBuilder::new( + config.vector_size, + Distance::Cosine, + )) + .quantization_config(quantization(&config.quantization)), + ) + .await + .map_err(|_| Error::Unavailable)?; + } + let _ = client + .create_field_index(CreateFieldIndexCollectionBuilder::new( + config.collection_name.clone(), + "litellm_cache_key".to_owned(), + FieldType::Keyword, + )) + .await; + Ok(Self { + client, + embedder, + codec, + config, + runtime, + }) + } + + pub fn collection_name(&self) -> &str { + &self.config.collection_name + } + + pub fn similarity_threshold(&self) -> f64 { + self.config.similarity_threshold + } + + pub fn vector_size(&self) -> u64 { + self.config.vector_size + } + + pub fn embedder(&self) -> &E { + &self.embedder + } + + fn prompt(context: &SemanticCacheContext) -> Result { + let Some(messages) = context.messages.as_ref().and_then(Value::as_array) else { + return Err(Error::MissingPrompt); + }; + if messages.is_empty() { + return Err(Error::MissingPrompt); + } + Ok(prompt_from_messages(messages)) + } + + async fn set( + &self, + key: &str, + value: C::Value, + context: &SemanticCacheContext, + ) -> Result<(), Error> { + let prompt = Self::prompt(context)?; + let vector = self.embedder.embed(&prompt).await?; + let response = + String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?; + let payload = Payload::try_from(json!({ + "litellm_cache_key": key, + "text": prompt, + "response": response, + })) + .map_err(|_| Error::InvalidEntry)?; + self.client + .upsert_points( + UpsertPointsBuilder::new( + self.collection_name(), + vec![PointStruct::new( + Uuid::new_v4().to_string(), + vector, + payload, + )], + ) + .wait(true), + ) + .await + .map_err(|_| Error::Unavailable)?; + Ok(()) + } + + async fn get( + &self, + key: &str, + context: &SemanticCacheContext, + ) -> Result, Error> { + let prompt = Self::prompt(context)?; + let vector = self.embedder.embed(&prompt).await?; + let result = self + .client + .search_points( + SearchPointsBuilder::new(self.collection_name(), vector, 1) + .with_payload(true) + .filter(Filter::must([Condition::matches( + "litellm_cache_key", + key.to_owned(), + )])) + .params( + SearchParamsBuilder::default().quantization( + QuantizationSearchParamsBuilder::default() + .ignore(false) + .rescore(true) + .oversampling(3.0), + ), + ), + ) + .await + .map_err(|_| Error::Unavailable)?; + let Some(point) = result.result.into_iter().next() else { + return Ok(None); + }; + let payload: Map = Payload::from(point.payload).into(); + if payload.get("litellm_cache_key").and_then(Value::as_str) != Some(key) { + return Ok(None); + } + if f64::from(point.score) < self.config.similarity_threshold { + return Ok(None); + } + let response = payload + .get("response") + .and_then(Value::as_str) + .ok_or(Error::InvalidEntry)?; + self.codec.decode(response.as_bytes()).map(Some) + } +} + +fn quantization(value: &Quantization) -> qdrant_client::qdrant::quantization_config::Quantization { + match value { + Quantization::Binary => BinaryQuantizationBuilder::new(false).into(), + Quantization::Scalar => ScalarQuantizationBuilder::default() + .quantile(0.99) + .always_ram(false) + .into(), + Quantization::Product => ProductQuantizationBuilder::new(CompressionRatio::X16.into()) + .always_ram(false) + .into(), + } +} + +impl BaseCache for QdrantSemanticCache { + type Value = C::Value; + type Context = SemanticCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + None + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + self.runtime.block_on(self.set(key, value, context)) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + self.runtime.block_on(self.get(key, context)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + self.set(key, value, &context).await + } + + async fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + self.get(key, context).await + } + + async fn async_set_cache_pipeline( + &self, + entries: Vec<(String, Self::Value)>, + context: Self::Context, + ) -> Result<(), Error> { + try_join_all(entries.into_iter().map(|(key, value)| { + let context = context.clone(); + async move { self.async_set_cache(&key, value, context).await } + })) + .await + .map(|_| ()) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs new file mode 100644 index 00000000000..6b09448fde8 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -0,0 +1,166 @@ +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::Error; +use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, OpenAiEmbedderConfig}; +use serde_json::Value; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, +}; + +struct TestHttpServer { + address: std::net::SocketAddr, + request: Arc>>>, + task: tokio::task::JoinHandle<()>, +} + +impl TestHttpServer { + async fn response(status: &str, body: &str) -> Self { + Self::response_after(status, body, Duration::ZERO).await + } + + async fn response_after(status: &str, body: &str, delay: Duration) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let request = Arc::new(Mutex::new(None)); + let captured = request.clone(); + let status = status.to_owned(); + let body = body.to_owned(); + let task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let request_bytes = read_request(&mut stream).await; + *captured.lock().unwrap() = Some(request_bytes); + tokio::time::sleep(delay).await; + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + stream.write_all(response.as_bytes()).await.unwrap(); + }); + Self { + address, + request, + task, + } + } + + fn base_url(&self) -> String { + format!("http://{}", self.address) + } +} + +impl Drop for TestHttpServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +async fn read_request(stream: &mut tokio::net::TcpStream) -> Vec { + let mut bytes = Vec::new(); + let header_end = loop { + let mut chunk = [0_u8; 1024]; + let count = stream.read(&mut chunk).await.unwrap(); + assert_ne!(count, 0); + bytes.extend_from_slice(&chunk[..count]); + if let Some(end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") { + break end + 4; + } + }; + let headers = String::from_utf8_lossy(&bytes[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + line.split_once(':') + .filter(|(name, _)| name.eq_ignore_ascii_case("content-length")) + .map(|(_, value)| value.trim()) + }) + .unwrap() + .parse::() + .unwrap(); + while bytes.len() < header_end + content_length { + let mut chunk = [0_u8; 1024]; + let count = stream.read(&mut chunk).await.unwrap(); + assert_ne!(count, 0); + bytes.extend_from_slice(&chunk[..count]); + } + bytes +} + +fn config(base: String, timeout: Option) -> OpenAiEmbedderConfig { + OpenAiEmbedderConfig { + api_base: base, + api_key: "test-key".to_owned(), + model: "test-model".to_owned(), + timeout, + } +} + +#[tokio::test] +async fn posts_embeddings_request_and_parses_vector() { + let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; + let embedder = OpenAiEmbedder::new( + reqwest::Client::new(), + config( + format!("{}/", server.base_url()), + Some(Duration::from_secs(1)), + ), + ); + assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]); + let request = server.request.lock().unwrap().clone().unwrap(); + let request_text = String::from_utf8(request).unwrap(); + assert!(request_text.starts_with("POST /embeddings HTTP/1.1\r\n")); + assert!(request_text.contains("\r\nauthorization: Bearer test-key\r\n")); + let body = request_text.split("\r\n\r\n").nth(1).unwrap(); + let body: Value = serde_json::from_str(body).unwrap(); + assert_eq!(body["model"], "test-model"); + assert_eq!(body["input"], "hello"); + assert_eq!(body["encoding_format"], "float"); +} + +#[tokio::test] +async fn status_and_timeout_errors_are_unavailable() { + let server = TestHttpServer::response("500 Internal Server Error", "{}").await; + let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), None)); + assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable)); + + let server = TestHttpServer::response_after( + "200 OK", + r#"{"data":[{"embedding":[0.1,0.2]}]}"#, + Duration::from_millis(500), + ) + .await; + let embedder = OpenAiEmbedder::new( + reqwest::Client::new(), + config(server.base_url(), Some(Duration::from_millis(200))), + ); + assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable)); + + let server = TestHttpServer::response_after( + "200 OK", + r#"{"data":[{"embedding":[0.1,0.2]}]}"#, + Duration::from_millis(100), + ) + .await; + let embedder = OpenAiEmbedder::new( + reqwest::Client::new(), + config(server.base_url(), Some(Duration::from_secs(1))), + ); + assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]); +} + +#[tokio::test] +async fn uses_the_injected_client() { + let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; + let client = reqwest::Client::builder() + .user_agent("litellm-embedder-test") + .build() + .unwrap(); + let embedder = OpenAiEmbedder::new(client, config(server.base_url(), None)); + assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]); + let request = server.request.lock().unwrap().clone().unwrap(); + let request_text = String::from_utf8(request).unwrap(); + assert!(request_text.contains("\r\nuser-agent: litellm-embedder-test\r\n")); +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs new file mode 100644 index 00000000000..38cd9e2f908 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs @@ -0,0 +1,38 @@ +use litellm_cache_qdrant_semantic::prompt_from_messages; +use serde_json::json; + +#[test] +fn prompt_matches_python_message_content_rules() { + let messages = vec![ + json!({"role": "user", "content": "hello"}), + json!({ + "role": "user", + "content": [ + {"type": "text", "text": "world"}, + {"type": "image_url", "image_url": {"url": "ignored"}}, + {"type": "text", "text": "!"}, + ], + }), + ]; + + assert_eq!(prompt_from_messages(&messages), "helloworld!"); +} + +#[test] +fn prompt_includes_search_result_text_and_compact_citations() { + let messages = vec![json!({ + "role": "tool", + "content": null, + "search_results": [{ + "source": "source", + "title": "title", + "content": [{"text": "body"}], + "citations": {"page": 1, "section": "intro"}, + }], + })]; + + assert_eq!( + prompt_from_messages(&messages), + r#"sourcetitlebody{"page":1,"section":"intro"}"# + ); +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs new file mode 100644 index 00000000000..c7522c0b313 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -0,0 +1,422 @@ +#[path = "support/mod.rs"] +mod support; + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use litellm_cache::{BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext}; +use litellm_cache_qdrant_semantic::{ + Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization, +}; +use litellm_cache_response::{ + CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, +}; +use qdrant_client::Payload; +use qdrant_client::{ + Qdrant, + qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams}, +}; +use serde_json::{Value as JsonValue, json}; + +use support::{FakeQdrant, FakeState, StoredPoint}; + +#[derive(Clone)] +struct FixedEmbedder { + vectors: Arc>>, +} + +impl FixedEmbedder { + fn new(vectors: impl IntoIterator)>) -> Self { + Self { + vectors: Arc::new( + vectors + .into_iter() + .map(|(prompt, vector)| (prompt.to_owned(), vector)) + .collect(), + ), + } + } +} + +impl Embedder for FixedEmbedder { + fn model(&self) -> &str { + "fixed" + } + + async fn embed(&self, input: &str) -> Result, Error> { + self.vectors.get(input).cloned().ok_or(Error::Unavailable) + } +} + +fn config(quantization: Quantization) -> QdrantSemanticConfig { + QdrantSemanticConfig { + collection_name: "semantic".to_owned(), + similarity_threshold: 0.9, + vector_size: 2, + quantization, + } +} + +fn context(prompt: &str) -> SemanticCacheContext { + SemanticCacheContext { + messages: Some(json!([{"role": "user", "content": prompt}])), + ..Default::default() + } +} + +fn value(response: JsonValue) -> CacheEntry { + CacheEntry { + timestamp: Some(1.0), + response, + } +} + +async fn connect( + server: &FakeQdrant, + vectors: impl IntoIterator)>, +) -> QdrantSemanticCache { + let client = Qdrant::from_url(&server.url()).build().unwrap(); + QdrantSemanticCache::connect( + client, + FixedEmbedder::new(vectors), + ResponseCacheCodec, + config(Quantization::Binary), + tokio::runtime::Handle::current(), + ) + .await + .unwrap() +} + +#[tokio::test(flavor = "multi_thread")] +#[expect( + deprecated, + reason = "the test verifies Qdrant's legacy always_ram quantization contract" +)] +async fn connect_sets_collection_quantization_and_index() { + for (quantization, expected) in [ + (Quantization::Binary, 0), + (Quantization::Scalar, 1), + (Quantization::Product, 2), + ] { + let server = FakeQdrant::start(FakeState::default()).await; + let client = Qdrant::from_url(&server.url()).build().unwrap(); + QdrantSemanticCache::connect( + client, + FixedEmbedder::new([]), + ResponseCacheCodec, + config(quantization), + tokio::runtime::Handle::current(), + ) + .await + .unwrap(); + let state = server.state.lock().unwrap(); + let request = &state.created_collections[0]; + let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) = + request + .vectors_config + .as_ref() + .and_then(|config| config.config.clone()) + else { + panic!("missing vector params"); + }; + assert_eq!(size, 2); + assert_eq!(distance, Distance::Cosine as i32); + let quantization_config = request + .quantization_config + .as_ref() + .unwrap() + .quantization + .unwrap(); + match (expected, quantization_config) { + (0, qdrant::quantization_config::Quantization::Binary(binary)) => { + assert_eq!(binary.always_ram, Some(false)); + } + (1, qdrant::quantization_config::Quantization::Scalar(scalar)) => { + assert_eq!(scalar.r#type, QuantizationType::Int8 as i32); + assert_eq!(scalar.quantile, Some(0.99)); + assert_eq!(scalar.always_ram, Some(false)); + } + (2, qdrant::quantization_config::Quantization::Product(product)) => { + assert_eq!(product.compression, CompressionRatio::X16 as i32); + assert_eq!(product.always_ram, Some(false)); + } + _ => panic!("unexpected quantization"), + } + assert!(state.index_creations >= 1); + assert_eq!(state.field_indexes[0].collection_name, "semantic"); + assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key"); + assert_eq!( + state.field_indexes[0].field_type, + Some(qdrant::FieldType::Keyword as i32) + ); + server.stop(); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn existing_collection_skips_create_and_index_failure_is_non_fatal() { + let server = FakeQdrant::start(FakeState { + collections: ["semantic".to_owned()].into_iter().collect(), + fail_field_index: true, + ..Default::default() + }) + .await; + let _cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; + let state = server.state.lock().unwrap(); + assert!(state.created_collections.is_empty()); + assert!(state.index_creations >= 1); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn async_and_sync_set_get_store_exact_payload() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await); + let ctx = context("hello"); + let entry = value(json!({"answer": 42})); + cache + .async_set_cache("key", entry.clone(), ctx.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache("key", &ctx).await.unwrap().as_ref(), + Some(&entry) + ); + { + let state = server.state.lock().unwrap(); + let payload = &state.points[0].payload; + let mut payload_keys = payload.keys().cloned().collect::>(); + payload_keys.sort(); + assert_eq!(payload_keys, ["litellm_cache_key", "response", "text"]); + assert_eq!(payload["litellm_cache_key"], Value::from("key")); + assert_eq!( + payload["response"], + Value::from(String::from_utf8(ResponseCacheCodec.encode(&entry).unwrap()).unwrap()) + ); + } + let sync_entry = entry.clone(); + let sync_cache = cache.clone(); + let sync_ctx = ctx.clone(); + tokio::task::spawn_blocking(move || { + sync_cache + .set_cache("sync", sync_entry.clone(), &sync_ctx) + .unwrap(); + assert_eq!( + sync_cache.get_cache("sync", &sync_ctx).unwrap(), + Some(sync_entry) + ); + }) + .await + .unwrap(); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn misses_and_payload_validation_are_safe() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect( + &server, + [("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])], + ) + .await; + let entry = value(json!({"answer": 1})); + cache + .async_set_cache("key", entry, context("hello")) + .await + .unwrap(); + assert_eq!( + cache + .async_get_cache("other", &context("hello")) + .await + .unwrap(), + None + ); + assert_eq!( + cache + .async_get_cache("key", &context("near")) + .await + .unwrap(), + None + ); + server.insert_point(StoredPoint { + id: Some(PointId::from(99_u64)), + vector: vec![1.0, 0.0], + payload: Payload::try_from(json!({ + "litellm_cache_key": 99, + "response": "{}", + })) + .unwrap() + .into(), + }); + assert_eq!( + cache + .async_get_cache("99", &context("hello")) + .await + .unwrap(), + None + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect(&server, [("one", vec![1.0, 0.0]), ("two", vec![0.0, 1.0])]).await; + let empty = SemanticCacheContext::default(); + assert_eq!( + cache + .async_set_cache("key", value(json!({})), empty.clone()) + .await, + Err(Error::MissingPrompt) + ); + assert_eq!( + cache.async_get_cache("key", &empty).await, + Err(Error::MissingPrompt) + ); + assert_eq!( + cache.async_get_cache("key", &context("unknown")).await, + Err(Error::Unavailable) + ); + cache + .async_set_cache( + "ttl", + value(json!({"ttl": true})), + context("one").with_ttl(Some(Duration::from_secs(1))), + ) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(1_100)).await; + assert!( + cache + .async_get_cache( + "ttl", + &context("one").with_ttl(Some(Duration::from_secs(1))), + ) + .await + .unwrap() + .is_some() + ); + cache + .async_set_cache_pipeline( + vec![ + ("one".to_owned(), value(json!({"n": 1}))), + ("two".to_owned(), value(json!({"n": 2}))), + ], + context("one"), + ) + .await + .unwrap(); + assert!( + cache + .async_get_cache("one", &context("one")) + .await + .unwrap() + .is_some() + ); + assert!( + cache + .async_get_cache("two", &context("one")) + .await + .unwrap() + .is_some() + ); + assert_eq!( + server.state.lock().unwrap().upsert_waits, + vec![Some(true), Some(true), Some(true)] + ); + assert_eq!(cache.get_ttl(&context("one")), None); + assert_eq!( + cache.test_connection().await, + Err(Error::UnsupportedOperation) + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn response_payloads_decode_and_invalid_entries_fail() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; + for (key, response) in [ + ("python", json!("{'timestamp': 1.0, 'response': {'a': 1}}")), + ("garbage", json!("not json")), + ("missing", json!("unused")), + ] { + let mut payload = serde_json::Map::new(); + payload.insert("litellm_cache_key".to_owned(), json!(key)); + if key != "missing" { + payload.insert("response".to_owned(), response); + } + server.insert_point(StoredPoint { + id: Some(PointId::from(key.len() as u64)), + vector: vec![1.0, 0.0], + payload: Payload::try_from(JsonValue::Object(payload)) + .unwrap() + .into(), + }); + } + assert_eq!( + cache + .async_get_cache("python", &context("hello")) + .await + .unwrap(), + Some(value(json!({"a": 1}))) + ); + assert_eq!( + cache.async_get_cache("garbage", &context("hello")).await, + Err(Error::InvalidEntry) + ); + assert_eq!( + cache.async_get_cache("missing", &context("hello")).await, + Err(Error::InvalidEntry) + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn response_cache_facade_turns_invalid_entry_into_miss() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await); + let request = ResponseCacheRequest::::new(CacheKeyInput { + preset: Some("key".to_owned()), + ..Default::default() + }) + .with_context(context("hello")); + let response = json!({"answer": 42}); + let facade = ResponseCache::new(cache.clone()); + facade + .async_store(&request, response.clone(), Duration::from_secs(1)) + .await + .unwrap(); + assert_eq!( + facade + .async_lookup(&request, Duration::from_secs(1)) + .await + .unwrap(), + Some(response) + ); + { + let mut state = server.state.lock().unwrap(); + state.points[0] + .payload + .insert("response".to_owned(), Value::from("not json")); + } + assert_eq!( + facade + .async_lookup(&request, Duration::from_secs(1)) + .await + .unwrap(), + None + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn stopped_qdrant_server_maps_to_unavailable() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; + server.stop(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!( + cache.async_get_cache("key", &context("hello")).await, + Err(Error::Unavailable) + ); +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs new file mode 100644 index 00000000000..9a556ae7df5 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs @@ -0,0 +1,342 @@ +use std::{ + collections::{HashMap, HashSet}, + net::SocketAddr, + sync::{Arc, Mutex}, +}; + +use qdrant_client::qdrant::collections_server::CollectionsServer; +use qdrant_client::qdrant::{ + self, CollectionExists, CollectionExistsRequest, CollectionExistsResponse, + CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId, + PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors, + collections_server::Collections, + points_server::{Points, PointsServer}, +}; +use tokio::sync::oneshot; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::{Request, Response, Status, transport::Server}; + +#[derive(Clone, Debug)] +pub struct StoredPoint { + pub id: Option, + pub vector: Vec, + pub payload: HashMap, +} + +#[derive(Default)] +pub struct FakeState { + pub collections: HashSet, + pub created_collections: Vec, + pub field_indexes: Vec, + pub points: Vec, + pub upsert_waits: Vec>, + pub index_creations: usize, + pub fail_field_index: bool, +} + +#[derive(Clone)] +pub struct FakeQdrant { + pub state: Arc>, + pub address: SocketAddr, + shutdown: Arc>>>, +} + +impl FakeQdrant { + pub async fn start(state: FakeState) -> Self { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let state = Arc::new(Mutex::new(state)); + let service = FakeService { + state: state.clone(), + }; + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + tokio::spawn(async move { + Server::builder() + .add_service(CollectionsServer::new(service.clone())) + .add_service(PointsServer::new(service)) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }) + .await + .unwrap(); + }); + Self { + state, + address, + shutdown: Arc::new(Mutex::new(Some(shutdown_tx))), + } + } + + pub fn url(&self) -> String { + format!("http://{}", self.address) + } + + pub fn stop(&self) { + self.shutdown + .lock() + .unwrap() + .take() + .unwrap() + .send(()) + .unwrap(); + } + + pub fn insert_point(&self, point: StoredPoint) { + self.state.lock().unwrap().points.push(point); + } +} + +#[derive(Clone)] +struct FakeService { + state: Arc>, +} + +macro_rules! unimplemented_collections { + ($($name:ident, $request:ty, $response:ty);* $(;)?) => { + $( + fn $name<'life0, 'async_trait>( + &'life0 self, + _: Request<$request>, + ) -> std::pin::Pin< + Box< + dyn std::future::Future< + Output = Result, Status>, + > + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async { Err(Status::unimplemented(stringify!($name))) }) + } + )* + }; +} + +macro_rules! unimplemented_points { + ($($name:ident, $request:ty, $response:ty);* $(;)?) => { + $( + fn $name<'life0, 'async_trait>( + &'life0 self, + _: Request<$request>, + ) -> std::pin::Pin< + Box< + dyn std::future::Future< + Output = Result, Status>, + > + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async { Err(Status::unimplemented(stringify!($name))) }) + } + )* + }; +} + +#[tonic::async_trait] +impl Collections for FakeService { + async fn create( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let mut state = self.state.lock().unwrap(); + state.collections.insert(request.collection_name.clone()); + state.created_collections.push(request); + Ok(Response::new(CollectionOperationResponse { + result: true, + ..Default::default() + })) + } + + async fn collection_exists( + &self, + request: Request, + ) -> Result, Status> { + let exists = self + .state + .lock() + .unwrap() + .collections + .contains(&request.into_inner().collection_name); + Ok(Response::new(CollectionExistsResponse { + result: Some(CollectionExists { exists }), + ..Default::default() + })) + } + + unimplemented_collections!( + get, qdrant::GetCollectionInfoRequest, qdrant::GetCollectionInfoResponse; + list, qdrant::ListCollectionsRequest, qdrant::ListCollectionsResponse; + update, qdrant::UpdateCollection, qdrant::CollectionOperationResponse; + delete, qdrant::DeleteCollection, qdrant::CollectionOperationResponse; + update_aliases, qdrant::ChangeAliases, qdrant::CollectionOperationResponse; + list_collection_aliases, qdrant::ListCollectionAliasesRequest, qdrant::ListAliasesResponse; + list_aliases, qdrant::ListAliasesRequest, qdrant::ListAliasesResponse; + collection_cluster_info, qdrant::CollectionClusterInfoRequest, qdrant::CollectionClusterInfoResponse; + update_collection_cluster_setup, qdrant::UpdateCollectionClusterSetupRequest, qdrant::UpdateCollectionClusterSetupResponse; + create_shard_key, qdrant::CreateShardKeyRequest, qdrant::CreateShardKeyResponse; + delete_shard_key, qdrant::DeleteShardKeyRequest, qdrant::DeleteShardKeyResponse; + list_shard_keys, qdrant::ListShardKeysRequest, qdrant::ListShardKeysResponse; + ); +} + +#[tonic::async_trait] +impl Points for FakeService { + async fn create_field_index( + &self, + request: Request, + ) -> Result, Status> { + let mut state = self.state.lock().unwrap(); + state.index_creations += 1; + state.field_indexes.push(request.into_inner()); + if state.fail_field_index { + return Err(Status::internal("field index failure")); + } + Ok(Response::new(PointsOperationResponse::default())) + } + + async fn upsert( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let mut state = self.state.lock().unwrap(); + state.upsert_waits.push(request.wait); + for point in request.points { + let stored = StoredPoint { + id: point.id.clone(), + vector: dense_vector(point.vectors)?, + payload: point.payload, + }; + if let Some(existing) = state + .points + .iter_mut() + .find(|existing| existing.id == stored.id) + { + *existing = stored; + } else { + state.points.push(stored); + } + } + Ok(Response::new(PointsOperationResponse::default())) + } + + async fn search( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let key_filter = keyword_filter(request.filter.as_ref()); + let state = self.state.lock().unwrap(); + let mut results = state + .points + .iter() + .filter(|point| { + key_filter.as_ref().is_none_or(|(field, expected)| { + point + .payload + .get(field) + .and_then(|value| { + let value: serde_json::Value = value.clone().into(); + value + .as_str() + .map(str::to_owned) + .or_else(|| value.as_i64().map(|value| value.to_string())) + }) + .is_some_and(|value| value == *expected) + }) + }) + .map(|point| ScoredPoint { + id: point.id.clone(), + payload: point.payload.clone(), + score: cosine(&request.vector, &point.vector), + ..Default::default() + }) + .collect::>(); + results.sort_by(|left, right| right.score.total_cmp(&left.score)); + results.truncate(request.limit as usize); + Ok(Response::new(SearchResponse { + result: results, + ..Default::default() + })) + } + + unimplemented_points!( + delete, qdrant::DeletePoints, qdrant::PointsOperationResponse; + get, qdrant::GetPoints, qdrant::GetResponse; + update_vectors, qdrant::UpdatePointVectors, qdrant::PointsOperationResponse; + delete_vectors, qdrant::DeletePointVectors, qdrant::PointsOperationResponse; + set_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse; + overwrite_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse; + delete_payload, qdrant::DeletePayloadPoints, qdrant::PointsOperationResponse; + clear_payload, qdrant::ClearPayloadPoints, qdrant::PointsOperationResponse; + delete_field_index, qdrant::DeleteFieldIndexCollection, qdrant::PointsOperationResponse; + create_vector_name, qdrant::CreateVectorNameRequest, qdrant::PointsOperationResponse; + delete_vector_name, qdrant::DeleteVectorNameRequest, qdrant::PointsOperationResponse; + search_batch, qdrant::SearchBatchPoints, qdrant::SearchBatchResponse; + search_groups, qdrant::SearchPointGroups, qdrant::SearchGroupsResponse; + scroll, qdrant::ScrollPoints, qdrant::ScrollResponse; + recommend, qdrant::RecommendPoints, qdrant::RecommendResponse; + recommend_batch, qdrant::RecommendBatchPoints, qdrant::RecommendBatchResponse; + recommend_groups, qdrant::RecommendPointGroups, qdrant::RecommendGroupsResponse; + discover, qdrant::DiscoverPoints, qdrant::DiscoverResponse; + discover_batch, qdrant::DiscoverBatchPoints, qdrant::DiscoverBatchResponse; + count, qdrant::CountPoints, qdrant::CountResponse; + update_batch, qdrant::UpdateBatchPoints, qdrant::UpdateBatchResponse; + query, qdrant::QueryPoints, qdrant::QueryResponse; + query_batch, qdrant::QueryBatchPoints, qdrant::QueryBatchResponse; + query_groups, qdrant::QueryPointGroups, qdrant::QueryGroupsResponse; + facet, qdrant::FacetCounts, qdrant::FacetResponse; + search_matrix_pairs, qdrant::SearchMatrixPoints, qdrant::SearchMatrixPairsResponse; + search_matrix_offsets, qdrant::SearchMatrixPoints, qdrant::SearchMatrixOffsetsResponse; + ); +} + +fn dense_vector(vectors: Option) -> Result, Status> { + let Some(Vectors { + vectors_options: + Some(qdrant::vectors::VectorsOptions::Vector(Vector { + vector: Some(qdrant::vector::Vector::Dense(qdrant::DenseVector { data })), + .. + })), + }) = vectors + else { + return Err(Status::invalid_argument("expected dense vector")); + }; + Ok(data) +} + +fn keyword_filter(filter: Option<&Filter>) -> Option<(String, String)> { + filter? + .must + .iter() + .find_map(|condition| match condition.condition_one_of.as_ref()? { + qdrant::condition::ConditionOneOf::Field(field) => { + let qdrant::r#match::MatchValue::Keyword(value) = + field.r#match.as_ref()?.match_value.as_ref()? + else { + return None; + }; + Some((field.key.clone(), value.clone())) + } + _ => None, + }) +} + +fn cosine(left: &[f32], right: &[f32]) -> f32 { + let dot = left + .iter() + .zip(right) + .map(|(left, right)| left * right) + .sum::(); + let left_norm = left.iter().map(|value| value * value).sum::().sqrt(); + let right_norm = right.iter().map(|value| value * value).sum::().sqrt(); + dot / (left_norm * right_norm) +} diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml new file mode 100644 index 00000000000..9a8755a189e --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-redis.workspace = true +litellm-cache-response.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } +r2d2 = "0.8.10" +serde_json.workspace = true +sha2.workspace = true +tokio.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" +serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs new file mode 100644 index 00000000000..e0ac31f3630 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -0,0 +1,618 @@ +use std::{ + future::Future, + sync::{Arc, OnceLock}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::{ + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + SemanticCacheContext, +}; +use litellm_cache_redis::{ + RedisTopology, + connection::{ConnectionRef, Connections}, +}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::prompt::prompt_from_context; + +const CACHE_KEY_FIELD: &str = "litellm_cache_key"; +const VECTOR_FIELD: &str = "prompt_vector"; + +pub trait Embedder: Send + Sync + 'static { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error>; + + fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> impl Future, Error>> + Send; +} + +#[derive(Clone, Debug)] +pub struct RedisSemanticConfig { + pub index_name: String, + pub similarity_threshold: f32, +} + +struct Inner { + index_name: String, + distance_threshold: f64, + resolved_index: OnceLock, + codec: ResponseCacheCodec, + clock: fn() -> f64, +} + +impl Inner { + fn new(config: RedisSemanticConfig) -> Self { + Self { + index_name: config.index_name, + distance_threshold: 1.0 - f64::from(config.similarity_threshold), + resolved_index: OnceLock::new(), + codec: ResponseCacheCodec, + clock: timestamp, + } + } + + fn ensure_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + if let Some(name) = self.resolved_index.get() { + return Ok(name.clone()); + } + let name = match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => match create_index(connection, &self.index_name, dims) { + Ok(()) => self.index_name.clone(), + Err(_) => match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => return Err(Error::Unavailable), + }, + }, + }; + let _ = self.resolved_index.set(name.clone()); + Ok(name) + } + + fn isolated_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + let name = format!("{}_isolated", self.index_name); + match index_compatible(connection, &name, dims)? { + Some(true) => Ok(name), + Some(false) => { + redis::cmd("FT.DROPINDEX") + .arg(&name) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + create_index(connection, &name, dims)?; + Ok(name) + } + None => { + create_index(connection, &name, dims)?; + Ok(name) + } + } + } + + fn store( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + value: &CacheEntry, + prompt: &str, + vector: &[f32], + ttl: Option, + ) -> Result<(), Error> { + let index = self.ensure_index(connection, vector.len())?; + let entry_id = entry_id(prompt, tag); + let hash_key = format!("{index}:{entry_id}"); + let response = self.codec.encode(value)?; + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(&entry_id) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg(VECTOR_FIELD) + .arg(vector_buffer(vector)) + .arg("inserted_at") + .arg(format!("{}", (self.clock)())) + .arg("updated_at") + .arg(format!("{}", (self.clock)())) + .arg(CACHE_KEY_FIELD) + .arg(tag) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + if let Some(ttl) = ttl { + redis::cmd("EXPIRE") + .arg(&hash_key) + .arg(ttl_seconds(ttl)) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + } + Ok(()) + } + + fn lookup( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + vector: &[f32], + ) -> Result, Error> { + let index = self.ensure_index(connection, vector.len())?; + let query = format!( + "(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]", + escape_tag(tag) + ); + let result = redis::cmd("FT.SEARCH") + .arg(&index) + .arg(query) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg(CACHE_KEY_FIELD) + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_buffer(vector)) + .query::(connection) + .map_err(|_| Error::Unavailable)?; + let Some(fields) = first_document(&result) else { + return Ok(None); + }; + if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) { + return Ok(None); + } + if number_field(fields, "vector_distance") + .is_none_or(|distance| distance > self.distance_threshold) + { + return Ok(None); + } + let Some(response) = bytes_field(fields, "response") else { + return Ok(None); + }; + self.codec.decode(&response).map(Some) + } +} + +pub struct RedisSemanticCache { + connections: Arc>, + embedder: E, + inner: Arc, +} + +impl RedisSemanticCache { + pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result { + Ok(Self { + connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), + embedder, + inner: Arc::new(Inner::new(config)), + }) + } +} + +impl RedisSemanticCache { + pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self { + Self { + connections: Arc::new(Connections::fixed(connection)), + embedder, + inner: Arc::new(Inner::new(config)), + } + } + + pub fn with_clock(self, clock: fn() -> f64) -> Self { + Self { + inner: Arc::new(Inner { + index_name: self.inner.index_name.clone(), + distance_threshold: self.inner.distance_threshold, + resolved_index: OnceLock::new(), + codec: self.inner.codec, + clock, + }), + ..self + } + } + + pub fn embedder(&self) -> &E { + &self.embedder + } + + pub fn index_name(&self) -> &str { + &self.inner.index_name + } + + pub fn similarity_threshold(&self) -> f32 { + (1.0 - self.inner.distance_threshold) as f32 + } + + fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { + context.scope.as_deref().unwrap_or(key) + } +} + +impl BaseCache + for RedisSemanticCache +{ + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(()); + }; + let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; + let tag = Self::tag(key, context).to_string(); + self.connections.execute(|connection| { + self.inner + .store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; + let tag = Self::tag(key, context).to_string(); + self.connections + .execute(|connection| self.inner.lookup(connection, &tag, &vector)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(&context) else { + return Ok(()); + }; + let vector = self + .embedder + .async_embed(&prompt, context.metadata.as_ref()) + .await?; + let tag = Self::tag(key, &context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + .await + } + + async fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self + .embedder + .async_embed(&prompt, context.metadata.as_ref()) + .await?; + let tag = Self::tag(key, context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.lookup(connection, &tag, &vector) + }) + .await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + match Connections::run_blocking(Arc::clone(&self.connections), |connection| { + Ok(match redis::cmd("PING").query::(connection) { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + }) + .await + { + Ok(result) => Ok(result), + Err(error) => Ok(CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }), + } + } +} + +fn timestamp() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or_default() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(CACHE_KEY_FIELD.as_bytes()); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn vector_buffer(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn escape_tag(value: &str) -> String { + value + .chars() + .flat_map(|ch| { + if matches!( + ch, + ',' | '.' + | '<' + | '>' + | '{' + | '}' + | '[' + | ']' + | '\\' + | '"' + | '\'' + | ':' + | ';' + | '!' + | '@' + | '#' + | '$' + | '%' + | '^' + | '&' + | '*' + | '(' + | ')' + | '-' + | '+' + | '=' + | '~' + | '|' + | '/' + | ' ' + | '?' + ) { + vec!['\\', ch] + } else { + vec![ch] + } + }) + .collect() +} + +fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> { + redis::cmd("FT.CREATE") + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg(VECTOR_FIELD) + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg(CACHE_KEY_FIELD) + .arg("TAG") + .arg("SEPARATOR") + .arg(",") + .query::<()>(connection) + .map_err(|_| Error::Unavailable) +} + +fn index_compatible( + connection: &mut ConnectionRef<'_>, + name: &str, + dims: usize, +) -> Result, Error> { + let info = match redis::cmd("FT.INFO") + .arg(name) + .query::(connection) + { + Ok(info) => info, + Err(error) if unknown_index(&error) => return Ok(None), + Err(_) => return Err(Error::Unavailable), + }; + Ok(Some(schema_compatible(&info, dims))) +} + +fn unknown_index(error: &redis::RedisError) -> bool { + let message = error.to_string().to_lowercase(); + message.contains("unknown") && message.contains("index") +} + +fn schema_compatible(info: &redis::Value, dims: usize) -> bool { + let redis::Value::Array(entries) = info else { + return false; + }; + let attributes = entries + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some("attributes")) + .map(|pair| &pair[1]); + let Some(redis::Value::Array(attributes)) = attributes else { + return false; + }; + let fields = attributes + .iter() + .map(|attribute| { + let redis::Value::Array(attribute) = attribute else { + return (None, None, None, None, None); + }; + let mut name = None; + let mut field_type = None; + let mut dim = None; + let mut data_type = None; + let mut distance_metric = None; + for pair in attribute.as_chunks::<2>().0 { + match string_value(&pair[0]).as_deref() { + Some("identifier") => name = string_value(&pair[1]), + Some("type") => field_type = string_value(&pair[1]), + Some("dim") => dim = number_value(&pair[1]), + Some("data_type") => data_type = string_value(&pair[1]), + Some("distance_metric") => distance_metric = string_value(&pair[1]), + _ => {} + } + } + (name, field_type, dim, data_type, distance_metric) + }) + .collect::>(); + let has_field = |name: &str, field_type: &str| { + fields + .iter() + .any(|(n, t, ..)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) + }; + has_field("prompt", "TEXT") + && has_field("response", "TEXT") + && has_field("inserted_at", "NUMERIC") + && has_field("updated_at", "NUMERIC") + && has_field(CACHE_KEY_FIELD, "TAG") + && fields.iter().any(|(n, t, d, data, metric)| { + n.as_deref() == Some(VECTOR_FIELD) + && t.as_deref() == Some("VECTOR") + && *d == Some(dims as f64) + && data + .as_deref() + .is_some_and(|data| data.eq_ignore_ascii_case("float32")) + && metric + .as_deref() + .is_some_and(|metric| metric.eq_ignore_ascii_case("cosine")) + }) +} + +fn string_value(value: &redis::Value) -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(text) => Some(text.clone()), + redis::Value::VerbatimString { text, .. } => Some(text.clone()), + _ => None, + } +} + +fn number_value(value: &redis::Value) -> Option { + match value { + redis::Value::Int(number) => Some(*number as f64), + redis::Value::Double(number) => Some(*number), + _ => string_value(value).and_then(|text| text.parse().ok()), + } +} + +fn first_document(result: &redis::Value) -> Option<&[redis::Value]> { + let redis::Value::Array(items) = result else { + return None; + }; + let [count, _document_id, fields, ..] = items.as_slice() else { + return None; + }; + if !matches!(count, redis::Value::Int(count) if *count > 0) { + return None; + } + match fields { + redis::Value::Array(fields) => Some(fields.as_slice()), + _ => None, + } +} + +fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> { + fields + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some(name)) + .map(|pair| &pair[1]) +} + +fn string_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(string_value) +} + +fn number_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(number_value) +} + +fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { + match field_value(fields, name)? { + redis::Value::BulkString(bytes) => Some(bytes.clone()), + redis::Value::SimpleString(text) => Some(text.clone().into_bytes()), + _ => None, + } +} + +fn ttl_seconds(ttl: Duration) -> u64 { + ttl.as_secs() + .saturating_add(u64::from(ttl.subsec_nanos() > 0)) + .max(1) +} diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs new file mode 100644 index 00000000000..51d0b4ba5f3 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/lib.rs @@ -0,0 +1,5 @@ +mod cache; +mod prompt; + +pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +pub use prompt::prompt_from_context; diff --git a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs new file mode 100644 index 00000000000..b9c38e98d77 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs @@ -0,0 +1,97 @@ +use litellm_cache::SemanticCacheContext; +use serde_json::Value; + +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if let Some(messages) = context.messages.as_ref().and_then(Value::as_array) + && !messages.is_empty() + { + return Some(messages_text(messages)); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = parts.join("\n").trim().to_string(); + (!prompt.is_empty()).then_some(prompt) +} + +fn messages_text(messages: &[Value]) -> String { + let mut text = String::new(); + for message in messages { + let Some(message) = message.as_object() else { + continue; + }; + match message.get("content") { + Some(Value::String(content)) => text.push_str(content), + Some(Value::Array(parts)) => { + for part in parts { + if let Some(text_content) = part.get("text").and_then(Value::as_str) { + text.push_str(text_content); + } + } + } + _ => {} + } + text.push_str(&search_results_text(message.get("search_results"))); + } + text +} + +fn search_results_text(search_results: Option<&Value>) -> String { + let Some(Value::Array(results)) = search_results else { + return String::new(); + }; + let mut text = String::new(); + for result in results { + let Some(result) = result.as_object() else { + continue; + }; + for key in ["source", "title"] { + if let Some(value) = result.get(key).and_then(Value::as_str) { + text.push_str(value); + } + } + if let Some(Value::Array(content)) = result.get("content") { + for block in content { + if let Some(value) = block.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + } + if let Some(citations) = result.get("citations") { + text.push_str(&citations.to_string()); + } + } + text +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(text) => { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + } + } + Value::Array(items) => { + for item in items { + collect_input_text(item, parts); + } + } + Value::Object(map) => { + if let Some(content) = map.get("content").filter(|content| !content.is_null()) { + collect_input_text(content, parts); + return; + } + for key in ["text", "output", "input_text", "output_text"] { + if let Some(Value::String(text)) = map.get(key) { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + return; + } + } + } + } + _ => {} + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs new file mode 100644 index 00000000000..233b87ec52f --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -0,0 +1,1003 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, Error, SemanticCacheContext}; +use litellm_cache_redis_semantic::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use redis_test::{MockCmd, MockRedisConnection}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; + +const INDEX: &str = "litellm_semantic_cache_index"; + +struct FakeEmbedder { + vectors: HashMap>, + calls: Arc>>, +} + +impl FakeEmbedder { + fn new(vectors: &[(&str, &[f32])]) -> (Self, Arc>>) { + let calls = Arc::new(Mutex::new(Vec::new())); + ( + Self { + vectors: vectors + .iter() + .map(|(prompt, vector)| (prompt.to_string(), vector.to_vec())) + .collect(), + calls: Arc::clone(&calls), + }, + calls, + ) + } +} + +impl Embedder for FakeEmbedder { + fn embed(&self, prompt: &str, _: Option<&Value>) -> Result, Error> { + self.calls.lock().unwrap().push(prompt.to_string()); + + Ok(self + .vectors + .get(prompt) + .cloned() + .unwrap_or_else(|| vec![0.1, 0.2, 0.3])) + } + + async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + self.embed(prompt, metadata) + } +} + +fn config() -> RedisSemanticConfig { + RedisSemanticConfig { + index_name: INDEX.into(), + similarity_threshold: 0.9, + } +} + +fn messages_context(messages: Vec) -> SemanticCacheContext { + SemanticCacheContext { + messages: Some(Value::Array(messages)), + ..Default::default() + } +} + +fn entry() -> CacheEntry { + CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "yes"}), + } +} + +fn encoded(entry: &CacheEntry) -> Vec { + ResponseCacheCodec.encode(entry).unwrap() +} + +fn vector_bytes(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(b"litellm_cache_key"); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn s(value: &str) -> redis::Value { + redis::Value::BulkString(value.as_bytes().to_vec()) +} + +fn unknown_index_error() -> redis::RedisError { + redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name")) +} + +fn attribute(name: &str, field_type: &str, extra: Vec) -> redis::Value { + let mut parts = vec![ + s("identifier"), + s(name), + s("attribute"), + s(name), + s("type"), + s(field_type), + ]; + parts.extend(extra); + redis::Value::Array(parts) +} + +fn index_info(attributes: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("index_name"), + s(INDEX), + s("attributes"), + redis::Value::Array(attributes), + ]) +} + +fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value { + attribute( + "prompt_vector", + "VECTOR", + vec![ + s("algorithm"), + s("FLAT"), + s("data_type"), + s(data_type), + s("dim"), + redis::Value::Int(dims), + s("distance_metric"), + s(distance_metric), + ], + ) +} + +fn vector_attribute(dims: i64) -> redis::Value { + vector_attribute_with(dims, "FLOAT32", "COSINE") +} + +fn info_with_vector(vector: redis::Value) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector, + attribute("litellm_cache_key", "TAG", vec![]), + ]) +} + +fn compatible_info(dims: i64) -> redis::Value { + info_with_vector(vector_attribute(dims)) +} + +fn unscoped_info(dims: i64) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector_attribute(dims), + ]) +} + +fn create_index_command(name: &str, dims: usize) -> redis::Cmd { + let mut command = redis::cmd("FT.CREATE"); + command + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg("prompt_vector") + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg("litellm_cache_key") + .arg("TAG") + .arg("SEPARATOR") + .arg(","); + command +} + +fn search_command(index: &str, tag: &str, vector: &[f32]) -> redis::Cmd { + let mut command = redis::cmd("FT.SEARCH"); + command + .arg(index) + .arg(format!( + "(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]" + )) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg("litellm_cache_key") + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_bytes(vector)); + command +} + +fn hit_fields(tag: &str, distance: &str, response: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("entry_id"), + s("stored-id"), + s("prompt"), + s("hello prompt"), + s("response"), + redis::Value::BulkString(response), + s("inserted_at"), + s("1700000000.5"), + s("updated_at"), + s("1700000000.5"), + s("litellm_cache_key"), + s(tag), + s("vector_distance"), + s(distance), + ]) +} + +fn search_result(fields: redis::Value) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::Int(1), + s("litellm_semantic_cache_index:stored-id"), + fields, + ]) +} + +fn empty_result() -> redis::Value { + redis::Value::Array(vec![redis::Value::Int(0)]) +} + +#[test] +fn store_creates_index_and_writes_hash_with_expire() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(5), Ok(1)), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(5)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + cache.set_cache(tag, value, &context).unwrap(); +} + +#[test] +fn store_without_ttl_skips_expire() { + let prompt = "hello prompt"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "key1"))) + .arg("entry_id") + .arg(entry_id(prompt, "key1")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("key1"), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + "key1", + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn lookup_returns_hit_below_distance_threshold() { + let vector = vec![0.1f32, 0.2, 0.3]; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let hit = cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]), + ) + .unwrap(); + assert_eq!(hit, Some(value)); +} + +#[test] +fn lookup_misses_above_distance_threshold_and_on_tag_mismatch() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.5", encoded(&entry())))), + ), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "other", + "0.05", + encoded(&entry()), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + let context = messages_context(vec![json!({"role": "user", "content": "hello prompt"})]); + + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); +} + +#[test] +fn lookup_returns_invalid_entry_on_malformed_response() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "key1", + "0.05", + b"not json!".to_vec(), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn missing_prompt_is_noop_and_never_embeds() { + let connection = MockRedisConnection::new(Vec::::new()).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let context = SemanticCacheContext::default(); + cache.set_cache("key1", entry(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert!(calls.lock().unwrap().is_empty()); +} + +#[test] +fn scope_overrides_key_as_filter_tag() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "scope-a"))) + .arg("entry_id") + .arg(entry_id(prompt, "scope-a")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("scope-a"), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, "scope\\-a", &vector), + Ok(search_result(hit_fields( + "scope-a", + "0.05", + encoded(&value), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = SemanticCacheContext { + scope: Some("scope-a".into()), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + + cache.set_cache("key1", value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(value)); +} + +#[test] +fn incompatible_schema_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(unscoped_info(3))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn create_index_race_rechecks_schema_and_stores() { + let prompt = "hello prompt"; + let tag = "key1"; + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new( + create_index_command(INDEX, 3), + Err::<&str, _>(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Index already exists", + ))), + ), + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn wrong_distance_metric_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2"))), + ), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn tag_special_characters_are_escaped_in_search_filter() { + let vector = vec![0.1f32, 0.2, 0.3]; + let tag = "a:b, c|d"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "a\\:b\\,\\ c\\|d", &vector), + Ok(empty_result()), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + tag, + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap(), + None + ); +} + +#[test] +fn prompt_extraction_matches_python_message_and_input_shapes() { + let vector = vec![0.1f32, 0.2, 0.3]; + let lookups = 5; + let mut commands = vec![MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(compatible_info(3)), + )]; + for _ in 0..lookups { + commands.push(MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(empty_result()), + )); + } + let connection = MockRedisConnection::new(commands).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + cache + .get_cache( + "key1", + &messages_context(vec![ + json!({"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}), + json!({"role": "assistant", "content": "reply"}), + ]), + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!(" plain input ")), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some( + json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"]), + ), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!({"output_text": " result text "})), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &messages_context(vec![json!({ + "role": "user", + "content": "question", + "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], + })]), + ) + .unwrap(); + + assert_eq!( + *calls.lock().unwrap(), + vec![ + "firstsecondreply", + "plain input", + "nested\ntail", + "result text", + "questionsrctfound{\"a\":1}", + ] + ); +} + +#[test] +fn ttl_passes_through_context_only() { + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection( + MockRedisConnection::new(Vec::::new()), + embedder, + config(), + ); + assert_eq!(cache.get_ttl(&SemanticCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&SemanticCacheContext { + ttl: Some(Duration::from_secs(9)), + ..Default::default() + }), + Some(Duration::from_secs(9)) + ); +} + +#[tokio::test] +async fn async_paths_embed_then_run_blocking_redis_work() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, tag, &vector), + Ok(search_result(hit_fields(tag, "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = messages_context(vec![json!({"role": "user", "content": prompt})]); + + cache + .async_set_cache(tag, value.clone(), context.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache(tag, &context).await.unwrap(), + Some(value) + ); +} + +#[test] +fn shared_base_index_across_dimensions_replaces_the_isolated_index() { + // Pins parity with Python's `_isolated` + overwrite=True flow. + let prompt = "shared prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let value = entry(); + let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); + let store_hash = |index: &str, vector: &[f32]| { + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{index}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ) + }; + + let vector_a = vec![0.1f32; 8]; + let connection_a = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 8), Ok("OK")), + store_hash(INDEX, &vector_a), + ]) + .assert_all_commands_consumed(); + let (embedder_a, _) = FakeEmbedder::new(&[(prompt, &vector_a)]); + let worker_a = RedisSemanticCache::with_connection(connection_a, embedder_a, config()) + .with_clock(|| 1700000000.5); + worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + + let vector_b = vec![0.2f32; 4]; + let connection_b = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 4), Ok("OK")), + store_hash(&isolated, &vector_b), + MockCmd::new( + search_command(&isolated, tag, &vector_b), + Ok(search_result(hit_fields(tag, "0.0", encoded(&value)))), + ), + MockCmd::new( + search_command(&isolated, tag, &vector_b), + Err::(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Vector dimension mismatch", + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder_b, _) = FakeEmbedder::new(&[(prompt, &vector_b)]); + let worker_b = RedisSemanticCache::with_connection(connection_b, embedder_b, config()) + .with_clock(|| 1700000000.5); + worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap(), + Some(value.clone()) + ); + + let vector_c = vec![0.3f32; 16]; + let connection_c = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), + MockCmd::new(redis::cmd("FT.INFO").arg(&isolated), Ok(compatible_info(4))), + MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")), + MockCmd::new(create_index_command(&isolated, 16), Ok("OK")), + store_hash(&isolated, &vector_c), + ]) + .assert_all_commands_consumed(); + let (embedder_c, _) = FakeEmbedder::new(&[(prompt, &vector_c)]); + let worker_c = RedisSemanticCache::with_connection(connection_c, embedder_c, config()) + .with_clock(|| 1700000000.5); + worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap_err(), + Error::Unavailable + ); +} + +#[test] +fn live_shared_index_is_replaced_across_dimensions() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + // Pins parity with Python's `_isolated` + overwrite=True flow. + let base = format!("rust_semantic_shared_{}", std::process::id()); + let isolated = format!("{base}_isolated"); + let prompt = "shared live prompt"; + let tag = "key1"; + let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); + let value = entry(); + let worker = |vector: Vec| { + let (embedder, _) = FakeEmbedder::new(&[(prompt, vector.as_slice())]); + RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: base.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap() + }; + + let worker_a = worker(vec![0.1f32; 8]); + worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + + let worker_b = worker(vec![0.2f32; 4]); + worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap(), + Some(value.clone()) + ); + + let worker_c = worker(vec![0.3f32; 16]); + worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap_err(), + Error::Unavailable + ); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + for index in [&base, &isolated] { + let _: Result<(), _> = redis::cmd("FT.DROPINDEX") + .arg(index) + .arg("DD") + .query(&mut connection); + } +} + +#[test] +fn live_store_lookup_and_ttl_against_redis_stack() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + let vector = vec![0.1f32, 0.2, 0.3, 0.4]; + let prompt = "rust semantic cache live prompt"; + let tag = "live-key"; + let index_name = format!("rust_semantic_test_{}", std::process::id()); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: index_name.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap(); + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(120)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + let value = entry(); + + cache.set_cache(tag, value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache(tag, &context).unwrap(), Some(value)); + assert_eq!(cache.get_cache("other-key", &context).unwrap(), None); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + let ttl: i64 = redis::Commands::ttl( + &mut connection, + format!("{index_name}:{}", entry_id(prompt, tag)), + ) + .unwrap(); + assert!( + ttl > 0, + "expected stored hash to carry an expiry, got {ttl}" + ); +} diff --git a/litellm-rust/crates/cache-redis/src/cache/connection.rs b/litellm-rust/crates/cache-redis/src/cache/connection.rs index 06364296992..013bf055f89 100644 --- a/litellm-rust/crates/cache-redis/src/cache/connection.rs +++ b/litellm-rust/crates/cache-redis/src/cache/connection.rs @@ -12,7 +12,7 @@ use redis::{ use super::REDIS_TIMEOUT; use crate::topology::RedisNode; -pub(super) struct PooledConnection { +pub struct PooledConnection { pub(super) connection: C, pub(super) failed: bool, } @@ -20,7 +20,7 @@ pub(super) struct PooledConnection { /// Pools connections without a checkout PING, which would double every operation's round trips. /// A timed-out command leaves its reply on the socket while redis still reports the connection /// open, so any connection whose operation failed is discarded instead of being reused. -pub(super) struct ConnectionManager(redis::Client); +pub struct ConnectionManager(redis::Client); impl ConnectionManager { pub(super) fn open(url: &str) -> Result { @@ -54,7 +54,7 @@ impl r2d2::ManageConnection for ConnectionManager { } } -pub(super) struct ClusterConnectionManager(ClusterClient); +pub struct ClusterConnectionManager(ClusterClient); impl ClusterConnectionManager { pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result { diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 46e561ddad1..d048afb69f8 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -36,6 +36,8 @@ Callers supply Unix time for response freshness. Backend TTL uses its own clock. The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API +The bridge also exposes a production-shaped response cache runtime selected through the Rust catalog. Its shipped rule set is empty, so current SDK, Router, and proxy calls stay on Python and do not construct native cache resources. Tests can inject a rule and build the native memory runtime from an ordinary Python `Cache` configuration without changing the legacy cache classes + Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs index 606c21410c7..68f2348e28b 100644 --- a/litellm-rust/crates/cache-response/src/buffer.rs +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -1,9 +1,9 @@ use std::{sync::Mutex, time::Duration}; -use litellm_cache::{BaseCache, Error, ExactCacheContext}; +use litellm_cache::Error; use serde_json::Value; -use crate::{CacheEntry, ResponseCache, ResponseCacheRequest}; +use crate::{ExactResponseCache, ResponseCacheRequest}; pub struct WriteBuffer { flush_size: usize, @@ -18,9 +18,9 @@ impl WriteBuffer { } } - pub async fn async_store>( + pub async fn async_store( &self, - cache: &ResponseCache, + cache: &dyn ExactResponseCache, request: &ResponseCacheRequest, response: Value, now: Duration, diff --git a/litellm-rust/crates/cache-response/src/exact.rs b/litellm-rust/crates/cache-response/src/exact.rs new file mode 100644 index 00000000000..f5e86b2598c --- /dev/null +++ b/litellm-rust/crates/cache-response/src/exact.rs @@ -0,0 +1,148 @@ +use std::{future::Future, pin::Pin, time::Duration}; + +use litellm_cache::{ + BaseCache, BatchCache, CacheConnectionResult, Error, ExactCacheContext, FlushCache, +}; +use serde_json::Value; + +use crate::{CacheEntry, PartialHits, ResponseCache, ResponseCacheRequest}; + +type BoxFuture<'a, T> = Pin + Send + 'a>>; + +/// Object-safe view of a `ResponseCache` over an exact-match backend, so hosts can hold every +/// exact backend behind one pointer without erasing which backend it is elsewhere. +pub trait ExactResponseCache: Send + Sync { + fn default_ttl(&self) -> Option; + + fn lookup(&self, request: &ResponseCacheRequest, now: Duration) + -> Result, Error>; + + fn store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error>; + + fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result; + + fn async_lookup<'a>( + &'a self, + request: &'a ResponseCacheRequest, + now: Duration, + ) -> BoxFuture<'a, Result, Error>>; + + fn async_store<'a>( + &'a self, + request: &'a ResponseCacheRequest, + response: Value, + now: Duration, + ) -> BoxFuture<'a, Result<(), Error>>; + + fn async_lookup_batch<'a>( + &'a self, + requests: &'a [ResponseCacheRequest], + now: Duration, + ) -> BoxFuture<'a, Result>; + + fn async_store_batch<'a>( + &'a self, + entries: Vec<(ResponseCacheRequest, Value)>, + now: Duration, + ) -> BoxFuture<'a, Result<(), Error>>; + + fn async_store_entries<'a>( + &'a self, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, + ) -> BoxFuture<'a, Result<(), Error>>; + + fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>>; + + fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result>; +} + +impl ExactResponseCache for ResponseCache +where + B: BaseCache + BatchCache + FlushCache, +{ + fn default_ttl(&self) -> Option { + ResponseCache::default_ttl(self) + } + + fn lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + ResponseCache::lookup(self, request, now) + } + + fn store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + ResponseCache::store(self, request, response, now) + } + + fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + ResponseCache::lookup_batch(self, requests, now) + } + + fn async_lookup<'a>( + &'a self, + request: &'a ResponseCacheRequest, + now: Duration, + ) -> BoxFuture<'a, Result, Error>> { + Box::pin(ResponseCache::async_lookup(self, request, now)) + } + + fn async_store<'a>( + &'a self, + request: &'a ResponseCacheRequest, + response: Value, + now: Duration, + ) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(ResponseCache::async_store(self, request, response, now)) + } + + fn async_lookup_batch<'a>( + &'a self, + requests: &'a [ResponseCacheRequest], + now: Duration, + ) -> BoxFuture<'a, Result> { + Box::pin(ResponseCache::async_lookup_batch(self, requests, now)) + } + + fn async_store_batch<'a>( + &'a self, + entries: Vec<(ResponseCacheRequest, Value)>, + now: Duration, + ) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(ResponseCache::async_store_batch(self, entries, now)) + } + + fn async_store_entries<'a>( + &'a self, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, + ) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(ResponseCache::async_store_entries(self, entries)) + } + + fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(ResponseCache::async_flush(self)) + } + + fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result> { + Box::pin(ResponseCache::test_connection(self)) + } +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index 91b36ebe24b..ab9867ac8db 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -2,6 +2,7 @@ mod buffer; mod caching; mod codec; mod embedding; +mod exact; mod response; pub use buffer::WriteBuffer; @@ -11,4 +12,5 @@ pub use caching::{ }; pub use codec::ResponseCacheCodec; pub use embedding::PartialHits; +pub use exact::ExactResponseCache; pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 2f949d511de..5088402f125 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -32,6 +32,17 @@ impl ResponseCacheRequest { } } +impl ResponseCacheRequest { + pub fn with_context(self, context: D) -> ResponseCacheRequest { + ResponseCacheRequest { + key: self.key, + controls: self.controls, + context, + max_age: self.max_age, + } + } +} + pub struct ResponseCache> where B::Context: Default + PartialEq, diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index e4f78dae8b2..dcfc0301148 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -1,12 +1,15 @@ use std::{ sync::{ - Arc, + Arc, Mutex, atomic::{AtomicU64, Ordering}, }, time::Duration, }; -use litellm_cache::{BaseCache, CacheCodec, Error}; +use litellm_cache::{ + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + SemanticCacheContext, +}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ @@ -30,6 +33,82 @@ fn request() -> ResponseCacheRequest { }) } +struct SemanticBackend { + entries: Mutex>, + contexts: Mutex>, +} + +impl BaseCache for SemanticBackend { + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + None + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + self.contexts.lock().unwrap().push(context.clone()); + self.entries.lock().unwrap().push((key.to_owned(), value)); + Ok(()) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + self.contexts.lock().unwrap().push(context.clone()); + Ok(self + .entries + .lock() + .unwrap() + .iter() + .find(|(entry_key, _)| entry_key == key) + .map(|(_, entry)| entry.clone())) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Ok(CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "ok".into(), + error: None, + }) + } +} + +#[test] +fn semantic_context_reaches_backend_for_store_and_lookup() { + let backend = Arc::new(SemanticBackend { + entries: Mutex::new(Vec::new()), + contexts: Mutex::new(Vec::new()), + }); + let cache = ResponseCache::new(backend.clone()); + let context = SemanticCacheContext { + messages: Some(json!([{"role": "user", "content": "hello"}])), + ..Default::default() + }; + let request = request().with_context(context.clone()); + let response = json!({"answer": 42}); + + cache + .store(&request, response.clone(), Duration::from_secs(100)) + .unwrap(); + + assert_eq!( + cache.lookup(&request, Duration::from_secs(100)).unwrap(), + Some(response) + ); + assert_eq!( + backend.contexts.lock().unwrap().as_slice(), + &[context.clone(), context] + ); +} + #[tokio::test] async fn sync_and_async_consumers_share_keys_ttls_and_freshness() { let clock = Arc::new(AtomicU64::new(100)); diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index 1a381d0afd8..79ef9cd18b1 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -8,4 +8,6 @@ pub enum Error { UnscopedFlush, #[error("operation is not supported by this cache")] UnsupportedOperation, + #[error("semantic cache requires request messages")] + MissingPrompt, } diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 9180ee9d0dc..36307ac9b33 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,7 +1,8 @@ use std::{sync::Mutex, time::Duration}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache, + BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext, + get_cache, }; struct TestCache { @@ -126,6 +127,24 @@ fn associated_context_preserves_backend_specific_lookup_inputs() { ); } +#[test] +fn semantic_context_with_ttl_preserves_lookup_inputs() { + let context = SemanticCacheContext { + input: Some(serde_json::json!("text")), + messages: Some(serde_json::json!([{"role": "user", "content": "hi"}])), + metadata: Some(serde_json::json!({"key": "value"})), + scope: Some("scope".into()), + ttl: None, + }; + let updated = context.with_ttl(Some(Duration::from_secs(30))); + assert_eq!(updated.ttl(), Some(Duration::from_secs(30))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, context.scope); + assert_eq!(context.with_ttl(None).ttl(), None); +} + #[tokio::test] async fn default_batch_operations_use_async_writes_and_stop_on_failure() { let cache = TestCache { diff --git a/litellm-rust/crates/core-utils/src/serde_compat.rs b/litellm-rust/crates/core-utils/src/serde_compat.rs index c767c709f50..fddab1d80e3 100644 --- a/litellm-rust/crates/core-utils/src/serde_compat.rs +++ b/litellm-rust/crates/core-utils/src/serde_compat.rs @@ -17,6 +17,11 @@ pub fn parse_str_bool(value: &str) -> Option { token.eq_ignore_ascii_case("false").then_some(false) } +/// `redis-py` string Booleans: only `1`, `true`, and `yes` (case-insensitive) are true. +pub fn parse_redis_bool(value: &str) -> bool { + value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes") +} + impl<'de> DeserializeAs<'de, i64> for LaxI64 { fn deserialize_as>(deserializer: D) -> Result { deserializer.deserialize_any(Self) diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 69ae8004d46..3bfc5bae925 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -36,6 +36,7 @@ url.workspace = true veil.workspace = true [dev-dependencies] +litellm-secrets.workspace = true litellm-auth-gcp.workspace = true litellm-llms = { workspace = true, features = ["test-support"] } rstest.workspace = true diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 19037e49033..c1265e1e91c 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -22,7 +22,12 @@ pub(crate) async fn perform_ocr_request( ) -> Result { request.response_format()?; let config = request.config; - let request = prepare_request(request, caller_document, client); + let secrets = client + .secret_source() + .resolve(&config.secret_names()) + .await + .map_err(|error| Error::Secret(std::sync::Arc::new(error)))?; + let request = prepare_request(request, caller_document, client, secrets); let hooks = OcrCallHooks::new(host.clone(), &request, config); config.ocr(client, &request, &hooks).await } diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 715aedc69df..54960256faa 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,7 +1,10 @@ use litellm_auth::{InputSource, SecretValue, Sourced}; -use litellm_llms::base_llm::ocr::{ - handler::OcrClient, - transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest}, +use litellm_llms::base_llm::{ + inference::secrets::Secrets, + ocr::{ + handler::OcrClient, + transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest}, + }, }; use super::provider_config::OcrProvider; @@ -11,6 +14,7 @@ pub(crate) fn prepare_request( request: ResolvedOcrRequest, caller_document: bool, client: &OcrClient, + secrets: Secrets, ) -> PreparedOcrRequest { let credentials = request.credentials.clone(); let (preferred_api_key_env, api_base_env) = match request.config.provider() { @@ -24,7 +28,7 @@ pub(crate) fn prepare_request( | OcrProvider::Reducto | OcrProvider::VertexAi => (None, None), }; - let secret = |name: &str| client.secrets().truthy(name); + let secret = |name: &str| secrets.truthy(name); let dynamic_api_key = credentials.dynamic_api_key.or_else(|| { credentials.api_key.clone().or_else(|| { preferred_api_key_env @@ -60,12 +64,7 @@ pub(crate) fn prepare_request( PreparedOcrRequest { model, document, - connection: OcrConnection::new( - resolved, - transport, - client.settings().clone(), - client.secrets().clone(), - ), + connection: OcrConnection::new(resolved, transport, client.settings().clone(), secrets), caller_document, optional_params, input_sources, @@ -79,6 +78,7 @@ pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedO request, true, &OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()), + std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment), ) } diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index d38d87b92cc..0b09e9be354 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -114,6 +114,10 @@ impl OcrConfigKind { with_config!(self, config => config.get_api_key_env_var()) } + pub(crate) fn secret_names(self) -> Vec<&'static str> { + with_config!(self, config => config.secret_names()) + } + pub(crate) fn get_health_check_document(self) -> OcrDocument { with_config!(self, config => config.get_health_check_document()) } @@ -213,6 +217,8 @@ fn is_document_intelligence_model(model: &str) -> bool { #[cfg(test)] mod tests { + use std::collections::HashSet; + use litellm_auth::{InputSource, Sourced}; use litellm_llms::{ base_llm::ocr::document::InlineDocument, cohere::ocr::transformation::validate_document, @@ -221,6 +227,27 @@ mod tests { use super::*; + #[rstest] + #[case(OcrConfigKind::AwsTextract)] + #[case(OcrConfigKind::AwsTextractAnalyze)] + #[case(OcrConfigKind::Cohere)] + #[case(OcrConfigKind::Mistral)] + #[case(OcrConfigKind::AzureAi)] + #[case(OcrConfigKind::AzureCohere)] + #[case(OcrConfigKind::AzureDocumentIntelligence)] + #[case(OcrConfigKind::ReductoLegacy)] + #[case(OcrConfigKind::ReductoV3)] + #[case(OcrConfigKind::VertexAi)] + #[case(OcrConfigKind::VertexDeepSeek)] + fn secret_names_include_api_keys_without_duplicates(#[case] config: OcrConfigKind) { + let names = config.secret_names(); + let unique = names.iter().collect::>(); + assert_eq!(names.len(), unique.len()); + if let Some(api_key) = config.get_api_key_env_var() { + assert!(names.contains(&api_key)); + } + } + #[rstest] #[case("cohere")] #[case("mistral")] diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 3aedc7b9023..d376f0df784 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,5 +1,6 @@ use std::sync::{Arc, Mutex}; +use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; use litellm_host::{ event::{CallEvent, MachineEvent, WireRequest}, @@ -10,11 +11,14 @@ use litellm_http::{ HttpClientPool, HttpSettings, Resolution, media::{PublicDnsResolver, UrlPolicy}, }; +use litellm_llms::base_llm::inference::secrets::{SecretSource, Secrets}; use litellm_llms::base_llm::ocr::{ error::Error as OcrError, handler::OcrClient, settings::OcrSettings, - transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig}, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, + }, }; use rstest::rstest; use serde_json::{Value, json}; @@ -27,6 +31,32 @@ use super::{ }; use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine}; +struct RecordingSecretSource { + names: Arc>>, + values: &'static [(&'static str, &'static str)], + api_base: String, +} + +impl SecretSource for RecordingSecretSource { + fn resolve<'a>( + &'a self, + names: &'a [&'static str], + ) -> BoxFuture<'a, Result> { + *self.names.lock().unwrap() = names.to_vec(); + let values = self.values; + let api_base = self.api_base.clone(); + Box::pin(async move { + Ok(Arc::new(move |name: &str| match name { + "MISTRAL_AZURE_API_BASE" => Some(api_base.clone()), + _ => values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()), + }) as Secrets) + }) + } +} + #[rstest] #[case::mistral("mistral/model", json!({}))] #[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] @@ -184,14 +214,11 @@ async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( #[case] expected_key: &str, ) { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let secret_base = base.clone(); - let client = ocr_client().with_secrets(Arc::new(move |name: &str| match name { - "MISTRAL_AZURE_API_BASE" => Some(secret_base.clone()), - "MISTRAL_API_BASE" => Some("http://127.0.0.1:9/never-read".into()), - _ => secrets - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| value.to_string()), + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: secrets, + api_base: base.clone(), })); let request = decode_request(OcrWireRequest { model: "mistral/model".into(), @@ -208,9 +235,47 @@ async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( crate::ocr::client::perform(&client, request).await.unwrap(); server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); } +#[tokio::test] +async fn mistral_ocr_resolves_provider_secrets_before_transformation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: &[("MISTRAL_API_KEY", "source-key")], + api_base: base.clone(), + })); + let request = decode_request(OcrWireRequest { + model: "mistral/mistral-ocr-latest".into(), + document: json!({ + "type":"document_url", + "document_url":"data:application/pdf;base64,YWJj" + }), + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); + assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); +} + #[tokio::test] async fn ocr_client_uses_the_injected_http_pool_configuration() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; @@ -224,7 +289,7 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() { UrlPolicy::default(), VertexAuth::default(), OcrSettings::default(), - Arc::new(litellm_core_utils::settings::ProcessEnvironment), + Arc::new(litellm_llms::base_llm::inference::secrets::EnvironmentSecrets), ) .unwrap(); crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) diff --git a/litellm-rust/crates/host-python/src/execution.rs b/litellm-rust/crates/host-python/src/execution.rs index 083c184e37e..b435bf241d2 100644 --- a/litellm-rust/crates/host-python/src/execution.rs +++ b/litellm-rust/crates/host-python/src/execution.rs @@ -29,7 +29,7 @@ pyo3::create_exception!( static FORK_GATE: ForkGate = ForkGate::new(); -/// Whether this process has started the Tokio runtime. +/// Whether this process has entered process-bound native execution. pub fn runtime_started() -> bool { FORK_GATE.started(std::process::id()) } @@ -40,9 +40,8 @@ pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> { FORK_GATE.reserve(std::process::id()) } -/// The only door to the Tokio runtime: every route reaches it through this module, which is -/// what lets the gate speak for the whole extension. `clippy.toml` disallows going around it. -fn enter_runtime() -> PyResult<()> { +/// Claims process-bound native state before runtime startup or tokenizer execution. +pub fn enter_native() -> PyResult<()> { FORK_GATE .enter(std::process::id()) .map_err(|refused| match refused { @@ -60,7 +59,7 @@ fn enter_runtime() -> PyResult<()> { #[expect(clippy::disallowed_methods, reason = "this is the gated door")] fn runtime() -> PyResult<&'static Runtime> { - enter_runtime()?; + enter_native()?; Ok(pyo3_async_runtimes::tokio::get_runtime()) } @@ -70,7 +69,7 @@ where F: Future> + Send + 'static, T: for<'py> IntoPyObject<'py> + Send + 'static, { - enter_runtime()?; + enter_native()?; pyo3_async_runtimes::tokio::future_into_py(py, future) } diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 7d164ab7535..4a33975a918 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -20,7 +20,7 @@ pub use argument::lookup; pub use callable::wrap_failure; pub use driver::run_call; pub use execution::{ - ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value, + ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value, reserve_process_for_forking, run_async, run_async_value, run_sync, run_sync_value, runtime_started, }; diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index f04b78feee1..ed15d9f7cdb 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -18,6 +18,7 @@ litellm-auth-gcp.workspace = true litellm-host.workspace = true litellm-framing.workspace = true litellm-http.workspace = true +litellm-secrets.workspace = true base64.workspace = true bytes.workspace = true data-url = "0.3.2" diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs index d476861e6e1..2ce1b0da51b 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs @@ -40,6 +40,10 @@ impl BaseOcrConfig for TextractAnalyzeDocumentConfig { type ProviderRequest = AnalyzeDocumentRequest; type Environment = TextractEnvironment; + fn secret_names(&self) -> Vec<&'static str> { + litellm_auth_aws::constants::SECRET_NAMES.to_vec() + } + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["feature_types"] } diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs index ad630a1ca4c..6eb195defaa 100644 --- a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs @@ -29,6 +29,10 @@ impl BaseOcrConfig for TextractDetectTextConfig { type ProviderRequest = DetectDocumentTextRequest; type Environment = TextractEnvironment; + fn secret_names(&self) -> Vec<&'static str> { + litellm_auth_aws::constants::SECRET_NAMES.to_vec() + } + fn get_health_check_document(&self) -> OcrDocument { health_check_document() } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs index 045d8744bc9..09639481cdf 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs @@ -28,6 +28,10 @@ impl BaseOcrConfig for AzureAICohereParseConfig { super::transformation::AzureAiOcrConfig.get_api_key_env_var() } + fn secret_names(&self) -> Vec<&'static str> { + super::transformation::AzureAiOcrConfig.secret_names() + } + fn get_health_check_document(&self) -> OcrDocument { CohereParseConfig.get_health_check_document() } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 9b27fdbb568..bfe0d76aab1 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -2,7 +2,7 @@ use std::{collections::BTreeSet, time::Duration}; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_auth::{InputSource, Sourced}; -use litellm_auth_azure::AzureAuthInputs; +use litellm_auth_azure::{AzureAuthInputs, SECRET_NAMES as AZURE_AUTH_SECRET_NAMES}; use litellm_core_utils::{ call_arguments::CallArguments, serde_compat::{FiniteF64, LaxI64}, @@ -141,6 +141,17 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { Some(AZURE_DI_API_KEY_ENV) } + fn secret_names(&self) -> Vec<&'static str> { + [ + [AZURE_DI_API_KEY_ENV, AZURE_DI_ENDPOINT_ENV].as_slice(), + AZURE_AUTH_SECRET_NAMES, + ] + .into_iter() + .flatten() + .copied() + .collect() + } + fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials { ResolvedOcrCredentials { api_key: inputs.api_key.and_then(|key| { diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 6df83e57eab..1fad860f757 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -1,5 +1,6 @@ use litellm_auth::{InputSource, Sourced}; use litellm_auth_azure::AzureAuthInputs; +use litellm_auth_azure::SECRET_NAMES as AZURE_AUTH_SECRET_NAMES; use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, url_utils::ApiUrl}; use serde_json::Value; @@ -37,6 +38,17 @@ impl BaseOcrConfig for AzureAiOcrConfig { Some(AZURE_AI_API_KEY_ENV) } + fn secret_names(&self) -> Vec<&'static str> { + [ + [AZURE_AI_API_KEY_ENV, AZURE_AI_API_BASE_ENV].as_slice(), + AZURE_AUTH_SECRET_NAMES, + ] + .into_iter() + .flatten() + .copied() + .collect() + } + fn map_ocr_params( &self, non_default_params: &CallArguments, diff --git a/litellm-rust/crates/llms/src/base_llm/inference/mod.rs b/litellm-rust/crates/llms/src/base_llm/inference/mod.rs new file mode 100644 index 00000000000..10c0454f947 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/inference/mod.rs @@ -0,0 +1 @@ +pub mod secrets; diff --git a/litellm-rust/crates/llms/src/base_llm/inference/secrets.rs b/litellm-rust/crates/llms/src/base_llm/inference/secrets.rs new file mode 100644 index 00000000000..eb13fe95116 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/inference/secrets.rs @@ -0,0 +1,19 @@ +use std::sync::Arc; + +use futures_util::future::BoxFuture; +use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; +use litellm_secrets::Error; + +pub type Secrets = Arc; + +pub trait SecretSource: Send + Sync { + fn resolve<'a>(&'a self, names: &'a [&'static str]) -> BoxFuture<'a, Result>; +} + +pub struct EnvironmentSecrets; + +impl SecretSource for EnvironmentSecrets { + fn resolve<'a>(&'a self, _names: &'a [&'static str]) -> BoxFuture<'a, Result> { + Box::pin(async { Ok(Arc::new(ProcessEnvironment) as Secrets) }) + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/mod.rs b/litellm-rust/crates/llms/src/base_llm/mod.rs index 8ed37da4573..9cced64b687 100644 --- a/litellm-rust/crates/llms/src/base_llm/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/mod.rs @@ -2,5 +2,6 @@ pub mod anthropic_messages; pub mod audio_transcription; pub mod base_model_iterator; pub mod chat; +pub mod inference; pub mod ocr; pub mod responses; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index e09842e2856..b3df8fc18c8 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -98,6 +98,8 @@ pub enum Error { "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" )] MissingReductoApiKey, + #[error("secret resolution failed: {0}")] + Secret(#[source] std::sync::Arc), #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 245261d9f92..3ec9de8197f 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; @@ -11,9 +13,10 @@ use litellm_http::{ use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; +use crate::base_llm::inference::secrets::SecretSource; use crate::base_llm::ocr::{ error::Error, - settings::{OcrSettings, Secrets}, + settings::OcrSettings, transformation::{ BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext, PreparedOcrRequest, decode_request_value, decode_response, @@ -35,7 +38,7 @@ pub struct OcrClient { document_fetcher: MediaFetcher, vertex_auth: VertexAuth, settings: OcrSettings, - secrets: Secrets, + secrets: Arc, } impl OcrClient { @@ -45,7 +48,7 @@ impl OcrClient { url_policy: UrlPolicy, vertex_auth: VertexAuth, settings: OcrSettings, - secrets: Secrets, + secrets: Arc, ) -> Result { Ok(Self { provider_http: pool.client(config, ClientVariant::Provider)?, @@ -77,7 +80,7 @@ impl OcrClient { &self.settings } - pub fn secrets(&self) -> &Secrets { + pub fn secret_source(&self) -> &Arc { &self.secrets } @@ -92,7 +95,7 @@ impl OcrClient { document_fetcher: MediaFetcher::for_test(document_http), vertex_auth: VertexAuth::default(), settings: OcrSettings::default(), - secrets: std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment), + secrets: Arc::new(crate::base_llm::inference::secrets::EnvironmentSecrets), } } @@ -102,7 +105,7 @@ impl OcrClient { } #[cfg(any(test, feature = "test-support"))] - pub fn with_secrets(self, secrets: Secrets) -> Self { + pub fn with_secrets(self, secrets: Arc) -> Self { Self { secrets, ..self } } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs index f5954599b43..87461cd36aa 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs @@ -1,9 +1,7 @@ -use std::{sync::Arc, time::Duration}; +use std::time::Duration; use litellm_core_utils::settings::Lookup; -pub type Secrets = Arc; - #[derive(Clone, Debug, PartialEq)] pub struct OcrSettings { pub request_timeout: Duration, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index e02a4b7f266..0506ff3d6df 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -14,10 +14,13 @@ use serde::{ use serde_json::{Map, Value}; use serde_with::serde_as; -use crate::base_llm::ocr::{ - error::Error, - handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body}, - settings::{OcrSettings, Secrets}, +use crate::base_llm::{ + inference::secrets::Secrets, + ocr::{ + error::Error, + handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body}, + settings::OcrSettings, + }, }; pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024; @@ -436,6 +439,8 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { None } + fn secret_names(&self) -> Vec<&'static str>; + fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials { ResolvedOcrCredentials { api_key: inputs diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index d141c68db38..c0bb4c60563 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -102,6 +102,10 @@ impl BaseOcrConfig for CohereParseConfig { Some(COHERE_API_KEY_ENV) } + fn secret_names(&self) -> Vec<&'static str> { + vec![COHERE_API_KEY_ENV] + } + fn get_health_check_document(&self) -> OcrDocument { OcrDocument::ImageUrl { image_url: COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI.into(), diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index 2b14372fbec..149e8056789 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -69,6 +69,14 @@ impl BaseOcrConfig for MistralOcrConfig { Some(MISTRAL_OCR_API_KEY_ENV_VAR) } + fn secret_names(&self) -> Vec<&'static str> { + vec![ + MISTRAL_OCR_API_KEY_ENV_VAR, + "MISTRAL_AZURE_API_KEY", + "MISTRAL_AZURE_API_BASE", + ] + } + fn map_ocr_params( &self, non_default_params: &CallArguments, diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 5272be97c24..f00259984ba 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -92,6 +92,10 @@ impl BaseOcrConfig for ReductoParseV3Config { type ProviderRequest = ReductoV3Request; type Environment = Vec<(String, String)>; + fn secret_names(&self) -> Vec<&'static str> { + vec![REDUCTO_API_KEY_ENV] + } + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["formatting", "retrieval", "settings"] } @@ -180,6 +184,10 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { type ProviderRequest = ReductoLegacyRequest; type Environment = Vec<(String, String)>; + fn secret_names(&self) -> Vec<&'static str> { + vec![REDUCTO_API_KEY_ENV] + } + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["enhance"] } diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs index 9a23deefb89..86231d50f9c 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs @@ -105,6 +105,10 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { VertexAiOcrConfig.get_api_key_env_var() } + fn secret_names(&self) -> Vec<&'static str> { + VertexAiOcrConfig.secret_names() + } + fn map_ocr_params( &self, _arguments: &CallArguments, diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index 2d505ba4342..c9342c87e9a 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -34,6 +34,10 @@ impl BaseOcrConfig for VertexAiOcrConfig { Some("VERTEX_AI_API_KEY") } + fn secret_names(&self) -> Vec<&'static str> { + litellm_auth_gcp::SECRET_NAMES.to_vec() + } + fn map_ocr_params( &self, non_default_params: &CallArguments, diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index d07a9839ebf..7846beef28a 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,7 +10,7 @@ name = "_native" crate-type = ["cdylib"] [features] -default = ["abi3", "fast"] +default = ["abi3", "fast", "huggingface", "tiktoken"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] panic-test = [] @@ -20,6 +20,7 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true +futures-util.workspace = true litellm-cache.workspace = true litellm-cache-azure-blob.workspace = true litellm-cache-memory.workspace = true @@ -27,7 +28,10 @@ litellm-cache-redis.workspace = true litellm-cache-s3.workspace = true litellm-cache-gcs.workspace = true litellm-cache-disk.workspace = true +litellm-cache-redis-semantic.workspace = true litellm-cache-response.workspace = true +litellm-cache-qdrant-semantic.workspace = true +qdrant-client.workspace = true litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" } serde.workspace = true litellm-auth.workspace = true @@ -38,16 +42,21 @@ litellm-core-utils.workspace = true litellm-auth-gcp.workspace = true litellm-http.workspace = true litellm-llms.workspace = true +litellm-secrets = { workspace = true, features = ["aws"] } +litellm-secrets-types.workspace = true litellm-types.workspace = true litellm-host-python.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true +reqwest.workspace = true redis = { version = "1.7.0", features = ["tls-rustls"] } serde_json.workspace = true -tokio = { workspace = true, features = ["sync"] } +url.workspace = true +tokio = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] +litellm-secrets-aws.workspace = true serde.workspace = true serde_with.workspace = true criterion.workspace = true @@ -55,6 +64,8 @@ futures-util.workspace = true rstest.workspace = true sha2.workspace = true tokio-tungstenite.workspace = true +wiremock = "0.6.5" +aws-sdk-secretsmanager = "1.117.0" [[bench]] name = "serialization" diff --git a/litellm-rust/crates/python-bridge/README.md b/litellm-rust/crates/python-bridge/README.md new file mode 100644 index 00000000000..faaca233f5a --- /dev/null +++ b/litellm-rust/crates/python-bridge/README.md @@ -0,0 +1,5 @@ +Native OCR uses `SecretSource` with `EnvironmentSecrets`, preserving process-environment reads. Readable Python secret managers still make OCR decline to the existing Python implementation. `ResolvedSecrets` and the separate `secret_manager_binding()` snapshot are inactive foundations for a later rollout + +Cache and secret-manager catalog entries remain Python-only, including when `LITELLM_RUST=1`. The new cache runtime is not connected to SDK or gateway caching + +OCR provider requests use the shared `litellm-http` pool. AWS and Google secret-manager SDK clients keep their SDK transports, which do not yet inherit the pool's proxy, TLS, certificate, timeout, or observability configuration. Preserve those SDK transports and configure them equivalently instead of forcing them through reqwest diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json deleted file mode 100644 index ea53d1d2025..00000000000 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ /dev/null @@ -1,154 +0,0 @@ -{ - "http_settings": { - "version": 1, - "fields": { - "ssl_verify": { - "adapter": "SslVerifyInput", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [ - "none", - "bool", - "str" - ], - "unsupported_live": "configuration_error" - }, - "ssl_certificate": { - "adapter": "OptionalStrictString", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "ssl_security_level": { - "adapter": "TuningString", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "ssl_ecdh_curve": { - "adapter": "TuningString", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "force_ipv4": { - "adapter": "Truthy", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "http2": { - "adapter": "ExactTrue", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "aiohttp_trust_env": { - "adapter": "Truthy", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "disable_aiohttp_trust_env": { - "adapter": "Truthy", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "disable_aiohttp_transport": { - "adapter": "ExactTrue", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "user_agent": { - "adapter": "StrictString", - "required": true, - "precedence": "accessor", - "sensitive": false, - "shapes": [], - "unsupported_live": null - } - } - }, - "url_policy": { - "version": 1, - "fields": { - "user_url_validation": { - "adapter": "Truthy", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - }, - "user_url_allowed_hosts": { - "adapter": "HostCollection", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - } - } - }, - "provider_defaults": { - "version": 1, - "fields": { - "vertex_project": { - "adapter": "FalsyOptionalString", - "required": true, - "precedence": "module_global", - "sensitive": true, - "shapes": [], - "unsupported_live": null - }, - "vertex_location": { - "adapter": "FalsyOptionalString", - "required": true, - "precedence": "module_global", - "sensitive": true, - "shapes": [], - "unsupported_live": null - }, - "enable_azure_ad_token_refresh": { - "adapter": "ExactTrue", - "required": true, - "precedence": "module_global", - "sensitive": false, - "shapes": [], - "unsupported_live": null - } - } - }, - "secret_manager": { - "version": 1, - "fields": { - "readable": { - "adapter": "StrictBool", - "required": true, - "precedence": "accessor", - "sensitive": false, - "shapes": [], - "unsupported_live": null - } - } - } -} diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index 2ff73238202..273d3f9ca4e 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -11,10 +11,12 @@ use serde_json::Value; use super::{ cache_error, callback::PythonCallback, + config::{CacheBackendConfig, CacheConfigProjection, NativeCacheConfig}, future::{ready_none, ready_value}, native::NativeResponseCache, request::{now, request, requests}, }; +use crate::errors::RustBridgeDeclined; pub(super) enum CacheBinding { Disabled, @@ -22,7 +24,7 @@ pub(super) enum CacheBinding { PythonCallback(PythonCallback), } -#[pyclass(frozen, name = "_CacheTestBinding")] +#[pyclass(frozen, name = "_ResponseCacheRuntime")] pub(crate) struct ResolvedCache { binding: CacheBinding, pid: u32, @@ -66,6 +68,33 @@ impl ResolvedCache { #[pymethods] impl ResolvedCache { + #[staticmethod] + fn from_cache(cache: &Bound<'_, PyAny>) -> PyResult { + let config = match NativeCacheConfig::project(cache)? { + CacheConfigProjection::Native(config) => *config, + CacheConfigProjection::Unsupported(reason) => { + return Err(RustBridgeDeclined::new_err(reason.message())); + } + }; + let service = match config.backend { + CacheBackendConfig::Memory(memory) => NativeResponseCache::memory( + memory.capacity, + memory.default_ttl, + memory.max_entry_bytes, + ), + _ => { + return Err(RustBridgeDeclined::new_err( + "native response cache activation is not implemented for this backend", + )); + } + }; + Ok(Self::new(CacheBinding::Native( + service + .with_scope(config.policy.semantic_cache_scope) + .with_redis_flush_size(config.policy.redis_flush_size), + ))) + } + #[getter] fn kind(&self) -> &'static str { match self.binding { diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 9bc666c4f2d..b6e08102e18 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -2,6 +2,7 @@ use std::{path::PathBuf, time::Duration}; use litellm_auth_aws::AwsAuthConfig; use litellm_cache::CacheType; +use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, QdrantSemanticConfig, Quantization}; use litellm_cache_redis::{RedisNode, RedisTopology}; use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; use pyo3::{ @@ -10,7 +11,7 @@ use pyo3::{ types::{PyAny, PyBool, PyDict, PyList, PyString}, }; -use super::{native::NativeResponseCache, request::duration}; +use super::{identity::BackendIdentity, native::NativeResponseCache, request::duration}; #[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct CachePolicy { @@ -88,6 +89,24 @@ pub(super) struct GcsCacheConfig { pub(super) path_service_account: Option, } +pub(super) struct AzureBlobCacheConfig { + pub(super) account_url: String, + pub(super) container: String, +} + +#[allow( + dead_code, + reason = "embedding settings are projected so drift falls back to Python" +)] +pub(super) struct RedisSemanticCacheConfig { + pub(super) redis_url: String, + pub(super) index_name: String, + pub(super) similarity_threshold: f64, + pub(super) embedding_model: String, + pub(super) embedding_max_input_tokens: Option, + pub(super) embedding_timeout: Option, +} + struct RedisClientProjection<'py> { topology: RedisTopology, host: String, @@ -107,9 +126,25 @@ pub(super) struct ValkeySemanticCacheConfig { pub(super) connection: RedisConnectionConfig, } -pub(super) struct AzureBlobCacheConfig { - pub(super) account_url: String, - pub(super) container: String, +pub(super) struct QdrantSemanticCacheConfig { + pub(super) grpc_url: String, + pub(super) api_key: Option, + pub(super) collection_name: String, + pub(super) similarity_threshold: f64, + pub(super) vector_size: u64, + pub(super) embedding: OpenAiEmbedderConfig, + pub(super) quantization: Quantization, +} + +impl QdrantSemanticCacheConfig { + pub(super) fn to_qdrant_config(&self) -> QdrantSemanticConfig { + QdrantSemanticConfig { + collection_name: self.collection_name.clone(), + similarity_threshold: self.similarity_threshold, + vector_size: self.vector_size, + quantization: self.quantization.clone(), + } + } } pub(super) enum CacheBackendConfig { @@ -120,6 +155,8 @@ pub(super) enum CacheBackendConfig { ValkeySemantic(Box), Disk(DiskCacheConfig), AzureBlob(AzureBlobCacheConfig), + RedisSemantic(Box), + QdrantSemantic(Box), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -139,6 +176,8 @@ pub(super) enum UnsupportedCacheConfig { S3Option, GcsBucket, DiskStore, + QdrantEndpoint, + SemanticEmbedding, } impl UnsupportedCacheConfig { @@ -154,6 +193,10 @@ impl UnsupportedCacheConfig { Self::S3Option => "native S3 configuration requires Python", Self::GcsBucket => "native GCS cache requires a configured bucket name", Self::DiskStore => "native disk cache requires the built-in diskcache store", + Self::QdrantEndpoint => { + "native Qdrant requires the default REST port so the gRPC port can be derived" + } + Self::SemanticEmbedding => "native semantic embedding requires Python", } } } @@ -224,138 +267,188 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::QdrantSemantic) => match project_qdrant_semantic(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::QdrantSemantic(Box::new(backend)), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| { CacheConfigProjection::Native(Box::new(Self { policy, backend: CacheBackendConfig::AzureBlob(backend), })) }), - Some(CacheType::RedisSemantic | CacheType::QdrantSemantic) | None => Ok( - CacheConfigProjection::Unsupported(UnsupportedCacheConfig::Backend), - ), + Some(CacheType::RedisSemantic) => project_redis_semantic(&backend).map(|backend| { + CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::RedisSemantic(Box::new(backend)), + })) + }), + None => Ok(CacheConfigProjection::Unsupported( + UnsupportedCacheConfig::Backend, + )), } } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - let default_ttl = match &self.backend { - CacheBackendConfig::Memory(config) => Some(config.default_ttl), - CacheBackendConfig::Redis(config) => Some(config.default_ttl), - CacheBackendConfig::S3(_) => None, - CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO), - CacheBackendConfig::Disk(_) - | CacheBackendConfig::AzureBlob(_) - | CacheBackendConfig::Gcs(_) => None, - }; - if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_)) - && service.default_ttl() != default_ttl - { - return Some("facade and native backend default TTLs must match"); - } - match &self.backend { - CacheBackendConfig::Memory(config) if service.kind() != "memory" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => { - Some("facade and native backend capacities must match") - } - CacheBackendConfig::Memory(config) - if service.max_entry_bytes() != Some(config.max_entry_bytes) => - { - Some("facade and native backend item limits must match") - } - CacheBackendConfig::Memory(_) => None, - CacheBackendConfig::Redis(_) if service.kind() != "redis" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::Redis(config) if service.topology() != Some(&config.topology) => { - Some("facade and native backend topologies must match") - } - CacheBackendConfig::Redis(config) => (service.namespace() - != config.namespace.as_deref()) - .then_some("facade and native backend namespaces must match"), - CacheBackendConfig::S3(_) if service.kind() != "s3" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::S3(config) if service.bucket() != Some(config.bucket.as_str()) => { - Some("facade and native backend buckets must match") - } - CacheBackendConfig::S3(config) - if service.key_prefix() != Some(config.key_prefix.as_str()) => - { - Some("facade and native backend key prefixes must match") - } - CacheBackendConfig::S3(config) if service.region() != Some(config.region.as_str()) => { - Some("facade and native backend regions must match") - } - CacheBackendConfig::S3(config) - if service.endpoint() - != config - .endpoint - .as_ref() - .map(|endpoint| endpoint.url.as_str()) => - { - Some("facade and native backend endpoints must match") - } - CacheBackendConfig::S3(_) => None, - CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::Gcs(config) - if service - .gcs_backend() - .is_none_or(|backend| backend.bucket_name() != config.bucket_name) => - { - Some("facade and native backend buckets must match") - } - CacheBackendConfig::Gcs(config) - if service - .gcs_backend() - .is_none_or(|backend| backend.key_prefix() != config.key_prefix) => - { - Some("facade and native backend key prefixes must match") - } - CacheBackendConfig::Gcs(config) - if service.gcs_backend().is_none_or(|backend| { - backend.path_service_account() != config.path_service_account.as_deref() - }) => - { - Some("facade and native backend credentials must match") - } - CacheBackendConfig::Gcs(_) => None, - CacheBackendConfig::ValkeySemantic(config) => { - if service.kind() != "valkey-semantic" { - return Some("facade and native backend types must match"); - } - let Some((threshold, index_name)) = service.semantic_config() else { - return Some("facade and native backend types must match"); - }; - (threshold != config.similarity_threshold || index_name != config.index_name) - .then_some("facade and native semantic settings must match") - } - CacheBackendConfig::Disk(_) if service.kind() != "disk" => { - Some("facade and native backend types must match") - } - CacheBackendConfig::Disk(config) => { - let Some(directory) = service.directory() else { - return Some("facade and native backend types must match"); - }; - let native = std::fs::canonicalize(directory).ok(); - let facade = std::fs::canonicalize(&config.directory).ok(); - (native != facade).then_some("facade and native backend directories must match") - } - CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() { - None => Some("facade and native backend types must match"), - Some((account_url, container)) - if account_url != config.account_url || container != config.container => - { - Some("facade and native backend containers must match") - } - Some(_) => None, + self.backend.identity().mismatch(&service.identity()) + } +} + +impl CacheBackendConfig { + /// The identity a native backend must have for this facade configuration to describe it. + pub(super) fn identity(&self) -> BackendIdentity { + match self { + Self::Memory(config) => BackendIdentity::Memory { + capacity: config.capacity, + max_entry_bytes: Some(config.max_entry_bytes), + default_ttl: Some(config.default_ttl), + }, + Self::Redis(config) => BackendIdentity::Redis { + topology: config.topology.clone(), + namespace: config.namespace.clone(), + default_ttl: Some(config.default_ttl), + }, + Self::S3(config) => BackendIdentity::S3 { + bucket: config.bucket.clone(), + key_prefix: config.key_prefix.clone(), + region: config.region.clone(), + endpoint: config + .endpoint + .as_ref() + .map(|endpoint| endpoint.url.clone()), + }, + Self::Gcs(config) => BackendIdentity::Gcs { + bucket_name: config.bucket_name.clone(), + key_prefix: config.key_prefix.clone(), + path_service_account: config.path_service_account.clone(), + }, + Self::ValkeySemantic(config) => BackendIdentity::ValkeySemantic { + index_name: config.index_name.clone(), + similarity_threshold: config.similarity_threshold, + }, + Self::Disk(config) => BackendIdentity::Disk { + directory: config.directory.clone(), + }, + Self::AzureBlob(config) => BackendIdentity::AzureBlob { + account_url: config.account_url.clone(), + container: config.container.clone(), + }, + Self::RedisSemantic(config) => BackendIdentity::RedisSemantic { + index_name: config.index_name.clone(), + similarity_threshold: config.similarity_threshold as f32, + }, + Self::QdrantSemantic(config) => BackendIdentity::QdrantSemantic { + collection_name: config.collection_name.clone(), + similarity_threshold: config.similarity_threshold, + vector_size: config.vector_size, + embedding_model: config.embedding.model.clone(), }, } } } +#[inline(never)] +fn project_qdrant_semantic( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let rest_url = backend.getattr("qdrant_api_base")?.extract::()?; + let parsed = match url::Url::parse(&rest_url) { + Ok(value) => value, + Err(_) => return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)), + }; + if !matches!(parsed.scheme(), "http" | "https") + || (!parsed.path().is_empty() && parsed.path() != "/") + || parsed.query().is_some() + || parsed.host_str().is_none() + || parsed.port() != Some(6333) + { + return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)); + } + let mut grpc_url = parsed; + if grpc_url.set_port(Some(6334)).is_err() { + return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)); + } + grpc_url.set_path(""); + grpc_url.set_query(None); + + if optional_attribute(backend, "embedding_max_input_tokens")? + .is_some_and(|value| !value.is_none()) + { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + let configured_model = backend.getattr("embedding_model")?.extract::()?; + let embedding_model = configured_model + .strip_prefix("openai/") + .unwrap_or(&configured_model) + .to_owned(); + if !embedding_model.starts_with("text-embedding-") { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + let proxy_server = py_sys_module(backend.py())?; + if let Some(proxy_server) = proxy_server { + let router = proxy_server.getattr("llm_router")?; + let model_list = proxy_server.getattr("llm_model_list")?; + let embedding_router = backend.py().import("litellm.caching._embedding_router")?; + if !embedding_router + .getattr("resolve_embedding_router")? + .call1((configured_model.as_str(), router, model_list))? + .is_none() + { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + } + let litellm = backend.py().import("litellm")?; + for name in ["api_key", "openai_key", "api_base"] { + if !litellm.getattr(name)?.is_none() { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + } + let Ok(embedding_api_key) = std::env::var("OPENAI_API_KEY") else { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + }; + if embedding_api_key.is_empty() { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + let embedding_api_base = std::env::var("OPENAI_BASE_URL") + .or_else(|_| std::env::var("OPENAI_API_BASE")) + .unwrap_or_else(|_| "https://api.openai.com/v1".to_owned()); + let timeout = optional_attribute(backend, "embedding_timeout")? + .map(|value| value.extract::>()) + .transpose()? + .flatten() + .map(duration) + .transpose()?; + Ok(Ok(QdrantSemanticCacheConfig { + grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(), + api_key: optional_string(backend.getattr("qdrant_api_key")?)?, + collection_name: backend.getattr("collection_name")?.extract()?, + similarity_threshold: backend.getattr("similarity_threshold")?.extract()?, + vector_size: backend.getattr("vector_size")?.extract::()?, + embedding: OpenAiEmbedderConfig { + api_base: embedding_api_base, + api_key: embedding_api_key, + model: embedding_model, + timeout, + }, + quantization: Quantization::Binary, + })) +} + +fn py_sys_module(py: Python<'_>) -> PyResult>> { + match py + .import("sys")? + .getattr("modules")? + .get_item("litellm.proxy.proxy_server") + { + Ok(module) => Ok(Some(module)), + Err(error) if error.is_instance_of::(py) => Ok(None), + Err(error) => Err(error), + } +} + #[inline(never)] fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult { let client = backend.getattr("container_client")?; @@ -371,6 +464,27 @@ fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult, +) -> PyResult { + Ok(RedisSemanticCacheConfig { + redis_url: backend.getattr("_redis_url")?.extract::()?, + index_name: backend + .getattr("_index_name")? + .extract::>()? + .unwrap_or_else(|| "litellm_semantic_cache_index".into()), + similarity_threshold: backend.getattr("similarity_threshold")?.extract::()?, + embedding_model: backend.getattr("embedding_model")?.extract::()?, + embedding_max_input_tokens: backend + .getattr("embedding_max_input_tokens")? + .extract::>()?, + embedding_timeout: backend + .getattr("embedding_timeout")? + .extract::>()?, + }) +} + #[inline(never)] fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; @@ -961,12 +1075,13 @@ mod tests { use pyo3::{prelude::*, types::PyDict}; use litellm_cache_redis::{RedisNode, RedisTopology}; + use litellm_cache_redis_semantic::RedisSemanticConfig; use super::{ - CacheBackendConfig, CacheConfigProjection, CertificateRequirement, GcsCacheConfig, - NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig, + CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement, + GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig, }; - use crate::cache::native::NativeResponseCache; + use crate::cache::{embedder::PythonEmbedder, native::NativeResponseCache}; fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> { facade( @@ -1039,6 +1154,49 @@ mod tests { }); } + #[test] + fn redis_semantic_service_mismatch_accepts_backend_precision_threshold() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(_redis_url='redis://127.0.0.1/', _index_name='semantic_idx', similarity_threshold=0.8, embedding_model='text-embedding-3-small', embedding_max_input_tokens=None, embedding_timeout=None)\n\ + facade = SimpleNamespace(type='redis-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let backend = facade.getattr("cache").unwrap(); + let embedder = PythonEmbedder::new(backend.clone().unbind()); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("Redis semantic cache should be supported"); + }; + let CacheBackendConfig::RedisSemantic(config) = config.backend else { + panic!("expected Redis semantic configuration"); + }; + let service = NativeResponseCache::redis_semantic( + &config.redis_url, + embedder, + RedisSemanticConfig { + index_name: config.index_name.clone(), + similarity_threshold: config.similarity_threshold as f32, + }, + ) + .unwrap(); + let matching_config = NativeCacheConfig { + policy: CachePolicy { + mode: "default-on".into(), + ttl: None, + namespace: None, + supported_call_types: None, + redis_flush_size: None, + semantic_cache_scope: "key".into(), + }, + backend: CacheBackendConfig::RedisSemantic(config), + }; + assert_eq!(matching_config.service_mismatch(&service), None); + }); + } + #[test] fn projects_resolved_redis_tls_configuration() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index 3de0ceb3b67..ffd72e33e1b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -1,63 +1,144 @@ -use std::{future::Future, sync::Arc}; +use std::future::Future; use litellm_cache::Error; -use litellm_cache_valkey_semantic::Embedder; use litellm_host_python::to_py; -use pyo3::{PyTraverseError, PyVisit, prelude::*}; +use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict}; use serde_json::Value; -#[derive(Clone)] -pub(super) struct PythonEmbedder { - sync_embed: Arc>, - async_embed_callable: Arc>, +tokio::task_local! { + static PREPARED_EMBEDDING: Result, Error>; +} + +/// Runs `future` with the vector the Python embedder already produced, so the backend's +/// `async_embed` never has to call back into Python from the runtime. +pub(super) fn with_prepared_embedding( + vector: Result, Error>, + future: F, +) -> impl Future { + PREPARED_EMBEDDING.scope(vector, future) +} + +/// The Python object that owns embedding for a semantic backend. +pub(super) struct PythonEmbedder(Py); + +impl Clone for PythonEmbedder { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } } impl PythonEmbedder { - pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult { - Ok(Self { - sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()), - async_embed_callable: Arc::new(backend.getattr("_get_async_embedding")?.unbind()), - }) + pub(super) fn new(object: Py) -> Self { + Self(object) } - pub(super) fn async_embed_awaitable<'py>( - &self, - py: Python<'py>, - prompt: &str, - metadata: &Option, - ) -> PyResult> { - let metadata = to_py(py, metadata)?; - self.async_embed_callable.bind(py).call1((prompt, metadata)) + pub(super) fn object(&self) -> &Py { + &self.0 } pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&*self.sync_embed)?; - visit.call(&*self.async_embed_callable) + visit.call(&self.0) + } + + fn metadata_kwargs<'py>( + py: Python<'py>, + metadata: Option<&Value>, + ) -> PyResult> { + let kwargs = PyDict::new(py); + kwargs.set_item("metadata", to_py(py, &metadata)?)?; + Ok(kwargs) + } + + /// The awaitable of `_get_async_embedding(prompt, metadata=...)`, to run in the caller's loop. + pub(super) fn async_embedding( + &self, + py: Python<'_>, + prompt: &str, + metadata: Option<&Value>, + ) -> PyResult> { + let kwargs = Self::metadata_kwargs(py, metadata)?; + self.0 + .bind(py) + .call_method("_get_async_embedding", (prompt,), Some(&kwargs)) + .map(Bound::unbind) + } + + pub(super) fn extract(vector: Bound<'_, PyAny>) -> PyResult> { + Ok(vector + .extract::>()? + .into_iter() + .map(|value| value as f32) + .collect()) + } + + fn embed_sync(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + Python::attach(|py| { + let kwargs = Self::metadata_kwargs(py, metadata)?; + Self::extract(self.0.bind(py).call_method( + "_get_embedding", + (prompt,), + Some(&kwargs), + )?) + }) + .map_err(|_| Error::Unavailable) + } + + fn seeded_embedding() -> Result, Error> { + PREPARED_EMBEDDING + .try_with(Clone::clone) + .unwrap_or(Err(Error::Unavailable)) } } -impl Embedder for PythonEmbedder { +impl litellm_cache_valkey_semantic::Embedder for PythonEmbedder { fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { - let result = Python::attach(|py| -> PyResult> { - let metadata = to_py(py, &metadata)?; - self.sync_embed - .bind(py) - .call1((prompt, metadata))? - .extract() - }) - .map_err(|_| Error::Unavailable)?; - Ok(result.into_iter().map(|value| value as f32).collect()) + self.embed_sync(prompt, metadata) } - #[expect( - clippy::manual_async_fn, - reason = "the shared Embedder trait uses an impl Future return" - )] fn async_embed( &self, _prompt: &str, _metadata: Option<&Value>, ) -> impl Future, Error>> + Send { - async { Err(Error::Unavailable) } + std::future::ready(Self::seeded_embedding()) + } +} + +impl litellm_cache_redis_semantic::Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + self.embed_sync(prompt, metadata) + } + + fn async_embed( + &self, + _prompt: &str, + _metadata: Option<&Value>, + ) -> impl Future, Error>> + Send { + std::future::ready(Self::seeded_embedding()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn async_embed_returns_the_seeded_vector_or_unavailable() { + Python::initialize(); + let object = Python::attach(|py| py.None()); + let embedder = PythonEmbedder::new(object); + let scoped_embedder = embedder.clone(); + let scoped = with_prepared_embedding(Ok(vec![0.25]), async move { + litellm_cache_redis_semantic::Embedder::async_embed(&scoped_embedder, "prompt", None) + .await + }); + assert_eq!(scoped.await, Ok(vec![0.25])); + let unscoped = + litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await; + assert_eq!(unscoped, Err(Error::Unavailable)); + let valkey = with_prepared_embedding(Ok(vec![0.5]), async move { + litellm_cache_valkey_semantic::Embedder::async_embed(&embedder, "prompt", None).await + }); + assert_eq!(valkey.await, Ok(vec![0.5])); } } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 4b4e3255cb4..17fa278ae5e 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -11,6 +11,7 @@ use serde_json::Value; use super::{ config::{CacheConfigProjection, NativeCacheConfig}, handle::CacheTestHandle, + identity::BackendIdentity, native::NativeResponseCache, }; @@ -352,37 +353,41 @@ impl FacadeGuard { facade: &Bound<'_, PyAny>, service: &NativeResponseCache, ) -> PyResult { - let kind = service.kind(); + let identity = service.identity(); + let kind = identity.kind(); let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; if !facade.get_type().is(&cache_type) { return Err(PyTypeError::new_err( "only exact built-in Cache facades can be registered", )); } - let cluster = matches!(service.topology(), Some(RedisTopology::Cluster { .. })); - let (module, name, cache_kind) = match (kind, cluster) { - ("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), - ("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"), - ("redis", true) => ( - "litellm.caching.redis_cluster_cache", - "RedisClusterCache", - "redis", + let cluster = matches!( + identity, + BackendIdentity::Redis { + topology: RedisTopology::Cluster { .. }, + .. + } + ); + let (module, name) = match (kind, cluster) { + ("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache"), + ("redis", false) => ("litellm.caching.redis_cache", "RedisCache"), + ("redis", true) => ("litellm.caching.redis_cluster_cache", "RedisClusterCache"), + ("redis_semantic", _) => ("litellm.caching.redis_semantic_cache", "RedisSemanticCache"), + ("qdrant_semantic", _) => ( + "litellm.caching.qdrant_semantic_cache", + "QdrantSemanticCache", ), - ("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"), - ("valkey-semantic", false) => ( + ("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache"), + ("valkey-semantic", _) => ( "litellm.caching.valkey_semantic_cache", "ValkeySemanticCache", - "valkey-semantic", ), - ("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"), - ("azure-blob", _) => ( - "litellm.caching.azure_blob_cache", - "AzureBlobCache", - "azure-blob", - ), - ("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"), + ("disk", _) => ("litellm.caching.disk_cache", "DiskCache"), + ("azure-blob", _) => ("litellm.caching.azure_blob_cache", "AzureBlobCache"), + ("s3", _) => ("litellm.caching.s3_cache", "S3Cache"), _ => unreachable!(), }; + let cache_kind = identity.cache_type(); let backend = facade.getattr("cache")?; if facade.getattr("type")?.extract::()? != cache_kind || !backend.get_type().is(&py.import(module)?.getattr(name)?) @@ -400,6 +405,15 @@ impl FacadeGuard { if let Some(message) = config.service_mismatch(service) { return Err(PyTypeError::new_err(message)); } + if kind == "redis_semantic" + && service + .embedder_object() + .is_none_or(|embedder| !backend.is(embedder.bind(py))) + { + return Err(PyTypeError::new_err( + "facade backend must be the native embedder", + )); + } Ok(Self { outer: ObjectGuard::capture( py, @@ -425,6 +439,17 @@ impl FacadeGuard { "redis_kwargs", "redis_flush_size", "similarity_threshold", + "distance_threshold", + "embedding_model", + "embedding_max_input_tokens", + "embedding_timeout", + "qdrant_api_base", + "qdrant_api_key", + "collection_name", + "vector_size", + "_index_name", + "_redis_url", + "similarity_threshold", "embedding_model", "index_name", "embedding_max_input_tokens", diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 769ad3548be..61993f42279 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,12 +1,25 @@ use litellm_auth_aws::AwsAuthConfig; use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; +use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization}; use litellm_cache_redis::{RedisNode, RedisTopology}; +use litellm_cache_redis_semantic::RedisSemanticConfig; use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; use litellm_host_python::{release_gil, run_sync_value}; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use litellm_http::ClientVariant; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyTypeError}, + prelude::*, + types::PyDict, +}; +use url::Url; use super::{ - cache_error, embedder::PythonEmbedder, facade::FacadeGuard, native::NativeResponseCache, + cache_error, + config::{QdrantSemanticCacheConfig, project_redis_semantic}, + embedder::PythonEmbedder, + facade::FacadeGuard, + native::NativeResponseCache, request::duration, }; @@ -141,6 +154,105 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (url, *, collection_name, similarity_threshold, vector_size, embedding_model="text-embedding-3-small", api_key=None, embedding_api_key=None, embedding_api_base=None, embedding_timeout_seconds=None, quantization="binary"))] + #[expect( + clippy::too_many_arguments, + reason = "the test handle exposes the complete Qdrant constructor" + )] + fn qdrant_semantic( + py: Python<'_>, + url: String, + collection_name: String, + similarity_threshold: f64, + vector_size: u64, + embedding_model: &str, + api_key: Option, + embedding_api_key: Option, + embedding_api_base: Option, + embedding_timeout_seconds: Option, + quantization: &str, + ) -> PyResult { + let parsed = Url::parse(&url).map_err(|_| { + pyo3::exceptions::PyValueError::new_err( + "native Qdrant requires the default REST port so the gRPC port can be derived", + ) + })?; + if !matches!(parsed.scheme(), "http" | "https") + || (!parsed.path().is_empty() && parsed.path() != "/") + || parsed.query().is_some() + || parsed.host_str().is_none() + || parsed.port() != Some(6333) + { + return Err(pyo3::exceptions::PyValueError::new_err( + "native Qdrant requires the default REST port so the gRPC port can be derived", + )); + } + let mut grpc_url = parsed; + grpc_url.set_port(Some(6334)).map_err(|_| { + pyo3::exceptions::PyValueError::new_err( + "native Qdrant requires the default REST port so the gRPC port can be derived", + ) + })?; + grpc_url.set_path(""); + grpc_url.set_query(None); + let embedding_api_key = embedding_api_key + .or_else(|| { + std::env::var("OPENAI_API_KEY") + .ok() + .filter(|value| !value.is_empty()) + }) + .ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err( + "native semantic embedding requires an OpenAI API key", + ) + })?; + let embedding_api_base = embedding_api_base.unwrap_or_else(|| { + std::env::var("OPENAI_BASE_URL") + .or_else(|_| std::env::var("OPENAI_API_BASE")) + .unwrap_or_else(|_| "https://api.openai.com/v1".to_owned()) + }); + let quantization = match quantization { + "binary" => Quantization::Binary, + "scalar" => Quantization::Scalar, + "product" => Quantization::Product, + _ => { + return Err(pyo3::exceptions::PyValueError::new_err( + "unsupported Qdrant quantization", + )); + } + }; + let config = QdrantSemanticCacheConfig { + grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(), + api_key, + collection_name, + similarity_threshold, + vector_size, + embedding: OpenAiEmbedderConfig { + api_base: embedding_api_base, + api_key: embedding_api_key, + model: embedding_model.to_owned(), + timeout: embedding_timeout_seconds.map(duration).transpose()?, + }, + quantization, + }; + let http_config = crate::http::call_config(py, &PyDict::new(py), true)?; + let client = crate::http::pool() + .client(&http_config, ClientVariant::Provider) + .map_err(crate::http::client_error)?; + let service = run_sync_value(py, async move { + let handle = tokio::runtime::Handle::current(); + NativeResponseCache::qdrant_semantic(config, client, handle) + .await + .map_err(cache_error) + })?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[staticmethod] #[pyo3(signature = (url, similarity_threshold, index_name, embedder))] fn valkey_semantic( @@ -149,7 +261,7 @@ impl CacheTestHandle { index_name: String, embedder: &Bound<'_, PyAny>, ) -> PyResult { - let python_embedder = PythonEmbedder::from_backend(embedder)?; + let python_embedder = PythonEmbedder::new(embedder.clone().unbind()); let service = NativeResponseCache::valkey_semantic( &url, similarity_threshold, @@ -179,6 +291,36 @@ impl CacheTestHandle { }) } + #[staticmethod] + fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult { + let class = py + .import("litellm.caching.redis_semantic_cache")? + .getattr("RedisSemanticCache")?; + if !backend.get_type().is(&class) { + return Err(PyTypeError::new_err( + "native redis-semantic handles require the built-in RedisSemanticCache", + )); + } + let config = project_redis_semantic(&backend)?; + let embedder = PythonEmbedder::new(backend.unbind()); + let service = release_gil(py, move || { + NativeResponseCache::redis_semantic( + &config.redis_url, + embedder, + RedisSemanticConfig { + index_name: config.index_name, + similarity_threshold: config.similarity_threshold as f32, + }, + ) + }) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() @@ -210,6 +352,7 @@ impl CacheTestHandle { } fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + self.service.traverse(&visit)?; if let Some(guard) = &self.guard { guard.traverse(visit)?; } diff --git a/litellm-rust/crates/python-bridge/src/cache/identity.rs b/litellm-rust/crates/python-bridge/src/cache/identity.rs new file mode 100644 index 00000000000..835bafd3ff1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/identity.rs @@ -0,0 +1,511 @@ +use std::{path::PathBuf, time::Duration}; + +use litellm_cache_redis::RedisTopology; + +/// What makes a native backend the one a Python facade describes: the configuration a user can +/// observe on the Python object, captured once so facade projection and native construction +/// compare plain data instead of reaching into each backend type. +#[derive(Clone, Debug, PartialEq)] +pub(super) enum BackendIdentity { + Memory { + capacity: usize, + max_entry_bytes: Option, + default_ttl: Option, + }, + Redis { + topology: RedisTopology, + namespace: Option, + default_ttl: Option, + }, + S3 { + bucket: String, + key_prefix: String, + region: String, + endpoint: Option, + }, + Gcs { + bucket_name: String, + key_prefix: String, + path_service_account: Option, + }, + Disk { + directory: PathBuf, + }, + AzureBlob { + account_url: String, + container: String, + }, + RedisSemantic { + index_name: String, + /// The backend stores the threshold as `f32`; a facade's `f64` is compared at that width. + similarity_threshold: f32, + }, + ValkeySemantic { + index_name: String, + similarity_threshold: f64, + }, + QdrantSemantic { + collection_name: String, + similarity_threshold: f64, + vector_size: u64, + embedding_model: String, + }, +} + +const TYPES: &str = "facade and native backend types must match"; + +impl BackendIdentity { + /// The native backend name reported to Python through `_CacheTestHandle.backend`. + pub(super) fn kind(&self) -> &'static str { + match self { + Self::Memory { .. } => "memory", + Self::Redis { .. } => "redis", + Self::S3 { .. } => "s3", + Self::Gcs { .. } => "gcs", + Self::ValkeySemantic { .. } => "valkey-semantic", + Self::RedisSemantic { .. } => "redis_semantic", + Self::QdrantSemantic { .. } => "qdrant_semantic", + Self::Disk { .. } => "disk", + Self::AzureBlob { .. } => "azure-blob", + } + } + + /// The `LiteLLMCacheType` value a facade of this backend carries in `Cache.type`. + pub(super) fn cache_type(&self) -> &'static str { + match self { + Self::Memory { .. } => "local", + Self::Redis { .. } => "redis", + Self::S3 { .. } => "s3", + Self::Gcs { .. } => "gcs", + Self::ValkeySemantic { .. } => "valkey-semantic", + Self::RedisSemantic { .. } => "redis-semantic", + Self::QdrantSemantic { .. } => "qdrant-semantic", + Self::Disk { .. } => "disk", + Self::AzureBlob { .. } => "azure-blob", + } + } + + /// The first difference between the facade's configuration (`self`) and the native + /// backend (`native`), in the order Python users see the attributes. + pub(super) fn mismatch(&self, native: &Self) -> Option<&'static str> { + let mut differences: Vec<(bool, &'static str)> = Vec::new(); + let mut differs = |condition: bool, message: &'static str| { + differences.push((condition, message)); + }; + match (self, native) { + ( + Self::Memory { + capacity, + max_entry_bytes, + default_ttl, + }, + Self::Memory { + capacity: native_capacity, + max_entry_bytes: native_max_entry_bytes, + default_ttl: native_default_ttl, + }, + ) => { + differs( + default_ttl != native_default_ttl, + "facade and native backend default TTLs must match", + ); + differs( + capacity != native_capacity, + "facade and native backend capacities must match", + ); + differs( + max_entry_bytes != native_max_entry_bytes, + "facade and native backend item limits must match", + ); + } + ( + Self::Redis { + topology, + namespace, + default_ttl, + }, + Self::Redis { + topology: native_topology, + namespace: native_namespace, + default_ttl: native_default_ttl, + }, + ) => { + differs( + default_ttl != native_default_ttl, + "facade and native backend default TTLs must match", + ); + differs( + topology != native_topology, + "facade and native backend topologies must match", + ); + differs( + namespace != native_namespace, + "facade and native backend namespaces must match", + ); + } + ( + Self::S3 { + bucket, + key_prefix, + region, + endpoint, + }, + Self::S3 { + bucket: native_bucket, + key_prefix: native_key_prefix, + region: native_region, + endpoint: native_endpoint, + }, + ) => { + differs( + bucket != native_bucket, + "facade and native backend buckets must match", + ); + differs( + key_prefix != native_key_prefix, + "facade and native backend key prefixes must match", + ); + differs( + region != native_region, + "facade and native backend regions must match", + ); + differs( + endpoint != native_endpoint, + "facade and native backend endpoints must match", + ); + } + ( + Self::Gcs { + bucket_name, + key_prefix, + path_service_account, + }, + Self::Gcs { + bucket_name: native_bucket_name, + key_prefix: native_key_prefix, + path_service_account: native_path_service_account, + }, + ) => { + differs( + bucket_name != native_bucket_name, + "facade and native backend buckets must match", + ); + differs( + key_prefix != native_key_prefix, + "facade and native backend key prefixes must match", + ); + differs( + path_service_account != native_path_service_account, + "facade and native backend credentials must match", + ); + } + ( + Self::Disk { directory }, + Self::Disk { + directory: native_directory, + }, + ) => { + let canonical = |path: &PathBuf| std::fs::canonicalize(path).ok(); + differs( + canonical(directory) != canonical(native_directory), + "facade and native backend directories must match", + ); + } + ( + Self::AzureBlob { + account_url, + container, + }, + Self::AzureBlob { + account_url: native_account_url, + container: native_container, + }, + ) => { + differs( + account_url != native_account_url || container != native_container, + "facade and native backend containers must match", + ); + } + ( + Self::RedisSemantic { + index_name, + similarity_threshold, + }, + Self::RedisSemantic { + index_name: native_index_name, + similarity_threshold: native_similarity_threshold, + }, + ) => { + differs( + index_name != native_index_name, + "facade and native backend index names must match", + ); + differs( + similarity_threshold != native_similarity_threshold, + "facade and native backend similarity thresholds must match", + ); + } + ( + Self::ValkeySemantic { + index_name, + similarity_threshold, + }, + Self::ValkeySemantic { + index_name: native_index_name, + similarity_threshold: native_similarity_threshold, + }, + ) => { + differs( + index_name != native_index_name + || similarity_threshold != native_similarity_threshold, + "facade and native semantic settings must match", + ); + } + ( + Self::QdrantSemantic { + collection_name, + similarity_threshold, + vector_size, + embedding_model, + }, + Self::QdrantSemantic { + collection_name: native_collection_name, + similarity_threshold: native_similarity_threshold, + vector_size: native_vector_size, + embedding_model: native_embedding_model, + }, + ) => { + differs( + collection_name != native_collection_name, + "facade and native backend collections must match", + ); + differs( + similarity_threshold != native_similarity_threshold, + "facade and native backend similarity thresholds must match", + ); + differs( + vector_size != native_vector_size, + "facade and native backend vector sizes must match", + ); + differs( + embedding_model != native_embedding_model, + "facade and native backend embedding models must match", + ); + } + _ => return Some(TYPES), + } + differences + .into_iter() + .find_map(|(condition, message)| condition.then_some(message)) + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use litellm_cache_redis::{RedisNode, RedisTopology}; + + use super::BackendIdentity; + + fn memory() -> BackendIdentity { + BackendIdentity::Memory { + capacity: 200, + max_entry_bytes: Some(1024), + default_ttl: Some(Duration::from_secs(60)), + } + } + + fn redis() -> BackendIdentity { + BackendIdentity::Redis { + topology: RedisTopology::Standalone, + namespace: Some("team".into()), + default_ttl: Some(Duration::from_secs(60)), + } + } + + fn s3() -> BackendIdentity { + BackendIdentity::S3 { + bucket: "bucket".into(), + key_prefix: "cache/".into(), + region: "us-east-1".into(), + endpoint: None, + } + } + + fn gcs() -> BackendIdentity { + BackendIdentity::Gcs { + bucket_name: "bucket".into(), + key_prefix: "cache/".into(), + path_service_account: Some("credentials.json".into()), + } + } + + fn azure() -> BackendIdentity { + BackendIdentity::AzureBlob { + account_url: "https://account.blob.core.windows.net".into(), + container: "cache".into(), + } + } + + fn redis_semantic() -> BackendIdentity { + BackendIdentity::RedisSemantic { + index_name: "idx".into(), + similarity_threshold: 0.8, + } + } + + #[test] + fn redis_semantic_thresholds_compare_at_backend_precision() { + let facade = BackendIdentity::RedisSemantic { + index_name: "idx".into(), + similarity_threshold: 0.8_f64 as f32, + }; + assert_eq!(facade.mismatch(&redis_semantic()), None); + } + + fn valkey_semantic() -> BackendIdentity { + BackendIdentity::ValkeySemantic { + index_name: "idx".into(), + similarity_threshold: 0.8, + } + } + + fn qdrant() -> BackendIdentity { + BackendIdentity::QdrantSemantic { + collection_name: "collection".into(), + similarity_threshold: 0.8, + vector_size: 1536, + embedding_model: "text-embedding-3-small".into(), + } + } + + #[test] + fn identical_identities_have_no_mismatch() { + for identity in [ + memory(), + redis(), + s3(), + gcs(), + azure(), + redis_semantic(), + valkey_semantic(), + qdrant(), + BackendIdentity::Disk { + directory: std::env::temp_dir(), + }, + ] { + assert_eq!(identity.mismatch(&identity), None, "{identity:?}"); + } + } + + #[test] + fn different_kinds_report_a_type_mismatch() { + assert_eq!( + memory().mismatch(&redis()), + Some("facade and native backend types must match") + ); + assert_eq!( + redis_semantic().mismatch(&valkey_semantic()), + Some("facade and native backend types must match") + ); + } + + #[test] + fn the_first_differing_field_names_the_mismatch() { + let BackendIdentity::Memory { capacity, .. } = memory() else { + unreachable!() + }; + assert_eq!( + memory().mismatch(&BackendIdentity::Memory { + capacity: capacity + 1, + max_entry_bytes: Some(1), + default_ttl: Some(Duration::from_secs(60)), + }), + Some("facade and native backend capacities must match") + ); + assert_eq!( + memory().mismatch(&BackendIdentity::Memory { + capacity, + max_entry_bytes: Some(1), + default_ttl: Some(Duration::from_secs(61)), + }), + Some("facade and native backend default TTLs must match") + ); + assert_eq!( + redis().mismatch(&BackendIdentity::Redis { + topology: RedisTopology::Cluster { + startup_nodes: vec![RedisNode { + host: "node".into(), + port: 7000, + }], + }, + namespace: None, + default_ttl: Some(Duration::from_secs(60)), + }), + Some("facade and native backend topologies must match") + ); + assert_eq!( + s3().mismatch(&BackendIdentity::S3 { + bucket: "bucket".into(), + key_prefix: "cache/".into(), + region: "us-east-1".into(), + endpoint: Some("http://localhost:9000".into()), + }), + Some("facade and native backend endpoints must match") + ); + assert_eq!( + gcs().mismatch(&BackendIdentity::Gcs { + bucket_name: "bucket".into(), + key_prefix: "cache/".into(), + path_service_account: None, + }), + Some("facade and native backend credentials must match") + ); + assert_eq!( + azure().mismatch(&BackendIdentity::AzureBlob { + account_url: "https://account.blob.core.windows.net".into(), + container: "other".into(), + }), + Some("facade and native backend containers must match") + ); + assert_eq!( + valkey_semantic().mismatch(&BackendIdentity::ValkeySemantic { + index_name: "idx".into(), + similarity_threshold: 0.9, + }), + Some("facade and native semantic settings must match") + ); + assert_eq!( + qdrant().mismatch(&BackendIdentity::QdrantSemantic { + collection_name: "collection".into(), + similarity_threshold: 0.8, + vector_size: 1536, + embedding_model: "text-embedding-3-large".into(), + }), + Some("facade and native backend embedding models must match") + ); + } + + #[test] + fn disk_directories_compare_canonically() { + let directory = std::env::temp_dir(); + let mut indirect = directory.clone(); + indirect.push("."); + assert_eq!( + BackendIdentity::Disk { + directory: directory.clone() + } + .mismatch(&BackendIdentity::Disk { + directory: indirect + }), + None + ); + assert_eq!( + BackendIdentity::Disk { directory }.mismatch(&BackendIdentity::Disk { + directory: "/definitely/missing".into() + }), + Some("facade and native backend directories must match") + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 278d3da1ff9..28dd6c3e798 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -5,10 +5,11 @@ mod embedder; mod facade; mod future; mod handle; +mod identity; mod native; mod request; mod resolver; -mod semantic_step; +mod semantic; use litellm_cache::Error; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 575dae45833..254b9cdea4d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,108 +1,84 @@ -use std::{path::Path, sync::Arc, time::Duration}; +use std::{sync::Arc, time::Duration}; -use litellm_cache::{ - CacheCodec, CacheConnectionResult, Error, ExactCacheContext, SemanticCacheContext, -}; +use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_azure_blob::AzureBlobCache; use litellm_cache_disk::DiskCache; use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource}; use litellm_cache_memory::InMemoryCache; +use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, QdrantSemanticCache}; use litellm_cache_redis::{RedisCache, RedisTopology}; +use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig}; use litellm_cache_response::{ - CacheEntry, CacheKeyField, PartialHits, ResponseCache, ResponseCacheCodec, - ResponseCacheRequest, WriteBuffer, + ExactResponseCache, PartialHits, ResponseCache, ResponseCacheCodec, WriteBuffer, }; use litellm_cache_s3::{S3Cache, S3CacheConfig}; use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig}; -use pyo3::prelude::*; +use pyo3::{PyTraverseError, PyVisit, prelude::*}; use serde_json::Value; use super::{ + config::QdrantSemanticCacheConfig, embedder::PythonEmbedder, - request::NativeRequest, - semantic_step::{SemanticEmbedExecution, drive_semantic}, + identity::BackendIdentity, + request::{NativeRequest, now}, + semantic::{EmbeddingFailure, SemanticExecution, SemanticOperation, drive}, }; -fn semantic_key(request: &NativeRequest, scope: &str) -> litellm_cache_response::CacheKeyInput { - let mut key = request.key.clone(); - if key.preset.is_some() { - return key; +/// What the Python embedder receives for one semantic request. +pub(super) struct EmbeddingInput { + pub(super) prompt: String, + pub(super) metadata: Option, +} + +/// An exact-match backend behind one pointer, with the identity its facade must reproduce. +pub(super) struct ExactService { + cache: Arc, + buffer: Option, + identity: BackendIdentity, +} + +impl ExactService { + fn new(cache: Arc, identity: BackendIdentity) -> Arc { + Arc::new(Self { + cache, + buffer: None, + identity, + }) } - key.fields - .retain(|field| !matches!(field.name.as_str(), "messages" | "prompt" | "input")); - const TENANT: [&str; 3] = [ - "user_api_key", - "user_api_key_team_id", - "user_api_key_org_id", - ]; - let end_user = (scope == "end_user").then_some("user_api_key_end_user_id"); - for name in TENANT.into_iter().chain(end_user) { - let sources = [ - request.metadata.as_ref(), - request.litellm_metadata.as_ref(), - request - .litellm_params - .as_ref() - .and_then(|params| params.get("metadata")), - request - .litellm_params - .as_ref() - .and_then(|params| params.get("litellm_metadata")), - ]; - let Some(value) = sources.into_iter().flatten().find_map(|source| { - source - .as_object() - .and_then(|values| values.get(name)) - .filter(|value| !value.is_null()) - }) else { - continue; - }; - let value = match value { - Value::Null => continue, - Value::String(text) => text.clone(), - other => other.to_string(), - }; - key.fields.push(CacheKeyField { - name: name.to_owned(), - value: Some(value), - api_parameter: true, - internal_parameter: false, - }); - } - key } #[derive(Clone)] pub(super) enum NativeResponseCache { - Memory(Arc>>), - Redis { - cache: Arc>>, - buffer: Option>, - }, - S3(Arc>>), - Gcs(Arc>>), + Exact(Arc), ValkeySemantic { cache: Arc>>, embedder: PythonEmbedder, scope: String, }, - Disk(Arc>>), - AzureBlob(Arc>>), + RedisSemantic { + cache: Arc>>, + embedder: PythonEmbedder, + }, + QdrantSemantic(Arc>>), } impl NativeResponseCache { pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { - Self::Memory(Arc::new(ResponseCache::new(Arc::new( - InMemoryCache::with_clock_and_size_measurement( - Some(capacity), - Some(ttl), - Some(max_entry_bytes), - Some(Arc::new(|entry| { - ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) - })), - super::request::now, - ), - )))) + let backend = InMemoryCache::with_clock_and_size_measurement( + Some(capacity), + Some(ttl), + Some(max_entry_bytes), + Some(Arc::new(|entry| { + ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) + })), + now, + ); + let identity = BackendIdentity::Memory { + capacity: backend.max_size_in_memory(), + max_entry_bytes: backend.max_entry_bytes(), + default_ttl: None, + }; + Self::exact(ResponseCache::new(Arc::new(backend)), identity) } pub fn redis( @@ -113,19 +89,97 @@ impl NativeResponseCache { ) -> Result { let backend = RedisCache::connect(url, topology, ttl, ResponseCacheCodec)?.with_namespace(namespace); - Ok(Self::Redis { - cache: Arc::new(ResponseCache::new(Arc::new(backend))), - buffer: None, - }) + let identity = BackendIdentity::Redis { + topology: backend.topology().clone(), + namespace: backend.namespace().map(str::to_owned), + default_ttl: None, + }; + Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) } pub async fn s3(config: S3CacheConfig) -> Self { let runtime = tokio::runtime::Handle::current(); - Self::S3(Arc::new(ResponseCache::new(Arc::new(S3Cache::new( - config, + let backend = S3Cache::new(config, ResponseCacheCodec, runtime); + let identity = BackendIdentity::S3 { + bucket: backend.bucket().to_owned(), + key_prefix: backend.key_prefix().to_owned(), + region: backend.region().to_owned(), + endpoint: backend.endpoint().map(str::to_owned), + }; + Self::exact(ResponseCache::new(Arc::new(backend)), identity) + } + + pub fn disk(directory: &str) -> Result { + let backend = DiskCache::open(directory, ResponseCacheCodec)?; + let identity = BackendIdentity::Disk { + directory: backend.directory().to_path_buf(), + }; + Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) + } + + pub fn gcs(config: GcsConfig, token: Option) -> Result { + let backend = match token { + Some(token) => GcsCache::with_token_source( + config, + ResponseCacheCodec, + Arc::new(StaticTokenSource(token)), + )?, + None => GcsCache::new(config, ResponseCacheCodec)?, + }; + let identity = BackendIdentity::Gcs { + bucket_name: backend.bucket_name().to_owned(), + key_prefix: backend.key_prefix().to_owned(), + path_service_account: backend.path_service_account().map(str::to_owned), + }; + Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) + } + + pub async fn azure_blob(account_url: &str, container: &str) -> Result { + let backend = AzureBlobCache::connect( + account_url, + container, ResponseCacheCodec, - runtime, - ))))) + tokio::runtime::Handle::current(), + ) + .await?; + let identity = BackendIdentity::AzureBlob { + account_url: backend.account_url().to_owned(), + container: backend.container_name().to_owned(), + }; + Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) + } + + /// Wraps a built exact backend; the TTL a facade must match comes from the built cache. + fn exact(cache: ResponseCache, identity: BackendIdentity) -> Self + where + ResponseCache: ExactResponseCache + 'static, + B: litellm_cache::BaseCache, + B::Context: Default + PartialEq, + { + let cache: Arc = Arc::new(cache); + let default_ttl = cache.default_ttl(); + let identity = match identity { + BackendIdentity::Memory { + capacity, + max_entry_bytes, + .. + } => BackendIdentity::Memory { + capacity, + max_entry_bytes, + default_ttl, + }, + BackendIdentity::Redis { + topology, + namespace, + .. + } => BackendIdentity::Redis { + topology, + namespace, + default_ttl, + }, + other => other, + }; + Self::Exact(ExactService::new(cache, identity)) } pub fn valkey_semantic( @@ -150,84 +204,76 @@ impl NativeResponseCache { }) } - pub fn disk(directory: &str) -> Result { - let cache = DiskCache::open(directory, ResponseCacheCodec)?; - Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) + pub fn redis_semantic( + url: &str, + embedder: PythonEmbedder, + config: RedisSemanticConfig, + ) -> Result { + let backend = RedisSemanticCache::new(url, embedder.clone(), config)?; + Ok(Self::RedisSemantic { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + embedder, + }) } - pub fn gcs(config: GcsConfig, token: Option) -> Result { - let backend = match token { - Some(token) => GcsCache::with_token_source( - config, - ResponseCacheCodec, - Arc::new(StaticTokenSource(token)), - )?, - None => GcsCache::new(config, ResponseCacheCodec)?, - }; - Ok(Self::Gcs(Arc::new(ResponseCache::new(Arc::new(backend))))) - } - - pub async fn azure_blob(account_url: &str, container: &str) -> Result { - let backend = AzureBlobCache::connect( - account_url, - container, + pub async fn qdrant_semantic( + config: QdrantSemanticCacheConfig, + client: reqwest::Client, + runtime: tokio::runtime::Handle, + ) -> Result { + let qdrant = qdrant_client::Qdrant::from_url(&config.grpc_url) + .skip_compatibility_check() + .api_key(config.api_key.as_deref()) + .build() + .map_err(|_| Error::Unavailable)?; + let qdrant_config = config.to_qdrant_config(); + let embedder = OpenAiEmbedder::new(client, config.embedding); + let cache = QdrantSemanticCache::connect( + qdrant, + embedder, ResponseCacheCodec, - tokio::runtime::Handle::current(), + qdrant_config, + runtime, ) .await?; - Ok(Self::AzureBlob(Arc::new(ResponseCache::new(Arc::new( - backend, - ))))) + Ok(Self::QdrantSemantic(Arc::new(ResponseCache::new( + Arc::new(cache), + )))) } - pub fn azure_blob_identity(&self) -> Option<(&str, &str)> { + pub fn identity(&self) -> BackendIdentity { match self { - Self::AzureBlob(cache) => Some(( - cache.backend().account_url(), - cache.backend().container_name(), - )), - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::Disk(_) - | Self::Gcs(_) => None, - } - } - - fn exact(request: &NativeRequest) -> ResponseCacheRequest { - ResponseCacheRequest { - key: request.key.clone(), - controls: request.controls, - context: ExactCacheContext { ttl: request.ttl }, - max_age: request.max_age, - } - } - - fn semantic( - request: &NativeRequest, - scope: &str, - ) -> ResponseCacheRequest { - ResponseCacheRequest { - key: semantic_key(request, scope), - controls: request.controls, - context: SemanticCacheContext { - input: request.input.clone(), - messages: request.messages.clone(), - metadata: request.metadata.clone(), - scope: Some(scope.to_owned()), - ttl: request.ttl, + Self::Exact(service) => service.identity.clone(), + Self::ValkeySemantic { cache, .. } => BackendIdentity::ValkeySemantic { + index_name: cache.backend().index_name().to_owned(), + similarity_threshold: cache.backend().similarity_threshold(), + }, + Self::RedisSemantic { cache, .. } => BackendIdentity::RedisSemantic { + index_name: cache.backend().index_name().to_owned(), + similarity_threshold: cache.backend().similarity_threshold(), + }, + Self::QdrantSemantic(cache) => BackendIdentity::QdrantSemantic { + collection_name: cache.backend().collection_name().to_owned(), + similarity_threshold: cache.backend().similarity_threshold(), + vector_size: cache.backend().vector_size(), + embedding_model: cache.backend().embedder().model().to_owned(), }, - max_age: request.max_age, } } + pub fn kind(&self) -> &'static str { + self.identity().kind() + } + pub fn with_redis_flush_size(self, flush_size: Option) -> Self { match self { - Self::Redis { cache, .. } => Self::Redis { - cache, - buffer: flush_size.map(|size| Arc::new(WriteBuffer::new(size))), - }, + Self::Exact(service) if matches!(service.identity, BackendIdentity::Redis { .. }) => { + Self::Exact(Arc::new(ExactService { + cache: Arc::clone(&service.cache), + buffer: flush_size.map(WriteBuffer::new), + identity: service.identity.clone(), + })) + } value => value, } } @@ -245,139 +291,56 @@ impl NativeResponseCache { } } - pub fn kind(&self) -> &'static str { + pub fn embedder_object(&self) -> Option<&Py> { match self { - Self::Memory(_) => "memory", - Self::Redis { .. } => "redis", - Self::S3(_) => "s3", - Self::Gcs(_) => "gcs", - Self::ValkeySemantic { .. } => "valkey-semantic", - Self::Disk(_) => "disk", - Self::AzureBlob(_) => "azure-blob", - } - } - - pub fn default_ttl(&self) -> Option { - match self { - Self::Memory(cache) => cache.default_ttl(), - Self::Redis { cache, .. } => cache.default_ttl(), - Self::S3(cache) => cache.default_ttl(), - Self::Gcs(cache) => cache.default_ttl(), - Self::ValkeySemantic { cache, .. } => cache.default_ttl(), - Self::Disk(cache) => cache.default_ttl(), - Self::AzureBlob(cache) => cache.default_ttl(), - } - } - - pub fn bucket(&self) -> Option<&str> { - match self { - Self::S3(cache) => Some(cache.backend().bucket()), + Self::RedisSemantic { embedder, .. } => Some(embedder.object()), _ => None, } } - pub fn key_prefix(&self) -> Option<&str> { - match self { - Self::S3(cache) => Some(cache.backend().key_prefix()), - _ => None, - } + /// The prompt and metadata this backend would embed for `request`, if it has a prompt. + pub(super) fn embedding_input(&self, request: &NativeRequest) -> Option { + let context = match self { + Self::ValkeySemantic { scope, .. } => request.scoped_semantic(scope).context, + Self::RedisSemantic { .. } => request.semantic().context, + Self::Exact(_) | Self::QdrantSemantic(_) => return None, + }; + let prompt = litellm_cache_redis_semantic::prompt_from_context(&context)?; + Some(EmbeddingInput { + prompt, + metadata: context.metadata, + }) } - pub fn region(&self) -> Option<&str> { - match self { - Self::S3(cache) => Some(cache.backend().region()), - _ => None, - } - } - - pub fn endpoint(&self) -> Option<&str> { - match self { - Self::S3(cache) => cache.backend().endpoint(), - _ => None, - } - } - - pub fn namespace(&self) -> Option<&str> { - match self { - Self::Memory(_) - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - Self::Redis { cache, .. } => cache.backend().namespace(), - } - } - - pub fn topology(&self) -> Option<&RedisTopology> { - match self { - Self::Memory(_) - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - Self::Redis { cache, .. } => Some(cache.backend().topology()), - } - } - - pub fn capacity(&self) -> Option { - match self { - Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - } - } - - pub fn max_entry_bytes(&self) -> Option { - match self { - Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - } - } - - pub fn directory(&self) -> Option<&Path> { - match self { - Self::Disk(cache) => Some(cache.backend().directory()), - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::ValkeySemantic { .. } - | Self::AzureBlob(_) - | Self::Gcs(_) => None, - } - } - - pub fn semantic_config(&self) -> Option<(f64, &str)> { - match self { - Self::ValkeySemantic { cache, .. } => Some(( - cache.backend().similarity_threshold(), - cache.backend().index_name(), - )), - _ => None, - } + /// Drives a semantic operation whose embedding comes from Python. + fn python_semantic<'py>( + &self, + py: Python<'py>, + operation: SemanticOperation, + ) -> PyResult> { + let (embedder, failure) = match self { + Self::ValkeySemantic { embedder, .. } => (embedder, EmbeddingFailure::Propagate), + Self::RedisSemantic { embedder, .. } => (embedder, EmbeddingFailure::Unavailable), + Self::Exact(_) | Self::QdrantSemantic(_) => { + return Err(pyo3::exceptions::PyRuntimeError::new_err( + "semantic execution requires a Python-embedded backend", + )); + } + }; + drive( + py, + SemanticExecution::new(self.clone(), embedder.clone(), failure, operation), + ) } pub fn lookup(&self, request: &NativeRequest, now: Duration) -> Result, Error> { match self { - Self::Memory(cache) => cache.lookup(&Self::exact(request), now), - Self::Redis { cache, .. } => cache.lookup(&Self::exact(request), now), - Self::S3(cache) => cache.lookup(&Self::exact(request), now), + Self::Exact(service) => service.cache.lookup(&request.exact(), now), Self::ValkeySemantic { cache, scope, .. } => { - cache.lookup(&Self::semantic(request, scope), now) + cache.lookup(&request.scoped_semantic(scope), now) } - Self::Gcs(cache) => cache.lookup(&Self::exact(request), now), - Self::Disk(cache) => cache.lookup(&Self::exact(request), now), - Self::AzureBlob(cache) => cache.lookup(&Self::exact(request), now), + Self::RedisSemantic { cache, .. } => cache.lookup(&request.semantic(), now), + Self::QdrantSemantic(cache) => cache.lookup(&request.semantic(), now), } } @@ -388,15 +351,12 @@ impl NativeResponseCache { now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.store(&Self::exact(request), response, now), - Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now), - Self::S3(cache) => cache.store(&Self::exact(request), response, now), + Self::Exact(service) => service.cache.store(&request.exact(), response, now), Self::ValkeySemantic { cache, scope, .. } => { - cache.store(&Self::semantic(request, scope), response, now) + cache.store(&request.scoped_semantic(scope), response, now) } - Self::Gcs(cache) => cache.store(&Self::exact(request), response, now), - Self::Disk(cache) => cache.store(&Self::exact(request), response, now), - Self::AzureBlob(cache) => cache.store(&Self::exact(request), response, now), + Self::RedisSemantic { cache, .. } => cache.store(&request.semantic(), response, now), + Self::QdrantSemantic(cache) => cache.store(&request.semantic(), response, now), } } @@ -406,26 +366,9 @@ impl NativeResponseCache { now: Duration, ) -> Result { match self { - Self::Memory(cache) => { - let requests = requests.iter().map(Self::exact).collect::>(); - cache.lookup_batch(&requests, now) - } - Self::Redis { cache, .. } => { - let requests = requests.iter().map(Self::exact).collect::>(); - cache.lookup_batch(&requests, now) - } - Self::S3(cache) => { - cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - } - Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), - Self::Gcs(cache) => { - cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - } - Self::Disk(cache) => { - cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - } - Self::AzureBlob(cache) => { - cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) + Self::Exact(service) => service.cache.lookup_batch(&exact_requests(requests), now), + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { + Err(Error::UnsupportedOperation) } } } @@ -436,17 +379,14 @@ impl NativeResponseCache { now: Duration, ) -> Result, Error> { match self { - Self::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await, - Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await, - Self::S3(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::Exact(service) => service.cache.async_lookup(&request.exact(), now).await, Self::ValkeySemantic { cache, scope, .. } => { cache - .async_lookup(&Self::semantic(request, scope), now) + .async_lookup(&request.scoped_semantic(scope), now) .await } - Self::Gcs(cache) => cache.async_lookup(&Self::exact(request), now).await, - Self::Disk(cache) => cache.async_lookup(&Self::exact(request), now).await, - Self::AzureBlob(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::RedisSemantic { cache, .. } => cache.async_lookup(&request.semantic(), now).await, + Self::QdrantSemantic(cache) => cache.async_lookup(&request.semantic(), now).await, } } @@ -456,31 +396,17 @@ impl NativeResponseCache { request: NativeRequest, ) -> PyResult> { match self { - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => { + Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); litellm_host_python::run_async( py, - async move { service.async_lookup(&request, super::request::now()).await }, + async move { service.async_lookup(&request, now()).await }, super::cache_error, ) } - Self::ValkeySemantic { - cache, - embedder, - scope, - } => drive_semantic( - py, - SemanticEmbedExecution::lookup( - Arc::clone(cache.backend_arc()), - embedder.clone(), - Self::semantic(&request, scope), - ), - ), + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + self.python_semantic(py, SemanticOperation::Lookup(request)) + } } } @@ -491,51 +417,29 @@ impl NativeResponseCache { now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await - } - Self::Redis { - cache, - buffer: None, - } => { - cache - .async_store(&Self::exact(request), response, now) - .await - } - Self::Redis { - cache, - buffer: Some(buffer), - } => { - buffer - .async_store(cache, &Self::exact(request), response, now) - .await - } - Self::S3(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await - } + Self::Exact(service) => match &service.buffer { + None => { + service + .cache + .async_store(&request.exact(), response, now) + .await + } + Some(buffer) => { + buffer + .async_store(service.cache.as_ref(), &request.exact(), response, now) + .await + } + }, Self::ValkeySemantic { cache, scope, .. } => { cache - .async_store(&Self::semantic(request, scope), response, now) + .async_store(&request.scoped_semantic(scope), response, now) .await } - Self::Gcs(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await + Self::RedisSemantic { cache, .. } => { + cache.async_store(&request.semantic(), response, now).await } - Self::Disk(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await - } - Self::AzureBlob(cache) => { - cache - .async_store(&Self::exact(request), response, now) - .await + Self::QdrantSemantic(cache) => { + cache.async_store(&request.semantic(), response, now).await } } } @@ -547,36 +451,17 @@ impl NativeResponseCache { response: Value, ) -> PyResult> { match self { - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => { + Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); litellm_host_python::run_async( py, - async move { - service - .async_store(&request, response, super::request::now()) - .await - }, + async move { service.async_store(&request, response, now()).await }, super::cache_error, ) } - Self::ValkeySemantic { - cache, - embedder, - scope, - } => drive_semantic( - py, - SemanticEmbedExecution::store( - Arc::clone(cache.backend_arc()), - embedder.clone(), - Self::semantic(&request, scope), - response, - ), - ), + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + self.python_semantic(py, SemanticOperation::Store(request, response)) + } } } @@ -586,34 +471,14 @@ impl NativeResponseCache { now: Duration, ) -> Result { match self { - Self::Memory(cache) => { - let requests = requests.iter().map(Self::exact).collect::>(); - cache.async_lookup_batch(&requests, now).await - } - Self::Redis { cache, .. } => { - let requests = requests.iter().map(Self::exact).collect::>(); - cache.async_lookup_batch(&requests, now).await - } - Self::S3(cache) => { - cache - .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) + Self::Exact(service) => { + service + .cache + .async_lookup_batch(&exact_requests(requests), now) .await } - Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), - Self::Gcs(cache) => { - cache - .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - .await - } - Self::Disk(cache) => { - cache - .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - .await - } - Self::AzureBlob(cache) => { - cache - .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) - .await + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { + Err(Error::UnsupportedOperation) } } } @@ -624,52 +489,25 @@ impl NativeResponseCache { now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => { + Self::Exact(service) => { let entries = entries .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) + .map(|(request, value)| (request.exact(), value)) .collect(); - cache.async_store_batch(entries, now).await - } - Self::Redis { cache, .. } => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) - .collect(); - cache.async_store_batch(entries, now).await - } - Self::S3(cache) => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) - .collect(); - cache.async_store_batch(entries, now).await + service.cache.async_store_batch(entries, now).await } Self::ValkeySemantic { cache, scope, .. } => { let entries = entries .into_iter() - .map(|(request, value)| (Self::semantic(&request, scope), value)) + .map(|(request, value)| (request.scoped_semantic(scope), value)) .collect(); cache.async_store_batch(entries, now).await } - Self::Gcs(cache) => { + Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation), + Self::QdrantSemantic(cache) => { let entries = entries .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) - .collect(); - cache.async_store_batch(entries, now).await - } - Self::Disk(cache) => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) - .collect(); - cache.async_store_batch(entries, now).await - } - Self::AzureBlob(cache) => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::exact(&request), value)) + .map(|(request, value)| (request.semantic(), value)) .collect(); cache.async_store_batch(entries, now).await } @@ -682,154 +520,54 @@ impl NativeResponseCache { entries: Vec<(NativeRequest, Value)>, ) -> PyResult> { match self { - Self::Memory(_) - | Self::Redis { .. } - | Self::S3(_) - | Self::Disk(_) - | Self::AzureBlob(_) - | Self::Gcs(_) => { + Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); litellm_host_python::run_async( py, - async move { - service - .async_store_batch(entries, super::request::now()) - .await - }, + async move { service.async_store_batch(entries, now()).await }, super::cache_error, ) } - Self::ValkeySemantic { - cache, - embedder, - scope, - } => { - let (requests, responses): (Vec<_>, Vec<_>) = entries - .into_iter() - .map(|(request, response)| (Self::semantic(&request, scope), response)) - .unzip(); - drive_semantic( - py, - SemanticEmbedExecution::store_batch( - Arc::clone(cache.backend_arc()), - embedder.clone(), - requests, - responses, - ), - ) + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + self.python_semantic(py, SemanticOperation::StoreBatch(entries.into())) } } } pub async fn async_flush(&self) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_flush().await, - Self::Redis { cache, buffer } => { - if let Some(buffer) = buffer { + Self::Exact(service) => { + if let Some(buffer) = &service.buffer { buffer.clear()?; } - cache.async_flush().await + service.cache.async_flush().await + } + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { + Err(Error::UnsupportedOperation) } - Self::S3(cache) => cache.async_flush().await, - Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), - Self::Gcs(cache) => cache.async_flush().await, - Self::Disk(cache) => cache.async_flush().await, - Self::AzureBlob(cache) => cache.async_flush().await, } } pub async fn test_connection(&self) -> Result { match self { - Self::Memory(cache) => cache.test_connection().await, - Self::Redis { cache, .. } => cache.test_connection().await, - Self::S3(cache) => cache.test_connection().await, + Self::Exact(service) => service.cache.test_connection().await, Self::ValkeySemantic { cache, .. } => cache.test_connection().await, - Self::Gcs(cache) => cache.test_connection().await, - Self::Disk(cache) => cache.test_connection().await, - Self::AzureBlob(cache) => cache.test_connection().await, + Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { + Err(Error::UnsupportedOperation) + } } } - pub fn gcs_backend(&self) -> Option<&GcsCache> { + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { match self { - Self::Gcs(cache) => Some(cache.backend()), - _ => None, + Self::ValkeySemantic { embedder, .. } | Self::RedisSemantic { embedder, .. } => { + embedder.traverse(visit) + } + Self::Exact(_) | Self::QdrantSemantic(_) => Ok(()), } } } -#[cfg(test)] -mod tests { - use litellm_cache_response::{CacheControls, CacheKeyInput, cache_key}; - use serde_json::json; - use sha2::{Digest, Sha256}; - - use super::*; - - fn native_request(key: CacheKeyInput, metadata: Value) -> NativeRequest { - NativeRequest { - key, - controls: CacheControls::default(), - ttl: None, - max_age: None, - messages: Some(json!([{"role": "user", "content": "prompt"}])), - input: None, - metadata: Some(metadata), - litellm_metadata: None, - litellm_params: None, - } - } - - #[test] - fn semantic_key_matches_python_scope_material() { - let key = CacheKeyInput { - fields: vec![ - CacheKeyField { - name: "model".to_owned(), - value: Some("gpt-4.1".to_owned()), - api_parameter: true, - internal_parameter: false, - }, - CacheKeyField { - name: "messages".to_owned(), - value: Some("prompt".to_owned()), - api_parameter: true, - internal_parameter: false, - }, - ], - ..Default::default() - }; - let request = native_request( - key, - json!({"user_api_key": "k1", "user_api_key_team_id": null}), - ); - let expected = format!("{:x}", Sha256::digest(b"model: gpt-4.1user_api_key: k1")); - assert_eq!(cache_key(&semantic_key(&request, "key")), expected); - - let end_user_request = native_request( - request.key.clone(), - json!({"user_api_key": "k1", "user_api_key_end_user_id": "u1"}), - ); - let expected = format!( - "{:x}", - Sha256::digest(b"model: gpt-4.1user_api_key: k1user_api_key_end_user_id: u1") - ); - assert_eq!( - cache_key(&semantic_key(&end_user_request, "end_user")), - expected - ); - - let preset_request = native_request( - CacheKeyInput { - preset: Some("preset-key".to_owned()), - ..Default::default() - }, - json!({"user_api_key": "k1"}), - ); - assert_eq!( - semantic_key(&preset_request, "end_user").preset.as_deref(), - Some("preset-key") - ); - assert!(semantic_key(&preset_request, "end_user").fields.is_empty()); - } +fn exact_requests(requests: &[NativeRequest]) -> Vec { + requests.iter().map(NativeRequest::exact).collect() } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 036951891a1..627bf9f1840 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,7 +1,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::ExactCacheContext; -use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; +use litellm_cache::{ExactCacheContext, SemanticCacheContext}; +use litellm_cache_response::{CacheControls, CacheKeyField, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; use serde::Deserialize; @@ -19,8 +19,10 @@ struct RequestInput { metadata: Option, litellm_metadata: Option, litellm_params: Option, + scope: Option, } +#[derive(Clone)] pub(super) struct NativeRequest { pub(super) key: CacheKeyInput, pub(super) controls: CacheControls, @@ -31,6 +33,100 @@ pub(super) struct NativeRequest { pub(super) metadata: Option, pub(super) litellm_metadata: Option, pub(super) litellm_params: Option, + pub(super) scope: Option, +} + +impl NativeRequest { + pub(super) fn exact(&self) -> ResponseCacheRequest { + ResponseCacheRequest { + key: self.key.clone(), + controls: self.controls, + context: ExactCacheContext { ttl: self.ttl }, + max_age: self.max_age, + } + } + + /// The request as a semantic backend that keys on the caller's scope sees it. + pub(super) fn semantic(&self) -> ResponseCacheRequest { + self.semantic_with(self.key.clone(), self.scope.clone()) + } + + /// The request keyed the way Python's Valkey semantic cache keys it: prompt fields drop out + /// and the tenant identifiers for `scope` join the key. + pub(super) fn scoped_semantic( + &self, + scope: &str, + ) -> ResponseCacheRequest { + self.semantic_with(semantic_key(self, scope), Some(scope.to_owned())) + } + + fn semantic_with( + &self, + key: CacheKeyInput, + scope: Option, + ) -> ResponseCacheRequest { + ResponseCacheRequest { + key, + controls: self.controls, + context: SemanticCacheContext { + input: self.input.clone(), + messages: self.messages.clone(), + metadata: self.metadata.clone(), + scope, + ttl: self.ttl, + }, + max_age: self.max_age, + } + } +} + +fn semantic_key(request: &NativeRequest, scope: &str) -> CacheKeyInput { + let mut key = request.key.clone(); + if key.preset.is_some() { + return key; + } + key.fields + .retain(|field| !matches!(field.name.as_str(), "messages" | "prompt" | "input")); + const TENANT: [&str; 3] = [ + "user_api_key", + "user_api_key_team_id", + "user_api_key_org_id", + ]; + let end_user = (scope == "end_user").then_some("user_api_key_end_user_id"); + for name in TENANT.into_iter().chain(end_user) { + let sources = [ + request.metadata.as_ref(), + request.litellm_metadata.as_ref(), + request + .litellm_params + .as_ref() + .and_then(|params| params.get("metadata")), + request + .litellm_params + .as_ref() + .and_then(|params| params.get("litellm_metadata")), + ]; + let Some(value) = sources.into_iter().flatten().find_map(|source| { + source + .as_object() + .and_then(|values| values.get(name)) + .filter(|value| !value.is_null()) + }) else { + continue; + }; + let value = match value { + Value::Null => continue, + Value::String(text) => text.clone(), + other => other.to_string(), + }; + key.fields.push(CacheKeyField { + name: name.to_owned(), + value: Some(value), + api_parameter: true, + internal_parameter: false, + }); + } + key } pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { @@ -52,6 +148,7 @@ fn request_input(input: RequestInput) -> PyResult { metadata: input.metadata, litellm_metadata: input.litellm_metadata, litellm_params: input.litellm_params, + scope: input.scope, }) } @@ -72,3 +169,90 @@ pub(super) fn now() -> Duration { .duration_since(UNIX_EPOCH) .unwrap_or_default() } + +#[cfg(test)] +mod tests { + use litellm_cache_response::{CacheControls, CacheKeyInput, cache_key}; + use serde_json::json; + use sha2::{Digest, Sha256}; + + use super::*; + + fn native_request(key: CacheKeyInput, metadata: Value) -> NativeRequest { + NativeRequest { + key, + controls: CacheControls::default(), + ttl: None, + max_age: None, + messages: Some(json!([{"role": "user", "content": "prompt"}])), + input: None, + metadata: Some(metadata), + litellm_metadata: None, + litellm_params: None, + scope: None, + } + } + + #[test] + fn semantic_key_matches_python_scope_material() { + let key = CacheKeyInput { + fields: vec![ + CacheKeyField { + name: "model".to_owned(), + value: Some("gpt-4.1".to_owned()), + api_parameter: true, + internal_parameter: false, + }, + CacheKeyField { + name: "messages".to_owned(), + value: Some("prompt".to_owned()), + api_parameter: true, + internal_parameter: false, + }, + ], + ..Default::default() + }; + let request = native_request( + key, + json!({"user_api_key": "k1", "user_api_key_team_id": null}), + ); + let expected = format!("{:x}", Sha256::digest(b"model: gpt-4.1user_api_key: k1")); + assert_eq!(cache_key(&semantic_key(&request, "key")), expected); + assert_eq!(cache_key(&request.scoped_semantic("key").key), expected); + + let end_user_request = native_request( + request.key.clone(), + json!({"user_api_key": "k1", "user_api_key_end_user_id": "u1"}), + ); + let expected = format!( + "{:x}", + Sha256::digest(b"model: gpt-4.1user_api_key: k1user_api_key_end_user_id: u1") + ); + assert_eq!( + cache_key(&semantic_key(&end_user_request, "end_user")), + expected + ); + + let preset_request = native_request( + CacheKeyInput { + preset: Some("preset-key".to_owned()), + ..Default::default() + }, + json!({"user_api_key": "k1"}), + ); + assert_eq!( + semantic_key(&preset_request, "end_user").preset.as_deref(), + Some("preset-key") + ); + assert!(semantic_key(&preset_request, "end_user").fields.is_empty()); + assert_eq!(preset_request.semantic().context.scope, None); + assert_eq!( + preset_request + .scoped_semantic("end_user") + .context + .scope + .as_deref(), + Some("end_user") + ); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs new file mode 100644 index 00000000000..9f4d18d45cd --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -0,0 +1,189 @@ +use std::{collections::VecDeque, time::Duration}; + +use litellm_cache::Error; +use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyException, PyRuntimeError}, + prelude::*, +}; +use serde_json::Value; + +use super::{ + cache_error, + embedder::{PythonEmbedder, with_prepared_embedding}, + native::NativeResponseCache, + request::{NativeRequest, now}, +}; + +pub(super) enum SemanticOperation { + Lookup(NativeRequest), + Store(NativeRequest, Value), + StoreBatch(VecDeque<(NativeRequest, Value)>), +} + +/// What an exception from the Python embedder means for the operation. +#[derive(Clone, Copy)] +pub(super) enum EmbeddingFailure { + /// Raise the Python exception unchanged. + Propagate, + /// Treat the embedding as unavailable and let the backend report that. + Unavailable, +} + +enum Phase { + Start, + AwaitingEmbedding, + AwaitingBackend, +} + +/// Runs a semantic cache operation whose embedding comes from Python: await the Python +/// embedder in the caller's event loop, seed the native backend with the vector, await the +/// backend, and repeat for each entry of a batch. +pub(super) struct SemanticExecution { + service: NativeResponseCache, + embedder: PythonEmbedder, + failure: EmbeddingFailure, + operation: SemanticOperation, + pending: Option<(NativeRequest, Option)>, + phase: Phase, + now: Duration, +} + +impl SemanticExecution { + pub(super) fn new( + service: NativeResponseCache, + embedder: PythonEmbedder, + failure: EmbeddingFailure, + operation: SemanticOperation, + ) -> Self { + Self { + service, + embedder, + failure, + operation, + pending: None, + phase: Phase::Start, + now: now(), + } + } + + /// Takes the next entry of the operation; `None` once a batch is exhausted. + fn next_pending(&mut self) -> Option<(NativeRequest, Option)> { + match &mut self.operation { + SemanticOperation::Lookup(request) => Some((request.clone(), None)), + SemanticOperation::Store(request, response) => { + Some((request.clone(), Some(std::mem::take(response)))) + } + SemanticOperation::StoreBatch(queue) => queue + .pop_front() + .map(|(request, response)| (request, Some(response))), + } + } + + fn start(&mut self, py: Python<'_>) -> PyResult { + let Some(pending) = self.next_pending() else { + return Ok(ExecutionStep::Return(py.None())); + }; + let (request, response) = &pending; + let enabled = match response { + None => request.controls.reads(), + Some(_) => request.controls.writes(), + }; + let input = enabled + .then(|| self.service.embedding_input(request)) + .flatten(); + self.pending = Some(pending); + let Some(input) = input else { + return self.backend_step(py, Err(Error::Unavailable)); + }; + let awaitable = + self.embedder + .async_embedding(py, &input.prompt, input.metadata.as_ref())?; + self.phase = Phase::AwaitingEmbedding; + Ok(ExecutionStep::Await(awaitable)) + } + + fn embedded(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { + let seed = match result { + Ok(vector) => { + PythonEmbedder::extract(vector.into_bound(py)).map_err(|_| Error::Unavailable) + } + Err(error) => match self.failure { + EmbeddingFailure::Propagate => return Err(error), + EmbeddingFailure::Unavailable if error.is_instance_of::(py) => { + Err(Error::Unavailable) + } + EmbeddingFailure::Unavailable => return Err(error), + }, + }; + self.backend_step(py, seed) + } + + fn backend_step( + &mut self, + py: Python<'_>, + seed: Result, Error>, + ) -> PyResult { + self.phase = Phase::AwaitingBackend; + let (request, response) = self.pending.take().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution resumed without a pending operation") + })?; + let service = self.service.clone(); + let now = self.now; + let future = async move { + match response { + None => service.async_lookup(&request, now).await, + Some(response) => service + .async_store(&request, response, now) + .await + .map(|_| None), + } + }; + let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?; + Ok(ExecutionStep::Await(awaitable.unbind())) + } + + fn resume_py( + &mut self, + py: Python<'_>, + result: Option>>, + ) -> PyResult { + match (&self.phase, result) { + (Phase::Start, None) => self.start(py), + (Phase::AwaitingEmbedding, Some(result)) => self.embedded(py, result), + (Phase::AwaitingBackend, Some(Err(error))) => Err(error), + (Phase::AwaitingBackend, Some(Ok(value))) => { + let more = matches!( + &self.operation, + SemanticOperation::StoreBatch(queue) if !queue.is_empty() + ); + if more { + self.phase = Phase::Start; + return self.start(py); + } + Ok(ExecutionStep::Return(value)) + } + _ => Err(PyRuntimeError::new_err( + "invalid semantic cache execution state", + )), + } + } +} + +impl ExecutionBody for SemanticExecution { + fn resume(&mut self, result: Option>>) -> PyResult { + Python::attach(|py| self.resume_py(py, result)) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.embedder.traverse(visit) + } +} + +pub(super) fn drive(py: Python<'_>, body: SemanticExecution) -> PyResult> { + let execution = Py::new(py, Execution::new(body))?; + py.import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs deleted file mode 100644 index 24caf3374d6..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs +++ /dev/null @@ -1,249 +0,0 @@ -use std::{sync::Arc, time::Duration}; - -use litellm_cache::SemanticCacheContext; -use litellm_cache_response::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; -use litellm_cache_valkey_semantic::{PreparedEmbedding, ValkeySemanticCache, prompt_from_context}; -use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; -use serde_json::Value; - -use super::{cache_error, embedder::PythonEmbedder}; - -pub(super) enum Op { - Lookup, - Store(Value), - StoreBatch(Vec), -} - -#[derive(Clone, Copy)] -enum State { - Start, - AwaitingEmbedding, - AwaitingStorage, - Done, -} - -pub(super) struct SemanticEmbedExecution { - backend: Arc>, - embedder: PythonEmbedder, - requests: Vec>, - op: Op, - now: Option, - prepared: Vec>>, - index: usize, - state: State, -} - -impl SemanticEmbedExecution { - pub(super) fn lookup( - backend: Arc>, - embedder: PythonEmbedder, - request: ResponseCacheRequest, - ) -> Self { - Self { - backend, - embedder, - requests: vec![request], - op: Op::Lookup, - now: None, - prepared: vec![None], - index: 0, - state: State::Start, - } - } - - pub(super) fn store( - backend: Arc>, - embedder: PythonEmbedder, - request: ResponseCacheRequest, - response: Value, - ) -> Self { - Self { - backend, - embedder, - requests: vec![request], - op: Op::Store(response), - now: None, - prepared: vec![None], - index: 0, - state: State::Start, - } - } - - pub(super) fn store_batch( - backend: Arc>, - embedder: PythonEmbedder, - requests: Vec>, - responses: Vec, - ) -> Self { - Self { - backend, - embedder, - prepared: vec![None; requests.len()], - requests, - op: Op::StoreBatch(responses), - now: None, - index: 0, - state: State::Start, - } - } - - fn start(&mut self, py: Python<'_>) -> PyResult { - if self.now.is_none() { - self.now = Some(super::request::now()); - } - while self.index < self.requests.len() { - let request = &self.requests[self.index]; - let enabled = match &self.op { - Op::Lookup => request.controls.reads(), - Op::Store(_) | Op::StoreBatch(_) => request.controls.writes(), - }; - if !enabled { - self.index += 1; - continue; - } - let Some(prompt) = prompt_from_context(&request.context) else { - self.index += 1; - continue; - }; - let metadata = request.context.metadata.clone(); - let awaitable = self - .embedder - .async_embed_awaitable(py, &prompt, &metadata)?; - self.state = State::AwaitingEmbedding; - return Ok(ExecutionStep::Await(awaitable.unbind())); - } - self.state = State::AwaitingStorage; - self.storage_step(py) - } - - fn storage_step(&self, py: Python<'_>) -> PyResult { - let requests = self.requests.clone(); - let prepared = self.prepared.clone(); - let backend = Arc::clone(&self.backend); - let now = self - .now - .ok_or_else(|| PyRuntimeError::new_err("semantic cache timestamp is unavailable"))?; - let awaitable = match &self.op { - Op::Lookup => { - let Some(request) = requests.into_iter().next() else { - return Err(PyRuntimeError::new_err( - "semantic lookup requires one request", - )); - }; - match prepared.into_iter().next().flatten() { - Some(values) => { - let backend = backend.with_embedder(PreparedEmbedding(values)); - let cache = Arc::new(ResponseCache::new(Arc::new(backend))); - run_async( - py, - async move { cache.async_lookup(&request, now).await }, - cache_error, - )? - } - None => { - let cache = Arc::new(ResponseCache::new(backend)); - run_async( - py, - async move { cache.async_lookup(&request, now).await }, - cache_error, - )? - } - } - } - Op::Store(response) => { - let Some(request) = requests.into_iter().next() else { - return Err(PyRuntimeError::new_err( - "semantic store requires one request", - )); - }; - let response = response.clone(); - match prepared.into_iter().next().flatten() { - Some(values) => { - let backend = backend.with_embedder(PreparedEmbedding(values)); - let cache = Arc::new(ResponseCache::new(Arc::new(backend))); - run_async( - py, - async move { cache.async_store(&request, response, now).await }, - cache_error, - )? - } - None => { - let cache = Arc::new(ResponseCache::new(backend)); - run_async( - py, - async move { cache.async_store(&request, response, now).await }, - cache_error, - )? - } - } - } - Op::StoreBatch(responses) => { - let responses = responses.clone(); - run_async( - py, - async move { - for ((request, response), prepared) in - requests.into_iter().zip(responses).zip(prepared) - { - let Some(values) = prepared else { - continue; - }; - let backend = backend.with_embedder(PreparedEmbedding(values)); - let cache = ResponseCache::new(Arc::new(backend)); - cache.async_store(&request, response, now).await?; - } - Ok(()) - }, - cache_error, - )? - } - }; - Ok(ExecutionStep::Await(awaitable.unbind())) - } - - fn resume_py( - &mut self, - py: Python<'_>, - result: Option>>, - ) -> PyResult { - match (self.state, result) { - (State::Start, None) => self.start(py), - (State::AwaitingEmbedding, Some(Ok(value))) => { - let values = value.bind(py).extract::>()?; - self.prepared[self.index] = - Some(values.into_iter().map(|value| value as f32).collect()); - self.index += 1; - self.start(py) - } - (State::AwaitingStorage, Some(Ok(value))) => { - self.state = State::Done; - Ok(ExecutionStep::Return(value)) - } - (_, Some(Err(error))) => Err(error), - _ => Err(PyRuntimeError::new_err( - "invalid semantic cache execution state", - )), - } - } -} - -impl ExecutionBody for SemanticEmbedExecution { - fn resume(&mut self, result: Option>>) -> PyResult { - Python::attach(|py| self.resume_py(py, result)) - } - - fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.embedder.traverse(visit) - } -} - -pub(super) fn drive_semantic<'py>( - py: Python<'py>, - body: SemanticEmbedExecution, -) -> PyResult> { - let execution = Py::new(py, Execution::new(body))?; - py.import("litellm.rust_bridge.lifecycle")? - .getattr("drive")? - .call1((execution,)) -} diff --git a/litellm-rust/crates/python-bridge/src/coercion.rs b/litellm-rust/crates/python-bridge/src/coercion.rs index bb5b8b2d454..1b0b073b3d1 100644 --- a/litellm-rust/crates/python-bridge/src/coercion.rs +++ b/litellm-rust/crates/python-bridge/src/coercion.rs @@ -1,7 +1,4 @@ -use std::collections::BTreeSet; - use litellm_core_utils::serde_compat::parse_str_bool; -use litellm_http::SslVerify; use pyo3::{ exceptions::{PyAttributeError, PyRuntimeError, PyValueError}, prelude::*, @@ -33,36 +30,52 @@ impl From for PyErr { } } -pub(crate) struct Truthy(pub bool); -pub(crate) struct ExactTrue(pub bool); -pub(crate) struct StrBool(pub Option); -pub(crate) struct OptionalStrictString(pub Option); -pub(crate) struct FalsyOptionalString(pub Option); -pub(crate) struct TuningString(pub Option); -pub(crate) struct StringCollection(pub Vec); -pub(crate) struct SslVerifyInput(pub Option); +pub(crate) struct FieldSpec { + name: &'static str, + decode: fn(&Field<'_>) -> Result, +} + +impl FieldSpec { + pub(crate) const fn new( + name: &'static str, + decode: fn(&Field<'_>) -> Result, + ) -> Self { + Self { name, decode } + } + + pub(crate) fn read( + &self, + snapshot: &Bound<'_, PyAny>, + group: &'static str, + ) -> Result { + (self.decode)(&Field::read(snapshot, group, self.name)?) + } +} pub(crate) struct Field<'py> { - path: &'static str, + group: &'static str, + name: &'static str, value: Bound<'py, PyAny>, } impl<'py> Field<'py> { - pub(crate) fn new(path: &'static str, value: Bound<'py, PyAny>) -> Self { - Self { path, value } + pub(crate) fn new(group: &'static str, name: &'static str, value: Bound<'py, PyAny>) -> Self { + Self { group, name, value } } + /// Reads `snapshot.`, distinguishing a field the accessor never declared from a + /// descriptor that raised `AttributeError`. pub(crate) fn read( snapshot: &Bound<'py, PyAny>, - path: &'static str, + group: &'static str, + name: &'static str, ) -> Result { - let name = path.rsplit('.').next().unwrap_or(path); match snapshot.getattr(name) { - Ok(value) => Ok(Self::new(path, value)), + Ok(value) => Ok(Self::new(group, name, value)), Err(error) if error.is_instance_of::(snapshot.py()) => { match Self::missing_field(snapshot, name) { Ok(true) => Err(ProjectionError::InternalSchemaFailure(format!( - "{path}: missing snapshot field" + "{group}.{name}: missing snapshot field" ))), _ => Err(error.into()), } @@ -84,27 +97,49 @@ impl<'py> Field<'py> { && getter.is(object.getattr("__getattribute__")?)) } - fn expected(&self, expected: &'static str) -> Result { + pub(crate) fn path(&self) -> String { + format!("{}.{}", self.group, self.name) + } + + /// A member of this field's collection, reported under the same path. + pub(crate) fn member(&self, value: Bound<'py, PyAny>) -> Self { + Self::new(self.group, self.name, value) + } + + pub(crate) fn expected(&self, expected: &str) -> Result { Ok(format!( "{}: expected {expected}, got {}", - self.path, + self.path(), self.value.get_type().name()? )) } - fn invalid(&self, expected: &'static str) -> ProjectionError { + pub(crate) fn invalid(&self, expected: &str) -> ProjectionError { match self.expected(expected) { Ok(message) => ProjectionError::InvalidConfiguration(message), Err(error) => error, } } - pub(crate) fn truthy(&self) -> Result { - Ok(Truthy(self.value.is_truthy()?)) + pub(crate) fn value(&self) -> &Bound<'py, PyAny> { + &self.value } - pub(crate) fn exact_true(&self) -> ExactTrue { - ExactTrue(self.value.is(PyBool::new(self.value.py(), true))) + pub(crate) fn truthy(&self) -> Result { + Ok(self.value.is_truthy()?) + } + + pub(crate) fn exact_true(&self) -> bool { + self.value.is(PyBool::new(self.value.py(), true)) + } + + pub(crate) fn schema_bool(&self) -> Result { + if !self.value.is_instance_of::() { + return Err(ProjectionError::InternalSchemaFailure( + self.expected("a Boolean")?, + )); + } + Ok(self.exact_true()) } pub(crate) fn strict_string(&self) -> Result { @@ -124,108 +159,366 @@ impl<'py> Field<'py> { self.strict_string() } - pub(crate) fn schema_bool(&self) -> Result { - if !self.value.is_instance_of::() { - return Err(ProjectionError::InternalSchemaFailure( - self.expected("a Boolean")?, - )); - } - Ok(self.exact_true().0) - } - - pub(crate) fn str_bool(&self) -> Result { + pub(crate) fn str_bool(&self) -> Result, ProjectionError> { if self.value.is_none() { - return Ok(StrBool(None)); + return Ok(None); } - Ok(StrBool(parse_str_bool(&self.strict_string()?))) + Ok(parse_str_bool(&self.strict_string()?)) } - pub(crate) fn optional_strict_string(&self) -> Result { + pub(crate) fn optional_strict_string(&self) -> Result, ProjectionError> { if self.value.is_none() { - return Ok(OptionalStrictString(None)); + return Ok(None); } - self.strict_string().map(Some).map(OptionalStrictString) + self.strict_string().map(Some) } - pub(crate) fn falsy_optional_string(&self) -> Result { - if !self.truthy()?.0 { - return Ok(FalsyOptionalString(None)); + pub(crate) fn falsy_optional_string(&self) -> Result, ProjectionError> { + if !self.truthy()? { + return Ok(None); } - self.strict_string().map(Some).map(FalsyOptionalString) + self.strict_string().map(Some) } - pub(crate) fn tuning_string(&self) -> Result { - if !self.truthy()?.0 || !self.value.is_instance_of::() { - return Ok(TuningString(None)); + pub(crate) fn tuning_string(&self) -> Result, ProjectionError> { + if !self.truthy()? || !self.value.is_instance_of::() { + return Ok(None); } - self.strict_string().map(Some).map(TuningString) + self.strict_string().map(Some) } - pub(crate) fn string_collection(&self) -> Result { - if !self.truthy()?.0 { - return Ok(StringCollection(Vec::new())); + pub(crate) fn string_collection(&self) -> Result, ProjectionError> { + if !self.truthy()? { + return Ok(Vec::new()); } if self.value.is_instance_of::() { - return self - .strict_string() - .map(|value| StringCollection(vec![value])); + return self.strict_string().map(|value| vec![value]); } - let values = self - .value + self.value .try_iter()? .filter_map(|item| { let member = match item { - Ok(value) => Self::new(self.path, value), + Ok(value) => self.member(value), Err(error) => return Some(Err(error.into())), }; match member.truthy() { - Ok(Truthy(false)) => None, - Ok(Truthy(true)) => Some(member.strict_string()), + Ok(false) => None, + Ok(true) => Some(member.strict_string()), Err(error) => Some(Err(error)), } }) - .collect::, ProjectionError>>()?; - Ok(StringCollection(values)) + .collect() } - pub(crate) fn host_collection(&self) -> Result { - let values = self - .string_collection()? - .0 - .into_iter() - .map(|host| litellm_http::media::normalize_host(&host)) - .collect::>(); - Ok(StringCollection(values.into_iter().collect())) - } - - pub(crate) fn ssl_verify(&self) -> Result { + pub(crate) fn optional_string_collection( + &self, + ) -> Result>, ProjectionError> { if self.value.is_none() { - return Ok(SslVerifyInput(None)); + return Ok(None); } - if self.value.is_instance_of::() { - return Ok(SslVerifyInput(Some(if self.exact_true().0 { - SslVerify::Enabled - } else { - SslVerify::Disabled - }))); - } - if self.value.is_instance_of::() { - let parsed = match self.str_bool()?.0 { - Some(true) => SslVerify::Enabled, - Some(false) => SslVerify::Disabled, - None => SslVerify::CaBundle(self.strict_string()?.into()), - }; - return Ok(SslVerifyInput(Some(parsed))); - } - let context = self.value.py().import("ssl")?.getattr("SSLContext")?; - if self.value.is_instance(&context)? { - return Err(ProjectionError::UnsupportedLiveObject(self.expected( - "a Boolean, Boolean string, CA path, or None; live SSLContext is unsupported", - )?)); - } - Err(self.invalid("a Boolean, Boolean string, CA path, or None")) + self.string_collection().map(Some) + } + + pub(crate) fn python_binding(&self) -> Option> { + (!self.value.is_none()).then(|| self.value.clone().unbind()) } } #[cfg(test)] -mod tests; +mod tests { + use std::ffi::CString; + + use pyo3::{ + exceptions::{PyLookupError, PyRuntimeError, PyValueError}, + types::PyDict, + }; + use rstest::rstest; + + use super::*; + + fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> { + py.eval(&CString::new(source).unwrap(), None, None).unwrap() + } + + #[rstest] + #[case("None", false, false)] + #[case("False", false, false)] + #[case("True", true, true)] + #[case("0", false, false)] + #[case("1", true, false)] + #[case("''", false, false)] + #[case("'false'", true, false)] + #[case("[]", false, false)] + #[case("[0]", true, false)] + #[case("{}", false, false)] + #[case("object()", true, false)] + fn boolean_operations_have_distinct_python_semantics( + #[case] source: &str, + #[case] truth: bool, + #[case] exact: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let value = evaluate(py, source); + let field = Field::new("test", "flag", value.clone()); + assert_eq!(field.truthy().unwrap(), truth); + assert_eq!(field.exact_true(), exact); + assert_eq!( + field.truthy().unwrap(), + py.import("builtins") + .unwrap() + .getattr("bool") + .unwrap() + .call1((value,)) + .unwrap() + .extract::() + .unwrap() + ); + }); + } + + #[rstest] + #[case("None", Ok(None), Ok(None), Ok(None))] + #[case("''", Ok(Some("")), Ok(None), Ok(None))] + #[case( + "' value '", + Ok(Some(" value ")), + Ok(Some(" value ")), + Ok(Some(" value ")) + )] + #[case("[]", Err(()), Ok(None), Ok(None))] + #[case("0", Err(()), Ok(None), Ok(None))] + #[case("1", Err(()), Err(()), Ok(None))] + #[case("object()", Err(()), Err(()), Ok(None))] + fn string_operations_do_not_conflate_absence_and_type_checks( + #[case] source: &str, + #[case] strict: Result, ()>, + #[case] fallback: Result, ()>, + #[case] tuning: Result, ()>, + ) { + Python::initialize(); + Python::attach(|py| { + let field = Field::new("test", "string", evaluate(py, source)); + let owned = + |expected: Result, ()>| expected.map(|value| value.map(str::to_owned)); + assert_eq!( + field.optional_strict_string().map_err(|_| ()), + owned(strict) + ); + assert_eq!( + field.falsy_optional_string().map_err(|_| ()), + owned(fallback) + ); + assert_eq!(field.tuning_string().map_err(|_| ()), owned(tuning)); + }); + } + + #[rstest] + #[case("None", None)] + #[case("' True '", Some(true))] + #[case("' fAlSe '", Some(false))] + #[case("'yes'", None)] + #[case("'1'", None)] + #[case("'unknown'", None)] + fn string_boolean_tokens_remain_separate_from_truthiness( + #[case] source: &str, + #[case] expected: Option, + ) { + Python::initialize(); + Python::attach(|py| { + assert_eq!( + Field::new("test", "flag", evaluate(py, source)) + .str_bool() + .unwrap(), + expected + ); + }); + } + + #[test] + fn protocol_errors_preserve_exception_identity_traceback_cause_and_context() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +failure = LookupError('protocol failed') +cause = ValueError('cause') +context = RuntimeError('context') +def fail(): + try: + raise context + except RuntimeError: + raise failure from cause +class Bool: + def __bool__(self): return fail() +class Length: + def __len__(self): return fail() +class Iter: + def __iter__(self): return fail() +class Next: + def __iter__(self): return self + def __next__(self): return fail() +class Descriptor: + @property + def flag(self): return fail() +values = (Bool(), Length(), Iter(), Next(), [Bool()]) +descriptor = Descriptor() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let values = locals.get_item("values").unwrap().unwrap(); + for value in values.try_iter().unwrap() { + let error = Field::new("test", "flag", value.unwrap()) + .string_collection() + .err() + .unwrap(); + let error = PyErr::from(error); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + assert!(error.is_instance_of::(py)); + assert!(error.traceback(py).is_some()); + assert!( + error + .value(py) + .getattr("__cause__") + .unwrap() + .is(locals.get_item("cause").unwrap().unwrap()) + ); + assert!( + error + .value(py) + .getattr("__context__") + .unwrap() + .is(locals.get_item("context").unwrap().unwrap()) + ); + } + let error = Field::read( + &locals.get_item("descriptor").unwrap().unwrap(), + "test", + "flag", + ) + .err() + .unwrap(); + assert!( + PyErr::from(error) + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn identity_and_string_contents_do_not_invoke_unrelated_protocols() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +class Hostile: + def __bool__(self): raise AssertionError('bool called') + def __eq__(self, other): raise AssertionError('eq called') + def __str__(self): raise AssertionError('str called') +class Text(str): + def __str__(self): raise AssertionError('str called') + def strip(self): raise AssertionError('strip called') + def lower(self): raise AssertionError('lower called') +hostile = Hostile() +text = Text(' False ') +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let hostile = Field::new("test", "flag", locals.get_item("hostile").unwrap().unwrap()); + assert!(!hostile.exact_true()); + assert!(matches!( + hostile.strict_string(), + Err(ProjectionError::InvalidConfiguration(_)) + )); + let text = Field::new("test", "flag", locals.get_item("text").unwrap().unwrap()); + assert_eq!(text.strict_string().unwrap(), " False "); + assert_eq!(text.str_bool().unwrap(), Some(false)); + }); + } + + #[test] + fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +failure = AttributeError('descriptor failed') +class Snapshot: + @property + def flag(self): raise failure +snapshot = Snapshot() +class Dynamic: + def __getattr__(self, name): raise failure +class Intercepted: + def __getattribute__(self, name): raise failure +dynamic = Dynamic() +intercepted = Intercepted() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let snapshot = locals.get_item("snapshot").unwrap().unwrap(); + let descriptor = PyErr::from(Field::read(&snapshot, "test", "flag").err().unwrap()); + assert!( + descriptor + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + for name in ["dynamic", "intercepted"] { + let value = locals.get_item(name).unwrap().unwrap(); + let error = PyErr::from(Field::read(&value, "test", "flag").err().unwrap()); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + } + let missing = PyErr::from(Field::read(&snapshot, "test", "missing").err().unwrap()); + assert!(missing.is_instance_of::(py)); + assert!(missing.to_string().contains("test.missing")); + }); + } + + #[test] + fn configuration_errors_name_fields_without_exposing_values() { + Python::initialize(); + Python::attach(|py| { + for source in [ + "{'secret': 'do-not-print'}", + "['host.test', {'secret': 'do-not-print'}]", + ] { + let field = Field::new("test", "setting", evaluate(py, source)); + let error = PyErr::from(field.falsy_optional_string().err().unwrap()); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("test.setting")); + assert!(!error.to_string().contains("do-not-print")); + } + let hosts = Field::new( + "url_policy", + "user_url_allowed_hosts", + evaluate(py, "['host.test', 1]"), + ); + assert!(matches!( + hosts.string_collection(), + Err(ProjectionError::InvalidConfiguration(_)) + )); + assert!(matches!( + Field::new("test", "flag", evaluate(py, "1")).str_bool(), + Err(ProjectionError::InvalidConfiguration(_)) + )); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/coercion/tests.rs b/litellm-rust/crates/python-bridge/src/coercion/tests.rs deleted file mode 100644 index 5ed237c3c64..00000000000 --- a/litellm-rust/crates/python-bridge/src/coercion/tests.rs +++ /dev/null @@ -1,372 +0,0 @@ -use std::ffi::CString; - -use pyo3::{ - exceptions::{PyLookupError, PyRuntimeError, PyValueError}, - types::PyDict, -}; -use rstest::rstest; - -use super::*; - -fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> { - py.eval(&CString::new(source).unwrap(), None, None).unwrap() -} - -#[rstest] -#[case("None", false, false)] -#[case("False", false, false)] -#[case("True", true, true)] -#[case("0", false, false)] -#[case("1", true, false)] -#[case("''", false, false)] -#[case("'false'", true, false)] -#[case("[]", false, false)] -#[case("[0]", true, false)] -#[case("{}", false, false)] -#[case("object()", true, false)] -fn boolean_operations_have_distinct_python_semantics( - #[case] source: &str, - #[case] truth: bool, - #[case] exact: bool, -) { - Python::initialize(); - Python::attach(|py| { - let value = evaluate(py, source); - let field = Field::new("test.flag", value.clone()); - assert_eq!(field.truthy().unwrap().0, truth); - assert_eq!(field.exact_true().0, exact); - assert_eq!( - field.truthy().unwrap().0, - py.import("builtins") - .unwrap() - .getattr("bool") - .unwrap() - .call1((value,)) - .unwrap() - .extract::() - .unwrap() - ); - }); -} - -#[rstest] -#[case("None", Ok(None), Ok(None), Ok(None))] -#[case("''", Ok(Some("")), Ok(None), Ok(None))] -#[case( - "' value '", - Ok(Some(" value ")), - Ok(Some(" value ")), - Ok(Some(" value ")) -)] -#[case("[]", Err(()), Ok(None), Ok(None))] -#[case("0", Err(()), Ok(None), Ok(None))] -#[case("1", Err(()), Err(()), Ok(None))] -#[case("object()", Err(()), Err(()), Ok(None))] -fn string_operations_do_not_conflate_absence_and_type_checks( - #[case] source: &str, - #[case] strict: Result, ()>, - #[case] fallback: Result, ()>, - #[case] tuning: Result, ()>, -) { - Python::initialize(); - Python::attach(|py| { - let field = Field::new("test.string", evaluate(py, source)); - let owned = - |expected: Result, ()>| expected.map(|value| value.map(str::to_owned)); - assert_eq!( - field - .optional_strict_string() - .map(|value| value.0) - .map_err(|_| ()), - owned(strict) - ); - assert_eq!( - field - .falsy_optional_string() - .map(|value| value.0) - .map_err(|_| ()), - owned(fallback) - ); - assert_eq!( - field.tuning_string().map(|value| value.0).map_err(|_| ()), - owned(tuning) - ); - }); -} - -#[rstest] -#[case("None", None)] -#[case("' True '", Some(true))] -#[case("' fAlSe '", Some(false))] -#[case("'yes'", None)] -#[case("'1'", None)] -#[case("'unknown'", None)] -fn string_boolean_tokens_remain_separate_from_truthiness( - #[case] source: &str, - #[case] expected: Option, -) { - Python::initialize(); - Python::attach(|py| { - assert_eq!( - Field::new("test.flag", evaluate(py, source)) - .str_bool() - .unwrap() - .0, - expected - ); - }); -} - -#[rstest] -#[case("'EXAMPLE.TEST.'", vec!["example.test"])] -#[case("['B.test', '', None, 0, [], 'A.test.', 'b.test']", vec!["a.test", "b.test"])] -#[case("('B.test', 'a.test')", vec!["a.test", "b.test"])] -#[case("{'B.test', 'a.test'}", vec!["a.test", "b.test"])] -#[case("(host for host in ['B.test', 'a.test'])", vec!["a.test", "b.test"])] -#[case("None", vec![])] -#[case("False", vec![])] -fn host_collection_is_owned_normalized_and_deterministic( - #[case] source: &str, - #[case] expected: Vec<&str>, -) { - Python::initialize(); - Python::attach(|py| { - assert_eq!( - Field::new("url_policy.user_url_allowed_hosts", evaluate(py, source)) - .host_collection() - .unwrap() - .0, - expected - ); - }); -} - -#[test] -fn protocol_errors_preserve_exception_identity_traceback_cause_and_context() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - c" -failure = LookupError('protocol failed') -cause = ValueError('cause') -context = RuntimeError('context') -def fail(): - try: - raise context - except RuntimeError: - raise failure from cause -class Bool: - def __bool__(self): return fail() -class Length: - def __len__(self): return fail() -class Iter: - def __iter__(self): return fail() -class Next: - def __iter__(self): return self - def __next__(self): return fail() -class Descriptor: - @property - def flag(self): return fail() -values = (Bool(), Length(), Iter(), Next(), [Bool()]) -descriptor = Descriptor() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let values = locals.get_item("values").unwrap().unwrap(); - for value in values.try_iter().unwrap() { - let error = Field::new("test.flag", value.unwrap()) - .host_collection() - .err() - .unwrap(); - let error = PyErr::from(error); - assert!( - error - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - assert!(error.is_instance_of::(py)); - assert!(error.traceback(py).is_some()); - assert!( - error - .value(py) - .getattr("__cause__") - .unwrap() - .is(locals.get_item("cause").unwrap().unwrap()) - ); - assert!( - error - .value(py) - .getattr("__context__") - .unwrap() - .is(locals.get_item("context").unwrap().unwrap()) - ); - } - let error = Field::read( - &locals.get_item("descriptor").unwrap().unwrap(), - "test.flag", - ) - .err() - .unwrap(); - assert!( - PyErr::from(error) - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - }); -} - -#[test] -fn identity_and_string_contents_do_not_invoke_unrelated_protocols() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - c" -class Hostile: - def __bool__(self): raise AssertionError('bool called') - def __eq__(self, other): raise AssertionError('eq called') - def __str__(self): raise AssertionError('str called') -class Text(str): - def __str__(self): raise AssertionError('str called') - def strip(self): raise AssertionError('strip called') - def lower(self): raise AssertionError('lower called') -hostile = Hostile() -text = Text(' False ') -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let hostile = Field::new("test.flag", locals.get_item("hostile").unwrap().unwrap()); - assert!(!hostile.exact_true().0); - assert!(matches!( - hostile.strict_string(), - Err(ProjectionError::InvalidConfiguration(_)) - )); - let text = Field::new("test.flag", locals.get_item("text").unwrap().unwrap()); - assert_eq!(text.strict_string().unwrap(), " False "); - assert_eq!(text.str_bool().unwrap().0, Some(false)); - }); -} - -#[test] -fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - c" -failure = AttributeError('descriptor failed') -class Snapshot: - @property - def flag(self): raise failure -snapshot = Snapshot() -class Dynamic: - def __getattr__(self, name): raise failure -class Intercepted: - def __getattribute__(self, name): raise failure -dynamic = Dynamic() -intercepted = Intercepted() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let snapshot = locals.get_item("snapshot").unwrap().unwrap(); - let descriptor = PyErr::from(Field::read(&snapshot, "test.flag").err().unwrap()); - assert!( - descriptor - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - for name in ["dynamic", "intercepted"] { - let value = locals.get_item(name).unwrap().unwrap(); - let error = PyErr::from(Field::read(&value, "test.flag").err().unwrap()); - assert!( - error - .value(py) - .is(locals.get_item("failure").unwrap().unwrap()) - ); - } - let missing = PyErr::from(Field::read(&snapshot, "test.missing").err().unwrap()); - assert!(missing.is_instance_of::(py)); - assert!(missing.to_string().contains("test.missing")); - }); -} - -#[test] -fn configuration_errors_name_fields_without_exposing_values() { - Python::initialize(); - Python::attach(|py| { - for source in [ - "{'secret': 'do-not-print'}", - "['host.test', {'secret': 'do-not-print'}]", - ] { - let field = Field::new("test.setting", evaluate(py, source)); - let error = PyErr::from(field.falsy_optional_string().err().unwrap()); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("test.setting")); - assert!(!error.to_string().contains("do-not-print")); - } - let hosts = Field::new( - "url_policy.user_url_allowed_hosts", - evaluate(py, "['host.test', 1]"), - ); - assert!(matches!( - hosts.host_collection(), - Err(ProjectionError::InvalidConfiguration(_)) - )); - assert!(matches!( - Field::new("test.flag", evaluate(py, "1")).str_bool(), - Err(ProjectionError::InvalidConfiguration(_)) - )); - }); -} - -#[test] -fn projection_releases_the_source_collection() { - Python::initialize(); - Python::attach(|py| { - let source = evaluate(py, "['A.test']"); - let projected = Field::new("test.hosts", source.clone()) - .host_collection() - .unwrap() - .0; - source.call_method1("append", ("b.test",)).unwrap(); - assert_eq!(projected, ["a.test"]); - assert_eq!( - Field::new("test.hosts", source) - .host_collection() - .unwrap() - .0, - ["a.test", "b.test"] - ); - }); -} - -#[rstest] -#[case("True", Some(true))] -#[case("False", Some(false))] -#[case("1", None)] -#[case("None", None)] -#[case("[]", None)] -fn accessor_booleans_are_strict_schema_values( - #[case] source: &str, - #[case] expected: Option, -) { - Python::initialize(); - Python::attach(|py| { - let result = Field::new("secret_manager.readable", evaluate(py, source)).schema_bool(); - match expected { - Some(expected) => assert_eq!(result.unwrap(), expected), - None => { - let error = PyErr::from(result.unwrap_err()); - assert!(error.is_instance_of::(py)); - assert!(error.to_string().contains("secret_manager.readable")); - } - } - }); -} diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 596a89a73d7..2515b409c54 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -1,5 +1,5 @@ use std::{ - collections::HashSet, + collections::{BTreeSet, HashSet}, path::{Path, PathBuf}, sync::{Arc, LazyLock, Mutex, PoisonError}, }; @@ -10,9 +10,75 @@ use litellm_http::{ TlsSource, Unsupported, media::{PublicDnsResolver, UrlPolicy}, }; -use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; +use pyo3::{ + exceptions::PyValueError, + prelude::*, + types::{PyBool, PyDict, PyString}, +}; -use crate::{coercion::Field, python_settings::PythonSettings}; +use crate::{ + coercion::{Field, FieldSpec, ProjectionError}, + python_settings::{PythonSettings, Snapshot}, +}; + +const SSL_VERIFY: FieldSpec> = FieldSpec::new("ssl_verify", decode_ssl_verify); +const SSL_CERTIFICATE: FieldSpec> = + FieldSpec::new("ssl_certificate", |field| field.optional_strict_string()); +const SSL_SECURITY_LEVEL: FieldSpec> = + FieldSpec::new("ssl_security_level", |field| field.tuning_string()); +const SSL_ECDH_CURVE: FieldSpec> = + FieldSpec::new("ssl_ecdh_curve", |field| field.tuning_string()); +const FORCE_IPV4: FieldSpec = FieldSpec::new("force_ipv4", |field| field.truthy()); +const HTTP2: FieldSpec = FieldSpec::new("http2", |field| Ok(field.exact_true())); +const AIOHTTP_TRUST_ENV: FieldSpec = + FieldSpec::new("aiohttp_trust_env", |field| field.truthy()); +const DISABLE_AIOHTTP_TRUST_ENV: FieldSpec = + FieldSpec::new("disable_aiohttp_trust_env", |field| field.truthy()); +const DISABLE_AIOHTTP_TRANSPORT: FieldSpec = + FieldSpec::new("disable_aiohttp_transport", |field| Ok(field.exact_true())); +const USER_AGENT: FieldSpec = FieldSpec::new("user_agent", |field| field.schema_string()); +const USER_URL_VALIDATION: FieldSpec = + FieldSpec::new("user_url_validation", |field| field.truthy()); +const USER_URL_ALLOWED_HOSTS: FieldSpec> = + FieldSpec::new("user_url_allowed_hosts", decode_hosts); + +fn decode_hosts(field: &Field<'_>) -> Result, ProjectionError> { + Ok(field + .string_collection()? + .into_iter() + .map(|host| litellm_http::media::normalize_host(&host)) + .collect::>() + .into_iter() + .collect()) +} + +fn decode_ssl_verify(field: &Field<'_>) -> Result, ProjectionError> { + let value = field.value(); + if value.is_none() { + return Ok(None); + } + if value.is_instance_of::() { + return Ok(Some(if field.exact_true() { + SslVerify::Enabled + } else { + SslVerify::Disabled + })); + } + if value.is_instance_of::() { + return Ok(Some(match field.str_bool()? { + Some(true) => SslVerify::Enabled, + Some(false) => SslVerify::Disabled, + None => SslVerify::CaBundle(field.strict_string()?.into()), + })); + } + let context = value.py().import("ssl")?.getattr("SSLContext")?; + if value.is_instance(&context)? { + return Err(ProjectionError::UnsupportedLiveObject(field.expected( + "a Boolean, Boolean string, CA path, or None; live SSLContext is unsupported", + )?)); + } + Err(field.invalid("a Boolean, Boolean string, CA path, or None")) +} static POOL: LazyLock = LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); @@ -80,20 +146,20 @@ pub(crate) fn url_policy(py: Python<'_>) -> PyResult { project_url_policy(&PythonSettings::UrlPolicy.read(py)?) } -fn project_url_policy(value: &Bound<'_, PyAny>) -> PyResult { +fn project_url_policy(snapshot: &Snapshot<'_>) -> PyResult { Ok(UrlPolicy { - validate: Field::read(value, "url_policy.user_url_validation")? - .truthy()? - .0, - allowed_hosts: Field::read(value, "url_policy.user_url_allowed_hosts")? - .host_collection()? - .0, + validate: snapshot.read(&USER_URL_VALIDATION)?, + allowed_hosts: snapshot.read(&USER_URL_ALLOWED_HOSTS)?, }) } fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult> { match kwargs.get_item("ssl_verify")? { - Some(value) => Ok(Field::new("request.ssl_verify", value).ssl_verify()?.0), + Some(value) => Ok(decode_ssl_verify(&Field::new( + "request", + "ssl_verify", + value, + ))?), None => Ok(None), } } @@ -106,39 +172,18 @@ fn for_call(call_ssl_verify: Option, asynchronous: bool) -> HttpSetti } } -fn configured(value: &Bound<'_, PyAny>) -> PyResult { +fn configured(snapshot: &Snapshot<'_>) -> PyResult { Ok(HttpSettingsLayer { - ssl_verify: Field::read(value, "http_settings.ssl_verify")? - .ssl_verify()? - .0, - ssl_certificate: Field::read(value, "http_settings.ssl_certificate")? - .optional_strict_string()? - .0 - .map(PathBuf::from), - ssl_security_level: Field::read(value, "http_settings.ssl_security_level")? - .tuning_string()? - .0, - ssl_ecdh_curve: Field::read(value, "http_settings.ssl_ecdh_curve")? - .tuning_string()? - .0, - force_ipv4: Some(Field::read(value, "http_settings.force_ipv4")?.truthy()?.0), - http2: Some(Field::read(value, "http_settings.http2")?.exact_true().0), - aiohttp_trust_env: Some( - Field::read(value, "http_settings.aiohttp_trust_env")? - .truthy()? - .0, - ), - disable_aiohttp_trust_env: Some( - Field::read(value, "http_settings.disable_aiohttp_trust_env")? - .truthy()? - .0, - ), - disable_aiohttp_transport: Some( - Field::read(value, "http_settings.disable_aiohttp_transport")? - .exact_true() - .0, - ), - user_agent: Some(Field::read(value, "http_settings.user_agent")?.schema_string()?), + ssl_verify: snapshot.read(&SSL_VERIFY)?, + ssl_certificate: snapshot.read(&SSL_CERTIFICATE)?.map(PathBuf::from), + ssl_security_level: snapshot.read(&SSL_SECURITY_LEVEL)?, + ssl_ecdh_curve: snapshot.read(&SSL_ECDH_CURVE)?, + force_ipv4: Some(snapshot.read(&FORCE_IPV4)?), + http2: Some(snapshot.read(&HTTP2)?), + aiohttp_trust_env: Some(snapshot.read(&AIOHTTP_TRUST_ENV)?), + disable_aiohttp_trust_env: Some(snapshot.read(&DISABLE_AIOHTTP_TRUST_ENV)?), + disable_aiohttp_transport: Some(snapshot.read(&DISABLE_AIOHTTP_TRANSPORT)?), + user_agent: Some(snapshot.read(&USER_AGENT)?), ..HttpSettingsLayer::default() }) } @@ -150,12 +195,15 @@ mod tests { use rstest::rstest; use super::*; - use crate::python_settings::CONTRACT; - fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> { + fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> { + py.eval(&std::ffi::CString::new(source).unwrap(), None, None) + .unwrap() + } + + fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Snapshot<'py> { let source = format!( " -import json import types defaults = dict( ssl_verify=True, @@ -170,14 +218,13 @@ defaults = dict( user_agent='litellm/test', ) defaults.update(dict({overrides})) -settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']['fields']}}) +settings = types.SimpleNamespace(**defaults) " ); let locals = PyDict::new(py); - locals.set_item("contract", CONTRACT).unwrap(); let source = std::ffi::CString::new(source).unwrap(); py.run(&source, Some(&locals), Some(&locals)).unwrap(); - locals.get_item("settings").unwrap().unwrap() + PythonSettings::Http.snapshot(locals.get_item("settings").unwrap().unwrap()) } #[test] @@ -395,7 +442,7 @@ user_agent='litellm/9.9.9', Python::attach(|py| { let value = py.eval(c"__import__('types').SimpleNamespace(user_url_validation=[], user_url_allowed_hosts=['B.test', 'a.test.', 'b.test'])", None, None).unwrap(); assert_eq!( - project_url_policy(&value).unwrap(), + project_url_policy(&PythonSettings::UrlPolicy.snapshot(value)).unwrap(), UrlPolicy { validate: false, allowed_hosts: vec!["a.test".into(), "b.test".into()], @@ -419,4 +466,44 @@ user_agent='litellm/9.9.9', let settings = HttpSettings::from_layers([for_call(None, asynchronous), opted_out]); assert_eq!(settings.trust_proxy_env, expected); } + #[rstest] + #[case("'EXAMPLE.TEST.'", vec!["example.test"])] + #[case("['B.test', '', None, 0, [], 'A.test.', 'b.test']", vec!["a.test", "b.test"])] + #[case("('B.test', 'a.test')", vec!["a.test", "b.test"])] + #[case("{'B.test', 'a.test'}", vec!["a.test", "b.test"])] + #[case("(host for host in ['B.test', 'a.test'])", vec!["a.test", "b.test"])] + #[case("None", vec![])] + #[case("False", vec![])] + fn host_collection_is_owned_normalized_and_deterministic( + #[case] source: &str, + #[case] expected: Vec<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + assert_eq!( + decode_hosts(&Field::new( + "url_policy", + "user_url_allowed_hosts", + evaluate(py, source) + )) + .unwrap(), + expected + ); + }); + } + + #[test] + fn projection_releases_the_source_collection() { + Python::initialize(); + Python::attach(|py| { + let source = evaluate(py, "['A.test']"); + let projected = decode_hosts(&Field::new("test", "hosts", source.clone())).unwrap(); + source.call_method1("append", ("b.test",)).unwrap(); + assert_eq!(projected, ["a.test"]); + assert_eq!( + decode_hosts(&Field::new("test", "hosts", source)).unwrap(), + ["a.test", "b.test"] + ); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index f13a3ad433f..b1fc5244d6f 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -7,7 +7,12 @@ mod http; mod marshal; mod python_settings; mod routes; -mod token_counter; +#[allow( + dead_code, + reason = "secret-manager foundations await rollout activation" +)] +mod secrets; +mod tokenizer; #[pymodule(gil_used = true)] mod _native { @@ -32,7 +37,12 @@ mod _native { #[pymodule_export] use crate::routes::responses::ResponsesWebSocketConnection; #[pymodule_export] - use crate::token_counter::TokenCounter; + use crate::routes::token_counter::TokenCounter; + #[cfg(feature = "huggingface")] + #[pymodule_export] + use crate::tokenizer::HuggingFaceEncoding; + #[pymodule_export] + use crate::tokenizer::Tokenizer; #[pymodule_export] use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; use pyo3::{prelude::*, types::PyModule}; @@ -43,7 +53,7 @@ mod _native { let dict = module.dict(); dict.set_item("_CacheTestHandle", py.get_type::())?; dict.set_item("_CacheTestResolver", py.get_type::())?; - dict.set_item("_CacheTestBinding", py.get_type::()) + dict.set_item("_ResponseCacheRuntime", py.get_type::()) } } @@ -78,10 +88,13 @@ mod tests { "achat_completions", "ResponsesWebSocketConnection", "TokenCounter", + "Tokenizer", "gil_stats", "process_state_started", "reserve_process_for_forking", ]; + #[cfg(feature = "huggingface")] + expected.push("HuggingFaceEncoding"); expected.sort_unstable(); let mut public_names: Vec = native_module(py) diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index bdc6d14356d..111ac3bc259 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -1,5 +1,7 @@ use pyo3::prelude::*; +use crate::coercion::{FieldSpec, ProjectionError}; + const MODULE: &str = "litellm.rust_bridge.settings"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -8,28 +10,39 @@ pub(crate) enum PythonSettings { UrlPolicy, ProviderDefaults, SecretManager, + SecretManagerBinding, +} + +pub(crate) struct Snapshot<'py> { + group: PythonSettings, + value: Bound<'py, PyAny>, +} + +impl Snapshot<'_> { + pub(crate) fn read(&self, spec: &FieldSpec) -> Result { + spec.read(&self.value, self.group.name()) + } } impl PythonSettings { - #[cfg(test)] - pub(crate) const ALL: [Self; 4] = [ - Self::Http, - Self::UrlPolicy, - Self::ProviderDefaults, - Self::SecretManager, - ]; - pub(crate) fn name(self) -> &'static str { match self { Self::Http => "http_settings", Self::UrlPolicy => "url_policy", Self::ProviderDefaults => "provider_defaults", Self::SecretManager => "secret_manager", + Self::SecretManagerBinding => "secret_manager_binding", } } - pub(crate) fn read(self, py: Python<'_>) -> PyResult> { - py.import(MODULE)?.getattr(self.name())?.call0() + pub(crate) fn read(self, py: Python<'_>) -> PyResult> { + let value = py.import(MODULE)?.getattr(self.name())?.call0()?; + Ok(Snapshot { group: self, value }) + } + + #[cfg(test)] + pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> { + Snapshot { group: self, value } } pub(crate) fn warn(py: Python<'_>, message: &str) -> PyResult<()> { @@ -38,209 +51,98 @@ impl PythonSettings { } } -#[cfg(test)] -pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); - #[cfg(test)] mod tests { - use super::{CONTRACT, PythonSettings}; - use pyo3::prelude::*; - use serde_json::{Value, json}; + use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict}; - struct SettingSpec { - group: &'static str, - name: &'static str, - adapter: &'static str, - precedence: &'static str, - sensitive: bool, - shapes: &'static [&'static str], - unsupported_live: Option<&'static str>, - } - - const SETTINGS: &[SettingSpec] = &[ - SettingSpec { - group: "http_settings", - name: "ssl_verify", - adapter: "SslVerifyInput", - precedence: "module_global", - sensitive: false, - shapes: &["none", "bool", "str"], - unsupported_live: Some("configuration_error"), - }, - SettingSpec { - group: "http_settings", - name: "ssl_certificate", - adapter: "OptionalStrictString", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "ssl_security_level", - adapter: "TuningString", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "ssl_ecdh_curve", - adapter: "TuningString", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "force_ipv4", - adapter: "Truthy", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "http2", - adapter: "ExactTrue", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "aiohttp_trust_env", - adapter: "Truthy", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "disable_aiohttp_trust_env", - adapter: "Truthy", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "disable_aiohttp_transport", - adapter: "ExactTrue", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "http_settings", - name: "user_agent", - adapter: "StrictString", - precedence: "accessor", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "url_policy", - name: "user_url_validation", - adapter: "Truthy", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "url_policy", - name: "user_url_allowed_hosts", - adapter: "HostCollection", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "provider_defaults", - name: "vertex_project", - adapter: "FalsyOptionalString", - precedence: "module_global", - sensitive: true, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "provider_defaults", - name: "vertex_location", - adapter: "FalsyOptionalString", - precedence: "module_global", - sensitive: true, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "provider_defaults", - name: "enable_azure_ad_token_refresh", - adapter: "ExactTrue", - precedence: "module_global", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - SettingSpec { - group: "secret_manager", - name: "readable", - adapter: "StrictBool", - precedence: "accessor", - sensitive: false, - shapes: &[], - unsupported_live: None, - }, - ]; + use super::PythonSettings; + use crate::coercion::FieldSpec; #[test] - fn settings_manifest_matches_the_semantic_contract() { - pyo3::Python::initialize(); - let manifest: Value = pyo3::Python::attach(|py| { - let value = py - .import("json") - .unwrap() - .call_method1("loads", (CONTRACT,)) - .unwrap(); - litellm_host_python::from_py(&value).unwrap() + fn declarations_select_the_decoder_and_read_only_the_requested_field() { + const TRUTHY: FieldSpec = FieldSpec::new("flag", |field| field.truthy()); + const EXACT: FieldSpec = FieldSpec::new("flag", |field| Ok(field.exact_true())); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +reads = [] +class Settings: + value = 1 + @property + def flag(self): + reads.append('flag') + return self.value + @property + def unrelated(self): + raise AssertionError('unrequested field') +settings = Settings() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let value = locals.get_item("settings").unwrap().unwrap(); + let snapshot = PythonSettings::Http.snapshot(value.clone()); + assert!(snapshot.read(&TRUTHY).unwrap()); + assert!(!snapshot.read(&EXACT).unwrap()); + value.setattr("value", true).unwrap(); + assert!(snapshot.read(&EXACT).unwrap()); + assert_eq!( + locals + .get_item("reads") + .unwrap() + .unwrap() + .extract::>() + .unwrap(), + ["flag", "flag", "flag"] + ); + }); + } + + #[test] + fn declared_reads_preserve_descriptor_and_decoder_failures_and_name_missing_fields() { + const FLAG: FieldSpec = FieldSpec::new("flag", |field| field.truthy()); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +from types import SimpleNamespace +failure = AttributeError('read failed') +class Descriptor: + @property + def flag(self): raise failure +class Truth: + def __bool__(self): raise failure +values = (Descriptor(), SimpleNamespace(flag=Truth())) +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let failure = locals.get_item("failure").unwrap().unwrap(); + for value in locals + .get_item("values") + .unwrap() + .unwrap() + .try_iter() + .unwrap() + { + let snapshot = PythonSettings::Http.snapshot(value.unwrap()); + let error = PyErr::from(snapshot.read(&FLAG).unwrap_err()); + assert!(error.value(py).is(&failure)); + assert!(error.traceback(py).is_some()); + } + let missing = PythonSettings::Http.snapshot(py.eval(c"object()", None, None).unwrap()); + let error = PyErr::from(missing.read(&FLAG).unwrap_err()); + assert!(error.is_instance_of::(py)); + assert!( + error + .to_string() + .contains("http_settings.flag: missing snapshot field") + ); }); - let expected: serde_json::Map = PythonSettings::ALL - .into_iter() - .map(|group| { - let fields: serde_json::Map = SETTINGS - .iter() - .filter(|spec| spec.group == group.name()) - .map(|spec| { - ( - spec.name.to_owned(), - json!({ - "adapter": spec.adapter, - "required": true, - "precedence": spec.precedence, - "sensitive": spec.sensitive, - "shapes": spec.shapes, - "unsupported_live": spec.unsupported_live, - }), - ) - }) - .collect(); - ( - group.name().to_owned(), - json!({"version": 1, "fields": fields}), - ) - }) - .collect(); - assert_eq!(manifest, Value::Object(expected)); } } diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 2d6b849a6b1..8a78a26423d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -3,6 +3,7 @@ pub(crate) mod chat_completions; pub(crate) mod messages; pub(crate) mod ocr; pub(crate) mod responses; +pub(crate) mod token_counter; #[cfg(test)] mod tests { diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 77c8d5d6641..325377e5285 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -130,6 +130,11 @@ impl RouteHost for OcrRouteHost { } fn classify(&self, py: Python<'_>, error: Error) -> PyResult { + if let Error::Secret(source) = &error + && let Some(original) = crate::secrets::callback::python_error(py, source) + { + return Ok(original); + } Ok(self.map_failure(py, ocr_error_to_pyerr(error))) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index d0b13e5056a..2dca6da66cd 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -10,16 +10,33 @@ use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy_python::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::route::ocr_machine; use litellm_core_utils::settings::ProcessEnvironment; -use litellm_llms::base_llm::ocr::{ - handler::OcrClient, - settings::{OcrSettings, Secrets}, +use litellm_llms::base_llm::{ + inference::secrets::{EnvironmentSecrets, SecretSource}, + ocr::{handler::OcrClient, settings::OcrSettings}, }; use pyo3::{ prelude::*, types::{PyDict, PyTuple}, }; -use crate::{coercion::Field, errors::RustBridgeDeclined, http, python_settings::PythonSettings}; +use crate::{ + coercion::FieldSpec, + errors::RustBridgeDeclined, + http, + python_settings::{PythonSettings, Snapshot}, +}; + +const SECRET_MANAGER_READABLE: FieldSpec = + FieldSpec::new("readable", |field| field.schema_bool()); + +const VERTEX_PROJECT: FieldSpec> = + FieldSpec::new("vertex_project", |field| field.falsy_optional_string()); +const VERTEX_LOCATION: FieldSpec> = + FieldSpec::new("vertex_location", |field| field.falsy_optional_string()); +const ENABLE_AZURE_AD_TOKEN_REFRESH: FieldSpec = + FieldSpec::new("enable_azure_ad_token_refresh", |field| { + Ok(field.exact_true()) + }); const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", @@ -62,33 +79,24 @@ fn run_ocr( ) } -fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult { - if Field::read(secret_manager, "secret_manager.readable")?.schema_bool()? { +fn process_environment_secrets(snapshot: &Snapshot<'_>) -> PyResult> { + if snapshot.read(&SECRET_MANAGER_READABLE)? { return Err(RustBridgeDeclined::new_err( "a readable secret manager is configured and the Rust route only reads the process environment", )); } - Ok(Arc::new(ProcessEnvironment)) + Ok(Arc::new(EnvironmentSecrets)) } fn ocr_settings(py: Python<'_>) -> PyResult { project_provider_defaults(&PythonSettings::ProviderDefaults.read(py)?) } -fn project_provider_defaults(value: &Bound<'_, PyAny>) -> PyResult { +fn project_provider_defaults(snapshot: &Snapshot<'_>) -> PyResult { Ok(OcrSettings { - vertex_project: Field::read(value, "provider_defaults.vertex_project")? - .falsy_optional_string()? - .0, - vertex_location: Field::read(value, "provider_defaults.vertex_location")? - .falsy_optional_string()? - .0, - enable_azure_ad_token_refresh: Field::read( - value, - "provider_defaults.enable_azure_ad_token_refresh", - )? - .exact_true() - .0, + vertex_project: snapshot.read(&VERTEX_PROJECT)?, + vertex_location: snapshot.read(&VERTEX_LOCATION)?, + enable_azure_ad_token_refresh: snapshot.read(&ENABLE_AZURE_AD_TOKEN_REFRESH)?, ..OcrSettings::from_environment(&ProcessEnvironment) }) } @@ -120,6 +128,8 @@ mod tests { use super::process_environment_secrets; use crate::errors::RustBridgeDeclined; + use crate::python_settings::PythonSettings; + fn secret_manager<'py>(py: Python<'py>, readable: bool) -> Bound<'py, PyAny> { let locals = PyDict::new(py); locals.set_item("readable", readable).unwrap(); @@ -132,12 +142,26 @@ mod tests { locals.get_item("manager").unwrap().unwrap() } + #[test] + fn a_readable_secret_manager_sends_the_call_back_to_python() { + Python::initialize(); + Python::attach(|py| { + let declined = process_environment_secrets( + &PythonSettings::SecretManager.snapshot(secret_manager(py, true)), + ) + .err() + .expect("the Rust route declines"); + assert!(declined.is_instance_of::(py)); + }); + } + #[test] fn provider_defaults_distinguish_falsey_values_and_exact_true() { Python::initialize(); Python::attach(|py| { let value = py.eval(c"__import__('types').SimpleNamespace(vertex_project=[], vertex_location=0, enable_azure_ad_token_refresh=1)", None, None).unwrap(); - let projected = super::project_provider_defaults(&value).unwrap(); + let snapshot = PythonSettings::ProviderDefaults.snapshot(value.clone()); + let projected = super::project_provider_defaults(&snapshot).unwrap(); assert_eq!(projected.vertex_project, None); assert_eq!(projected.vertex_location, None); assert!(!projected.enable_azure_ad_token_refresh); @@ -146,12 +170,12 @@ mod tests { value .setattr("enable_azure_ad_token_refresh", true) .unwrap(); - let next = super::project_provider_defaults(&value).unwrap(); + let next = super::project_provider_defaults(&snapshot).unwrap(); assert_eq!(next.vertex_project.as_deref(), Some("project")); assert_eq!(next.vertex_location.as_deref(), Some("region")); assert!(next.enable_azure_ad_token_refresh); value.setattr("vertex_project", 1).unwrap(); - let error = super::project_provider_defaults(&value).err().unwrap(); + let error = super::project_provider_defaults(&snapshot).err().unwrap(); assert!(error.is_instance_of::(py)); assert!( error @@ -160,28 +184,4 @@ mod tests { ); }); } - - #[test] - fn a_readable_secret_manager_sends_the_call_back_to_python() { - Python::initialize(); - Python::attach(|py| { - let declined = process_environment_secrets(&secret_manager(py, true)) - .err() - .expect("the Rust route declines"); - assert!(declined.is_instance_of::(py)); - }); - } - - #[test] - fn without_a_readable_secret_manager_secrets_are_the_process_environment() { - Python::initialize(); - Python::attach(|py| { - let secrets = process_environment_secrets(&secret_manager(py, false)).unwrap(); - assert_eq!( - secrets.get("LITELLM_RUST_BRIDGE_UNSET_VARIABLE_FOR_TEST"), - None - ); - assert_eq!(secrets.get("PATH"), std::env::var("PATH").ok()); - }); - } } diff --git a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs new file mode 100644 index 00000000000..168c4883b0b --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs @@ -0,0 +1,87 @@ +use std::sync::Arc; +use std::{num::NonZero, thread::available_parallelism}; + +use litellm_host_python::{enter_native, run_async}; +use litellm_token_counter::{ + CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, +}; +use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + types::PyAny, +}; +use tokio::sync::Semaphore; + +use crate::errors::RustBridgeDeclined; +use crate::tokenizer::Tokenizer; + +/// Counts the input tokens of a raw request body off the Python event loop with +/// the GIL released. Python owns which requests get here and what to do with +/// the count. At most one encode per core runs at a time; the rest wait in the +/// async task, where a cancelled Python awaiter drops them before any blocking +/// work is scheduled. +#[pyclass(frozen)] +pub(crate) struct TokenCounter { + inner: Arc, + encode_slots: Arc, +} + +#[pymethods] +impl TokenCounter { + #[staticmethod] + #[pyo3(signature = (tokenizer, fast = false))] + fn from_tokenizer(py: Python<'_>, tokenizer: &Tokenizer, fast: bool) -> PyResult { + enter_native()?; + let inner = CoreTokenCounter::new(tokenizer.counter(py, fast)); + Ok(Self { + inner: Arc::new(inner), + encode_slots: Arc::new(Semaphore::new(encode_parallelism())), + }) + } + + fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult> { + let counter = Arc::clone(&self.inner); + let encode_slots = Arc::clone(&self.encode_slots); + let body = body.to_vec(); + run_async( + py, + async move { + let _slot = encode_slots + .acquire_owned() + .await + .map_err(|error| Error::Task(error.to_string()))?; + tokio::task::spawn_blocking(move || count_body(&counter, &body)) + .await + .map_err(|error| Error::Task(error.to_string()))? + }, + token_count_error_to_pyerr, + ) + } +} + +fn encode_parallelism() -> usize { + available_parallelism().map_or(1, NonZero::get) +} + +fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result { + let request = CountableRequest::parse(body)?; + counter.count_request(&request) +} + +pub(crate) fn token_count_error_to_pyerr(error: Error) -> PyErr { + let message = error.to_string(); + match error { + Error::Load(_) + | Error::Ranks(_) + | Error::UnicodeClasses + | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message), + Error::RequestParse(_) + | Error::MissingInput + | Error::FloatText + | Error::ContentBlock + | Error::ArrayItems + | Error::JsonSerialization(_) + | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message), + Error::Encode(_) | Error::Decode(_) | Error::Task(_) => PyRuntimeError::new_err(message), + } +} diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs new file mode 100644 index 00000000000..2b98081acef --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -0,0 +1,349 @@ +use std::{fmt, future::Future, pin::Pin}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets::{ + Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, +}; +use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; + +const HANDLER_MODULE: &str = "litellm.secret_managers.secret_manager_handler"; + +struct PythonSecretError(Py); + +impl fmt::Debug for PythonSecretError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("PythonSecretError") + } +} + +impl fmt::Display for PythonSecretError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("Python secret manager failed") + } +} + +impl std::error::Error for PythonSecretError {} + +pub(crate) fn python_error(py: Python<'_>, error: &Error) -> Option { + let Error::ExternalManager(source) = error else { + return None; + }; + source + .downcast_ref::() + .map(|error| PyErr::from_value(error.0.clone_ref(py).into_bound(py).into_any())) +} + +/// A secret manager whose reads execute in Python: a custom manager, a legacy compatible +/// client, or a manually assigned SDK client. +pub(crate) struct PythonSecretManager { + client: Py, + system: Option, + /// The `key_manager` name Python's handler dispatches on. + key_manager: &'static str, + settings: Option>, +} + +impl PythonSecretManager { + pub(crate) fn new( + client: Py, + system: Option, + settings: Option>, + ) -> Self { + Self { + client, + system, + key_manager: system.map_or("local", python_name), + settings, + } + } + + fn read(&self, py: Python<'_>, name: &str) -> PyResult> { + let client = self.client.bind(py); + if self.system == Some(KeyManagementSystem::Custom) + || (self.system.is_none() && client.hasattr("sync_read_secret")?) + { + let kwargs = PyDict::new(py); + kwargs.set_item("secret_name", name)?; + if self.system == Some(KeyManagementSystem::Custom) { + let optional_params = self + .settings + .as_ref() + .map(|settings| settings.bind(py).call_method0("model_dump")) + .transpose()?; + kwargs.set_item("optional_params", optional_params)?; + } + return client + .call_method("sync_read_secret", (), Some(&kwargs))? + .extract(); + } + let kwargs = PyDict::new(py); + kwargs.set_item("client", client)?; + kwargs.set_item("key_manager", self.key_manager)?; + kwargs.set_item("secret_name", name)?; + kwargs.set_item( + "key_management_settings", + self.settings + .as_ref() + .map_or_else(|| py.None(), |settings| settings.clone_ref(py)), + )?; + py.import(HANDLER_MODULE)? + .getattr("get_secret_from_manager")? + .call((), Some(&kwargs))? + .extract() + } +} + +/// The `KeyManagementSystem` value as Python spells it. +fn python_name(system: KeyManagementSystem) -> &'static str { + match system { + KeyManagementSystem::GoogleKms => "google_kms", + KeyManagementSystem::AzureKeyVault => "azure_key_vault", + KeyManagementSystem::AwsSecretManager => "aws_secret_manager", + KeyManagementSystem::GoogleSecretManager => "google_secret_manager", + KeyManagementSystem::HashicorpVault => "hashicorp_vault", + KeyManagementSystem::Cyberark => "cyberark", + KeyManagementSystem::Local => "local", + KeyManagementSystem::AwsKms => "aws_kms", + KeyManagementSystem::Custom => "custom", + } +} + +impl ExternalSecretManager for PythonSecretManager { + fn system(&self) -> KeyManagementSystem { + self.system.unwrap_or(KeyManagementSystem::Custom) + } + + fn read_secret<'a>( + &'a self, + name: &'a str, + _settings: &'a KeyManagementSettings, + _environment: &'a (dyn Lookup + Send + Sync), + ) -> Pin, Error>> + Send + 'a>> { + Box::pin(async move { + Python::attach(|py| { + self.read(py, name) + .map(|value| value.map(SecretValue::new).map(Secret::String)) + .map_err(|error| { + Error::ExternalManager(Box::new(PythonSecretError(error.into_value(py)))) + }) + }) + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use litellm_secrets::{ + FailurePolicy, KeyManagementSettings, KeyManagementSystem, OidcResolver, SecretManager, + SecretManagerState, SecretResolver, + }; + use pyo3::{prelude::*, types::PyDict}; + + use super::{HANDLER_MODULE, PythonSecretManager, python_error, python_name}; + + #[tokio::test] + async fn callback_failures_preserve_python_exceptions_even_with_environment_fallback() { + Python::initialize(); + for failure_type in ["ValueError", "asyncio.CancelledError"] { + for fallback in [None, Some("environment-key")] { + let (reader, locals) = Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("failure_type", failure_type).unwrap(); + py.run( + c" +import asyncio +failure = eval(failure_type)('secret manager failed') +cause = RuntimeError('original cause') +context = RuntimeError('original context') +failure.__cause__ = cause +failure.__context__ = context +class Manager: + def sync_read_secret(self, secret_name): + raise failure +manager = Manager() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let reader = PythonSecretManager::new( + locals.get_item("manager").unwrap().unwrap().unbind(), + None, + None, + ); + (reader, locals.unbind()) + }); + let resolver = SecretResolver::new( + Arc::new(SecretManagerState::new( + SecretManager::External(Arc::new(reader)), + KeyManagementSettings::default(), + )), + Arc::new(move |_: &str| fallback.map(str::to_owned)), + OidcResolver::default(), + ) + .with_failure_policy(FailurePolicy::EnvironmentFallback); + let error = resolver.get_secret("API_KEY", None).await.unwrap_err(); + Python::attach(|py| { + let original = python_error(py, &error).unwrap(); + let locals = locals.bind(py); + assert!( + original + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + for (attribute, name) in [("__cause__", "cause"), ("__context__", "context")] { + assert!( + original + .value(py) + .getattr(attribute) + .unwrap() + .is(locals.get_item(name).unwrap().unwrap()) + ); + } + assert!(original.traceback(py).is_some()); + }); + } + } + } + + /// Installs a fake `get_secret_from_manager` that records its kwargs, runs `body`, and + /// removes the fake modules again. + fn with_fake_handler<'py>(py: Python<'py>, body: impl FnOnce(&Bound<'py, PyDict>)) { + let locals = PyDict::new(py); + py.run( + c" +import sys, types +calls = [] +def get_secret_from_manager(**kwargs): + calls.append(kwargs) + return 'handled-' + kwargs['secret_name'] +handler = types.ModuleType('litellm.secret_managers.secret_manager_handler') +handler.get_secret_from_manager = get_secret_from_manager +installed = {} +for name in ('litellm', 'litellm.secret_managers'): + if name not in sys.modules: + sys.modules[name] = types.ModuleType(name) + installed[name] = True +sys.modules['litellm.secret_managers.secret_manager_handler'] = handler +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + body(&locals); + py.run( + c" +sys.modules.pop('litellm.secret_managers.secret_manager_handler', None) +for name in installed: + sys.modules.pop(name, None) +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + } + + #[test] + fn python_names_round_trip_through_serde() { + for system in [ + KeyManagementSystem::GoogleKms, + KeyManagementSystem::AzureKeyVault, + KeyManagementSystem::AwsSecretManager, + KeyManagementSystem::GoogleSecretManager, + KeyManagementSystem::HashicorpVault, + KeyManagementSystem::Cyberark, + KeyManagementSystem::Local, + KeyManagementSystem::AwsKms, + KeyManagementSystem::Custom, + ] { + assert_eq!( + serde_json::to_value(system).unwrap(), + serde_json::Value::String(python_name(system).to_owned()) + ); + } + } + + #[test] + fn custom_readers_without_a_system_are_called_directly() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +class Manager: + def __init__(self): + self.names = [] + def sync_read_secret(self, secret_name, optional_params=None, timeout=None): + self.names.append(secret_name) + return 'direct-' + secret_name +manager = Manager() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let manager = locals.get_item("manager").unwrap().unwrap(); + let reader = PythonSecretManager::new(manager.clone().unbind(), None, None); + assert_eq!( + reader.read(py, "API_KEY").unwrap().as_deref(), + Some("direct-API_KEY") + ); + assert_eq!( + manager + .getattr("names") + .unwrap() + .extract::>() + .unwrap(), + ["API_KEY"] + ); + }); + } + + #[test] + fn configured_systems_dispatch_through_the_python_handler_with_the_original_settings() { + Python::initialize(); + Python::attach(|py| { + with_fake_handler(py, |locals| { + let client = py.eval(c"object()", None, None).unwrap(); + let settings = py.eval(c"object()", None, None).unwrap(); + let reader = PythonSecretManager::new( + client.clone().unbind(), + Some(KeyManagementSystem::AzureKeyVault), + Some(settings.clone().unbind()), + ); + assert_eq!( + reader.read(py, "API_KEY").unwrap().as_deref(), + Some("handled-API_KEY") + ); + assert!(py.import(HANDLER_MODULE).is_ok()); + let calls = locals.get_item("calls").unwrap().unwrap(); + let call = calls.get_item(0).unwrap().cast_into::().unwrap(); + assert!(call.get_item("client").unwrap().unwrap().is(&client)); + assert!( + call.get_item("key_management_settings") + .unwrap() + .unwrap() + .is(&settings) + ); + assert_eq!( + call.get_item("key_manager") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + "azure_key_vault" + ); + assert_eq!( + call.get_item("secret_name") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + "API_KEY" + ); + }); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs new file mode 100644 index 00000000000..6fd380fe40c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -0,0 +1,338 @@ +use std::sync::Arc; + +use litellm_secrets::{SecretManager, SecretManagerState}; +use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, SecretValue}; +use pyo3::prelude::*; +use serde_json::Value; + +use super::callback::PythonSecretManager; +use crate::{ + coercion::{Field, FieldSpec, ProjectionError}, + python_settings::{PythonSettings, Snapshot}, +}; + +const SYSTEM: FieldSpec> = + FieldSpec::new("system", parse_optional_system); +const ACCESS_MODE: FieldSpec = FieldSpec::new("access_mode", parse_access_mode); +const HOSTED_KEYS: FieldSpec>> = + FieldSpec::new("hosted_keys", |field| field.optional_string_collection()); +const STORE_VIRTUAL_KEYS: FieldSpec = + FieldSpec::new("store_virtual_keys", |field| field.truthy()); +const PREFIX_FOR_STORED_VIRTUAL_KEYS: FieldSpec = + FieldSpec::new("prefix_for_stored_virtual_keys", |field| { + field.strict_string() + }); +const PRIMARY_SECRET_NAME: FieldSpec> = + FieldSpec::new("primary_secret_name", |field| field.falsy_optional_string()); +const KMS_KEY_ID: FieldSpec> = + FieldSpec::new("kms_key_id", |field| field.falsy_optional_string()); +const CUSTOM_SECRET_MANAGER: FieldSpec> = + FieldSpec::new("custom_secret_manager", |field| { + field.falsy_optional_string() + }); +const AWS_REGION_NAME: FieldSpec> = + FieldSpec::new("aws_region_name", |field| field.falsy_optional_string()); +const AWS_ROLE_NAME: FieldSpec> = + FieldSpec::new("aws_role_name", |field| field.falsy_optional_string()); +const AWS_SESSION_NAME: FieldSpec> = + FieldSpec::new("aws_session_name", |field| field.falsy_optional_string()); +const AWS_EXTERNAL_ID: FieldSpec> = + FieldSpec::new("aws_external_id", |field| field.falsy_optional_string()); +const AWS_PROFILE_NAME: FieldSpec> = + FieldSpec::new("aws_profile_name", |field| field.falsy_optional_string()); +const AWS_WEB_IDENTITY_TOKEN: FieldSpec> = + FieldSpec::new("aws_web_identity_token", |field| { + field.falsy_optional_string() + }); +const AWS_STS_ENDPOINT: FieldSpec> = + FieldSpec::new("aws_sts_endpoint", |field| field.falsy_optional_string()); +const REPLICA_REGIONS: FieldSpec>> = + FieldSpec::new("replica_regions", |field| { + field.optional_string_collection() + }); +const CLIENT: FieldSpec>> = + FieldSpec::new("client", |field| Ok(field.python_binding())); +const SETTINGS_OBJECT: FieldSpec>> = + FieldSpec::new("settings_object", |field| Ok(field.python_binding())); + +/// `litellm.secret_manager_client` as the bridge classifies it. +#[derive(Debug)] +pub(crate) enum SecretManagerClient { + /// `None`: reads come from the process environment. + Local, + /// A custom manager, legacy compatible client, or manually assigned SDK client that keeps + /// executing in Python. + PythonCallback(Py), +} + +/// One operation-local capture of the secret manager globals, taken while attached to Python. +#[derive(Debug)] +pub(crate) struct SecretManagerSnapshot { + pub(crate) client: SecretManagerClient, + pub(crate) system: Option, + /// Typed settings that drive native routing: access mode and hosted keys. + pub(crate) settings: KeyManagementSettings, + /// The original `KeyManagementSettings` object, handed back to Python callbacks unchanged. + pub(crate) settings_object: Option>, +} + +impl SecretManagerSnapshot { + pub(crate) fn into_state(self) -> Arc { + match self.client { + SecretManagerClient::Local => Arc::new(SecretManagerState::default()), + SecretManagerClient::PythonCallback(client) => Arc::new(SecretManagerState::new( + SecretManager::External(Arc::new(PythonSecretManager::new( + client, + self.system, + self.settings_object, + ))), + self.settings, + )), + } + } +} + +/// Reads and projects the secret manager settings group in one attached operation. +pub(crate) fn read(py: Python<'_>) -> PyResult { + Ok(project(&PythonSettings::SecretManagerBinding.read(py)?)?) +} + +pub(crate) fn project(snapshot: &Snapshot<'_>) -> Result { + let system = snapshot.read(&SYSTEM)?; + let access_mode = snapshot.read(&ACCESS_MODE)?; + let settings = KeyManagementSettings { + hosted_keys: snapshot.read(&HOSTED_KEYS)?, + store_virtual_keys: Some(snapshot.read(&STORE_VIRTUAL_KEYS)?), + prefix_for_stored_virtual_keys: snapshot.read(&PREFIX_FOR_STORED_VIRTUAL_KEYS)?, + access_mode, + primary_secret_name: snapshot.read(&PRIMARY_SECRET_NAME)?, + kms_key_id: snapshot.read(&KMS_KEY_ID)?, + custom_secret_manager: snapshot.read(&CUSTOM_SECRET_MANAGER)?, + aws_region_name: snapshot.read(&AWS_REGION_NAME)?, + aws_role_name: snapshot.read(&AWS_ROLE_NAME)?, + aws_session_name: snapshot.read(&AWS_SESSION_NAME)?, + aws_external_id: snapshot.read(&AWS_EXTERNAL_ID)?.map(SecretValue::new), + aws_profile_name: snapshot.read(&AWS_PROFILE_NAME)?, + aws_web_identity_token: snapshot + .read(&AWS_WEB_IDENTITY_TOKEN)? + .map(SecretValue::new), + aws_sts_endpoint: snapshot.read(&AWS_STS_ENDPOINT)?, + replica_regions: snapshot.read(&REPLICA_REGIONS)?, + ..KeyManagementSettings::default() + }; + let client = match snapshot.read(&CLIENT)? { + None => SecretManagerClient::Local, + Some(client) => SecretManagerClient::PythonCallback(client), + }; + Ok(SecretManagerSnapshot { + client, + system, + settings, + settings_object: snapshot.read(&SETTINGS_OBJECT)?, + }) +} + +fn parse_optional_system( + field: &Field<'_>, +) -> Result, ProjectionError> { + let Some(value) = field.falsy_optional_string()? else { + return Ok(None); + }; + serde_json::from_value(Value::String(value)) + .map(Some) + .map_err(|error| { + ProjectionError::InvalidConfiguration(format!("secret manager system: {error}")) + }) +} + +fn parse_access_mode(field: &Field<'_>) -> Result { + let value = field.strict_string()?; + serde_json::from_value(Value::String(value)).map_err(|error| { + ProjectionError::InvalidConfiguration(format!("secret manager access mode: {error}")) + }) +} + +#[cfg(test)] +mod tests { + use pyo3::{ + prelude::*, + types::{PyDict, PyTuple}, + }; + + use super::{SecretManagerClient, project}; + use crate::python_settings::PythonSettings; + + fn snapshot<'py>( + py: Python<'py>, + system: &str, + access_mode: &str, + store_virtual_keys: Bound<'py, PyAny>, + hosted_keys: Bound<'py, PyAny>, + ) -> crate::python_settings::Snapshot<'py> { + snapshot_with_client( + py, + system, + access_mode, + store_virtual_keys, + hosted_keys, + py.None().into_bound(py), + ) + } + + fn snapshot_with_client<'py>( + py: Python<'py>, + system: &str, + access_mode: &str, + store_virtual_keys: Bound<'py, PyAny>, + hosted_keys: Bound<'py, PyAny>, + client: Bound<'py, PyAny>, + ) -> crate::python_settings::Snapshot<'py> { + let locals = PyDict::new(py); + locals.set_item("client", client).unwrap(); + locals.set_item("system", system).unwrap(); + locals.set_item("access_mode", access_mode).unwrap(); + locals + .set_item("store_virtual_keys", store_virtual_keys) + .unwrap(); + locals.set_item("hosted_keys", hosted_keys).unwrap(); + py.run( + cr#" +from dataclasses import dataclass +from types import SimpleNamespace + +@dataclass(frozen=True, slots=True) +class SecretManager: + system: object + access_mode: object + hosted_keys: object + primary_secret_name: object + store_virtual_keys: object + prefix_for_stored_virtual_keys: object + kms_key_id: object + custom_secret_manager: object + aws_region_name: object + aws_role_name: object + aws_session_name: object + aws_external_id: object + aws_profile_name: object + aws_web_identity_token: object + aws_sts_endpoint: object + replica_regions: object + client: object + settings_object: object + +root = SimpleNamespace(secret_manager=SecretManager( + system=system, + access_mode=access_mode, + hosted_keys=hosted_keys, + primary_secret_name=None, + store_virtual_keys=store_virtual_keys, + prefix_for_stored_virtual_keys="litellm/", + kms_key_id=None, + custom_secret_manager=None, + aws_region_name=None, + aws_role_name=None, + aws_session_name=None, + aws_external_id=None, + aws_profile_name=None, + aws_web_identity_token=None, + aws_sts_endpoint=None, + replica_regions=None, + client=client, + settings_object=None, +)) +"#, + Some(&locals), + Some(&locals), + ) + .unwrap(); + PythonSettings::SecretManagerBinding.snapshot( + locals + .get_item("root") + .unwrap() + .unwrap() + .getattr("secret_manager") + .unwrap(), + ) + } + + #[rstest::rstest] + #[case::string_true(Some("true"), false, true)] + #[case::string_one(Some("1"), false, true)] + #[case::true_value(None, true, true)] + #[case::false_value(None, false, false)] + #[case::string_false(Some("false"), false, true)] + fn python_compatible_boolean_coercion( + #[case] string_value: Option<&str>, + #[case] bool_value: bool, + #[case] expected: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let store_virtual_keys = match string_value { + Some(value) => value.into_pyobject(py).unwrap().into_any(), + None => bool_value.into_pyobject(py).unwrap().to_owned().into_any(), + }; + let hosted_keys = PyTuple::new(py, ["ONE"]).unwrap().into_any(); + let projected = project(&snapshot( + py, + "local", + "read_only", + store_virtual_keys, + hosted_keys, + )) + .unwrap(); + assert_eq!(projected.settings.store_virtual_keys, Some(expected)); + }); + } + + #[test] + fn unknown_system_is_rejected() { + Python::initialize(); + Python::attach(|py| { + let error = project(&snapshot( + py, + "unknown", + "read_only", + false.into_pyobject(py).unwrap().to_owned().into_any(), + PyTuple::empty(py).into_any(), + )) + .unwrap_err(); + let error: PyErr = error.into(); + assert!(error.is_instance_of::(py)); + }); + } + + #[test] + fn client_identity_selects_local_or_python_callback() { + Python::initialize(); + Python::attach(|py| { + let falsy = false.into_pyobject(py).unwrap().to_owned().into_any(); + let local = project(&snapshot( + py, + "local", + "read_only", + falsy.clone(), + PyTuple::empty(py).into_any(), + )) + .unwrap(); + assert!(matches!(local.client, SecretManagerClient::Local)); + assert!(local.settings_object.is_none()); + + let manager = py.eval(c"object()", None, None).unwrap(); + let custom = project(&snapshot_with_client( + py, + "custom", + "read_only", + falsy, + PyTuple::empty(py).into_any(), + manager.clone(), + )) + .unwrap(); + let SecretManagerClient::PythonCallback(client) = custom.client else { + panic!("a live client must stay a Python callback"); + }; + assert!(client.bind(py).is(&manager)); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/secrets/mod.rs b/litellm-rust/crates/python-bridge/src/secrets/mod.rs new file mode 100644 index 00000000000..f6ca57b08d1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/secrets/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod callback; +pub(crate) mod config; +pub(crate) mod resolved; diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs new file mode 100644 index 00000000000..877c429169a --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -0,0 +1,252 @@ +use std::{collections::HashMap, sync::Arc}; + +use futures_util::{future::BoxFuture, future::try_join_all}; +use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; +use litellm_llms::base_llm::inference::secrets::{SecretSource, Secrets}; +use litellm_secrets::{ + Error, FailurePolicy, OidcResolver, Secret, SecretManagerState, SecretResolver, +}; + +use super::config::SecretManagerSnapshot; + +pub(crate) struct ResolvedSecrets { + resolver: SecretResolver, +} + +impl ResolvedSecrets { + pub(crate) fn new(snapshot: SecretManagerSnapshot) -> Self { + Self::from_state(snapshot.into_state()) + } + + fn from_state(state: Arc) -> Self { + Self { + resolver: SecretResolver::new( + state, + Arc::new(ProcessEnvironment), + OidcResolver::default(), + ) + .with_failure_policy(FailurePolicy::EnvironmentFallback), + } + } +} + +impl SecretSource for ResolvedSecrets { + fn resolve<'a>(&'a self, names: &'a [&'static str]) -> BoxFuture<'a, Result> { + Box::pin(async move { + let values = try_join_all(names.iter().map(|name| async move { + self.resolver + .get_secret(name, None) + .await + .map(|secret| secret.map(|secret| ((*name).to_owned(), secret_value(secret)))) + })) + .await? + .into_iter() + .flatten() + .collect::>(); + Ok(Arc::new(ResolvedLookup { values }) as Secrets) + }) + } +} + +struct ResolvedLookup { + values: HashMap, +} + +impl Lookup for ResolvedLookup { + fn get(&self, name: &str) -> Option { + self.values + .get(name) + .cloned() + .or_else(|| ProcessEnvironment.get(name)) + } +} + +fn secret_value(secret: Secret) -> String { + match secret { + Secret::String(value) => value.expose().to_owned(), + Secret::Bool(value) => if value { "True" } else { "False" }.to_owned(), + Secret::Json(value) => value.to_string(), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use aws_sdk_secretsmanager::Client; + use aws_sdk_secretsmanager::config::{ + BehaviorVersion, Credentials, Region, retry::RetryConfig, + }; + use litellm_secrets::{AccessMode, KeyManagementSettings, SecretManager, SecretManagerState}; + use litellm_secrets_aws::AwsSecretsManagerV2; + use serde_json::json; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, header}, + }; + + use super::ResolvedSecrets; + use litellm_llms::base_llm::inference::secrets::SecretSource; + + fn state(server: &MockServer, settings: KeyManagementSettings) -> Arc { + let client = Client::from_conf( + aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(), + ); + Arc::new(SecretManagerState::new( + SecretManager::AwsSecretsManagerV2(AwsSecretsManagerV2::new( + client, + (&settings).into(), + )), + settings, + )) + } + + async fn resolve(state: Arc, name: &'static str) -> Option { + ResolvedSecrets::from_state(state) + .resolve(&[name]) + .await + .unwrap() + .get(name) + } + + #[tokio::test] + async fn hosted_key_miss_falls_back_to_environment() { + let name = "LITELLM_RUST_BRIDGE_HOSTED_KEY_MISS"; + unsafe { std::env::set_var(name, "env-key") }; + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .and(body_partial_json(json!({"SecretId": name}))) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})), + ) + .expect(0) + .mount(&server) + .await; + let result = resolve( + state( + &server, + KeyManagementSettings { + hosted_keys: Some(vec!["OTHER".into()]), + ..Default::default() + }, + ), + name, + ) + .await; + unsafe { std::env::remove_var(name) }; + assert_eq!(result.as_deref(), Some("env-key")); + assert_eq!(server.received_requests().await.unwrap().len(), 0); + } + + #[tokio::test] + async fn manager_failure_falls_back_to_environment() { + let name = "LITELLM_RUST_BRIDGE_MANAGER_FAILURE"; + unsafe { std::env::set_var(name, "env-key") }; + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&server) + .await; + let result = resolve(state(&server, KeyManagementSettings::default()), name).await; + unsafe { std::env::remove_var(name) }; + assert_eq!(result.as_deref(), Some("env-key")); + assert_eq!(server.received_requests().await.unwrap().len(), 1); + + let missing_server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&missing_server) + .await; + let missing = + ResolvedSecrets::from_state(state(&missing_server, KeyManagementSettings::default())) + .resolve(&["LITELLM_RUST_BRIDGE_MANAGER_FAILURE_MISSING"]) + .await; + assert!(matches!(missing, Err(litellm_secrets::Error::Aws(_)))); + } + + #[tokio::test] + async fn write_only_mode_never_consults_the_manager() { + let name = "LITELLM_RUST_BRIDGE_WRITE_ONLY"; + unsafe { std::env::set_var(name, "env-key") }; + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})), + ) + .expect(0) + .mount(&server) + .await; + let result = resolve( + state( + &server, + KeyManagementSettings { + access_mode: AccessMode::WriteOnly, + ..Default::default() + }, + ), + name, + ) + .await; + unsafe { std::env::remove_var(name) }; + assert_eq!(result.as_deref(), Some("env-key")); + assert_eq!(server.received_requests().await.unwrap().len(), 0); + } + + #[tokio::test] + async fn read_only_mode_resolves_from_the_manager() { + let name = "LITELLM_RUST_BRIDGE_READ_ONLY"; + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .and(body_partial_json(json!({"SecretId": name}))) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"SecretString": "manager-key"})), + ) + .expect(1) + .mount(&server) + .await; + assert_eq!( + resolve(state(&server, KeyManagementSettings::default()), name) + .await + .as_deref(), + Some("manager-key") + ); + assert_eq!(server.received_requests().await.unwrap().len(), 1); + } + + #[tokio::test] + async fn oidc_failures_are_not_converted_to_missing_secrets() { + let result = ResolvedSecrets::from_state(Arc::new(SecretManagerState::default())) + .resolve(&["oidc/"]) + .await; + assert!(matches!(result, Err(litellm_secrets::Error::InvalidOidc))); + } + + #[tokio::test] + async fn undeclared_names_still_read_the_process_environment() { + let name = "LITELLM_RUST_BRIDGE_UNDECLARED"; + unsafe { std::env::set_var(name, "env-key") }; + let server = MockServer::start().await; + let result = resolve( + state( + &server, + KeyManagementSettings { + hosted_keys: Some(vec!["OTHER".into()]), + ..Default::default() + }, + ), + name, + ) + .await; + unsafe { std::env::remove_var(name) }; + assert_eq!(result.as_deref(), Some("env-key")); + assert_eq!(server.received_requests().await.unwrap().len(), 0); + } +} diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs deleted file mode 100644 index 244401e6696..00000000000 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ /dev/null @@ -1,158 +0,0 @@ -use std::sync::Arc; - -#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] -use std::{num::NonZero, thread::available_parallelism}; - -#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] -use litellm_host_python::release_gil; -use litellm_host_python::run_async; -use litellm_token_counter::{ - CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, -}; -use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, - prelude::*, - types::PyAny, -}; -use tokio::sync::Semaphore; - -use crate::errors::RustBridgeDeclined; - -/// Counts the input tokens of a raw request body off the Python event loop with -/// the GIL released. Python owns which requests get here and what to do with -/// the count. At most one encode per core runs at a time; the rest wait in the -/// async task, where a cancelled Python awaiter drops them before any blocking -/// work is scheduled. -#[pyclass(frozen)] -pub(crate) struct TokenCounter { - inner: Arc, - encode_slots: Arc, -} - -#[pymethods] -impl TokenCounter { - #[new] - fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult { - #[cfg(feature = "fast")] - { - Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json)) - } - #[cfg(all(not(feature = "fast"), feature = "huggingface"))] - { - Self::load(py, || CoreTokenCounter::from_json(tokenizer_json)) - } - #[cfg(not(any(feature = "fast", feature = "huggingface")))] - { - let _ = (py, tokenizer_json); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the fast or huggingface feature", - )) - } - } - - #[staticmethod] - fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { - #[cfg(feature = "fast")] - { - Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file)) - } - #[cfg(not(feature = "fast"))] - { - let _ = (py, rank_file); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the fast feature", - )) - } - } - - #[staticmethod] - fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { - #[cfg(feature = "fast")] - { - Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file)) - } - #[cfg(not(feature = "fast"))] - { - let _ = (py, rank_file); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the fast feature", - )) - } - } - - #[staticmethod] - fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { - #[cfg(feature = "tiktoken")] - { - Self::load(py, || CoreTokenCounter::from_tiktoken(encoding)) - } - #[cfg(not(feature = "tiktoken"))] - { - let _ = (py, encoding); - Err(RustBridgeDeclined::new_err( - "tokenizer backend requires the tiktoken feature", - )) - } - } - - fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult> { - let counter = Arc::clone(&self.inner); - let encode_slots = Arc::clone(&self.encode_slots); - let body = body.to_vec(); - run_async( - py, - async move { - let _slot = encode_slots - .acquire_owned() - .await - .map_err(|error| Error::Task(error.to_string()))?; - tokio::task::spawn_blocking(move || count_body(&counter, &body)) - .await - .map_err(|error| Error::Task(error.to_string()))? - }, - token_count_error_to_pyerr, - ) - } -} - -impl TokenCounter { - #[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] - fn load( - py: Python<'_>, - load: impl FnOnce() -> Result + Send, - ) -> PyResult { - let inner = release_gil(py, load).map_err(token_count_error_to_pyerr)?; - Ok(Self { - inner: Arc::new(inner), - encode_slots: Arc::new(Semaphore::new(encode_parallelism())), - }) - } -} - -#[cfg(any(feature = "fast", feature = "huggingface", feature = "tiktoken"))] -fn encode_parallelism() -> usize { - available_parallelism().map_or(1, NonZero::get) -} - -fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result { - let request = CountableRequest::parse(body)?; - counter.count_request(&request) -} - -fn token_count_error_to_pyerr(error: Error) -> PyErr { - let message = error.to_string(); - match error { - Error::Load(_) - | Error::Ranks(_) - | Error::UnicodeClasses - | Error::UnsupportedTokenizer(_) => PyValueError::new_err(message), - Error::RequestParse(_) - | Error::MissingInput - | Error::FloatText - | Error::ContentBlock - | Error::ArrayItems - | Error::JsonSerialization(_) - | Error::JsonUtf8(_) => RustBridgeDeclined::new_err(message), - Error::Encode(_) | Error::Task(_) => PyRuntimeError::new_err(message), - } -} diff --git a/litellm-rust/crates/python-bridge/src/tokenizer.rs b/litellm-rust/crates/python-bridge/src/tokenizer.rs new file mode 100644 index 00000000000..df219d55eb6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/tokenizer.rs @@ -0,0 +1,713 @@ +//! The Python face of the text codecs: one `Tokenizer` class over the tiktoken and Hugging +//! Face backends, carrying the read-only surface of `tiktoken.Encoding` and +//! `tokenizers.Tokenizer` that `litellm/litellm_core_utils/tokenizer.py` wraps. +use std::borrow::Cow; +#[cfg(any(feature = "tiktoken", feature = "huggingface"))] +use std::collections::HashMap; +use std::sync::Arc; +#[cfg(feature = "fast")] +use std::sync::OnceLock; + +use litellm_host_python::{enter_native, release_gil}; +#[cfg(feature = "fast")] +use litellm_token_counter::fast::{FastCounter, FastTokenizer}; +use litellm_token_counter::{Error, TextCodec}; +use pyo3::{exceptions::PyUnicodeEncodeError, prelude::*, types::PyString}; + +#[cfg(any(feature = "tiktoken", feature = "huggingface"))] +use pyo3::exceptions::PyValueError; +#[cfg(feature = "huggingface")] +use pyo3::{exceptions::PyIOError, types::PyDict}; +#[cfg(feature = "tiktoken")] +use pyo3::{ + exceptions::{PyKeyError, PyRuntimeError}, + types::PyBytes, +}; + +#[cfg(not(all(feature = "tiktoken", feature = "huggingface")))] +use crate::errors::RustBridgeDeclined; +use crate::routes::token_counter::token_count_error_to_pyerr; + +#[cfg(feature = "huggingface")] +use litellm_token_counter::huggingface::{ + EncodeInput, Encoding, HuggingFaceTokenizer, InputSequence, PaddingDirection, PaddingStrategy, + TruncationDirection, encoding_from_json, encoding_to_json, +}; +#[cfg(feature = "tiktoken")] +use litellm_token_counter::tiktoken::{TiktokenTokenizer, Vocabulary}; + +#[cfg(feature = "tiktoken")] +pub(crate) fn load_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { + enter_native()?; + let resource: std::path::PathBuf = + PyModule::import(py, "litellm.litellm_core_utils.tokenizers")? + .getattr("__file__")? + .extract()?; + release_gil(py, || { + TiktokenTokenizer::from_cached_ranks(encoding, |file| { + std::fs::read_to_string(resource.with_file_name(file)) + }) + }) + .map_err(|error| token_count_error_to_pyerr(error.into())) +} + +pub(crate) enum Codec { + #[cfg(feature = "tiktoken")] + Tiktoken(TiktokenTokenizer), + #[cfg(feature = "huggingface")] + HuggingFace(HuggingFaceTokenizer), +} + +impl Codec { + pub(crate) fn codec(&self) -> &dyn TextCodec { + match *self { + #[cfg(feature = "tiktoken")] + Self::Tiktoken(ref tokenizer) => tokenizer, + #[cfg(feature = "huggingface")] + Self::HuggingFace(ref tokenizer) => tokenizer, + } + } + + #[cfg(feature = "fast")] + fn fast_counter(&self) -> Option { + match *self { + #[cfg(feature = "tiktoken")] + Self::Tiktoken(ref tokenizer) => tokenizer.fast_counter(), + #[cfg(feature = "huggingface")] + Self::HuggingFace(ref tokenizer) => tokenizer.fast_counter(), + } + } +} + +/// The loaded model is shared: `TokenCounter::from_tokenizer` counts with the same parse, +/// and the opt-in count-only counter is derived from it once, on first use. +#[pyclass(frozen, module = "litellm.rust_bridge._native")] +pub(crate) struct Tokenizer { + inner: Arc, + #[cfg(feature = "fast")] + fast: OnceLock>>, +} + +#[pymethods] +impl Tokenizer { + #[staticmethod] + fn from_tiktoken(py: Python<'_>, encoding: &str) -> PyResult { + #[cfg(feature = "tiktoken")] + { + let tokenizer = load_tiktoken(py, encoding)?; + Ok(Self::new(Codec::Tiktoken(tokenizer))) + } + #[cfg(not(feature = "tiktoken"))] + { + let _ = (py, encoding); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the tiktoken feature", + )) + } + } + + #[staticmethod] + fn from_json(py: Python<'_>, tokenizer_json: &str) -> PyResult { + #[cfg(feature = "huggingface")] + { + enter_native()?; + let tokenizer = release_gil(py, || HuggingFaceTokenizer::from_json(tokenizer_json)) + .map_err(|error| token_count_error_to_pyerr(error.into()))?; + Ok(Self::new(Codec::HuggingFace(tokenizer))) + } + #[cfg(not(feature = "huggingface"))] + { + let _ = (py, tokenizer_json); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the huggingface feature", + )) + } + } + + #[staticmethod] + #[pyo3(signature = (identifier, revision = "main", token = None))] + fn from_pretrained( + py: Python<'_>, + identifier: &str, + revision: &str, + token: Option<&str>, + ) -> PyResult { + #[cfg(feature = "huggingface")] + { + enter_native()?; + let kwargs = PyDict::new(py); + kwargs.set_item("repo_id", identifier)?; + kwargs.set_item("filename", "tokenizer.json")?; + kwargs.set_item("revision", revision)?; + kwargs.set_item("token", token)?; + let path: String = PyModule::import(py, "huggingface_hub")? + .getattr("hf_hub_download")? + .call((), Some(&kwargs))? + .extract()?; + let json = + release_gil(py, || std::fs::read_to_string(path)).map_err(PyIOError::new_err)?; + Self::from_json(py, &json) + } + #[cfg(not(feature = "huggingface"))] + { + let _ = (py, identifier, revision, token); + Err(RustBridgeDeclined::new_err( + "tokenizer backend requires the huggingface feature", + )) + } + } + + fn encode(&self, py: Python<'_>, text: &Bound<'_, PyString>) -> PyResult> { + enter_native()?; + let text = self.text(text)?; + release_gil(py, || self.inner.codec().encode(&text)).map_err(token_count_error_to_pyerr) + } + + #[pyo3(signature = (ids, skip_special_tokens = true))] + fn decode(&self, py: Python<'_>, ids: Vec, skip_special_tokens: bool) -> PyResult { + enter_native()?; + release_gil(py, || self.inner.codec().decode(&ids, skip_special_tokens)) + .map_err(token_count_error_to_pyerr) + } + + #[pyo3(signature = (text, fast = false))] + fn count(&self, py: Python<'_>, text: &Bound<'_, PyString>, fast: bool) -> PyResult { + enter_native()?; + let text = self.text(text)?; + let counter = self.counter(py, fast); + release_gil(py, || { + litellm_token_counter::Tokenizer::count_tokens(&counter, &text) + }) + .map_err(token_count_error_to_pyerr) + } + + #[getter] + fn name(&self) -> &str { + self.inner.codec().name() + } + + // ---- tiktoken: the `tiktoken.Encoding` surface ------------------------------------------ + + #[cfg(feature = "tiktoken")] + fn encode_special( + &self, + py: Python<'_>, + text: &Bound<'_, PyString>, + allowed: Vec, + ) -> PyResult> { + enter_native()?; + let tokenizer = self.tiktoken()?; + let text = self.text(text)?; + release_gil(py, || tokenizer.encode_special(&text, &allowed)) + .map_err(PyRuntimeError::new_err) + } + + /// tiktoken's `encode_with_unstable`: `(stable_tokens, completions)`. + #[cfg(feature = "tiktoken")] + fn encode_with_unstable( + &self, + py: Python<'_>, + text: &Bound<'_, PyString>, + allowed: Vec, + ) -> PyResult<(Vec, Vec>)> { + enter_native()?; + let tokenizer = self.tiktoken()?; + let text = self.text(text)?; + Ok(release_gil(py, || { + tokenizer.encode_with_unstable(&text, &allowed) + })) + } + + /// The special tokens by text: tiktoken's `_special_tokens`. + #[cfg(feature = "tiktoken")] + fn special_tokens(&self) -> PyResult> { + Ok(self + .vocabulary()? + .special_tokens() + .map(|(token, rank)| (token.to_owned(), rank)) + .collect()) + } + + #[cfg(feature = "tiktoken")] + fn max_token_value(&self) -> PyResult { + Ok(self.vocabulary()?.max_token_value()) + } + + #[cfg(feature = "tiktoken")] + fn is_special_token(&self, token: u32) -> PyResult { + Ok(self.vocabulary()?.is_special_token(token)) + } + + /// Every mergeable token's bytes, sorted bytewise like tiktoken's `token_byte_values`. + #[cfg(feature = "tiktoken")] + fn token_byte_values<'py>(&self, py: Python<'py>) -> PyResult>> { + let vocabulary = self.vocabulary()?; + let values = release_gil(py, || vocabulary.token_byte_values()); + Ok(values.iter().map(|value| PyBytes::new(py, value)).collect()) + } + + /// The token of one whole piece; `KeyError` when it is not in the vocabulary. + #[cfg(feature = "tiktoken")] + fn encode_single_token(&self, py: Python<'_>, piece: Vec) -> PyResult { + self.vocabulary()? + .encode_single_token(&piece) + .ok_or_else(|| PyKeyError::new_err(PyBytes::new(py, &piece).unbind())) + } + + #[cfg(feature = "tiktoken")] + fn decode_bytes<'py>(&self, py: Python<'py>, ids: Vec) -> PyResult> { + enter_native()?; + let tokenizer = self.tiktoken()?; + let bytes = + release_gil(py, || tokenizer.decode_bytes(&ids)).map_err(PyKeyError::new_err)?; + Ok(PyBytes::new(py, &bytes)) + } + + // ---- Hugging Face: the `tokenizers.Tokenizer` surface ----------------------------------- + + #[cfg(feature = "huggingface")] + #[pyo3(signature = (sequence, pair = None, is_pretokenized = false, add_special_tokens = true, fast = false))] + fn encode_huggingface( + &self, + py: Python<'_>, + sequence: Sequence, + pair: Option, + is_pretokenized: bool, + add_special_tokens: bool, + fast: bool, + ) -> PyResult { + enter_native()?; + let tokenizer = self.huggingface()?; + let sequence = sequence.input(is_pretokenized)?; + let input = match pair { + Some(pair) => EncodeInput::Dual(sequence, pair.input(is_pretokenized)?), + None => EncodeInput::Single(sequence), + }; + release_gil(py, || { + tokenizer.encode_result(input, add_special_tokens, fast) + }) + .map(|inner| HuggingFaceEncoding { inner }) + .map_err(|error| token_count_error_to_pyerr(Error::from(error))) + } + + #[cfg(feature = "huggingface")] + #[pyo3(signature = (inputs, is_pretokenized = false, add_special_tokens = true, fast = false))] + fn encode_batch_huggingface( + &self, + py: Python<'_>, + inputs: Vec<(Sequence, Option)>, + is_pretokenized: bool, + add_special_tokens: bool, + fast: bool, + ) -> PyResult> { + enter_native()?; + let tokenizer = self.huggingface()?; + let inputs = inputs + .into_iter() + .map(|(sequence, pair)| { + let sequence = sequence.input(is_pretokenized)?; + match pair { + Some(pair) => Ok(EncodeInput::Dual(sequence, pair.input(is_pretokenized)?)), + None => Ok(EncodeInput::Single(sequence)), + } + }) + .collect::>>()?; + release_gil(py, || { + tokenizer.encode_batch_result(inputs, add_special_tokens, fast) + }) + .map(|encodings| { + encodings + .into_iter() + .map(|inner| HuggingFaceEncoding { inner }) + .collect() + }) + .map_err(|error| token_count_error_to_pyerr(Error::from(error))) + } + + #[cfg(feature = "huggingface")] + #[pyo3(signature = (pretty = false))] + fn to_json(&self, py: Python<'_>, pretty: bool) -> PyResult { + enter_native()?; + let tokenizer = self.huggingface()?; + release_gil(py, || tokenizer.to_json(pretty)) + .map_err(|error| token_count_error_to_pyerr(Error::from(error))) + } + + #[cfg(feature = "huggingface")] + fn token_to_id(&self, token: &str) -> PyResult> { + Ok(self.huggingface()?.token_to_id(token)) + } + + #[cfg(feature = "huggingface")] + fn id_to_token(&self, id: u32) -> PyResult> { + Ok(self.huggingface()?.id_to_token(id)) + } + + #[cfg(feature = "huggingface")] + #[pyo3(signature = (with_added_tokens = true))] + fn get_vocab(&self, py: Python<'_>, with_added_tokens: bool) -> PyResult> { + let tokenizer = self.huggingface()?; + Ok(release_gil(py, || tokenizer.vocab(with_added_tokens))) + } + + #[cfg(feature = "huggingface")] + #[pyo3(signature = (with_added_tokens = true))] + fn get_vocab_size(&self, with_added_tokens: bool) -> PyResult { + Ok(self.huggingface()?.vocab_size(with_added_tokens)) + } + + /// The added tokens by id as `(id, (content, single_word, lstrip, rstrip, normalized, + /// special))`, for Python to rebuild as `tokenizers.AddedToken`. + #[cfg(feature = "huggingface")] + fn added_tokens_decoder(&self) -> PyResult> { + Ok(self + .huggingface()? + .added_tokens_decoder() + .into_iter() + .map(|(id, token)| { + ( + id, + ( + token.content, + token.single_word, + token.lstrip, + token.rstrip, + token.normalized, + token.special, + ), + ) + }) + .collect()) + } + + /// The padding parameters as `tokenizers.Tokenizer.padding` reports them. + #[cfg(feature = "huggingface")] + fn padding<'py>(&self, py: Python<'py>) -> PyResult>> { + let Some(params) = self.huggingface()?.padding() else { + return Ok(None); + }; + let padding = PyDict::new(py); + padding.set_item( + "length", + match params.strategy { + PaddingStrategy::BatchLongest => None, + PaddingStrategy::Fixed(length) => Some(length), + }, + )?; + padding.set_item("pad_to_multiple_of", params.pad_to_multiple_of)?; + padding.set_item("pad_id", params.pad_id)?; + padding.set_item("pad_type_id", params.pad_type_id)?; + padding.set_item("pad_token", ¶ms.pad_token)?; + padding.set_item("direction", params.direction.as_ref())?; + Ok(Some(padding)) + } + + /// The truncation parameters as `tokenizers.Tokenizer.truncation` reports them. + #[cfg(feature = "huggingface")] + fn truncation<'py>(&self, py: Python<'py>) -> PyResult>> { + let Some(params) = self.huggingface()?.truncation() else { + return Ok(None); + }; + let truncation = PyDict::new(py); + truncation.set_item("max_length", params.max_length)?; + truncation.set_item("stride", params.stride)?; + truncation.set_item("strategy", params.strategy.as_ref())?; + truncation.set_item("direction", params.direction.as_ref())?; + Ok(Some(truncation)) + } + + #[cfg(feature = "huggingface")] + fn num_special_tokens_to_add(&self, is_pair: bool) -> PyResult { + Ok(self.huggingface()?.num_special_tokens_to_add(is_pair)) + } + + #[cfg(feature = "huggingface")] + fn encode_special_tokens(&self) -> PyResult { + Ok(self.huggingface()?.encode_special_tokens()) + } +} + +#[cfg(feature = "huggingface")] +type AddedTokenFields = (String, bool, bool, bool, bool, bool); + +impl Tokenizer { + fn new(inner: Codec) -> Self { + Self { + inner: Arc::new(inner), + #[cfg(feature = "fast")] + fast: OnceLock::new(), + } + } + + pub(crate) fn counter(&self, py: Python<'_>, fast: bool) -> SharedCounter { + #[cfg(feature = "fast")] + if fast { + let counter = self.fast.get().unwrap_or_else(|| { + release_gil(py, || { + self.fast + .get_or_init(|| self.inner.fast_counter().map(Arc::new)) + }) + }); + if let Some(counter) = counter { + return SharedCounter::Fast(Arc::clone(counter)); + } + } + #[cfg(not(feature = "fast"))] + let _ = (py, fast); + SharedCounter::Codec(Arc::clone(&self.inner)) + } + + /// A Python `str` as UTF-8. tiktoken replaces lone surrogates the way its Python `encode` + /// does; `tokenizers` rejects them, so that backend keeps the encode error. + fn text<'a>(&self, text: &'a Bound<'_, PyString>) -> PyResult> { + match text.to_cow() { + Ok(text) => Ok(text), + Err(error) => match *self.inner { + #[cfg(feature = "tiktoken")] + Codec::Tiktoken(_) if error.is_instance_of::(text.py()) => { + text.call_method1("encode", ("utf-16", "surrogatepass"))? + .call_method1("decode", ("utf-16", "replace"))? + .extract::() + .map(Cow::Owned) + } + _ => Err(error), + }, + } + } + + #[cfg(feature = "tiktoken")] + fn tiktoken(&self) -> PyResult<&TiktokenTokenizer> { + match *self.inner { + Codec::Tiktoken(ref tokenizer) => Ok(tokenizer), + #[cfg(feature = "huggingface")] + Codec::HuggingFace(_) => Err(PyValueError::new_err("requires a tiktoken encoding")), + } + } + + #[cfg(feature = "tiktoken")] + fn vocabulary(&self) -> PyResult<&Vocabulary> { + self.tiktoken()?.vocabulary().ok_or_else(|| { + PyRuntimeError::new_err("this encoding was built without its vocabulary") + }) + } + + #[cfg(feature = "huggingface")] + fn huggingface(&self) -> PyResult<&HuggingFaceTokenizer> { + match *self.inner { + Codec::HuggingFace(ref tokenizer) => Ok(tokenizer), + #[cfg(feature = "tiktoken")] + Codec::Tiktoken(_) => Err(PyValueError::new_err("requires a Hugging Face tokenizer")), + } + } +} + +pub(crate) enum SharedCounter { + Codec(Arc), + #[cfg(feature = "fast")] + Fast(Arc), +} + +impl litellm_token_counter::Tokenizer for SharedCounter { + fn count_tokens(&self, text: &str) -> Result { + match self { + Self::Codec(codec) => codec.codec().count_tokens(text), + #[cfg(feature = "fast")] + Self::Fast(counter) => counter.count_tokens(text).map_err(Error::from), + } + } +} + +#[cfg(feature = "huggingface")] +#[derive(FromPyObject)] +pub(crate) enum Sequence { + Text(String), + Words(Vec), +} + +#[cfg(feature = "huggingface")] +impl Sequence { + fn input(self, is_pretokenized: bool) -> PyResult> { + match (self, is_pretokenized) { + (Self::Text(text), false) => Ok(text.into()), + (Self::Words(words), true) => Ok(words.into()), + _ => Err(pyo3::exceptions::PyTypeError::new_err( + "input must match is_pretokenized", + )), + } + } +} + +#[cfg(feature = "huggingface")] +fn direction(value: &str, left: T, right: T, what: &str) -> PyResult { + match value { + "left" => Ok(left), + "right" => Ok(right), + other => Err(PyValueError::new_err(format!( + "invalid {what} direction {other:?}: expected 'left' or 'right'" + ))), + } +} + +/// `tokenizers.Encoding`, mutable like the original: `pad`, `truncate` and `set_sequence_id` +/// change it in place. +#[cfg(feature = "huggingface")] +#[pyclass(module = "litellm.rust_bridge._native")] +pub(crate) struct HuggingFaceEncoding { + inner: Encoding, +} + +#[cfg(feature = "huggingface")] +#[pymethods] +impl HuggingFaceEncoding { + #[new] + #[pyo3(signature = (json = None))] + fn new(json: Option<&str>) -> PyResult { + let inner = match json { + Some(json) => encoding_from_json(json) + .map_err(|error| PyValueError::new_err(error.to_string()))?, + None => Encoding::default(), + }; + Ok(Self { inner }) + } + + #[staticmethod] + #[pyo3(signature = (encodings, growing_offsets = true))] + fn merge(encodings: Vec>, growing_offsets: bool) -> Self { + Self { + inner: Encoding::merge( + encodings.iter().map(|encoding| encoding.inner.clone()), + growing_offsets, + ), + } + } + + fn __reduce__<'py>( + &self, + py: Python<'py>, + ) -> PyResult<(Bound<'py, pyo3::types::PyType>, (String,))> { + let json = encoding_to_json(&self.inner) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + Ok((py.get_type::(), (json,))) + } + + fn __repr__(&self) -> String { + format!( + "Encoding(num_tokens={}, attributes=[ids, type_ids, tokens, offsets, \ + attention_mask, special_tokens_mask, overflowing])", + self.inner.len() + ) + } + + fn __len__(&self) -> usize { + self.inner.len() + } + #[getter] + fn ids(&self) -> Vec { + self.inner.get_ids().to_vec() + } + #[getter] + fn tokens(&self) -> Vec { + self.inner.get_tokens().to_vec() + } + #[getter] + fn offsets(&self) -> Vec<(usize, usize)> { + self.inner.get_offsets().to_vec() + } + #[getter] + fn type_ids(&self) -> Vec { + self.inner.get_type_ids().to_vec() + } + #[getter] + fn attention_mask(&self) -> Vec { + self.inner.get_attention_mask().to_vec() + } + #[getter] + fn special_tokens_mask(&self) -> Vec { + self.inner.get_special_tokens_mask().to_vec() + } + #[getter] + fn word_ids(&self) -> Vec> { + self.inner.get_word_ids().to_vec() + } + #[getter] + fn sequence_ids(&self) -> Vec> { + self.inner.get_sequence_ids() + } + #[getter] + fn overflowing(&self) -> Vec { + self.inner + .get_overflowing() + .iter() + .cloned() + .map(|inner| Self { inner }) + .collect() + } + #[getter] + fn n_sequences(&self) -> usize { + self.inner.n_sequences() + } + + #[pyo3(signature = (word_index, sequence_index = 0))] + fn word_to_tokens(&self, word_index: u32, sequence_index: usize) -> Option<(usize, usize)> { + self.inner.word_to_tokens(word_index, sequence_index) + } + #[pyo3(signature = (word_index, sequence_index = 0))] + fn word_to_chars(&self, word_index: u32, sequence_index: usize) -> Option<(usize, usize)> { + self.inner.word_to_chars(word_index, sequence_index) + } + fn token_to_sequence(&self, token_index: usize) -> Option { + self.inner.token_to_sequence(token_index) + } + fn token_to_chars(&self, token_index: usize) -> Option<(usize, usize)> { + self.inner + .token_to_chars(token_index) + .map(|(_, offsets)| offsets) + } + fn token_to_word(&self, token_index: usize) -> Option { + self.inner.token_to_word(token_index).map(|(_, word)| word) + } + #[pyo3(signature = (char_pos, sequence_index = 0))] + fn char_to_token(&self, char_pos: usize, sequence_index: usize) -> Option { + self.inner.char_to_token(char_pos, sequence_index) + } + #[pyo3(signature = (char_pos, sequence_index = 0))] + fn char_to_word(&self, char_pos: usize, sequence_index: usize) -> Option { + self.inner.char_to_word(char_pos, sequence_index) + } + + fn set_sequence_id(&mut self, sequence_id: usize) { + self.inner.set_sequence_id(sequence_id); + } + + #[pyo3(signature = (length, direction = "right", pad_id = 0, pad_type_id = 0, pad_token = "[PAD]"))] + fn pad( + &mut self, + length: usize, + direction: &str, + pad_id: u32, + pad_type_id: u32, + pad_token: &str, + ) -> PyResult<()> { + let direction = self::direction( + direction, + PaddingDirection::Left, + PaddingDirection::Right, + "padding", + )?; + self.inner + .pad(length, pad_id, pad_type_id, pad_token, direction); + Ok(()) + } + + #[pyo3(signature = (max_length, stride = 0, direction = "right"))] + fn truncate(&mut self, max_length: usize, stride: usize, direction: &str) -> PyResult<()> { + let direction = self::direction( + direction, + TruncationDirection::Left, + TruncationDirection::Right, + "truncation", + )?; + self.inner.truncate(max_length, stride, direction); + Ok(()) + } +} diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs index 36d319311a3..44acf512224 100644 --- a/litellm-rust/crates/secrets-types/src/config.rs +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -4,7 +4,7 @@ use serde::{Deserialize, Serialize}; use crate::SecretValue; -#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] #[serde(rename_all = "snake_case")] pub enum KeyManagementSystem { GoogleKms, @@ -18,7 +18,7 @@ pub enum KeyManagementSystem { Custom, } -#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] #[serde(rename_all = "snake_case")] pub enum AccessMode { #[default] @@ -33,7 +33,7 @@ impl AccessMode { } } -#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)] #[serde(default)] pub struct KeyManagementSettings { pub hosted_keys: Option>, diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index 1be0adc2bf5..de325ff4981 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -24,6 +24,8 @@ pub enum Error { OidcFile, #[error("secret cannot be converted to {expected}")] TypeMismatch { expected: &'static str }, + #[error("external secret manager failed")] + ExternalManager(#[source] Box), #[cfg(feature = "aws")] #[error(transparent)] Aws(#[from] litellm_secrets_aws::Error), diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 5ab2caa75d4..8762b8e7405 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -1,10 +1,24 @@ +use std::{future::Future, pin::Pin, sync::Arc}; + use litellm_core_utils::settings::Lookup; use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue}; +pub trait ExternalSecretManager: Send + Sync { + fn system(&self) -> KeyManagementSystem; + + fn read_secret<'a>( + &'a self, + name: &'a str, + settings: &'a KeyManagementSettings, + environment: &'a (dyn Lookup + Send + Sync), + ) -> Pin, Error>> + Send + 'a>>; +} + #[derive(Clone)] pub enum SecretManager { Local, + External(Arc), #[cfg(feature = "aws")] AwsKms(crate::aws::AwsKms), #[cfg(feature = "aws")] @@ -25,6 +39,7 @@ impl SecretManager { pub fn system(&self) -> KeyManagementSystem { match self { Self::Local => KeyManagementSystem::Local, + Self::External(manager) => manager.system(), #[cfg(feature = "aws")] Self::AwsKms(_) => KeyManagementSystem::AwsKms, #[cfg(feature = "aws")] @@ -54,6 +69,11 @@ pub async fn get_secret_from_manager( .get(secret_name) .map(SecretValue::new) .map(Secret::String)), + SecretManager::External(manager) => { + manager + .read_secret(secret_name, _settings, environment) + .await + } #[cfg(feature = "aws")] SecretManager::AwsKms(client) => { let ciphertext = environment diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index 1acb5269e66..58aba8494fd 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -7,7 +7,7 @@ mod resolver; mod state; pub use error::Error; -pub use handler::{SecretManager, get_secret_from_manager}; +pub use handler::{ExternalSecretManager, SecretManager, get_secret_from_manager}; pub use litellm_secrets_types::{ AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, }; diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs index 89439893852..597ca11b171 100644 --- a/litellm-rust/crates/secrets/src/resolver.rs +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -72,6 +72,7 @@ impl SecretResolver { Ok(value) => Ok(value .or_else(|| self.environment_secret(name)) .or(default_value)), + Err(error @ Error::ExternalManager(_)) => Err(error), Err(error) => match self.failure_policy { FailurePolicy::Propagate => Err(error), FailurePolicy::EnvironmentFallback => self diff --git a/litellm-rust/crates/token-counter-fast/src/lib.rs b/litellm-rust/crates/token-counter-fast/src/lib.rs index ce91af642ea..157ee847658 100644 --- a/litellm-rust/crates/token-counter-fast/src/lib.rs +++ b/litellm-rust/crates/token-counter-fast/src/lib.rs @@ -8,6 +8,8 @@ mod scanner; mod tiktoken; mod unicode_classes; +use std::sync::Arc; + use byte_level::ByteLevelCounter; use scanner::{SplitPattern, TiktokenCounter}; @@ -15,22 +17,29 @@ pub use error::Error; enum Encoder { HuggingFace { - tokenizer: Box, + tokenizer: Arc, byte_level: Option, }, Tiktoken(TiktokenCounter), } +/// A count-only tokenizer. Its model tables are immutable, so one built from an already +/// loaded model (`from_shared`, `from_*_pairs`) adds only the count-specific tables. pub struct FastTokenizer(Encoder); impl FastTokenizer { pub fn from_json(json: &str) -> Result { let tokenizer = json.parse::().map_err(Error::Load)?; + Ok(Self::from_shared(Arc::new(tokenizer))) + } + + /// Counts with a Hugging Face model another codec already holds; nothing is re-parsed. + pub fn from_shared(tokenizer: Arc) -> Self { let byte_level = ByteLevelCounter::detect(&tokenizer); - Ok(Self(Encoder::HuggingFace { - tokenizer: Box::new(tokenizer), + Self(Encoder::HuggingFace { + tokenizer, byte_level, - })) + }) } pub fn from_cl100k_ranks(ranks: &str) -> Result { @@ -41,12 +50,36 @@ impl FastTokenizer { Self::from_ranks(SplitPattern::O200k, ranks) } + /// `cl100k_base` from ranks another loader already parsed. + pub fn from_cl100k_pairs<'a>( + pairs: impl IntoIterator, + ) -> Result { + Self::from_pairs(SplitPattern::Cl100k, pairs) + } + + /// `o200k_base` (and `o200k_harmony`, whose ordinary tokens are the same) from ranks + /// another loader already parsed. + pub fn from_o200k_pairs<'a>( + pairs: impl IntoIterator, + ) -> Result { + Self::from_pairs(SplitPattern::O200k, pairs) + } + fn from_ranks(split: SplitPattern, ranks: &str) -> Result { TiktokenCounter::from_ranks(split, ranks) .map(Encoder::Tiktoken) .map(Self) } + fn from_pairs<'a>( + split: SplitPattern, + pairs: impl IntoIterator, + ) -> Result { + TiktokenCounter::from_pairs(split, pairs) + .map(Encoder::Tiktoken) + .map(Self) + } + pub fn count_tokens(&self, text: &str) -> Result { match &self.0 { Encoder::Tiktoken(counter) => Ok(counter.count(text)), diff --git a/litellm-rust/crates/token-counter-fast/src/scanner.rs b/litellm-rust/crates/token-counter-fast/src/scanner.rs index c2c3057aeeb..882ea91db81 100644 --- a/litellm-rust/crates/token-counter-fast/src/scanner.rs +++ b/litellm-rust/crates/token-counter-fast/src/scanner.rs @@ -39,8 +39,19 @@ pub(super) struct TiktokenCounter { impl TiktokenCounter { pub(super) fn from_ranks(split: SplitPattern, rank_file: &str) -> Result { + Self::new(split, MergeRanks::parse(rank_file)?) + } + + pub(super) fn from_pairs<'a>( + split: SplitPattern, + pairs: impl IntoIterator, + ) -> Result { + Self::new(split, MergeRanks::from_pairs(pairs)?) + } + + fn new(split: SplitPattern, ranks: MergeRanks) -> Result { Ok(Self { - ranks: MergeRanks::parse(rank_file)?, + ranks, piece_len: split.piece_len(), unicode_classes: UnicodeClasses::get().ok_or(Error::UnicodeClasses)?, }) diff --git a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs index 16172b7a688..68f09b14a25 100644 --- a/litellm-rust/crates/token-counter-fast/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter-fast/src/tiktoken.rs @@ -22,11 +22,25 @@ pub(super) struct MergeRanks(FxHashMap, Rank>); impl MergeRanks { pub(super) fn parse(text: &str) -> Result { - let ranks = text - .lines() - .filter(|line| !line.is_empty()) - .map(parse_line) - .collect::, _>>()?; + Self::from_entries(text.lines().filter(|line| !line.is_empty()).map(parse_line)) + } + + /// The same table from ranks another loader already parsed. + pub(super) fn from_pairs<'a>( + pairs: impl IntoIterator, + ) -> Result { + Self::from_entries(pairs.into_iter().map(|(bytes, rank)| { + if rank == NO_RANK { + return Err(Error::Ranks(format!("rank {rank} is reserved"))); + } + Ok((Box::from(bytes), rank)) + })) + } + + fn from_entries( + entries: impl Iterator, Rank), Error>>, + ) -> Result { + let ranks = entries.collect::, _>>()?; if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) { return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token"))); } diff --git a/litellm-rust/crates/token-counter-huggingface/Cargo.toml b/litellm-rust/crates/token-counter-huggingface/Cargo.toml index 6d8cb85e524..a5c2b1bb160 100644 --- a/litellm-rust/crates/token-counter-huggingface/Cargo.toml +++ b/litellm-rust/crates/token-counter-huggingface/Cargo.toml @@ -6,5 +6,6 @@ license.workspace = true repository.workspace = true [dependencies] +serde_json.workspace = true thiserror.workspace = true tokenizers.workspace = true diff --git a/litellm-rust/crates/token-counter-huggingface/src/error.rs b/litellm-rust/crates/token-counter-huggingface/src/error.rs index adc4551886f..e7f6260321b 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/error.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/error.rs @@ -6,4 +6,6 @@ pub enum Error { Load(#[source] tokenizers::Error), #[error("tokenization failed: {0}")] Encode(#[source] tokenizers::Error), + #[error("token decoding failed: {0}")] + Decode(#[source] tokenizers::Error), } diff --git a/litellm-rust/crates/token-counter-huggingface/src/lib.rs b/litellm-rust/crates/token-counter-huggingface/src/lib.rs index 8e05c2cca46..170a36aea05 100644 --- a/litellm-rust/crates/token-counter-huggingface/src/lib.rs +++ b/litellm-rust/crates/token-counter-huggingface/src/lib.rs @@ -2,22 +2,243 @@ mod error; -pub use error::Error; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; -pub struct HuggingFaceTokenizer(Box); +pub use error::Error; +use tokenizers::PostProcessor; +pub use tokenizers::{ + AddedToken, EncodeInput, Encoding, InputSequence, PaddingDirection, PaddingParams, + PaddingStrategy, TruncationDirection, TruncationParams, +}; + +pub fn encoding_from_json(json: &str) -> Result { + serde_json::from_str(json).map_err(|error| Error::Load(error.into())) +} + +pub fn encoding_to_json(encoding: &Encoding) -> Result { + serde_json::to_string(encoding).map_err(|error| Error::Load(error.into())) +} + +pub struct HuggingFaceTokenizer { + tokenizer: Arc, + special_token_ids: HashSet, +} impl HuggingFaceTokenizer { pub fn from_json(json: &str) -> Result { json.parse::() - .map(Box::new) - .map(Self) + .map(Self::new) .map_err(Error::Load) } + fn new(tokenizer: tokenizers::Tokenizer) -> Self { + let special_token_ids: HashSet = tokenizer + .get_added_tokens_decoder() + .into_iter() + .filter_map(|(id, token)| token.special.then_some(id)) + .collect(); + Self { + tokenizer: Arc::new(tokenizer), + special_token_ids, + } + } + + /// The parsed model, for a count-only counter to share instead of parsing it again. + pub fn shared(&self) -> Arc { + Arc::clone(&self.tokenizer) + } + pub fn count_tokens(&self, text: &str) -> Result { - self.0 + self.tokenizer .encode_fast(text, true) .map(|encoding| encoding.len()) .map_err(Error::Encode) } + + pub fn encode(&self, text: &str) -> Result, Error> { + self.tokenizer + .encode_fast(text, true) + .map(|encoding| encoding.get_ids().to_vec()) + .map_err(Error::Encode) + } + + pub fn encode_result<'a>( + &self, + input: EncodeInput<'a>, + add_special_tokens: bool, + fast: bool, + ) -> Result { + if fast { + return self + .tokenizer + .encode_fast(input, add_special_tokens) + .map_err(Error::Encode); + } + self.tokenizer + .encode_char_offsets(input, add_special_tokens) + .map_err(Error::Encode) + } + + pub fn encode_batch_result<'a>( + &self, + inputs: Vec>, + add_special_tokens: bool, + fast: bool, + ) -> Result, Error> { + if fast { + return self + .tokenizer + .encode_batch_fast(inputs, add_special_tokens) + .map_err(Error::Encode); + } + self.tokenizer + .encode_batch_char_offsets(inputs, add_special_tokens) + .map_err(Error::Encode) + } + + pub fn to_json(&self, pretty: bool) -> Result { + self.tokenizer.to_string(pretty).map_err(Error::Load) + } + + pub fn token_to_id(&self, token: &str) -> Option { + self.tokenizer.token_to_id(token) + } + + pub fn id_to_token(&self, id: u32) -> Option { + self.tokenizer.id_to_token(id) + } + + pub fn vocab(&self, with_added_tokens: bool) -> HashMap { + self.tokenizer.get_vocab(with_added_tokens) + } + + pub fn vocab_size(&self, with_added_tokens: bool) -> usize { + self.tokenizer.get_vocab_size(with_added_tokens) + } + + /// The added tokens by id, in id order. + pub fn added_tokens_decoder(&self) -> Vec<(u32, AddedToken)> { + let mut added: Vec<(u32, AddedToken)> = self + .tokenizer + .get_added_tokens_decoder() + .into_iter() + .collect(); + added.sort_unstable_by_key(|(id, _)| *id); + added + } + + pub fn padding(&self) -> Option<&PaddingParams> { + self.tokenizer.get_padding() + } + + pub fn truncation(&self) -> Option<&TruncationParams> { + self.tokenizer.get_truncation() + } + + /// How many special tokens the post-processor adds to a single sequence or a pair. + pub fn num_special_tokens_to_add(&self, is_pair: bool) -> usize { + self.tokenizer + .get_post_processor() + .map_or(0, |processor| processor.added_tokens(is_pair)) + } + + pub fn encode_special_tokens(&self) -> bool { + self.tokenizer.get_encode_special_tokens() + } + + pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result { + if !skip_special_tokens { + return self.tokenizer.decode(ids, false).map_err(Error::Decode); + } + let filtered_ids: Vec = ids + .iter() + .copied() + .filter(|id| !self.special_token_ids.contains(id)) + .collect(); + self.tokenizer + .decode(&filtered_ids, true) + .map_err(Error::Decode) + } + + pub fn name(&self) -> &str { + "huggingface" + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn codecs_round_trip_and_skip_special_tokens() { + let json = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" + )); + let tokenizer = HuggingFaceTokenizer::from_json(json).unwrap(); + let ids = tokenizer.encode("hello").unwrap(); + + assert!(tokenizer.decode(&ids, false).unwrap().contains("")); + assert_eq!(tokenizer.decode(&ids, true).unwrap(), "hello"); + } + + #[test] + fn decode_filters_special_added_tokens() { + let json = r#"{ + "version": "1.0", + "truncation": null, + "padding": null, + "added_tokens": [ + { + "id": 1, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + ], + "normalizer": null, + "pre_tokenizer": {"type": "Whitespace"}, + "post_processor": null, + "decoder": null, + "model": { + "type": "WordLevel", + "vocab": {"": 0, "": 1, "hello": 2}, + "unk_token": "" + } + }"#; + let tokenizer = HuggingFaceTokenizer::from_json(json).unwrap(); + + assert!(!tokenizer.decode(&[1, 2], true).unwrap().contains("")); + assert!(tokenizer.decode(&[1, 2], false).unwrap().contains("")); + } + + #[test] + fn vocabulary_lookups_mirror_the_tokenizers_api() { + let json = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../litellm/litellm_core_utils/tokenizers/anthropic_tokenizer.json" + )); + let tokenizer = HuggingFaceTokenizer::from_json(json).unwrap(); + let ids = tokenizer.encode("hello").unwrap(); + + let token = tokenizer.id_to_token(ids[0]).unwrap(); + assert_eq!(tokenizer.token_to_id(&token), Some(ids[0])); + assert_eq!(tokenizer.id_to_token(u32::MAX), None); + assert_eq!(tokenizer.vocab(true).len(), tokenizer.vocab_size(true)); + assert!(tokenizer.vocab_size(true) >= tokenizer.vocab_size(false)); + let added = tokenizer.added_tokens_decoder(); + assert!(added.windows(2).all(|pair| pair[0].0 < pair[1].0)); + assert!(added.iter().any(|(_, token)| token.special)); + assert!(tokenizer.padding().is_none()); + assert!(tokenizer.truncation().is_none()); + assert!(!tokenizer.encode_special_tokens()); + assert_eq!( + tokenizer.num_special_tokens_to_add(false), + tokenizer.encode("").unwrap().len() + ); + } } diff --git a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml index 494a9233e69..2fb3103e0c8 100644 --- a/litellm-rust/crates/token-counter-tiktoken/Cargo.toml +++ b/litellm-rust/crates/token-counter-tiktoken/Cargo.toml @@ -6,5 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +base64.workspace = true +once_cell = "1.21.3" +rustc-hash = "2.1.3" thiserror.workspace = true tiktoken-rs.workspace = true diff --git a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs index ecdb3946eee..f049e90a1cb 100644 --- a/litellm-rust/crates/token-counter-tiktoken/src/lib.rs +++ b/litellm-rust/crates/token-counter-tiktoken/src/lib.rs @@ -1,27 +1,125 @@ #![forbid(unsafe_code)] mod error; +mod ranks; + +use std::collections::HashSet; pub use error::UnsupportedTokenizer; +pub use ranks::{LoadError, Vocabulary}; -pub struct TiktokenTokenizer(&'static tiktoken_rs::CoreBPE); +pub struct TiktokenTokenizer { + encoder: &'static tiktoken_rs::CoreBPE, + /// Present for encodings built from a rank file; the embedded tiktoken-rs singletons + /// behind [`from_name`](Self::from_name) keep their ranks private. + vocabulary: Option<&'static Vocabulary>, + name: &'static str, +} impl TiktokenTokenizer { + /// Builds `name` from its packaged rank file (read through `load`), once per process. + /// The tokenizer reports the requested name, so `gpt2` stays `gpt2` like tiktoken does. + pub fn from_cached_ranks( + name: &str, + load: impl FnOnce(&str) -> std::io::Result, + ) -> Result { + let (loaded, name) = ranks::load(name, load)?; + Ok(Self { + encoder: &loaded.bpe, + vocabulary: Some(&loaded.vocabulary), + name, + }) + } + + /// The encodings tiktoken-rs embeds, for hosts without the packaged rank files. pub fn from_name(name: &str) -> Result { - let tokenizer = match name { - "cl100k_base" => tiktoken_rs::cl100k_base_singleton(), - "o200k_base" => tiktoken_rs::o200k_base_singleton(), - "o200k_harmony" => tiktoken_rs::o200k_harmony_singleton(), - "p50k_base" => tiktoken_rs::p50k_base_singleton(), - "p50k_edit" => tiktoken_rs::p50k_edit_singleton(), - "r50k_base" | "gpt2" => tiktoken_rs::r50k_base_singleton(), + let (encoder, name) = match name { + "cl100k_base" => (tiktoken_rs::cl100k_base_singleton(), "cl100k_base"), + "o200k_base" => (tiktoken_rs::o200k_base_singleton(), "o200k_base"), + "o200k_harmony" => (tiktoken_rs::o200k_harmony_singleton(), "o200k_harmony"), + "p50k_base" => (tiktoken_rs::p50k_base_singleton(), "p50k_base"), + "p50k_edit" => (tiktoken_rs::p50k_edit_singleton(), "p50k_edit"), + "r50k_base" => (tiktoken_rs::r50k_base_singleton(), "r50k_base"), + "gpt2" => (tiktoken_rs::r50k_base_singleton(), "gpt2"), _ => return Err(UnsupportedTokenizer(name.to_owned())), }; - Ok(Self(tokenizer)) + Ok(Self { + encoder, + vocabulary: None, + name, + }) + } + + pub fn vocabulary(&self) -> Option<&Vocabulary> { + self.vocabulary } pub fn count_tokens(&self, text: &str) -> usize { - self.0.count_ordinary(text) + self.encoder.count_ordinary(text) + } + + pub fn encode(&self, text: &str) -> Vec { + self.encoder.encode_ordinary(text) + } + + pub fn encode_special(&self, text: &str, allowed: &[String]) -> Result, String> { + let allowed = allowed.iter().map(String::as_str).collect(); + self.encoder + .encode(text, &allowed) + .map(|(ids, _)| ids) + .map_err(|error| error.to_string()) + } + + pub fn special_tokens(&self) -> HashSet { + self.encoder + .special_tokens() + .into_iter() + .map(str::to_owned) + .collect() + } + + /// tiktoken's `encode_with_unstable`: the stable prefix of `text`'s tokens and every + /// token sequence the unstable tail could still become, sorted for a stable order. + pub fn encode_with_unstable( + &self, + text: &str, + allowed: &[String], + ) -> (Vec, Vec>) { + let allowed = allowed.iter().map(String::as_str).collect(); + let (stable, completions) = self.encoder._encode_unstable_native(text, &allowed); + let mut completions: Vec> = completions.into_iter().collect(); + completions.sort_unstable(); + (stable, completions) + } + + pub fn decode_bytes(&self, ids: &[u32]) -> Result, String> { + self.encoder + .decode_bytes(ids) + .map_err(|error| error.to_string()) + } + + pub fn decode(&self, ids: &[u32]) -> Result { + self.encoder + .decode_bytes(ids) + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .map_err(|error| error.to_string()) + } + + pub fn name(&self) -> &str { + self.name + } +} + +pub fn encoding_for_model(model: &str) -> Option<&'static str> { + match tiktoken_rs::tokenizer::get_tokenizer(model)? { + tiktoken_rs::tokenizer::Tokenizer::Cl100kBase => Some("cl100k_base"), + tiktoken_rs::tokenizer::Tokenizer::O200kBase => Some("o200k_base"), + tiktoken_rs::tokenizer::Tokenizer::O200kHarmony => Some("o200k_harmony"), + tiktoken_rs::tokenizer::Tokenizer::P50kBase => Some("p50k_base"), + tiktoken_rs::tokenizer::Tokenizer::P50kEdit => Some("p50k_edit"), + tiktoken_rs::tokenizer::Tokenizer::R50kBase | tiktoken_rs::tokenizer::Tokenizer::Gpt2 => { + Some("r50k_base") + } } } @@ -66,5 +164,75 @@ mod tests { panic!("unknown encoding must be rejected"); }; assert_eq!(name, "unknown-encoding"); + assert_eq!(TiktokenTokenizer::from_name("gpt2").unwrap().name(), "gpt2"); + } + + #[test] + fn codecs_round_trip_named_encodings() { + let encodings = [ + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", + "gpt2", + ]; + let texts = ["hello world", "café 漢字 مرحبا 🙂", "line one\nline two"]; + for name in encodings { + let tokenizer = TiktokenTokenizer::from_name(name).unwrap(); + for text in texts { + assert_eq!( + tokenizer.decode(&tokenizer.encode(text)).unwrap(), + text, + "{name}: {text:?}", + ); + } + } + } + + #[test] + fn decoding_token_prefixes_replaces_incomplete_utf8() { + let tokenizer = TiktokenTokenizer::from_name("cl100k_base").unwrap(); + let reference = tiktoken_rs::cl100k_base_singleton(); + let ids = tokenizer.encode("🙂漢字"); + for end in 1..=ids.len() { + let bytes = reference.decode_bytes(&ids[..end]).unwrap(); + assert_eq!( + tokenizer.decode(&ids[..end]).unwrap(), + String::from_utf8_lossy(&bytes), + ); + } + assert!(tokenizer.decode(&[u32::MAX]).is_err()); + } + + #[test] + fn unstable_encoding_prefixes_stay_consistent_with_full_encoding() { + let tokenizer = TiktokenTokenizer::from_name("cl100k_base").unwrap(); + let text = "hello fanta"; + let (stable, completions) = tokenizer.encode_with_unstable(text, &[]); + assert!( + text.as_bytes() + .starts_with(&tokenizer.decode_bytes(&stable).unwrap()) + ); + assert!(!completions.is_empty()); + for completion in &completions { + let mut ids = stable.clone(); + ids.extend(completion); + assert!( + tokenizer + .decode_bytes(&ids) + .unwrap() + .starts_with(text.as_bytes()) + ); + } + assert!(completions.windows(2).all(|pair| pair[0] < pair[1])); + } + + #[test] + fn encoding_for_model_maps_known_models() { + assert_eq!(encoding_for_model("gpt-4o"), Some("o200k_base")); + assert_eq!(encoding_for_model("text-davinci-003"), Some("p50k_base")); + assert_eq!(encoding_for_model("unknown-model"), None); } } diff --git a/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs new file mode 100644 index 00000000000..1f5e5262de5 --- /dev/null +++ b/litellm-rust/crates/token-counter-tiktoken/src/ranks.rs @@ -0,0 +1,340 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +use once_cell::sync::OnceCell; +use rustc_hash::FxHashMap; +use thiserror::Error; +use tiktoken_rs::{CoreBPE, O200K_BASE_PAT_STR, Rank}; + +use crate::UnsupportedTokenizer; + +const CL100K: &str = "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"; +const O200K: &str = "fb374d419588a4632f3f557e76b4b70aebbca790"; +const P50K: &str = "ec7223a39ce59f226a68acc30dc1af2788490e15"; +const LEGACY_PATTERN: &str = + r"'(?:[sdmt]|ll|ve|re)| ?\p{L}++| ?\p{N}++| ?[^\s\p{L}\p{N}]++|\s++$|\s+(?!\S)|\s"; +const CL100K_PATTERN: &str = r"'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}++|\p{N}{1,3}+| ?[^\s\p{L}\p{N}]++[\r\n]*+|\s++$|\s*[\r\n]|\s+(?!\S)|\s"; + +static CL100K_ENCODER: OnceCell = OnceCell::new(); +static O200K_ENCODER: OnceCell = OnceCell::new(); +static HARMONY_ENCODER: OnceCell = OnceCell::new(); +static P50K_ENCODER: OnceCell = OnceCell::new(); +static EDIT_ENCODER: OnceCell = OnceCell::new(); +static R50K_ENCODER: OnceCell = OnceCell::new(); + +/// One encoding built from a rank file: the BPE engine plus the vocabulary it was built +/// from, kept because `CoreBPE` does not expose its ranks and tiktoken's Python API does +/// (`token_byte_values`, `encode_single_token`, `max_token_value`, `_special_tokens`). +pub(super) struct Loaded { + pub(super) bpe: CoreBPE, + pub(super) vocabulary: Vocabulary, +} + +/// The byte-level vocabulary of a tiktoken encoding. +pub struct Vocabulary { + ranks: FxHashMap, Rank>, + special_tokens: FxHashMap, + max_token_value: Rank, +} + +impl Vocabulary { + /// Every mergeable token's bytes, sorted bytewise like tiktoken's `token_byte_values`. + pub fn token_byte_values(&self) -> Vec> { + let mut values: Vec> = self.ranks.keys().cloned().collect(); + values.sort_unstable(); + values + } + + /// The rank of one whole token: a mergeable piece first, then a special token's text. + pub fn encode_single_token(&self, piece: &[u8]) -> Option { + if let Some(rank) = self.ranks.get(piece) { + return Some(*rank); + } + std::str::from_utf8(piece) + .ok() + .and_then(|text| self.special_tokens.get(text).copied()) + } + + pub fn max_token_value(&self) -> Rank { + self.max_token_value + } + + /// Every mergeable token with its rank, for building other tables from one parse. + pub fn ranks(&self) -> impl Iterator + '_ { + self.ranks + .iter() + .map(|(bytes, rank)| (bytes.as_slice(), *rank)) + } + + /// The special tokens with their ranks, tiktoken's `_special_tokens`. + pub fn special_tokens(&self) -> impl Iterator + '_ { + self.special_tokens + .iter() + .map(|(token, rank)| (token.as_str(), *rank)) + } + + pub fn is_special_token(&self, rank: Rank) -> bool { + self.special_tokens.values().any(|special| *special == rank) + } +} + +#[derive(Debug, Error)] +pub enum LoadError { + #[error(transparent)] + Unsupported(#[from] UnsupportedTokenizer), + #[error("failed to load tiktoken ranks: {0}")] + Ranks(String), +} + +/// Loads `name` once per process. The returned name is the one requested (`gpt2` stays +/// `gpt2`, as `tiktoken.get_encoding("gpt2").name` does), while `gpt2` and `r50k_base` share +/// one cached encoder. +pub(super) fn load( + name: &str, + load_file: impl FnOnce(&str) -> std::io::Result, +) -> Result<(&'static Loaded, &'static str), LoadError> { + let (requested, canonical, file, cache) = match name { + "cl100k_base" => ("cl100k_base", "cl100k_base", CL100K, &CL100K_ENCODER), + "o200k_base" => ("o200k_base", "o200k_base", O200K, &O200K_ENCODER), + "o200k_harmony" => ("o200k_harmony", "o200k_harmony", O200K, &HARMONY_ENCODER), + "p50k_base" => ("p50k_base", "p50k_base", P50K, &P50K_ENCODER), + "p50k_edit" => ("p50k_edit", "p50k_edit", P50K, &EDIT_ENCODER), + "r50k_base" => ("r50k_base", "r50k_base", P50K, &R50K_ENCODER), + "gpt2" => ("gpt2", "r50k_base", P50K, &R50K_ENCODER), + _ => return Err(UnsupportedTokenizer(name.to_owned()).into()), + }; + let loaded = cache.get_or_try_init(|| { + let ranks = load_file(file).map_err(|error| LoadError::Ranks(error.to_string()))?; + build(canonical, &ranks) + })?; + Ok((loaded, requested)) +} + +fn build(name: &str, ranks: &str) -> Result { + let parsed = ranks + .lines() + .map(parse_rank) + .collect::, _>>()?; + let encoder: FxHashMap<_, _> = parsed + .into_iter() + .filter(|(_, rank)| name != "r50k_base" || *rank < 50256) + .collect(); + if encoder + .values() + .collect::>() + .len() + != encoder.len() + || (0..=u8::MAX).any(|byte| !encoder.contains_key(&[byte][..])) + { + return Err(LoadError::Ranks("invalid vocabulary ranks".into())); + } + let (pattern, specials): (&str, &[(&str, Rank)]) = match name { + "cl100k_base" => ( + CL100K_PATTERN, + &[ + ("<|endoftext|>", 100257), + ("<|fim_prefix|>", 100258), + ("<|fim_middle|>", 100259), + ("<|fim_suffix|>", 100260), + ("<|endofprompt|>", 100276), + ], + ), + "o200k_base" => ( + O200K_BASE_PAT_STR, + &[("<|endoftext|>", 199999), ("<|endofprompt|>", 200018)], + ), + "o200k_harmony" => ( + O200K_BASE_PAT_STR, + &[ + ("<|startoftext|>", 199998), + ("<|endoftext|>", 199999), + ("<|reserved_200000|>", 200000), + ("<|reserved_200001|>", 200001), + ("<|return|>", 200002), + ("<|constrain|>", 200003), + ("<|reserved_200004|>", 200004), + ("<|channel|>", 200005), + ("<|start|>", 200006), + ("<|end|>", 200007), + ("<|message|>", 200008), + ("<|reserved_200009|>", 200009), + ("<|reserved_200010|>", 200010), + ("<|reserved_200011|>", 200011), + ("<|call|>", 200012), + ], + ), + "p50k_edit" => ( + LEGACY_PATTERN, + &[ + ("<|endoftext|>", 50256), + ("<|fim_prefix|>", 50281), + ("<|fim_middle|>", 50282), + ("<|fim_suffix|>", 50283), + ], + ), + _ => (LEGACY_PATTERN, &[("<|endoftext|>", 50256)]), + }; + let reserved = (200013..=201087) + .filter(|_| name == "o200k_harmony") + .map(|rank| (format!("<|reserved_{rank}|>"), rank)); + let special_tokens: FxHashMap = specials + .iter() + .map(|(token, rank)| ((*token).to_owned(), *rank)) + .chain(reserved) + .collect(); + let max_token_value = encoder + .values() + .chain(special_tokens.values()) + .copied() + .max() + .ok_or_else(|| LoadError::Ranks("empty vocabulary".into()))?; + let bpe = CoreBPE::new(encoder.clone(), special_tokens.clone(), pattern) + .map_err(|error| LoadError::Ranks(error.to_string()))?; + Ok(Loaded { + bpe, + vocabulary: Vocabulary { + ranks: encoder, + special_tokens, + max_token_value, + }, + }) +} + +fn parse_rank(line: &str) -> Result<(Vec, Rank), LoadError> { + let (token, rank) = line + .split_once(' ') + .ok_or_else(|| LoadError::Ranks("missing rank".into()))?; + let bytes = STANDARD + .decode(token) + .map_err(|error| LoadError::Ranks(error.to_string()))?; + let rank = rank + .parse() + .map_err(|error: std::num::ParseIntError| LoadError::Ranks(error.to_string()))?; + Ok((bytes, rank)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::TiktokenTokenizer; + + fn read_packaged_ranks(file: &str) -> std::io::Result { + std::fs::read_to_string( + std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../../../litellm/litellm_core_utils/tokenizers") + .join(file), + ) + } + + #[test] + fn packaged_encodings_match_embedded_encodings_and_reuse_successful_loads() { + for name in [ + "cl100k_base", + "o200k_base", + "o200k_harmony", + "p50k_base", + "p50k_edit", + "r50k_base", + "gpt2", + ] { + if name != "gpt2" { + assert!( + TiktokenTokenizer::from_cached_ranks(name, |_| { + Err(std::io::Error::other("unreadable vocabulary")) + }) + .is_err() + ); + } + let loads = std::sync::atomic::AtomicUsize::new(0); + let barrier = std::sync::Barrier::new(4); + let encoders = std::thread::scope(|scope| { + let tasks: Vec<_> = (0..4) + .map(|_| { + scope.spawn(|| { + barrier.wait(); + TiktokenTokenizer::from_cached_ranks(name, |file| { + loads.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + read_packaged_ranks(file) + }) + .unwrap() + }) + }) + .collect(); + tasks + .into_iter() + .map(|task| task.join().unwrap()) + .collect::>() + }); + assert_eq!(loads.into_inner(), usize::from(name != "gpt2")); + let actual = &encoders[0]; + let expected = TiktokenTokenizer::from_name(name).unwrap(); + assert_eq!(actual.special_tokens(), expected.special_tokens()); + let specials: Vec<_> = expected.special_tokens().into_iter().collect(); + let special_text = specials.join(" "); + assert_eq!( + actual.encode_special(&special_text, &specials).unwrap(), + expected.encode_special(&special_text, &specials).unwrap() + ); + for text in [ + "", + "café 漢字 ع 🙂", + "a\r\nb\t ", + " hello 123456789", + &special_text, + ] { + let ids = expected.encode(text); + assert_eq!(actual.encode(text), ids, "{name}: {text:?}"); + assert_eq!(actual.count_tokens(text), ids.len(), "{name}: {text:?}"); + assert_eq!( + actual.decode_bytes(&ids).unwrap(), + expected.decode_bytes(&ids).unwrap() + ); + } + let cached = TiktokenTokenizer::from_cached_ranks(name, |_| { + panic!("reloaded cached vocabulary") + }) + .unwrap(); + assert_eq!(cached.encode("cached"), expected.encode("cached")); + assert_eq!(cached.name(), name); + assert!(expected.vocabulary().is_none()); + assert_vocabulary_lookups(name, actual); + } + } + + /// The token-level lookups tiktoken's Python `Encoding` exposes, checked against the + /// encoder itself and against the known vocabulary sizes. + fn assert_vocabulary_lookups(name: &str, tokenizer: &TiktokenTokenizer) { + let max_token_value = match name { + "cl100k_base" => 100_276, + "o200k_base" => 200_018, + "o200k_harmony" => 201_087, + "p50k_base" => 50_280, + "p50k_edit" => 50_283, + "r50k_base" | "gpt2" => 50_256, + _ => unreachable!("{name}"), + }; + let vocabulary = tokenizer.vocabulary().unwrap(); + assert_eq!(vocabulary.max_token_value(), max_token_value, "{name}"); + let values = vocabulary.token_byte_values(); + assert!(values.windows(2).all(|pair| pair[0] < pair[1]), "{name}"); + for piece in values.iter().step_by(997) { + let rank = vocabulary.encode_single_token(piece).unwrap(); + assert_eq!(tokenizer.decode_bytes(&[rank]).unwrap(), *piece, "{name}"); + assert!(!vocabulary.is_special_token(rank), "{name}"); + } + for (token, rank) in vocabulary.special_tokens() { + assert_eq!(vocabulary.encode_single_token(token.as_bytes()), Some(rank)); + assert!(vocabulary.is_special_token(rank), "{name}: {token}"); + } + assert_eq!(vocabulary.encode_single_token(b"<|not-a-token|>"), None); + } + + #[test] + fn malformed_ranks_return_errors_instead_of_panicking() { + for ranks in ["", "IQ==", "IQ== x", "!!! 1", "IQ== 1"] { + assert!(build("cl100k_base", ranks).is_err()); + } + let repeated_rank = (0..=u8::MAX) + .map(|byte| format!("{} 0\n", STANDARD.encode([byte]))) + .collect::(); + assert!(build("cl100k_base", &repeated_rank).is_err()); + } +} diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md index a6b2b50aac0..3c6381fcea7 100644 --- a/litellm-rust/crates/token-counter/README.md +++ b/litellm-rust/crates/token-counter/README.md @@ -1,6 +1,10 @@ # Token counting -`Tokenizer` is the text-counting interface. `TokenCounter` applies LiteLLM request, message, and tool accounting using any implementation of that interface +`Tokenizer` is the text-counting interface. `TextCodec` adds encoding, decoding, and a name. `TokenCounter` applies LiteLLM request, message, and tool accounting using any `Tokenizer` + +Counts follow the codec: tiktoken treats special-token spellings as ordinary text, while Hugging Face applies its added tokens, post-processing, padding, and truncation. `fast=True` preserves those semantics and requests acceleration where available. Unsupported configurations use the normal codec, including tiktoken encodings without a scanner and builds without the `fast` feature. Invalid input and process-guard errors still propagate. Runtime request counting currently uses the normal codec; the custom accelerator is retained for explicit use and testing + +`FastCounter: TextCodec` exposes an optional accelerator over a loaded codec. `None` means callers should use that codec. The Python bridge caches this selection per immutable tokenizer, shares it with request counters, and initializes it with the GIL released. Hugging Face can also choose the full encoder per input when added tokens require it The `fast` feature provides `fast::FastTokenizer` from `litellm-token-counter-fast`. `TokenCounter::from_json_fast` uses this implementation @@ -8,7 +12,9 @@ The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through t The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2` -All three backends are enabled by default. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend +All three backends are enabled by default in this crate and the Python extension. With `default-features = false`, Rust callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend + +Python `tiktoken` and `tokenizers` remain runtime dependencies and the default implementations. The catalog independently selects the tokenizer and request-counting routes. Enabling Rust changes factory dispatch; existing tokenizer objects keep their backend. Native Hugging Face wrappers provide an immutable encoding and decoding API, while training and mutable configuration remain available through the Python backend Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits diff --git a/litellm-rust/crates/token-counter/src/error.rs b/litellm-rust/crates/token-counter/src/error.rs index b05ce007e46..6a94b0e39b8 100644 --- a/litellm-rust/crates/token-counter/src/error.rs +++ b/litellm-rust/crates/token-counter/src/error.rs @@ -32,6 +32,8 @@ pub enum Error { JsonUtf8(#[source] FromUtf8Error), #[error("tokenization failed: {0}")] Encode(#[source] Box), + #[error("token decoding failed: {0}")] + Decode(String), #[error("token counting task failed: {0}")] Task(String), } diff --git a/litellm-rust/crates/token-counter/src/fast.rs b/litellm-rust/crates/token-counter/src/fast.rs index de3f86abd68..f601c3cbcca 100644 --- a/litellm-rust/crates/token-counter/src/fast.rs +++ b/litellm-rust/crates/token-counter/src/fast.rs @@ -1,7 +1,31 @@ use litellm_token_counter_fast::Error as BackendError; pub use litellm_token_counter_fast::FastTokenizer; -use crate::{Error, TokenCounter, Tokenizer}; +use crate::{Error, TextCodec, TokenCounter, Tokenizer}; + +pub trait FastCounter: TextCodec { + fn fast_counter(&self) -> Option; +} + +#[cfg(feature = "huggingface")] +impl FastCounter for crate::huggingface::HuggingFaceTokenizer { + fn fast_counter(&self) -> Option { + Some(FastTokenizer::from_shared(self.shared())) + } +} + +#[cfg(feature = "tiktoken")] +impl FastCounter for crate::tiktoken::TiktokenTokenizer { + fn fast_counter(&self) -> Option { + let vocabulary = self.vocabulary()?; + match self.name() { + "cl100k_base" => FastTokenizer::from_cl100k_pairs(vocabulary.ranks()), + "o200k_base" | "o200k_harmony" => FastTokenizer::from_o200k_pairs(vocabulary.ranks()), + _ => return None, + } + .ok() + } +} impl TokenCounter { pub fn from_json_fast(tokenizer_json: &str) -> Result { @@ -39,3 +63,65 @@ impl From for Error { } } } + +#[cfg(all(test, feature = "huggingface", feature = "tiktoken"))] +mod tests { + use super::*; + use crate::huggingface::HuggingFaceTokenizer; + use crate::tiktoken::TiktokenTokenizer; + + const TEXTS: [&str; 4] = [ + "", + "hello world <|endoftext|>", + "café 漢字 ع 🙂 line\r\n indented 123456789", + "system a\u{301} fi", + ]; + + fn packaged(file: &str) -> String { + std::fs::read_to_string(format!( + "{}/../../../litellm/litellm_core_utils/tokenizers/{file}", + env!("CARGO_MANIFEST_DIR") + )) + .unwrap() + } + + #[test] + fn fast_counters_derived_from_codecs_count_like_the_codecs() { + let huggingface = + HuggingFaceTokenizer::from_json(&packaged("anthropic_tokenizer.json")).unwrap(); + let fast = huggingface.fast_counter().unwrap(); + for text in TEXTS { + assert_eq!( + fast.count_tokens(text).unwrap(), + Tokenizer::count_tokens(&huggingface, text).unwrap(), + "{text:?}" + ); + } + + for name in ["cl100k_base", "o200k_base", "o200k_harmony"] { + let tiktoken = + TiktokenTokenizer::from_cached_ranks(name, |file| Ok(packaged(file))).unwrap(); + let fast = tiktoken.fast_counter().unwrap(); + for text in TEXTS { + assert_eq!( + fast.count_tokens(text).unwrap(), + tiktoken.count_tokens(text), + "{name}: {text:?}" + ); + } + } + } + + #[test] + fn encodings_without_a_fast_scanner_keep_the_codec() { + let tiktoken = + TiktokenTokenizer::from_cached_ranks("p50k_base", |file| Ok(packaged(file))).unwrap(); + assert!(tiktoken.fast_counter().is_none()); + assert!( + TiktokenTokenizer::from_name("cl100k_base") + .unwrap() + .fast_counter() + .is_none() + ); + } +} diff --git a/litellm-rust/crates/token-counter/src/huggingface.rs b/litellm-rust/crates/token-counter/src/huggingface.rs index fb7683b373e..43613fac8ea 100644 --- a/litellm-rust/crates/token-counter/src/huggingface.rs +++ b/litellm-rust/crates/token-counter/src/huggingface.rs @@ -1,7 +1,11 @@ use litellm_token_counter_huggingface::Error as BackendError; -pub use litellm_token_counter_huggingface::HuggingFaceTokenizer; +pub use litellm_token_counter_huggingface::{ + AddedToken, EncodeInput, Encoding, HuggingFaceTokenizer, InputSequence, PaddingDirection, + PaddingParams, PaddingStrategy, TruncationDirection, TruncationParams, encoding_from_json, + encoding_to_json, +}; -use crate::{Error, TokenCounter, Tokenizer}; +use crate::{Error, TextCodec, TokenCounter, Tokenizer}; impl TokenCounter { pub fn from_json(tokenizer_json: &str) -> Result { @@ -17,11 +21,26 @@ impl Tokenizer for HuggingFaceTokenizer { } } +impl TextCodec for HuggingFaceTokenizer { + fn encode(&self, text: &str) -> Result, Error> { + HuggingFaceTokenizer::encode(self, text).map_err(Error::from) + } + + fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result { + HuggingFaceTokenizer::decode(self, ids, skip_special_tokens).map_err(Error::from) + } + + fn name(&self) -> &str { + HuggingFaceTokenizer::name(self) + } +} + impl From for Error { fn from(error: BackendError) -> Self { match error { BackendError::Load(source) => Self::Load(source), BackendError::Encode(source) => Self::Encode(source), + BackendError::Decode(source) => Self::Decode(source.to_string()), } } } diff --git a/litellm-rust/crates/token-counter/src/lib.rs b/litellm-rust/crates/token-counter/src/lib.rs index 446c91049de..84bf35ca2ba 100644 --- a/litellm-rust/crates/token-counter/src/lib.rs +++ b/litellm-rust/crates/token-counter/src/lib.rs @@ -20,5 +20,5 @@ pub mod tiktoken; pub use counter::{InputTokenCount, TokenCounter}; pub use error::Error; -pub use tokenizer::Tokenizer; +pub use tokenizer::{TextCodec, Tokenizer}; pub use types::CountableRequest; diff --git a/litellm-rust/crates/token-counter/src/tiktoken.rs b/litellm-rust/crates/token-counter/src/tiktoken.rs index 07c1c9f5b73..5883a629ea5 100644 --- a/litellm-rust/crates/token-counter/src/tiktoken.rs +++ b/litellm-rust/crates/token-counter/src/tiktoken.rs @@ -1,7 +1,7 @@ -pub use litellm_token_counter_tiktoken::TiktokenTokenizer; -use litellm_token_counter_tiktoken::UnsupportedTokenizer; +use litellm_token_counter_tiktoken::{LoadError, UnsupportedTokenizer}; +pub use litellm_token_counter_tiktoken::{TiktokenTokenizer, Vocabulary, encoding_for_model}; -use crate::{Error, TokenCounter, Tokenizer}; +use crate::{Error, TextCodec, TokenCounter, Tokenizer}; impl TokenCounter { pub fn from_tiktoken(encoding: &str) -> Result { @@ -17,8 +17,31 @@ impl Tokenizer for TiktokenTokenizer { } } +impl TextCodec for TiktokenTokenizer { + fn encode(&self, text: &str) -> Result, Error> { + Ok(TiktokenTokenizer::encode(self, text)) + } + + fn decode(&self, ids: &[u32], _skip_special_tokens: bool) -> Result { + TiktokenTokenizer::decode(self, ids).map_err(|error| Error::Decode(error.to_string())) + } + + fn name(&self) -> &str { + TiktokenTokenizer::name(self) + } +} + impl From for Error { fn from(error: UnsupportedTokenizer) -> Self { Self::UnsupportedTokenizer(error.0) } } + +impl From for Error { + fn from(error: LoadError) -> Self { + match error { + LoadError::Unsupported(error) => error.into(), + LoadError::Ranks(message) => Self::Ranks(message), + } + } +} diff --git a/litellm-rust/crates/token-counter/src/tokenizer.rs b/litellm-rust/crates/token-counter/src/tokenizer.rs index 88c29c672a7..146ac5b4d0d 100644 --- a/litellm-rust/crates/token-counter/src/tokenizer.rs +++ b/litellm-rust/crates/token-counter/src/tokenizer.rs @@ -4,6 +4,12 @@ pub trait Tokenizer: Send + Sync { fn count_tokens(&self, text: &str) -> Result; } +pub trait TextCodec: Tokenizer { + fn encode(&self, text: &str) -> Result, Error>; + fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result; + fn name(&self) -> &str; +} + #[cfg(test)] mod tests { use super::*; diff --git a/litellm/__init__.py b/litellm/__init__.py index 44515472648..a044676a843 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -2146,6 +2146,10 @@ if TYPE_CHECKING: from .llms.edenai.videos.transformation import ( EdenAIVideoConfig as EdenAIVideoConfig, ) + from .llms.fal_ai.chat.transformation import ( + FalAIChatConfig as FalAIChatConfig, + FalAIError as FalAIError, + ) from .llms.ovhcloud.chat.transformation import ( OVHCloudChatConfig as OVHCloudChatConfig, ) diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index fe3c7c264ee..29fb46fa125 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -58,7 +58,8 @@ from ._lazy_imports_registry import ( if TYPE_CHECKING: import httpx - from tiktoken import Encoding + + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer def get_litellm_globals() -> dict[str, object]: @@ -89,26 +90,11 @@ def _get_module_level_client_timeout(litellm_globals: Mapping[str, Any]) -> "flo # These are special lazy loaders for things that are used internally # They're separate from the main lazy import system because they have specific use cases -# Lazy loader for default encoding - avoids importing heavy tiktoken library at startup -_default_encoding: "Encoding | None" = None +def _get_default_encoding() -> "Tokenizer": + from litellm.rust_bridge.tokenizer import get_encoding -def _get_default_encoding() -> "Encoding": - """ - Lazily load and cache the default OpenAI encoding. - - This avoids importing `litellm.litellm_core_utils.default_encoding` (and thus tiktoken) - at `litellm` import time. The encoding is cached after the first import. - - This is used internally by utils.py functions that need the encoding but shouldn't - trigger its import during module load. - """ - global _default_encoding - if _default_encoding is None: - from litellm.litellm_core_utils.default_encoding import encoding - - _default_encoding = encoding - return _default_encoding + return get_encoding("cl100k_base") # Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index db4eb8bdb33..9a53273c9d5 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -335,6 +335,8 @@ LLM_CONFIG_NAMES: Final = ( "EdenAITextToSpeechConfig", "EdenAIImageGenerationConfig", "EdenAIVideoConfig", + "FalAIChatConfig", + "FalAIError", "OVHCloudChatConfig", "OVHCloudEmbeddingConfig", "CometAPIEmbeddingConfig", @@ -1251,6 +1253,8 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "EdenAITextToSpeechConfig": (".llms.edenai.text_to_speech.transformation", "EdenAITextToSpeechConfig"), "EdenAIImageGenerationConfig": (".llms.edenai.image_generation.transformation", "EdenAIImageGenerationConfig"), "EdenAIVideoConfig": (".llms.edenai.videos.transformation", "EdenAIVideoConfig"), + "FalAIChatConfig": (".llms.fal_ai.chat.transformation", "FalAIChatConfig"), + "FalAIError": (".llms.fal_ai.chat.transformation", "FalAIError"), "OVHCloudChatConfig": (".llms.ovhcloud.chat.transformation", "OVHCloudChatConfig"), "OVHCloudEmbeddingConfig": ( ".llms.ovhcloud.embedding.transformation", diff --git a/litellm/anthropic_interface/exceptions/exceptions.py b/litellm/anthropic_interface/exceptions/exceptions.py index 91bcf82f455..b48cd2fee6f 100644 --- a/litellm/anthropic_interface/exceptions/exceptions.py +++ b/litellm/anthropic_interface/exceptions/exceptions.py @@ -25,6 +25,7 @@ class AnthropicErrorDetail(TypedDict): type: AnthropicErrorType message: str provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]] + litellm_call_id: NotRequired[ReadOnly[str]] class AnthropicErrorResponse(TypedDict, total=False): diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 058cc8a1579..b99023c07fd 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -12,6 +12,7 @@ import ast import asyncio import json import os +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm @@ -36,6 +37,8 @@ from ._embedding_router import ( ) from .base_cache import BaseCache +_WAIT_FOR_INDEXING: Final = MappingProxyType({"wait": "true"}) + if TYPE_CHECKING: from litellm.router import Router @@ -313,6 +316,7 @@ class QdrantSemanticCache(BaseCache): self.sync_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, + params=_WAIT_FOR_INDEXING, json=data, ) @@ -422,6 +426,7 @@ class QdrantSemanticCache(BaseCache): await self.async_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, + params=_WAIT_FOR_INDEXING, json=data, ) diff --git a/litellm/chat_completions/dispatch.py b/litellm/chat_completions/dispatch.py index d36c0343988..8e4b7f82eb8 100644 --- a/litellm/chat_completions/dispatch.py +++ b/litellm/chat_completions/dispatch.py @@ -4,7 +4,7 @@ from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable from litellm import main -from litellm.rust_bridge.catalog import Context, Delivery, Route +from litellm.rust_bridge.catalog import Delivery, Route, RouteContext from litellm.rust_bridge.chat_completions.entrypoints import ( NATIVE_ACOMPLETION, NATIVE_COMPLETION, @@ -72,8 +72,8 @@ def _public_request( ) -def _context(request: LiteLLMChatCompletionsRequest) -> Context: - return Context( +def _context(request: LiteLLMChatCompletionsRequest) -> RouteContext: + return RouteContext( Route.CHAT_COMPLETIONS, provider=request.custom_llm_provider, model=request.model, diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b317e356e1d..7b02a5ede00 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -115,6 +115,7 @@ from litellm.types.utils import ( LlmProviders, LlmProvidersSet, ModelInfo, + ModelInfoBase, PromptTokensDetailsWrapper, ServiceTier, StandardBuiltInToolsParams, @@ -322,6 +323,48 @@ class OCRPricing(TypedDict, total=False): annotation_cost_per_page: ReadOnly[float | None] +_WALL_CLOCK_PRICED_MODES: Final = frozenset({"chat", "completion", "embedding", "responses"}) + + +def _has_token_or_tiered_pricing(model_info: ModelInfoBase) -> bool: + return ( + (model_info.get("input_cost_per_token") or 0.0) > 0 + or (model_info.get("output_cost_per_token") or 0.0) > 0 + or model_info.get("tiered_pricing") is not None + ) + + +def _bills_wall_clock_seconds(model_info: ModelInfoBase) -> bool: + mode: Final = model_info.get("mode") + return mode is None or mode in _WALL_CLOCK_PRICED_MODES + + +def _per_second_pricing_cost( + model: str, + custom_llm_provider: str | None, + response_time_ms: float | None, +) -> tuple[float, float] | None: + try: + model_info: Final = _cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # the lookup raises plain Exception for an unmapped model + return None + if _has_token_or_tiered_pricing(model_info) or not _bills_wall_clock_seconds(model_info): + return None + input_cost_per_second: Final = model_info.get("input_cost_per_second") + output_cost_per_second: Final = model_info.get("output_cost_per_second") + if input_cost_per_second is None and output_cost_per_second is None: + return None + seconds: Final = (response_time_ms or 0.0) / 1000 + verbose_logger.debug( + "For model=%s - input_cost_per_second: %s; output_cost_per_second: %s; response time: %s", + model, + input_cost_per_second, + output_cost_per_second, + response_time_ms, + ) + return (input_cost_per_second or 0.0) * seconds, (output_cost_per_second or 0.0) * seconds + + def cost_per_token( model: str = "", prompt_tokens: int = 0, @@ -448,9 +491,6 @@ def cost_per_token( if response_cost is not None: return response_cost[0], response_cost[1] - # given - prompt_tokens_cost_usd_dollar: float = 0 - completion_tokens_cost_usd_dollar: float = 0 model_cost_ref: Final = litellm.model_cost # Only callers that explicitly pass `custom_llm_provider` get the # dedup/prefix-join treatment. When provider is omitted, preserve legacy @@ -611,6 +651,14 @@ def cost_per_token( number_of_queries=number_of_queries or 1, optional_params=(getattr(response, "_hidden_params", None) if response else None), ) + elif ( + per_second_cost := _per_second_pricing_cost( + model=model, + custom_llm_provider=custom_llm_provider, + response_time_ms=response_time_ms, + ) + ) is not None: + return per_second_cost elif custom_llm_provider == "vertex_ai": cost_router: Final = google_cost_router( model=model_without_prefix, @@ -685,12 +733,7 @@ def cost_per_token( ) else: model_info: Final = _cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) - - if ( - (model_info.get("input_cost_per_token") or 0.0) > 0 - or (model_info.get("output_cost_per_token") or 0.0) > 0 - or model_info.get("tiered_pricing") is not None - ): + if _has_token_or_tiered_pricing(model_info): return generic_cost_per_token( model=model, usage=usage_block, @@ -698,36 +741,8 @@ def cost_per_token( service_tier=service_tier, data_residency=data_residency, ) - - input_cost_per_second: Final = model_info.get("input_cost_per_second") - if input_cost_per_second is not None and response_time_ms is not None: - verbose_logger.debug( - "For model=%s - input_cost_per_second: %s; response time: %s", - model, - input_cost_per_second, - response_time_ms, - ) - ## COST PER SECOND ## - prompt_tokens_cost_usd_dollar = input_cost_per_second * response_time_ms / 1000 - - output_cost_per_second: Final = model_info.get("output_cost_per_second") - if output_cost_per_second is not None and response_time_ms is not None: - verbose_logger.debug( - "For model=%s - output_cost_per_second: %s; response time: %s", - model, - output_cost_per_second, - response_time_ms, - ) - ## COST PER SECOND ## - completion_tokens_cost_usd_dollar = output_cost_per_second * response_time_ms / 1000 - - verbose_logger.debug( - "Returned custom cost for model=%s - prompt_tokens_cost_usd_dollar: %s, completion_tokens_cost_usd_dollar: %s", - model, - prompt_tokens_cost_usd_dollar, - completion_tokens_cost_usd_dollar, - ) - return prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar + verbose_logger.debug("No per-token, tiered, or per-second pricing for model=%s; cost is 0", model) + return 0.0, 0.0 def get_replicate_completion_pricing(completion_response: dict, total_time=0.0): @@ -948,7 +963,7 @@ def _extract_service_tier(source: object) -> str | None: return None -def _get_usage_object( +def get_usage_object( completion_response: object, ) -> Usage | None: usage_obj: Final = cast( @@ -1222,6 +1237,21 @@ def _split_responses_ws_logging_object_by_service_tier( ) +def _response_time_ms_for_cost( + completion_response: object, + litellm_logging_obj: LitellmLoggingObject | None, + total_time: float | None, +) -> float: + stamped: Final = getattr(completion_response, "_response_ms", None) + if isinstance(stamped, (int, float)): + return float(stamped) + if total_time: + return total_time + if litellm_logging_obj is not None: + return litellm_logging_obj.get_response_ms() + return 0.0 + + def completion_cost( completion_response: object | None = None, model: str | None = None, @@ -1336,7 +1366,7 @@ def completion_cost( cache_creation_input_tokens: int | None = None cache_read_input_tokens: int | None = None audio_transcription_file_duration: float = 0.0 - provider_usage_object: Final = _get_usage_object(completion_response=completion_response) + provider_usage_object: Final = get_usage_object(completion_response=completion_response) cost_per_token_usage_object: Final[Usage | None] = ( _without_provider_stated_cost(provider_usage_object) if custom_pricing else provider_usage_object ) @@ -1443,8 +1473,6 @@ def completion_cost( prompt_tokens_details = _usage.get("prompt_tokens_details") or {} cache_read_input_tokens = prompt_tokens_details.get("cached_tokens", 0) - total_time = getattr(completion_response, "_response_ms", 0) - hidden_params = getattr(completion_response, "_hidden_params", None) if hidden_params is not None: custom_llm_provider = hidden_params.get("custom_llm_provider", custom_llm_provider or None) @@ -1676,6 +1704,11 @@ def completion_cost( ) return MCPCostCalculator.calculate_mcp_tool_call_cost(litellm_logging_obj=litellm_logging_obj) + response_time_ms = _response_time_ms_for_cost( + completion_response=completion_response, + litellm_logging_obj=litellm_logging_obj, + total_time=total_time, + ) # Calculate cost based on prompt_tokens, completion_tokens if ( "togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai" @@ -1686,7 +1719,7 @@ def completion_cost( # see https://replicate.com/pricing elif (model in litellm.replicate_models or "replicate" in model) and model not in litellm.model_cost: # for unmapped replicate model, default to replicate's time tracking logic - return get_replicate_completion_pricing(completion_response, total_time) + return get_replicate_completion_pricing(completion_response, response_time_ms) if model is None: raise ValueError( @@ -1718,7 +1751,7 @@ def completion_cost( prompt_tokens=prompt_tokens or 0, completion_tokens=completion_tokens or 0, custom_llm_provider=custom_llm_provider, - response_time_ms=total_time, + response_time_ms=response_time_ms, region_name=None if explicit_pricing else region_name, custom_cost_per_second=custom_cost_per_second, custom_cost_per_token=custom_cost_per_token, @@ -2033,6 +2066,45 @@ def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelIn return None +def _raw_cost_map_entry(key: str) -> Mapping[str, object] | None: + raw_entry: Final = litellm.model_cost.get(key) + return raw_entry if isinstance(raw_entry, Mapping) else None + + +def pricing_entry_for_cost_calc( + model: str | None, + completion_response: object | None, + custom_llm_provider: str | None, + custom_pricing: bool | None, + base_model: str | None, + router_model_id: str | None, + region_name: str | None, + litellm_logging_obj: LitellmLoggingObject | None, +) -> tuple[str, Mapping[str, object]] | None: + deployment_entry: Final = _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id) + deployment_key: Final = router_model_id or model + if deployment_entry is not None and deployment_key is not None: + registered_entry: Final = _raw_cost_map_entry(router_model_id) if router_model_id is not None else None + return deployment_key, registered_entry or deployment_entry + selected_model: Final = _select_model_name_for_cost_calc( + model=model, + completion_response=completion_response, + base_model=base_model, + custom_pricing=custom_pricing, + custom_llm_provider=custom_llm_provider, + router_model_id=router_model_id, + region_name=region_name, + ) + candidates: Final = (selected_model, _get_response_model(completion_response), model) + resolved: Final = next( + (info for info in (_cost_map_model_info(name, custom_llm_provider) for name in candidates if name) if info), + None, + ) + if resolved is None: + return None + return resolved["key"], _raw_cost_map_entry(resolved["key"]) or resolved + + def ocr_cost( model: str, custom_llm_provider: str | None, diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index d3ad7234d93..7ecbeea255c 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -427,7 +427,7 @@ class LLMCallSpanData: # plain ``.get`` — no repeated ``isinstance`` guards. raw_response: Final = payload.get("response") response: Final = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {}) - choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response) or _ocr_choices(response) + choices_out: Final = _output_choices(response) # ``finish_reasons`` is metadata, not content, so derive it from # ``choices_out`` before gating. The raw message/choice bodies are only # retained when content capture is enabled (see ``capture_span_content``); @@ -752,20 +752,99 @@ def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: return (choice,) -def _ocr_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: - markdowns: Final = tuple( - text for page in _dicts(response.get("pages")) if (text := as_str(page.get("markdown"))) is not None +def _output_choices(response: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + """The response output as chat-shaped choices; images and binary bodies become size summaries, never bytes.""" + return ( + _completion_choices(response) + or _responses_choices(response) + or _ocr_choices(response) + or _transcription_choices(response) + or _moderation_choices(response) + or _image_choices(response) + or _binary_choices(response) ) - if not markdowns: + + +def _text_choice(content: str, finish_reason: str | None = None) -> _Choice: + message: Final[_AssistantMessage] = {"role": "assistant", "content": content, "refusal": None, "tool_calls": None} + return {"message": message, "finish_reason": finish_reason} + + +def _joined_choice(parts: tuple[str, ...]) -> tuple[_Choice, ...]: + return (_text_choice("\n\n".join(parts)),) if parts else () + + +def _completion_choices(response: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return tuple( + _text_choice(text, as_str(choice.get("finish_reason"))) + if "message" not in choice and isinstance(text := choice.get("text"), str) + else choice + for choice in _dicts(response.get("choices")) + ) + + +def _ocr_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + return _joined_choice( + tuple(text for page in _dicts(response.get("pages")) if (text := as_str(page.get("markdown"))) is not None) + ) + + +def _transcription_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + text: Final = response.get("text") + return (_text_choice(text),) if isinstance(text, str) and text else () + + +def _moderation_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + return _joined_choice( + tuple( + _moderation_verdict(flagged, result.get("categories")) + for result in _dicts(response.get("results")) + if isinstance(flagged := result.get("flagged"), bool) + ) + ) + + +def _moderation_verdict(flagged: bool, categories: object) -> str: + if not flagged: + return "not flagged" + hits: Final = ( + tuple(name for name, hit in cast(Mapping[str, object], categories).items() if hit is True) + if isinstance(categories, dict) + else () + ) + return f"flagged: {', '.join(hits)}" if hits else "flagged" + + +def _image_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + return _joined_choice( + tuple(summary for item in _dicts(response.get("data")) if (summary := _image_summary(item)) is not None) + ) + + +def _image_summary(item: Mapping[str, object]) -> str | None: + location: Final = _image_location(item) + if location is None: + return None + revised: Final = as_str(item.get("revised_prompt")) + return f"{revised}\n{location}" if revised else location + + +def _image_location(item: Mapping[str, object]) -> str | None: + url: Final = as_str(item.get("url")) + if url is not None: + return url + encoded: Final = item.get("b64_json") + if not isinstance(encoded, str): + return None + return f"b64_json image ({len(encoded) * 3 // 4 - encoded[-2:].count('=')} bytes)" + + +def _binary_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + size: Final = as_int(response.get("num_bytes")) + if size is None: return () - message: Final[_AssistantMessage] = { - "role": "assistant", - "content": "\n\n".join(markdowns), - "refusal": None, - "tool_calls": None, - } - choice: Final[_Choice] = {"message": message, "finish_reason": None} - return (choice,) + content_type: Final = as_str(response.get("content_type")) + return (_text_choice(f"{content_type} ({size} bytes)" if content_type else f"{size} bytes"),) def _responses_parts_text(parts: tuple[Mapping[str, object], ...], part_type: str, field: str) -> str | None: diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 37b7344917e..28ac9f5cdae 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -10,6 +10,7 @@ import sys from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import replace from datetime import datetime, timedelta +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast from pydantic import BaseModel @@ -66,6 +67,7 @@ from litellm.types.proxy.carried_budget_state import ( from litellm.types.utils import ( StandardLoggingGuardrailInformation, StandardLoggingPayload, + StandardLoggingZeroCostDiagnostic, ) if TYPE_CHECKING: @@ -713,6 +715,15 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_requests_metric"), ) + self.litellm_zero_cost_requests_total = self._counter_factory( + name="litellm_zero_cost_requests_total", + documentation=( + "Requests that carried usage but were logged at $0 on a model whose pricing entry " + "has a non-zero rate, by reason (missing_pricing_key, pricing_not_applied, cost_calculation_error)" + ), + labelnames=self.get_labels_for_metric("litellm_zero_cost_requests_total"), + ) + # Cache metrics self.litellm_cache_hits_metric = self._counter_factory( name="litellm_cache_hits_metric", @@ -1410,6 +1421,11 @@ class PrometheusLogger(CustomLogger): enum_values=enum_values, label_context=label_context, ) + self._increment_zero_cost_requests_metric( + zero_cost_diagnostic=standard_logging_payload.get("zero_cost_diagnostic"), + enum_values=enum_values, + label_context=label_context, + ) # input, output, total token metrics self._increment_token_metrics( @@ -1983,6 +1999,30 @@ class PrometheusLogger(CustomLogger): amount=float(response_cost), ) + def _increment_zero_cost_requests_metric( + self, + zero_cost_diagnostic: StandardLoggingZeroCostDiagnostic | None, + enum_values: UserAPIKeyLabelValues, + label_context: PrometheusLabelFactoryContext, + ) -> None: + if zero_cost_diagnostic is None: + return + supported_labels: Final = self.get_labels_for_metric("litellm_zero_cost_requests_total") + reason_label: Final = ( + MappingProxyType({ZERO_COST_REASON_LABEL: zero_cost_diagnostic["reason"]}) + if ZERO_COST_REASON_LABEL in supported_labels + else MappingProxyType({}) + ) + labels: Final = MappingProxyType( + { + **prometheus_label_factory( + supported_enum_labels=supported_labels, enum_values=enum_values, label_context=label_context + ), + **reason_label, + } + ) + self.litellm_zero_cost_requests_total.labels(**labels).inc() + @staticmethod def _get_remaining_from_v3_rate_limit_headers( standard_logging_payload: StandardLoggingPayload | None, @@ -2333,6 +2373,8 @@ class PrometheusLogger(CustomLogger): team_alias=user_api_team_alias, user=user_id, model_id=standard_logging_payload.get("model_id", ""), + requested_model=standard_logging_payload.get("model_group"), + api_provider=standard_logging_payload.get("custom_llm_provider"), custom_metadata_labels=get_custom_labels_from_metadata( metadata=_get_combined_custom_metadata_from_standard_logging_payload( standard_logging_payload=standard_logging_payload @@ -2345,6 +2387,11 @@ class PrometheusLogger(CustomLogger): "litellm_llm_api_failed_requests_metric", enum_values, ) + self._increment_zero_cost_requests_metric( + zero_cost_diagnostic=standard_logging_payload.get("zero_cost_diagnostic"), + enum_values=enum_values, + label_context=PrometheusLabelFactoryContext(enum_values), + ) self.set_llm_deployment_failure_metrics(kwargs) await self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, diff --git a/litellm/litellm_core_utils/README.md b/litellm/litellm_core_utils/README.md index b61c8982762..a5f5e8326b9 100644 --- a/litellm/litellm_core_utils/README.md +++ b/litellm/litellm_core_utils/README.md @@ -6,8 +6,9 @@ Core files: - `streaming_handler.py`: The core streaming logic + streaming related helper utils - `core_helpers.py`: code used in `types/` - e.g. `map_finish_reason`. - `exception_mapping_utils.py`: utils for mapping exceptions to openai-compatible error types. -- `default_encoding.py`: code for loading the default encoding (tiktoken) +- `default_encoding.py`: code for loading the default Python tokenizer and bundled cache - `get_llm_provider_logic.py`: code for inferring the LLM provider from a given model name. - `duration_parser.py`: code for parsing durations - e.g. "1d", "1mo", "10s" - `api_route_to_call_types.py`: mapping of API routes to their corresponding CallTypes (e.g., `/chat/completions` -> [acompletion, completion]) +Tokenizer factories return Python tokenizer objects by default. Set `LITELLM_RUST=1` or call `litellm.rust(True)` before constructing tokenizers to select the Rust backend through `Route.TOKENIZER` in the Rust catalog. Missing native bindings or unsupported native features fall back to Python. Existing tokenizer objects keep their selected backend. Rust-backed tokenizer objects carry the read-only `tiktoken.Encoding` / `tokenizers.Tokenizer` surface and are immutable: `enable_padding`, `enable_truncation` and `add_tokens` stay on the Python tokenizer. diff --git a/litellm/litellm_core_utils/agentic_followup_kwargs.py b/litellm/litellm_core_utils/agentic_followup_kwargs.py new file mode 100644 index 00000000000..50ec19f62c4 --- /dev/null +++ b/litellm/litellm_core_utils/agentic_followup_kwargs.py @@ -0,0 +1,32 @@ +from collections.abc import Collection, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final + + +def build_agentic_followup_kwargs( + *, + request_kwargs: Mapping[str, object], + patch_kwargs: Mapping[str, object], + request_params: Collection[str], + depth: int, + max_loops: int, + fingerprints: Sequence[str], + fingerprint: str, +) -> Mapping[str, object]: + """Kwargs for an agentic follow-up call: the request's kwargs overlaid by the plan's, never repeating a key already sent as a request param""" + seen: Final = [*fingerprints, fingerprint] # mutable-ok: the chat loop's settings reader only accepts a list + return MappingProxyType( + { + key: value + for key, value in chain( + ((k, v) for k, v in request_kwargs.items() if k not in request_params), + ((k, v) for k, v in patch_kwargs.items() if k not in request_params), + ( + ("_agentic_loop_depth", depth + 1), + ("max_agentic_loops", max_loops), + ("_agentic_loop_fingerprints", seen), + ), + ) + } + ) diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index c8e9e2583ba..e0bd85a7937 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -2,10 +2,13 @@ import json from collections.abc import Mapping +from itertools import chain +from types import MappingProxyType from typing import Final, cast from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, @@ -117,13 +120,25 @@ def _wrap_response_as_fake_stream( ) -def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None: - metadata = kwargs_for_followup.get("litellm_metadata") - metadata = dict(metadata) if isinstance(metadata, dict) else {} - for key, value in kwargs_for_followup.items(): - if key.startswith("_agentic_loop") or key == "max_agentic_loops" or is_interception_internal_key(key): - metadata[key] = value - kwargs_for_followup["litellm_metadata"] = metadata +def _with_agentic_loop_metadata(kwargs_for_followup: Mapping[str, object]) -> Mapping[str, object]: + metadata: Final = kwargs_for_followup.get("litellm_metadata") + return MappingProxyType( + { + **kwargs_for_followup, + "litellm_metadata": dict( # mutable-ok: the follow-up call's logging and proxy hooks write into litellm_metadata in place + chain( + metadata.items() if isinstance(metadata, dict) else (), + ( + (key, value) + for key, value in kwargs_for_followup.items() + if key.startswith("_agentic_loop") + or key == "max_agentic_loops" + or is_interception_internal_key(key) + ), + ) + ), + } + ) def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]: @@ -165,14 +180,17 @@ async def _execute_chat_completion_agentic_plan( if "tool_choice" not in patch.optional_params: optional_params_for_followup.pop("tool_choice", None) - kwargs_for_followup: Final = _filter_followup_kwargs(kwargs) - kwargs_for_followup.update( - {k: v for k, v in _filter_followup_kwargs(patch.kwargs).items() if k not in optional_params_for_followup} + kwargs_for_followup: Final = _with_agentic_loop_metadata( + build_agentic_followup_kwargs( + request_kwargs=_filter_followup_kwargs(kwargs), + patch_kwargs=_filter_followup_kwargs(patch.kwargs), + request_params=frozenset((*optional_params_for_followup, "model", "messages")), + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + ) ) - kwargs_for_followup["_agentic_loop_depth"] = depth + 1 - kwargs_for_followup["max_agentic_loops"] = max_loops - kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] - _add_agentic_loop_metadata(kwargs_for_followup) try: response_followup = await litellm.acompletion( diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index ce6f77f78a0..5bcde688521 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -359,6 +359,46 @@ def get_litellm_metadata_from_kwargs(kwargs: dict): return {} +def _budget_reservation_on_auth_object(user_api_key_auth: object) -> object: + if isinstance(user_api_key_auth, Mapping): + return user_api_key_auth.get("budget_reservation") + return getattr(user_api_key_auth, "budget_reservation", None) + + +def budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict | None: + stamped: Final = metadata.get("user_api_key_budget_reservation") + if isinstance(stamped, dict): + return stamped + on_auth_object: Final = _budget_reservation_on_auth_object(metadata.get("user_api_key_auth")) + return on_auth_object if isinstance(on_auth_object, dict) else None + + +def _stamp_budget_reservation_callback_bound(litellm_params: Mapping[str, object], callback_bound: bool) -> None: + for metadata_variable_name in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_variable_name) + if not isinstance(metadata, Mapping): + continue + budget_reservation = budget_reservation_from_metadata(metadata) + if budget_reservation is not None: + budget_reservation["callback_bound"] = callback_bound + + +def bind_budget_reservation_to_callbacks(litellm_params: Mapping[str, object]) -> None: + """Mark the request's budget reservation as owned by the success callbacks of this call. + + The proxy releases any reservation still unbound when the request ends; one bound here + is left for the cost callback, which may finish after the response has been sent. Bind + only where a success handler is guaranteed to run: a logging object merely existing is + not that, since the proxy builds one for every route before calling anything. + """ + _stamp_budget_reservation_callback_bound(litellm_params, True) + + +def unbind_budget_reservation_from_callbacks(litellm_params: Mapping[str, object]) -> None: + """Hand a failed call's reservation back to the request-end release: failure handlers never settle it.""" + _stamp_budget_reservation_callback_bound(litellm_params, False) + + def reconstruct_model_name( model_name: str, custom_llm_provider: str | None, diff --git a/litellm/litellm_core_utils/default_encoding.py b/litellm/litellm_core_utils/default_encoding.py index 71b30614d8d..c3b6a008411 100644 --- a/litellm/litellm_core_utils/default_encoding.py +++ b/litellm/litellm_core_utils/default_encoding.py @@ -1,5 +1,4 @@ import os -from pathlib import Path from typing import Final import litellm @@ -15,20 +14,6 @@ except (ImportError, AttributeError): filename = pkg_resources.resource_filename(__name__, "litellm_core_utils/tokenizers") -CL100K_BASE_RANK_FILE: Final = "9b5ad71b2ce5302211f9c61530b329a4922fc6a4" -O200K_BASE_RANK_FILE: Final = "fb374d419588a4632f3f557e76b4b70aebbca790" - - -def cl100k_base_rank_file() -> str: - """The vendored tiktoken `cl100k_base` rank file (`base64(token) rank` lines).""" - return Path(filename, CL100K_BASE_RANK_FILE).read_text(encoding="ascii") - - -def o200k_base_rank_file() -> str: - """The vendored tiktoken `o200k_base` rank file (`base64(token) rank` lines).""" - return Path(filename, O200K_BASE_RANK_FILE).read_text(encoding="ascii") - - # Always default TIKTOKEN_CACHE_DIR to the bundled tokenizers directory # unless the user explicitly overrides it via CUSTOM_TIKTOKEN_CACHE_DIR. # This keeps tiktoken fully offline-capable by default (see #1071). diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 9b2db9aad18..36fd7fa4e61 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -46,6 +46,8 @@ OPTIONAL_KWARGS_KEYS: Final = ( "bucket_name", "s3_endpoint_url", "s3_region_name", + "s3_access_key_id", + "s3_secret_access_key", "vertex_credentials", "vertex_project", "vertex_location", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 5868e79323a..192679957b1 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -859,6 +859,9 @@ def _get_openai_compatible_provider_info( elif custom_llm_provider == "edenai": api_base = litellm.EdenAIChatConfig.get_api_base(api_base) # rebind-ok: chain resolves in place dynamic_api_key = litellm.EdenAIChatConfig.get_api_key(api_key) # rebind-ok: chain resolves in place + elif custom_llm_provider == "fal_ai": + api_base = litellm.FalAIChatConfig.get_api_base(api_base) # rebind-ok: chain resolves in place + dynamic_api_key = litellm.FalAIChatConfig.get_api_key(api_key) # rebind-ok: chain resolves in place elif custom_llm_provider == "aiml": ( api_base, diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 91a22144805..5471fe50d5f 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -17,14 +17,16 @@ import random import sys import threading import time -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, replace from datetime import datetime, timezone from importlib.resources import files from pathlib import Path +from types import MappingProxyType from typing import Final, Protocol import httpx +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm import verbose_logger @@ -37,6 +39,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( ) FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" +_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) @@ -88,6 +91,18 @@ class GetModelCostMap: """Load the local backup model cost map bundled with the package.""" return GetModelCostMap.load_local_model_cost_map_with_revision().model_cost_map + _loaded_catalog: Mapping[str, Mapping[str, object]] = MappingProxyType({}) + + @classmethod + def loaded_model_cost_map(cls) -> Mapping[str, Mapping[str, object]]: + """The catalog as last loaded (bundled or remote), untouched by ``register_model`` or router registrations.""" + return cls._loaded_catalog + + @classmethod + def _snapshot_loaded_catalog(cls, model_cost: Mapping[str, object]) -> None: + raw: Final = _CATALOG_ADAPTER.validate_python(model_cost) + cls._loaded_catalog = MappingProxyType({key: MappingProxyType(entry) for key, entry in raw.items()}) + @classmethod def _get_backup_model_count(cls) -> int: """Return the number of models in the local backup (cached int).""" @@ -533,7 +548,9 @@ def _finalize_model_cost_map(model_cost: dict) -> dict: def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded: _cost_map_source_info.source_revision = loaded.revision _cost_map_source_info.etag = loaded.etag - return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map)) + finalized: Final = _finalize_model_cost_map(loaded.model_cost_map) + GetModelCostMap._snapshot_loaded_catalog(finalized) # pyright: ignore[reportPrivateUsage] # same module + return replace(loaded, model_cost_map=finalized) def adopt_model_cost_map( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b34f1b3aafd..d06780dda53 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -50,6 +50,8 @@ from litellm.cost_calculator import ( RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, _select_model_name_for_cost_calc, + get_usage_object, + pricing_entry_for_cost_calc, ) from litellm.exceptions import ( BudgetExceededError, @@ -89,6 +91,10 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( InteractionsUsageObjectTransformation, ) +from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import ( + diagnose_zero_cost, + zero_cost_warning, +) from litellm.litellm_core_utils.logging_utils import ( truncate_base64_in_messages, truncate_base64_in_messages_async, @@ -157,6 +163,7 @@ from litellm.types.utils import ( StandardLoggingPayloadStatusFields, StandardLoggingPromptManagementMetadata, StandardLoggingVectorStoreRequest, + StandardLoggingZeroCostDiagnostic, TextCompletionResponse, TranscriptionResponse, Usage, @@ -489,6 +496,14 @@ def mask_api_base_credentials(api_base: str) -> str: return api_base[:key_end] + "*" * 5 + api_base[-4:] +def _timestamp_seconds(moment: object) -> float | None: + if isinstance(moment, datetime.datetime): + return moment.timestamp() + if isinstance(moment, (int, float)): + return float(moment) + return None + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -614,6 +629,7 @@ class Logging(LiteLLMLoggingBaseClass): self.truncated_messages_for_logging: str | list | dict | None = None # mutable-ok: logged messages shape ## TIME TO FIRST TOKEN LOGGING ## self.completion_start_time: datetime.datetime | None = None + self.zero_cost_warned: bool = False self._llm_caching_handler: LLMCachingHandler | None = None # INITIAL LITELLM_PARAMS @@ -1626,10 +1642,12 @@ class Logging(LiteLLMLoggingBaseClass): return response.mcp_tool_call_response def get_response_ms(self) -> float: - return ( - self.model_call_details.get("end_time", datetime.datetime.now()) - - self.model_call_details.get("start_time", datetime.datetime.now()) - ).total_seconds() * 1000 + now: Final = datetime.datetime.now() + start_seconds: Final = _timestamp_seconds(self.model_call_details.get("start_time", now)) + end_seconds: Final = _timestamp_seconds(self.model_call_details.get("end_time", now)) + if start_seconds is None or end_seconds is None: + return 0.0 + return (end_seconds - start_seconds) * 1000 def set_cost_breakdown( self, @@ -1757,17 +1775,26 @@ class Logging(LiteLLMLoggingBaseClass): if transformed_result is not None: result = transformed_result - result_hidden_params: Final = getattr(result, "_hidden_params", None) or MappingProxyType({}) - result_additional_headers: Final = ( - result_hidden_params.get("additional_headers") - if isinstance(result_hidden_params, dict) - else getattr(result_hidden_params, "additional_headers", None) + priced_result: Final = ( + result.response + if isinstance(result, (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent)) + else result ) - if isinstance(result, (BaseModel, HttpxBinaryResponseContent)) and hasattr(result, "_hidden_params"): + + result_hidden_params: Final = getattr(priced_result, "_hidden_params", None) or MappingProxyType({}) + if isinstance(priced_result, (BaseModel, HttpxBinaryResponseContent)) and hasattr( + priced_result, "_hidden_params" + ): hidden_params: Final = result_hidden_params if ( "response_cost" in hidden_params and hidden_params["response_cost"] is not None ): # use cost if already calculated + self._record_zero_cost_diagnostic( + priced_result, + hidden_params["response_cost"], + litellm_model_name=litellm_model_name, + router_model_id=router_model_id or hidden_params.get("model_id"), + ) return hidden_params["response_cost"] elif router_model_id is None and "model_id" in hidden_params: # use model_id if not already set router_model_id = hidden_params["model_id"] @@ -1779,18 +1806,7 @@ class Logging(LiteLLMLoggingBaseClass): router_model_id = self.get_router_model_id() ## RESPONSE COST ## - spilled_over: Final = is_spilled_over_ptu_request( - model_info=_deployment_model_info(self.litellm_params if hasattr(self, "litellm_params") else None), - response_headers=self.model_call_details.get("response_headers"), - additional_headers=result_additional_headers, - ) - custom_pricing: Final = ( - False - if spilled_over - else use_custom_pricing_for_model( - litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None) - ) - ) + custom_pricing: Final = self._custom_pricing_for(priced_result) prompt = self._prompt_for_cost_calculation() @@ -1799,7 +1815,7 @@ class Logging(LiteLLMLoggingBaseClass): try: response_cost_calculator_kwargs: Final = { - "response_object": result, + "response_object": priced_result, "model": litellm_model_name or self.model, "cache_hit": cache_hit, "custom_llm_provider": self.model_call_details.get("custom_llm_provider", None), @@ -1842,9 +1858,18 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug("response_cost: %s", response_cost) additional_response_cost: Final[object] = self.model_call_details.get("additional_response_cost") - if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0: - return (response_cost or 0.0) + additional_response_cost - return response_cost + total_response_cost: Final = ( + (response_cost or 0.0) + additional_response_cost + if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0 + else response_cost + ) + self._record_zero_cost_diagnostic( + priced_result, + total_response_cost, + litellm_model_name=litellm_model_name, + router_model_id=router_model_id, + ) + return total_response_cost except Exception as e: # error calculating cost debug_info = StandardLoggingModelCostFailureDebugInformation( error_str=str(e), @@ -1858,9 +1883,108 @@ class Logging(LiteLLMLoggingBaseClass): ) verbose_logger.debug("response_cost_failure_debug_information: %s", debug_info) self.model_call_details["response_cost_failure_debug_information"] = debug_info + self._record_zero_cost_diagnostic( + priced_result, + None, + calculation_failed=True, + litellm_model_name=litellm_model_name, + router_model_id=router_model_id, + ) return None + def _record_zero_cost_diagnostic( + self, + result: object, + response_cost: float | None, + *, + calculation_failed: bool = False, + litellm_model_name: str | None = None, + router_model_id: str | None = None, + ) -> None: + if response_cost is None and not calculation_failed: + return + if self.model_call_details.get("cache_hit") is True: + self.model_call_details["zero_cost_diagnostic"] = None + return + try: + finding: Final = self._zero_cost_finding( + result, + response_cost, + calculation_failed=calculation_failed, + litellm_model_name=litellm_model_name, + router_model_id=router_model_id, + ) + except Exception as e: # noqa: BLE001 # the pricing helpers raise plain Exception and a diagnostic must never break cost tracking + verbose_logger.debug("zero_cost_diagnostic skipped: %s", e) + return + self.model_call_details["zero_cost_diagnostic"] = finding[0] if finding is not None else None + if finding is None or self.zero_cost_warned: + return + self.zero_cost_warned = True + verbose_logger.warning(finding[1]) + + def _zero_cost_finding( + self, + result: object, + response_cost: float | None, + *, + calculation_failed: bool, + litellm_model_name: str | None, + router_model_id: str | None, + ) -> tuple[StandardLoggingZeroCostDiagnostic, str] | None: + metadata: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(self.litellm_params) + if response_cost or is_unbilled_non_inference_call(self.call_type, metadata, result): + return None + usage: Final = get_usage_object(completion_response=result) + if usage is None: + return None + model: Final = litellm_model_name or self.model + custom_llm_provider: Final = self.model_call_details.get("custom_llm_provider") + pricing: Final = pricing_entry_for_cost_calc( + model=model, + completion_response=result, + custom_llm_provider=custom_llm_provider, + custom_pricing=self._custom_pricing_for(result), + base_model=_get_base_model_from_metadata(model_call_details=self.model_call_details), + router_model_id=router_model_id or self.get_router_model_id(), + region_name=_resolve_mantle_region_for_cost( + custom_llm_provider=custom_llm_provider, + litellm_params=self.model_call_details.get("litellm_params"), + ), + litellm_logging_obj=self, + ) + if pricing is None: + return None + diagnostic: Final = diagnose_zero_cost( + usage=usage, pricing_model=pricing[0], pricing_entry=pricing[1], calculation_failed=calculation_failed + ) + if diagnostic is None: + return None + model_group: Final = metadata.get("model_group") + return diagnostic, zero_cost_warning( + diagnostic, + model_group=model_group if isinstance(model_group, str) else None, + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + ) + + def _custom_pricing_for(self, result: object) -> bool: + litellm_params: Final = getattr(self, "litellm_params", None) + result_hidden_params: Final = getattr(result, "_hidden_params", None) or MappingProxyType({}) + additional_headers: Final = ( + result_hidden_params.get("additional_headers") + if isinstance(result_hidden_params, dict) + else getattr(result_hidden_params, "additional_headers", None) + ) + spilled_over: Final = is_spilled_over_ptu_request( + model_info=_deployment_model_info(litellm_params), + response_headers=self.model_call_details.get("response_headers"), + additional_headers=additional_headers, + ) + return False if spilled_over else use_custom_pricing_for_model(litellm_params=litellm_params) + def _prompt_for_cost_calculation(self) -> str: """ The raw input string is only priced directly for text-to-speech, which bills per character. @@ -2205,6 +2329,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = 0.0 elif "response_cost" in hidden_params: self.model_call_details["response_cost"] = hidden_params["response_cost"] + self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"]) elif (existing_cost := self.model_call_details.get("response_cost")) is not None and existing_cost != 0: # Preserve response_cost if already calculated (e.g., by pass-through # handlers like Gemini/Vertex which call completion_cost directly). @@ -5378,7 +5503,7 @@ def request_model_access_groups_from_litellm_params(litellm_params: Mapping[str, """Access groups the auth layer stamped onto this request, from whichever metadata field carries them. Detached internal sub-calls only inherit the identity keys, so the auth object is the - fallback there, exactly as _get_budget_reservation_from_metadata does for reservations. + fallback there, exactly as budget_reservation_from_metadata does for reservations. """ for metadata_variable_name in ("metadata", "litellm_metadata"): metadata = litellm_params.get(metadata_variable_name) @@ -6162,6 +6287,8 @@ def _extract_response_obj_and_hidden_params( hidden_params = getattr(init_response_obj, "_hidden_params", None) elif isinstance(init_response_obj, dict): response_obj = init_response_obj + elif isinstance(init_response_obj, HttpxBinaryResponseContent): + response_obj = dict(init_response_obj.logging_summary()) else: response_obj = {} @@ -6499,6 +6626,7 @@ def get_standard_logging_object_payload( error_str=error_str, error_information=error_information, response_cost_failure_debug_info=kwargs.get("response_cost_failure_debug_information"), + zero_cost_diagnostic=kwargs.get("zero_cost_diagnostic"), guardrail_information=metadata.get("standard_logging_guardrail_information", None), standard_built_in_tools_params=standard_built_in_tools_params, ) @@ -6677,6 +6805,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: response_cost=response_cost, autorouter_savings=None, response_cost_failure_debug_info=None, + zero_cost_diagnostic=None, status="success", total_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT), prompt_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT), diff --git a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py new file mode 100644 index 00000000000..6331d815bdc --- /dev/null +++ b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py @@ -0,0 +1,146 @@ +from collections.abc import Mapping +from functools import reduce +from typing import Final + +from pydantic import TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.types.utils import StandardLoggingZeroCostDiagnostic, Usage + +ZERO_COST_COUNTER_NAME: Final = "litellm_zero_cost_requests_total" + +_TEXT_INPUT_RATE: Final = "input_cost_per_token" +_AUDIO_INPUT_RATE: Final = "input_cost_per_audio_token" +_TEXT_OUTPUT_RATE: Final = "output_cost_per_token" +_AUDIO_OUTPUT_RATE: Final = "output_cost_per_audio_token" +_RATE_KEY_MARKERS: Final = ("cost", "pricing") +_NESTED_PRICING: Final = TypeAdapter(Mapping[str, object] | tuple[object, ...]) +_MAX_PRICING_DEPTH: Final = 4 + + +def _audio_tokens(details: object) -> int: + audio_tokens: Final = getattr(details, "audio_tokens", None) + return audio_tokens if isinstance(audio_tokens, int) and audio_tokens > 0 else 0 + + +def _tokens(value: object) -> int: + return value if isinstance(value, int) and value > 0 else 0 + + +def used_pricing_keys(usage: Usage) -> tuple[str, ...]: + prompt_audio: Final = _audio_tokens(usage.prompt_tokens_details) + completion_audio: Final = _audio_tokens(usage.completion_tokens_details) + prompt_text: Final = _tokens(usage.prompt_tokens) - prompt_audio + completion_text: Final = _tokens(usage.completion_tokens) - completion_audio + components: Final = ( + (_TEXT_INPUT_RATE, prompt_text), + (_AUDIO_INPUT_RATE, prompt_audio), + (_TEXT_OUTPUT_RATE, completion_text), + (_AUDIO_OUTPUT_RATE, completion_audio), + ) + return tuple(key for key, count in components if count > 0) + + +def _nested_pricing(value: object) -> Mapping[str, object] | tuple[object, ...] | None: + try: + return _NESTED_PRICING.validate_python(value) + except ValidationError: + return None + + +def _is_rate_key(key: str) -> bool: + return any(marker in key for marker in _RATE_KEY_MARKERS) + + +def _rate_values(value: object) -> tuple[object, ...]: + nested: Final = _nested_pricing(value) + if isinstance(nested, Mapping): + return tuple(child for key, child in nested.items() if _is_rate_key(key)) + if nested is None: + return (value,) + return nested + + +def _expand_rate_values(values: tuple[object, ...], _depth: int) -> tuple[object, ...]: + return tuple(nested for value in values for nested in _rate_values(value)) + + +def _is_positive_number(value: object) -> bool: + return not isinstance(value, bool) and isinstance(value, (int, float)) and value > 0 + + +def _declares_a_rate(pricing_entry: Mapping[str, object]) -> bool: + leaves: Final = reduce(_expand_rate_values, range(_MAX_PRICING_DEPTH), (pricing_entry,)) + return any(_is_positive_number(leaf) for leaf in leaves) + + +def _is_explicit_zero(value: object) -> bool: + return not isinstance(value, bool) and isinstance(value, (int, float)) and value == 0 + + +def diagnose_zero_cost( + usage: Usage, + pricing_model: str, + pricing_entry: Mapping[str, object], + calculation_failed: bool, +) -> StandardLoggingZeroCostDiagnostic | None: + used_keys: Final = used_pricing_keys(usage) + if not used_keys: + return None + missing_keys: Final = tuple(key for key in used_keys if pricing_entry.get(key) is None) + if not missing_keys and all(_is_explicit_zero(pricing_entry[key]) for key in used_keys): + return None + if not _declares_a_rate(pricing_entry): + return None + if calculation_failed: + return StandardLoggingZeroCostDiagnostic( + reason="cost_calculation_error", pricing_model=pricing_model, missing_pricing_keys=() + ) + if missing_keys: + return StandardLoggingZeroCostDiagnostic( + reason="missing_pricing_key", pricing_model=pricing_model, missing_pricing_keys=missing_keys + ) + return StandardLoggingZeroCostDiagnostic( + reason="pricing_not_applied", pricing_model=pricing_model, missing_pricing_keys=() + ) + + +def _cause(diagnostic: StandardLoggingZeroCostDiagnostic) -> str: + reason: Final = diagnostic["reason"] + match reason: + case "missing_pricing_key": + return ( + f"pricing entry '{diagnostic['pricing_model']}' has no {', '.join(diagnostic['missing_pricing_keys'])}. " + "Set the missing rate in the deployment's model_info or in the model cost map, " + "or set every rate to 0 to mark the model free" + ) + case "pricing_not_applied": + return ( + f"pricing entry '{diagnostic['pricing_model']}' declares non-zero rates for this usage, " + "but the cost calculator returned $0" + ) + case "cost_calculation_error": + return ( + f"cost calculation raised for pricing entry '{diagnostic['pricing_model']}', " + "see response_cost_failure_debug_information" + ) + case _: + return assert_never(reason) + + +def zero_cost_warning( + diagnostic: StandardLoggingZeroCostDiagnostic, + *, + model_group: str | None, + model: str, + custom_llm_provider: str | None, + usage: Usage, +) -> str: + request: Final = ( + f"model_group={model_group or model} model={model} provider={custom_llm_provider or 'unknown'} " + f"prompt_tokens={_tokens(usage.prompt_tokens)} completion_tokens={_tokens(usage.completion_tokens)}" + ) + return ( + f"Billable request priced at $0 and logged as such ({request}): {_cause(diagnostic)}. " + f'Counted in {ZERO_COST_COUNTER_NAME}{{reason="{diagnostic["reason"]}"}}' + ) diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 4a3c8de78c5..503814cc143 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -251,6 +251,6 @@ def update_response_metadata( return metadata: Final = ResponseMetadata(result) - metadata.set_hidden_params(logging_obj, model, kwargs) metadata.set_timing_metrics(start_time, end_time, logging_obj, include_overhead) + metadata.set_hidden_params(logging_obj, model, kwargs) metadata.apply() diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 4ba7c3966c0..1e96e20a03b 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2261,6 +2261,37 @@ def system_messages_first( ] +def _system_content_as_text_parts(content: object) -> tuple[object, ...]: + if isinstance(content, str): + return (ChatCompletionTextObject(type="text", text=content),) + return tuple(cast(Sequence[object], content)) # cast-ok: non-str system content is a list of content parts + + +def _merge_system_message_run(run: Sequence[AllMessageValues]) -> AllMessageValues: + if len(run) == 1: + return run[0] + contents: Final = tuple(content for content in (message.get("content") for message in run) if content is not None) + if not contents: + return run[0] + if all(isinstance(content, str) for content in contents): + joined_text: Final = "\n\n".join(cast(tuple[str, ...], contents)) # cast-ok: every content is a str + return cast(AllMessageValues, {**run[0], "content": joined_text}) # cast-ok: dict spread keeps message shape + merged_parts: Final = [ # mutable-ok: chat message content must stay a json list + part for content in contents for part in _system_content_as_text_parts(content) + ] + return cast(AllMessageValues, {**run[0], "content": merged_parts}) # cast-ok: dict spread keeps message shape + + +def merge_consecutive_system_messages( + messages: list[AllMessageValues], # mutable-ok: message pipelines type messages as mutable lists +) -> list[AllMessageValues]: # mutable-ok: message pipelines type messages as mutable lists + return [ # mutable-ok: pipelines mutate message lists + merged + for is_system_run, run in groupby(messages, key=lambda message: message.get("role") == "system") + for merged in ((_merge_system_message_run(tuple(run)),) if is_system_run else run) + ] + + def _attempt_json_repair(s: str) -> object | None: """ Attempt to repair truncated JSON produced by LLM tool calls. diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index b409b181a79..15bccf0301e 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -245,6 +245,8 @@ def _redact_model_response_dict_choices(choices, redacted_str: str): if "audio" in choice["delta"]: choice["delta"]["audio"] = None _redact_tool_calls_dict(choice["delta"]) + elif choice.get("text") is not None: + choice["text"] = redacted_str else: _redact_choice_content(choice) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index bf37b1be2e4..5d7956059e4 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -10,7 +10,6 @@ import anyio import anyio.lowlevel import httpx import tiktoken -from tokenizers import Tokenizer from typing_extensions import ParamSpec, TypeVar import litellm @@ -30,8 +29,10 @@ from litellm.constants import ( TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace, HuggingFaceTokenizer, OpenAIEncoding from litellm.litellm_core_utils.url_utils import safe_get from litellm.llms.custom_httpx.http_handler import _get_httpx_client +from litellm.rust_bridge.tokenizer import get_encoding from litellm.types.llms.anthropic import ( AnthropicContentParamSource, AnthropicContentParamSourceFileId, @@ -622,9 +623,11 @@ def _get_exact_count_function( if model is not None or custom_tokenizer is not None: tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model) if tokenizer_json["type"] == "huggingface_tokenizer": - tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"] + tokenizer: Final[HuggingFace] = tokenizer_json["tokenizer"] def count_tokens(text: str) -> int: + if isinstance(tokenizer, HuggingFaceTokenizer): + return tokenizer.count(text) return len(tokenizer.encode_batch_fast([text])[0]) return count_tokens @@ -632,31 +635,43 @@ def _get_exact_count_function( encoding: Final = openai_tokenizer_encoding(model) def encode_length(text: str) -> int: - return len(encoding.encode(text, disallowed_special=())) + return _encoding_count(encoding, text) return _get_tiktoken_count_function(encode_length) else: raise ValueError("Unsupported tokenizer type") else: + default_encoding: Final = _get_default_encoding() def encode_length(text: str) -> int: - return len(_get_default_encoding().encode(text, disallowed_special=())) + return _encoding_count(default_encoding, text) return _get_tiktoken_count_function(encode_length) -def openai_tokenizer_encoding(model: str) -> tiktoken.Encoding: - """The tiktoken encoding `token_counter` uses for a model on the `openai_tokenizer` path.""" +def _encoding_count(encoding: Encoding, text: str) -> int: + if isinstance(encoding, OpenAIEncoding): + return encoding.count(text) + return len(encoding.encode(text, disallowed_special=())) + + +def openai_tokenizer_encoding(model: str) -> Encoding: + """The encoding `token_counter` uses for a model on the `openai_tokenizer` path.""" + return get_encoding(openai_tokenizer_encoding_name(model)) + + +def openai_tokenizer_encoding_name(model: str) -> str: + """The tiktoken encoding name for `model`, without loading the encoding.""" from litellm.utils import print_verbose model_to_use: Final = _fix_model_name(model) if "gpt-4o" in model_to_use: - return tiktoken.get_encoding("o200k_base") + return "o200k_base" try: - return tiktoken.encoding_for_model(model_to_use) + return tiktoken.encoding_name_for_model(model_to_use) except KeyError: print_verbose("Warning: model not found. Using cl100k_base encoding.") - return tiktoken.get_encoding("cl100k_base") + return "cl100k_base" def uses_legacy_message_accounting(model: str) -> bool: diff --git a/litellm/litellm_core_utils/tokenizer.py b/litellm/litellm_core_utils/tokenizer.py new file mode 100644 index 00000000000..aea187fa08e --- /dev/null +++ b/litellm/litellm_core_utils/tokenizer.py @@ -0,0 +1,402 @@ +"""Python faces of the Rust text codecs. + +``OpenAIEncoding`` mirrors ``tiktoken.Encoding`` and ``HuggingFaceTokenizer`` mirrors +``tokenizers.Tokenizer``, so a caller holding ``litellm.encoding`` or the object returned by +``litellm.create_tokenizer`` sees the same read-only surface whichever backend the Rust catalog +selected. Both wrappers are immutable: ``tokenizers`` mutators (``enable_padding``, +``enable_truncation``, ``add_tokens``) stay on the Python tokenizer. +""" + +from __future__ import annotations + +from collections.abc import Callable, Collection, Mapping, Sequence, Set +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, runtime_checkable + +import tiktoken +from tokenizers import AddedToken +from tokenizers import Tokenizer as PythonHuggingFaceTokenizer + +if TYPE_CHECKING: + import numpy as np + import numpy.typing as npt + + from litellm.rust_bridge._native import HuggingFaceEncoding + from litellm.rust_bridge._native import Tokenizer as NativeTokenizer + +SpecialTokens: TypeAlias = Literal["all"] | Collection[str] +AllowedSpecial: TypeAlias = Literal["all"] | Set[str] +HuggingFaceInput: TypeAlias = str | list[str] | tuple[str, ...] +HuggingFaceBatchInput: TypeAlias = HuggingFaceInput | tuple[HuggingFaceInput, HuggingFaceInput] | list[HuggingFaceInput] + + +@dataclass(frozen=True, slots=True) +class OpenAIEncoding: + """``tiktoken.Encoding`` over the Rust tiktoken codec.""" + + _native: NativeTokenizer + _special_tokens: Mapping[str, int] + + @staticmethod + def wrap(native: NativeTokenizer) -> OpenAIEncoding: + return OpenAIEncoding(native, MappingProxyType(native.special_tokens())) + + @staticmethod + def from_tiktoken(encoding: str) -> OpenAIEncoding: + from litellm.rust_bridge._native import Tokenizer as NativeTokenizer + + return OpenAIEncoding.wrap(NativeTokenizer.from_tiktoken(encoding)) + + def __repr__(self) -> str: + return f"" + + @property + def name(self) -> str: + return self._native.name + + @property + def max_token_value(self) -> int: + return self._native.max_token_value() + + @property + def n_vocab(self) -> int: + """For backwards compatibility. Prefer to use `enc.max_token_value + 1`.""" + return self.max_token_value + 1 + + @property + def eot_token(self) -> int: + return self._special_tokens["<|endoftext|>"] + + @property + def special_tokens_set(self) -> set[str]: # mutable-ok: [LIT001, LIT002] SDK return type + return set(self._special_tokens) + + def is_special_token(self, token: int) -> bool: + return self._native.is_special_token(token) + + # ---- encoding ------------------------------------------------------------------------- + + def encode_ordinary(self, text: str) -> list[int]: # mutable-ok: [LIT001, LIT002] SDK return type + return self._native.encode(text) + + def encode( + self, + text: str, + *, + allowed_special: AllowedSpecial = frozenset(), + disallowed_special: SpecialTokens = "all", + ) -> list[int]: # mutable-ok: [LIT001, LIT002] SDK return type + allowed: Final = self._allowed(text, allowed_special, disallowed_special) + if not allowed: + return self.encode_ordinary(text) + return self._native.encode_special(text, tuple(allowed)) + + def encode_to_numpy( + self, + text: str, + *, + allowed_special: AllowedSpecial = frozenset(), + disallowed_special: SpecialTokens = "all", + ) -> npt.NDArray[np.uint32]: + import numpy + + return numpy.asarray( + self.encode(text, allowed_special=allowed_special, disallowed_special=disallowed_special), + dtype=numpy.uint32, + ) + + def encode_ordinary_batch( + self, text: Sequence[str], *, num_threads: int = 8 + ) -> list[list[int]]: # mutable-ok: [LIT001, LIT002] SDK return type + with ThreadPoolExecutor(num_threads) as executor: + return list( # mutable-ok: [LIT002] SDK returns a list + executor.map(self.encode_ordinary, text) + ) + + def encode_batch( + self, + text: Sequence[str], + *, + num_threads: int = 8, + allowed_special: AllowedSpecial = frozenset(), + disallowed_special: SpecialTokens = "all", + ) -> list[list[int]]: # mutable-ok: [LIT001, LIT002] SDK return type + encode: Final = partial(self.encode, allowed_special=allowed_special, disallowed_special=disallowed_special) + with ThreadPoolExecutor(num_threads) as executor: + return list( # mutable-ok: [LIT002] SDK returns a list + executor.map(encode, text) + ) + + def encode_with_unstable( + self, + text: str, + *, + allowed_special: AllowedSpecial = frozenset(), + disallowed_special: SpecialTokens = "all", + ) -> tuple[list[int], list[list[int]]]: # mutable-ok: [LIT001, LIT002] SDK return type + """The stable tokens of `text` and every completion its unstable tail could become. + + Completions come back sorted; tiktoken returns them in hash order.""" + allowed: Final = self._allowed(text, allowed_special, disallowed_special) + return self._native.encode_with_unstable(text, tuple(allowed)) + + def encode_single_token(self, text_or_bytes: str | bytes) -> int: + """The token of one whole piece, special tokens included. Raises `KeyError` otherwise.""" + piece: Final = text_or_bytes.encode("utf-8") if isinstance(text_or_bytes, str) else text_or_bytes + return self._native.encode_single_token(piece) + + def count(self, text: str, fast: bool = False) -> int: + """Count ordinary text; `fast` accelerates supported encodings and otherwise counts normally.""" + return self._native.count(text, fast) + + # ---- decoding ------------------------------------------------------------------------- + + def decode_bytes(self, tokens: Sequence[int]) -> bytes: + return self._native.decode_bytes(tokens) + + def decode(self, tokens: Sequence[int], errors: str = "replace") -> str: + return self.decode_bytes(tokens).decode("utf-8", errors=errors) + + def decode_single_token_bytes(self, token: int) -> bytes: + return self.decode_bytes((token,)) + + def decode_tokens_bytes(self, tokens: Sequence[int]) -> list[bytes]: # mutable-ok: [LIT001, LIT002] SDK return type + return [ # mutable-ok: [LIT002] SDK returns a list + self.decode_single_token_bytes(token) for token in tokens + ] + + def decode_with_offsets( + self, tokens: Sequence[int] + ) -> tuple[str, list[int]]: # mutable-ok: [LIT001, LIT002] SDK return type + """The decoded text and, per token, the index of the first character holding its bytes. + + Like tiktoken, raises `UnicodeDecodeError` when the tokens do not decode to valid UTF-8.""" + token_bytes: Final = self.decode_tokens_bytes(tokens) + text_len = 0 + offsets: Final[list[int]] = [] # mutable-ok: [LIT001] local accumulator + for token in token_bytes: + offsets.append(max(0, text_len - (0x80 <= token[0] < 0xC0))) + text_len += sum(1 for c in token if not 0x80 <= c < 0xC0) + return b"".join(token_bytes).decode("utf-8", errors="strict"), offsets + + def decode_batch( + self, batch: Sequence[Sequence[int]], *, errors: str = "replace", num_threads: int = 8 + ) -> list[str]: # mutable-ok: [LIT001, LIT002] SDK return type + with ThreadPoolExecutor(num_threads) as executor: + return list( # mutable-ok: [LIT002] SDK returns a list + executor.map(partial(self.decode, errors=errors), batch) + ) + + def decode_bytes_batch( + self, batch: Sequence[Sequence[int]], *, num_threads: int = 8 + ) -> list[bytes]: # mutable-ok: [LIT001, LIT002] SDK return type + with ThreadPoolExecutor(num_threads) as executor: + return list( # mutable-ok: [LIT002] SDK returns a list + executor.map(self.decode_bytes, batch) + ) + + def token_byte_values(self) -> list[bytes]: # mutable-ok: [LIT001, LIT002] SDK return type + return self._native.token_byte_values() + + def __reduce__(self) -> tuple[Callable[[str], OpenAIEncoding], tuple[str]]: + return (OpenAIEncoding.from_tiktoken, (self.name,)) + + # ---- private -------------------------------------------------------------------------- + + def _allowed(self, text: str, allowed_special: AllowedSpecial, disallowed_special: SpecialTokens) -> frozenset[str]: + """tiktoken's special-token policy: which specials `text` may encode, after rejecting + any it must not contain.""" + allowed: Final = frozenset(self._special_tokens) if allowed_special == "all" else frozenset(allowed_special) + disallowed: Final = ( + frozenset(self._special_tokens) - allowed if disallowed_special == "all" else frozenset(disallowed_special) + ) + for token in disallowed: + if token in text: + raise ValueError( + f"Encountered text corresponding to disallowed special token {token!r}.\n" + "If you want this text to be encoded as a special token, " + f"pass it to `allowed_special`, e.g. `allowed_special={{{token!r}, ...}}`.\n" + "If you want this text to be encoded as normal text, disable the check for this token " + f"by passing `disallowed_special=(enc.special_tokens_set - {{{token!r}}})`.\n" + "To disable this check for all special tokens, pass `disallowed_special=()`.\n" + ) + return allowed + + +@dataclass(frozen=True, slots=True) +class HuggingFaceTokenizer: + """The read-only ``tokenizers.Tokenizer`` surface over the Rust Hugging Face codec.""" + + _native: NativeTokenizer + + @staticmethod + def from_str(json: str) -> HuggingFaceTokenizer: + from litellm.rust_bridge._native import Tokenizer as NativeTokenizer + + return HuggingFaceTokenizer(NativeTokenizer.from_json(json)) + + from_json = from_str + + @staticmethod + def from_buffer(buffer: bytes) -> HuggingFaceTokenizer: + return HuggingFaceTokenizer.from_str(buffer.decode("utf-8")) + + @staticmethod + def from_file(path: str) -> HuggingFaceTokenizer: + return HuggingFaceTokenizer.from_str(Path(path).read_text(encoding="utf-8")) + + @staticmethod + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> HuggingFaceTokenizer: + from litellm.rust_bridge._native import Tokenizer as NativeTokenizer + + return HuggingFaceTokenizer(NativeTokenizer.from_pretrained(identifier, revision=revision, token=token)) + + def to_str(self, pretty: bool = False) -> str: + return self._native.to_json(pretty) + + def save(self, path: str, pretty: bool = True) -> None: + Path(path).write_text(self.to_str(pretty), encoding="utf-8") + + @property + def name(self) -> str: + return self._native.name + + # ---- vocabulary ----------------------------------------------------------------------- + + def token_to_id(self, token: str) -> int | None: + return self._native.token_to_id(token) + + def id_to_token(self, id: int) -> str | None: + return self._native.id_to_token(id) + + def get_vocab( + self, with_added_tokens: bool = True + ) -> dict[str, int]: # mutable-ok: [LIT001, LIT002] SDK return type + return self._native.get_vocab(with_added_tokens) + + def get_vocab_size(self, with_added_tokens: bool = True) -> int: + return self._native.get_vocab_size(with_added_tokens) + + def get_added_tokens_decoder(self) -> dict[int, AddedToken]: # mutable-ok: [LIT001, LIT002] SDK return type + return { # mutable-ok: [LIT002] SDK returns a dict + token_id: AddedToken( + content, single_word=single_word, lstrip=lstrip, rstrip=rstrip, normalized=normalized, special=special + ) + for token_id, ( + content, + single_word, + lstrip, + rstrip, + normalized, + special, + ) in self._native.added_tokens_decoder() + } + + def num_special_tokens_to_add(self, is_pair: bool) -> int: + return self._native.num_special_tokens_to_add(is_pair) + + @property + def padding(self) -> dict[str, object] | None: # mutable-ok: [LIT001, LIT002] SDK return type + return self._native.padding() + + @property + def truncation(self) -> dict[str, object] | None: # mutable-ok: [LIT001, LIT002] SDK return type + return self._native.truncation() + + @property + def encode_special_tokens(self) -> bool: + return self._native.encode_special_tokens() + + # ---- encoding and decoding ------------------------------------------------------------ + + def encode( + self, + sequence: HuggingFaceInput, + pair: HuggingFaceInput | None = None, + is_pretokenized: bool = False, + add_special_tokens: bool = True, + ) -> HuggingFaceEncoding: + return self._native.encode_huggingface(sequence, pair, is_pretokenized, add_special_tokens) + + def encode_batch( + self, + input: Sequence[HuggingFaceBatchInput], + is_pretokenized: bool = False, + add_special_tokens: bool = True, + ) -> list[HuggingFaceEncoding]: # mutable-ok: [LIT001, LIT002] SDK return type + return self._encode_batch(input, is_pretokenized, add_special_tokens, fast=False) + + def encode_batch_fast( + self, + input: Sequence[HuggingFaceBatchInput], + is_pretokenized: bool = False, + add_special_tokens: bool = True, + ) -> list[HuggingFaceEncoding]: # mutable-ok: [LIT001, LIT002] SDK return type + return self._encode_batch(input, is_pretokenized, add_special_tokens, fast=True) + + def _encode_batch( + self, input: Sequence[HuggingFaceBatchInput], is_pretokenized: bool, add_special_tokens: bool, fast: bool + ) -> list[HuggingFaceEncoding]: # mutable-ok: [LIT001, LIT002] SDK return type + sequences: Final = tuple(_batch_input(item, is_pretokenized) for item in input) + return self._native.encode_batch_huggingface(sequences, is_pretokenized, add_special_tokens, fast) + + def count(self, text: str, fast: bool = False) -> int: + """Count with this tokenizer's configuration; `fast` uses acceleration where supported.""" + return self._native.count(text, fast) + + def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str: + return self._native.decode(ids, skip_special_tokens=skip_special_tokens) + + def decode_batch( + self, sequences: Sequence[Sequence[int]], skip_special_tokens: bool = True + ) -> list[str]: # mutable-ok: [LIT001, LIT002] SDK return type + return [ # mutable-ok: [LIT002] SDK returns a list + self.decode(ids, skip_special_tokens=skip_special_tokens) for ids in sequences + ] + + def __reduce__(self) -> tuple[Callable[[str], HuggingFaceTokenizer], tuple[str]]: + return (HuggingFaceTokenizer.from_str, (self.to_str(),)) + + +def _batch_input( + item: HuggingFaceBatchInput, is_pretokenized: bool +) -> tuple[HuggingFaceInput, HuggingFaceInput | None]: + if isinstance(item, str): + return (item, None) + if is_pretokenized and all(isinstance(word, str) for word in item): + return (tuple(word for word in item if isinstance(word, str)), None) + if len(item) != 2: + raise TypeError("batch input must be a sequence or a pair of sequences") + return (item[0], item[1]) + + +Encoding: TypeAlias = tiktoken.Encoding | OpenAIEncoding +HuggingFace: TypeAlias = PythonHuggingFaceTokenizer | HuggingFaceTokenizer +Tokenizer: TypeAlias = Encoding | HuggingFace + + +class _AddedToken(Protocol): + @property + def special(self) -> bool: ... + + +@runtime_checkable +class _AddedTokenDecoder(Protocol): + def get_added_tokens_decoder(self) -> Mapping[int, _AddedToken]: ... + + +def strip_special_tokens(tokenizer: object, tokens: Sequence[int]) -> Sequence[int]: + """Drop the special added tokens before a Python `tokenizers` decode; the Rust codec's + `decode(skip_special_tokens=True)` already does this itself.""" + if isinstance(tokenizer, HuggingFaceTokenizer) or not isinstance(tokenizer, _AddedTokenDecoder): + return tokens + try: + added: Final = tokenizer.get_added_tokens_decoder() + except Exception: # noqa: BLE001 # optional metadata failures historically fall back to decoding + return tokens + special_ids: Final = frozenset(token_id for token_id, token in added.items() if token.special) + return tuple(token for token in tokens if token not in special_ids) diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index 77f26b65de0..c5a71daaba5 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -23,9 +23,8 @@ from ..common_utils import ( from .streaming_iterator import A2AModelResponseIterator if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer _REGISTRY_PARAMS_KEPT_OUT_OF_OPTIONAL_PARAMS: Final = ( @@ -292,7 +291,7 @@ class A2AConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index 4f4cd074165..b585d35a2ac 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/litellm/llms/aiml/image_generation/transformation.py @@ -14,9 +14,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -171,7 +170,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/aiohttp_openai/chat/transformation.py b/litellm/llms/aiohttp_openai/chat/transformation.py index 530896bf9b0..a06c670e3f1 100644 --- a/litellm/llms/aiohttp_openai/chat/transformation.py +++ b/litellm/llms/aiohttp_openai/chat/transformation.py @@ -16,9 +16,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, ModelResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -68,7 +67,7 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/amazon_nova/chat/transformation.py b/litellm/llms/amazon_nova/chat/transformation.py index 7551fb28c21..a93fbf1e933 100644 --- a/litellm/llms/amazon_nova/chat/transformation.py +++ b/litellm/llms/amazon_nova/chat/transformation.py @@ -17,7 +17,7 @@ from litellm.types.utils import ModelResponse from ...openai_like.chat.transformation import OpenAILikeChatConfig if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class AmazonNovaChatConfig(OpenAILikeChatConfig): @@ -86,7 +86,7 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 4f4d39f09b0..7dee7513538 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -14,9 +14,8 @@ from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest from litellm.types.utils import LiteLLMBatch, LlmProviders, ModelResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LoggingClass = LiteLLMLoggingObj else: @@ -290,7 +289,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 545e920156e..c221e9f1505 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -100,9 +100,8 @@ from ..common_utils import ( ) if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LoggingClass = LiteLLMLoggingObj else: @@ -2688,7 +2687,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py index b15b0159bd9..46ab27ab0c7 100644 --- a/litellm/llms/anthropic/completion/transformation.py +++ b/litellm/llms/anthropic/completion/transformation.py @@ -33,7 +33,7 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class AnthropicTextError(BaseLLMException): @@ -185,7 +185,7 @@ class AnthropicTextConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 54d10837d74..116f96cf00c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -27,6 +27,7 @@ from litellm.llms.anthropic.experimental_pass_through.utils import ( from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) +from litellm.types.llms.openai import OpenAIWebSearchOptions from litellm.types.utils import ModelResponse from litellm.utils import get_model_info @@ -383,6 +384,49 @@ class LiteLLMMessagesToCompletionTransformationHandler: updated_reasoning_effort["summary"] = effective_summary completion_kwargs["reasoning_effort"] = updated_reasoning_effort + @staticmethod + def _plain_effort_for_chat_target( + completion_kwargs: _CompletionKwargs, + *, + thinking: Mapping[str, object] | None, + ) -> str | None: + reasoning_effort: Final = completion_kwargs.get("reasoning_effort") + if not thinking or not isinstance(reasoning_effort, dict) or "summary" not in reasoning_effort: + return None + effort: Final = reasoning_effort.get("effort") + model: Final = completion_kwargs.get("model") + if not isinstance(effort, str) or not isinstance(model, str) or not model: + return None + custom_llm_provider: Final = completion_kwargs.get("custom_llm_provider") + api_base: Final = completion_kwargs.get("api_base") + api_key: Final = completion_kwargs.get("api_key") + try: + local_model, resolved_provider, _, resolved_api_base = litellm.utils.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + api_base=api_base if isinstance(api_base, str) else None, + api_key=api_key if isinstance(api_key, str) else None, + ) + except Exception: + return None + if resolved_provider == "litellm_proxy": + return None + from litellm.main import responses_api_bridge_check + + web_search_options: Final = completion_kwargs.get("web_search_options") + tools: Final = completion_kwargs.get("tools") + model_info, _ = responses_api_bridge_check( + model=local_model, + custom_llm_provider=resolved_provider, + web_search_options=( + cast(OpenAIWebSearchOptions, web_search_options) if isinstance(web_search_options, dict) else None + ), + tools=cast("list[dict[str, object]]", tools) if isinstance(tools, list) else None, + reasoning_effort=reasoning_effort, + api_base=resolved_api_base, + ) + return None if model_info.get("mode") == "responses" else effort + @staticmethod def _normalize_reasoning_effort( completion_kwargs: _CompletionKwargs, @@ -547,6 +591,13 @@ class LiteLLMMessagesToCompletionTransformationHandler: thinking=thinking, ) + plain_effort: Final = LiteLLMMessagesToCompletionTransformationHandler._plain_effort_for_chat_target( + completion_kwargs, + thinking=thinking, + ) + if plain_effort is not None: + completion_kwargs["reasoning_effort"] = plain_effort + return completion_kwargs, tool_name_mapping @staticmethod diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index ebd0342b10c..f23f2602ba8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -208,9 +208,12 @@ async def _check_summary_model_access( (``ProxyException`` from ``_can_object_call_model`` / ``can_*_model``). Unexpected errors during an access check fail closed but are logged separately so operators can distinguish them from a real access-denied - response. DB-lookup failures (object missing from cache or DB) skip the - corresponding scope — matching ``common_checks``, which only enforces a - scope when its backing object can be loaded. + response. User and project lookup failures (object missing from cache or + DB) skip the corresponding scope — matching ``common_checks``, which only + enforces a scope when its backing object can be loaded. A failed team + membership read (a database outage) fails closed instead, since a member + whose limits cannot be read must not have the summary model invoked with + those limits dropped. """ if user_api_key_auth is None: return True @@ -346,13 +349,12 @@ async def _check_summary_model_access( proxy_logging_obj=proxy_logging_obj, ) except Exception as e: - verbose_logger.debug( - "compact_20260112: team membership lookup failed for " - "summary_model=%s access check; skipping member-level scope: %s", + verbose_logger.warning( + "compact_20260112: team membership lookup failed for summary_model=%s access check; denying access: %s", summary_model, e, ) - team_membership = None + return False member_allowed_models: Final = ( team_membership.litellm_budget_table.allowed_models if team_membership is not None and team_membership.litellm_budget_table is not None diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 449319c1b95..2e7e7bb0c9d 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -481,7 +481,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): additional_args={"complete_input_dict": data}, original_response=str(e), ) - raise AzureOpenAIError(status_code=500, message=str(e)) + raise except Exception as e: message: Final = getattr(e, "message", str(e)) body: Final = getattr(e, "body", None) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 424422612db..355714c0daf 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -29,9 +29,8 @@ from ...base_llm.chat.transformation import BaseConfig from ..common_utils import AzureOpenAIError if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LoggingClass = LiteLLMLoggingObj else: @@ -304,7 +303,7 @@ class AzureOpenAIConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/azure_ai/agents/transformation.py b/litellm/llms/azure_ai/agents/transformation.py index 60ce81a23c7..baba3149963 100644 --- a/litellm/llms/azure_ai/agents/transformation.py +++ b/litellm/llms/azure_ai/agents/transformation.py @@ -34,9 +34,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -297,7 +296,7 @@ class AzureAIAgentsConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py index 9e35e396e15..0e4c8ca0d15 100644 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ b/litellm/llms/azure_ai/azure_model_router/transformation.py @@ -15,7 +15,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class AzureModelRouterConfig(AzureAIStudioConfig): @@ -59,7 +59,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 00e1c1e25ba..779a86629e2 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -30,7 +30,7 @@ from litellm.types.utils import ModelResponse, ProviderField from litellm.utils import _add_path_to_api_base, supports_tool_choice if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class AzureFoundryErrorStrings(str, enum.Enum): @@ -305,7 +305,7 @@ class AzureAIStudioConfig(OpenAIConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 67b1a8bcab3..18b4b6f456a 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -12,9 +12,10 @@ from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object if TYPE_CHECKING: - import tiktoken from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): """Azure AI Foundry MAI image generation (e.g. MAI-Image-2.5).""" @@ -245,7 +246,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/base_llm/audio_transcription/transformation.py b/litellm/llms/base_llm/audio_transcription/transformation.py index 2296909cfe1..e4bf148abf3 100644 --- a/litellm/llms/base_llm/audio_transcription/transformation.py +++ b/litellm/llms/base_llm/audio_transcription/transformation.py @@ -12,9 +12,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import FileTypes, ModelResponse, TranscriptionResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -121,7 +120,7 @@ class BaseAudioTranscriptionConfig(BaseConfig, ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/base_llm/bridges/completion_transformation.py b/litellm/llms/base_llm/bridges/completion_transformation.py index 87b55152d09..a5c03705088 100644 --- a/litellm/llms/base_llm/bridges/completion_transformation.py +++ b/litellm/llms/base_llm/bridges/completion_transformation.py @@ -7,10 +7,10 @@ from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Union if TYPE_CHECKING: - import tiktoken from pydantic import BaseModel from litellm import LiteLLMLoggingObj, ModelResponse + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import AllMessageValues @@ -39,7 +39,7 @@ class CompletionTransformationBridge(ABC): messages: list["AllMessageValues"], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> "ModelResponse": diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 7bfc87a30d6..7decf1b4186 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -21,9 +21,8 @@ from litellm.types.llms.openai import ( ) if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.types.utils import ModelResponse from ..base_utils import ( @@ -344,7 +343,7 @@ class BaseConfig(ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> "ModelResponse": diff --git a/litellm/llms/base_llm/completion/transformation.py b/litellm/llms/base_llm/completion/transformation.py index fb472dfa63b..b8ebbfed12f 100644 --- a/litellm/llms/base_llm/completion/transformation.py +++ b/litellm/llms/base_llm/completion/transformation.py @@ -8,9 +8,8 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser from litellm.types.utils import ModelResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -68,7 +67,7 @@ class BaseTextCompletionConfig(BaseConfig, ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/base_llm/embedding/transformation.py b/litellm/llms/base_llm/embedding/transformation.py index da87dcc7f98..46ac3ccf433 100644 --- a/litellm/llms/base_llm/embedding/transformation.py +++ b/litellm/llms/base_llm/embedding/transformation.py @@ -8,9 +8,8 @@ from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse, ModelResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -80,7 +79,7 @@ class BaseEmbeddingConfig(BaseConfig, ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 254995c028f..5ecf033fb2c 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -21,9 +21,8 @@ from litellm.types.utils import LlmProviders, ModelResponse from ..chat.transformation import BaseConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.router import Router as _Router from litellm.types.llms.openai import HttpxBinaryResponseContent @@ -231,7 +230,7 @@ class BaseFilesConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/base_llm/image_generation/transformation.py b/litellm/llms/base_llm/image_generation/transformation.py index 4616441133e..7ac440c6f48 100644 --- a/litellm/llms/base_llm/image_generation/transformation.py +++ b/litellm/llms/base_llm/image_generation/transformation.py @@ -11,9 +11,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -93,7 +92,7 @@ class BaseImageGenerationConfig(ABC): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/base_llm/image_variations/transformation.py b/litellm/llms/base_llm/image_variations/transformation.py index d3e02139e0e..15a4e0f243c 100644 --- a/litellm/llms/base_llm/image_variations/transformation.py +++ b/litellm/llms/base_llm/image_variations/transformation.py @@ -17,9 +17,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -82,7 +81,7 @@ class BaseImageVariationConfig(BaseConfig, ABC): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, ) -> ImageResponse: pass @@ -98,7 +97,7 @@ class BaseImageVariationConfig(BaseConfig, ABC): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, ) -> ImageResponse: pass @@ -125,7 +124,7 @@ class BaseImageVariationConfig(BaseConfig, ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index 8f35b8eac7a..172316027a8 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -5,7 +5,7 @@ import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.rust_bridge import runtime -from litellm.rust_bridge.catalog import Context, Route +from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.timeouts import timeout_to_seconds from litellm.rust_bridge.transcription.native import ( NATIVE_ATRANSCRIPTION, @@ -74,7 +74,7 @@ class BedrockAudioTranscriptionRustDispatch: ) return runtime.run( - Context(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), + RouteContext(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), binding=NATIVE_TRANSCRIPTION, native=native, python=_no_python_implementation, @@ -107,7 +107,7 @@ class BedrockAudioTranscriptionRustDispatch: ) return await runtime.arun( - Context(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), + RouteContext(Route.TRANSCRIPTION, provider=custom_llm_provider, model=model), binding=NATIVE_ATRANSCRIPTION, native=native, python=_no_async_python_implementation, diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index dd62cdb424a..badb76d00c7 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -524,6 +524,17 @@ class BaseAWSLLM(SignsRequestsWithAWS): aws_session_tags=_canonical_aws_session_tags(auth_params.aws_session_tags), ) + def resolve_s3_credentials(self, params: Mapping[str, object], aws_region_name: str | None) -> Credentials: + """S3 signing identity: the s3_* static pair as-is when both are set, otherwise the resolved aws_* params.""" + from botocore.credentials import Credentials + + from litellm.llms.bedrock.common_utils import s3_static_key_pair + + s3_pair: Final = s3_static_key_pair(params) + if s3_pair is None: + return self.resolve_credentials(AwsAuthParams.model_validate(params), aws_region_name) + return Credentials(access_key=s3_pair[0], secret_key=s3_pair[1]) + def _get_aws_region_from_model_arn(self, model: str | None) -> str | None: try: # First check if the string contains the expected prefix diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index e1a9a807abc..29133bcfaf9 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -40,9 +40,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -990,7 +989,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 4f9b1f56a5b..21830eb0d8e 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -99,7 +99,7 @@ from ..common_utils import ( ) if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer # Computer use tool prefixes supported by Bedrock BEDROCK_COMPUTER_USE_TOOLS: Final = [ @@ -1920,7 +1920,7 @@ class AmazonConverseConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index d489e47c3b5..6877af74494 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -37,9 +37,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -438,7 +437,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py index 5a3f4f17b8b..5699f94d084 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py @@ -25,7 +25,7 @@ from litellm.types.utils import ( from .amazon_llama_transformation import AmazonLlamaConfig if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class AmazonDeepSeekR1Config(AmazonLlamaConfig): @@ -39,7 +39,7 @@ class AmazonDeepSeekR1Config(AmazonLlamaConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index 5d39b68d9d5..f8b730b6cd0 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -21,9 +21,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.types.utils import ModelResponse LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -198,7 +197,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> "ModelResponse": diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py index bc97551d57a..8d1ff1d2bd9 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py @@ -28,7 +28,7 @@ from ..converse_transformation import AmazonConverseConfig from .base_invoke_transformation import AmazonInvokeConfig if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer _CachePointCarrier = TypeVar("_CachePointCarrier", SystemContentBlock, ContentBlock) _INJECTION_POINTS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) @@ -128,7 +128,7 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py index c78375c37bb..67364ccfda0 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -21,7 +21,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, Usage if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class AmazonQwen2Config(AmazonQwen3Config): @@ -44,7 +44,7 @@ class AmazonQwen2Config(AmazonQwen3Config): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py index e251fb15725..2e19dbf77af 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py @@ -19,7 +19,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, Usage if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): @@ -170,7 +170,7 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index 39cded4ed64..fe287111fdd 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -25,9 +25,8 @@ from litellm.types.utils import ModelResponse, Usage from litellm.utils import get_base64_str if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -190,7 +189,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 1326dc22ca0..a8b94fb5703 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -34,9 +34,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -359,7 +358,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 90a2692f68a..dcc5e249d8a 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -34,9 +34,8 @@ from litellm.types.utils import ModelResponse, Usage from litellm.utils import CustomStreamWrapper if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -288,7 +287,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index f1066643874..d9fc813a594 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -112,6 +112,17 @@ def merge_bedrock_aws_request_params( return request_params +def s3_static_key_pair(params: Mapping[str, object]) -> tuple[str, str] | None: + """The s3_access_key_id / s3_secret_access_key pair when both are set, otherwise None.""" + s3_access_key_id: Final = params.get("s3_access_key_id") + s3_secret_access_key: Final = params.get("s3_secret_access_key") + if not isinstance(s3_access_key_id, str) or not s3_access_key_id: + return None + if not isinstance(s3_secret_access_key, str) or not s3_secret_access_key: + return None + return s3_access_key_id, s3_secret_access_key + + # Lazy import cache to avoid circular imports and performance impact _get_model_info = None @@ -776,7 +787,7 @@ def is_bedrock_application_inference_profile_arn(model: str) -> bool: def strip_bedrock_routing_prefix(model: str) -> str: """Strip LiteLLM routing prefixes from model name.""" - for prefix in ["bedrock/", "converse/", "invoke/", "openai/", "nova-2/", "nova/"]: + for prefix in ["bedrock/", "converse/", "invoke/", "openai/", "mantle/", "nova-2/", "nova/"]: if model.startswith(prefix): model = model.split("/", 1)[1] return model @@ -839,6 +850,7 @@ def get_bedrock_base_model(model: str) -> str: Handle model names like: - "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1" - "bedrock/converse/model" -> "model" + - "bedrock/mantle/anthropic.claude-sonnet-5" -> "anthropic.claude-sonnet-5" - "anthropic.claude-3-5-sonnet-20241022-v2:0:51k" -> "anthropic.claude-3-5-sonnet-20241022-v2:0" - "bedrock/nova-2/arn:aws:..." -> "amazon.nova-2-custom" - "bedrock/nova/arn:aws:..." -> "amazon.nova-custom" diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 0b75474ba1b..3d23b69f846 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -11,7 +11,6 @@ from litellm.litellm_core_utils.cloud_storage_security import ( validate_managed_cloud_file_id, ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.types.llms.bedrock import AwsAuthParams from litellm.types.llms.openai import ( FileContentRequest, HttpxBinaryResponseContent, @@ -103,9 +102,7 @@ class BedrockFilesHandler(BaseAWSLLM): ) aws_region_name: Final = self._get_aws_region_name(optional_params=optional_params, model="") - credentials: Final[Credentials] = self.resolve_credentials( - AwsAuthParams.model_validate(optional_params), aws_region_name - ) + credentials: Final[Credentials] = self.resolve_s3_credentials(optional_params, aws_region_name) # Create S3 client s3_client: Final = boto3.client( diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index a7486dd4de0..43faa7d79ea 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -63,7 +63,11 @@ from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id +from ..common_utils import ( + BedrockError, + merge_bedrock_aws_request_params, + resolve_s3_encryption_key_id, +) S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" @@ -148,6 +152,8 @@ class _BedrockS3RequestParams(AwsAuthParams): aws_region_name: str | None = None s3_region_name: str | None = None s3_endpoint_url: str | None = None + s3_access_key_id: str | None = None + s3_secret_access_key: str | None = None @dataclass(frozen=True, slots=True) @@ -1147,7 +1153,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") aws_region_name: Final = self._get_aws_region_name(optional_params=optional_params, model="") - credentials: Final = self.resolve_credentials(AwsAuthParams.model_validate(optional_params), aws_region_name) + credentials: Final = self.resolve_s3_credentials(optional_params, aws_region_name) # Calculate SHA256 hash of the content (REQUIRED for S3) content_hash: Final = hashlib.sha256(content.encode("utf-8")).hexdigest() @@ -1494,7 +1500,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - credentials: Final = self.resolve_credentials(request_params, aws_region_name) + credentials: Final = self.resolve_s3_credentials(request_params.model_dump(exclude_none=True), aws_region_name) empty_body_hash: Final = hashlib.sha256(b"").hexdigest() aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 1f37fafde01..f46edc766c7 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -81,6 +81,10 @@ class AmazonAnthropicClaudeMessagesConfig( def custom_llm_provider(self) -> str | None: return "bedrock" + @property + def beta_headers_provider(self) -> str: + return self.custom_llm_provider or "bedrock" + BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys()) def get_error_class( @@ -552,7 +556,7 @@ class AmazonAnthropicClaudeMessagesConfig( if "tool-search-tool-2025-10-19" in beta_set: beta_set.add("tool-examples-2025-10-29") - beta_provider: Final = self.custom_llm_provider or "bedrock" + beta_provider: Final = self.beta_headers_provider filtered_betas: Final = sorted( filter_and_transform_beta_headers( beta_headers=list(beta_set), diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index 7c8758960ad..052eb90a833 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -2,16 +2,20 @@ Transformation for Bedrock Mantle (Claude Mythos Preview) - /messages endpoint Inherits all Messages API request/response transformations from -AmazonAnthropicClaudeMessagesConfig. Overrides only the URL and model-prefix -stripping that are specific to the bedrock-mantle endpoint. +AmazonAnthropicClaudeMessagesConfig. Overrides the URL, the model-prefix +stripping, and the anthropic-version / anthropic-beta placement (headers, +never the body) that are specific to the bedrock-mantle endpoint. """ -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + DEFAULT_ANTHROPIC_API_VERSION, AnthropicMessagesConfig, ) from litellm.llms.bedrock.common_utils import build_mantle_messages_url @@ -31,6 +35,18 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_BODY_FIELDS_MANTLE_READS_FROM_HEADERS: Final = frozenset({"anthropic_version", "anthropic_beta"}) +_ANTHROPIC_BETAS: Final = TypeAdapter(tuple[str, ...]) +_MANTLE_REQUEST: Final = TypeAdapter(dict[str, object]) + + +def _move_betas_into_header(request: Mapping[str, object], headers: dict[str, str]) -> None: + betas: Final = _ANTHROPIC_BETAS.validate_python(request.get("anthropic_beta") or ()) + if betas: + headers["anthropic-beta"] = ",".join(betas) # rebind-ok: the handler signs and sends this same dict + return + headers.pop("anthropic-beta", None) # rebind-ok: a caller header Mantle rejects in full must not reach it + class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): """ @@ -40,6 +56,13 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): model ID in the request body (unlike Bedrock Invoke which puts it in the URL). """ + @property + def beta_headers_provider(self) -> str: + return "bedrock_mantle" + + def should_filter_anthropic_beta_headers(self) -> bool: + return False + def get_complete_url( self, api_base: str | None, @@ -66,7 +89,7 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): api_key: str | None = None, api_base: str | None = None, ) -> tuple[dict, str | None]: - headers, api_base = super().validate_anthropic_messages_environment( + merged_headers, resolved_api_base = super().validate_anthropic_messages_environment( headers=headers, model=model, messages=messages, @@ -76,9 +99,21 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): api_base=api_base, ) project_id: Final = litellm_params.get("aws_bedrock_project_id") - if project_id: - headers["anthropic-workspace"] = project_id - return headers, api_base + has_version: Final = any(name.lower() == "anthropic-version" for name in merged_headers) + mantle_headers: Final = MappingProxyType( + { + name: value + for name, value in ( + ("anthropic-workspace", project_id), + ("anthropic-version", None if has_version else DEFAULT_ANTHROPIC_API_VERSION), + ) + if value + } + ) + return { # mutable-ok: the base class contract returns a dict the handler signs into in place + **merged_headers, + **mantle_headers, + }, resolved_api_base def transform_anthropic_messages_request( self, @@ -88,25 +123,28 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> dict: - # Strip "mantle/" routing prefix to get the real model ID model_id: Final = model.replace("mantle/", "", 1) - - request: Final = super().transform_anthropic_messages_request( - model=model_id, - messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - litellm_params=litellm_params, - headers=headers, + request: Final = _MANTLE_REQUEST.validate_python( + super().transform_anthropic_messages_request( + model=model_id, + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ), ) - - # Parent (AmazonAnthropicClaudeMessagesConfig) removes "model" and - # "stream" from the body (Bedrock Invoke puts the model in the URL and - # streams via a dedicated endpoint). The mantle endpoint (Messages API) - # requires both in the request body. - stream_fields: Final[dict[str, bool]] = ( - {"stream": True} if anthropic_messages_optional_request_params.get("stream") is True else {} + _move_betas_into_header(request, headers) + body: Final = MappingProxyType( + {key: value for key, value in request.items() if key not in _BODY_FIELDS_MANTLE_READS_FROM_HEADERS} ) - return {**request, "model": model_id, **stream_fields} + streaming: Final = anthropic_messages_optional_request_params.get("stream") is True + mantle_fields: Final = MappingProxyType( + {key: value for key, value in (("model", model_id), ("stream", streaming)) if value} + ) + return { # mutable-ok: the base class contract returns the dict the handler serializes as the body + **body, + **mantle_fields, + } def transform_anthropic_messages_response( self, diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py index 6e975d072ed..78881153e91 100644 --- a/litellm/llms/bedrock_mantle/messages/transformation.py +++ b/litellm/llms/bedrock_mantle/messages/transformation.py @@ -2,11 +2,6 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final -from pydantic import TypeAdapter - -from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( - DEFAULT_ANTHROPIC_API_VERSION, -) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import MANTLE_MESSAGES_PATH from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig @@ -17,7 +12,6 @@ from litellm.llms.bedrock_mantle.common_utils import ( ) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES -from litellm.types.router import GenericLiteLLMParams _BASE_SUFFIXES_TO_STRIP: Final = ( MANTLE_MESSAGES_PATH, @@ -27,9 +21,6 @@ _BASE_SUFFIXES_TO_STRIP: Final = ( "/openai/v1", "/v1", ) -_BODY_FIELDS_MANTLE_READS_FROM_HEADERS: Final = frozenset({"anthropic_version", "anthropic_beta"}) -_ANTHROPIC_BETAS: Final = TypeAdapter(tuple[str, ...]) -_MANTLE_REQUEST: Final = TypeAdapter(dict[str, object]) def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mapping[str, object]) -> str: @@ -74,54 +65,3 @@ class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleM stream: bool | None = None, ) -> str: return build_mantle_native_messages_url(api_base=api_base, litellm_params=litellm_params) - - def validate_anthropic_messages_environment( - self, - headers: dict, - model: str, - messages: list[dict], - optional_params: dict, - litellm_params: dict, - api_key: str | None = None, - api_base: str | None = None, - ) -> tuple[dict, str | None]: - merged_headers, resolved_api_base = super().validate_anthropic_messages_environment( - headers=headers, - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - api_key=api_key, - api_base=api_base, - ) - if any(name.lower() == "anthropic-version" for name in merged_headers): - return merged_headers, resolved_api_base - return { # mutable-ok: the base class contract returns a dict the handler signs into in place - **merged_headers, - "anthropic-version": DEFAULT_ANTHROPIC_API_VERSION, - }, resolved_api_base - - def transform_anthropic_messages_request( - self, - model: str, - messages: list[dict], - anthropic_messages_optional_request_params: dict, - litellm_params: GenericLiteLLMParams, - headers: dict, - ) -> dict: - request: Final = _MANTLE_REQUEST.validate_python( - super().transform_anthropic_messages_request( - model=model, - messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - litellm_params=litellm_params, - headers=headers, - ), - ) - betas: Final = request.get("anthropic_beta") - if betas is not None: - header_betas: Final = ",".join(_ANTHROPIC_BETAS.validate_python(betas)) - headers["anthropic-beta"] = header_betas # rebind-ok: the handler signs and sends this same dict - return { # mutable-ok: the base class contract returns the dict the handler serializes as the body - key: value for key, value in request.items() if key not in _BODY_FIELDS_MANTLE_READS_FROM_HEADERS - } diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py index 119ffff1c34..e5c2bdf59e9 100644 --- a/litellm/llms/black_forest_labs/image_generation/transformation.py +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -29,9 +29,8 @@ from ..common_utils import ( ) if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -258,7 +257,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/brave/search/__init__.py b/litellm/llms/brave/search/__init__.py index cc1168d7ef8..de70c62e040 100644 --- a/litellm/llms/brave/search/__init__.py +++ b/litellm/llms/brave/search/__init__.py @@ -1,7 +1,7 @@ -""" -Brave Search API module. -""" - -from litellm.llms.brave.search.transformation import BraveSearchConfig - -__all__ = ["BraveSearchConfig"] +""" +Brave Search API module. +""" + +from litellm.llms.brave.search.transformation import BraveSearchConfig + +__all__ = ["BraveSearchConfig"] diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index 7977db0f056..e622761dd7f 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -23,9 +23,8 @@ from litellm.utils import CustomStreamWrapper, ModelResponse, Usage from ..common_utils import API_BASE, BytezError if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -187,7 +186,7 @@ class BytezChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/clarifai/chat/transformation.py b/litellm/llms/clarifai/chat/transformation.py index 76d35467497..a0946254de0 100644 --- a/litellm/llms/clarifai/chat/transformation.py +++ b/litellm/llms/clarifai/chat/transformation.py @@ -13,9 +13,8 @@ from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -87,7 +86,7 @@ class ClarifaiConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/cohere/chat/transformation.py b/litellm/llms/cohere/chat/transformation.py index fa46bd7f6cf..a26cdc81695 100644 --- a/litellm/llms/cohere/chat/transformation.py +++ b/litellm/llms/cohere/chat/transformation.py @@ -15,9 +15,8 @@ from ..common_utils import ModelResponseIterator as CohereModelResponseIterator from ..common_utils import validate_environment as cohere_validate_environment if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -227,7 +226,7 @@ class CohereChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index 4252e7d02e9..37c53640d18 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -20,9 +20,8 @@ from ..common_utils import CohereError, CohereV2ModelResponseIterator from ..common_utils import validate_environment as cohere_validate_environment if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -191,7 +190,7 @@ class CohereV2ChatConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/cohere/embed/handler.py b/litellm/llms/cohere/embed/handler.py index 3cebf6b9a90..6d9543d7c30 100644 --- a/litellm/llms/cohere/embed/handler.py +++ b/litellm/llms/cohere/embed/handler.py @@ -20,7 +20,7 @@ from litellm.types.utils import EmbeddingResponse from .v1_transformation import CohereEmbeddingConfig if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer def validate_environment(api_key, headers: dict): @@ -60,7 +60,7 @@ async def async_embedding( api_base: str, api_key: str | None, headers: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", client: AsyncHTTPHandler | None = None, ): ## LOGGING @@ -122,7 +122,7 @@ def embedding( logging_obj: LiteLLMLoggingObj, optional_params: dict, headers: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", data: dict | CohereEmbeddingRequest | None = None, complete_api_base: str | None = None, api_key: str | None = None, diff --git a/litellm/llms/cometapi/image_generation/transformation.py b/litellm/llms/cometapi/image_generation/transformation.py index 03c820de198..4432c151a64 100644 --- a/litellm/llms/cometapi/image_generation/transformation.py +++ b/litellm/llms/cometapi/image_generation/transformation.py @@ -13,9 +13,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -132,7 +131,7 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/compactifai/chat/transformation.py b/litellm/llms/compactifai/chat/transformation.py index 3c5a889ce63..a02db2338d8 100644 --- a/litellm/llms/compactifai/chat/transformation.py +++ b/litellm/llms/compactifai/chat/transformation.py @@ -17,9 +17,8 @@ from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -66,7 +65,7 @@ class CompactifAIChatConfig(OpenAIGPTConfig): messages: Sequence[AllMessageValues], optional_params: Mapping[str, object], litellm_params: Mapping[str, object], - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 0809ef5274f..034c9514092 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -26,9 +26,8 @@ from litellm.types.utils import HttpHandlerRequestFields, ImageResponse, LlmProv from litellm.utils import CustomStreamWrapper, ModelResponse, ProviderConfigManager if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -268,7 +267,7 @@ class BaseLLMAIOHTTPHandler: messages: list, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, client: ClientSession | None = None, ): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index db821f42a90..333ce523e34 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4,7 +4,6 @@ import ssl from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from functools import lru_cache -from itertools import chain from types import MappingProxyType, ModuleType from typing import ( TYPE_CHECKING, @@ -34,6 +33,7 @@ from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.files.types import FileContentStreamingResult +from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, @@ -182,22 +182,22 @@ from litellm.utils import ( def _rust_responses_websocket_enabled( custom_llm_provider: str | None, ) -> bool: - from litellm.rust_bridge.catalog import Context, Delivery, Route, decision + from litellm.rust_bridge.catalog import Delivery, Route, RouteContext, decision from litellm.rust_bridge.configuration import Decision - context: Final = Context(Route.RESPONSES, provider=custom_llm_provider, delivery=Delivery.WEBSOCKET) + context: Final = RouteContext(Route.RESPONSES, provider=custom_llm_provider, delivery=Delivery.WEBSOCKET) return decision(context) is not Decision.PYTHON from .http_handler import get_shared_realtime_ssl_context if TYPE_CHECKING: - import tiktoken from aiohttp import ClientSession from websockets.asyncio.client import ClientConnection from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( FakeAnthropicMessagesStreamIterator, ) @@ -493,7 +493,7 @@ class BaseLLMHTTPHandler: messages: list, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, client: AsyncHTTPHandler | None = None, json_mode: bool = False, @@ -559,7 +559,7 @@ class BaseLLMHTTPHandler: api_base: str | None, custom_llm_provider: str, model_response: ModelResponse, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", logging_obj: LiteLLMLoggingObj, optional_params: dict, timeout: float | httpx.Timeout, @@ -5614,28 +5614,22 @@ class BaseLLMHTTPHandler: } internal_keys: Final = {"litellm_logging_obj"} - kwargs_for_followup: Final = MappingProxyType( - { - key: value - for key, value in chain( - ( - (k, v) - for k, v in kwargs.items() - if not is_interception_internal_key( - k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES - ) - and k != "_code_interpreter_interception_converted_stream" - and k not in internal_keys - and k not in optional_params - ), - ((k, v) for k, v in patch.kwargs.items() if k not in optional_params), - ( - ("_agentic_loop_depth", depth + 1), - ("max_agentic_loops", max_loops), - ("_agentic_loop_fingerprints", fingerprints + [fingerprint]), - ), - ) - } + kwargs_for_followup: Final = build_agentic_followup_kwargs( + request_kwargs=MappingProxyType( + { + k: v + for k, v in kwargs.items() + if not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES) + and k != "_code_interpreter_interception_converted_stream" + and k not in internal_keys + } + ), + patch_kwargs=patch.kwargs, + request_params=frozenset((*optional_params, "model", "input")), + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, ) try: @@ -5756,17 +5750,23 @@ class BaseLLMHTTPHandler: "stream_response", "custom_prompt_dict", } - kwargs_for_followup: Final = { - k: v - for k, v in kwargs.items() - if not k.startswith("_websearch_interception") - and not k.startswith("_compression_interception") - and k not in internal_params - } - kwargs_for_followup.update(patch.kwargs) - kwargs_for_followup["_agentic_loop_depth"] = depth + 1 - kwargs_for_followup["max_agentic_loops"] = max_loops - kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] + kwargs_for_followup: Final = build_agentic_followup_kwargs( + request_kwargs=MappingProxyType( + { + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") + and not k.startswith("_compression_interception") + and k not in internal_params + } + ), + patch_kwargs=patch.kwargs, + request_params=frozenset((*optional_params_for_followup, "model", "messages")), + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + ) return await litellm.acompletion( model=full_model_name, diff --git a/litellm/llms/dashscope/image_generation/transformation.py b/litellm/llms/dashscope/image_generation/transformation.py index c0e278a96ef..ffa60a3d9bc 100644 --- a/litellm/llms/dashscope/image_generation/transformation.py +++ b/litellm/llms/dashscope/image_generation/transformation.py @@ -38,9 +38,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -165,7 +164,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index dd257cd68b0..538904b34e6 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -16,6 +16,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( _extract_reasoning_content, # pyright: ignore[reportPrivateUsage] # same import as the OpenAI transformation + merge_consecutive_system_messages, strip_litellm_internal_message_fields, strip_name_from_message, ) @@ -148,9 +149,8 @@ def _split_parallel_tool_calls(messages: list[AllMessageValues]) -> list[AllMess if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -188,7 +188,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): return "databricks" @classmethod - def get_config(cls): + def get_config(cls, *, model: str | None = None): return super().get_config() def get_required_params(self) -> list[ProviderField]: @@ -465,7 +465,9 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): new_messages.append(_message) if "claude" not in model: - new_messages = _split_parallel_tool_calls(cast(list[AllMessageValues], new_messages)) + new_messages = _split_parallel_tool_calls( + merge_consecutive_system_messages(cast(list[AllMessageValues], new_messages)) + ) if is_async: return super()._transform_messages(messages=new_messages, model=model, is_async=cast(Literal[True], True)) @@ -648,7 +650,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/deprecated_providers/aleph_alpha.py b/litellm/llms/deprecated_providers/aleph_alpha.py index 2ad9ce4edc8..8f5c35f32f2 100644 --- a/litellm/llms/deprecated_providers/aleph_alpha.py +++ b/litellm/llms/deprecated_providers/aleph_alpha.py @@ -277,12 +277,7 @@ def completion( ## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here. prompt_tokens: Final = len(encoding.encode(prompt)) - completion_tokens: Final = len( - encoding.encode( - model_response["choices"][0]["message"]["content"], - disallowed_special=(), - ) - ) + completion_tokens: Final = len(encoding.encode(model_response["choices"][0]["message"]["content"])) model_response.created = int(time.time()) model_response.model = model diff --git a/litellm/llms/edenai/chat/transformation.py b/litellm/llms/edenai/chat/transformation.py index 4fd5a9d550b..67d308d9e38 100644 --- a/litellm/llms/edenai/chat/transformation.py +++ b/litellm/llms/edenai/chat/transformation.py @@ -25,9 +25,8 @@ from litellm.types.utils import ModelResponse, ModelResponseStream, Usage from ..common_utils import EdenAIException, reported_cost, resolve_api_base, resolve_api_key if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding _OPTIONAL_MAPPING: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) @@ -97,7 +96,7 @@ class EdenAIChatConfig(OpenAIGPTConfig): messages: list[AllMessageValues], # mutable-ok: inherited contract optional_params: dict[str, object], # mutable-ok: inherited contract litellm_params: dict[str, object], # mutable-ok: inherited contract - encoding: "tiktoken.Encoding | None", + encoding: "Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/edenai/image_generation/transformation.py b/litellm/llms/edenai/image_generation/transformation.py index 2965c4041de..7f729cd7fbb 100644 --- a/litellm/llms/edenai/image_generation/transformation.py +++ b/litellm/llms/edenai/image_generation/transformation.py @@ -19,9 +19,8 @@ from litellm.utils import convert_to_model_response_object from ..common_utils import EdenAIException, endpoint_url, json_headers, pick, reported_cost if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding _SUPPORTED_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = ( "background", @@ -94,7 +93,7 @@ class EdenAIImageGenerationConfig(BaseImageGenerationConfig): request_data: dict[str, object], # mutable-ok: inherited contract optional_params: dict[str, object], # mutable-ok: inherited contract litellm_params: dict[str, object], # mutable-ok: inherited contract - encoding: "tiktoken.Encoding | None", + encoding: "Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/chat/__init__.py b/litellm/llms/fal_ai/chat/__init__.py new file mode 100644 index 00000000000..b2a4a006aef --- /dev/null +++ b/litellm/llms/fal_ai/chat/__init__.py @@ -0,0 +1,3 @@ +from .transformation import FalAIChatConfig, FalAIError + +__all__ = ("FalAIChatConfig", "FalAIError") diff --git a/litellm/llms/fal_ai/chat/transformation.py b/litellm/llms/fal_ai/chat/transformation.py new file mode 100644 index 00000000000..2c5af6538ab --- /dev/null +++ b/litellm/llms/fal_ai/chat/transformation.py @@ -0,0 +1,244 @@ +""" +Support for `/v1/chat/completions` on Fal AI model endpoints, e.g. fal-ai/moondream3-preview/query. + +These endpoints are not OpenAI-compatible: the request body is a flat ``{"prompt", "image_url"}`` +object and the response is ``{"output", "reasoning", "finish_reason", "usage_info"}``. +""" + +import time +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import Message, ModelResponse, Usage + +if TYPE_CHECKING: + import tiktoken + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +DEFAULT_BASE_URL: Final[str] = "https://fal.run" +PROVIDER_PREFIX: Final[str] = "fal_ai/" +PASSTHROUGH_PARAMS: Final[frozenset[str]] = frozenset(("reasoning", "temperature", "top_p")) +REASONING_DISABLED_EFFORTS: Final[frozenset[str]] = frozenset(("none", "minimal")) +REASONING_ENABLED_EFFORTS: Final[frozenset[str]] = frozenset(("low", "medium", "high")) + + +class _FalUsage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + input_tokens: int + output_tokens: int + + +class _FalChatResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + output: str + usage_info: _FalUsage + reasoning: str | None = None + finish_reason: str | None = None + + +_CHAT_RESPONSE: Final = TypeAdapter(_FalChatResponse) + + +class FalAIError(BaseLLMException): + def __init__( + self, + status_code: int, + message: str, + headers: dict | httpx.Headers | None = None, # mutable-ok: BaseLLMException header contract + ) -> None: + super().__init__(status_code=status_code, message=message, headers=headers) + + +def _image_part_url(part: Mapping[str, object]) -> str | None: + image_url: Final = part.get("image_url") + if isinstance(image_url, str): + return image_url + if isinstance(image_url, Mapping): + url: Final = image_url.get("url") + return url if isinstance(url, str) else None + return None + + +def _prompt_and_image(messages: Sequence[AllMessageValues]) -> tuple[str, str]: + if len(messages) != 1 or messages[0].get("role") != "user": + raise FalAIError( + status_code=400, + message="fal_ai chat completions accept exactly one user message; system prompts and multi-turn history are not supported", + ) + content: Final = messages[0].get("content") + if isinstance(content, str): + if not content: + raise FalAIError(status_code=400, message="fal_ai chat completions require text in the user message") + raise FalAIError( + status_code=400, + message="fal_ai chat completions require exactly one image_url content part in the user message", + ) + parts: Final[tuple[Mapping[str, object], ...]] = ( + tuple(part for part in content if isinstance(part, Mapping)) if isinstance(content, Sequence) else () + ) + prompt: Final = "\n".join( + text for part in parts if part.get("type") == "text" and isinstance((text := part.get("text")), str) and text + ) + image_urls: Final = tuple( + url for part in parts if part.get("type") == "image_url" and (url := _image_part_url(part)) is not None + ) + if not prompt: + raise FalAIError(status_code=400, message="fal_ai chat completions require text in the user message") + if len(image_urls) != 1: + raise FalAIError( + status_code=400, + message="fal_ai chat completions require exactly one image_url content part in the user message", + ) + return prompt, image_urls[0] + + +class FalAIChatConfig(BaseConfig): + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return api_key or get_secret_str("FAL_AI_API_KEY") + + @staticmethod + def get_api_base(api_base: str | None = None) -> str: + return (api_base or get_secret_str("FAL_AI_API_BASE") or DEFAULT_BASE_URL).rstrip("/") + + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: inherited contract returns a list + return list(("reasoning_effort", "temperature", "top_p")) # mutable-ok: inherited contract returns a list + + def _map_reasoning_effort(self, value: object, model: str, drop_params: bool) -> bool | None: + if value in REASONING_DISABLED_EFFORTS: + return False + if value in REASONING_ENABLED_EFFORTS: + return True + if drop_params: + return None + raise FalAIError(status_code=400, message=f"Unsupported reasoning_effort '{value}' for {model}") + + def _translate_param(self, param: str, value: object, model: str, drop_params: bool) -> tuple[str, object] | None: + if param in ("temperature", "top_p"): + return param, value + if param == "reasoning_effort": + reasoning: Final = self._map_reasoning_effort(value, model, drop_params) + return ("reasoning", reasoning) if reasoning is not None else None + return None + + def map_openai_params( + self, + non_default_params: dict, # mutable-ok: inherited contract + optional_params: dict, # mutable-ok: inherited contract + model: str, + drop_params: bool, + ) -> dict: # mutable-ok: inherited contract returns a dict + mapped: Final = { # mutable-ok: intermediate translation map, folded into the returned dict + translated[0]: translated[1] + for param, value in non_default_params.items() + if (translated := self._translate_param(param, value, model, drop_params)) is not None + } + return {**optional_params, **mapped} # mutable-ok: inherited contract returns a dict + + def validate_environment( + self, + headers: dict, # mutable-ok: inherited contract + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict, # mutable-ok: inherited contract + litellm_params: dict, # mutable-ok: inherited contract + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: # mutable-ok: inherited contract returns a dict + final_api_key: Final = self.get_api_key(api_key) + if not final_api_key: + raise ValueError("FAL_AI_API_KEY is not set") + return { # mutable-ok: inherited contract returns a dict + "content-type": "application/json", + **(headers or {}), # mutable-ok: empty default for the inherited contract's headers + "Authorization": f"Key {final_api_key}", + } + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict, # mutable-ok: inherited contract + litellm_params: dict, # mutable-ok: inherited contract + stream: bool | None = None, + ) -> str: + return f"{self.get_api_base(api_base)}/{model.removeprefix(PROVIDER_PREFIX)}" + + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict, # mutable-ok: inherited contract + litellm_params: dict, # mutable-ok: inherited contract + headers: dict, # mutable-ok: inherited contract + ) -> dict: # mutable-ok: inherited contract returns a dict + if optional_params.get("stream"): + raise FalAIError(status_code=400, message="fal_ai chat completions do not support streaming") + prompt, image_url = _prompt_and_image(messages) + return { # mutable-ok: JSON request body + "prompt": prompt, + "image_url": image_url, + **{ # mutable-ok: JSON request body + key: value for key, value in optional_params.items() if key in PASSTHROUGH_PARAMS and value is not None + }, + } + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: dict, # mutable-ok: inherited contract + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict, # mutable-ok: inherited contract + litellm_params: dict, # mutable-ok: inherited contract + encoding: "tiktoken.Encoding | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + try: + completion_response: Final = _CHAT_RESPONSE.validate_json(raw_response.content) + except ValueError: + raise FalAIError( + status_code=422, + message=f"fal_ai returned an unexpected response body: {raw_response.text}", + headers=raw_response.headers, + ) + + message: Final = Message( + content=completion_response.output, + role="assistant", + reasoning_content=completion_response.reasoning, + ) + model_response.choices[0].message = message # rebind-ok: ModelResponse populated in place per contract + model_response.choices[0].finish_reason = map_finish_reason( # rebind-ok: same contract + completion_response.finish_reason or "stop" + ) + model_response.created = int(time.time()) # rebind-ok: same contract + model_response.model = model # rebind-ok: same contract + model_response.usage = Usage( # rebind-ok: same contract + prompt_tokens=completion_response.usage_info.input_tokens, + completion_tokens=completion_response.usage_info.output_tokens, + total_tokens=completion_response.usage_info.input_tokens + completion_response.usage_info.output_tokens, + ) + return model_response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict | httpx.Headers, # mutable-ok: inherited contract + ) -> BaseLLMException: + return FalAIError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index fd7d82d314f..31f0995bf9f 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -1,3 +1,4 @@ +import os from collections.abc import Mapping from math import ceil from types import MappingProxyType @@ -9,7 +10,8 @@ import litellm from litellm.types.utils import ImageObject, ImageResponse FAL_KEYED_PRICING_DEFAULT_QUALITY: Final[str] = "high" -FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = "1024-x-768" +_DEFAULT_KEYED_DIMENSIONS: Final[tuple[int, int]] = (1024, 768) +FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = f"{_DEFAULT_KEYED_DIMENSIONS[0]}-x-{_DEFAULT_KEYED_DIMENSIONS[1]}" FAL_PIXELS_PER_MEGAPIXEL: Final[int] = 1_048_576 FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( { @@ -24,6 +26,12 @@ FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( _OBJECT_MAP: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +FAL_AI_QUEUE_DEFAULT_BASE: Final[str] = "https://queue.fal.run" + + +def fal_ai_queue_base() -> str: + return os.getenv("FAL_AI_QUEUE_API_BASE") or FAL_AI_QUEUE_DEFAULT_BASE + def _keyed_size(optional_params: Mapping[str, object]) -> str | None: image_size: Final = optional_params.get("image_size") @@ -55,37 +63,56 @@ def _image_dimensions(image: object) -> tuple[int, int] | None: return width, height -def _response_size(image: object) -> str | None: - dimensions: Final = _image_dimensions(image) - if dimensions is None: - return None - width, height = dimensions - return f"{width}-x-{height}" - - def _keyed_quality(optional_params: Mapping[str, object]) -> str: raw_quality: Final = optional_params.get("quality") return raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY +def _parse_keyed_dimensions(size: str | None) -> tuple[int, int] | None: + if size is None: + return None + parts: Final = tuple(size.split("-x-")) + if len(parts) != 2: + return None + try: + width, height = (int(part) for part in parts) + except ValueError: + return None + return (width, height) if width > 0 and height > 0 else None + + +def _keyed_rows(model: str, quality: str) -> tuple[tuple[int, int, float], ...]: + prefix: Final = f"fal_ai/{quality}/" + suffix: Final = f"/{model}" + return tuple( + (width, height, float(raw_cost)) + for key in litellm.model_cost + if isinstance(key, str) and key.startswith(prefix) and key.endswith(suffix) + for size in (key[len(prefix) : -len(suffix)],) + for dimensions in (_parse_keyed_dimensions(size),) + if dimensions is not None + for entry in (_entry(key),) + if entry is not None + for raw_cost in (entry.get("output_cost_per_image"),) + if isinstance(raw_cost, (int, float)) + for width, height in (dimensions,) + ) + + def _keyed_cost_per_image( model: str, image: object, optional_params: Mapping[str, object], ) -> float | None: quality: Final = _keyed_quality(optional_params) - request_size: Final = _keyed_size(optional_params) or FAL_TEXT_TO_IMAGE_DEFAULT_SIZE - sizes: Final = (_response_size(image), request_size, FAL_TEXT_TO_IMAGE_DEFAULT_SIZE) - for size in sizes: - if size is None: - continue - keyed_entry = _entry(f"fal_ai/{quality}/{size}/{model}") - if keyed_entry is None: - continue - keyed_cost = keyed_entry.get("output_cost_per_image") - if isinstance(keyed_cost, (int, float)): - return float(keyed_cost) - return None + rows: Final = _keyed_rows(model, quality) + if not rows: + return None + target_dimensions: Final = ( + _image_dimensions(image) or _parse_keyed_dimensions(_keyed_size(optional_params)) or _DEFAULT_KEYED_DIMENSIONS + ) + target_pixels: Final = target_dimensions[0] * target_dimensions[1] + return min(rows, key=lambda row: (abs(row[0] * row[1] - target_pixels), row[0] * row[1]))[2] def _flat_cost_per_image( @@ -108,6 +135,16 @@ def _entry(key: str) -> Mapping[str, object] | None: return _OBJECT_MAP.validate_python(raw_entry) +def fal_ai_passthrough_cost(model: str, request_body: Mapping[str, object]) -> float | None: + entry: Final = _entry(f"{litellm.LlmProviders.FAL_AI.value}/{model}") + if entry is None: + return None + resolution: Final = request_body.get("resolution") + keyed_cost: Final = entry.get(f"output_cost_per_image_{resolution}") if isinstance(resolution, int) else None + cost: Final = keyed_cost if isinstance(keyed_cost, (int, float)) else entry.get("output_cost_per_image") + return float(cost) if isinstance(cost, (int, float)) else None + + def cost_calculator( model: str, image_response: object, @@ -129,7 +166,7 @@ def cost_calculator( ) for image in images ) - if all(cost is not None for cost in keyed_costs): + if not any(cost is None for cost in keyed_costs): return sum(cost for cost in keyed_costs if cost is not None) model_info: Final = litellm.get_model_info( model=normalized_model, @@ -144,10 +181,12 @@ def cost_calculator( float(raw_output_cost_per_pixel) if isinstance(raw_output_cost_per_pixel, (int, float)) else None ) return sum( - _flat_cost_per_image( + keyed_cost + if keyed_cost is not None + else _flat_cost_per_image( image=image, output_cost_per_image=output_cost_per_image, output_cost_per_pixel=output_cost_per_pixel, ) - for image in images + for image, keyed_cost in zip(images, keyed_costs) ) diff --git a/litellm/llms/fal_ai/image_edit/__init__.py b/litellm/llms/fal_ai/image_edit/__init__.py index c2f0f311f8c..60775ef2c8d 100644 --- a/litellm/llms/fal_ai/image_edit/__init__.py +++ b/litellm/llms/fal_ai/image_edit/__init__.py @@ -1,3 +1,24 @@ +from typing import Final + +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig + +from .flux_lora_depth_transformation import FalAIFluxLoraDepthEditConfig from .transformation import FalAIImageEditConfig -__all__ = ("FalAIImageEditConfig",) +__all__ = ("FalAIFluxLoraDepthEditConfig", "FalAIImageEditConfig") + + +def get_fal_ai_image_edit_config(model: str) -> BaseImageEditConfig: + """ + Get the appropriate Fal AI image edit configuration based on the model. + + Args: + model: The Fal AI model name (e.g., "openai/gpt-image-2.5/flare/edit", "fal-ai/flux-lora-depth") + + Returns: + The appropriate configuration class for the specified model + """ + model_lower: Final = model.lower() + if "flux-lora-depth" in model_lower: + return FalAIFluxLoraDepthEditConfig() + return FalAIImageEditConfig() diff --git a/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py new file mode 100644 index 00000000000..fa469d638d2 --- /dev/null +++ b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py @@ -0,0 +1,75 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from httpx._types import RequestFiles + +from litellm.secret_managers.main import get_secret_str +from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import FileTypes + +from .transformation import DEFAULT_BASE_URL, FalAIImageEditConfig, to_data_url + +FLUX_LORA_DEPTH_ENDPOINT: Final[str] = "fal-ai/flux-lora-depth" +SUPPORTED_OPENAI_PARAMS: Final[tuple[str, ...]] = ("n", "size") +PARAM_TRANSLATION: Final[Mapping[str, str]] = MappingProxyType({"n": "num_images", "size": "image_size"}) + + +class FalAIFluxLoraDepthEditConfig(FalAIImageEditConfig): + """ + FLUX.1 [dev] depth LoRA edit endpoint served through Fal AI. + + Unlike the openai gpt-image ``/edit`` endpoints, this endpoint takes a single ``image_url`` + control image and has no ``/edit`` path suffix. + """ + + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a list + return list(SUPPORTED_OPENAI_PARAMS) # mutable-ok: base class contract returns a list + + def map_openai_params( # mutable-ok: base class contract returns a dict + self, + image_edit_optional_params: ImageEditOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict: # mutable-ok: base class contract returns a dict + return { # mutable-ok: base class contract returns a dict + PARAM_TRANSLATION.get(key, key): self._translate_value(key, value, model) + for key, value in image_edit_optional_params.items() + if value is not None and key in PARAM_TRANSLATION + } + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, # mutable-ok: base class contract + ) -> str: + base_url: Final = (api_base or get_secret_str("FAL_AI_API_BASE") or DEFAULT_BASE_URL).rstrip("/") + return f"{base_url}/{FLUX_LORA_DEPTH_ENDPOINT}" + + def transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | None, + image_edit_optional_request_params: dict, # mutable-ok: base class contract + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: base class contract + ) -> tuple[dict, RequestFiles]: # mutable-ok: base class contract returns a dict + images: Final = tuple(img for img in (image if isinstance(image, list) else (image,)) if img is not None) + if not images: + raise ValueError("Fal AI image edit requires at least one input image") + if len(images) > 1: + raise ValueError(f"{FLUX_LORA_DEPTH_ENDPOINT} accepts exactly one control image") + provider_params: Final[Mapping[str, object]] = MappingProxyType( + { + key: value for key, value in image_edit_optional_request_params.items() if key != "mask" + } # mutable-ok: frozen by MappingProxyType + ) + request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict + "prompt": prompt, + "image_url": to_data_url(next(iter(images))), + **provider_params, + } + return request_body, () diff --git a/litellm/llms/fal_ai/image_edit/transformation.py b/litellm/llms/fal_ai/image_edit/transformation.py index 70b5d0612f2..6e6a872839a 100644 --- a/litellm/llms/fal_ai/image_edit/transformation.py +++ b/litellm/llms/fal_ai/image_edit/transformation.py @@ -61,7 +61,7 @@ def _read_image_bytes(image: object) -> bytes: raise ValueError(f"Unsupported image type for Fal AI image edit: {type(image).__name__}") -def _to_data_url(image: object) -> str: +def to_data_url(image: object) -> str: if isinstance(image, str): return image image_bytes: Final = _read_image_bytes(image) @@ -143,7 +143,7 @@ class FalAIImageEditConfig(BaseImageEditConfig): raise ValueError("Fal AI image edit requires at least one input image") mask: Final = _first(image_edit_optional_request_params.get("mask")) mask_field: Final[Mapping[str, str]] = ( - MappingProxyType({"mask_url": _to_data_url(mask)}) if mask is not None else MappingProxyType({}) + MappingProxyType({"mask_url": to_data_url(mask)}) if mask is not None else MappingProxyType({}) ) provider_params: Final[Mapping[str, object]] = MappingProxyType( { @@ -152,7 +152,7 @@ class FalAIImageEditConfig(BaseImageEditConfig): ) request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict "prompt": prompt, - "image_urls": tuple(_to_data_url(img) for img in images), + "image_urls": tuple(to_data_url(img) for img in images), **mask_field, **provider_params, } diff --git a/litellm/llms/fal_ai/image_generation/bria_transformation.py b/litellm/llms/fal_ai/image_generation/bria_transformation.py index c528550811a..53da48e62ca 100644 --- a/litellm/llms/fal_ai/image_generation/bria_transformation.py +++ b/litellm/llms/fal_ai/image_generation/bria_transformation.py @@ -8,9 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -187,7 +186,7 @@ class FalAIBriaConfig(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py index 6b8558b8124..7c63e1077f1 100644 --- a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py @@ -8,9 +8,8 @@ from litellm.types.utils import ImageResponse from .transformation import FalAIBaseConfig, fal_images_to_image_objects if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -194,7 +193,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py index 04b4f426878..ad1852a622b 100644 --- a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py @@ -8,9 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -150,7 +149,7 @@ class FalAIIdeogramV3Config(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py index 8a6665b2585..3624b76a4a3 100644 --- a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py +++ b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py @@ -8,9 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -182,7 +181,7 @@ class FalAIImagen4Config(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py index 4880dfec7e3..934ce420d53 100644 --- a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py @@ -8,9 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -172,7 +171,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py index bc3a4d07282..79d8800773b 100644 --- a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py +++ b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py @@ -8,9 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -208,7 +207,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index fd8e280da1c..8f081c7228d 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -16,9 +16,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -117,7 +116,7 @@ class FalAIBaseConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 28ebb39a303..196022b3558 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -46,7 +46,7 @@ from ..common_utils import ( ) if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer def _map_reasoning_effort(value: object) -> object: @@ -708,7 +708,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index d009fe4cd72..bb8d7455031 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -24,9 +24,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -173,7 +172,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index d9250ea8836..c047dc0c881 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -26,9 +26,8 @@ from ..authenticator import get_access_token from ..file_handler import upload_file_sync if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -416,7 +415,7 @@ class GigaChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: tiktoken.Encoding | None, + encoding: Tokenizer | None, api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index 41a2df17c6f..1da7ad0a7b5 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -27,7 +27,7 @@ from litellm.types.utils import ModelResponse, ModelResponseStream, ServerToolUs from ...openai_like.chat.transformation import OpenAILikeChatConfig if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer GROQ_COMPOUND_MODELS: Final = frozenset({"compound", "compound-mini"}) @@ -286,7 +286,7 @@ class GroqChatConfig(OpenAILikeChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 57d1357ee46..60917c68221 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -1,6 +1,5 @@ import json import os -from collections.abc import Sequence from typing import Final, Literal, Protocol, get_args import httpx @@ -32,7 +31,7 @@ hf_tasks_embeddings: Final = ( class _SupportsTokenEncode(Protocol): """Token encoder handle. Only ``encode`` is ever called on it here.""" - def encode(self, text: str, *, disallowed_special: tuple[str, ...]) -> Sequence[int]: ... + def encode(self, text: str) -> list[int]: ... def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> str | None: @@ -214,7 +213,7 @@ class HuggingFaceEmbedding(BaseLLM): model_response.model = model input_tokens = 0 for text in input: - input_tokens += len(encoding.encode(text, disallowed_special=())) + input_tokens += len(encoding.encode_ordinary(text)) setattr( model_response, diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 33b0e21e326..3fdd4abda73 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -25,9 +25,8 @@ from litellm.utils import token_counter from ..common_utils import HuggingFaceError, hf_task_list, hf_tasks, output_parser if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LoggingClass = LiteLLMLoggingObj else: @@ -479,7 +478,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/langflow/chat/transformation.py b/litellm/llms/langflow/chat/transformation.py index 17ae7017cf6..a53887b36af 100644 --- a/litellm/llms/langflow/chat/transformation.py +++ b/litellm/llms/langflow/chat/transformation.py @@ -14,9 +14,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import CustomStreamWrapper @@ -225,7 +224,7 @@ class LangFlowConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index 84d79e6bd31..c9388ee472f 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -23,9 +23,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import CustomStreamWrapper @@ -415,7 +414,7 @@ class LangGraphConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index c01ad2a0edd..341e8dd2e12 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -19,7 +19,7 @@ from litellm.types.utils import ModelResponse from ...openai_like.chat.transformation import OpenAILikeChatConfig if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class LemonadeChatConfig(OpenAILikeChatConfig): @@ -231,7 +231,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index f77e828b59a..33b567e9710 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -32,7 +32,7 @@ from litellm.types.utils import ModelResponse, ModelResponseStream from litellm.utils import convert_to_model_response_object, supports_reasoning if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer def _accepted_reasoning_effort(model: str, requested: str, custom_llm_provider: str) -> str: @@ -580,7 +580,7 @@ class MistralConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/nlp_cloud/chat/transformation.py b/litellm/llms/nlp_cloud/chat/transformation.py index 17c547618d3..2ff48894619 100644 --- a/litellm/llms/nlp_cloud/chat/transformation.py +++ b/litellm/llms/nlp_cloud/chat/transformation.py @@ -14,9 +14,8 @@ from litellm.utils import ModelResponse, Usage from ..common_utils import NLPCloudError if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LoggingClass = LiteLLMLoggingObj else: @@ -175,7 +174,7 @@ class NLPCloudConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 8e4e41b4ac1..ecff823a18d 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -65,9 +65,8 @@ from litellm.types.utils import ( from litellm.utils import supports_reasoning if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -603,7 +602,7 @@ class OCIChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index bcde8a041a6..cb3080e6534 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -31,9 +31,8 @@ from litellm.types.utils import ModelResponse, ModelResponseStream from ..common_utils import OllamaError if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -321,7 +320,7 @@ class OllamaChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 3eb2c833094..0fc1cd926b8 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -35,9 +35,8 @@ from litellm.types.utils import ( from ..common_utils import OllamaError, OllamaModelInfo, _convert_image if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -252,7 +251,7 @@ class OllamaConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/oobabooga/chat/transformation.py b/litellm/llms/oobabooga/chat/transformation.py index 43d627102b6..05383a35389 100644 --- a/litellm/llms/oobabooga/chat/transformation.py +++ b/litellm/llms/oobabooga/chat/transformation.py @@ -11,9 +11,8 @@ from litellm.types.utils import ModelResponse, Usage from ..common_utils import OobaboogaError if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LoggingClass = LiteLLMLoggingObj else: @@ -39,7 +38,7 @@ class OobaboogaConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 9dbcf0cc089..b63684db782 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -60,9 +60,8 @@ from litellm.utils import convert_to_model_response_object from ..common_utils import OpenAIError if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.types.llms.openai import ChatCompletionToolParam @@ -671,7 +670,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openai/image_generation/dall_e_2_transformation.py b/litellm/llms/openai/image_generation/dall_e_2_transformation.py index 74936cf1895..ffc6f1d5fe9 100644 --- a/litellm/llms/openai/image_generation/dall_e_2_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_2_transformation.py @@ -10,9 +10,10 @@ from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object if TYPE_CHECKING: - import tiktoken from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + class DallE2ImageGenerationConfig(BaseImageGenerationConfig): """ @@ -52,7 +53,7 @@ class DallE2ImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/openai/image_generation/dall_e_3_transformation.py b/litellm/llms/openai/image_generation/dall_e_3_transformation.py index 5c561d011a9..90b7eaedf2f 100644 --- a/litellm/llms/openai/image_generation/dall_e_3_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_3_transformation.py @@ -10,9 +10,10 @@ from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object if TYPE_CHECKING: - import tiktoken from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + class DallE3ImageGenerationConfig(BaseImageGenerationConfig): """ @@ -52,7 +53,7 @@ class DallE3ImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index 8dc4d8953ea..c3a826616ed 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -10,9 +10,10 @@ from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object if TYPE_CHECKING: - import tiktoken from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + class GPTImageGenerationConfig(BaseImageGenerationConfig): """ @@ -61,7 +62,7 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/openai/image_variations/transformation.py b/litellm/llms/openai/image_variations/transformation.py index afd2909b697..73b44d5ea6a 100644 --- a/litellm/llms/openai/image_variations/transformation.py +++ b/litellm/llms/openai/image_variations/transformation.py @@ -12,7 +12,7 @@ from ...base_llm.image_variations.transformation import BaseImageVariationConfig from ..common_utils import OpenAIError if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class OpenAIImageVariationConfig(BaseImageVariationConfig): @@ -53,7 +53,7 @@ class OpenAIImageVariationConfig(BaseImageVariationConfig): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, ) -> ImageResponse: return model_response @@ -68,7 +68,7 @@ class OpenAIImageVariationConfig(BaseImageVariationConfig): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, ) -> ImageResponse: return model_response diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 869ad387c5a..7ac0d988074 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -6,9 +6,10 @@ from typing import TYPE_CHECKING, Final, Literal, Optional, cast import httpx if TYPE_CHECKING: - import tiktoken from aiohttp import ClientSession + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + import openai from openai import AsyncOpenAI, OpenAI from openai._base_client import make_request_options @@ -277,7 +278,7 @@ class OpenAIConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openai/videos/guardrail_translation/__init__.py b/litellm/llms/openai/videos/guardrail_translation/__init__.py new file mode 100644 index 00000000000..7bd869612d6 --- /dev/null +++ b/litellm/llms/openai/videos/guardrail_translation/__init__.py @@ -0,0 +1,23 @@ +"""OpenAI Video Generation handler for Unified Guardrails.""" + +from typing import Final + +from litellm.llms.openai.videos.guardrail_translation.handler import ( + OpenAIVideoGenerationHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings: Final = { # mutable-ok: discover_guardrail_translation_mappings only accepts isinstance(mappings, dict) + CallTypes.video_generation: OpenAIVideoGenerationHandler, + CallTypes.avideo_generation: OpenAIVideoGenerationHandler, + CallTypes.create_video: OpenAIVideoGenerationHandler, + CallTypes.acreate_video: OpenAIVideoGenerationHandler, + CallTypes.video_remix: OpenAIVideoGenerationHandler, + CallTypes.avideo_remix: OpenAIVideoGenerationHandler, + CallTypes.video_edit: OpenAIVideoGenerationHandler, + CallTypes.avideo_edit: OpenAIVideoGenerationHandler, + CallTypes.video_extension: OpenAIVideoGenerationHandler, + CallTypes.avideo_extension: OpenAIVideoGenerationHandler, +} + +__all__ = ("OpenAIVideoGenerationHandler", "guardrail_translation_mappings") diff --git a/litellm/llms/openai/videos/guardrail_translation/handler.py b/litellm/llms/openai/videos/guardrail_translation/handler.py new file mode 100644 index 00000000000..49a8d05100c --- /dev/null +++ b/litellm/llms/openai/videos/guardrail_translation/handler.py @@ -0,0 +1,48 @@ +from typing import TYPE_CHECKING, Final + +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth + + +class OpenAIVideoGenerationHandler(BaseTranslation): + async def process_input_messages( + self, + data: dict[str, object], # mutable-ok: BaseTranslation contract passes the proxy's request dict through + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> dict[str, object]: # mutable-ok: BaseTranslation contract returns the proxy's request dict + prompt: Final = data.get("prompt") + if not isinstance(prompt, str): + return data + + model: Final = data.get("model") + texts: Final = [prompt] # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str] + inputs: Final = ( + GenericGuardrailAPIInputs(texts=texts, model=model) + if isinstance(model, str) + else GenericGuardrailAPIInputs(texts=texts) + ) + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # request_data is a bare dict + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + guardrailed_texts: Final = guardrailed_inputs.get("texts") + guardrailed_prompt: Final = guardrailed_texts[0] if guardrailed_texts else prompt + return {**data, "prompt": guardrailed_prompt} # mutable-ok: BaseTranslation contract returns a dict + + async def process_output_response( + self, + response: object, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, + request_data: dict[str, object] | None = None, # mutable-ok: BaseTranslation contract + ) -> object: + return response diff --git a/litellm/llms/openai_like/chat/transformation.py b/litellm/llms/openai_like/chat/transformation.py index 030710c8b2d..e5d6cbb7e5e 100644 --- a/litellm/llms/openai_like/chat/transformation.py +++ b/litellm/llms/openai_like/chat/transformation.py @@ -13,9 +13,8 @@ from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -131,7 +130,7 @@ class OpenAILikeChatConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index 77a902149d9..08d43c169f4 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -23,9 +23,8 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig from ..common_utils import OpenRouterException if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class CacheControlSupportedModels(str, Enum): @@ -182,7 +181,7 @@ class OpenrouterConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index 6bbda324336..67d90d027ec 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -50,9 +50,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer else: LiteLLMLoggingObj = Any @@ -319,7 +318,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index 354f7692fd5..dca2f9857b8 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -15,7 +15,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionAnnotation from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class PerplexityChatConfig(OpenAIGPTConfig): @@ -75,7 +75,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/petals/completion/transformation.py b/litellm/llms/petals/completion/transformation.py index 3e0de14a7b2..ee20c2b12d5 100644 --- a/litellm/llms/petals/completion/transformation.py +++ b/litellm/llms/petals/completion/transformation.py @@ -14,7 +14,7 @@ from litellm.types.utils import ModelResponse from ..common_utils import PetalsError if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class PetalsConfig(BaseConfig): @@ -112,7 +112,7 @@ class PetalsConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 2a63c489395..69924396a1e 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -18,9 +18,8 @@ from litellm.types.utils import Choices, Message, ModelResponse, Usage from ..common_utils import PredibaseError if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -150,7 +149,7 @@ class PredibaseConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/recraft/image_generation/transformation.py b/litellm/llms/recraft/image_generation/transformation.py index 3a04e0a62b4..f65bf1e7292 100644 --- a/litellm/llms/recraft/image_generation/transformation.py +++ b/litellm/llms/recraft/image_generation/transformation.py @@ -14,9 +14,8 @@ from litellm.types.llms.recraft import RecraftImageGenerationRequestParams from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -122,7 +121,7 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/replicate/chat/transformation.py b/litellm/llms/replicate/chat/transformation.py index 769160c6ced..f7e09b7bec0 100644 --- a/litellm/llms/replicate/chat/transformation.py +++ b/litellm/llms/replicate/chat/transformation.py @@ -19,9 +19,8 @@ from litellm.utils import token_counter from ..common_utils import ReplicateError if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LoggingClass = LiteLLMLoggingObj else: @@ -237,7 +236,7 @@ class ReplicateConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/runwayml/image_generation/transformation.py b/litellm/llms/runwayml/image_generation/transformation.py index 5913709c8a0..e5e988328d8 100644 --- a/litellm/llms/runwayml/image_generation/transformation.py +++ b/litellm/llms/runwayml/image_generation/transformation.py @@ -22,9 +22,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -308,7 +307,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: @@ -383,7 +382,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py index 576018f0046..e1bc496a82d 100644 --- a/litellm/llms/sagemaker/completion/transformation.py +++ b/litellm/llms/sagemaker/completion/transformation.py @@ -24,9 +24,8 @@ from litellm.utils import token_counter from ..common_utils import SagemakerError if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -198,7 +197,7 @@ class SagemakerConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index d64d7a57281..4c73ccacc16 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -15,9 +15,8 @@ from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -57,7 +56,7 @@ def validate_dict(data: dict, model) -> dict: return model(**data).model_dump(by_alias=True, exclude_unset=True) -def _messages_to_sap_template(messages: list[dict[str, str]]) -> list: +def _messages_to_sap_template(messages: list[AllMessageValues]) -> list: template: Final = [] for message in messages: if message["role"] == "user": @@ -311,7 +310,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): def transform_request( self, model: str, - messages: list[dict[str, str]], + messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, headers: dict, @@ -383,7 +382,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/stability/image_generation/transformation.py b/litellm/llms/stability/image_generation/transformation.py index cf3576a9404..656ffe395c8 100644 --- a/litellm/llms/stability/image_generation/transformation.py +++ b/litellm/llms/stability/image_generation/transformation.py @@ -26,9 +26,8 @@ from litellm.types.llms.stability import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -207,7 +206,7 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/topaz/image_variations/transformation.py b/litellm/llms/topaz/image_variations/transformation.py index f4753c8ba17..94f60d29cb9 100644 --- a/litellm/llms/topaz/image_variations/transformation.py +++ b/litellm/llms/topaz/image_variations/transformation.py @@ -23,7 +23,7 @@ from ...base_llm.image_variations.transformation import BaseImageVariationConfig from ..common_utils import TopazException, TopazModelInfo if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): @@ -139,7 +139,7 @@ class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, ) -> ImageResponse: image_content: Final = await raw_response.read() @@ -158,7 +158,7 @@ class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, ) -> ImageResponse: image_content: Final = raw_response.content diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index 3c868b3a96f..b37bbd78f2b 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -29,7 +29,7 @@ from litellm.types.utils import ( from ..common_utils import TritonError if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class TritonConfig(BaseConfig): @@ -95,7 +95,7 @@ class TritonConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -215,7 +215,7 @@ class TritonGenerateConfig(TritonConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -280,7 +280,7 @@ class TritonInferConfig(TritonConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index c5ca9f38144..b37bf473731 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -29,9 +29,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import CustomStreamWrapper @@ -285,7 +284,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index b2c52c53580..17ccf16837b 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -24,9 +24,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -284,7 +283,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 8faf7b0d484..ae6e08611ae 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -20,9 +20,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -214,7 +213,7 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 508f68b3eca..2fab6f438f6 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -14,7 +14,7 @@ from ....anthropic.chat.transformation import AnthropicConfig from .output_params_utils import sanitize_vertex_anthropic_output_params if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class VertexAIError(Exception): @@ -197,7 +197,7 @@ class VertexAIAnthropicConfig(AnthropicConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py index 89a5b8a570e..f2d2c0896d2 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py @@ -21,7 +21,7 @@ from litellm.types.utils import ( from ...common_utils import VertexAIError if TYPE_CHECKING: - import tiktoken + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer class VertexAILlama3Config(OpenAIGPTConfig): @@ -112,7 +112,7 @@ class VertexAILlama3Config(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 67b01c2dc43..2ae8b4cd188 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -24,9 +24,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.base_llm.base_model_iterator import MockResponseIterator @@ -276,7 +275,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): litellm_params: dict, client: HTTPHandler | httpx.Client | None = None, timeout: float | httpx.Timeout | None = None, - encoding: "tiktoken.Encoding | None" = None, + encoding: "Tokenizer | None" = None, ): """Synchronous completion request""" from litellm.utils import convert_to_model_response_object @@ -366,7 +365,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): litellm_params: dict, client: AsyncHTTPHandler | httpx.AsyncClient | None = None, timeout: float | httpx.Timeout | None = None, - encoding: "tiktoken.Encoding | None" = None, + encoding: "Tokenizer | None" = None, ): """Asynchronous completion request""" from litellm.utils import convert_to_model_response_object diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index 2be007336b4..2fec8485cf9 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -21,9 +21,8 @@ from ..common_utils import ( ) if TYPE_CHECKING: - import tiktoken - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -280,7 +279,7 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: "tiktoken.Encoding | None", + encoding: "Tokenizer | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/main.py b/litellm/main.py index 6704358e3ea..7d231a1bf7a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -37,7 +37,6 @@ if TYPE_CHECKING: import dotenv import httpx import openai -import tiktoken from pydantic import BaseModel from typing_extensions import overload @@ -100,6 +99,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) +from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer from litellm.llms.azure_ai.common_utils import ( azure_ai_supports_native_responses, foundry_chat_rejects_function_tools_while_reasoning, @@ -3594,6 +3594,37 @@ def _complete_edenai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResu return response +def _complete_fal_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + if ctx.stream: + raise litellm.FalAIError( + status_code=400, + message="fal_ai chat completions do not support streaming", + ) + api_base: Final = litellm.FalAIChatConfig.get_api_base(ctx.api_base) + api_key: Final = litellm.FalAIChatConfig.get_api_key(ctx.api_key or litellm.api_key) + response: Final = base_llm_http_handler.completion( + model=ctx.model, + messages=ctx.messages, + api_base=api_base, + custom_llm_provider="fal_ai", + model_response=ctx.model_response, + encoding=_get_encoding(), + logging_obj=ctx.logging, + optional_params=ctx.optional_params, + timeout=ctx.timeout, + litellm_params=ctx.litellm_params, + shared_session=ctx.shared_session, + acompletion=ctx.acompletion, + stream=ctx.stream, + api_key=api_key, + headers=ctx.headers or litellm.headers, + client=_dispatch_client_http(ctx), + provider_config=ctx.provider_config, + ) + ctx.logging.post_call(input=ctx.messages, api_key=api_key, original_response=response) + return response + + def _complete_vertex_ai_beta( ctx: _CompletionDispatchContext, ) -> _CompletionDispatchResult: @@ -5799,6 +5830,8 @@ def completion( response = _complete_hosted_vllm(_dispatch_ctx) elif custom_llm_provider == "edenai": response = _complete_edenai(_dispatch_ctx) # rebind-ok: dispatch chain binds response per branch + elif custom_llm_provider == "fal_ai": + response = _complete_fal_ai(_dispatch_ctx) # rebind-ok: dispatch chain binds response per branch elif ( # A known OpenAI model name only decides the route when nothing else # resolved a provider. get_llm_provider() already maps these names to @@ -7456,7 +7489,9 @@ def text_completion( if isinstance(prompt, list): import concurrent.futures - tokenizer: Final = tiktoken.encoding_for_model("text-davinci-003") + from litellm.rust_bridge.tokenizer import get_encoding + + tokenizer: Final = get_encoding("p50k_base") ## if it's a 2d list - each element in the list is a text_completion() request if len(prompt) > 0 and isinstance(prompt[0], list): responses: Final = [None for x in prompt] # init responses @@ -9226,7 +9261,7 @@ async def acount_tokens( except Exception as e: verbose_logger.debug("Provider token counting failed for model=%s, falling back to local: %s", model, e) - # Fallback to local tiktoken-based token counting + # Fallback to local token counting fallback_messages = messages or [] if system and fallback_messages: fallback_messages = [{"role": "system", "content": system}] + fallback_messages @@ -9245,16 +9280,16 @@ async def acount_tokens( # Cache for encoding to avoid repeated __getattr__ calls -_encoding_cache: tiktoken.Encoding | None = None +_encoding_cache: Tokenizer | None = None -def _load_module_encoding() -> tiktoken.Encoding: +def _load_module_encoding() -> Tokenizer: import sys return sys.modules[__name__].encoding -def _get_encoding() -> tiktoken.Encoding: +def _get_encoding() -> Tokenizer: """Get encoding, loading it lazily if needed.""" global _encoding_cache if _encoding_cache is None: @@ -9263,18 +9298,15 @@ def _get_encoding() -> tiktoken.Encoding: return _encoding_cache -def _load_default_encoding() -> tiktoken.Encoding: +def _load_default_encoding() -> Tokenizer: from litellm._lazy_imports import _get_default_encoding return _get_default_encoding() -def __getattr__(name: str) -> tiktoken.Encoding: +def __getattr__(name: str) -> Tokenizer: """Lazy import handler for main module""" if name == "encoding": - # Use _get_default_encoding which properly sets TIKTOKEN_CACHE_DIR - # before loading tiktoken, ensuring the local cache is used - # instead of downloading from the internet _encoding: Final = _load_default_encoding() # Cache it in the module's __dict__ for subsequent accesses import sys diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index c75f6564d1b..a0c791a136c 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -4,7 +4,7 @@ from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable from litellm.llms.anthropic.experimental_pass_through.messages import handler as main -from litellm.rust_bridge.catalog import Context, Delivery, Route +from litellm.rust_bridge.catalog import Delivery, Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch, call_hook from litellm.rust_bridge.messages.entrypoints import ( NATIVE_AMESSAGES, @@ -71,8 +71,8 @@ def _public_request( ) -def _context(request: LiteLLMMessagesRequest) -> Context: - return Context( +def _context(request: LiteLLMMessagesRequest) -> RouteContext: + return RouteContext( Route.MESSAGES, provider=request.custom_llm_provider, model=request.model, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3ca9958ad0d..77cada25d25 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24966,6 +24966,52 @@ "/v1/images/generations" ] }, + "fal_ai/fal-ai/trellis": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.02, + "source": "https://fal.ai/models/fal-ai/trellis", + "metadata": { + "comment": "image-to-3D, returns a GLB mesh; served through the /fal_ai pass-through route" + } + }, + "fal_ai/fal-ai/trellis-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.3, + "output_cost_per_image_512": 0.25, + "output_cost_per_image_1024": 0.3, + "output_cost_per_image_1536": 0.35, + "source": "https://fal.ai/models/fal-ai/trellis-2", + "metadata": { + "comment": "image-to-3D, returns a GLB mesh; priced by the request's resolution field (default 1024); served through the /fal_ai pass-through route" + } + }, + "fal_ai/fal-ai/flux-lora-depth": { + "litellm_provider": "fal_ai", + "metadata": { + "notes": "fal bills fal-ai/flux-lora-depth at $0.035 per megapixel, rounding each image up to the nearest megapixel. The per-pixel rate is used when Fal reports the output size, and the flat per-image price prices the default 1 MP output like the sibling flux entries" + }, + "mode": "image_generation", + "output_cost_per_image": 0.035, + "output_cost_per_pixel": 3.337860107421875e-08, + "source": "https://fal.ai/models/fal-ai/flux-lora-depth", + "supported_endpoints": [ + "/v1/images/edits" + ] + }, + "fal_ai/fal-ai/moondream3-preview/query": { + "input_cost_per_token": 4e-07, + "litellm_provider": "fal_ai", + "mode": "chat", + "output_cost_per_token": 3.5e-06, + "source": "https://fal.ai/models/fal-ai/moondream3-preview/query", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_reasoning": true, + "supports_vision": true + }, "featherless_ai/featherless-ai/Qwerky-72B": { "litellm_provider": "featherless_ai", "max_input_tokens": 32768, @@ -43037,21 +43083,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.83746e-07, + "input_cost_per_token": 9.5526e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.767492e-06, + "output_cost_per_token": 1.91052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.36455e-08, + "cache_read_input_token_cost": 7.9605e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -44812,26 +44858,6 @@ "max_tokens": 128000, "mode": "chat" }, - "openrouter/stealth/union-alpha": { - "deprecation_date": "2098-12-31", - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "openrouter", - "max_input_tokens": 262144, - "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", - "source": "https://openrouter.ai/api/v1/models", - "supports_audio_input": false, - "supports_function_calling": true, - "supports_pdf_input": false, - "supports_prompt_caching": false, - "supports_reasoning": false, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_vision": true, - "supports_web_search": false - }, "ovhcloud/DeepSeek-R1-Distill-Llama-70B": { "input_cost_per_token": 6.7e-07, "litellm_provider": "ovhcloud", @@ -54217,7 +54243,7 @@ "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54227,6 +54253,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true, "supported_endpoints": [ "/v1/responses" @@ -54241,7 +54268,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54251,6 +54278,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-0309-reasoning": { @@ -54262,7 +54290,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -54271,6 +54299,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, @@ -54283,7 +54312,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, @@ -54292,11 +54321,13 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.3": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "xai", @@ -54306,7 +54337,7 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54318,6 +54349,7 @@ "xai/grok-4.3-latest": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "xai", @@ -54327,7 +54359,7 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54339,6 +54371,7 @@ "xai/grok-4.5": { "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54348,7 +54381,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54360,6 +54393,7 @@ "xai/grok-4.5-latest": { "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54369,7 +54403,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54381,6 +54415,7 @@ "xai/grok-build-latest": { "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54390,7 +54425,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/developers/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54402,6 +54437,7 @@ "xai/grok-4.6": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54411,7 +54447,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/developers/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54423,6 +54459,7 @@ "xai/grok-4.7": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54432,7 +54469,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/developers/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54450,7 +54487,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54460,7 +54497,8 @@ "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, "supports_vision": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_image_token": 1e-06 }, "xai/grok-code-fast-1": { "cache_read_input_token_cost": 2e-07, @@ -54471,7 +54509,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54481,7 +54519,8 @@ "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, "supports_vision": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_image_token": 1e-06 }, "xai/grok-code-fast-1-0825": { "cache_read_input_token_cost": 2e-07, @@ -54492,7 +54531,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54502,7 +54541,8 @@ "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, "supports_vision": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_image_token": 1e-06 }, "zai.glm-4.7": { "input_cost_per_token": 6e-07, @@ -61978,7 +62018,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, @@ -61987,6 +62027,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-0309": { @@ -61998,7 +62039,7 @@ "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": false, "supports_prompt_caching": true, "supports_reasoning": true, @@ -62008,6 +62049,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true, "supported_endpoints": [ "/v1/responses" @@ -62022,7 +62064,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -62030,6 +62072,7 @@ "input_cost_per_token_above_200k_tokens": 2e-06, "output_cost_per_token_above_200k_tokens": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1e-06, "supports_response_schema": true, "supports_vision": true }, @@ -65225,7 +65268,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -65234,6 +65277,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, @@ -65246,7 +65290,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -65255,6 +65299,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, @@ -65267,7 +65312,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -65276,6 +65321,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, @@ -65519,7 +65565,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, @@ -65528,6 +65574,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-non-reasoning-latest": { @@ -65539,7 +65586,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, @@ -65548,6 +65595,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent": { @@ -65559,7 +65607,7 @@ "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" ], @@ -65572,6 +65620,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-latest": { @@ -65583,7 +65632,7 @@ "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" ], @@ -65596,6 +65645,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "groq/qwen/qwen3.8-27b": { @@ -68252,9 +68302,9 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 9.1e-07, - "output_cost_per_token": 2.86e-06, - "cache_read_input_token_cost": 1.69e-07, + "input_cost_per_token": 8.4e-07, + "output_cost_per_token": 2.64e-06, + "cache_read_input_token_cost": 1.56e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -68941,9 +68991,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 5.544e-08, - "output_cost_per_token": 1.1088e-07, - "cache_read_input_token_cost": 1.1088e-08, + "input_cost_per_token": 8.8606e-08, + "output_cost_per_token": 1.77212e-07, + "cache_read_input_token_cost": 1.77212e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -72854,6 +72904,16 @@ "supports_reasoning": true, "supports_vision": true }, + "openrouter/typesafe/jev-1.13": { + "input_cost_per_token": 4.2e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 32000, + "max_output_tokens": 28800, + "max_tokens": 28800, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://openrouter.ai/typesafe/jev-1.13" + }, "typesafe/jev-1.13.0": { "input_cost_per_token": 4.2e-08, "litellm_provider": "typesafe", @@ -73272,14 +73332,14 @@ "supports_web_search": false }, "openrouter/~z-ai/glm-latest": { - "cache_read_input_token_cost": 1.69e-07, - "input_cost_per_token": 9.1e-07, + "cache_read_input_token_cost": 1.56e-07, + "input_cost_per_token": 8.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.86e-06, + "output_cost_per_token": 2.64e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -75083,6 +75143,7 @@ "supports_web_search": false }, "openrouter/nex-agi/nex-n2.5-mini:free": { + "deprecation_date": "2026-09-25", "input_cost_per_token": 0.0, "litellm_provider": "openrouter", "max_input_tokens": 262144, @@ -75102,6 +75163,7 @@ "supports_web_search": false }, "openrouter/nex-agi/nex-n2.5-pro:free": { + "deprecation_date": "2026-09-25", "input_cost_per_token": 0.0, "litellm_provider": "openrouter", "max_input_tokens": 262144, @@ -77044,5 +77106,467 @@ "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true + }, + "xai/grok-4.20-0309": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-0309": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-latest-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-latest-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-0304": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-0304-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-0304-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-non-reasoning-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-reasoning-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-multi-agent-beta-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "responses", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "xai/grok-4.20-multi-agent-experimental-beta-0304": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "responses", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "xai/grok-4.20-multi-agent-experimental-beta-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "responses", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "xai/grok-4.20-non-reasoning-gv2": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-reasoning-gv2": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "openrouter/nex-agi/nex-n2.5-mini": { + "cache_read_input_token_cost": 2.5e-09, + "input_cost_per_token": 2.5e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 235929, + "max_tokens": 235929, + "mode": "chat", + "output_cost_per_token": 1e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": false, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": false + }, + "openrouter/nex-agi/nex-n2.5-pro": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 7.5e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 235929, + "max_tokens": 235929, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": false, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": false } } diff --git a/litellm/ocr/dispatch.py b/litellm/ocr/dispatch.py index 55b19458b7a..b26175c943b 100644 --- a/litellm/ocr/dispatch.py +++ b/litellm/ocr/dispatch.py @@ -6,7 +6,7 @@ import httpx from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.ocr import main from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type -from litellm.rust_bridge.catalog import Context, Route +from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch, call_hook from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR, LiteLLMOcrRequest @@ -52,10 +52,10 @@ _PYTHON_AOCR: Final = cast( # cast-ok: forward the original call shape through ) -def _context(request: LiteLLMOcrRequest) -> Context: +def _context(request: LiteLLMOcrRequest) -> RouteContext: prefix, separator, _ = request.model.partition("/") provider: Final = request.custom_llm_provider or (prefix if separator else None) - return Context(Route.OCR, provider=provider, model=request.model) + return RouteContext(Route.OCR, provider=provider, model=request.model) _DISPATCH: Final = PublicDispatch( diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index b0640e4f0dd..60dc91a69cc 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -44,6 +44,11 @@ from litellm.proxy._types import ( UserAPIKeyAuth, user_api_key_has_admin_view, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + CeilingResolver, + resolve_agent_access_group_ceiling, +) +from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.user_api_key_auth import ( _get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth @@ -184,6 +189,21 @@ def _has_client_supplied_mcp_auth( return bool(mcp_auth_header) or bool(mcp_server_auth_headers) +def _agent_capped_servers( + allowed_mcp_servers: Sequence[str], + agent_servers: Sequence[str], + agent_access_group_servers: frozenset[str] | None, +) -> tuple[str, ...] | None: + if not agent_servers and agent_access_group_servers is None: + return None + return tuple( + s + for s in allowed_mcp_servers + if (not agent_servers or s in agent_servers) + and (agent_access_group_servers is None or s in agent_access_group_servers) + ) + + def _is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bool: """True when this auth is a keyless subject admitted by the gateway session / bridge user path, as opposed to a JWT or other keyless auth that merely lacks a ``team_id``. @@ -946,10 +966,11 @@ class MCPRequestHandler: on top of these direct grants, each source bounded by ITS OWN org, so a user spanning organizations cannot leak one org's servers past another's ceiling. - Error handling: ``get_user_object`` catches every DB failure and re-raises a bare ``ValueError``, so a - missing user and a real outage look identical (the cause survives only as ``__context__``). - ``_raise_503_if_db_unavailable`` walks the cause chain so an outage stays a retryable 503 while any - other failure fails closed as 401, not an opaque 500; the object-permission load shares that boundary.""" + Error handling: ``get_user_object`` lets a database outage propagate as-is and re-raises every other + DB failure as a bare ``ValueError`` (the cause surviving only as ``__context__``). + ``_raise_503_if_db_unavailable`` walks the cause chain so an outage stays a retryable 503 whichever + shape it arrives in, while any other failure fails closed as 401, not an opaque 500; the + object-permission load shares that boundary.""" from litellm.proxy.auth.auth_checks import get_object_permission, get_user_object from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -1096,9 +1117,8 @@ class MCPRequestHandler: (401) or surface as an opaque 500; the caller retries. Mirrors ``UserAPIKeyAuthExceptionHandler``, which renders a service-unavailable database error as 503 on the standard pipeline. - Classifies across the ``__cause__``/``__context__`` chain, not just ``e`` itself: ``get_user_object`` - re-raises every DB failure as a bare ``ValueError``, so a type-based check on the top exception - would miss a real outage wrapped inside it.""" + Classifies across the ``__cause__``/``__context__`` chain, not just ``e`` itself, so an outage a + caller re-raised inside a domain exception is still recognized.""" from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(e) @@ -1546,25 +1566,33 @@ class MCPRequestHandler: # Check agent permissions if agent_id is set on the key ######################################################### if user_api_key_auth and user_api_key_auth.agent_id: - allowed_mcp_servers_for_agent: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_agent( - user_api_key_auth + agent_capped: Final = _agent_capped_servers( + allowed_mcp_servers, + await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth), + await MCPRequestHandler._get_agent_access_group_server_ceiling(user_api_key_auth), ) - if len(allowed_mcp_servers_for_agent) > 0: + if agent_capped is not None: has_lower_level_mcp_restrictions = True - # Intersect: agent can only use servers allowed by BOTH key/team AND agent config - allowed_mcp_servers = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_agent] + allowed_mcp_servers = list(agent_capped) verbose_logger.debug( "Applied agent intersection filter. Final allowed servers: %s", allowed_mcp_servers ) + ######################################################### + # Cap an agent key at what the user and team that invoked the agent may reach + ######################################################### + caller_capped, caller_restricts = await MCPRequestHandler._apply_agent_caller_ceiling( + allowed_mcp_servers, user_api_key_auth + ) + ######################################################### # Apply the internal user's own ceiling (the entitlement attached to the human) ######################################################### capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling( - allowed_mcp_servers, user_api_key_auth, keyless_source=keyless_source + caller_capped, user_api_key_auth, keyless_source=keyless_source ) allowed_mcp_servers = list(capped) - has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or user_restricts + has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or caller_restricts or user_restricts ######################################################### # Apply org-level ceiling if org_id is set @@ -2907,6 +2935,28 @@ class MCPRequestHandler: verbose_logger.debug("Applied user ceiling filter. Final allowed servers: %s", capped) return capped, True + @staticmethod + async def _apply_agent_caller_ceiling( + allowed_mcp_servers: Sequence[str], + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> tuple[tuple[str, ...], bool]: + """Narrow an agent key's servers to those the invoking user and team (echoed back by the agent + as ``x-litellm-user-id`` / ``x-litellm-team-id``) may reach: the echoed team's grants when it + names any, then the echoed user's own entitlement. Raises like the user ceiling when that + entitlement is known but unreadable, so the resolver denies rather than widens.""" + caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None + if caller_auth is None: + return tuple(allowed_mcp_servers), False + team_servers: Final = frozenset(await MCPRequestHandler._get_allowed_mcp_servers_for_team(caller_auth)) + team_capped: Final = ( + tuple(server for server in allowed_mcp_servers if server in team_servers) + if team_servers + else tuple(allowed_mcp_servers) + ) + user_capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling(team_capped, caller_auth) + verbose_logger.debug("Applied agent caller ceiling. Final allowed servers: %s", user_capped) + return user_capped, bool(team_servers) or user_restricts + @staticmethod async def _user_places_mcp_ceiling(user_api_key_auth: UserAPIKeyAuth | None = None) -> bool: """Whether this human's own entitlement bounds their MCP access at all. @@ -3137,6 +3187,27 @@ class MCPRequestHandler: verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e) return [] + @staticmethod + async def _get_agent_access_group_server_ceiling( + user_api_key_auth: UserAPIKeyAuth, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + ) -> frozenset[str] | None: + """ + Server IDs the agent's attached unified access groups (``LiteLLM_AgentsTable.access_group_ids``) + allow, or None when the agent has none attached. Unlike the object_permission path above, an + attached group set that names no servers is an empty ceiling and denies every server. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + if not user_api_key_auth.agent_id: + return None + ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id) + if ceiling is None: + return None + return frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids))) + @staticmethod async def _get_agent_tool_permissions_for_server( server_id: str, diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 37a893973e3..6ae33cc1629 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -276,9 +276,9 @@ async def load_active_user_by_id( database-service-unavailable error is a retryable outage (``unavailable``). Everything else fails closed as ``no_active_key`` (the caller maps it to invalid_grant): a ``ProxyException`` / ``HTTPException``, a SCIM-deactivated user, and, unlike the key path, a missing user. ``get_user_object`` - catches every DB failure and re-raises a bare ``ValueError`` (a deleted user and a real outage look - identical, the original error surviving only as ``__context__``), so the outage check walks the cause - chain, and a missing user falls through to ``no_active_key`` rather than an opaque gateway fault. + lets a real outage propagate as-is and re-raises any other DB failure as a bare ``ValueError`` (the + original error surviving only as ``__context__``), so the outage check walks the cause chain, and a + missing user falls through to ``no_active_key`` rather than an opaque gateway fault. ``source="database"`` reads the row from the database, never the cache, so the credential mint refuses a user that a writer deactivated or deleted without evicting the cached row, and it leaves the fresh row in the cache for the requests the credential makes next. Every other caller keeps the cache read, diff --git a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py index 80868296b50..a437df17e6a 100644 --- a/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/idp_token_exchange.py @@ -196,9 +196,8 @@ def _check_unavailable_description(outage: GatewayOutage) -> str: def _gateway_could_not_verify(denied: Exception) -> GatewayOutage | None: - """A database fault anywhere in the chain (``get_user_object`` wraps prisma failures in a - bare ``ValueError``) or a 5xx from JWT auth (the IdP's JWKS unreachable with no cached - copy) is the gateway failing, not the token. A fault retrying cannot clear (a missing or + """A database fault anywhere in the chain or a 5xx from JWT auth (the IdP's JWKS + unreachable with no cached copy) is the gateway failing, not the token. A fault retrying cannot clear (a missing or version-skewed query engine) is named as such, the way the mint path words it, so the client is not told to wait on a deployment that needs repair.""" fault: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(denied) diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 17c85bbdbca..5be87a8bf4d 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -203,6 +203,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( "/cursor/", "/deepgram/", "/eu.assemblyai/", + "/fal_ai/", "/gemini/", "/gigachat/", "/milvus/", @@ -210,8 +211,10 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( "/nvidia_nim/", "/openai/", "/openai_passthrough/", + "/tinyfish/", "/transcribe", "/typesafe/", + "/openrouter/", "/vertex-ai/", "/vertex_ai/", "/vllm/", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 6f9a2d8c96d..03122f25870 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -2357,6 +2357,20 @@ }, "AgentConfig": { "properties": { + "access_group_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Group Ids" + }, "agent_card_params": { "$ref": "#/components/schemas/AgentCard" }, @@ -2683,6 +2697,20 @@ }, "AgentResponse": { "properties": { + "access_group_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Group Ids" + }, "agent_card_params": { "additionalProperties": true, "title": "Agent Card Params", @@ -3506,6 +3534,20 @@ }, "PatchAgentRequest": { "properties": { + "access_group_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Group Ids" + }, "agent_card_params": { "$ref": "#/components/schemas/AgentCard" }, @@ -15382,14 +15424,31 @@ "title": "Jwt Issuer" }, "key": { - "title": "Key", - "type": "string" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key" + }, + "token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token" } }, "required": [ "jwt_claim_name", - "jwt_claim_value", - "key" + "jwt_claim_value" ], "title": "CreateJWTKeyMappingRequest", "type": "object" @@ -15553,6 +15612,17 @@ } ], "title": "Key" + }, + "token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token" } }, "required": [ @@ -18448,6 +18518,223 @@ ] } }, + "/fal_ai/{endpoint}": { + "delete": { + "operationId": "fal_ai_proxy_route_fal_ai__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Fal Ai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "operationId": "fal_ai_proxy_route_fal_ai__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Fal Ai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "operationId": "fal_ai_proxy_route_fal_ai__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Fal Ai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "operationId": "fal_ai_proxy_route_fal_ai__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Fal Ai Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "operationId": "fal_ai_proxy_route_fal_ai__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Fal Ai Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, "/gemini/{endpoint}": { "delete": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio)", @@ -20659,6 +20946,313 @@ ] } }, + "/openrouter/{endpoint}": { + "delete": { + "operationId": "openrouter_proxy_route_openrouter__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openrouter Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "operationId": "openrouter_proxy_route_openrouter__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openrouter Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "operationId": "openrouter_proxy_route_openrouter__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openrouter Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "operationId": "openrouter_proxy_route_openrouter__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openrouter Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "operationId": "openrouter_proxy_route_openrouter__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Openrouter Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, + "/tinyfish/{endpoint}": { + "get": { + "description": "Pass-through for the TinyFish Agent API (goal-based web automation).\n\nForwarded endpoints:\n- POST /v1/automation/run \u2014 run to completion (blocking)\n- POST /v1/automation/run-async \u2014 submit a run, poll GET /v1/runs/{id} for the result\n- POST /v1/automation/run-sse \u2014 run with SSE progress events\n- GET /v1/runs/{id} \u2014 run status / result\n- POST /v1/runs/{id}/cancel \u2014 cancel a run\n\nEvery other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs\nlisting, which would let any caller discover other callers' run ids) returns 403: all\nproxy callers share one upstream key.\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. TINYFISH_API_KEY environment variable\n\n[Docs](https://docs.litellm.ai/docs/pass_through/tinyfish)", + "operationId": "tinyfish_proxy_route_tinyfish__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Tinyfish Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Pass-through for the TinyFish Agent API (goal-based web automation).\n\nForwarded endpoints:\n- POST /v1/automation/run \u2014 run to completion (blocking)\n- POST /v1/automation/run-async \u2014 submit a run, poll GET /v1/runs/{id} for the result\n- POST /v1/automation/run-sse \u2014 run with SSE progress events\n- GET /v1/runs/{id} \u2014 run status / result\n- POST /v1/runs/{id}/cancel \u2014 cancel a run\n\nEvery other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs\nlisting, which would let any caller discover other callers' run ids) returns 403: all\nproxy callers share one upstream key.\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. TINYFISH_API_KEY environment variable\n\n[Docs](https://docs.litellm.ai/docs/pass_through/tinyfish)", + "operationId": "tinyfish_proxy_route_tinyfish__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Tinyfish Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, "/transcribe": { "post": { "description": "AWS-SDK-shaped pass-through for Amazon Transcribe: point the SDK's `endpoint_url`\nat `/transcribe` and the operation is read from the `X-Amz-Target` header, per the\nAWS JSON 1.1 protocol.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/transcribe)", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 03caca399c3..8b2c81fea77 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -27,6 +27,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_langfuse_span_scope_value, validate_no_callback_env_reference, ) +from litellm.types.agents import AgentCaller from litellm.types.integrations.compression_interception import ( CompressionSavingsMetadata, ) @@ -490,14 +491,17 @@ class LiteLLMRoutes(enum.Enum): "/openai_passthrough", "/assemblyai", "/eu.assemblyai", + "/tinyfish", "/vllm", "/mistral", "/typesafe", + "/openrouter", "/milvus", "/gigachat", "/watsonx", "/nvidia_nim", "/deepgram", + "/fal_ai", ] ######################################################### @@ -2647,6 +2651,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="opt-in to RFC 8628 verification_uri_complete for the CLI SSO device flow, pre-filling the user_code in the browser. Off by default; intended for same-host clients where the device that starts the flow and the browser run on the same machine", ) + include_call_id_in_error_body: bool | None = Field( + None, + description="opt-in to copy the x-litellm-call-id response header's value into JSON error bodies, as error.litellm_call_id on the OpenAI-shaped and /v1/messages routes and as a top-level litellm_call_id on pass-through routes, so an error a client prints names the request to look up. Off by default", + ) enable_claude_code_gateway: bool | None = Field( None, description="serve the Claude Code gateway protocol (https://code.claude.com/docs/en/claude-apps-gateway) under /claude_code_gateway: OAuth device-flow sign-in reusing proxy SSO, plus managed settings and OTLP telemetry ingestion. Off by default", @@ -3274,6 +3282,15 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob "user id." ), ) + agent_caller: AgentCaller | None = Field( + default=None, + exclude=True, + description=( + "Set per request from the x-litellm-user-id / x-litellm-team-id headers an agent echoes back on " + "calls made with its own key. Every check treats it as a ceiling, so a forged value can only " + "narrow the agent's access." + ), + ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True) user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) @@ -3306,6 +3323,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob values.pop("mcp_session_resource_server_id", None) values.pop("mcp_toolset_id", None) values.pop("via_virtual_key", None) + values.pop("agent_caller", None) if values.get("api_key") is not None: values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) if isinstance(values.get("api_key"), str): @@ -4261,6 +4279,11 @@ class ProxyErrorTypes(str, enum.Enum): Project does not have access to the model """ + agent_model_access_denied = "agent_model_access_denied" + """ + The agent behind the key does not have access to the model + """ + model_cost_map_missing = "model_cost_map_missing" expired_key = "expired_key" @@ -4335,7 +4358,7 @@ class ProxyErrorTypes(str, enum.Enum): @classmethod def get_model_access_error_type_for_object( - cls, object_type: Literal["key", "user", "team", "org", "project"] + cls, object_type: Literal["key", "user", "team", "org", "project", "agent"] ) -> "ProxyErrorTypes": """ Get the model access error type for object_type @@ -4350,6 +4373,8 @@ class ProxyErrorTypes(str, enum.Enum): return cls.org_model_access_denied elif object_type == "project": return cls.project_model_access_denied + elif object_type == "agent": + return cls.agent_model_access_denied @classmethod def get_vector_store_access_error_type_for_object( @@ -4715,7 +4740,8 @@ class KeyHealthResponse(TypedDict, total=False): class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): jwt_claim_name: str jwt_claim_value: str - key: str + key: str | None = None + token: str | None = None jwt_issuer: str | None = None description: str | None = None @@ -4723,6 +4749,7 @@ class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): id: str key: str | None = None + token: str | None = None jwt_issuer: str | None = None description: str | None = None is_active: bool | None = None diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 834c16ba6dc..2a189a76545 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -146,12 +146,17 @@ def _validate_push_notification_url(url: str) -> None: def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, str]: + """The human behind this call. An agent key acting for an invoking user forwards that user, not + itself, so a chain of agents stays capped at what the original caller may reach.""" + caller: Final = user_api_key_dict.agent_caller + user_id: Final = caller.user_id if caller is not None else user_api_key_dict.user_id + team_id: Final = caller.team_id if caller is not None else user_api_key_dict.team_id return MappingProxyType( { name: value for name, value in ( - ("X-LiteLLM-User-Id", user_api_key_dict.user_id), - ("X-LiteLLM-Team-Id", user_api_key_dict.team_id), + ("X-LiteLLM-User-Id", user_id), + ("X-LiteLLM-Team-Id", team_id), ) if value } diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index c7b6bca72cf..d6b12e830e1 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -7,6 +7,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypedDict from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly import litellm from litellm.constants import REDACTED_BY_LITELM_STRING @@ -37,6 +38,7 @@ class AgentRecordDump(TypedDict): agent_card_params: dict[str, object] static_headers: dict[str, str] | None extra_headers: list[str] | None + access_group_ids: ReadOnly[Sequence[str] | None] object_permission: dict[str, object] | None spend: float tpm_limit: int | None @@ -65,6 +67,9 @@ class AgentRecord(Protocol): @property def object_permission(self) -> AgentObjectPermissionRecord | None: ... + @property + def access_group_ids(self) -> Sequence[str] | None: ... + @property def spend(self) -> float: ... @@ -284,6 +289,12 @@ def _resolved_agent_param_value( return _MISSING_AGENT_PARAM +def _patched_access_group_ids(agent: PatchAgentRequest) -> Mapping[str, object]: + if "access_group_ids" not in agent: + return MappingProxyType({}) + return MappingProxyType({"access_group_ids": tuple(dict.fromkeys(agent.get("access_group_ids") or ()))}) + + def _restore_redacted_litellm_params( incoming: Mapping[str, object], existing: Mapping[str, object], @@ -516,6 +527,7 @@ class AgentRegistry: static_headers_val: Final[str | None] = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None extra_headers_val: Final = agent.get("extra_headers") + access_group_ids_val: Final = agent.get("access_group_ids") create_data: Final[dict[str, object]] = { "agent_name": agent_name, @@ -532,6 +544,8 @@ class AgentRegistry: create_data["static_headers"] = static_headers_val if extra_headers_val is not None: create_data["extra_headers"] = extra_headers_val + if access_group_ids_val is not None: + create_data["access_group_ids"] = tuple(dict.fromkeys(access_group_ids_val)) if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id @@ -601,7 +615,7 @@ class AgentRegistry: existing_agent: Final[Mapping[str, object]] = dict(existing_record) augment_agent: Final = {**existing_agent, **agent} - update_data: Final[dict[str, object]] = {} + update_data: Final[dict[str, object]] = {**_patched_access_group_ids(agent)} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if "litellm_params" in agent: @@ -703,6 +717,7 @@ class AgentRegistry: safe_dumps(dict(static_headers_obj_u)) if static_headers_obj_u is not None else safe_dumps({}) ) extra_headers_val_u: Final = agent.get("extra_headers") or [] + access_group_ids_val_u: Final = tuple(dict.fromkeys(agent.get("access_group_ids") or ())) update_data: Final[dict[str, object]] = { "agent_name": agent_name, @@ -710,6 +725,7 @@ class AgentRegistry: "agent_card_params": agent_card_params, "static_headers": static_headers_val_u, "extra_headers": extra_headers_val_u, + "access_group_ids": access_group_ids_val_u, "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), } diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py new file mode 100644 index 00000000000..49e5407ff88 --- /dev/null +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -0,0 +1,75 @@ +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Final, TypeAlias + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AccessGroupTable + +AccessGroupIds: TypeAlias = tuple[str, ...] +AccessGroupIdsLoader: TypeAlias = Callable[[str], Awaitable[AccessGroupIds]] # mutable-ok: Callable params +LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None +AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LoadedAccessGroup]] # mutable-ok: Callable parameter syntax + + +@dataclass(frozen=True, slots=True) +class AgentAccessGroupCeiling: + """Everything the agent's attached access groups allow. An empty set denies that resource kind.""" + + access_group_ids: AccessGroupIds + models: frozenset[str] + mcp_server_ids: frozenset[str] + agent_ids: frozenset[str] + + +CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] # mutable-ok: Callable params + + +async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + agent: Final = await get_agent_with_read_through(agent_id) + return tuple(agent.access_group_ids or ()) if agent is not None else () + + +async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: + from litellm.proxy.auth.auth_checks import get_access_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) + return None + try: + return await get_access_object( + access_group_id=access_group_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException as e: + verbose_proxy_logger.warning( + "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail + ) + return None + + +async def resolve_agent_access_group_ceiling( + agent_id: str, + load_access_group_ids: AccessGroupIdsLoader = _registry_access_group_ids, + load_access_group: AccessGroupLoader = _load_access_group, +) -> AgentAccessGroupCeiling | None: + """``None`` when the agent has no access groups attached, so nothing is capped.""" + access_group_ids: Final = await load_access_group_ids(agent_id) + if not access_group_ids: + return None + + loaded: Final = await asyncio.gather(*(load_access_group(group_id) for group_id in access_group_ids)) + groups: Final = tuple(group for group in loaded if group is not None) + return AgentAccessGroupCeiling( + access_group_ids=access_group_ids, + models=frozenset(model for group in groups for model in group.access_model_names), + mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids), + agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids), + ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_caller.py b/litellm/proxy/agent_endpoints/auth/agent_caller.py new file mode 100644 index 00000000000..47d43e8f71b --- /dev/null +++ b/litellm/proxy/agent_endpoints/auth/agent_caller.py @@ -0,0 +1,87 @@ +"""The human behind an agent's own proxy calls. + +``/a2a/{agent}`` forwards the invoking key's ``X-LiteLLM-User-Id`` / ``X-LiteLLM-Team-Id`` to the +agent backend. When the agent echoes them back on requests made with its own key, the proxy caps +that key at what the invoking user and team may reach. The cap is intersected with, never +substituted for, the agent key's own grants and the agent's access group ceiling, so the headers +can only narrow access and need no trust. +""" + +from collections.abc import Mapping +from typing import Final + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.types.agents import ( + AGENT_CALLER_TEAM_ID_HEADER, + AGENT_CALLER_USER_ID_HEADER, + AgentCaller, +) + + +def _header(headers: Mapping[str, str], name: str) -> str | None: + value: Final = next((raw for key, raw in headers.items() if key.lower() == name), None) + return value.strip() or None if value is not None else None + + +def agent_caller_from_headers(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth) -> AgentCaller | None: + """The caller an agent key is acting for, or ``None`` when the key is not an agent's or no id was echoed.""" + if not user_api_key_auth.agent_id: + return None + user_id: Final = _header(headers, AGENT_CALLER_USER_ID_HEADER) + team_id: Final = _header(headers, AGENT_CALLER_TEAM_ID_HEADER) + if user_id is None and team_id is None: + return None + return AgentCaller(user_id=user_id, team_id=team_id) + + +def agent_caller_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | None: + """A minimal auth context standing for the invoking user and team, so the shared key/team/user + resolvers can be reused unchanged to compute what the caller may reach.""" + caller: Final = user_api_key_auth.agent_caller + if caller is None: + return None + return UserAPIKeyAuth( + user_id=caller.user_id, + team_id=caller.team_id, + parent_otel_span=user_api_key_auth.parent_otel_span, + ) + + +async def load_agent_caller_team(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + """The invoking team's row, or ``None`` when no team id was echoed. Raises when the id names a team + that cannot be loaded, since a caller we cannot resolve must not be treated as unrestricted.""" + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + caller: Final = user_api_key_auth.agent_caller + if caller is None or caller.team_id is None: + return None + return await get_team_object( + team_id=caller.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def load_agent_caller_user(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_UserTable | None: + """The invoking user's row, or ``None`` when no user id was echoed or the row does not exist.""" + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + caller: Final = user_api_key_auth.agent_caller + if caller is None or caller.user_id is None: + return None + user_object: Final = await get_user_object( + user_id=caller.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if user_object is None: + verbose_proxy_logger.debug("agent caller user %r not found; no user ceiling applied", caller.user_id) + return user_object diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index e4dd77e2f82..9fe74bfee3f 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -19,6 +19,11 @@ from litellm.proxy._types import ( LitellmUserRoles, UserAPIKeyAuth, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + CeilingResolver, + resolve_agent_access_group_ceiling, +) +from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentResponse @@ -44,6 +49,22 @@ def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]: return frozenset(global_agent_registry.stable_agent_id(agent_id) for agent_id in agent_ids) +def _restricted_ids(access: AgentAccess) -> frozenset[str] | None: + if isinstance(access, UnrestrictedAgentAccess): + return None + return _to_stable_ids(access.agent_ids) + + +def _intersect_agent_access(key_access: AgentAccess, team_access: AgentAccess) -> AgentAccess: + key_ids: Final = _restricted_ids(key_access) + team_ids: Final = _restricted_ids(team_access) + if key_ids is None: + return UnrestrictedAgentAccess() if team_ids is None else RestrictedAgentAccess(team_ids) + if team_ids is None: + return RestrictedAgentAccess(key_ids) + return RestrictedAgentAccess(key_ids & team_ids) + + class AgentRequestHandler: """ Class to handle agent permission checking, including: @@ -61,35 +82,56 @@ class AgentRequestHandler: @staticmethod async def resolve_agent_access( user_api_key_auth: UserAPIKeyAuth | None = None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> AgentAccess: - """ - Resolve the agents the given user/key may reach. + """Agents the key may reach: key and team grants, intersected with the agent's access group ceiling + and, for an agent key acting on behalf of an invoking user, with that user's team grants.""" + key_team_access: Final = await AgentRequestHandler._resolve_key_team_agent_access(user_api_key_auth) + caller_access: Final = await AgentRequestHandler._agent_caller_access(user_api_key_auth) + own_access: Final = _intersect_agent_access(key_team_access, caller_access) + agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth, resolve_ceiling) + if agent_ceiling is None: + return own_access + if isinstance(own_access, UnrestrictedAgentAccess): + return RestrictedAgentAccess(agent_ceiling) + return RestrictedAgentAccess(own_access.agent_ids & agent_ceiling) - ``UnrestrictedAgentAccess`` is only returned when neither the key nor its team - carries any grant. Grants that intersect to nothing stay restricted, so - narrowing a caller can never widen what it reaches. - """ + @staticmethod + async def _agent_caller_access(user_api_key_auth: UserAPIKeyAuth | None) -> AgentAccess: + caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None + if caller_auth is None: + return UnrestrictedAgentAccess() + return await AgentRequestHandler._get_allowed_agents_for_team(caller_auth) + + @staticmethod + async def _resolve_key_team_agent_access( + user_api_key_auth: UserAPIKeyAuth | None, + ) -> AgentAccess: try: key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) - - match (key_access, team_access): - case (UnrestrictedAgentAccess(), UnrestrictedAgentAccess()): - return UnrestrictedAgentAccess() - case (UnrestrictedAgentAccess(), RestrictedAgentAccess(team_ids)): - return RestrictedAgentAccess(_to_stable_ids(team_ids)) - case (RestrictedAgentAccess(key_ids), UnrestrictedAgentAccess()): - return RestrictedAgentAccess(_to_stable_ids(key_ids)) - case (RestrictedAgentAccess(key_ids), RestrictedAgentAccess(team_ids)): - return RestrictedAgentAccess(_to_stable_ids(key_ids) & _to_stable_ids(team_ids)) except Exception as e: verbose_logger.warning("Failed to get allowed agents: %s", e) return UnrestrictedAgentAccess() + return _intersect_agent_access(key_access, team_access) + + @staticmethod + async def _agent_access_group_ceiling( + user_api_key_auth: UserAPIKeyAuth | None, + resolve_ceiling: CeilingResolver, + ) -> frozenset[str] | None: + if user_api_key_auth is None or not user_api_key_auth.agent_id: + return None + ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id) + if ceiling is None: + return None + return _to_stable_ids(ceiling.agent_ids) @staticmethod async def is_agent_allowed( agent_id: str, user_api_key_auth: UserAPIKeyAuth | None = None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> bool: """ Check if a specific agent is allowed for the given user/key. @@ -103,7 +145,7 @@ class AgentRequestHandler: """ from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - match await AgentRequestHandler.resolve_agent_access(user_api_key_auth): + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth, resolve_ceiling): case UnrestrictedAgentAccess(): return True case RestrictedAgentAccess(allowed_agent_ids): diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 644778bcb9f..d9558b86e95 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -8,7 +8,11 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse import litellm -from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping +from litellm.anthropic_interface.exceptions import ( + AnthropicErrorDetail, + AnthropicErrorResponse, + AnthropicExceptionMapping, +) from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.anthropic.experimental_pass_through.context_management import ( AnthropicContextManagementError, @@ -25,8 +29,10 @@ from litellm.proxy.common_request_processing import ( proxy_exception_from_http_exception, resolve_litellm_call_id, ) +from litellm.proxy.common_utils.error_body_call_id import error_body_call_id from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, error_status_code, openai_error_param, openai_error_type, @@ -37,9 +43,29 @@ from litellm.types.utils import TokenCountResponse router: Final = APIRouter() +def _with_provider_specific_fields(exc: ProxyException, detail: AnthropicErrorDetail) -> AnthropicErrorDetail: + if not exc.provider_specific_fields: + return detail + with_fields: Final[AnthropicErrorDetail] = {**detail, "provider_specific_fields": exc.provider_specific_fields} + return with_fields + + +def _anthropic_error_detail( + exc: ProxyException, detail: AnthropicErrorDetail, call_id: str | None +) -> AnthropicErrorDetail: + if call_id is None: + return _with_provider_specific_fields(exc, detail) + with_call_id: Final[AnthropicErrorDetail] = { + **_with_provider_specific_fields(exc, detail), + "litellm_call_id": call_id, + } + return with_call_id + + def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSONResponse: from litellm.proxy.proxy_server import ( _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does + general_settings_view, ) status_code: Final = int(exc.code) if exc.code is not None and exc.code.isdigit() else 500 @@ -49,11 +75,10 @@ def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSO raw_message=exc.message, request_id=request.headers.get("x-request-id"), ) - if not exc.provider_specific_fields: - return JSONResponse(status_code=status_code, content=envelope, headers=exc.headers) + body_call_id: Final = error_body_call_id(general_settings_view(), exc.headers.get(LITELLM_CALL_ID_HEADER)) content: Final[AnthropicErrorResponse] = { **envelope, - "error": {**envelope["error"], "provider_specific_fields": exc.provider_specific_fields}, + "error": _anthropic_error_detail(exc, envelope["error"], body_call_id), } return JSONResponse(status_code=status_code, content=content, headers=exc.headers) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c2279fb2fe1..f0588b3fa79 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,7 +15,7 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias from fastapi import HTTPException, Request, status from pydantic import BaseModel, TypeAdapter @@ -68,6 +68,15 @@ from litellm.proxy._types import ( SpecialModelNames, UserAPIKeyAuth, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + CeilingResolver, + resolve_agent_access_group_ceiling, +) +from litellm.proxy.agent_endpoints.auth.agent_caller import ( + agent_caller_auth, + load_agent_caller_team, + load_agent_caller_user, +) from litellm.proxy.auth.budget_throttle import ( budget_throttle_percentage, should_throttle_budget_exceeded, @@ -1006,6 +1015,16 @@ async def common_checks( code=status.HTTP_400_BAD_REQUEST, ) + await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router) + await _check_agent_caller_model_access( + model=_model, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ## 2.1 If user can call model (if personal key) if _model and team_object is None and user_object is not None: with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"): @@ -2678,9 +2697,15 @@ async def get_user_object( raise except Exception as e: _log_budget_lookup_failure("user", e) - raise ValueError( - f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call. Got error - {e}" - ) + raise _user_read_failure(user_id=user_id, error=e) + + +def _user_read_failure(user_id: str, error: Exception) -> Exception: + if PrismaDBExceptionHandler.is_database_service_unavailable_error(error): + return error + return ValueError( + f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call. Got error - {error}" + ) async def _cache_management_object( @@ -4251,7 +4276,7 @@ def _can_object_call_model( models: list[str], team_model_aliases: dict[str, str] | None = None, team_id: str | None = None, - object_type: Literal["user", "team", "key", "org", "project"] = "user", + object_type: Literal["user", "team", "key", "org", "project", "agent"] = "user", fallback_depth: int = 0, ) -> Literal[True]: """ @@ -4317,6 +4342,82 @@ def _can_object_call_model( ) +async def _check_agent_access_group_model_access( + model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] + valid_token: UserAPIKeyAuth | None, + llm_router: Router | None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, +) -> Literal[True]: + """Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows.""" + if not model or valid_token is None or not valid_token.agent_id: + return True + ceiling: Final = await resolve_ceiling(valid_token.agent_id) + if ceiling is None: + return True + if not ceiling.models: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=model), + internal_message=f"agent {valid_token.agent_id} access groups {ceiling.access_group_ids} grant no models", + type=ProxyErrorTypes.agent_model_access_denied, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + return _can_object_call_model( + model=model, + llm_router=llm_router, + models=sorted(ceiling.models), + team_id=valid_token.team_id, + object_type="agent", + ) + + +LoadedCallerTeam: TypeAlias = LiteLLM_TeamTable | None +LoadedCallerUser: TypeAlias = LiteLLM_UserTable | None +CallerTeamLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerTeam]] # mutable-ok: Callable params +CallerUserLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerUser]] # mutable-ok: Callable params + + +async def _check_agent_caller_model_access( + model: str | list[str] | None, # mutable-ok: the model checks it delegates to take list[str] + valid_token: UserAPIKeyAuth | None, + llm_router: Router | None, + prisma_client: Optional["PrismaClient"], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> None: + """An agent key acting for an invoking user may call only what that user's own key could: the + invoking team's models (and per-member scope) when a team was echoed, else the user's models.""" + if not model or valid_token is None: + return + caller_auth: Final = agent_caller_auth(valid_token) + if caller_auth is None: + return + caller_team: Final = await load_team(valid_token) + if caller_team is not None: + await can_team_access_model( + model=model, + team_object=caller_team, + llm_router=llm_router, + prisma_client=prisma_client, + ) + await _check_team_member_model_access( + model=model, + team_object=caller_team, + valid_token=caller_auth, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + return + caller_user: Final = await load_user(valid_token) + if caller_user is None: + return + await can_user_call_model(model=model, llm_router=llm_router, user_object=caller_user) + + def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool: """ Returns True if `model` being accessed is an alias of a team model diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 010a1b4536e..07d8d00d202 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -2099,8 +2099,11 @@ class JWTAuthManager: spend / metadata can be attributed correctly. Returns (team_id, team_object, team_membership_object). - Any DB error is debug-logged and the tuple is (None, None, None) — no - exception ever propagates from this helper. + A team that cannot be loaded (HTTPException from get_team_object) is + debug-logged and the tuple is (None, None, None), the same as the DB + team fallback. A failed membership read propagates, so a database + outage surfaces as the 503 the rest of auth answers with instead of + serving the request with the member's limits dropped. """ if user_object is None or not user_object.teams or len(user_object.teams) != 1: return None, None, None @@ -2115,28 +2118,28 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=team_id_upsert, ) - if team_row is None: - return None, None, None - - if not user_id: - return _tid, team_row, None - - team_membership: Final = await get_team_membership( - user_id=user_id, - team_id=_tid, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - return _tid, team_row, team_membership - except Exception: + except HTTPException: verbose_proxy_logger.debug( - "JWT single-team fallback error, skipping. team_id=%s", + "JWT single-team fallback: team could not be loaded, skipping. team_id=%s", _tid, exc_info=True, ) return None, None, None + if team_row is None: + return None, None, None + + if not user_id: + return _tid, team_row, None + + team_membership: Final = await get_team_membership( + user_id=user_id, + team_id=_tid, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return _tid, team_row, team_membership @staticmethod async def _resolve_db_team_fallback( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4371ce4fda8..a6c0792a86f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -39,6 +39,7 @@ from litellm.integrations.otel.runtime import phase_span, seed_request_identity from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * +from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, TeamNotFoundError, @@ -649,6 +650,7 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str "type": "http", "headers": scope_headers, "path": ws_scope.get("path", ""), + "state": ws_scope.setdefault("state", {}), # mutable-ok: Starlette's socket state, shared with the request } for key in ("root_path", "app_root_path"): if key in ws_scope: @@ -3085,31 +3087,30 @@ async def _reserve_budget_after_common_checks( request: Request | None = None, ) -> None: user_api_key_auth_obj.budget_reservation = None - if skip_budget_checks: - return - if general_settings.get("disable_budget_reservation") is True: - return + if not skip_budget_checks and general_settings.get("disable_budget_reservation") is not True: + from litellm.proxy.spend_tracking.budget_reservation import ( + reserve_budget_for_request, + ) - from litellm.proxy.spend_tracking.budget_reservation import ( - reserve_budget_for_request, - ) - - user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request( - request_body=request_data, - route=route, - llm_router=llm_router, - valid_token=user_api_key_auth_obj, - team_object=team_object, - user_object=user_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - end_user_id=end_user_id, - end_user_object=end_user_object, - apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, - fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, - raw_body=await read_raw_json_body(request=request), - ) + user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request( + request_body=request_data, + route=route, + llm_router=llm_router, + valid_token=user_api_key_auth_obj, + team_object=team_object, + user_object=user_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + end_user_id=end_user_id, + end_user_object=end_user_object, + apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, + fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, + raw_body=await read_raw_json_body(request=request), + ) + if request is not None: + reservation: Final = user_api_key_auth_obj.budget_reservation + request.state.budget_reservation = reservation # rebind-ok: read by the release middleware def _should_skip_budget_checks( @@ -3330,6 +3331,9 @@ async def user_api_key_auth( raise body_parse_exception raise user_api_key_auth_obj.budget_reservation = None + user_api_key_auth_obj.agent_caller = agent_caller_from_headers( + _safe_get_request_headers(request), user_api_key_auth_obj + ) _seed_request_destinations(user_api_key_auth_obj, request) # A body that never parsed is authenticated (so the trace carries identity diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index f4bebc4a4cb..b3fdc4695cb 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -22,8 +22,12 @@ from pathlib import Path from types import MappingProxyType from typing import Final, TypeAlias +import click +from filelock import FileLock +from packaging.version import InvalidVersion, Version from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError +from litellm._version import version as litellm_version from litellm.litellm_core_utils.private_json import ( commit_staged_json, discard_staged_json, @@ -75,6 +79,7 @@ BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json" CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json" STATUSLINE_SCRIPT_PATH: Final = Path.home() / ".litellm" / "statusline.py" +STATUSLINE_VERSION_PREFIX: Final = b"# litellm-statusline-version: " @dataclass(frozen=True, slots=True) @@ -305,11 +310,54 @@ def statusline_command(script_path: Path, platform: str = sys.platform) -> str: return " ".join(quote(token) for token in (sys.executable, str(script_path))) -def install_statusline_script(script_path: Path | None = None) -> str: +def _statusline_version(value: str) -> Version | None: + try: + return Version(value) + except InvalidVersion: + return None + + +def _installed_statusline_version(target: Path) -> Version | None: + try: + with target.open("rb") as script: + header: Final = script.readline(256) + except FileNotFoundError: + return None + if not header.startswith(STATUSLINE_VERSION_PREFIX): + return None + try: + return _statusline_version(header.removeprefix(STATUSLINE_VERSION_PREFIX).decode("ascii").strip()) + except UnicodeDecodeError: + return None + + +def install_statusline_script( + script_path: Path | None = None, + *, + package_version: str = litellm_version, + write: Callable[[str, bytes], None] = write_private_bytes, +) -> str: target: Final = script_path or STATUSLINE_SCRIPT_PATH try: ensure_private_dir(target.parent) - write_private_bytes(str(target), Path(statusline_script.__file__).read_bytes()) + bundled_version: Final = _statusline_version(package_version) + with FileLock(str(target) + ".lock", timeout=10, mode=0o600): + installed_version: Final = _installed_statusline_version(target) + if installed_version is not None and (bundled_version is None or installed_version > bundled_version): + cli_version: Final = str(bundled_version) if bundled_version is not None else "unknown" + click.echo( + f"Keeping the status line from LiteLLM {installed_version}; this CLI is {cli_version}. " + "Upgrade the CLI to refresh it.", + err=True, + ) + return statusline_command(target) + source: Final = Path(statusline_script.__file__).read_bytes() + header: Final = ( + STATUSLINE_VERSION_PREFIX + str(bundled_version).encode("ascii") + b"\n" + if bundled_version is not None + else b"" + ) + write(str(target), header + source) except OSError as e: raise ClaudeSettingsError(f"Could not install the status line script at {target}: {e}") from e return statusline_command(target) diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 09dd062c888..d16160b1ab8 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -1,7 +1,7 @@ """Claude Code status line and Codex Stop hook for auto-routed sessions. -`lite` copies this file verbatim to ~/.litellm/statusline.py and registers it as Claude -Code's `statusLine` command and as Codex's `[[hooks.Stop]]` command, so it must stay +`lite` copies this file to ~/.litellm/statusline.py with a CLI version header when known and registers +it as Claude Code's `statusLine` command and as Codex's `[[hooks.Stop]]` command, so it must stay standard-library only and must never import litellm. Claude Code re-runs it on every status refresh (about every 300ms while typing), so the proxy is asked at most once per TTL per session and every other refresh is served from a small on-disk cache that holds diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9484fd7c723..f6e0d56127f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -75,11 +75,13 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, ) +from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( get_client_requested_model, get_tags_from_request_body, ) from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, attribute_of, error_status_code, openai_error_param, @@ -946,6 +948,9 @@ async def _resolve_stream_headers( return headers +_NO_GENERAL_SETTINGS: Final[Mapping[str, object]] = MappingProxyType({}) + + async def create_response( generator: AsyncGenerator[str, None], media_type: str, @@ -953,6 +958,7 @@ async def create_response( default_status_code: int = status.HTTP_200_OK, request: Request | None = None, refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None = None, + general_settings: Mapping[str, object] = _NO_GENERAL_SETTINGS, ) -> StreamingResponse | JSONResponse: """ Create streaming response, checking if the first chunk is an error. @@ -960,7 +966,8 @@ async def create_response( Otherwise, return StreamingResponse and stream all content. ``refresh_headers`` is consulted once the first chunk has been buffered, for - callers whose headers can only be known then. + callers whose headers can only be known then. ``general_settings`` decides whether + the first-chunk error body also carries the ``x-litellm-call-id`` header's value. """ first_chunk_value: str | None = None final_status_code = default_status_code @@ -987,7 +994,10 @@ async def create_response( ) # Parse error content - error_dict: Final = _extract_error_from_sse_chunk(first_chunk_value) + error_dict: Final = with_call_id( + JSON_OBJECT.validate_python(_extract_error_from_sse_chunk(first_chunk_value)), + error_body_call_id(general_settings, resolved_headers.get(LITELLM_CALL_ID_HEADER)), + ) # Consume and close generator (avoid resource leak) try: @@ -2738,6 +2748,7 @@ class ProxyBaseLLMRequestProcessing: headers=custom_headers, request=request, refresh_headers=refresh_stream_headers, + general_settings=general_settings, ) ### CALL HOOKS ### - modify outgoing data diff --git a/litellm/proxy/common_utils/error_body_call_id.py b/litellm/proxy/common_utils/error_body_call_id.py new file mode 100644 index 00000000000..f50be5df509 --- /dev/null +++ b/litellm/proxy/common_utils/error_body_call_id.py @@ -0,0 +1,20 @@ +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter + +INCLUDE_CALL_ID_IN_ERROR_BODY_SETTING: Final = "include_call_id_in_error_body" +LITELLM_CALL_ID_BODY_KEY: Final = "litellm_call_id" +JSON_OBJECT: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) # mutable-ok: JSONResponse input + + +def error_body_call_id(general_settings: Mapping[str, object], call_id: str | None) -> str | None: + if general_settings.get(INCLUDE_CALL_ID_IN_ERROR_BODY_SETTING) is not True: + return None + return call_id if call_id else None + + +def with_call_id(error: dict[str, object], call_id: str | None) -> dict[str, object]: # mutable-ok: JSONResponse input + if call_id is None: + return error + return {**error, LITELLM_CALL_ID_BODY_KEY: call_id} # mutable-ok: JSONResponse input diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index f25c2787252..99167c1275c 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -398,11 +398,8 @@ class PrismaDBExceptionHandler: ``is_database_service_unavailable_error`` classifies a single exception by type, which a caller that catches a raw DB failure and re-raises a - domain exception of a different type defeats. ``get_user_object`` in - ``litellm/proxy/auth/auth_checks.py`` is the concrete case: it wraps - every DB error, a genuine outage included, in a bare ``ValueError`` - whose original error survives only as ``__context__``. A type check on - the ``ValueError`` misses the outage, so the caller would mistake an + domain exception of a different type defeats. A type check on the + wrapper misses the outage, so the caller would mistake an infrastructure fault for an auth failure. Walking the chain recovers the real signal, which is the PEP 3134 way to inspect a wrapped cause. diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 70ea21320ee..fe91d6d7a28 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -11,7 +11,7 @@ import asyncio import json import threading -from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Sequence from contextlib import asynccontextmanager from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypedDict, cast @@ -39,6 +39,11 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.anthropic_sse import ( + anthropic_sse_chunks_from_response, + assemble_anthropic_sse_stream, + model_response_text, +) from litellm.types.guardrails import ( GuardrailEventHooks, LitellmParams, @@ -1327,30 +1332,44 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) return response - async def _stream_apply_output_masking( - self, - response: AsyncIterable[object], - request_data: dict, - ) -> AsyncGenerator[ModelResponseStream | bytes, None]: - """Apply Presidio masking to streaming output (apply_to_output=True path).""" + async def _mask_buffered_model_response_stream( + self, all_chunks: Sequence[ModelResponseStream], request_data: dict + ) -> tuple[object, ...]: from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, ) from litellm.main import stream_chunk_builder from litellm.types.utils import ModelResponse + assembled: Final = stream_chunk_builder(chunks=list(all_chunks), messages=request_data.get("messages")) + if not isinstance(assembled, ModelResponse): + return tuple(all_chunks) + await self._process_response_for_pii(response=assembled, request_data=request_data, mode="mask") + return (convert_model_response_to_streaming(assembled),) + + async def _stream_apply_output_masking( + self, + response: AsyncIterable[object], + request_data: dict, + ) -> AsyncGenerator[object, None]: + """Apply Presidio masking to streaming output (apply_to_output=True path).""" all_chunks: list[ModelResponseStream] = [] passthrough_due_to_unknown_stream_shape = False try: - async for chunk in response: + stream: Final = response.__aiter__() + async for chunk in stream: if isinstance(chunk, ModelResponseStream): if passthrough_due_to_unknown_stream_shape: yield chunk else: all_chunks.append(chunk) elif isinstance(chunk, bytes): - yield chunk - continue + if passthrough_due_to_unknown_stream_shape or all_chunks: + yield chunk + continue + for masked_chunk in await self._mask_anthropic_sse_stream(chunk, stream, request_data): + yield masked_chunk + return else: if all_chunks: # Flush buffered chunks and switch to transparent passthrough for this stream shape. @@ -1375,33 +1394,39 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): if not all_chunks: verbose_proxy_logger.warning( "Presidio apply_to_output: streaming response contained no " - "ModelResponseStream chunks (e.g. raw SSE bytes or an empty " - "upstream stream). Output PII masking was skipped for this " - "response." + "ModelResponseStream chunks (an empty upstream stream). " + "Output PII masking was skipped for this response." ) return - assembled_model_response = stream_chunk_builder(chunks=all_chunks, messages=request_data.get("messages")) - - if not isinstance(assembled_model_response, ModelResponse): - for chunk in all_chunks: - yield chunk - return - - await self._process_response_for_pii( - response=assembled_model_response, - request_data=request_data, - mode="mask", - ) - - mock_response_stream: Final = convert_model_response_to_streaming(assembled_model_response) - yield mock_response_stream + for masked_chunk in await self._mask_buffered_model_response_stream(all_chunks, request_data): + yield masked_chunk except Exception as e: + if not all_chunks or isinstance(e, BlockedPiiEntityError): + raise verbose_proxy_logger.error("Error masking streaming PII output: %s", e) for chunk in all_chunks: yield chunk + async def _mask_anthropic_sse_stream( + self, first_chunk: bytes, rest: AsyncIterator[object], request_data: dict + ) -> tuple[object, ...]: + rest_chunks: Final = [chunk async for chunk in rest] # mutable-ok: tuple() cannot consume an async iterator + chunks: Final = (first_chunk, *rest_chunks) + assembled: Final = assemble_anthropic_sse_stream(chunks, restore_identity=True) + if assembled is None: + verbose_proxy_logger.warning( + "Presidio apply_to_output: raw SSE stream could not be assembled into a response. " + "Output PII masking was skipped for this response." + ) + return chunks + original_text: Final = model_response_text(assembled) + await self._process_response_for_pii(response=assembled, request_data=request_data, mode="mask") + if model_response_text(assembled) == original_text: + return chunks + return anthropic_sse_chunks_from_response(assembled) + @staticmethod def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str]) -> bytes: try: @@ -1460,7 +1485,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self, response: AsyncIterable[object], request_data: dict, - ) -> AsyncGenerator[ModelResponseStream | bytes, None]: + ) -> AsyncGenerator[object, None]: """Apply PII unmasking to streaming output (output_parse_pii=True path).""" from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, @@ -1536,7 +1561,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, response: AsyncIterable[object], request_data: dict, - ) -> AsyncGenerator[ModelResponseStream | bytes, None]: + ) -> AsyncGenerator[object, None]: """ Process streaming response chunks to unmask PII tokens when needed. diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 356eb7c96c6..7bafad26569 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -165,6 +165,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail) -> apply_to_output=True, event_hook=GuardrailEventHooks.post_call.value, output_parse_pii=False, + mask_response_content=True, ) if run_output else None diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 08e8e4f8c10..81894a5ff12 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -11,6 +11,7 @@ from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, + budget_reservation_from_metadata, get_litellm_metadata_from_kwargs, ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -630,17 +631,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]: def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: - metadata_budget_reservation: Final = metadata.get("user_api_key_budget_reservation") - if isinstance(metadata_budget_reservation, dict): - return metadata_budget_reservation - - user_api_key_auth_obj: Final = metadata.get("user_api_key_auth") - if user_api_key_auth_obj is None: - return None - if isinstance(user_api_key_auth_obj, dict): - budget_reservation: Final = user_api_key_auth_obj.get("budget_reservation") - return budget_reservation if isinstance(budget_reservation, dict) else None - return getattr(user_api_key_auth_obj, "budget_reservation", None) + return budget_reservation_from_metadata(metadata) def _get_request_tags_for_cost_tracking( diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index a6cc5140b15..b4923b0a2dc 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -5,6 +5,7 @@ from types import MappingProxyType from typing import Final, Protocol from fastapi import APIRouter, Depends, HTTPException, status +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager @@ -109,6 +110,36 @@ class _KeyTable(Protocol): async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... +class _AgentRecord(Protocol): + @property + def agent_id(self) -> str: ... + + @property + def access_group_ids(self) -> Sequence[str] | None: ... + + +class _AgentTable(Protocol): + async def find_many(self, where: Mapping[str, object]) -> Sequence[_AgentRecord]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + + +class _HasSomeFilter(TypedDict): + hasSome: ReadOnly[Sequence[str]] + + +class _AgentAccessGroupsWhere(TypedDict): + access_group_ids: ReadOnly[_HasSomeFilter] + + +class _AgentIdWhere(TypedDict): + agent_id: ReadOnly[str] + + +class _AgentAccessGroupsData(TypedDict): + access_group_ids: ReadOnly[Sequence[str]] + + class _AccessGroupTx(Protocol): @property def litellm_accessgrouptable(self) -> _AccessGroupTable: ... @@ -119,6 +150,9 @@ class _AccessGroupTx(Protocol): @property def litellm_verificationtoken(self) -> _KeyTable: ... + @property + def litellm_agentstable(self) -> _AgentTable: ... + def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: @@ -324,6 +358,41 @@ async def _sync_remove_access_group_from_keys(tx: _AccessGroupTx, key_tokens: li ) +def _without_access_group(access_group_ids: Sequence[str] | None, access_group_id: str) -> tuple[str, ...]: + return tuple(ag for ag in (access_group_ids or ()) if ag != access_group_id) + + +async def _detach_access_group_from_agents(tx: _AccessGroupTx, access_group_id: str) -> tuple[str, ...]: + agents_with_group: Final = await tx.litellm_agentstable.find_many( + where=_AgentAccessGroupsWhere(access_group_ids=_HasSomeFilter(hasSome=(access_group_id,))) + ) + for agent in agents_with_group: + await tx.litellm_agentstable.update( + where=_AgentIdWhere(agent_id=agent.agent_id), + data=_AgentAccessGroupsData( + access_group_ids=_without_access_group(agent.access_group_ids, access_group_id) + ), + ) + return tuple(agent.agent_id for agent in agents_with_group) + + +def _detach_access_group_from_agent_registry(agent_ids: Sequence[str], access_group_id: str) -> None: + registered: Final = tuple( + agent + for agent in (global_agent_registry.get_agent_by_id(agent_id) for agent_id in agent_ids) + if agent is not None + ) + for agent in registered: + global_agent_registry.deregister_agent(agent_name=agent.agent_name) + global_agent_registry.register_agent( + agent_config=agent.model_copy( + update=_AgentAccessGroupsData( + access_group_ids=_without_access_group(agent.access_group_ids, access_group_id) + ) + ) + ) + + # --------------------------------------------------------------------------- # Cache patch helpers # --------------------------------------------------------------------------- @@ -705,11 +774,14 @@ async def delete_access_group( 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) + detached_agent_ids: Final = await _detach_access_group_from_agents(tx, access_group_id) + await tx.litellm_accessgrouptable.delete(where={"access_group_id": access_group_id}) from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache await invalidate_access_group_cache(access_group_id) + _detach_access_group_from_agent_registry(detached_agent_ids, access_group_id) await _patch_team_caches_remove_access_group( affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj ) diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 07234883062..292cec1346d 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,3 +1,4 @@ +import re from collections.abc import Mapping, Sequence from datetime import datetime from typing import Final, Protocol @@ -21,6 +22,51 @@ from litellm.repositories.table_repositories import JWTKeyMappingRepository router: Final = APIRouter() +_TOKEN_HASH_PATTERN: Final = re.compile(r"[0-9a-f]{64}") + + +def _validated_token_hash(token: str) -> str: + """Guards a plaintext key from being stored as a hash of a hash, which would never match.""" + if _TOKEN_HASH_PATTERN.fullmatch(token) is None: + raise HTTPException( + status_code=400, + detail=( + "`token` must be the SHA-256 hash of a virtual key " + "(64 lowercase hex characters). Pass the plaintext as `key` instead." + ), + ) + return token + + +_EXACTLY_ONE_IDENTIFIER: Final = ( + "Provide exactly one of `key` (the plaintext virtual key) or `token` (its SHA-256 hash)." +) +_AT_MOST_ONE_IDENTIFIER: Final = ( + "Provide at most one of `key` (the plaintext virtual key) or `token` (its SHA-256 hash)." +) + + +def _token_hash_for_create(data: CreateJWTKeyMappingRequest) -> str: + """Resolve the token hash to store, from either the plaintext key or its hash.""" + if data.key is not None and data.token is not None: + raise HTTPException(status_code=400, detail=_EXACTLY_ONE_IDENTIFIER) + if data.token is not None: + return _validated_token_hash(data.token) + if data.key is not None: + return hash_token(data.key) + raise HTTPException(status_code=400, detail=_EXACTLY_ONE_IDENTIFIER) + + +def _token_hash_for_update(data: UpdateJWTKeyMappingRequest) -> str | None: + """Resolve the token hash to store, or None to leave the mapped key alone.""" + if data.key is not None and data.token is not None: + raise HTTPException(status_code=400, detail=_AT_MOST_ONE_IDENTIFIER) + if data.token is not None: + return _validated_token_hash(data.token) + if data.key is not None: + return hash_token(data.key) + return None + class _JWTKeyMappingRecord(Protocol): """A ``LiteLLM_JWTKeyMapping`` row, viewed through the columns these endpoints read.""" @@ -111,7 +157,7 @@ async def create_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: - hashed_key: Final = hash_token(data.key) + hashed_key: Final = _token_hash_for_create(data) create_data: Final = { "jwt_issuer": data.jwt_issuer or "", "jwt_claim_name": data.jwt_claim_name, @@ -166,9 +212,10 @@ async def update_jwt_key_mapping( if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") - update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key"}) - if data.key is not None: - update_data["token"] = hash_token(data.key) + update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key", "token"}) + token_hash: Final = _token_hash_for_update(data) + if token_hash is not None: + update_data["token"] = token_hash if "jwt_issuer" in update_data: # DB column is NOT NULL (see schema.prisma); "" is the global/unscoped sentinel. update_data["jwt_issuer"] = update_data["jwt_issuer"] or "" diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index ee0e4db291b..9ad78876043 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -983,7 +983,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers=None, ) tools: Final = listing.tools - dumped_tools: Final = [dict(tool) for tool in tools] + dumped_tools: Final = [tool.model_dump(by_alias=True) for tool in tools] return {"tools": dumped_tools} diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 10a0a2f3104..fcadcfe2cae 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -14,12 +14,12 @@ import asyncio import datetime import json from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence -from contextlib import AbstractAsyncContextManager, asynccontextmanager +from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress from dataclasses import dataclass from fnmatch import fnmatchcase from json import JSONDecodeError from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast, runtime_checkable +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator @@ -29,6 +29,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.ptu_pricing import ( CUSTOM_PRICING_FIELDS, PTU_EMPTIED_PRICING_FIELDS, @@ -139,7 +140,12 @@ from litellm.types.router import ( updateDeployment, updateLiteLLMParams, ) -from litellm.types.utils import echoed_cost_map_pricing_fields, without_server_derived_pricing +from litellm.types.utils import ( + COST_MAP_LOOKUP_KEY, + echoed_cost_map_fields, + echoed_cost_map_pricing_fields, + without_server_derived_pricing, +) from litellm.utils import get_utc_datetime if TYPE_CHECKING: @@ -928,7 +934,33 @@ def _ptu_priced_deployment(model_params: Deployment) -> Deployment: ) -def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: +def _cost_map_entry(db_model: Deployment, incoming_model_info: Mapping[str, object]) -> Mapping[str, object]: + base_model: Final = incoming_model_info.get("base_model") + lookup: Final = base_model if isinstance(base_model, str) else _decrypted_model(db_model.litellm_params.model) + if lookup is None: + return MappingProxyType({}) + with suppress(Exception): + return MappingProxyType(dict(litellm.get_model_info(model=lookup))) + return MappingProxyType({}) + + +LoadedCatalog: TypeAlias = Callable[[], Mapping[str, Mapping[str, object]]] # mutable-ok: Callable parameter syntax + + +def _loaded_catalog_entry( + incoming_model_info: Mapping[str, object], loaded_catalog: LoadedCatalog +) -> Mapping[str, object]: + catalog_key: Final = incoming_model_info.get(COST_MAP_LOOKUP_KEY) + if not isinstance(catalog_key, str): + return MappingProxyType({}) + return loaded_catalog().get(catalog_key, MappingProxyType({})) + + +def update_db_model( + db_model: Deployment, + updated_patch: updateDeployment, + loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, +) -> PrismaCompatibleUpdateDBModel: if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) merged_model_name: Final = updated_patch.model_name or db_model.model_name @@ -955,7 +987,24 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr # update model info if updated_patch.model_info: - merged_model_info.update(without_server_derived_pricing(updated_patch.model_info.model_dump(exclude_none=True))) + incoming_model_info: Final = updated_patch.model_info.model_dump(exclude_none=True) + echoed_fields: Final = echoed_cost_map_fields( + incoming_model_info, + _cost_map_entry(db_model, incoming_model_info), + _loaded_catalog_entry(incoming_model_info, loaded_catalog), + ) + merged_model_info.update( + MappingProxyType( + dict( + (k, v) + for k, v in without_server_derived_pricing(incoming_model_info).items() + if k not in echoed_fields + ) + ) + ) + for k in echoed_fields: + if k in merged_model_info and merged_model_info[k] != incoming_model_info[k]: + del merged_model_info[k] # Honor explicit-null clears LAST, after both merges, so a model_info blob a client # passes through cannot silently undo a litellm_params clear via .update(). diff --git a/litellm/proxy/middleware/budget_reservation_release_middleware.py b/litellm/proxy/middleware/budget_reservation_release_middleware.py new file mode 100644 index 00000000000..f7ac885274e --- /dev/null +++ b/litellm/proxy/middleware/budget_reservation_release_middleware.py @@ -0,0 +1,33 @@ +from collections.abc import Awaitable, Callable, Mapping +from typing import Final + +from starlette.types import ASGIApp, Receive, Scope, Send + +_SCOPES_AUTH_STAMPS: Final = frozenset({"http", "websocket"}) + + +class BudgetReservationReleaseMiddleware: + """Releases the budget reservation auth made for a request once no callback owns it. + + Auth stamps the reservation on the request or socket state; a call that starts + claims it for the cost callbacks, which settle it on success or failure. When the + response has been sent or the socket has closed and the reservation is still + unclaimed, nothing else ever would, so it is released here instead of pinning the + spend counter until its TTL. + """ + + def __init__(self, app: ASGIApp, release: Callable[[Mapping[str, object]], Awaitable[None]]) -> None: + self.app = app + self.release = release + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] not in _SCOPES_AUTH_STAMPS: + await self.app(scope, receive, send) + return + try: + await self.app(scope, receive, send) + finally: + state: Final = scope.get("state") + budget_reservation: Final = state.get("budget_reservation") if isinstance(state, Mapping) else None + if isinstance(budget_reservation, Mapping): + await self.release(budget_reservation) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 44d9f11360d..b1960b9a046 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -14,6 +14,7 @@ import json import os import posixpath import re +import sys from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass from functools import partial @@ -52,6 +53,7 @@ from litellm.llms.deepgram.common_utils import ( deepgram_listen_requested_model, deepgram_listen_websocket_target, ) +from litellm.llms.fal_ai.cost_calculator import fal_ai_queue_base from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse @@ -100,6 +102,12 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, ) +from litellm.types.passthrough_endpoints.tinyfish import ( + TINYFISH_AUTHENTICATED_RUN_FIELDS, + TINYFISH_PASSTHROUGH_TIMEOUT_SECONDS, + TINYFISH_REJECTED_ENVELOPE_FIELDS, + is_allowed_tinyfish_endpoint, +) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials from litellm.types.router import LiteLLMParamsTypedDict from litellm.types.utils import LlmProviders @@ -421,6 +429,56 @@ async def cohere_proxy_route( return received_value +def _fal_target(endpoint: str) -> httpx.URL: + base_target_url: Final = fal_ai_queue_base() + encoded_endpoint: Final = httpx.URL(endpoint).path + normalized_endpoint: Final = encoded_endpoint if encoded_endpoint.startswith("/") else f"/{encoded_endpoint}" + base_url: Final = httpx.URL(base_target_url) + return base_url.copy_with( + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, normalized_endpoint), + ) + + +@router.api_route( + "/fal_ai/{endpoint:path}", + methods=["GET", "POST", "PUT", "DELETE", "PATCH"], # mutable-ok: FastAPI route metadata requires a list + tags=["Fal AI Pass-through", "pass-through"], # mutable-ok: FastAPI route metadata requires a list +) +async def fal_ai_proxy_route( + endpoint: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +): + updated_url: Final = _fal_target(endpoint) + fal_ai_api_key: Final = passthrough_endpoint_router.get_credentials( + custom_llm_provider="fal_ai", + region_name=None, + ) + if fal_ai_api_key is None: + raise HTTPException( + status_code=401, + detail="FAL_AI_API_KEY is not set and no fal_ai pass-through deployment credentials are configured", + ) + if "/requests/" not in endpoint: + priced_model: Final = f"fal_ai/{endpoint}" + if priced_model not in (litellm.model_cost or {}): + raise HTTPException( + status_code=400, + detail=f"{priced_model} has no pricing entry; only priced Fal endpoints can be submitted through /fal_ai", + ) + endpoint_func: Final = create_pass_through_route( + endpoint=endpoint, + target=str(updated_url), + custom_headers={ + "Authorization": f"Key {fal_ai_api_key}" + }, # mutable-ok: pass-through request headers require a mutable mapping + custom_llm_provider="fal_ai", + is_streaming_request=False, + ) + return await endpoint_func(request, fastapi_response, user_api_key_dict) + + @router.api_route( "/vllm/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -579,6 +637,42 @@ async def typesafe_proxy_route( return await endpoint_func(request, fastapi_response, user_api_key_dict) +@router.api_route( + "/openrouter/{endpoint:path}", + methods=["GET", "POST", "PUT", "DELETE", "PATCH"], # mutable-ok: FastAPI route metadata requires a list + tags=["OpenRouter Pass-through", "pass-through"], # mutable-ok: FastAPI route metadata requires a list +) +async def openrouter_proxy_route( + endpoint: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +): + base_target_url: Final = get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" + api_root: Final = base_target_url.removesuffix("/").removesuffix("/v1") + encoded_endpoint: Final = httpx.URL(endpoint).path + normalized_endpoint: Final = encoded_endpoint if encoded_endpoint.startswith("/") else f"/{encoded_endpoint}" + base_url: Final = httpx.URL(api_root) + updated_url: Final = base_url.copy_with( + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, normalized_endpoint), + ) + openrouter_api_key: Final = passthrough_endpoint_router.get_credentials( + custom_llm_provider="openrouter", + region_name=None, + ) + endpoint_func: Final = create_pass_through_route( + endpoint=endpoint, + target=str(updated_url), + custom_headers={ # mutable-ok: pass-through request headers require a mutable mapping + "Authorization": f"Bearer {openrouter_api_key}", + "Content-Type": "application/json", + }, + custom_llm_provider="openrouter", + is_streaming_request=False, + ) + return await endpoint_func(request, fastapi_response, user_api_key_dict) + + @router.api_route( "/milvus/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -3270,6 +3364,138 @@ async def cursor_proxy_route( return received_value +TINYFISH_JSON_OBJECT_BODY_DETAIL: Final = ( + "TinyFish requests must be a JSON object body sent with Content-Type: application/json." +) + + +async def _tinyfish_json_object_field_names(request: Request) -> frozenset[str] | None: + content_type: Final = request.headers.get("content-type", "") + if content_type and not is_json_content_type(content_type): + return None + raw_body: Final = await request.body() + if not raw_body: + return frozenset() + try: + parsed: Final[object] = json.loads(raw_body) # any-ok: json.loads -> Any + except (json.JSONDecodeError, UnicodeDecodeError): + return None + return frozenset(parsed) if isinstance(parsed, dict) else None + + +def _tinyfish_route_timeout() -> float | None: + # only raise the 600s default to cover legal 1200s runs; an operator's configured timeout still wins + proxy_server: Final = sys.modules.get("litellm.proxy.proxy_server") + operator_settings: Final = getattr(proxy_server, "general_settings", None) + operator_timeout: Final = ( + operator_settings.get("pass_through_request_timeout") if isinstance(operator_settings, Mapping) else None + ) + return None if operator_timeout is not None else TINYFISH_PASSTHROUGH_TIMEOUT_SECONDS + + +@router.api_route( + "/tinyfish/{endpoint:path}", + methods=["GET", "POST"], # mutable-ok: fastapi api_route requires List[str] + tags=["TinyFish Pass-through", "pass-through"], # mutable-ok: fastapi api_route requires a list +) +async def tinyfish_proxy_route( + endpoint: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection +) -> Response: + """ + Pass-through for the TinyFish Agent API (goal-based web automation). + + Forwarded endpoints: + - POST /v1/automation/run — run to completion (blocking) + - POST /v1/automation/run-async — submit a run, poll GET /v1/runs/{id} for the result + - POST /v1/automation/run-sse — run with SSE progress events + - GET /v1/runs/{id} — run status / result + - POST /v1/runs/{id}/cancel — cancel a run + + Every other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs + listing, which would let any caller discover other callers' run ids) returns 403: all + proxy callers share one upstream key. + + Credential lookup order: + 1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through) + 2. TINYFISH_API_KEY environment variable + + [Docs](https://docs.litellm.ai/docs/pass_through/tinyfish) + """ + from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + resolve_tinyfish_agent_api_base, + ) + + raw_endpoint_path: Final = httpx.URL(endpoint).path + encoded_endpoint: Final = raw_endpoint_path if raw_endpoint_path.startswith("/") else f"/{raw_endpoint_path}" + + if not is_allowed_tinyfish_endpoint(request.method, encoded_endpoint): + raise HTTPException( + status_code=403, + detail=f"{request.method} {encoded_endpoint} is not an allowed TinyFish Agent passthrough endpoint. " + "Allowed: POST /v1/automation/run, POST /v1/automation/run-async, POST /v1/automation/run-sse, " + "GET /v1/runs/{id}, POST /v1/runs/{id}/cancel.", + ) + + if request.method == "POST": + body_fields: Final = await _tinyfish_json_object_field_names(request) + if body_fields is None: + raise HTTPException(status_code=400, detail=TINYFISH_JSON_OBJECT_BODY_DETAIL) + envelope_fields: Final = tuple(sorted(body_fields & TINYFISH_REJECTED_ENVELOPE_FIELDS)) + if envelope_fields: + raise HTTPException( + status_code=400, + detail=f"Request fields [{', '.join(envelope_fields)}] are LiteLLM pass-through envelope controls " + "and are not accepted on the TinyFish route. Send the native TinyFish request body; streaming is " + "determined by the endpoint.", + ) + blocked_fields: Final = tuple(sorted(body_fields & TINYFISH_AUTHENTICATED_RUN_FIELDS)) + if ( + blocked_fields + and encoded_endpoint.startswith("/v1/automation/") + and str_to_bool(os.getenv("TINYFISH_ALLOW_AUTHENTICATED_RUNS")) is not True + ): + raise HTTPException( + status_code=403, + detail=f"Request fields [{', '.join(blocked_fields)}] run with the shared TinyFish account's saved " + "credentials and are disabled on this proxy. Ask the proxy admin to set " + "TINYFISH_ALLOW_AUTHENTICATED_RUNS=true to allow them.", + ) + + tinyfish_api_key: Final = passthrough_endpoint_router.get_credentials( + custom_llm_provider="tinyfish", + region_name=None, + ) + if tinyfish_api_key is None: + raise HTTPException( + status_code=401, + detail="TinyFish API key not found. Set the TINYFISH_API_KEY environment variable or add a " + "deployment with use_in_pass_through: true.", + ) + + base_url: Final = httpx.URL(resolve_tinyfish_agent_api_base()) + updated_url: Final = base_url.copy_with( + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) + ) + + endpoint_func: Final = create_pass_through_route( + endpoint=endpoint, + target=str(updated_url), + custom_headers=MappingProxyType({"X-API-Key": tinyfish_api_key}), + custom_llm_provider="tinyfish", + timeout=_tinyfish_route_timeout(), + ) + received_value: Final = await endpoint_func( + request, + fastapi_response, + user_api_key_dict, + ) + + return received_value + + VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON: Final = ( "Vertex AI auth failed: set a use_in_pass_through vertex model, default_vertex_config, or DEFAULT_VERTEXAI_* env" ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/fal_ai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/fal_ai_passthrough_logging_handler.py new file mode 100644 index 00000000000..3d1fad90e03 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/fal_ai_passthrough_logging_handler.py @@ -0,0 +1,71 @@ +from collections.abc import Mapping, Sequence +from typing import Final +from urllib.parse import urlparse + +import httpx + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.fal_ai.cost_calculator import fal_ai_passthrough_cost, fal_ai_queue_base +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.types.utils import ImageObject, ImageResponse + +FAL_AI_PROVIDER: Final[str] = litellm.LlmProviders.FAL_AI.value + + +def _url_parts(value: object) -> tuple[Mapping[str, object], ...]: + if isinstance(value, Mapping): + return (value,) if isinstance(value.get("url"), str) else () + if isinstance(value, Sequence) and not isinstance(value, str): + return tuple(item for item in value if isinstance(item, Mapping) and isinstance(item.get("url"), str)) + return () + + +class FalAIPassthroughLoggingHandler: + @staticmethod + def is_fal_ai_route(url_route: str, custom_llm_provider: str | None) -> bool: + return custom_llm_provider == FAL_AI_PROVIDER + + def fal_ai_passthrough_handler( + self, + response_body: Mapping[str, object], + request_body: Mapping[str, object], + logging_obj: LiteLLMLoggingObj, + url_route: str, + kwargs: Mapping[str, object], + ) -> PassThroughEndpointLoggingTypedDict: + base_path: Final = httpx.URL(fal_ai_queue_base()).path.strip("/") + raw_path: Final = urlparse(url_route).path.strip("/") + upstream_path: Final = raw_path.removeprefix(f"{base_path}/") if base_path else raw_path + model: Final = upstream_path.partition("/requests/")[0] + is_submit: Final = "/requests/" not in upstream_path + response: Final = ImageResponse( + data=tuple( + ImageObject(url=url) + for value in response_body.values() + for part in _url_parts(value) + if isinstance((url := part.get("url")), str) + ) + ) + response_cost: Final = fal_ai_passthrough_cost(model, request_body) if is_submit else None + response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads a precomputed cost off the response's hidden params + logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object + logging_obj.model_call_details["model"] = model # rebind-ok: same shared logging object + logging_obj.model_call_details["custom_llm_provider"] = FAL_AI_PROVIDER # rebind-ok: same shared logging object + logging_obj.model_call_details["response_cost"] = response_cost # rebind-ok: same shared logging object + verbose_proxy_logger.debug( + "Fal AI passthrough cost tracking: model %s, cost %s", + model, + response_cost, + ) + logging_result: Final[PassThroughEndpointLoggingTypedDict] = { + "result": response, + "kwargs": { + **kwargs, + "model": model, + "custom_llm_provider": FAL_AI_PROVIDER, + "response_cost": response_cost, + }, + } + return logging_result diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py new file mode 100644 index 00000000000..a6c3cb669a6 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py @@ -0,0 +1,425 @@ +import asyncio +import json +import os +import time +import urllib.parse +from collections.abc import Mapping, Sequence +from datetime import datetime +from types import MappingProxyType +from typing import Final, NamedTuple +from urllib.parse import urlparse + +import httpx +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, +) +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, +) +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.types.llms.custom_http import httpxSpecialProvider +from litellm.types.passthrough_endpoints.tinyfish import ( + TINYFISH_AGENT_DEFAULT_API_BASE, + TINYFISH_DEFAULT_COST_PER_STEP, + TINYFISH_MAX_CONSECUTIVE_POLL_FAILURES, + TINYFISH_MAX_POLLING_SECONDS, + TINYFISH_MODEL_NAME, + TINYFISH_POLLING_INTERVAL_SECONDS, + TINYFISH_TERMINAL_RUN_STATUSES, + TinyfishRun, +) +from litellm.types.utils import StandardPassThroughResponseObject + +_RUN_ADAPTER: Final = TypeAdapter(TinyfishRun) + +_EMPTY_KWARGS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _TinyfishLoggingPayload(NamedTuple): + result: StandardPassThroughResponseObject + kwargs: Mapping[str, object] + + def as_handler_result(self) -> PassThroughEndpointLoggingTypedDict: + handler_result: Final[PassThroughEndpointLoggingTypedDict] = { + "result": self.result, + "kwargs": {**self.kwargs}, + } + return handler_result + + +# asyncio tasks are weakly referenced by the loop; hold them until done or they can vanish mid-poll +_BACKGROUND_BILLING_TASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: task registry + + +def _register_billing_task(task: "asyncio.Task[None]") -> None: + _BACKGROUND_BILLING_TASKS.add(task) + task.add_done_callback(_BACKGROUND_BILLING_TASKS.discard) + task.add_done_callback(_warn_if_cancelled) + + +def _warn_if_cancelled(task: "asyncio.Task[None]") -> None: + # CancelledError bypasses the poller's exception handler, so shutdown-time charge loss must be logged here + if task.cancelled(): + verbose_proxy_logger.warning("TinyFish passthrough: billing poller cancelled mid-poll; the run may go unbilled") + + +_SSE_POLLER_SPAWNED_KEY: Final = "tinyfish_sse_poller_spawned" + + +def mark_sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> None: + logging_obj.model_call_details[_SSE_POLLER_SPAWNED_KEY] = True # rebind-ok: request-scoped scratch dict + + +def sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> bool: + return logging_obj.model_call_details.get(_SSE_POLLER_SPAWNED_KEY) is True + + +def run_id_from_sse_frames(frames: bytes) -> str | None: + return _run_id_from_sse_chunks(frames.decode("utf-8", errors="replace").splitlines()) + + +def resolve_tinyfish_agent_api_base() -> str: + raw: Final = (os.getenv("TINYFISH_AGENT_API_BASE") or TINYFISH_AGENT_DEFAULT_API_BASE).rstrip("/") + # a schemeless override would silently break both routing and billing (urlparse hostname becomes None) + return raw if "://" in raw else f"https://{raw}" + + +def resolve_tinyfish_cost_per_step() -> float: + raw: Final = os.getenv("TINYFISH_COST_PER_STEP") + if raw is None: + return TINYFISH_DEFAULT_COST_PER_STEP + try: + return float(raw) + except ValueError: + verbose_proxy_logger.warning( + "TINYFISH_COST_PER_STEP=%r is not a number; using the default rate %s", + raw, + TINYFISH_DEFAULT_COST_PER_STEP, + ) + return TINYFISH_DEFAULT_COST_PER_STEP + + +def is_tinyfish_agent_url(url: str) -> bool: + hostname: Final = urlparse(url).hostname + return hostname is not None and hostname == urlparse(resolve_tinyfish_agent_api_base()).hostname + + +def _parse_run(payload: object) -> TinyfishRun | None: + try: + return _RUN_ADAPTER.validate_python(payload) + except ValidationError as e: + verbose_proxy_logger.warning("TinyFish passthrough: unexpected run object shape: %s", e) + return None + + +def _run_cost(run: TinyfishRun | None) -> float | None: + if run is None: + return None + # TinyFish only invoices COMPLETED runs, so FAILED/CANCELLED runs must charge the team $0 + if run.get("status") != "COMPLETED": + return None + num_of_steps: Final = run.get("num_of_steps") + if num_of_steps is None: + return None + return num_of_steps * resolve_tinyfish_cost_per_step() + + +class TinyFishPassthroughLoggingHandler: + @staticmethod + def should_log_request(request_method: str, url_route: str) -> bool: + """Only run submissions are billed; GET /v1/runs* polling and cancels never write spend rows.""" + return request_method == "POST" and "/v1/automation/" in urlparse(url_route).path + + @staticmethod + def is_run_async_route(url_route: str) -> bool: + return urlparse(url_route).path.endswith("/v1/automation/run-async") + + @staticmethod + def tinyfish_passthrough_handler( + httpx_response: httpx.Response, + response_body: Mapping[str, object] | None, + logging_obj: LiteLLMLoggingObj, + url_route: str, + result: str, + start_time: datetime, + end_time: datetime, + cache_hit: bool, + request_body: Mapping[str, object], + **kwargs: object, # kwargs-ok: the passthrough logging dispatch forwards shared logging kwargs to every handler + ) -> PassThroughEndpointLoggingTypedDict: + """Bill a blocking POST /v1/automation/run: the response is the terminal run object.""" + try: + run: Final = _parse_run(response_body) if response_body is not None else None + handler_payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload( + run=run, + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + kwargs=kwargs, + ).as_handler_result() + except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request + verbose_proxy_logger.exception("Error in TinyFish passthrough logging handler: %s", e) + fallback_payload: Final[PassThroughEndpointLoggingTypedDict] = { + "result": StandardPassThroughResponseObject(response=result), + "kwargs": kwargs, + } + return fallback_payload + return handler_payload + + @staticmethod + def start_async_run_billing( + response_body: Mapping[str, object] | None, + logging_obj: LiteLLMLoggingObj, + result: str, + start_time: datetime, + cache_hit: bool, + **kwargs: object, # kwargs-ok: shared logging kwargs, replayed into _handle_logging when the run finishes + ) -> None: + """Bill POST /v1/automation/run-async once, when the polled run turns terminal.""" + submitted: Final = _parse_run(response_body) if response_body is not None else None + run_id: Final = submitted.get("run_id") if submitted is not None else None + if not run_id: + verbose_proxy_logger.warning( + "TinyFish passthrough: run-async response carried no run_id; logging the request without cost" + ) + task: Final = asyncio.create_task( + TinyFishPassthroughLoggingHandler._poll_and_log( + run_id=run_id, + logging_obj=logging_obj, + result=result, + start_time=start_time, + cache_hit=cache_hit, + kwargs=kwargs, + ) + ) + _register_billing_task(task) + + @staticmethod + def start_sse_run_billing( + run_id: str, + litellm_logging_obj: LiteLLMLoggingObj, + start_time: datetime, + client: AsyncHTTPHandler | None = None, + ) -> None: + """Bill POST /v1/automation/run-sse once via a detached poller that outlives client disconnects.""" + mark_sse_poller_spawned(litellm_logging_obj) + task: Final = asyncio.create_task( + TinyFishPassthroughLoggingHandler._poll_and_log( + run_id=run_id, + logging_obj=litellm_logging_obj, + result="", + start_time=start_time, + cache_hit=litellm_logging_obj.model_call_details.get("cache_hit") is True, + kwargs=_EMPTY_KWARGS, + client=client, + ) + ) + _register_billing_task(task) + + @staticmethod + async def _poll_and_log( + run_id: str | None, + logging_obj: LiteLLMLoggingObj, + result: str, + start_time: datetime, + cache_hit: bool, + kwargs: Mapping[str, object], + client: AsyncHTTPHandler | None = None, + ) -> None: + from ..pass_through_endpoints import pass_through_endpoint_logging + + try: + run: Final = ( + await TinyFishPassthroughLoggingHandler._poll_until_terminal(run_id, client) if run_id else None + ) + run_end_time: Final = datetime.now() # noqa: DTZ005 # naive to match the start_time stamped by pass_through_request + payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload( + run=run, + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=run_end_time, + kwargs=kwargs, + ) + await pass_through_endpoint_logging._handle_logging( # pyright: ignore[reportPrivateUsage] # shared passthrough logging dispatcher, same access as the assemblyai handler + logging_obj=logging_obj, + standard_logging_response_object=payload.result, + result=result, + start_time=start_time, + end_time=run_end_time, + cache_hit=cache_hit, + **payload.kwargs, + ) + except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request + verbose_proxy_logger.exception("[Non blocking logging error] TinyFish run-async billing failed: %s", e) + + @staticmethod + async def _poll_until_terminal( + run_id: str, + client: AsyncHTTPHandler | None = None, + poll_interval_seconds: float = TINYFISH_POLLING_INTERVAL_SECONDS, + ) -> TinyfishRun | None: + deadline: Final = time.monotonic() + TINYFISH_MAX_POLLING_SECONDS + last_run: TinyfishRun | None = None # rebind-ok: poll-loop state + consecutive_failures = 0 # rebind-ok: poll-loop state + while time.monotonic() < deadline: + run = await TinyFishPassthroughLoggingHandler._fetch_run(run_id, client) + if run is None: + # a single transient poll failure must not drop the run's charge + consecutive_failures += 1 + if consecutive_failures >= TINYFISH_MAX_CONSECUTIVE_POLL_FAILURES: + verbose_proxy_logger.warning( + "TinyFish passthrough: giving up on run %s after %s consecutive poll failures; " + "logging the request without cost", + run_id, + consecutive_failures, + ) + return last_run + else: + consecutive_failures = 0 + last_run = run + if (run.get("status") or "") in TINYFISH_TERMINAL_RUN_STATUSES: + return run + await asyncio.sleep(poll_interval_seconds) + verbose_proxy_logger.warning( + "TinyFish passthrough: run %s not terminal after %ss; logging the request without cost", + run_id, + TINYFISH_MAX_POLLING_SECONDS, + ) + return last_run + + @staticmethod + async def _fetch_run(run_id: str, client: AsyncHTTPHandler | None = None) -> TinyfishRun | None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + passthrough_endpoint_router, + ) + + api_key: Final = passthrough_endpoint_router.get_credentials(custom_llm_provider="tinyfish", region_name=None) + if api_key is None: + verbose_proxy_logger.warning("TinyFish passthrough: no API key available to poll run %s", run_id) + return None + if any(c in run_id for c in ("/", "\\", "#", "?")) or ".." in run_id: + verbose_proxy_logger.warning("TinyFish passthrough: invalid run_id %r", run_id) + return None + safe_run_id: Final = urllib.parse.quote(run_id, safe="") + resolved_client: Final = client or get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint, + params={"timeout": 30.0}, # mutable-ok: get_async_httpx_client takes a plain dict of client params + ) + try: + # screenshots=none keeps the poll payload small (no per-step screenshot URLs needed) + response: Final = await resolved_client.get( + f"{resolve_tinyfish_agent_api_base()}/v1/runs/{safe_run_id}?screenshots=none", + headers={"X-API-Key": api_key}, # mutable-ok: httpx headers= takes a plain dict + ) + if not (200 <= response.status_code < 300): + verbose_proxy_logger.warning( + "TinyFish passthrough: GET /v1/runs/%s returned %s", safe_run_id, response.status_code + ) + return None + payload: Final[object] = response.json() # any-ok: httpx Response.json() -> Any + return _parse_run(payload) + except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request + verbose_proxy_logger.warning("[Non blocking logging error] TinyFish run fetch failed: %s", e) + return None + + @staticmethod + async def handle_logging_tinyfish_collected_chunks( + litellm_logging_obj: LiteLLMLoggingObj, + url_route: str, + start_time: datetime, + all_chunks: Sequence[str], + end_time: datetime, + client: AsyncHTTPHandler | None = None, + ) -> PassThroughEndpointLoggingTypedDict: + """Fallback for run-sse streams with no poller: logs the request, pricing via one GET if a run_id parses.""" + try: + run_id: Final = _run_id_from_sse_chunks(all_chunks) + if run_id is None: + verbose_proxy_logger.warning( + "TinyFish passthrough: no run_id in SSE stream; logging the request without cost" + ) + run: Final = await TinyFishPassthroughLoggingHandler._fetch_run(run_id, client) if run_id else None + payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload( + run=run, + logging_obj=litellm_logging_obj, + result="", + start_time=start_time, + end_time=end_time, + kwargs=_EMPTY_KWARGS, + ).as_handler_result() + except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request + verbose_proxy_logger.exception("Error in TinyFish SSE passthrough logging handler: %s", e) + fallback_payload: Final[PassThroughEndpointLoggingTypedDict] = { + "result": StandardPassThroughResponseObject(response=""), + "kwargs": {}, + } + return fallback_payload + return payload + + @staticmethod + def _build_logging_payload( + run: TinyfishRun | None, + logging_obj: LiteLLMLoggingObj, + result: str, + start_time: datetime, + end_time: datetime, + kwargs: Mapping[str, object], + ) -> _TinyfishLoggingPayload: + response_cost: Final = _run_cost(run) + updated_kwargs: Final = { # mutable-ok: the logging pipeline requires a plain kwargs dict + **kwargs, + "model": TINYFISH_MODEL_NAME, + "custom_llm_provider": "tinyfish", + "response_cost": response_cost, + # spend rows key on this as request_id; without it every poller-billed row is a NULL-key collision + "litellm_call_id": logging_obj.litellm_call_id, + # the poller paths pass no request kwargs, so SLO attribution (key hash, team, tags) needs the stored params + "litellm_params": kwargs.get("litellm_params") + or logging_obj.model_call_details.get("litellm_params") + or {}, # mutable-ok: the logging pipeline requires a plain kwargs dict + } + logging_obj.model_call_details.update( + model=TINYFISH_MODEL_NAME, + custom_llm_provider="tinyfish", + response_cost=response_cost, + ) + + logged_response: Final = StandardPassThroughResponseObject( + response=json.dumps(run) if run is not None else result + ) + standard_logging_object: Final = get_standard_logging_object_payload( + kwargs=updated_kwargs, + init_response_obj=logged_response, + start_time=start_time, + end_time=end_time, + logging_obj=logging_obj, + status="success", + ) + return _TinyfishLoggingPayload( + result=logged_response, + kwargs=MappingProxyType({**updated_kwargs, "standard_logging_object": standard_logging_object}), + ) + + +def _run_id_from_sse_chunks(all_chunks: Sequence[str]) -> str | None: + for line in all_chunks: + if not line.startswith("data:"): + continue + try: + event_payload: object = json.loads(line[5:].strip()) # any-ok: json.loads -> Any + except json.JSONDecodeError: + continue + event = _parse_run(event_payload) + if event is None: + continue + run_id = event.get("run_id") + if run_id: + return run_id + return None diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/typesafe_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/typesafe_passthrough_logging_handler.py index 9b196660c2c..887d17a7a20 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/typesafe_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/typesafe_passthrough_logging_handler.py @@ -65,6 +65,7 @@ class TypeSafePassthroughLoggingHandler: end_time: datetime, cache_hit: bool, request_body: Mapping[str, object], + custom_llm_provider: str, **kwargs: object, ) -> PassThroughEndpointLoggingTypedDict: response: Final = _parse_typesafe_response(response_body) @@ -72,12 +73,12 @@ class TypeSafePassthroughLoggingHandler: request_model_value: Final = request_body.get("model") request_model: Final = request_model_value if isinstance(request_model_value, str) else None logged_model: Final = response_model or request_model or "unknown" - model_name: Final = f"typesafe/{logged_model}" + model_name: Final = f"{custom_llm_provider}/{logged_model}" usage: Final = response.usage or _TypeSafeUsage() input_tokens: Final = usage.input_tokens output_tokens: Final = usage.output_tokens candidate_model_keys: Final = tuple( - f"typesafe/{model}" for model in (response_model, request_model) if model is not None + f"{custom_llm_provider}/{model}" for model in (response_model, request_model) if model is not None ) pricing: Final = _pricing_for(candidate_model_keys) response_cost: Final = ( @@ -91,13 +92,13 @@ class TypeSafePassthroughLoggingHandler: updated_kwargs: Final = { # mutable-ok: pass-through logging contract requires mutable kwargs **kwargs, "model": model_name, - "custom_llm_provider": "typesafe", + "custom_llm_provider": custom_llm_provider, "response_cost": response_cost, "combined_usage_object": usage_object, } logging_obj.model_call_details.update( model=model_name, - custom_llm_provider="typesafe", + custom_llm_provider=custom_llm_provider, response_cost=response_cost, ) standard_logging_object: Final = get_standard_logging_object_payload( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 79a328f5199..c2874ac948f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -47,6 +47,7 @@ from litellm.constants import ( from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( + bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) @@ -78,11 +79,13 @@ from litellm.proxy.common_request_processing import ( open_sse_before_first_byte, resolve_litellm_call_id, ) +from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, ) from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, error_status_code, litellm_call_id_headers, openai_error_param, @@ -110,6 +113,9 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( ) from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage +from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + is_tinyfish_agent_url, +) from .streaming_handler import PassThroughStreamingHandler from .success_handler import PassThroughEndpointLogging from .upstream_usage_headers import ( @@ -380,6 +386,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): or (parsed_url.hostname and "openai.com" in parsed_url.hostname) ): return EndpointType.OPENAI + elif is_tinyfish_agent_url(url): + return EndpointType.TINYFISH return EndpointType.GENERIC @staticmethod @@ -1127,6 +1135,9 @@ async def pass_through_request( from litellm.proxy.proxy_server import ( general_settings as proxy_general_settings, ) + from litellm.proxy.proxy_server import ( + general_settings_view, + ) _managed_id_provider: Final = resolve_passthrough_managed_id_provider(custom_llm_provider) @@ -1632,6 +1643,7 @@ async def pass_through_request( **kwargs, ) ) + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) ## CUSTOM HEADERS - `x-litellm-*` custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -1656,11 +1668,24 @@ async def pass_through_request( headers=response.headers, custom_headers=custom_headers, ) + emitted_call_id: Final = ( + JSON_OBJECT.validate_python(response_headers).get(LITELLM_CALL_ID_HEADER) + if response.status_code >= 400 + else None + ) + error_call_id: Final = ( + error_body_call_id(general_settings_view(), emitted_call_id) if isinstance(emitted_call_id, str) else None + ) + relayed_content: Final = ( + json.dumps(with_call_id(JSON_OBJECT.validate_python(response_body), error_call_id)).encode("utf-8") + if error_call_id is not None and isinstance(response_body, dict) + else content + ) if _content_modified: response_headers.pop("content-length", None) return Response( - content=content, + content=relayed_content, status_code=response.status_code, headers=response_headers, ) @@ -2543,6 +2568,7 @@ async def websocket_passthrough_request( **success_kwargs, ) ) + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) # Call the proxy logging success hook if proxy_logging_obj: @@ -2714,6 +2740,7 @@ async def _relay_passthrough_response_bytes( **success_handler_kwargs, ) ) + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) def _extract_model_from_vertex_ai_setup(setup_response: Mapping[str, object]) -> str | None: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index fe9e104789b..e1f13f2bee0 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -9,6 +9,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy._types import PassThroughEndpointLoggingResultValues @@ -26,6 +27,11 @@ from .llm_provider_handlers.gemini_passthrough_logging_handler import ( from .llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, ) +from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + TinyFishPassthroughLoggingHandler, + run_id_from_sse_frames, + sse_poller_spawned, +) from .llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) @@ -69,6 +75,9 @@ class PassThroughStreamingHandler: exception: Exception, stream_context: PassThroughStreamContext | None = None, ) -> None: + # the tinyfish poller writes the one authoritative row; a failure row here would collide on its request_id + if endpoint_type == EndpointType.TINYFISH and sse_poller_spawned(litellm_logging_obj): + return await asyncify(PassThroughStreamingHandler._record_partial_usage_for_failure)( litellm_logging_obj=litellm_logging_obj, endpoint_type=endpoint_type, @@ -178,12 +187,25 @@ class PassThroughStreamingHandler: ) ) ) + # TinyFish SSE bills via a detached poller spawned on the first run_id frame, so disconnects can't lose the charge + tinyfish_scan_active = endpoint_type == EndpointType.TINYFISH # rebind-ok: scan stops once the poller spawns + tinyfish_pending = b"" # rebind-ok: SSE frame reassembly buffer across transport chunks try: if not cost_injection_active: # Hot path: just buffer for end-of-stream logging and forward. async for chunk in response.aiter_bytes(): raw_bytes.append(chunk) PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj) + if tinyfish_scan_active: + complete_frames, tinyfish_pending = split_complete_sse_frames(tinyfish_pending + chunk) + run_id = run_id_from_sse_frames(complete_frames) if b"run_id" in complete_frames else None + if run_id: + TinyFishPassthroughLoggingHandler.start_sse_run_billing( + run_id=run_id, + litellm_logging_obj=litellm_logging_obj, + start_time=start_time, + ) + tinyfish_scan_active = False yield chunk else: # ``cost_injection_active`` already requires ``model_name`` to @@ -218,6 +240,7 @@ class PassThroughStreamingHandler: and response.status_code < 400 ): logging_scheduled = True + bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params) litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),) except Exception as e: verbose_proxy_logger.error("Error in chunk_processor: %s", e) @@ -250,6 +273,8 @@ class PassThroughStreamingHandler: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine()) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) + else: + bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params) @staticmethod async def _route_streaming_logging_to_handler( @@ -290,6 +315,37 @@ class PassThroughStreamingHandler: and not _is_provider_error_chunk(complete_frames) ) try: + # TinyFish billing is owned by the detached poller; the $0 fallback below is only for streams with no run_id + if endpoint_type == EndpointType.TINYFISH: + if sse_poller_spawned(litellm_logging_obj): + return + late_run_id: Final = run_id_from_sse_frames(b"".join(raw_bytes)) + if late_run_id: + # the run_id arrived in an unterminated frame; poll to terminal instead of mispricing a RUNNING run + TinyFishPassthroughLoggingHandler.start_sse_run_billing( + run_id=late_run_id, + litellm_logging_obj=litellm_logging_obj, + start_time=start_time, + ) + return + tinyfish_payload: Final = ( + await TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + url_route=url_route, + start_time=start_time, + all_chunks=PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes), + end_time=end_time, + ) + ) + await litellm_logging_obj.dispatch_success_handlers( + result=tinyfish_payload["result"], + start_time=start_time, + end_time=end_time, + cache_hit=litellm_logging_obj.model_call_details.get("cache_hit") is True, + prefer_async_handlers=True, + **tinyfish_payload["kwargs"], + ) + return ( standard_logging_response_object, kwargs, diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index de1a8ae1d93..6bba879b6c1 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -29,9 +29,16 @@ from .llm_provider_handlers.cursor_passthrough_logging_handler import ( from .llm_provider_handlers.deepgram_listen_passthrough_logging_handler import ( DeepgramListenPassthroughLoggingHandler, ) +from .llm_provider_handlers.fal_ai_passthrough_logging_handler import ( + FalAIPassthroughLoggingHandler, +) from .llm_provider_handlers.gemini_passthrough_logging_handler import ( GeminiPassthroughLoggingHandler, ) +from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + TinyFishPassthroughLoggingHandler, + is_tinyfish_agent_url, +) from .llm_provider_handlers.transcribe_passthrough_logging_handler import ( TRANSCRIBE_CUSTOM_LLM_PROVIDER, PassThroughLogDispatch, @@ -278,6 +285,22 @@ class PassThroughEndpointLogging: ) standard_logging_response_object = comprehend_medical_handler_result["result"] # rebind-ok: elif-chain kwargs = comprehend_medical_handler_result["kwargs"] # rebind-ok: elif-chain contract + elif self.is_tinyfish_route(url_route, custom_llm_provider): + tinyfish_handler_result: Final = TinyFishPassthroughLoggingHandler.tinyfish_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body if isinstance(response_body, dict) else None, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, + ) + standard_logging_response_object = tinyfish_handler_result["result"] # rebind-ok: elif-chain + kwargs = tinyfish_handler_result["kwargs"] # rebind-ok: elif-chain contract + elif self.is_azure_speech_route(custom_llm_provider): from .llm_provider_handlers.azure_speech_passthrough_logging_handler import ( AzureSpeechPassthroughLoggingHandler, @@ -311,7 +334,9 @@ class PassThroughEndpointLogging: ) standard_logging_response_object = transcribe_handler_result["result"] # rebind-ok: elif-chain kwargs = transcribe_handler_result["kwargs"] # rebind-ok: elif-chain contract - elif self.is_typesafe_route(custom_llm_provider): + elif self.is_typesafe_route(custom_llm_provider) or self.is_openrouter_decisions_route( + url_route, custom_llm_provider + ): from .llm_provider_handlers.typesafe_passthrough_logging_handler import ( TypeSafePassthroughLoggingHandler, ) @@ -326,10 +351,12 @@ class PassThroughEndpointLogging: end_time=end_time, cache_hit=cache_hit, request_body=request_body, + custom_llm_provider=custom_llm_provider or "", **kwargs, ) standard_logging_response_object = typesafe_handler_result["result"] kwargs = typesafe_handler_result["kwargs"] + elif self.is_vertex_ai_live_route(url_route): from .llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import ( VertexAILivePassthroughLoggingHandler, @@ -367,6 +394,16 @@ class PassThroughEndpointLogging: ) standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract + elif FalAIPassthroughLoggingHandler.is_fal_ai_route(url_route, custom_llm_provider): + fal_ai_handler_result: Final = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler( + response_body=response_body if isinstance(response_body, dict) else MappingProxyType({}), + request_body=request_body, + logging_obj=logging_obj, + url_route=url_route, + kwargs=kwargs, + ) + standard_logging_response_object = fal_ai_handler_result["result"] # rebind-ok: elif-chain + kwargs = fal_ai_handler_result["kwargs"] # rebind-ok: elif-chain contract return_dict["standard_logging_response_object"] = standard_logging_response_object return_dict["kwargs"] = kwargs @@ -389,6 +426,20 @@ class PassThroughEndpointLogging: ): standard_logging_response_object: PassThroughEndpointLoggingResultValues | None = None logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload + if self.is_tinyfish_route(url_route, custom_llm_provider): + # polls and cancels never write spend rows; run-async bills once from the background poller + if not TinyFishPassthroughLoggingHandler.should_log_request(httpx_response.request.method, url_route): + return + if TinyFishPassthroughLoggingHandler.is_run_async_route(url_route): + TinyFishPassthroughLoggingHandler.start_async_run_billing( + response_body=response_body if isinstance(response_body, dict) else None, + logging_obj=logging_obj, + result=result, + start_time=start_time, + cache_hit=cache_hit, + **kwargs, + ) + return if self.is_assemblyai_route(url_route) and not self.is_azure_speech_route(custom_llm_provider): if AssemblyAIPassthroughLoggingHandler._should_log_request(httpx_response.request.method) is not True: return @@ -496,6 +547,9 @@ class PassThroughEndpointLogging: def is_comprehend_medical_route(self, custom_llm_provider: str | None) -> bool: return custom_llm_provider == "comprehendmedical" + def is_tinyfish_route(self, url_route: str, custom_llm_provider: str | None) -> bool: + return custom_llm_provider == "tinyfish" or is_tinyfish_agent_url(url_route) + def is_azure_speech_route(self, custom_llm_provider: str | None) -> bool: return custom_llm_provider == AZURE_SPEECH_CUSTOM_LLM_PROVIDER @@ -505,6 +559,9 @@ class PassThroughEndpointLogging: def is_typesafe_route(self, custom_llm_provider: str | None) -> bool: return custom_llm_provider == "typesafe" + def is_openrouter_decisions_route(self, url_route: str, custom_llm_provider: str | None) -> bool: + return custom_llm_provider == "openrouter" and urlparse(url_route).path.endswith("/alpha/decisions") + def is_langfuse_route(self, url_route: str): parsed_url: Final = urlparse(url_route) for route in self.TRACKED_LANGFUSE_ROUTES: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 56ae450ef63..0c4442072f2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -393,6 +393,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id from litellm.proxy.common_utils.healthy_model_filter import ( get_hidden_unhealthy_model_names, is_healthy_only_listing_default, @@ -418,6 +419,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, headers_with_litellm_call_id, litellm_call_id_headers, with_litellm_call_id, @@ -652,6 +654,9 @@ from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableRequestMetricsMiddleware, BillingRecorder, ) +from litellm.proxy.middleware.budget_reservation_release_middleware import ( + BudgetReservationReleaseMiddleware, +) from litellm.proxy.plugin_routes import ( register_plugins_from_config, ) @@ -730,7 +735,10 @@ from litellm.proxy.shutdown.scheduled_jobs import ( from litellm.proxy.spend_tracking.background_interaction_settlement import ( install_background_interaction_settlement, ) -from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.budget_reservation import ( + get_budget_window_start, + release_unbound_budget_reservation, +) from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( run_scheduled_daily_global_spend_reconcile, ) @@ -1817,7 +1825,10 @@ async def openai_exception_handler(request: Request, exc: ProxyException): # NOTE: DO NOT MODIFY THIS, its crucial to map to Openai exceptions _log_model_access_denial(exc) headers: Final = exc.headers - error_dict: Final = exc.to_dict() + error_dict: Final = with_call_id( + JSON_OBJECT.validate_python(exc.to_dict()), + error_body_call_id(general_settings_view(), headers.get(LITELLM_CALL_ID_HEADER)), + ) status_code: Final = int(exc.code) if exc.code else status.HTTP_500_INTERNAL_SERVER_ERROR _close_dangling_otel_server_span(request, status_code, exc=exc) return JSONResponse( @@ -2357,6 +2368,7 @@ app.add_middleware( # it sees prisma_client as of the first request rather than import time. sink_factory=lambda: gateway_request_accumulator if prisma_client is not None else None, ) +app.add_middleware(BudgetReservationReleaseMiddleware, release=release_unbound_budget_reservation) app.add_middleware(InFlightRequestsMiddleware) app.add_middleware(SecurityHeadersMiddleware) @@ -2481,6 +2493,13 @@ heuristic_v1_tuning_baselines: Mapping[str, str] | None = None # second ProxyConfig instance must not get its own independent lock over it. MODEL_RECONCILE_LOCK: Final = asyncio.Lock() general_settings: dict = {} +_GENERAL_SETTINGS_VIEW: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def general_settings_view() -> Mapping[str, object]: + return _GENERAL_SETTINGS_VIEW.validate_python(general_settings) + + config_passthrough_endpoints: list[dict[str, Any]] | None = None log_file: Final = "api_log.json" worker_config: Final = None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f802a141d41..c55456b2a40 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -73,6 +73,7 @@ model LiteLLM_AgentsTable { static_headers Json? @default("{}") extra_headers String[] @default([]) agent_access_groups String[] @default([]) + access_group_ids String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index f9e5c4ff1e4..e28fa2c06a4 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -36,10 +36,10 @@ from litellm.proxy.common_utils.user_api_key_cache import ( tag_cache_key, team_membership_reservation_cache_key, ) +from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router -from litellm.rust_bridge.token_counter import RustTokenizer, count_input_tokens, rust_tokenizer from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.types.router import DeploymentTypedDict @@ -366,6 +366,7 @@ async def reserve_budget_for_request( "reserved_cost": reservation_cost, "entries": applied_entries, "finalized": False, + "callback_bound": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), "input_tokens": max(input_token_counts.values(), default=None), } @@ -474,6 +475,19 @@ async def release_or_invalidate_budget_reservation( budget_reservation["finalized"] = True +async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object]) -> None: + """Release a reservation no logging callback took ownership of, once the request ended. + + A handler whose litellm call never builds a logging object (batch cancel, file + content, anything without the client decorator) runs no cost callback, so nothing + else would ever reconcile its reservation. A bound reservation is left alone: its + success or failure handler settles it, possibly after the response has been sent. + """ + if not isinstance(budget_reservation, dict) or budget_reservation.get("callback_bound") is True: + return + await release_or_invalidate_budget_reservation(budget_reservation=budget_reservation) + + async def _get_budget_counters( request_body: dict, valid_token: UserAPIKeyAuth, @@ -1475,9 +1489,6 @@ def _get_request_models( return (model,) if isinstance(model, str) else tuple(model) -TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 - - async def count_request_input_tokens( request_body: dict, route: str, @@ -1486,105 +1497,11 @@ async def count_request_input_tokens( ) -> Mapping[str, int]: """Input-token count per candidate model, counted once per request. - Tokenizing is the reservation path's dominant CPU cost and is O(prompt), so - counting a large prompt inline stalls every other request on the worker. - Models whose tokenizer the Rust bridge ports (Anthropic, tiktoken cl100k_base - and o200k_base) are counted from the raw body by the bridge when it is enabled, once per - distinct tokenizer, which parses and tokenizes with the GIL released. - Everything it declines is counted in Python, large prompts in a worker - thread. The counts are reused by both the max-cost and the input-cost - estimate. - """ + The counts are reused by both the max-cost and the input-cost estimate.""" models: Final = _get_request_models(request_body=request_body, route=route, llm_router=llm_router) if not models: return MappingProxyType({}) - tokenizers: Final[Mapping[str, RustTokenizer | None]] = MappingProxyType( - {model: rust_tokenizer(model) for model in models} - ) - distinct_tokenizers: Final[tuple[RustTokenizer, ...]] = tuple( - dict.fromkeys(tokenizer for tokenizer in tokenizers.values() if tokenizer is not None) - ) - rust_counts_by_tokenizer: Final[Mapping[RustTokenizer, int]] = MappingProxyType( - { - tokenizer: count.input_tokens - for tokenizer in distinct_tokenizers - if raw_body is not None and (count := await count_input_tokens(raw_body, tokenizer)) is not None - } - ) - rust_counts: Final = MappingProxyType( - { - model: rust_counts_by_tokenizer[tokenizer] - for model, tokenizer in tokenizers.items() - if tokenizer is not None and tokenizer in rust_counts_by_tokenizer - } - ) - python_models: Final = tuple(model for model in models if model not in rust_counts) - python_counts: Final = ( - MappingProxyType({}) - if not python_models - else _count_input_tokens_for_models(request_body=request_body, models=python_models) - if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS - else await asyncio.to_thread( - _count_input_tokens_for_models, - request_body=request_body, - models=python_models, - ) - ) - verbose_proxy_logger.debug("input token counts: rust=%s python=%s", dict(rust_counts), dict(python_counts)) - return MappingProxyType({**rust_counts, **python_counts}) - - -def _count_input_tokens_for_models( - request_body: dict, - models: Sequence[str], -) -> Mapping[str, int]: - return MappingProxyType( - { - model: tokens - for model in models - if (tokens := _count_input_tokens(request_body=request_body, model=model)) is not None - } - ) - - -_INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice") - - -def _approximate_input_size(request_body: Mapping[str, object]) -> int: - """Length of the request's input text, a cheap stand-in for tokenizing cost. - - Every field _count_input_tokens hands the tokenizer is sized here, and - rendering rather than walking keeps mapping keys in the total, which a tool - schema's property names are.""" - return sum(len(str(request_body.get(field, ""))) for field in _INPUT_SIZE_FIELDS) - - -def _count_input_tokens(request_body: dict, model: str) -> int | None: - try: - if "messages" in request_body: - try: - return litellm.token_counter( - model=model, - messages=request_body.get("messages") or (), - tools=request_body.get("tools"), - tool_choice=request_body.get("tool_choice"), - ) - except ValueError: - return _count_text_tokens(model=model, text=request_body.get("messages")) - if "prompt" in request_body: - return _count_text_tokens(model=model, text=request_body.get("prompt")) - if "input" in request_body: - return _count_text_tokens(model=model, text=request_body.get("input")) - if "query" in request_body or "documents" in request_body: - query_tokens: Final = _count_text_tokens(model=model, text=request_body.get("query")) - document_tokens: Final = _count_text_tokens( - model=model, - text=request_body.get("documents"), - ) - return query_tokens + document_tokens - except Exception: - verbose_proxy_logger.debug("Unable to count input tokens for budget reservation", exc_info=True) - return None + return await count_input_tokens(request_body=request_body, raw_body=raw_body, models=models) def _estimate_input_tokens( @@ -1595,7 +1512,9 @@ def _estimate_input_tokens( input_tokens: int | None = None, ) -> int | None: counted: Final = ( - input_tokens if input_tokens is not None else _count_input_tokens(request_body=request_body, model=model) + input_tokens + if input_tokens is not None + else count_input_tokens_for_model(request_body=request_body, model=model) ) if counted is not None: return counted @@ -1644,26 +1563,6 @@ def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None: return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None) -def _count_text_tokens(model: str, text: object) -> int: - if text is None: - return 0 - - token_count = 0 - stack: Final = [text] - while stack: - item = stack.pop() - if item is None: - continue - if isinstance(item, list): - stack.extend(item) - continue - if isinstance(item, dict): - token_count += litellm.token_counter(model=model, text=json.dumps(item)) - continue - token_count += litellm.token_counter(model=model, text=str(item)) - return token_count - - def _get_output_multiplier(request_body: dict) -> int: output_multiplier = 1 for key in ("n", "best_of"): diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py new file mode 100644 index 00000000000..6c7083fb6db --- /dev/null +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -0,0 +1,173 @@ +"""Input-token counting for the budget reservation path. + +Tokenizing is the reservation path's dominant CPU cost and is O(prompt), so +counting a large prompt inline stalls every other request on the worker. +Models whose tokenizer the Rust bridge ports (Anthropic, tiktoken cl100k_base +and o200k_base) are counted from the raw body by the bridge, once per distinct +tokenizer, which parses and tokenizes with the GIL released. Everything it +declines, and every model with no Rust tokenizer, is counted in Python, large +prompts in a worker thread. +""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.rust_bridge import runtime +from litellm.rust_bridge.catalog import Route, RouteContext +from litellm.rust_bridge.token_counter import ( + TOKEN_COUNTER, + RustTokenCounterFactory, + RustTokenizer, + native_count, + rust_tokenizer, +) + +TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000 + +_INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice") + + +def _approximate_input_size(request_body: Mapping[str, object]) -> int: + """Length of the request's input text, a cheap stand-in for tokenizing cost. + + Every field count_input_tokens_for_model hands the tokenizer is sized here, + and rendering rather than walking keeps mapping keys in the total, which a + tool schema's property names are.""" + return sum(len(str(request_body.get(field, ""))) for field in _INPUT_SIZE_FIELDS) + + +async def count_input_tokens( + request_body: dict, + raw_body: bytes | None, + models: Sequence[str], +) -> Mapping[str, int]: + """Input-token count per model, sharing one native count across models that + select the same tokenizer.""" + tokenizers: Final[tuple[tuple[str, RustTokenizer | None], ...]] = tuple( + (model, rust_tokenizer(model)) for model in models + ) + groups: Final[tuple[RustTokenizer | None, ...]] = tuple(dict.fromkeys(tokenizer for _, tokenizer in tokenizers)) + group_counts: Final = [ + await _count_group( + request_body=request_body, + raw_body=raw_body, + tokenizer=tokenizer, + models=tuple(model for model, selected in tokenizers if selected == tokenizer), + ) + for tokenizer in groups + ] + counts: Final = MappingProxyType({model: tokens for group in group_counts for model, tokens in group.items()}) + verbose_proxy_logger.debug("input token counts: %s", dict(counts)) + return counts + + +async def _count_group( + request_body: dict, + raw_body: bytes | None, + tokenizer: RustTokenizer | None, + models: tuple[str, ...], +) -> Mapping[str, int]: + async def python() -> Mapping[str, int]: + if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: + return _count_input_tokens_for_models(request_body=request_body, models=models) + return await asyncio.to_thread( + _count_input_tokens_for_models, + request_body=request_body, + models=models, + ) + + if tokenizer is None or raw_body is None: + return await python() + try: + return await runtime.arun( + RouteContext(Route.TOKEN_COUNTER, provider=tokenizer), + binding=TOKEN_COUNTER, + native=lambda factory: _native_counts(factory, tokenizer, raw_body, models), + python=python, + ) + except (RuntimeError, ValueError) as error: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking + + if isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): + raise + verbose_proxy_logger.debug("Rust token counter (%s) failed, counting in Python: %s", tokenizer, error) + return await python() + + +async def _native_counts( + factory: RustTokenCounterFactory, + tokenizer: RustTokenizer, + raw_body: bytes, + models: tuple[str, ...], +) -> Mapping[str, int]: + count: Final = await native_count(factory, tokenizer, raw_body) + verbose_proxy_logger.debug("Rust token counter (%s) counted %d input tokens", tokenizer, count.input_tokens) + return MappingProxyType({model: count.input_tokens for model in models}) + + +def _count_input_tokens_for_models( + request_body: dict, + models: Sequence[str], +) -> Mapping[str, int]: + return MappingProxyType( + { + model: tokens + for model in models + if (tokens := count_input_tokens_for_model(request_body=request_body, model=model)) is not None + } + ) + + +def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: + try: + if "messages" in request_body: + try: + return litellm.token_counter( + model=model, + messages=request_body.get("messages") or (), + tools=request_body.get("tools"), + tool_choice=request_body.get("tool_choice"), + ) + except ValueError: + return _count_text_tokens(model=model, text=request_body.get("messages")) + if "prompt" in request_body: + return _count_text_tokens(model=model, text=request_body.get("prompt")) + if "input" in request_body: + return _count_text_tokens(model=model, text=request_body.get("input")) + if "query" in request_body or "documents" in request_body: + query_tokens: Final = _count_text_tokens(model=model, text=request_body.get("query")) + document_tokens: Final = _count_text_tokens( + model=model, + text=request_body.get("documents"), + ) + return query_tokens + document_tokens + except Exception: + verbose_proxy_logger.debug("Unable to count input tokens for budget reservation", exc_info=True) + return None + + +def _count_text_tokens(model: str, text: object) -> int: + if text is None: + return 0 + + token_count = 0 + stack: Final = [text] + while stack: + item = stack.pop() + if item is None: + continue + if isinstance(item, list): + stack.extend(item) + continue + if isinstance(item, dict): + token_count += litellm.token_counter(model=model, text=json.dumps(item)) + continue + token_count += litellm.token_counter(model=model, text=str(item)) + return token_count diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index de5f545f109..bfd412db1a7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -148,6 +148,7 @@ from litellm.proxy._types import ( Member, UserAPIKeyAuth, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change @@ -8261,6 +8262,51 @@ async def _get_access_group_models( return tuple(dict.fromkeys((*team_group_models, *key_group_models))) +async def _agent_access_group_visible_models( + user_api_key_dict: "UserAPIKeyAuth", + llm_router: "Router | None", + include_model_access_groups: bool, + return_wildcard_routes: bool, + team_id: str | None, + resolve_agent_ceiling: CeilingResolver, +) -> frozenset[str] | None: + """Models an agent key may still list once its attached access groups cap it, ``None`` when + nothing caps it, so ``/v1/models`` never advertises a model the same key would be denied on.""" + from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models + + if not user_api_key_dict.agent_id: + return None + ceiling: Final = await resolve_agent_ceiling(user_api_key_dict.agent_id) + if ceiling is None: + return None + if llm_router is None: + return ceiling.models + proxy_model_list: Final = llm_router.get_model_names() + model_access_groups: Final = llm_router.get_model_access_groups() + granted: Final = get_team_models( + team_models=sorted(ceiling.models), + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + ) + if not granted: + return frozenset() + return frozenset( + get_complete_model_list( + key_models=granted, + team_models=(), + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=return_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + team_id=team_id, + ) + ) + + async def get_available_models_for_user( user_api_key_dict: "UserAPIKeyAuth", llm_router: Optional["Router"], @@ -8273,6 +8319,7 @@ async def get_available_models_for_user( only_model_access_groups: bool = False, return_wildcard_routes: bool = False, user_api_key_cache: Optional["UserApiKeyCache"] = None, + resolve_agent_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> list[str]: """ Get the list of models available to a user based on their API key and team permissions. @@ -8376,7 +8423,18 @@ async def get_available_models_for_user( team_id=effective_team_id, ) - return all_models + agent_visible: Final = await _agent_access_group_visible_models( + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + include_model_access_groups=include_model_access_groups, + return_wildcard_routes=return_wildcard_routes, + team_id=effective_team_id, + resolve_agent_ceiling=resolve_agent_ceiling, + ) + if agent_visible is None: + return all_models + capped: Final = [m for m in all_models if m in agent_visible] # mutable-ok: callers expect the list all_models is + return capped def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None: diff --git a/litellm/responses/dispatch.py b/litellm/responses/dispatch.py index b2748fca4b6..d240356805c 100644 --- a/litellm/responses/dispatch.py +++ b/litellm/responses/dispatch.py @@ -5,7 +5,7 @@ from typing import Final, TypeAlias, cast # noqa: TID251 # native binding sele from litellm.responses import main from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator -from litellm.rust_bridge.catalog import Context, Delivery, Route +from litellm.rust_bridge.catalog import Delivery, Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch, call_hook from litellm.rust_bridge.public_call import bind, optional_bool, optional_mapping, optional_str, signature from litellm.rust_bridge.responses.entrypoints import ( @@ -64,8 +64,8 @@ def _public_request( ) -def _context(request: LiteLLMResponsesRequest) -> Context: - return Context( +def _context(request: LiteLLMResponsesRequest) -> RouteContext: + return RouteContext( Route.RESPONSES, provider=request.custom_llm_provider, model=request.model, diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 5a4a08b760c..c5032536df4 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -433,13 +433,18 @@ def _bridges_to_chat_completions( return responses_api_provider_config is None or use_chat_completions_api is True +_RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED: Final = frozenset({"client_metadata"}) + + def _bridge_kwargs( kwargs: Mapping[str, object], responses_api_provider_config: BaseResponsesAPIConfig | None, allowed_openai_params: Sequence[str] | None, ) -> Mapping[str, object]: if responses_api_provider_config is None: - return kwargs + return MappingProxyType( + {key: value for key, value in kwargs.items() if key not in _RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED} + ) forwarded_keys: Final = frozenset( ( *litellm.OPENAI_CHAT_COMPLETION_PARAMS, @@ -448,7 +453,7 @@ def _bridge_kwargs( *GenericLiteLLMParams.model_fields, *(allowed_openai_params or ()), ) - ) + ).difference(_RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED) return MappingProxyType({key: value for key, value in kwargs.items() if key in forwarded_keys}) diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index c60020ab979..3b5cb85862d 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -91,16 +91,6 @@ async def create_mcp_list_tools_events( # Use the pre-processed MCP tools that were already fetched, filtered, and deduplicated by the parent filtered_mcp_tools: Final = pre_processed_mcp_tools - # Convert tools to dict format for the event - _mcp_tools_dict: Final = [ - tool.model_dump() - if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump", None)) - else tool.__dict__ - if hasattr(tool, "__dict__") - else {"name": getattr(tool, "name", str(tool))} - for tool in filtered_mcp_tools - ] - # Emit list tools completed event completed_event: Final = MCPListToolsCompletedEvent( type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 27affc09337..f0fef3974f4 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1779,7 +1779,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing: bool = True, tier_litellm_params: Mapping[str, object] | None = None, context_escalation_original_tier: ComplexityTier | str | None = None, - heuristic_v2_forecast: StandardLoggingHeuristicV2Forecast | None = None, + previous_decision: StandardLoggingRoutingDecision | None = None, ) -> StandardLoggingRoutingDecision: """Assemble the per-request provenance record for this router's decision. @@ -1839,8 +1839,15 @@ class ComplexityRouter(CustomLogger): masked_tier_litellm_params: Final = mask_credentials_in_payload(tier_litellm_params) if isinstance(masked_tier_litellm_params, Mapping): decision["tier_litellm_params"] = masked_tier_litellm_params - return ( - decision if heuristic_v2_forecast is None else {**decision, "heuristic_v2_forecast": heuristic_v2_forecast} + forecast_fields: Final = MappingProxyType( + { + field: value + for field, value in (previous_decision.items() if previous_decision is not None else ()) + if field.startswith("classifier_") or field == "heuristic_v2_forecast" + } + ) + return cast( # cast-ok: retaining optional keys from a typed decision preserves their declared values + StandardLoggingRoutingDecision, {**forecast_fields, **decision} ) async def aclassify( @@ -3595,7 +3602,7 @@ class ComplexityRouter(CustomLogger): context_escalation_original_tier=( decision.get("context_escalation_original_tier") if decision is not None else None ), - heuristic_v2_forecast=decision.get("heuristic_v2_forecast") if decision is not None else None, + previous_decision=decision, ) from litellm.types.router import PreRoutingHookResponse as HookResponse @@ -3775,7 +3782,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing=bool(decision.get("conversation_continuing", True)), tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model), context_escalation_original_tier=decision.get("context_escalation_original_tier"), - heuristic_v2_forecast=decision.get("heuristic_v2_forecast"), + previous_decision=decision, ) return response.model_copy( update={ # mutable-ok: model_copy types update as a plain dict @@ -3820,7 +3827,7 @@ class ComplexityRouter(CustomLogger): conversation_continuing=bool(decision.get("conversation_continuing", True)), tier_litellm_params=self._litellm_params_for_model(None, default_model), context_escalation_original_tier=decision.get("context_escalation_original_tier"), - heuristic_v2_forecast=decision.get("heuristic_v2_forecast"), + previous_decision=decision, ) return response.model_copy( update={ # mutable-ok: model_copy types update as a plain dict diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index a41df18b55f..c0d0d1de8e3 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -165,6 +165,7 @@ class HttpJevClassifierClient: end_time=end_time, cache_hit=False, request_body=MappingProxyType({"model": request.model}), + custom_llm_provider="typesafe", litellm_params=params, ) success_handlers: Final = logging_obj.dispatch_success_handlers( diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 978e770b4be..61e597bf674 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -94,7 +94,9 @@ class ResponsesWebSocketConnection: def close(self) -> Future[None]: ... @final -class _CacheTestBinding: +class _ResponseCacheRuntime: + @staticmethod + def from_cache(cache: object) -> _ResponseCacheRuntime: ... @property def kind(self) -> str: ... def lookup( @@ -167,8 +169,24 @@ class _CacheTestHandle: @staticmethod def disk(directory: str) -> _CacheTestHandle: ... @staticmethod + def qdrant_semantic( + url: str, + *, + collection_name: str, + similarity_threshold: float, + vector_size: int, + embedding_model: str = "text-embedding-3-small", + api_key: str | None = None, + embedding_api_key: str | None = None, + embedding_api_base: str | None = None, + embedding_timeout_seconds: float | None = None, + quantization: str = "binary", + ) -> _CacheTestHandle: ... + @staticmethod def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ... @staticmethod + def redis_semantic(backend: object) -> _CacheTestHandle: ... + @staticmethod def valkey_semantic( url: str, similarity_threshold: float, @@ -202,30 +220,125 @@ class _CacheTestHandle: @final class _CacheTestResolver: def __new__(cls, namespace: object) -> _CacheTestResolver: ... - def resolve(self) -> _CacheTestBinding: ... + def resolve(self) -> _ResponseCacheRuntime: ... @final class TokenCounter: - def __new__(cls, tokenizer_json: str) -> TokenCounter: ... @staticmethod - def from_cl100k_ranks(rank_file: str) -> TokenCounter: ... - @staticmethod - def from_o200k_ranks(rank_file: str) -> TokenCounter: ... - @staticmethod - def from_tiktoken(encoding: str) -> TokenCounter: ... + def from_tokenizer(tokenizer: Tokenizer, fast: bool = False) -> TokenCounter: ... def acount_request(self, body: bytes) -> Future[dict[str, object]]: ... +@final +class Tokenizer: + @staticmethod + def from_tiktoken(encoding: str) -> Tokenizer: ... + @staticmethod + def from_json(tokenizer_json: str) -> Tokenizer: ... + @staticmethod + def from_pretrained( + identifier: str, + revision: str = "main", + token: str | None = None, + ) -> Tokenizer: ... + def encode(self, text: str) -> list[int]: ... + def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str: ... + def count(self, text: str, fast: bool = False) -> int: ... + # tiktoken encodings + def encode_special(self, text: str, allowed: Sequence[str]) -> list[int]: ... + def encode_with_unstable(self, text: str, allowed: Sequence[str]) -> tuple[list[int], list[list[int]]]: ... + def encode_single_token(self, piece: bytes) -> int: ... + def special_tokens(self) -> dict[str, int]: ... + def max_token_value(self) -> int: ... + def is_special_token(self, token: int) -> bool: ... + def token_byte_values(self) -> list[bytes]: ... + def decode_bytes(self, ids: Sequence[int]) -> bytes: ... + # Hugging Face tokenizers + def to_json(self, pretty: bool = False) -> str: ... + def token_to_id(self, token: str) -> int | None: ... + def id_to_token(self, id: int) -> str | None: ... + def get_vocab(self, with_added_tokens: bool = True) -> dict[str, int]: ... + def get_vocab_size(self, with_added_tokens: bool = True) -> int: ... + def added_tokens_decoder(self) -> list[tuple[int, tuple[str, bool, bool, bool, bool, bool]]]: ... + def padding(self) -> dict[str, object] | None: ... + def truncation(self) -> dict[str, object] | None: ... + def num_special_tokens_to_add(self, is_pair: bool) -> int: ... + def encode_special_tokens(self) -> bool: ... + def encode_huggingface( + self, + sequence: str | Sequence[str], + pair: str | Sequence[str] | None = None, + is_pretokenized: bool = False, + add_special_tokens: bool = True, + fast: bool = False, + ) -> HuggingFaceEncoding: ... + def encode_batch_huggingface( + self, + inputs: Sequence[tuple[str | Sequence[str], str | Sequence[str] | None]], + is_pretokenized: bool = False, + add_special_tokens: bool = True, + fast: bool = False, + ) -> list[HuggingFaceEncoding]: ... + @property + def name(self) -> str: ... + +@final +class HuggingFaceEncoding: + def __new__(cls, json: str | None = None) -> HuggingFaceEncoding: ... + @staticmethod + def merge(encodings: Sequence[HuggingFaceEncoding], growing_offsets: bool = True) -> HuggingFaceEncoding: ... + def __len__(self) -> int: ... + def __reduce__(self) -> tuple[type[HuggingFaceEncoding], tuple[str]]: ... + def word_to_tokens(self, word_index: int, sequence_index: int = 0) -> tuple[int, int] | None: ... + def word_to_chars(self, word_index: int, sequence_index: int = 0) -> tuple[int, int] | None: ... + def token_to_sequence(self, token_index: int) -> int | None: ... + def token_to_chars(self, token_index: int) -> tuple[int, int] | None: ... + def token_to_word(self, token_index: int) -> int | None: ... + def char_to_token(self, char_pos: int, sequence_index: int = 0) -> int | None: ... + def char_to_word(self, char_pos: int, sequence_index: int = 0) -> int | None: ... + def set_sequence_id(self, sequence_id: int) -> None: ... + def pad( + self, + length: int, + direction: str = "right", + pad_id: int = 0, + pad_type_id: int = 0, + pad_token: str = "[PAD]", + ) -> None: ... + def truncate(self, max_length: int, stride: int = 0, direction: str = "right") -> None: ... + @property + def ids(self) -> list[int]: ... + @property + def tokens(self) -> list[str]: ... + @property + def offsets(self) -> list[tuple[int, int]]: ... + @property + def type_ids(self) -> list[int]: ... + @property + def attention_mask(self) -> list[int]: ... + @property + def special_tokens_mask(self) -> list[int]: ... + @property + def word_ids(self) -> list[int | None]: ... + @property + def sequence_ids(self) -> list[int | None]: ... + @property + def overflowing(self) -> list[HuggingFaceEncoding]: ... + @property + def n_sequences(self) -> int: ... + def gil_stats() -> dict[str, int]: ... def process_state_started() -> bool: ... def reserve_process_for_forking() -> None: ... __all__ = [ "ForkedAfterNativeRuntimeStarted", + "HuggingFaceEncoding", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", "RustUpstreamError", "TokenCounter", + "Tokenizer", "achat_completions", "amessages", "aocr", diff --git a/litellm/rust_bridge/callbacks_legacy_python.py b/litellm/rust_bridge/callbacks_legacy_python.py index 30aa1d97bfc..6bbf2ffed6b 100644 --- a/litellm/rust_bridge/callbacks_legacy_python.py +++ b/litellm/rust_bridge/callbacks_legacy_python.py @@ -59,9 +59,17 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - return CallSetup(supplied, arguments) + return _claim_budget_reservation(CallSetup(supplied, arguments), asynchronous) logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments) - return CallSetup(logger, prepared) + return _claim_budget_reservation(CallSetup(logger, prepared), asynchronous) + + +def _claim_budget_reservation(call_setup: CallSetup, asynchronous: bool) -> CallSetup: + from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks + + if asynchronous and not is_internal_call(): + bind_budget_reservation_to_callbacks(call_setup.logger.litellm_params) + return call_setup def check_limits(kwargs: Mapping[str, object]) -> None: @@ -96,6 +104,9 @@ def finalize( class LoggingSurface(Protocol): + @property + def litellm_params(self) -> Mapping[str, object]: ... + def update_from_kwargs( self, kwargs: dict[str, object], @@ -236,8 +247,12 @@ def sync_success_for_async_call( def failure_handler( logger: LoggingSurface, error: Exception, start: datetime.datetime, end: datetime.datetime, asynchronous: bool ) -> Coroutine[object, object, None] | None: + from litellm.litellm_core_utils.core_helpers import unbind_budget_reservation_from_callbacks + trace: Final = "".join(traceback.format_exception(error)) if asynchronous: + if not is_internal_call(): + unbind_budget_reservation_from_callbacks(logger.litellm_params) return logger.async_failure_handler(error, trace, start, end) logger.failure_handler(error, trace, start, end) return None diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 8794ff2db95..d7479631e04 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -1,9 +1,7 @@ -"""Declarative Rust/Python selection for routes with Rust integration. +"""Ordered rollout policy for routes, cache backends, and secret managers. -Rules are static data matched top to bottom; the first match wins and a -context with no matching rule stays on Python. Whether the Rust core can serve -a specific request body is not decided here: that is Rust admission, which -signals ``RustBridgeDeclined`` before any provider I/O. +The first matching rule wins; unmatched contexts stay on Python. Native +admission separately decides whether the selected implementation can execute. """ from __future__ import annotations @@ -14,6 +12,8 @@ from typing import Final, TypeAlias from litellm.rust_bridge.configuration import Decision, Rollout from litellm.rust_bridge.configuration import decision as _decision +from litellm.types.caching import LiteLLMCacheType +from litellm.types.secret_managers.main import KeyManagementSystem class Route(str, Enum): @@ -22,6 +22,8 @@ class Route(str, Enum): RESPONSES = "responses" TRANSCRIPTION = "transcription" OCR = "ocr" + TOKEN_COUNTER = "token_counter" + TOKENIZER = "tokenizer" class Delivery(Enum): @@ -31,7 +33,7 @@ class Delivery(Enum): @dataclass(frozen=True, slots=True) -class Context: +class RouteContext: route: Route provider: str | None = None model: str | None = None @@ -39,7 +41,7 @@ class Context: @dataclass(frozen=True, slots=True) -class Rule: +class RouteRule: route: Route rollout: Rollout providers: frozenset[str] | None = None @@ -48,26 +50,78 @@ class Rule: def matches(self, context: Context) -> bool: return ( - context.route is self.route + isinstance(context, RouteContext) + and context.route is self.route and (self.providers is None or context.provider in self.providers) and (self.models is None or context.model in self.models) and (self.deliveries is None or context.delivery in self.deliveries) ) +@dataclass(frozen=True, slots=True) +class CacheContext: + backend: str + + +@dataclass(frozen=True, slots=True) +class CacheRule: + rollout: Rollout + backends: frozenset[str] | None = None + + def matches(self, context: Context) -> bool: + return isinstance(context, CacheContext) and (self.backends is None or context.backend in self.backends) + + +@dataclass(frozen=True, slots=True) +class SecretManagerContext: + system: str + + +@dataclass(frozen=True, slots=True) +class SecretManagerRule: + rollout: Rollout + systems: frozenset[str] | None = None + + def matches(self, context: Context) -> bool: + return isinstance(context, SecretManagerContext) and (self.systems is None or context.system in self.systems) + + +Context: TypeAlias = RouteContext | CacheContext | SecretManagerContext +Rule: TypeAlias = RouteRule | CacheRule | SecretManagerRule Rules: TypeAlias = tuple[Rule, ...] RULES: Final[Rules] = ( - Rule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), - Rule(Route.OCR, Rollout.RUST_OPT_OUT), - Rule(Route.MESSAGES, Rollout.RUST_OPT_IN), - Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), + RouteRule(Route.OCR, Rollout.RUST_OPT_OUT), + RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), + RouteRule(Route.TOKEN_COUNTER, Rollout.RUST_OPT_IN), + RouteRule(Route.TOKENIZER, Rollout.RUST_OPT_IN), + RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.LOCAL})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.REDIS})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.REDIS_SEMANTIC})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.VALKEY_SEMANTIC})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.S3})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.DISK})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.QDRANT_SEMANTIC})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.AZURE_BLOB})), + CacheRule(Rollout.PYTHON_ONLY, backends=frozenset({LiteLLMCacheType.GCS})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.GOOGLE_KMS.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.AZURE_KEY_VAULT.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.AWS_SECRET_MANAGER.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.GOOGLE_SECRET_MANAGER.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.HASHICORP_VAULT.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.CYBERARK.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.LOCAL.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.AWS_KMS.value})), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({KeyManagementSystem.CUSTOM.value})), ) -def rollout(context: Context, rules: Rules = RULES) -> Rollout: - return next((rule.rollout for rule in rules if rule.matches(context)), Rollout.PYTHON_ONLY) +def rollout(context: Context, rules: Rules | None = None) -> Rollout: + selected_rules: Final = RULES if rules is None else rules + return next((rule.rollout for rule in selected_rules if rule.matches(context)), Rollout.PYTHON_ONLY) -def decision(context: Context, rules: Rules = RULES) -> Decision: +def decision(context: Context, rules: Rules | None = None) -> Decision: return _decision(rollout(context, rules)) diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 791e13a51d0..152ba632996 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -2,6 +2,7 @@ from __future__ import annotations import os from enum import Enum, auto +from functools import lru_cache from typing import Final from pydantic import TypeAdapter, ValidationError @@ -32,7 +33,9 @@ class _RustConfiguration: _CONFIGURATION: Final = _RustConfiguration() +@lru_cache(maxsize=16) def _parse_env_bool(value: str | None) -> bool | None: + """`LITELLM_RUST` as a bool; cached by raw value because `decision` runs per tokenizer call.""" if value is None: return None try: @@ -84,7 +87,7 @@ def reset_rust_configuration() -> None: def rust(enabled: bool | None) -> None: """Set the process override for optional Rust paths. - ``PYTHON_ONLY`` and ``RUST_REQUIRED`` routes in the catalog ignore this switch, + ``PYTHON_ONLY`` and ``RUST_REQUIRED`` entries in the catalog ignore this switch, and an explicit ``LITELLM_RUST`` environment value wins over it. """ _CONFIGURATION.override = enabled diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index 7ddc903df58..076b7759c6d 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/litellm/rust_bridge/dispatch.py @@ -6,7 +6,7 @@ from typing import Final, Generic, TypeVar from litellm.rust_bridge import catalog, runtime from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Context, Route, Rules +from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule, Rules from litellm.rust_bridge.configuration import Decision from litellm.rust_bridge.configuration import decision as rollout_decision @@ -30,12 +30,12 @@ def call_hook( class PublicDispatch(Generic[RequestT]): route: Route request: Callable[[tuple[object, ...], Mapping[str, object]], RequestT | None] - context: Callable[[RequestT], Context] + context: Callable[[RequestT], RouteContext] bypass: Callable[[RequestT], bool] | None = None def _requires_projection(self, rules: Rules) -> bool: for rule in rules: - if rule.route is not self.route: + if not isinstance(rule, RouteRule) or rule.route is not self.route: continue if rule.providers is not None or rule.models is not None or rule.deliveries is not None: if rollout_decision(rule.rollout) is not Decision.PYTHON: diff --git a/litellm/rust_bridge/response_cache.py b/litellm/rust_bridge/response_cache.py new file mode 100644 index 00000000000..82d27fce27b --- /dev/null +++ b/litellm/rust_bridge/response_cache.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import math +from collections.abc import Awaitable, Mapping, Sequence +from dataclasses import dataclass +from typing import Final, Protocol, cast + +from typing_extensions import ReadOnly, Required, TypedDict, assert_never + +from litellm.rust_bridge.bindings import NativeBinding, native_exception_types +from litellm.rust_bridge.catalog import CacheContext, Rules, decision +from litellm.rust_bridge.configuration import Decision + + +class CacheFacade(Protocol): + @property + def type(self) -> object: ... + + @property + def ttl(self) -> float | None: ... + + @property + def semantic_cache_scope(self) -> str: ... + + def get_cache_key(self, **kwargs: object) -> str: ... # kwargs-ok: mirrors the legacy cache facade contract + + +class NativeCacheKey(TypedDict): + preset: ReadOnly[str] + + +class NativeCacheRequest(TypedDict, total=False): + key: Required[ReadOnly[NativeCacheKey]] + ttl_seconds: ReadOnly[float | None] + max_age_seconds: ReadOnly[float | None] + messages: ReadOnly[object | None] + input: ReadOnly[object | None] + metadata: ReadOnly[object | None] + litellm_metadata: ReadOnly[object | None] + litellm_params: ReadOnly[object | None] + scope: ReadOnly[str] + + +class NativeResponseCacheRuntime(Protocol): + @property + def kind(self) -> str: ... + + def lookup(self, request: NativeCacheRequest) -> object: ... + def store(self, request: NativeCacheRequest, response: object) -> None: ... + def lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> object: ... + def async_lookup(self, request: NativeCacheRequest) -> Awaitable[object]: ... + def async_store(self, request: NativeCacheRequest, response: object) -> Awaitable[None]: ... + def async_lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> Awaitable[object]: ... + def async_store_batch( + self, + requests: Sequence[NativeCacheRequest], + responses: Sequence[object], + ) -> Awaitable[object]: ... + def async_flush(self) -> Awaitable[None]: ... + def ping(self) -> Awaitable[object]: ... + + +class NativeResponseCacheRuntimeFactory(Protocol): + @staticmethod + def from_cache(cache: CacheFacade) -> NativeResponseCacheRuntime: ... + + +def _runtime_factory(value: object) -> NativeResponseCacheRuntimeFactory | None: + return cast(NativeResponseCacheRuntimeFactory, value) if callable(getattr(value, "from_cache", None)) else None + + +_RUNTIME: Final = NativeBinding("_ResponseCacheRuntime", validate=_runtime_factory) + + +@dataclass(frozen=True, slots=True) +class ResponseCacheRuntime: + native: NativeResponseCacheRuntime + + @property + def kind(self) -> str: + return self.native.kind + + def request(self, cache: CacheFacade, kwargs: Mapping[str, object]) -> NativeCacheRequest | None: + key_value: Final = kwargs.get("cache_key") + key: Final = key_value if isinstance(key_value, str) else cache.get_cache_key(**dict(kwargs)) + if not key: + return None + control_value: Final = kwargs.get("cache") + control: Final = _string_mapping(control_value) + configured_ttl: Final = cache.ttl if cache.ttl is not None else _duration(kwargs.get("ttl")) + control_ttl: Final = _duration(control.get("ttl")) + current_max_age: Final = _duration(control.get("s-max-age")) + legacy_max_age: Final = _duration(control.get("s-maxage")) + ttl: Final = configured_ttl if control_ttl is None else control_ttl + max_age: Final = legacy_max_age if current_max_age is None else current_max_age + return NativeCacheRequest( + key=NativeCacheKey(preset=key), + ttl_seconds=ttl, + max_age_seconds=max_age, + messages=kwargs.get("messages"), + input=kwargs.get("input"), + metadata=kwargs.get("metadata"), + litellm_metadata=kwargs.get("litellm_metadata"), + litellm_params=kwargs.get("litellm_params"), + scope=cache.semantic_cache_scope, + ) + + def lookup(self, request: NativeCacheRequest) -> object: + return self.native.lookup(request) + + def store(self, request: NativeCacheRequest, response: object) -> None: + self.native.store(request, response) + + def lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> object: + return self.native.lookup_batch(requests) + + async def async_lookup(self, request: NativeCacheRequest) -> object: + return await self.native.async_lookup(request) + + async def async_store(self, request: NativeCacheRequest, response: object) -> None: + await self.native.async_store(request, response) + + async def async_lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> object: + return await self.native.async_lookup_batch(requests) + + async def async_store_batch( + self, + requests: Sequence[NativeCacheRequest], + responses: Sequence[object], + ) -> object: + return await self.native.async_store_batch(requests, responses) + + async def ping(self) -> object: + return await self.native.ping() + + async def async_flush(self) -> None: + await self.native.async_flush() + + +def resolve_response_cache( + cache: CacheFacade, + rules: Rules | None = None, +) -> ResponseCacheRuntime | None: + backend_value: Final = cache.type + backend: Final = str.__str__(backend_value) if isinstance(backend_value, str) else str(backend_value) + selected: Final = decision(CacheContext(backend=backend), rules) + match selected: + case Decision.PYTHON: + return None + case Decision.RUST_WITH_FALLBACK | Decision.RUST_REQUIRED: + factory: Final = _RUNTIME.load() + if factory is None: + if selected is Decision.RUST_REQUIRED: + raise RuntimeError("Rust response cache runtime is unavailable") + return None + try: + return ResponseCacheRuntime(factory.from_cache(cache)) + except Exception as error: + exceptions: Final = native_exception_types() + if exceptions is None or not isinstance(error, exceptions[0]): + raise + if selected is Decision.RUST_REQUIRED: + raise RuntimeError(f"Rust response cache runtime declined the cache: {error}") from error + return None + case _: + assert_never(selected) + + +def _duration(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + duration: Final = float(value) + return duration if math.isfinite(duration) and duration >= 0 else None + + +def _string_mapping(value: object) -> Mapping[str, object]: + if not isinstance(value, Mapping): + return {} + source: Final = cast(Mapping[object, object], value) + return {key: item for key, item in source.items() if isinstance(key, str)} diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index 1fcde1bf555..cf02a33eab6 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -8,7 +8,7 @@ from typing_extensions import assert_never from litellm.exceptions import APIError from litellm.rust_bridge.bindings import NativeBinding, native_exception_types -from litellm.rust_bridge.catalog import RULES, Context, Rules, decision +from litellm.rust_bridge.catalog import RouteContext, Rules, decision from litellm.rust_bridge.configuration import Decision from litellm.rust_bridge.response_metadata import mark_rust_response @@ -42,14 +42,14 @@ class BridgeErrorContext: def run( - context: Context, + context: RouteContext, *, binding: NativeBinding[NativeT], native: Callable[[NativeT], ResultT], python: Callable[[], ResultT], rules: Rules | None = None, ) -> ResultT: - selected: Final = decision(context, RULES if rules is None else rules) + selected: Final = decision(context, rules) match selected: case Decision.PYTHON: return python() @@ -70,14 +70,14 @@ def run( async def arun( - context: Context, + context: RouteContext, *, binding: NativeBinding[NativeT], native: Callable[[NativeT], Awaitable[ResultT]], python: Callable[[], Awaitable[ResultT]], rules: Rules | None = None, ) -> ResultT: - selected: Final = decision(context, RULES if rules is None else rules) + selected: Final = decision(context, rules) match selected: case Decision.PYTHON: return await python() @@ -101,7 +101,7 @@ def _identity(value: ResultT) -> ResultT: return value -def _error_context(context: Context) -> BridgeErrorContext: +def _error_context(context: RouteContext) -> BridgeErrorContext: return BridgeErrorContext(route=context.route.value, provider=context.provider or "", model=context.model or "") diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 9a5cf49f298..866f2fce989 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Final @dataclass(frozen=True, slots=True) @@ -35,6 +36,28 @@ class SecretManager: readable: bool +@dataclass(frozen=True, slots=True) +class SecretManagerBinding: + system: object + access_mode: object + hosted_keys: object + primary_secret_name: object + store_virtual_keys: object + prefix_for_stored_virtual_keys: object + kms_key_id: object + custom_secret_manager: object + aws_region_name: object + aws_role_name: object + aws_session_name: object + aws_external_id: object + aws_profile_name: object + aws_web_identity_token: object + aws_sts_endpoint: object + replica_regions: object + client: object + settings_object: object + + def warn(message: str) -> None: from litellm._logging import verbose_logger @@ -49,6 +72,42 @@ def secret_manager() -> SecretManager: return SecretManager(readable=_should_read_secret_from_secret_manager()) +def secret_manager_binding() -> SecretManagerBinding: + import litellm + from litellm.types.secret_managers.main import KeyManagementSettings + + configured_system: Final = ( + litellm._key_management_system # pyright: ignore[reportPrivateUsage] # canonical key management globals are private + ) + configured_settings: Final = ( + litellm._key_management_settings # pyright: ignore[reportPrivateUsage] # canonical key management globals are private + ) + settings: Final = configured_settings or KeyManagementSettings() + system: Final = ( + configured_system.value if litellm.secret_manager_client is not None and configured_system is not None else None + ) + return SecretManagerBinding( + system=system, + access_mode=settings.access_mode, + hosted_keys=settings.hosted_keys, + primary_secret_name=settings.primary_secret_name, + store_virtual_keys=settings.store_virtual_keys, + prefix_for_stored_virtual_keys=settings.prefix_for_stored_virtual_keys, + kms_key_id=settings.kms_key_id, + custom_secret_manager=settings.custom_secret_manager, + aws_region_name=settings.aws_region_name, + aws_role_name=settings.aws_role_name, + aws_session_name=settings.aws_session_name, + aws_external_id=settings.aws_external_id, + aws_profile_name=settings.aws_profile_name, + aws_web_identity_token=settings.aws_web_identity_token, + aws_sts_endpoint=settings.aws_sts_endpoint, + replica_regions=settings.replica_regions, + client=litellm.secret_manager_client, + settings_object=configured_settings, + ) + + def provider_defaults() -> ProviderDefaults: import litellm diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index d36234f56c1..250ad18d44c 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -5,18 +5,19 @@ from __future__ import annotations from collections.abc import Awaitable from dataclasses import dataclass from functools import lru_cache -from typing import Final, Literal, Protocol, cast # noqa: TID251 # native extension exposes untyped callables +from typing import TYPE_CHECKING, Final, Literal, Protocol, cast # noqa: TID251 # PyO3 binding validation from pydantic import TypeAdapter +from typing_extensions import assert_never import litellm -from litellm._logging import verbose_logger -from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file -from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting +from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding_name, uses_legacy_message_accounting +from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.configuration import rust_enabled -from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt -from litellm.utils import claude_json_str, huggingface_tokenizer_kind +from litellm.utils import huggingface_tokenizer_kind + +if TYPE_CHECKING: + from litellm.rust_bridge._native import Tokenizer as NativeTokenizer RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"] @@ -27,13 +28,7 @@ class RustTokenCounter(Protocol): class RustTokenCounterFactory(Protocol): - def __call__(self, tokenizer_json: str) -> RustTokenCounter: - raise NotImplementedError - - def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter: - raise NotImplementedError - - def from_o200k_ranks(self, rank_file: str) -> RustTokenCounter: + def from_tokenizer(self, tokenizer: NativeTokenizer, fast: bool = False) -> RustTokenCounter: raise NotImplementedError @@ -51,7 +46,7 @@ def _as_factory(value: object) -> RustTokenCounterFactory | None: cast( # cast-ok: native extension protocol is runtime-defined RustTokenCounterFactory, value ) - if callable(value) + if callable(getattr(value, "from_tokenizer", None)) else None ) @@ -73,7 +68,7 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: return "anthropic" if kind is not None or uses_legacy_message_accounting(model): return None - match openai_tokenizer_encoding(model).name: + match openai_tokenizer_encoding_name(model): case "cl100k_base": return "cl100k_base" case "o200k_base": @@ -84,31 +79,26 @@ def rust_tokenizer(model: str) -> RustTokenizer | None: @lru_cache(maxsize=4) def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter: + return factory.from_tokenizer(_native_tokenizer(tokenizer)) + + +def _native_tokenizer(tokenizer: RustTokenizer) -> NativeTokenizer: match tokenizer: case "anthropic": - return factory(claude_json_str) - case "cl100k_base": - return factory.from_cl100k_ranks(cl100k_base_rank_file()) - case "o200k_base": - return factory.from_o200k_ranks(o200k_base_rank_file()) + native = tokenizer_dispatch.native_anthropic() + case "cl100k_base" | "o200k_base": + native = tokenizer_dispatch.native_encoding(tokenizer) + case _: + assert_never(tokenizer) + if native is None: + raise RuntimeError(f"native {tokenizer} tokenizer is unavailable") + return native -async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None: - if not rust_enabled(): - return None - factory: Final = TOKEN_COUNTER.load() - if factory is None: - return None - try: - attempt: Final = await aattempt( - native_call=lambda: _counter(factory, tokenizer).acount_request(body), - adapt=_INPUT_TOKEN_COUNT.validate_python, - context=BridgeErrorContext(route="token_counter", provider=tokenizer, model=""), - ) - except (RuntimeError, ValueError) as error: - verbose_logger.debug("Rust token counter (%s) failed, counting in Python: %s", tokenizer, error) - return None - if not isinstance(attempt, RustHandled): - return None - verbose_logger.debug("Rust token counter (%s) counted %d input tokens", tokenizer, attempt.value.input_tokens) - return attempt.value +async def native_count(factory: RustTokenCounterFactory, tokenizer: RustTokenizer, body: bytes) -> InputTokenCount: + """One native count, validated into the public shape. + + ``RustBridgeDeclined`` and upstream errors propagate so the caller's route + runner can map them onto its fallback policy; other failures (RuntimeError, + ValueError) propagate as-is.""" + return _INPUT_TOKEN_COUNT.validate_python(await _counter(factory, tokenizer).acount_request(body)) diff --git a/litellm/rust_bridge/tokenizer.py b/litellm/rust_bridge/tokenizer.py new file mode 100644 index 00000000000..5ed89813620 --- /dev/null +++ b/litellm/rust_bridge/tokenizer.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # native class is validated at the binding boundary + +import tiktoken +from tokenizers import Tokenizer as PythonHuggingFaceTokenizer + +from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace, HuggingFaceTokenizer, OpenAIEncoding +from litellm.rust_bridge import runtime +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.catalog import Route, RouteContext + +if TYPE_CHECKING: + from litellm.rust_bridge._native import Tokenizer as NativeTokenizer + + +def _as_factory(value: object) -> type[NativeTokenizer] | None: + return ( + cast(type["NativeTokenizer"], value) # cast-ok: PyO3 class validated at the native boundary + if isinstance(value, type) + else None + ) + + +TOKENIZER: Final = NativeBinding("Tokenizer", validate=_as_factory) + +# The catalog contexts the tokenizer factories dispatch on. Callers that cache a tokenizer per +# backend key their cache on `decision(...)` of the same context, so key and dispatch agree. +TIKTOKEN_CONTEXT: Final = RouteContext(Route.TOKENIZER, provider="tiktoken") +HUGGINGFACE_CONTEXT: Final = RouteContext(Route.TOKENIZER, provider="huggingface") + + +@lru_cache(maxsize=8) +def _native_tiktoken(factory: type[NativeTokenizer], name: str) -> NativeTokenizer: + return factory.from_tiktoken(name) + + +@lru_cache(maxsize=1) +def _native_anthropic(factory: type[NativeTokenizer]) -> NativeTokenizer: + from litellm.utils import claude_json_str + + return factory.from_json(claude_json_str) + + +@lru_cache(maxsize=8) +def _native_encoding(factory: type[NativeTokenizer], name: str) -> OpenAIEncoding: + return OpenAIEncoding.wrap(_native_tiktoken(factory, name)) + + +def native_encoding(name: str) -> NativeTokenizer | None: + """The native tiktoken encoding behind `get_encoding(name)`, for a Rust route that counts + with the same loaded model; `None` without the extension.""" + factory: Final = TOKENIZER.load() + return None if factory is None else _native_tiktoken(factory, name) + + +def native_anthropic() -> NativeTokenizer | None: + """The native packaged Anthropic tokenizer behind `anthropic()`, parsed once per process.""" + factory: Final = TOKENIZER.load() + return None if factory is None else _native_anthropic(factory) + + +def _python_encoding(name: str) -> tiktoken.Encoding: + from litellm.litellm_core_utils.default_encoding import encoding + + return encoding if name == encoding.name else tiktoken.get_encoding(name) + + +def get_encoding(name: str) -> Encoding: + return runtime.run( + TIKTOKEN_CONTEXT, + binding=TOKENIZER, + native=lambda factory: _native_encoding(factory, name), + python=lambda: _python_encoding(name), + ) + + +def anthropic() -> HuggingFace: + """The packaged Anthropic tokenizer on the selected backend.""" + from litellm.utils import claude_json_str + + return runtime.run( + HUGGINGFACE_CONTEXT, + binding=TOKENIZER, + native=lambda factory: HuggingFaceTokenizer(_native_anthropic(factory)), + python=lambda: PythonHuggingFaceTokenizer.from_str(claude_json_str), + ) + + +def from_str(json: str) -> HuggingFace: + return runtime.run( + HUGGINGFACE_CONTEXT, + binding=TOKENIZER, + native=lambda factory: HuggingFaceTokenizer(factory.from_json(json)), + python=lambda: PythonHuggingFaceTokenizer.from_str(json), + ) + + +def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> HuggingFace: + return runtime.run( + HUGGINGFACE_CONTEXT, + binding=TOKENIZER, + native=lambda factory: HuggingFaceTokenizer( + factory.from_pretrained(identifier, revision=revision, token=token) + ), + python=lambda: PythonHuggingFaceTokenizer.from_pretrained(identifier, revision=revision, token=token), + ) diff --git a/litellm/types/agents.py b/litellm/types/agents.py index dbaaab62d86..7f8d8c6af66 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal -from pydantic import BaseModel, PrivateAttr, StrictInt +from pydantic import BaseModel, ConfigDict, PrivateAttr, StrictInt from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -189,6 +189,7 @@ class AgentConfig(TypedDict, total=False): session_rpm_limit: int | None static_headers: dict[str, str] | None extra_headers: list[str] | None + access_group_ids: ReadOnly[Sequence[str] | None] class PatchAgentRequest(TypedDict, total=False): @@ -202,6 +203,21 @@ class PatchAgentRequest(TypedDict, total=False): session_rpm_limit: int | None static_headers: dict[str, str] | None extra_headers: list[str] | None + access_group_ids: ReadOnly[Sequence[str] | None] + + +AGENT_CALLER_USER_ID_HEADER: Final = "x-litellm-user-id" +AGENT_CALLER_TEAM_ID_HEADER: Final = "x-litellm-team-id" + + +class AgentCaller(BaseModel): + """The user and team that invoked an agent, echoed back by the agent on its own proxy calls. + Only ever narrows what the agent's key may do.""" + + model_config = ConfigDict(frozen=True) + + user_id: str | None = None + team_id: str | None = None # Request/Response models for CRUD endpoints @@ -226,6 +242,7 @@ class AgentResponse(BaseModel): session_rpm_limit: int | None = None static_headers: dict[str, str] | None = None extra_headers: list[str] | None = None + access_group_ids: Sequence[str] | None = None keys: list[AgentKeySummary] | None = None search_score: float | None = None created_at: datetime | None = None diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index f4893e857d1..c929ee2ee79 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -131,6 +131,7 @@ EXCEPTION_STATUS: Final = "exception_status" EXCEPTION_CLASS: Final = "exception_class" RATE_LIMIT_CATEGORY: Final = "rate_limit_category" RATE_LIMIT_TYPE: Final = "rate_limit_type" +ZERO_COST_REASON_LABEL: Final = "reason" STATUS_CODE: Final = "status_code" EXCEPTION_LABELS: Final = [EXCEPTION_STATUS, EXCEPTION_CLASS] LATENCY_BUCKETS: Final = ( @@ -279,6 +280,7 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_guardrail_latency_seconds", "litellm_guardrail_errors_total", "litellm_guardrail_requests_total", + "litellm_zero_cost_requests_total", # Cache metrics "litellm_cache_hits_metric", "litellm_cache_misses_metric", @@ -590,6 +592,14 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.SERVICE_TIER.value, ] + litellm_zero_cost_requests_total = ( + UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.MODEL_ID.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + ZERO_COST_REASON_LABEL, + ) + litellm_input_tokens_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index bc44eb5b5b7..599b1a76249 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -110,6 +110,14 @@ FileTypes = ( EmbeddingInput = str | list[str] +class BinaryResponseSummary(TypedDict): + """What logging keeps of a binary response (speech audio, file content): size and media type, never the bytes.""" + + object: ReadOnly[Literal["binary"]] + content_type: ReadOnly[str | None] + num_bytes: ReadOnly[int] + + class HttpxBinaryResponseContent(_HttpxBinaryResponseContent): _hidden_params: dict @@ -117,6 +125,19 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent): super().__init__(response) self._hidden_params = {} # mutable-ok: mutable-dict contract shared with ModelResponse logging consumers + def logging_summary(self) -> BinaryResponseSummary: + return { + "object": "binary", + "content_type": self.response.headers.get("content-type"), + "num_bytes": self._num_bytes(), + } + + def _num_bytes(self) -> int: + try: + return len(self.response.content) + except httpx.ResponseNotRead: + return self.response.num_bytes_downloaded + def set_response_cost(self, response_cost: float | None) -> None: if response_cost is None: self._hidden_params.pop("response_cost", None) diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index e47acf9d68b..619001a5791 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -28,6 +28,7 @@ class EndpointType(str, Enum): GEMINI = "gemini" ANTHROPIC = "anthropic" OPENAI = "openai" + TINYFISH = "tinyfish" GENERIC = "generic" diff --git a/litellm/types/passthrough_endpoints/tinyfish.py b/litellm/types/passthrough_endpoints/tinyfish.py new file mode 100644 index 00000000000..365eaef77a6 --- /dev/null +++ b/litellm/types/passthrough_endpoints/tinyfish.py @@ -0,0 +1,55 @@ +from typing import Final + +from typing_extensions import ReadOnly, TypedDict + +TINYFISH_AGENT_DEFAULT_API_BASE: Final = "https://agent.tinyfish.ai" +TINYFISH_AGENT_DOCS_URL: Final = "https://docs.tinyfish.ai/agent-api" +# TinyFish's published Agent API rate (USD per run step); override with env TINYFISH_COST_PER_STEP +TINYFISH_DEFAULT_COST_PER_STEP: Final = 0.016 +TINYFISH_MODEL_NAME: Final = "tinyfish/automation-run" +TINYFISH_POLLING_INTERVAL_SECONDS: Final = 5.0 +# the Agent API caps runs at 1200s but queue wait extends wall time, so billing polls with generous headroom +TINYFISH_MAX_POLLING_SECONDS: Final = 3600.0 +# at the 5s interval this tolerates a ~60s upstream outage before abandoning the charge +TINYFISH_MAX_CONSECUTIVE_POLL_FAILURES: Final = 12 + +TINYFISH_TERMINAL_RUN_STATUSES: Final = frozenset({"COMPLETED", "FAILED", "CANCELLED"}) + +# these fields use the shared account's saved logins/vault, so they 403 unless TINYFISH_ALLOW_AUTHENTICATED_RUNS=true +TINYFISH_AUTHENTICATED_RUN_FIELDS: Final = frozenset({"use_profile", "profile_id", "use_vault", "credential_item_ids"}) + +# litellm's pass-through envelope controls; rejected here or custom_body smuggles past the field gate and +# a caller stream flag flips the billing mode away from what the endpoint dictates +TINYFISH_REJECTED_ENVELOPE_FIELDS: Final = frozenset({"custom_body", "stream", "query_params"}) + +# covers the upstream 1200s max run duration plus response headroom for blocking runs +TINYFISH_PASSTHROUGH_TIMEOUT_SECONDS: Final = 1500.0 + +_RUN_SUBMIT_PATHS: Final = frozenset( + {("v1", "automation", "run"), ("v1", "automation", "run-async"), ("v1", "automation", "run-sse")} +) + + +class TinyfishRun(TypedDict, total=False): + """Run objects are null-heavy until terminal, so every field must tolerate None.""" + + run_id: ReadOnly[str | None] + status: ReadOnly[str | None] + num_of_steps: ReadOnly[int | None] + result: ReadOnly[object] + # left untyped on purpose: a strict error shape would fail whole-run validation on upstream drift and drop the charge + error: ReadOnly[object] + type: ReadOnly[str | None] + + +def is_allowed_tinyfish_endpoint(method: str, path: str) -> bool: + """The host also serves vault/wallet/profile management under the same key, so only run endpoints forward.""" + segments: Final = tuple(path.split("/")[1:]) + if not path.startswith("/") or any(segment in ("", ".", "..") for segment in segments): + return False + if method == "POST" and segments in _RUN_SUBMIT_PATHS: + return True + # no GET /v1/runs listing: run ids are unguessable, so blocking the list keeps teams out of each other's runs + if method == "GET" and len(segments) == 3 and segments[:2] == ("v1", "runs"): + return True + return method == "POST" and len(segments) == 4 and segments[:2] == ("v1", "runs") and segments[3] == "cancel" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e1d43b7fccb..65d66677b27 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1,7 +1,7 @@ import json import re import time -from collections.abc import Mapping, Sequence +from collections.abc import Collection, Mapping, Sequence from enum import Enum from types import MappingProxyType from typing import ( @@ -36,12 +36,14 @@ from pydantic import ( BaseModel, ConfigDict, Field, + FieldSerializationInfo, JsonValue, PrivateAttr, SkipValidation, field_serializer, field_validator, ) +from pydantic.main import IncEx from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm._logging import verbose_logger @@ -80,7 +82,30 @@ from .llms.openai import ( ) from .rerank import RerankResponse as RerankResponse + +def _nested_selector( + selector: IncEx | None, + index: int, + count: int, + is_include: bool, +) -> tuple[bool, IncEx | None]: + if selector is None: + return True, None + if isinstance(selector, Mapping): + value: Final = selector.get(index, selector.get(index - count, selector.get("__all__"))) + keep: Final = value is not None if is_include else value is not True + per_item_selector: Final = None if value is True or value is None else value + return keep, per_item_selector + if isinstance(selector, Collection) and not isinstance(selector, (str, bytes)): + if all(isinstance(item, int) for item in selector): + addressed: Final = index in selector or index - count in selector + return (addressed if is_include else not addressed), None + return True, selector + + if TYPE_CHECKING: + from litellm.litellm_core_utils.tokenizer import Tokenizer + from .vector_stores import VectorStoreSearchResponse else: VectorStoreSearchResponse = Any @@ -333,6 +358,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_second_768p: ReadOnly[float | None] output_cost_per_second_2k: ReadOnly[float | None] output_cost_per_second_4k: ReadOnly[float | None] + output_cost_per_image_512: ReadOnly[float | None] + output_cost_per_image_1024: ReadOnly[float | None] + output_cost_per_image_1536: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models ocr_cost_per_page_batches: ReadOnly[float | None] ocr_cost_per_credit: float | None # for OCR models priced by credit @@ -459,6 +487,8 @@ class CallTypes(str, Enum): ######################################################### create_video = "create_video" acreate_video = "acreate_video" + video_generation = "video_generation" + avideo_generation = "avideo_generation" avideo_retrieve = "avideo_retrieve" video_retrieve = "video_retrieve" avideo_content = "avideo_content" @@ -2555,8 +2585,35 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): model_config = ConfigDict(extra="allow", protected_namespaces=()) @field_serializer("data") - def _serialize_image_data(self, data: Sequence[OpenAIImage] | None) -> Sequence[Mapping[str, object]] | None: - return None if data is None else [image.model_dump() for image in data] + def _serialize_image_data( + self, + data: Sequence[OpenAIImage] | None, + info: FieldSerializationInfo, + ) -> Sequence[Mapping[str, object]] | None: + if data is None: + return None + include: Final = info.include + exclude: Final = info.exclude + + def _serialize_image(index: int, image: OpenAIImage) -> Mapping[str, object] | None: + include_keep, include_selector = _nested_selector(include, index, len(data), is_include=True) + exclude_keep, exclude_selector = _nested_selector(exclude, index, len(data), is_include=False) + if not include_keep or not exclude_keep: + return None + return image.model_dump( + mode=info.mode, + include=include_selector, + exclude=exclude_selector, + context=info.context, + exclude_none=info.exclude_none, + exclude_unset=info.exclude_unset, + exclude_defaults=info.exclude_defaults, + round_trip=info.round_trip, + by_alias=info.by_alias, + ) + + serialized_images: Final = tuple(_serialize_image(index, image) for index, image in enumerate(data)) + return [image for image in serialized_images if image is not None] def __init__( self, @@ -3154,6 +3211,15 @@ class StandardLoggingModelCostFailureDebugInformation(TypedDict, total=False): custom_pricing: bool | None +ZeroCostReason = Literal["missing_pricing_key", "pricing_not_applied", "cost_calculation_error"] + + +class StandardLoggingZeroCostDiagnostic(TypedDict): + reason: ReadOnly[ZeroCostReason] + pricing_model: ReadOnly[str] + missing_pricing_keys: ReadOnly[tuple[str, ...]] + + class StandardLoggingPayloadErrorInformation(TypedDict, total=False): error_code: str | None error_class: str | None @@ -3472,6 +3538,7 @@ class StandardLoggingPayload(ClassifierAudit): autorouter_savings_estimate: ReadOnly[Mapping[str, JsonValue] | None] autorouter_baseline_observation: ReadOnly[str | None] response_cost_failure_debug_info: StandardLoggingModelCostFailureDebugInformation | None + zero_cost_diagnostic: NotRequired[ReadOnly[StandardLoggingZeroCostDiagnostic | None]] status: StandardLoggingPayloadStatus status_fields: StandardLoggingPayloadStatusFields custom_llm_provider: str | None @@ -3620,6 +3687,9 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_second_768p: float | None = None output_cost_per_second_2k: float | None = None output_cost_per_second_4k: float | None = None + output_cost_per_image_512: float | None = None + output_cost_per_image_1024: float | None = None + output_cost_per_image_1536: float | None = None input_cost_per_pixel: float | None = None output_cost_per_pixel: float | None = None @@ -3790,6 +3860,24 @@ def echoed_cost_map_pricing_fields(model_info: Mapping[str, Any]) -> tuple[str, return tuple(sorted(k for k in model_info if is_server_derived_pricing_key(k))) +def echoed_cost_map_fields( + model_info: Mapping[str, object], *cost_map_entries: Mapping[str, object] +) -> tuple[str, ...]: + """Fields a ``/model/info`` echo copied from the cost map unchanged. + + Only ``litellm.get_model_info`` emits ``key``, so a blob carrying it is an echo of that + response. Anything in it that still equals a resolved cost-map entry is a display value + nobody typed; a value the operator edited differs from every entry and stays a real override. + Callers pass both the live entry, which the router rewrites with each deployment's own + overrides, and the catalog entry as loaded, so a reset to the catalog value reads as an echo either way. + """ + if COST_MAP_LOOKUP_KEY not in model_info: + return () + return tuple( + sorted(k for k, v in model_info.items() if any(k in entry and entry[k] == v for entry in cost_map_entries)) + ) + + def pricing_override_fields(*sources: Mapping[str, Any]) -> tuple[str, ...]: return tuple( sorted( @@ -4320,7 +4408,7 @@ class ProviderSpecificHeader(TypedDict): class SelectTokenizerResponse(TypedDict): type: Literal["openai_tokenizer", "huggingface_tokenizer"] - tokenizer: Any + tokenizer: ReadOnly["Tokenizer"] class LiteLLMFineTuningJob(FineTuningJob): diff --git a/litellm/utils.py b/litellm/utils.py index 709f3f6d1dd..e088a5988c8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -40,14 +40,11 @@ from types import MappingProxyType import dotenv import httpx import openai -import tiktoken from httpx import Proxy from httpx._utils import get_environment_proxies from openai.lib import _parsing, _pydantic from openai.types.chat.completion_create_params import ResponseFormat from pydantic import BaseModel -from tiktoken import Encoding -from tokenizers import Tokenizer import litellm import litellm.litellm_core_utils @@ -81,12 +78,20 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + bind_budget_reservation_to_callbacks, + normalize_drop_params, + unbind_budget_reservation_from_callbacks, +) from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload +from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace, strip_special_tokens +from litellm.rust_bridge import tokenizer as tokenizer_dispatch +from litellm.rust_bridge.catalog import decision +from litellm.rust_bridge.configuration import Decision _CachingHandlerResponse = None _LLMCachingHandler = None @@ -285,6 +290,8 @@ import importlib.metadata from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args +from typing_extensions import assert_never + from litellm import utils as litellm_utils # These are lazy loaded via __getattr__ @@ -1880,6 +1887,8 @@ def client(original_function): # Type assertion: logging_obj is guaranteed to be non-None after function_setup assert logging_obj is not None, "logging_obj should not be None after function_setup" + if not _is_litellm_internal_call: + bind_budget_reservation_to_callbacks(logging_obj.litellm_params) kwargs["litellm_logging_obj"] = logging_obj modified_kwargs: Final = await async_pre_call_deployment_hook(kwargs, call_type) @@ -2081,6 +2090,7 @@ def client(original_function): # the failure hook ran, so a slow callback doesn't inflate the reported duration. end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with if logging_obj and not _is_litellm_internal_call: + unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) try: logging_obj.failure_handler( e, traceback_exception, start_time, end_time @@ -2247,17 +2257,27 @@ def _select_tokenizer(model: str, custom_tokenizer: CustomHuggingfaceTokenizer | identifier=custom_tokenizer["identifier"], revision=custom_tokenizer["revision"], auth_token=custom_tokenizer["auth_token"], + backend=_huggingface_tokenizer_backend(), ) return _select_tokenizer_helper(model=model) +def _huggingface_tokenizer_backend() -> Decision: + """The backend `tokenizer_dispatch.from_str` / `from_pretrained` will select right now. + + Cached HuggingFace tokenizers are keyed on it, so flipping `LITELLM_RUST` or + `litellm.rust(...)` reaches a fresh object instead of the other backend's.""" + return decision(tokenizer_dispatch.HUGGINGFACE_CONTEXT) + + @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) -def _select_custom_tokenizer_helper(identifier: str, revision: str, auth_token: str | None) -> SelectTokenizerResponse: +def _select_custom_tokenizer_helper( + identifier: str, revision: str, auth_token: str | None, backend: Decision +) -> SelectTokenizerResponse: verbose_logger.debug("Loading custom HuggingFace tokenizer %s (revision %s)", identifier, revision) return create_pretrained_tokenizer(identifier=identifier, revision=revision, auth_token=auth_token) -@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse: if litellm.disable_hf_tokenizer_download is True: return _return_openai_tokenizer(model) @@ -2267,6 +2287,10 @@ def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse: if result is not None: return result except Exception as e: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking + + if isinstance(e, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): + raise verbose_logger.debug("Error selecting tokenizer: %s", e) # default - tiktoken @@ -2301,19 +2325,26 @@ def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None: kind: Final = huggingface_tokenizer_kind(model) if kind is None: return None - return {"type": "huggingface_tokenizer", "tokenizer": _load_huggingface_tokenizer(kind)} + return { + "type": "huggingface_tokenizer", + "tokenizer": _load_huggingface_tokenizer(kind, _huggingface_tokenizer_backend()), + } -def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind) -> Tokenizer: +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decision) -> HuggingFace: + """One tokenizer per kind and backend; `backend` is the cache key, the dispatch re-derives it.""" match kind: case "cohere": - return Tokenizer.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer") + return tokenizer_dispatch.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer") case "anthropic": - return Tokenizer.from_str(claude_json_str) + return tokenizer_dispatch.anthropic() case "llama2": - return Tokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") + return tokenizer_dispatch.from_pretrained("hf-internal-testing/llama-tokenizer") case "llama3": - return Tokenizer.from_pretrained("Xenova/llama-3-tokenizer") + return tokenizer_dispatch.from_pretrained("Xenova/llama-3-tokenizer") + case _: + assert_never(kind) def encode(model="", text="", custom_tokenizer: dict | None = None): @@ -2329,15 +2360,13 @@ def encode(model="", text="", custom_tokenizer: dict | None = None): enc: The encoded text. """ tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model=model) - if isinstance(tokenizer_json["tokenizer"], Encoding): - enc = tokenizer_json["tokenizer"].encode(text, disallowed_special=()) - else: - enc = tokenizer_json["tokenizer"].encode(text) - # Normalize: HuggingFace Tokenizer.encode() returns an Encoding object; - # extract .ids so the return type is always List[int]. - if hasattr(enc, "ids"): - return enc.ids - return enc + if tokenizer_json["type"] == "openai_tokenizer": + openai_tokenizer: Final = cast( # cast-ok: [LIT006] caller's explicit type tag selects this interface + Encoding, tokenizer_json["tokenizer"] + ) + return openai_tokenizer.encode(text, disallowed_special=()) + encoded: Final = tokenizer_json["tokenizer"].encode(text) + return encoded.ids if hasattr(encoded, "ids") else encoded def decode( @@ -2356,26 +2385,12 @@ def decode( """ tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model=model) if tokenizer_json["type"] == "huggingface_tokenizer": - if skip_special_tokens: - tokens = _strip_huggingface_special_token_ids(tokenizer_json["tokenizer"], tokens) - dec = tokenizer_json["tokenizer"].decode(tokens, skip_special_tokens=skip_special_tokens) - return dec - dec = tokenizer_json["tokenizer"].decode(tokens) - return dec - - -def _strip_huggingface_special_token_ids(tokenizer: Tokenizer, tokens: Sequence[int]) -> Sequence[int]: - try: - added_tokens_decoder: Final = tokenizer.get_added_tokens_decoder() - except Exception: - return tokens - - special_token_ids: Final = { - token_id for token_id, added_token in added_tokens_decoder.items() if getattr(added_token, "special", False) - } - if not special_token_ids: - return tokens - return [token for token in tokens if token not in special_token_ids] + ids: Final = strip_special_tokens(tokenizer_json["tokenizer"], tokens) if skip_special_tokens else tokens + hf_tokenizer: Final = cast( # cast-ok: [LIT006] caller's explicit type tag selects this interface + HuggingFace, tokenizer_json["tokenizer"] + ) + return hf_tokenizer.decode(ids, skip_special_tokens=skip_special_tokens) + return tokenizer_json["tokenizer"].decode(tokens) def create_pretrained_tokenizer(identifier: str, revision="main", auth_token: str | None = None): @@ -2391,7 +2406,7 @@ def create_pretrained_tokenizer(identifier: str, revision="main", auth_token: st dict: A dictionary with the tokenizer and its type. """ - tokenizer: Final = Tokenizer.from_pretrained(identifier, revision=revision, token=auth_token) + tokenizer: Final = tokenizer_dispatch.from_pretrained(identifier, revision=revision, token=auth_token) return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} @@ -2406,7 +2421,7 @@ def create_tokenizer(json: str): dict: A dictionary with the tokenizer and its type. """ - tokenizer: Final = Tokenizer.from_str(json) + tokenizer: Final = tokenizer_dispatch.from_str(json) return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} @@ -8340,6 +8355,7 @@ class ProviderConfigManager: False, ), LlmProviders.EDENAI: (litellm.EdenAIChatConfig, False), + LlmProviders.FAL_AI: (litellm.FalAIChatConfig, False), LlmProviders.COMETAPI: (lambda: litellm.CometAPIConfig(), False), LlmProviders.DATAROBOT: (lambda: litellm.DataRobotConfig(), False), LlmProviders.GEMINI: (lambda: litellm.GoogleAIStudioGeminiConfig(), False), @@ -9569,9 +9585,9 @@ class ProviderConfigManager: return BlackForestLabsImageEditConfig() elif LlmProviders.FAL_AI == provider: - from litellm.llms.fal_ai.image_edit import FalAIImageEditConfig + from litellm.llms.fal_ai.image_edit import get_fal_ai_image_edit_config - return FalAIImageEditConfig() + return get_fal_ai_image_edit_config(model) elif LlmProviders.AZURE_AI == provider: from litellm.llms.azure_ai.image_edit import get_azure_ai_image_edit_config diff --git a/migrations/Dockerfile b/migrations/Dockerfile index c6d1b0cc46e..f34940c0ce0 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -67,6 +67,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --python python3.13 +RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/ + COPY migrations/run.py /app/run.py # Pre-warm the Prisma binary cache so the Job pod doesn't reach the diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3ca9958ad0d..77cada25d25 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24966,6 +24966,52 @@ "/v1/images/generations" ] }, + "fal_ai/fal-ai/trellis": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.02, + "source": "https://fal.ai/models/fal-ai/trellis", + "metadata": { + "comment": "image-to-3D, returns a GLB mesh; served through the /fal_ai pass-through route" + } + }, + "fal_ai/fal-ai/trellis-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.3, + "output_cost_per_image_512": 0.25, + "output_cost_per_image_1024": 0.3, + "output_cost_per_image_1536": 0.35, + "source": "https://fal.ai/models/fal-ai/trellis-2", + "metadata": { + "comment": "image-to-3D, returns a GLB mesh; priced by the request's resolution field (default 1024); served through the /fal_ai pass-through route" + } + }, + "fal_ai/fal-ai/flux-lora-depth": { + "litellm_provider": "fal_ai", + "metadata": { + "notes": "fal bills fal-ai/flux-lora-depth at $0.035 per megapixel, rounding each image up to the nearest megapixel. The per-pixel rate is used when Fal reports the output size, and the flat per-image price prices the default 1 MP output like the sibling flux entries" + }, + "mode": "image_generation", + "output_cost_per_image": 0.035, + "output_cost_per_pixel": 3.337860107421875e-08, + "source": "https://fal.ai/models/fal-ai/flux-lora-depth", + "supported_endpoints": [ + "/v1/images/edits" + ] + }, + "fal_ai/fal-ai/moondream3-preview/query": { + "input_cost_per_token": 4e-07, + "litellm_provider": "fal_ai", + "mode": "chat", + "output_cost_per_token": 3.5e-06, + "source": "https://fal.ai/models/fal-ai/moondream3-preview/query", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_reasoning": true, + "supports_vision": true + }, "featherless_ai/featherless-ai/Qwerky-72B": { "litellm_provider": "featherless_ai", "max_input_tokens": 32768, @@ -43037,21 +43083,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.83746e-07, + "input_cost_per_token": 9.5526e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.767492e-06, + "output_cost_per_token": 1.91052e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.36455e-08, + "cache_read_input_token_cost": 7.9605e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -44812,26 +44858,6 @@ "max_tokens": 128000, "mode": "chat" }, - "openrouter/stealth/union-alpha": { - "deprecation_date": "2098-12-31", - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "openrouter", - "max_input_tokens": 262144, - "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", - "source": "https://openrouter.ai/api/v1/models", - "supports_audio_input": false, - "supports_function_calling": true, - "supports_pdf_input": false, - "supports_prompt_caching": false, - "supports_reasoning": false, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_vision": true, - "supports_web_search": false - }, "ovhcloud/DeepSeek-R1-Distill-Llama-70B": { "input_cost_per_token": 6.7e-07, "litellm_provider": "ovhcloud", @@ -54217,7 +54243,7 @@ "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54227,6 +54253,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true, "supported_endpoints": [ "/v1/responses" @@ -54241,7 +54268,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54251,6 +54278,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-0309-reasoning": { @@ -54262,7 +54290,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -54271,6 +54299,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, @@ -54283,7 +54312,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, @@ -54292,11 +54321,13 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.3": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "xai", @@ -54306,7 +54337,7 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54318,6 +54349,7 @@ "xai/grok-4.3-latest": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "litellm_provider": "xai", @@ -54327,7 +54359,7 @@ "mode": "chat", "output_cost_per_token": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54339,6 +54371,7 @@ "xai/grok-4.5": { "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54348,7 +54381,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54360,6 +54393,7 @@ "xai/grok-4.5-latest": { "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54369,7 +54403,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54381,6 +54415,7 @@ "xai/grok-build-latest": { "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54390,7 +54425,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/developers/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54402,6 +54437,7 @@ "xai/grok-4.6": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54411,7 +54447,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/developers/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54423,6 +54459,7 @@ "xai/grok-4.7": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_image_token": 2e-06, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -54432,7 +54469,7 @@ "mode": "chat", "output_cost_per_token": 6e-06, "output_cost_per_token_above_200k_tokens": 1.2e-05, - "source": "https://docs.x.ai/developers/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54450,7 +54487,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54460,7 +54497,8 @@ "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, "supports_vision": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_image_token": 1e-06 }, "xai/grok-code-fast-1": { "cache_read_input_token_cost": 2e-07, @@ -54471,7 +54509,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54481,7 +54519,8 @@ "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, "supports_vision": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_image_token": 1e-06 }, "xai/grok-code-fast-1-0825": { "cache_read_input_token_cost": 2e-07, @@ -54492,7 +54531,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -54502,7 +54541,8 @@ "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, "supports_vision": true, - "deprecation_date": "2026-05-15" + "deprecation_date": "2026-05-15", + "input_cost_per_image_token": 1e-06 }, "zai.glm-4.7": { "input_cost_per_token": 6e-07, @@ -61978,7 +62018,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, @@ -61987,6 +62027,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-0309": { @@ -61998,7 +62039,7 @@ "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": false, "supports_prompt_caching": true, "supports_reasoning": true, @@ -62008,6 +62049,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true, "supported_endpoints": [ "/v1/responses" @@ -62022,7 +62064,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 2e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -62030,6 +62072,7 @@ "input_cost_per_token_above_200k_tokens": 2e-06, "output_cost_per_token_above_200k_tokens": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1e-06, "supports_response_schema": true, "supports_vision": true }, @@ -65225,7 +65268,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -65234,6 +65277,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, @@ -65246,7 +65290,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -65255,6 +65299,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, @@ -65267,7 +65312,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -65276,6 +65321,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_prompt_caching": true, "supports_response_schema": true }, @@ -65519,7 +65565,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, @@ -65528,6 +65574,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-non-reasoning-latest": { @@ -65539,7 +65586,7 @@ "max_tokens": 1000000, "mode": "chat", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, @@ -65548,6 +65595,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent": { @@ -65559,7 +65607,7 @@ "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" ], @@ -65572,6 +65620,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "xai/grok-4.20-multi-agent-latest": { @@ -65583,7 +65632,7 @@ "max_tokens": 1000000, "mode": "responses", "output_cost_per_token": 2.5e-06, - "source": "https://docs.x.ai/docs/models", + "source": "https://api.x.ai/v1/language-models", "supported_endpoints": [ "/v1/responses" ], @@ -65596,6 +65645,7 @@ "input_cost_per_token_above_200k_tokens": 2.5e-06, "output_cost_per_token_above_200k_tokens": 5e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, "supports_response_schema": true }, "groq/qwen/qwen3.8-27b": { @@ -68252,9 +68302,9 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 9.1e-07, - "output_cost_per_token": 2.86e-06, - "cache_read_input_token_cost": 1.69e-07, + "input_cost_per_token": 8.4e-07, + "output_cost_per_token": 2.64e-06, + "cache_read_input_token_cost": 1.56e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -68941,9 +68991,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 5.544e-08, - "output_cost_per_token": 1.1088e-07, - "cache_read_input_token_cost": 1.1088e-08, + "input_cost_per_token": 8.8606e-08, + "output_cost_per_token": 1.77212e-07, + "cache_read_input_token_cost": 1.77212e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -72854,6 +72904,16 @@ "supports_reasoning": true, "supports_vision": true }, + "openrouter/typesafe/jev-1.13": { + "input_cost_per_token": 4.2e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 32000, + "max_output_tokens": 28800, + "max_tokens": 28800, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://openrouter.ai/typesafe/jev-1.13" + }, "typesafe/jev-1.13.0": { "input_cost_per_token": 4.2e-08, "litellm_provider": "typesafe", @@ -73272,14 +73332,14 @@ "supports_web_search": false }, "openrouter/~z-ai/glm-latest": { - "cache_read_input_token_cost": 1.69e-07, - "input_cost_per_token": 9.1e-07, + "cache_read_input_token_cost": 1.56e-07, + "input_cost_per_token": 8.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 2.86e-06, + "output_cost_per_token": 2.64e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -75083,6 +75143,7 @@ "supports_web_search": false }, "openrouter/nex-agi/nex-n2.5-mini:free": { + "deprecation_date": "2026-09-25", "input_cost_per_token": 0.0, "litellm_provider": "openrouter", "max_input_tokens": 262144, @@ -75102,6 +75163,7 @@ "supports_web_search": false }, "openrouter/nex-agi/nex-n2.5-pro:free": { + "deprecation_date": "2026-09-25", "input_cost_per_token": 0.0, "litellm_provider": "openrouter", "max_input_tokens": 262144, @@ -77044,5 +77106,467 @@ "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true + }, + "xai/grok-4.20-0309": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-0309": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-latest-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-latest-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-beta-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-0304": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-0304-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-0304-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-non-reasoning-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-experimental-beta-reasoning-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "xai/grok-4.20-multi-agent-beta-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "responses", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "xai/grok-4.20-multi-agent-experimental-beta-0304": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "responses", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "xai/grok-4.20-multi-agent-experimental-beta-latest": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "responses", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "xai/grok-4.20-non-reasoning-gv2": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_response_schema": true + }, + "xai/grok-4.20-reasoning-gv2": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://api.x.ai/v1/language-models", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_image_token": 1.25e-06, + "supports_prompt_caching": true, + "supports_response_schema": true + }, + "openrouter/nex-agi/nex-n2.5-mini": { + "cache_read_input_token_cost": 2.5e-09, + "input_cost_per_token": 2.5e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 235929, + "max_tokens": 235929, + "mode": "chat", + "output_cost_per_token": 1e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": false, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": false + }, + "openrouter/nex-agi/nex-n2.5-pro": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 7.5e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 235929, + "max_tokens": 235929, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": false, + "supports_pdf_input": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false, + "supports_vision": true, + "supports_web_search": false } } diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 0509516ac32..f3a4e614f59 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -586,6 +586,18 @@ "type": "number", "minimum": 0 }, + "output_cost_per_image_1024": { + "type": "number", + "minimum": 0 + }, + "output_cost_per_image_1536": { + "type": "number", + "minimum": 0 + }, + "output_cost_per_image_512": { + "type": "number", + "minimum": 0 + }, "output_cost_per_image_token": { "type": "number", "minimum": 0 diff --git a/pyproject.toml b/pyproject.toml index 95da93df41e..f447343ff33 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,13 +15,17 @@ dependencies = [ # When changing a floor, verify it installs + imports on every supported # Python with: `uv pip install --resolution=lowest-direct .` "fastuuid>=0.14.0,<1.0", + "filelock>=3.16.1,<4.0", "httpx[http2]>=0.28.0,<1.0", "openai>=2.20.0,<3.0.0", "python-dotenv>=1.0.0,<2.0", + "pyyaml>=6.0.3,<7.0", + "packaging>=24.0", + "importlib-metadata>=8.0.0,<9.0", "tiktoken>=0.8.0,<1.0; python_version < '3.14'", "tiktoken>=0.12.0,<1.0; python_version >= '3.14'", - "importlib-metadata>=8.0.0,<9.0", "tokenizers>=0.21.0,<1.0", + "huggingface-hub>=0.34.0,<2.0", "click>=8.0.0,<9.0", "jinja2>=3.1.6,<4.0", "aiohttp>=3.14.2,<4.0", @@ -189,6 +193,7 @@ litellm-proxy = "litellm.proxy.client.cli:litellm_proxy_cli" [dependency-groups] dev = [ + "numpy>=1.26.0,<3.0", "diff-cover==9.7.2", "hypothesis==6.165.10", "reportlab==5.0.1", @@ -288,6 +293,12 @@ healthcheck = [ "httpx==0.28.1", "pyyaml==6.0.3", ] +benchmarks = [ + "pytest==9.0.3", + "pytest-codspeed==4.3.0", + "mcp>=2.2.0,<3", + "a2a-sdk==1.1.0", +] [build-system] requires = ["maturin==1.15.0"] diff --git a/schema.prisma b/schema.prisma index f802a141d41..c55456b2a40 100644 --- a/schema.prisma +++ b/schema.prisma @@ -73,6 +73,7 @@ model LiteLLM_AgentsTable { static_headers Json? @default("{}") extra_headers String[] @default([]) agent_access_groups String[] @default([]) + access_group_ids String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index ee5b42fe0b7..8f40ed6dfb7 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -18,6 +18,15 @@ longer signal it. - **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them - **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement +- `litellm_jwt_key_mapping` accepts `token_id` as an alternative to `key`, so a + mapping can name its virtual key by the SHA-256 hash the proxy stores instead + of by the plaintext. Exactly one of the two is required. This is what lets a + mapping reference a key managed in the same configuration + (`token_id = litellm_key.foo.token_id`), which `key` cannot do, because + `litellm_key` marks its generated key write-only and referencing it fails at + plan time. `POST /jwt/key/mapping/new` and `/jwt/key/mapping/update` gained a + matching `token` field, validated as 64 lowercase hex characters so a + plaintext key sent by mistake is rejected instead of hashed twice - **jwt_key_mapping**: New `litellm_jwt_key_mapping` resource for the proxy's JWT to virtual key mappings, so JWT clients identified by a claim (`client_id`, `azp`, `sub`) map to virtual keys and inherit their models, budgets and rate limits. Supports `description` and `is_active`, rotating the mapped key in place, and forces replacement when the claim name or value changes - **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it - **user**: New `litellm_user` resource and `litellm_user` / `litellm_users` data sources for managing internal users diff --git a/terraform/provider/docs/resources/jwt_key_mapping.md b/terraform/provider/docs/resources/jwt_key_mapping.md index fbc30947113..726c4b16021 100644 --- a/terraform/provider/docs/resources/jwt_key_mapping.md +++ b/terraform/provider/docs/resources/jwt_key_mapping.md @@ -65,7 +65,8 @@ resource "litellm_jwt_key_mapping" "developer" { - `jwt_claim_name` - (Required, ForceNew) Name of the JWT claim to match on, for example `client_id`, `azp` or `sub`. Must match `virtual_key_claim_field` in the proxy JWT config - `jwt_claim_value` - (Required, ForceNew) Value of the claim identifying the JWT client. Unique together with `jwt_claim_name`, so a second mapping for the same pair fails with a 409 -- `key` - (Required, Sensitive) The virtual key this claim value maps to. It has to exist already, otherwise the proxy rejects the mapping with `The provided key does not match an existing virtual key` +- `key` - (Optional, Sensitive) The virtual key this claim value maps to, as plaintext. It has to exist already, otherwise the proxy rejects the mapping with `The provided key does not match an existing virtual key`. Exactly one of `key` or `token_id` is required. `litellm_key` marks its generated `key` write-only, so this cannot reference a `litellm_key` resource -- use `token_id` for that, or supply the plaintext from a variable or a secret manager +- `token_id` - (Optional) The SHA-256 hash of the virtual key this claim value maps to, which is what the proxy stores. `litellm_key` exposes it as `token_id`, so unlike `key` it can be referenced directly from a `litellm_key` resource. Not a secret, so it is not marked sensitive. Exactly one of `key` or `token_id` is required - `description` - (Optional) Description of the mapping - `is_active` - (Optional) Whether the mapping is active. Inactive mappings are ignored during JWT auth. Defaults to `true` diff --git a/terraform/provider/litellm/resource_jwt_key_mapping.go b/terraform/provider/litellm/resource_jwt_key_mapping.go index e606e865737..ea968821527 100644 --- a/terraform/provider/litellm/resource_jwt_key_mapping.go +++ b/terraform/provider/litellm/resource_jwt_key_mapping.go @@ -29,10 +29,17 @@ func resourceLiteLLMJWTKeyMapping() *schema.Resource { Description: "Value of the claim identifying the JWT client. Unique together with jwt_claim_name", }, "key": { - Type: schema.TypeString, - Required: true, - Sensitive: true, - Description: "The virtual key this claim value maps to. The proxy stores only a hash of it and never returns it, so drift on this attribute cannot be detected and Terraform tracks the configured value", + Type: schema.TypeString, + Optional: true, + Sensitive: true, + ExactlyOneOf: []string{"key", "token_id"}, + Description: "The virtual key this claim value maps to, as plaintext. The proxy stores only a hash of it and never returns it, so drift on this attribute cannot be detected and Terraform tracks the configured value. litellm_key marks its generated key write-only, so this cannot reference a litellm_key resource; use token_id for that, or supply the plaintext from a variable or a secret manager", + }, + "token_id": { + Type: schema.TypeString, + Optional: true, + ExactlyOneOf: []string{"key", "token_id"}, + Description: "The SHA-256 hash of the virtual key this claim value maps to, which is what the proxy stores. litellm_key exposes it as token_id, so unlike key it can be referenced directly from a litellm_key resource. Not a secret, so it is not marked sensitive", }, "description": { Type: schema.TypeString, diff --git a/terraform/provider/litellm/resource_jwt_key_mapping_crud.go b/terraform/provider/litellm/resource_jwt_key_mapping_crud.go index 725235305f6..6c3883de2e1 100644 --- a/terraform/provider/litellm/resource_jwt_key_mapping_crud.go +++ b/terraform/provider/litellm/resource_jwt_key_mapping_crud.go @@ -19,6 +19,7 @@ func resourceLiteLLMJWTKeyMappingCreate(d *schema.ResourceData, m interface{}) e JWTClaimName: d.Get("jwt_claim_name").(string), JWTClaimValue: d.Get("jwt_claim_value").(string), Key: d.Get("key").(string), + Token: d.Get("token_id").(string), Description: d.Get("description").(string), } @@ -95,6 +96,7 @@ func resourceLiteLLMJWTKeyMappingUpdate(d *schema.ResourceData, m interface{}) e client := m.(*Client) oldKey, _ := d.GetChange("key") + oldTokenID, _ := d.GetChange("token_id") oldDescription, _ := d.GetChange("description") oldIsActive, _ := d.GetChange("is_active") @@ -104,6 +106,7 @@ func resourceLiteLLMJWTKeyMappingUpdate(d *schema.ResourceData, m interface{}) e // attempting to resync, so a failed refresh can't leave the rejected // values persisted into state. d.Set("key", oldKey) + d.Set("token_id", oldTokenID) d.Set("description", oldDescription) d.Set("is_active", oldIsActive) if readErr := resourceLiteLLMJWTKeyMappingRead(d, m); readErr != nil { @@ -146,6 +149,7 @@ func updateJWTKeyMapping(d *schema.ResourceData, client *Client) error { updateRequest := JWTKeyMappingUpdateRequest{ ID: d.Id(), Key: d.Get("key").(string), + Token: d.Get("token_id").(string), Description: d.Get("description").(string), IsActive: d.Get("is_active").(bool), } diff --git a/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go b/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go index 8007d1d4e08..27b75849fe6 100644 --- a/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go +++ b/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go @@ -628,3 +628,99 @@ func TestJWTKeyMappingCreateDoesNotLeakKeyInErrors(t *testing.T) { t.Fatalf("the virtual key must be redacted in errors, got %v", err) } } + +func TestJWTKeyMappingCreateSendsTokenIDAndOmitsKey(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + const tokenHash = "1923314ae0efc8b2523c7d421bac5a7cf88df291273b139948b526d396974a41" + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "token_id": tokenHash, + "is_active": true, + }) + + if err := resourceLiteLLMJWTKeyMappingCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + create := (*calls)[0] + if create.Body["token"] != tokenHash { + t.Fatalf("token hash not sent: %v", create.Body["token"]) + } + if _, sent := create.Body["key"]; sent { + t.Fatalf("key must be omitted when token_id is used, got: %v", create.Body) + } +} + +func TestJWTKeyMappingCreateOmitsTokenWhenKeyIsUsed(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + "is_active": true, + }) + + if err := resourceLiteLLMJWTKeyMappingCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + create := (*calls)[0] + if create.Body["key"] != "sk-abc123" { + t.Fatalf("virtual key not sent: %v", create.Body["key"]) + } + if _, sent := create.Body["token"]; sent { + t.Fatalf("token must be omitted when key is used, got: %v", create.Body) + } +} + +func TestJWTKeyMappingUpdateSendsTokenIDAndOmitsKey(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + const oldHash = "1111111111111111111111111111111111111111111111111111111111111111" + const newHash = "2222222222222222222222222222222222222222222222222222222222222222" + + client := NewClient(srv.URL, "test-key", true) + d := resourceDataWithChange(t, + map[string]string{ + "id": "map-abc-123", + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "token_id": oldHash, + "is_active": "true", + }, + map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "token_id": newHash, + "is_active": true, + }) + + if err := resourceLiteLLMJWTKeyMappingUpdate(d, client); err != nil { + t.Fatalf("update failed: %v", err) + } + + var update *jwtKeyMappingCall + for i := range *calls { + if (*calls)[i].Path == "/jwt/key/mapping/update" { + update = &(*calls)[i] + } + } + if update == nil { + t.Fatalf("expected an update call, got %v", *calls) + } + if update.Body["token"] != newHash { + t.Fatalf("new token hash not sent: %v", update.Body["token"]) + } + if _, sent := update.Body["key"]; sent { + t.Fatalf("key must be omitted when token_id is used, got: %v", update.Body) + } +} diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 7bef44409fd..8bcf7dc4fe3 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -276,13 +276,15 @@ type VectorStoreInfoRequest struct { type JWTKeyMappingRequest struct { JWTClaimName string `json:"jwt_claim_name"` JWTClaimValue string `json:"jwt_claim_value"` - Key string `json:"key"` + Key string `json:"key,omitempty"` + Token string `json:"token,omitempty"` Description string `json:"description,omitempty"` } type JWTKeyMappingUpdateRequest struct { ID string `json:"id"` Key string `json:"key,omitempty"` + Token string `json:"token,omitempty"` Description string `json:"description"` IsActive bool `json:"is_active"` } diff --git a/tests/benchmarks/conftest.py b/tests/benchmarks/conftest.py index c9b31cfb7d7..309ecab9991 100644 --- a/tests/benchmarks/conftest.py +++ b/tests/benchmarks/conftest.py @@ -8,8 +8,11 @@ flipping results between runs. Running the executor inline keeps each benchmark's cost self-contained and deterministic. """ +import os +import sys from collections.abc import Callable, Iterator from concurrent.futures import Future +from pathlib import Path from typing import ParamSpec, TypeVar import pytest @@ -20,6 +23,21 @@ P = ParamSpec("P") R = TypeVar("R") +def pytest_configure(config: pytest.Config) -> None: + if os.environ.get("LITELLM_REQUIRE_INSTALLED_WHEEL") != "1": + return + + import litellm + import litellm.rust_bridge._native as native + + prefix = Path(sys.prefix).resolve() + for name, module_file in (("litellm", litellm.__file__), ("litellm.rust_bridge._native", native.__file__)): + path = Path(module_file).resolve() + if not path.is_relative_to(prefix): + raise pytest.UsageError(f"{name} resolved outside the benchmark environment: {path}") + print(f"{name}: {path}") # noqa: T201 # provenance evidence must be visible in CI logs + + def _submit_inline(fn: Callable[P, R], /, *args: P.args, **kwargs: P.kwargs) -> Future[R]: future: Future[R] = Future() try: diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 78e6562a4a8..2b37ab14e7f 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -212,6 +212,11 @@ def select_tests(changed: tuple[str, ...]) -> tuple[str, ...]: (("tests/e2e/batches/test_managed_files_enforcement_e2e.py",), ()), (("tests/e2e/guardrails/test_presidio_masking_e2e.py",), ()), (("tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py",), ()), + (("tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e.py",), ()), + ( + ("tests/e2e/logging/test_team_langfuse_callback_e2e.py",), + ("tests/e2e/logging/test_team_langfuse_callback_e2e.py",), + ), ( ("tests/e2e/llm_translation/realtime/test_realtime_e2e.py",), ("tests/e2e/llm_translation/realtime/test_realtime_e2e.py",), diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 2b59e8770b5..15c2d6763d6 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -105,7 +105,7 @@ A couple of logging destinations are configured on the proxy rather than by the ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite and both JWT suites as canaries, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Keycloak, Jaeger, and TLS cluster-mode Valkey. Realm-only edits also trigger these canaries. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, and `guardrails/test_presidio_masking_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start. The Redis chaos test under `load/` needs a proxy it can pause the Redis of on the same host (`gateway/redis_chaos_ci_config.yml`), which `.github/workflows/test-e2e-redis-chaos.yml` boots, and which the Buildkite `e2e-redis-chaos` step in project-releaser runs co-located with Postgres and Valkey in one pod; it is deselected unless `E2E_REDIS_CHAOS` is set +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite and both JWT suites as canaries, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Keycloak, Jaeger, and TLS cluster-mode Valkey. Realm-only edits also trigger these canaries. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, and `guardrails/test_presidio_masking_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start. `logging/test_otel_v2_langfuse_generation_output_e2e.py` is marked `otel_v2` and deselects itself unless `E2E_OTEL_V2` is set, because it needs a gateway booted with `LITELLM_OTEL_V2=true` and Langfuse credentials, neither of which this stack provides, so run it with `E2E_OTEL_V2=1` against a local OTel v2 proxy. The Redis chaos test under `load/` needs a proxy it can pause the Redis of on the same host (`gateway/redis_chaos_ci_config.yml`), which `.github/workflows/test-e2e-redis-chaos.yml` boots, and which the Buildkite `e2e-redis-chaos` step in project-releaser runs co-located with Postgres and Valkey in one pod; it is deselected unless `E2E_REDIS_CHAOS` is set Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A file whose tests are all marked skip therefore cannot pass this check, so unskip at least one of them, or add the file to `UNSUPPORTED` in `select_tests.py` with the reason, before changing one. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index d18bed6c088..1731b4c620d 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -21,6 +21,7 @@ failures are hard test failures (see `tests/e2e/AGENTS.md`). | Vertex AI | yes | yes | yes | yes | yes (provider-transformed) | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | | Bedrock | yes (unified only) | yes | yes | yes (unfiltered managed list) | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | | Bedrock GovCloud (`us-gov-west-1`) | yes (unified only) | yes | no | no | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` on model, resolved from `AWS_GOVCLOUD_ACCESS_KEY_ID` / `AWS_GOVCLOUD_SECRET_ACCESS_KEY` / `AWS_GOVCLOUD_BATCH_S3_BUCKET` / `AWS_GOVCLOUD_BATCH_ROLE_ARN`) | +| Bedrock split S3 identity | no | no | no | no | yes (file upload, content, delete) | S3 signed with `s3_access_key_id` / `s3_secret_access_key` (`AWS_S3_ONLY_ACCESS_KEY_ID` / `AWS_S3_ONLY_SECRET_ACCESS_KEY`, object rights on `AWS_BATCH_S3_BUCKET` only) while `aws_*` is `AWS_BEDROCK_ONLY_ACCESS_KEY_ID` / `AWS_BEDROCK_ONLY_SECRET_ACCESS_KEY`, an identity with no S3 rights on that bucket | Bedrock cancel maps to `StopModelInvocationJob` and comes back `cancelling`; the lifecycle asserts it the same way it does for OpenAI (`_CANCEL_ASSERTED_PROVIDERS`). diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 9b3c06d9a1b..9bb6d05bec8 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -1024,6 +1024,72 @@ class TestBedrockBatchAssumeRole: assert fetched.id == batch.id +def _split_s3_identity_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=ASSUME_ROLE_RAW_MODEL, + aws_access_key_id="os.environ/AWS_BEDROCK_ONLY_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_BEDROCK_ONLY_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + s3_region_name="os.environ/AWS_REGION", + s3_bucket_name="os.environ/AWS_BATCH_S3_BUCKET", + s3_access_key_id="os.environ/AWS_S3_ONLY_ACCESS_KEY_ID", + s3_secret_access_key="os.environ/AWS_S3_ONLY_SECRET_ACCESS_KEY", + aws_batch_role_arn="os.environ/AWS_BATCH_ROLE_ARN", + ) + + +class TestBedrockBatchSplitS3Credentials: + """Bedrock batch deployment whose aws_* identity cannot touch the bucket. + + AWS_BEDROCK_ONLY_* is an IAM user with no S3 rights on AWS_BATCH_S3_BUCKET; + AWS_S3_ONLY_* is an IAM user with object rights on that bucket only. Every + S3 call the proxy signs (PutObject on upload, GetObject on content, + DeleteObject on delete) must use the s3_* pair, otherwise S3 answers 403. + """ + + @pytest.mark.covers( + "llm.files.bedrock.split_s3_credentials.nonstream.works", + exercised_on=["files"], + ) + def test_file_lifecycle_signs_s3_with_s3_credentials( + self, client: BatchClient, resources: ResourceManager + ) -> None: + model_name = batch_model_name("bedrock-split-s3-batch") + model_id = client.create_model(model_name, _split_s3_identity_params()) + resources.defer(lambda: client.delete_model(model_id)) + key = resources.key() + + uploaded = client.upload_file( + content=render_jsonl(ASSUME_ROLE_RAW_MODEL), + form=FileUploadForm(purpose="batch", target_model_names=model_name), + key=key, + ) + assert isinstance(uploaded, Success), ( + f"upload must sign the S3 PutObject with s3_access_key_id, got {uploaded!r}" + ) + file = uploaded.data + resources.defer(lambda: cleanup_file(client, file.id, key=key)) + assert_file_object(file, provider="bedrock") + + downloaded = client.proxy.transport.download( + f"/v1/files/{file.id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"content must sign the S3 GetObject with s3_access_key_id, " + f"got {downloaded.status_code}: {downloaded.body[:300]}" + ) + assert all(json.loads(line) for line in downloaded.body.strip().splitlines()), ( + f"content download returned non-JSONL body: {downloaded.body[:200]}" + ) + + deleted = client.delete_file(file.id, key=key) + assert isinstance(deleted, Success), ( + f"delete must sign the S3 DeleteObject with s3_access_key_id, got {deleted!r}" + ) + assert deleted.data.id == file.id, f"delete confirmed a different file: {deleted.data!r}" + + GOVCLOUD_REGION: Final = "us-gov-west-1" GOVCLOUD_RAW_MODEL: Final = "bedrock/amazon.nova-lite-v1:0" diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 8776d00d502..ca0fbd84c35 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -30,6 +30,7 @@ from e2e_config import ( FIXTURE_MODE_RAW, MANAGED_FILES_OPT_IN_ENV, MCP_OAUTH_LIVE_OPT_IN_ENV, + OTEL_V2_OPT_IN_ENV, PROMPT_CACHING_OPT_IN_ENV, PROVIDER_EDGE_HOST_OPT_IN_ENV, PROXY_BASE_URL, @@ -61,6 +62,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, "mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV, "provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV, + "otel_v2": OTEL_V2_OPT_IN_ENV, } ) @@ -150,6 +152,10 @@ def pytest_configure(config: pytest.Config) -> None: "provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the " "gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set", ) + config.addinivalue_line( + "markers", + "otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index 81832bebf49..4fe9f949fae 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -2,10 +2,12 @@ # Rolls up into the "Logging & Guardrails" dashboard module together with logging.* - {id: guardrail.presidio.pre_call.masks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "PII masking pre-call; data-leak blast radius"} - {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"} +- {id: guardrail.presidio.post_call.masks_generated_output, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, chat_completions_stream, anthropic_messages_stream], source: "guardrail_hooks/presidio.py", rationale: "Mask model-generated credit-card output with the UI default scope"} - {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"} - {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"} - {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"} - {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"} +- {id: guardrail.litellm_content_filter.pre_call.blocks_video, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [videos], source: "test_key_guardrail_video_e2e.py", fail_before_fix: proven, rationale: "A content-filter guardrail attached to a key (metadata.guardrails) blocks a banned prompt on POST /v1/videos before the provider is called; before the fix the route's call type was unknown to the unified guardrail hook and the prompt went to the provider unscanned (LIT-6685)"} - {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"} - {id: guardrail.litellm_content_filter.apply_endpoint.blocks, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail blocks banned content for customers that call the apply surface directly"} - {id: guardrail.litellm_content_filter.apply_endpoint.allows, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [allows], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail returns clean text for allowed input"} diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 64335aa560c..49d4d92ff0b 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -89,9 +89,6 @@ - {id: llm.chat_completions.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip"} - {id: llm.chat_completions.together_ai.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together cost header and spend row match the registry price"} - {id: llm.chat_completions.together_ai.thinking.nonstream.effort_none_disables, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [effort_none_disables], source: "llm_translation/test_together_ai_e2e.py", rationale: "reasoning_effort=none maps to Together's reasoning disable toggle on hybrid models"} -- {id: llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "Native MiMo v2.6 rows price the cost header and spend row from the cost map"} -- {id: llm.chat_completions.xiaomi_mimo.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo reasoning deltas stream as reasoning_content"} -- {id: llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo tool calls are not dropped"} - {id: llm.chat_completions.together_ai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: structured_output, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "response_format json_schema reaches Together and constrains the reply"} - {id: llm.chat_completions.together_ai.prompt_cache_5m.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: prompt_cache_5m, streaming: nonstream, assertions: [cache_hit, cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together prefix-cache reads bill at cache_read_input_token_cost, not full input price"} - {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"} diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 50f9b9808b2..c58c8af44ff 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -47,6 +47,7 @@ - {id: llm.files.vertex.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:52", rationale: "Vertex file upload to GCS"} - {id: llm.files.bedrock.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:59", rationale: "Bedrock file upload to S3"} - {id: llm.files.bedrock.govcloud_partition.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: govcloud_partition, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock file upload to an S3 bucket in the us-gov-west-1 partition"} +- {id: llm.files.bedrock.split_s3_credentials.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: split_s3_credentials, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-8297", rationale: "Bedrock file upload, content and delete sign S3 with s3_access_key_id / s3_secret_access_key when they differ from the aws_* identity"} - {id: llm.files.gemini.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: gemini, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Gemini Files API upload via proxy"} - {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"} - {id: llm.files.openai.require_managed_files_upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_managed_files_enforcement_e2e.py / LIT-5902", rationale: "With require_managed_files enabled, an upload without target_model_names and an upload carrying a model param are both rejected 400; runs only in the sequential managed-files stack phase (E2E_MANAGED_FILES_STACK)"} diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 65354100f58..334780eda53 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -12,6 +12,7 @@ - {id: reliability.cooldown.429.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "429", assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:69", rationale: "Cools on 429, avoids hammering exhausted provider"} - {id: reliability.cooldown.auth.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: auth, assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:74", rationale: "Cools on 401 auth error"} - {id: reliability.cooldown.timeout.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: timeout, assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:77", rationale: "Cools on 408 timeout"} +- {id: reliability.cooldown.client_disconnect.stays_healthy, module: reliability, tier: P0, behavior: cooldown, variant: client_disconnect, assertions: [stays_healthy], exercised_on: [chat_completions], source: "llms/azure/azure.py:484", fail_before_fix: proven, rationale: "A client hanging up mid-request under cancel_on_disconnect never benches the Azure deployment it was talking to: the cancellation used to surface as a fake 500 that tripped the cooldown and sent every caller behind it to billed fallbacks (GitHub issues #35329 and #42222)"} - {id: reliability.routing.simple_shuffle.picks_healthy_deployment, module: reliability, tier: P1, behavior: routing, variant: simple_shuffle, assertions: [picks_healthy_deployment], exercised_on: [chat_completions], source: "router_strategy/simple_shuffle.py", rationale: "Baseline weighted/uniform pick"} - {id: reliability.routing.latency_based.picks_lowest_latency, module: reliability, tier: P1, behavior: routing, variant: latency_based, assertions: [picks_lowest_latency], exercised_on: [chat_completions], source: "router_strategy/lowest_latency.py", rationale: "Routes to lowest-latency deployment"} - {id: reliability.routing.cost_based.picks_lowest_cost, module: reliability, tier: P1, behavior: routing, variant: cost_based, assertions: [picks_lowest_cost], exercised_on: [chat_completions], source: "router_strategy/lowest_cost.py", rationale: "Spend-aware routing"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index d9c20d5c588..f3ac1ef8a83 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -67,6 +67,7 @@ LlmCapability = Literal[ "batch_deployment", "count_tokens", "govcloud_partition", + "split_s3_credentials", "input_validation", "long_context_1m", "mid_conversation_system", diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 14de4619664..311c944eeb4 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -147,6 +147,7 @@ REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE" PROVIDER_EDGE_HOST_OPT_IN_ENV: Final = "E2E_PROVIDER_EDGE_HOST_REACHABLE" +OTEL_V2_OPT_IN_ENV: Final = "E2E_OTEL_V2" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index d4978601b20..97f1e1671f8 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -676,6 +676,35 @@ def send( return streaming_outcome(resp, stream, sent_at=sent_at) +class AbandonedRequest(BaseModel): + """A non-streaming request whose socket the client closed ``after`` seconds in, + before the proxy had answered.""" + + kind: Literal["abandoned"] = "abandoned" + after: float + + +def abandon( + url: URL, *, headers: BaseModel, json: BaseModel, after: float, connect_timeout: float = 10.0 +) -> AbandonedRequest | StreamingResponse: + """POST and close the connection ``after`` seconds if no response head has arrived + by then; returns the response instead when the proxy answered first.""" + sent_at: Final = time.monotonic() + session: Final = requests.Session() + try: + resp = session.post( + str(url), + headers=_headers(headers), + json=wire_body(json), + timeout=(connect_timeout, after), + ) + except requests.exceptions.ReadTimeout: + return AbandonedRequest(after=after) + finally: + session.close() + return streaming_outcome(resp, False, sent_at=sent_at) + + def stream(url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0) -> StreamingResponse: """Streaming (SSE) call: consumes the stream counting events, and captures the x-litellm-call-id + content-type headers. Body is elided.""" diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 352caddf588..2d02fedddae 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -10,6 +10,7 @@ general_settings: store_prompts_in_spend_logs: true database_connection_pool_limit: 10 forward_client_headers_to_llm_api: false + cancel_on_disconnect: true maximum_spend_logs_retention_period: "60d" maximum_spend_logs_cleanup_cron: "0 1 * * *" proxy_budget_rescheduler_min_time: 15 diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index ed112a79b9b..97ecac0290f 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -20,6 +20,7 @@ from models import ( ChatResponse, ChatTool, KeyGenerateBody, + KeyMetadata, LiteLLMParamsBody, TeamDeleteBody, TeamInfoParams, @@ -27,6 +28,8 @@ from models import ( TeamMetadata, TeamNewBody, TeamNewResponse, + VideoCreateBody, + VideoCreateResponse, ) from proxy_client import ProxyClient from pydantic import BaseModel @@ -43,7 +46,7 @@ class BlockedWordBody(BaseModel): class GuardrailParamsBase(BaseModel): - mode: GuardrailMode + mode: GuardrailMode | list[GuardrailMode] default_on: bool @@ -151,12 +154,12 @@ class _ResponsesGuardrailBody(BaseModel): class GuardrailsClient: proxy: ProxyClient - def create_content_filter_guardrail(self, name: str, blocked_keyword: str) -> str: + def create_content_filter_guardrail(self, name: str, blocked_keyword: str, *, default_on: bool = True) -> str: return self.register( name, ContentFilterParamsBody( mode="pre_call", - default_on=True, + default_on=default_on, blocked_words=[BlockedWordBody(keyword=blocked_keyword, action="BLOCK")], ), ) @@ -266,6 +269,21 @@ class GuardrailsClient: def create_key_in_team(self, team_id: str) -> str: return self.proxy.generate_key(KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user")) + def create_key_with_guardrails(self, resources: ResourceManager, guardrails: list[str]) -> str: + key = self.proxy.generate_key( + KeyGenerateBody(user_id="e2e-guardrails-user", metadata=KeyMetadata(guardrails=guardrails)) + ) + resources.defer(lambda: self.proxy.delete_key(key)) + return key + + def create_video(self, key: str, model: str, prompt: str) -> Result[VideoCreateResponse]: + return self.proxy.transport.post( + "/v1/videos", + headers=self.proxy.transport.bearer(key), + json=VideoCreateBody(model=model, prompt=prompt, seconds="4"), + response_type=VideoCreateResponse, + ) + def chat( self, key: str, @@ -363,6 +381,26 @@ class GuardrailsClient: ), ) + def messages_stream_raw( + self, + key: str, + model: str, + text: str, + *, + guardrails: list[str] | None = None, + max_tokens: int = 64, + ) -> StreamingResponse: + return self.proxy.messages_stream( + key, + AnthropicMessagesBody( + model=model, + messages=[ChatMessage(role="user", content=text)], + max_tokens=max_tokens, + stream=True, + guardrails=guardrails, + ), + ) + def responses( self, key: str, diff --git a/tests/e2e/guardrails/test_key_guardrail_video_e2e.py b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py new file mode 100644 index 00000000000..5f318e141a5 --- /dev/null +++ b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import pytest +from e2e_config import unique_marker +from e2e_http import Success, UnknownApiError +from guardrails_client import GuardrailsClient, poll_until_blocked +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +CHAT_MODEL = "gemini-2.5-flash" +VIDEO_BACKEND = "vertex_ai/veo-3.1-fast-generate-001" + + +def _video_prompt_with(banned_keyword: str) -> str: + return f"A short clip of a paper boat floating down a stream. {banned_keyword}" + + +def _create_video_model(client: GuardrailsClient, resources: ResourceManager) -> str: + model_name = f"e2e-guard-video-{unique_marker()}" + model_id = client.proxy.create_model( + model_name, + LiteLLMParamsBody( + model=VIDEO_BACKEND, + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location="os.environ/VERTEXAI_LOCATION", + vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", + ), + provider_live=True, + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model_name + + +class TestKeyAttachedGuardrailOnVideos: + @pytest.mark.covers( + "guardrail.litellm_content_filter.pre_call.blocks_video", + exercised_on=["videos"], + ) + def test_key_attached_content_filter_blocks_banned_video_prompt( + self, client: GuardrailsClient, resources: ResourceManager + ) -> None: + banned = unique_marker() + guardrail_name = f"e2e-video-filter-{banned}" + guardrail_id = client.create_content_filter_guardrail(guardrail_name, banned, default_on=False) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + key = client.create_key_with_guardrails(resources, [guardrail_name]) + model = _create_video_model(client, resources) + + synced = poll_until_blocked(lambda: client.chat(key, CHAT_MODEL, _video_prompt_with(banned))) + assert isinstance(synced, UnknownApiError) and synced.status_code == 400, ( + f"key guardrail {guardrail_name!r} never synced to the proxy on /chat/completions: {synced}" + ) + + result = client.create_video(key, model, _video_prompt_with(banned)) + match result: + case UnknownApiError(status_code=status, body=body): + assert status == 400, f"expected a 400 guardrail block, got {status}: {body[:300]}" + assert "content blocked" in body.lower() or banned in body, ( + f"block response missing content-filter reason: {body[:300]}" + ) + case Success(data=video): + pytest.fail( + f"key-attached guardrail {guardrail_name!r} was skipped on /v1/videos: " + f"the banned prompt reached the provider and started video job {video.id}" + ) + case _: + pytest.fail(f"unexpected /v1/videos outcome for a banned prompt: {result}") diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py index 49d698938ce..d47d64be9e3 100644 --- a/tests/e2e/guardrails/test_presidio_masking_e2e.py +++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py @@ -30,6 +30,7 @@ this suite deliberately requires the detected-entity details to remain visible. from __future__ import annotations import os +import re import time from collections.abc import Callable from typing import Final, Literal @@ -65,9 +66,13 @@ GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS = 5.0 # angle brackets; the logged payload keeps the placeholder verbatim. MASKED_EMAIL_TOKEN = "EMAIL_ADDRESS" MASKED_PHONE_TOKEN = "PHONE_NUMBER" +MASKED_CREDIT_CARD_TOKEN = "CREDIT_CARD" # Fictional NANP 555 number; a standard format Presidio's phone recognizer detects. FAKE_PHONE = "+1 415-555-0134" +FAKE_VISA_TEST_CARD = "4111 1111 1111 1111" + +_CARD_DIGIT_RUN: Final = re.compile(r"(?:\d[ -]?){13,19}") def _presidio_bases() -> tuple[str, str]: @@ -86,8 +91,8 @@ def _register_presidio( resources: ResourceManager, *, name: str, - mode: GuardrailMode = "pre_call", - filter_scope: Literal["input", "output", "both"] = "input", + mode: GuardrailMode | list[GuardrailMode] = "pre_call", + filter_scope: Literal["input", "output", "both"] | None = "input", entities: dict[PiiEntity, PiiAction] | None = None, ) -> None: analyzer, anonymizer = _presidio_bases() @@ -123,6 +128,74 @@ def _first_content(response: ChatResponse) -> str: return (message.content if message else None) or "" +class _StreamDelta(BaseModel): + content: str | None = None + + +class _StreamChoice(BaseModel): + delta: _StreamDelta + + +class _StreamChunk(BaseModel): + choices: tuple[_StreamChoice, ...] = () + + +class _AnthropicStreamDelta(BaseModel): + type: str | None = None + text: str | None = None + + +class _AnthropicStreamEvent(BaseModel): + type: str + delta: _AnthropicStreamDelta | None = None + + +def _credit_card_prompt(marker: str) -> str: + return ( + f"{marker} Reply with only the well known Visa sandbox test card number that starts with 4111, " + "the 16 digits grouped in fours separated by spaces, and nothing else." + ) + + +def _passes_luhn(digits: str) -> bool: + checksum = sum( + digit if position % 2 == 0 else (digit * 2 - 9 if digit * 2 > 9 else digit * 2) + for position, digit in enumerate(int(char) for char in reversed(digits)) + ) + return checksum % 10 == 0 + + +def _contains_card_number(text: str) -> bool: + """Presidio's CREDIT_CARD recognizer only reports Luhn-valid digit runs, so a + Luhn-invalid number the model hallucinates is not something masking can catch.""" + return any( + 13 <= len(digits) <= 19 and _passes_luhn(digits) + for digits in (re.sub(r"[ -]", "", match.group()) for match in _CARD_DIGIT_RUN.finditer(text)) + ) + + +def _stream_content(result: StreamingResponse) -> str: + return "".join( + choice.delta.content + for event in result.stream_events + if event != "[DONE]" + for choice in _StreamChunk.model_validate_json(event).choices[:1] + if choice.delta.content + ) + + +def _anthropic_stream_content(result: StreamingResponse) -> str: + return "".join( + event.delta.text + for payload in result.stream_events + for event in [_AnthropicStreamEvent.model_validate_json(payload)] + if event.type == "content_block_delta" + and event.delta is not None + and event.delta.type == "text_delta" + and event.delta.text + ) + + def _messages_text(response: AnthropicMessagesResponse) -> str: """The text of a /v1/messages answer, whichever shape the proxy produced (Anthropic-native content blocks or OpenAI-normalized choices).""" @@ -290,6 +363,124 @@ class TestPresidioPostCallMasking: time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) +def _assert_eventually_masks_generated_card(fetch: Callable[[], str | None]) -> None: + deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS + last: str = "" + while True: + content = fetch() + if content is not None: + last = content + if _contains_card_number(content): + pytest.fail( + "the post_call output masking let a card number through: " + f"{content[:300]!r}" + ) + if MASKED_CREDIT_CARD_TOKEN in content: + return + if time.monotonic() >= deadline: + pytest.fail( + "presidio post_call output masking never masked the generated card within " + f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}" + ) + time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) + + +class TestPresidioCreditCardOutputMasking: + """Proves the UI-default Presidio scope masks model-generated card output.""" + + @pytest.mark.covers( + "guardrail.presidio.post_call.masks_generated_output", + exercised_on=["chat_completions"], + ) + def test_ui_default_scope_masks_a_card_number_the_model_generates_on_chat_completions( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + name = f"e2e-presidio-card-chat-{unique_marker()}" + _register_presidio( + client, + resources, + name=name, + mode=["pre_call", "post_call"], + filter_scope=None, + entities={"CREDIT_CARD": "MASK"}, + ) + prompt: Final = _credit_card_prompt(unique_marker()) + + def fetch() -> str | None: + result: Final = client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=512) + match result: + case Success(data=data): + return _first_content(data) + case _: + return None + + _assert_eventually_masks_generated_card(fetch) + + @pytest.mark.covers( + "guardrail.presidio.post_call.masks_generated_output", + exercised_on=["chat_completions_stream"], + ) + def test_ui_default_scope_masks_a_card_number_the_model_generates_on_streaming_chat_completions( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + name = f"e2e-presidio-card-stream-{unique_marker()}" + _register_presidio( + client, + resources, + name=name, + mode=["pre_call", "post_call"], + filter_scope=None, + entities={"CREDIT_CARD": "MASK"}, + ) + prompt: Final = _credit_card_prompt(unique_marker()) + + def fetch() -> str | None: + result: Final = client.chat_stream_raw( + scoped_key, + MODEL, + prompt, + guardrails=[name], + max_tokens=512, + ) + if not result.ok or result.stream_error: + return None + return _stream_content(result) + + _assert_eventually_masks_generated_card(fetch) + + @pytest.mark.covers( + "guardrail.presidio.post_call.masks_generated_output", + exercised_on=["anthropic_messages_stream"], + ) + def test_ui_default_scope_masks_a_card_number_the_model_generates_on_streaming_anthropic_messages( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + name = f"e2e-presidio-card-messages-stream-{unique_marker()}" + _register_presidio( + client, + resources, + name=name, + mode=["pre_call", "post_call"], + filter_scope=None, + entities={"CREDIT_CARD": "MASK"}, + ) + prompt: Final = _credit_card_prompt(unique_marker()) + + def fetch() -> str | None: + result: Final = client.messages_stream_raw( + scoped_key, + MODEL, + prompt, + guardrails=[name], + max_tokens=512, + ) + if not result.ok or result.stream_error: + return None + return _anthropic_stream_content(result) + + _assert_eventually_masks_generated_card(fetch) + + _LOGGED_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"} _ENTITY_LIST_ADAPTER: Final = TypeAdapter(list[GuardrailEntityMatch]) diff --git a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py index 87bd32d8dab..363b2a7e02e 100644 --- a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py @@ -45,6 +45,10 @@ pytestmark = pytest.mark.e2e COHERE_BACKEND = "cohere/command-r-08-2024" GEMINI_BACKEND = "gemini/gemini-2.5-flash" +VERTEX_BACKEND: Final = "vertex_ai/gemini-2.5-flash" +AZURE_OPENAI_BACKEND: Final = "azure/gpt-5.4-nano" +AZURE_OPENAI_API_VERSION: Final = "v1" +AZURE_FOUNDRY_BACKEND: Final = "azure_ai/claude-haiku-4-5" OPENAI_BACKEND = "openai/gpt-5.6" ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5-20251001" BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" @@ -108,7 +112,7 @@ def _assert_describes_cat(response: ChatResponse) -> None: assert response.choices, f"vision returned no choices: {response}" message = response.choices[0].message content = (message.content if message else None) or "" - assert "cat" in content.lower() or "feline" in content.lower(), ( + assert any(term in content.lower() for term in ("cat", "feline", "kitten", "kitty")), ( f"vision response did not describe the image: {content[:200]}" ) @@ -208,7 +212,6 @@ class TestChatCompletionsRegression: @pytest.mark.covers( "llm.chat_completions.openai.basic.nonstream.works", "llm.chat_completions.anthropic.basic.nonstream.works", - "llm.chat_completions.vertex.basic.nonstream.works", exercised_on=[], ) def test_chat_returns_real_completion( @@ -336,6 +339,232 @@ class TestGeminiChatCompletions: assert row.status == "success", f"gemini chat spend status={row.status!r}" +class TestVertexChatCompletions: + def _register(self, client: PassthroughClient, resources: ResourceManager, prefix: str) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=VERTEX_BACKEND, + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location="us-central1", + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + @pytest.mark.covers( + "llm.chat_completions.vertex.basic.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_vertex_chat_returns_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-vertex-chat") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=f"Reply with the single word pong. {unique_marker()}", + ) + ], + max_tokens=32, + ), + ) + ) + assert response.choices, f"vertex chat returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content and content.strip(), f"vertex chat returned empty content: {response}" + + @pytest.mark.covers( + "llm.chat_completions.vertex.tool_use.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_vertex_chat_returns_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-vertex-tool") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content="What is the weather in San Francisco? Use the get_weather tool.", + ) + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + ), + ) + ) + _assert_weather_tool_call(response) + + @pytest.mark.covers( + "llm.chat_completions.vertex.vision.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_vertex_chat_vision_describes_image( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-vertex-vision") + key = resources.key() + + response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32))) + _assert_describes_cat(response) + + @pytest.mark.covers( + "llm.chat_completions.vertex.basic.stream.works", + exercised_on=["chat_completions"], + ) + def test_vertex_chat_streams_real_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-vertex-stream") + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=f"Count from 1 to 5, one number per line. {unique_marker()}", + ) + ], + max_tokens=64, + stream=True, + ), + ) + _assert_streamed_completion(result) + + +class TestAzureOpenAIChatCompletions: + def _register(self, client: PassthroughClient, resources: ResourceManager, prefix: str) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=AZURE_OPENAI_BACKEND, + api_base="os.environ/AZURE_API_BASE", + api_key="os.environ/AZURE_API_KEY", + api_version=AZURE_OPENAI_API_VERSION, + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + @pytest.mark.covers( + "llm.chat_completions.azure_openai.basic.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_azure_openai_chat_returns_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-azure-openai-chat") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=f"Reply with the single word pong. {unique_marker()}", + ) + ], + max_tokens=32, + ), + ) + ) + assert response.choices, f"azure openai chat returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content and content.strip(), f"azure openai chat returned empty content: {response}" + + @pytest.mark.covers( + "llm.chat_completions.azure_openai.tool_use.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_azure_openai_chat_returns_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-azure-openai-tool") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content="What is the weather in San Francisco? Use the get_weather tool.", + ) + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + ), + ) + ) + _assert_weather_tool_call(response) + + +class TestAzureFoundryChatCompletions: + @pytest.mark.covers( + "llm.chat_completions.azure_foundry.basic.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_azure_foundry_chat_returns_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-azure-foundry-chat-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=AZURE_FOUNDRY_BACKEND, + api_base="os.environ/AZURE_AI_API_BASE", + api_key="os.environ/AZURE_AI_API_KEY", + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=f"Reply with the single word pong. {unique_marker()}", + ) + ], + max_tokens=32, + ), + ) + ) + assert response.choices, f"azure foundry chat returned no choices: {response}" + content = response.choices[0].message.content if response.choices[0].message else None + assert content and content.strip(), f"azure foundry chat returned empty content: {response}" + + class TestHostedVllmChat: """hosted_vllm (self-hosted OpenAI-compatible server) via /chat/completions.""" @@ -764,6 +993,90 @@ class TestAnthropicChatCompletions: resources.defer(lambda: client.proxy.delete_model(model_id)) return model + @pytest.mark.covers( + "llm.chat_completions.anthropic.structured_output.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_anthropic_chat_structured_output_conforms_to_schema( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-anthropic-schema") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content="Extract the person. John Doe is 42 years old.", + ) + ], + response_format=_PERSON_SCHEMA, + max_tokens=128, + ), + ) + ) + assert response.choices, f"anthropic structured output returned no choices: {response}" + message = response.choices[0].message + content = message.content if message else None + assert content, f"anthropic structured output returned empty content: {response}" + person = _Person.model_validate_json(content) + assert person.name.strip() and person.age == 42, f"anthropic schema output was wrong: {person}" + + @pytest.mark.covers( + "llm.chat_completions.anthropic.thinking.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_anthropic_chat_returns_thinking_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-anthropic-thinking") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage( + role="user", + content=( + "Prove that the sum of two odd integers is even, then find the smallest prime " + "greater than 100 such that p+2 is also prime." + ), + ) + ], + thinking=ThinkingParam(type="enabled", budget_tokens=1024), + max_tokens=2048, + ), + ) + ) + assert response.choices, f"anthropic thinking returned no choices: {response}" + message = response.choices[0].message + assert message and message.content and message.content.strip(), ( + f"anthropic thinking returned no answer content: {response}" + ) + assert message.reasoning_content and message.reasoning_content.strip(), ( + f"anthropic thinking returned no reasoning content: {response}" + ) + + @pytest.mark.covers( + "llm.chat_completions.anthropic.vision.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_anthropic_chat_vision_describes_image( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-anthropic-vision") + key = resources.key() + + response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32))) + _assert_describes_cat(response) + @pytest.mark.covers( "llm.chat_completions.anthropic.basic.stream.works", exercised_on=["chat_completions"], diff --git a/tests/e2e/llm_translation/test_responses_e2e.py b/tests/e2e/llm_translation/test_responses_e2e.py index 9cc70da63b0..6fa77694eb8 100644 --- a/tests/e2e/llm_translation/test_responses_e2e.py +++ b/tests/e2e/llm_translation/test_responses_e2e.py @@ -45,6 +45,9 @@ class _OptionalResponsesBody(BaseModel): BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +VERTEX_BACKEND: Final = "vertex_ai/gemini-2.5-flash" +AZURE_OPENAI_BACKEND: Final = "azure/gpt-5.4-nano" +AZURE_OPENAI_API_VERSION: Final = "v1" INSTRUCTIONS = "You are a helpful assistant" CAT_IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg" BEDROCK_EDGE_REGION: Final = "us-east-1" @@ -105,6 +108,23 @@ def _bedrock_params() -> LiteLLMParamsBody: ) +def _vertex_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=VERTEX_BACKEND, + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location="us-central1", + ) + + +def _azure_openai_params() -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=AZURE_OPENAI_BACKEND, + api_base="os.environ/AZURE_API_BASE", + api_key="os.environ/AZURE_API_KEY", + api_version=AZURE_OPENAI_API_VERSION, + ) + + def _register( proxy: ProxyClient, resources: ResourceManager, params: LiteLLMParamsBody, prefix: str = "e2e-responses" ) -> str: @@ -291,6 +311,66 @@ class TestResponses: ) _assert_weather_call(response) + @pytest.mark.covers("llm.responses.vertex.basic.nonstream.works") + def test_responses_vertex_returns_completion( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model = _register(proxy, resources, _vertex_params(), prefix="e2e-responses-vertex") + client = sdk.openai(resources.key()) + + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS, extra_body=NO_PROXY_CACHE + ) + assert response.output_text.strip(), f"/responses over vertex returned no output text: {response.output!r}" + + @pytest.mark.covers("llm.responses.vertex.tool_use.nonstream.works") + def test_responses_vertex_returns_function_call( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model = _register(proxy, resources, _vertex_params(), prefix="e2e-responses-vertex-tool") + client = sdk.openai(resources.key()) + + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], + tool_choice="required", + extra_body=NO_PROXY_CACHE, + ) + _assert_weather_call(response) + + @pytest.mark.covers("llm.responses.azure_openai.basic.nonstream.works") + def test_responses_azure_openai_returns_completion( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model = _register(proxy, resources, _azure_openai_params(), prefix="e2e-responses-azure-openai") + client = sdk.openai(resources.key()) + + response = client.responses.create( + model=model, input="reply with one word", instructions=INSTRUCTIONS, extra_body=NO_PROXY_CACHE + ) + assert response.output_text.strip(), ( + f"/responses over azure openai returned no output text: {response.output!r}" + ) + + @pytest.mark.covers("llm.responses.azure_openai.tool_use.nonstream.works") + def test_responses_azure_openai_returns_function_call( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model = _register(proxy, resources, _azure_openai_params(), prefix="e2e-responses-azure-openai-tool") + client = sdk.openai(resources.key()) + + response = client.responses.create( + model=model, + input="What is the weather in San Francisco? Use the get_weather tool.", + instructions=INSTRUCTIONS, + tools=[WEATHER_TOOL], + tool_choice="required", + extra_body=NO_PROXY_CACHE, + ) + _assert_weather_call(response) + @pytest.mark.provider_edge_host @pytest.mark.parametrize("endpoint", ["/v1/responses", "/v1/chat/completions"]) def test_bedrock_forwards_allowed_safety_identifier_as_additional_model_request_field( diff --git a/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py b/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py deleted file mode 100644 index efca216634b..00000000000 --- a/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py +++ /dev/null @@ -1,214 +0,0 @@ -"""Live e2e: Xiaomi MiMo v2.6 through the gateway on /chat/completions. - -Both native ``xiaomi_mimo/`` v2.6 rows (pro and flash) are registered via -``/model/new`` and driven against Xiaomi's own endpoint. What the gateway owes -us is that the reasoning chain surfaces as ``reasoning_content``, tool calls -survive translation, and the cost header plus spend row follow the proxy's own -cost-map price for the row (read back from ``/model/info``, never pinned here). -Requires XIAOMI_MIMO_API_KEY on the proxy; no skip gate. -""" - -from __future__ import annotations - -from typing import Final - -import pytest -from e2e_config import unique_marker -from e2e_http import StreamingResponse, require_successful_call, unwrap -from lifecycle import ResourceManager -from models import ( - ChatBody, - ChatMessage, - ChatResponse, - ChatTool, - ChatToolFunction, - CostMapEntry, - LiteLLMParamsBody, - OutMessage, - SpendLogRow, -) -from passthrough_client import PassthroughClient -from pydantic import BaseModel - -pytestmark = pytest.mark.e2e - -BACKENDS: Final = ("xiaomi_mimo/mimo-v2.6-pro", "xiaomi_mimo/mimo-v2.6-flash") -ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number." -WEATHER_PROMPT = "What is the weather in Paris? Use the tool." -COUNTING_PROMPT = "Count from 1 to 50, one number per line." - -WEATHER_TOOL = ChatTool( - function=ChatToolFunction( - name="get_weather", - description="Get the current weather for a location.", - parameters={ - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"], - }, - ) -) - - -class _WeatherArgs(BaseModel): - location: str - - -class _StreamDelta(BaseModel): - content: str | None = None - reasoning_content: str | None = None - - -class _StreamChoice(BaseModel): - delta: _StreamDelta | None = None - - -class _StreamChunk(BaseModel): - choices: list[_StreamChoice] = [] - - -def _approx_equal(actual: float, expected: float) -> bool: - return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) - - -@pytest.fixture(scope="module") -def registry(client: PassthroughClient) -> dict[str, CostMapEntry]: - return client.proxy.model_cost_map() - - -def _register(client: PassthroughClient, resources: ResourceManager, backend: str) -> tuple[str, str]: - model = f"e2e-xiaomi-{unique_marker()}" - model_id = client.proxy.create_model( - model, LiteLLMParamsBody(model=backend, api_key="os.environ/XIAOMI_MIMO_API_KEY") - ) - resources.defer(lambda: client.proxy.delete_model(model_id)) - return model, resources.key() - - -def _message(response: ChatResponse) -> OutMessage: - assert response.choices, f"Xiaomi returned no choices: {response}" - message = response.choices[0].message - assert message is not None, f"Xiaomi choice has no message: {response}" - return message - - -def _deltas(result: StreamingResponse) -> list[_StreamDelta]: - require_successful_call(result) - assert result.is_streaming, f"response was not streamed: {result.headers}" - assert not result.stream_error, f"stream errored: {result.stream_error}" - assert result.stream_done, f"stream never reached [DONE]: {result.stream_events[-3:]}" - return [ - choice.delta - for event in result.stream_events - for choice in _StreamChunk.model_validate_json(event).choices - if choice.delta is not None - ] - - -@pytest.mark.parametrize("backend", BACKENDS) -class TestXiaomiMimoChatCompletions: - @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged") - def test_cost_header_and_spend_row_match_the_registry_price( - self, - client: PassthroughClient, - resources: ResourceManager, - registry: dict[str, CostMapEntry], - backend: str, - ) -> None: - price = registry.get(backend) - assert price is not None, f"{backend} has no row in the proxy's cost map, so native calls would bill $0" - assert price.litellm_provider == "xiaomi_mimo", f"{backend} is filed under the wrong provider: {price}" - assert price.input_cost_per_token and price.output_cost_per_token, f"{backend} carries no price: {price}" - model, key = _register(client, resources, backend) - - result = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(key), - json=ChatBody( - model=model, - messages=[ChatMessage(role="user", content=f"{ARITHMETIC_PROMPT} {unique_marker()}")], - max_tokens=1024, - ), - ) - require_successful_call(result) - response = ChatResponse.model_validate_json(result.body) - message = _message(response) - assert message.content and "43" in message.content, f"answer lost: {message}" - assert message.reasoning_content, f"{backend} reasons, but no reasoning_content came back: {message}" - - usage = response.usage - assert usage is not None and usage.prompt_tokens and usage.completion_tokens, ( - f"response carries no usage, so the cost cannot be real: {result.body[:300]}" - ) - header_cost = result.response_cost - assert header_cost is not None and header_cost > 0, ( - f"x-litellm-response-cost header missing or non-positive: {result.headers}" - ) - cached = (usage.prompt_tokens_details.cached_tokens or 0) if usage.prompt_tokens_details else 0 - expected = ( - (usage.prompt_tokens - cached) * price.input_cost_per_token - + cached * (price.cache_read_input_token_cost or 0.0) - + usage.completion_tokens * price.output_cost_per_token - ) - assert _approx_equal(header_cost, expected), ( - f"header cost {header_cost} disagrees with the registry price for {backend} at {usage}: expected {expected}" - ) - - def _priced(rows: list[SpendLogRow]) -> bool: - return any(row.spend is not None and row.spend > 0 for row in rows) - - rows = client.proxy.poll_logs_for_key(key, predicate=_priced) - priced = [row for row in rows if row.spend is not None and row.spend > 0] - assert priced, f"no priced spend row landed for key {key}; got {rows}" - row = priced[0] - assert row.custom_llm_provider == "xiaomi_mimo", f"spend row misattributed: {row}" - assert row.spend is not None and _approx_equal(row.spend, header_cost), ( - f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}" - ) - - @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.thinking.stream.works") - def test_reasoning_and_answer_stream_as_deltas( - self, client: PassthroughClient, resources: ResourceManager, backend: str - ) -> None: - model, key = _register(client, resources, backend) - - deltas = _deltas( - client.proxy.chat_stream( - key, - ChatBody( - model=model, - messages=[ChatMessage(role="user", content=COUNTING_PROMPT)], - max_tokens=2048, - stream=True, - ), - ) - ) - reasoning = "".join(delta.reasoning_content or "" for delta in deltas) - content = "".join(delta.content or "" for delta in deltas) - assert reasoning, f"stream carried no reasoning_content deltas: {deltas[:5]}" - assert "50" in content, f"streamed answer lost: {content[:300]!r}" - - @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works") - def test_tool_call_is_returned(self, client: PassthroughClient, resources: ResourceManager, backend: str) -> None: - model, key = _register(client, resources, backend) - - message = _message( - unwrap( - client.proxy.chat( - key, - ChatBody( - model=model, - messages=[ChatMessage(role="user", content=WEATHER_PROMPT)], - tools=[WEATHER_TOOL], - max_tokens=1024, - ), - ) - ) - ) - assert message.tool_calls, f"{backend} dropped the tool call: {message}" - call = message.tool_calls[0] - assert call.id, f"tool call carries no id, so a tool result cannot answer it: {call}" - assert call.function.name == "get_weather", f"wrong tool called: {call}" - assert call.function.arguments, f"tool call carries no arguments: {call}" - args = _WeatherArgs.model_validate_json(call.function.arguments) - assert "paris" in args.location.lower(), f"tool arguments lost the location: {args}" diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index 66dfa233ec4..c2f987ea33d 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -146,6 +146,29 @@ class LangfuseObservationList(BaseModel): data: list[LangfuseObservation] = [] +class LangfuseOtelMetadata(BaseModel): + """Langfuse stores every OTel span attribute under metadata.attributes.""" + + model_config = ConfigDict(extra="ignore") + + attributes: dict[str, str] = {} + + +def otel_attributes(obs: LangfuseObservation) -> dict[str, str]: + try: + return LangfuseOtelMetadata.model_validate(obs.metadata).attributes + except ValidationError: + return {} + + +def is_otel_v2_generation(obs: LangfuseObservation, *, key_alias: str) -> bool: + attributes = otel_attributes(obs) + return ( + attributes.get("langfuse.observation.type") == "generation" + and attributes.get("litellm.metadata.user_api_key_alias") == key_alias + ) + + class LangfuseListParams(BaseModel): model_config = ConfigDict(populate_by_name=True) @@ -630,6 +653,18 @@ class LoggingClient: time.sleep(POLL_INTERVAL) return last + def poll_langfuse_generation( + self, creds: LangfuseCreds, *, key_alias: str, from_start_time: str + ) -> LangfuseObservation | None: + """The OTel v2 generation the proxy exported for one key alias since from_start_time.""" + deadline = time.monotonic() + POLL_TIMEOUT + while time.monotonic() < deadline: + for obs in self.list_langfuse_observations(creds, from_start_time=from_start_time): + if is_otel_v2_generation(obs, key_alias=key_alias): + return obs + time.sleep(POLL_INTERVAL) + return None + def poll_langfuse_trace_observations( self, creds: LangfuseCreds, diff --git a/tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e.py b/tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e.py new file mode 100644 index 00000000000..71008d94557 --- /dev/null +++ b/tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e.py @@ -0,0 +1,210 @@ +"""Live e2e: the OTel v2 Langfuse generation carries output for every non-chat endpoint (LIT-8309). + +With LITELLM_OTEL_V2=true the proxy exports one generation per request to the +team's Langfuse destination. Chat, Responses, embeddings and OCR already fill +its output; this file pins the remaining five families. Each test registers a +real OpenAI deployment, drives the endpoint through the shared transport, then +reads the generation back from Langfuse and asserts its output reflects what +the caller received: the completion text, the transcript, the moderation +verdict, and for images and speech a bounded summary that never carries the +raw base64 or audio bytes. +""" + +from __future__ import annotations + +import base64 +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from logging_client import LangfuseCreds, LangfuseObservation, LoggingClient, load_langfuse_creds +from models import ( + CompletionBody, + CompletionResponse, + ImageGenerationBody, + ImageGenerationResponse, + LiteLLMParamsBody, + ModerationBody, + ModerationResponse, + SpeechBody, + TranscriptionForm, + TranscriptionResponse, +) +from pydantic import BaseModel, TypeAdapter, ValidationError + +pytestmark = [pytest.mark.e2e, pytest.mark.otel_v2] + +WEATHER_WAV: Final = ( + Path(__file__).resolve().parent.parent / "llm_translation" / "realtime" / "fixtures" / "weather_question_24k.wav" +) +BOUNDED_OUTPUT_CHARS: Final = 1024 + + +class _OutputMessage(BaseModel): + """One assistant message of the Langfuse generation output; only the text is read.""" + + content: str = "" + + +_OUTPUT_MESSAGES: Final = TypeAdapter(list[_OutputMessage]) + + +@pytest.fixture(scope="session") +def langfuse_creds() -> LangfuseCreds: + return load_langfuse_creds() + + +def _langfuse_key( + client: LoggingClient, creds: LangfuseCreds, resources: ResourceManager, params: LiteLLMParamsBody +) -> tuple[str, str, str]: + """A model registered for this run plus a key on a team whose Langfuse callback is `creds`.""" + model: Final = f"e2e-otel-out-{unique_marker()}" + model_id: Final = client.proxy.create_model(model, params) + resources.defer(lambda: client.proxy.delete_model(model_id)) + team_id: Final = client.create_team(f"otel-out-team-{unique_marker()}", models=[model]) + resources.defer(lambda: client.delete_team(team_id)) + client.add_team_langfuse_callback(team_id, creds) + alias: Final = f"otel-out-key-{unique_marker()}" + key: Final = client.key_with_alias(alias, models=[model], team_id=team_id) + resources.defer(lambda: client.delete_key(key)) + return model, key, alias + + +def _generation(client: LoggingClient, creds: LangfuseCreds, *, alias: str, started: datetime) -> LangfuseObservation: + since: Final = (started - timedelta(seconds=5)).isoformat() + observation: Final = client.poll_langfuse_generation(creds, key_alias=alias, from_start_time=since) + assert observation is not None, f"no OTel v2 generation reached Langfuse for key alias {alias!r}" + return observation + + +def _output_text(observation: LangfuseObservation) -> str: + assert observation.output not in (None, "", [], {}), f"generation output is empty: {observation!r}" + try: + messages: Final = _OUTPUT_MESSAGES.validate_python(observation.output) + except ValidationError: + pytest.fail(f"generation output is not a list of assistant messages: {observation!r}") + assert messages, f"generation output is empty: {observation!r}" + return "\n".join(message.content for message in messages) + + +def _openai(model: str) -> LiteLLMParamsBody: + return LiteLLMParamsBody(model=model, api_key="os.environ/OPENAI_API_KEY") + + +class TestOtelV2LangfuseGenerationOutput: + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["completions"]) + def test_completions_output_is_the_completion_text( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + model, key, alias = _langfuse_key(client, langfuse_creds, resources, _openai("openai/gpt-3.5-turbo-instruct")) + started: Final = datetime.now(timezone.utc) + response: Final = unwrap( + client.proxy.transport.post( + "/v1/completions", + headers=client.proxy.transport.bearer(key), + json=CompletionBody(model=model, prompt=f"Repeat exactly: {unique_marker()}", n=2), + response_type=CompletionResponse, + ) + ) + texts: Final = tuple(choice.text.strip() for choice in response.choices) + assert len(texts) == 2 and all(texts), f"/v1/completions returned no text: {response!r}" + + output: Final = _output_text(_generation(client, langfuse_creds, alias=alias, started=started)) + assert all(text in output for text in texts), ( + f"generation output lacks the completion texts {texts!r}: {output!r}" + ) + + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["images_generations"]) + def test_images_output_is_a_bounded_summary_without_base64( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + model, key, alias = _langfuse_key(client, langfuse_creds, resources, _openai("openai/gpt-image-1-mini")) + started: Final = datetime.now(timezone.utc) + response: Final = unwrap( + client.proxy.transport.post( + "/v1/images/generations", + headers=client.proxy.transport.bearer(key), + json=ImageGenerationBody(model=model, prompt=f"a plain red square {unique_marker()}"), + response_type=ImageGenerationResponse, + timeout=180.0, + ) + ) + assert response.data, f"/v1/images/generations returned no data: {response!r}" + encoded: Final = response.data[0].b64_json or "" + assert encoded, f"expected a b64_json image from gpt-image-1-mini: {response.data[0].url!r}" + image_bytes: Final = len(base64.b64decode(encoded)) + + output: Final = _output_text(_generation(client, langfuse_creds, alias=alias, started=started)) + assert len(output) <= BOUNDED_OUTPUT_CHARS, f"image generation output is not bounded ({len(output)} chars)" + assert encoded[:64] not in output, "image generation output leaks the raw base64 payload" + assert output == f"b64_json image ({image_bytes} bytes)", ( + f"image generation output does not report the {image_bytes} decoded bytes: {output!r}" + ) + + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["audio_speech"]) + def test_speech_output_is_a_bounded_summary_without_audio_bytes( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + model, key, alias = _langfuse_key(client, langfuse_creds, resources, _openai("openai/gpt-4o-mini-tts")) + started: Final = datetime.now(timezone.utc) + audio: Final = client.proxy.transport.stream_binary( + "/v1/audio/speech", + headers=client.proxy.transport.bearer(key), + json=SpeechBody(model=model, input=f"hello {unique_marker()}"), + ) + assert audio.ok and audio.total_bytes > 0, f"/v1/audio/speech returned no audio: {audio!r}" + + output: Final = _output_text(_generation(client, langfuse_creds, alias=alias, started=started)) + assert len(output) <= BOUNDED_OUTPUT_CHARS, f"speech output is not bounded ({len(output)} chars)" + assert output.endswith(f" ({audio.total_bytes} bytes)"), ( + f"speech output does not report the {audio.total_bytes} audio bytes the caller received: {output!r}" + ) + + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["audio_transcriptions"]) + def test_transcription_output_is_the_transcript( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + model, key, alias = _langfuse_key(client, langfuse_creds, resources, _openai("openai/gpt-4o-mini-transcribe")) + started: Final = datetime.now(timezone.utc) + response: Final = unwrap( + client.proxy.transport.upload( + "/v1/audio/transcriptions", + headers=client.proxy.transport.bearer(key), + form=TranscriptionForm(model=model), + filename=WEATHER_WAV.name, + content=WEATHER_WAV.read_bytes(), + file_content_type="audio/wav", + response_type=TranscriptionResponse, + ) + ) + transcript: Final = response.text.strip() + assert transcript, f"/v1/audio/transcriptions returned no text: {response!r}" + + output: Final = _output_text(_generation(client, langfuse_creds, alias=alias, started=started)) + assert transcript in output, f"generation output lacks the transcript {transcript!r}: {output!r}" + + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["moderations"]) + def test_moderations_output_is_the_verdict( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + model, key, alias = _langfuse_key(client, langfuse_creds, resources, _openai("openai/omni-moderation-latest")) + started: Final = datetime.now(timezone.utc) + response: Final = unwrap( + client.proxy.transport.post( + "/v1/moderations", + headers=client.proxy.transport.bearer(key), + json=ModerationBody(model=model, input=f"I will find you and hurt you badly {unique_marker()}"), + response_type=ModerationResponse, + ) + ) + assert response.results, f"/v1/moderations returned no results: {response!r}" + verdict: Final = "flagged: " if response.results[0].flagged else "not flagged" + + output: Final = _output_text(_generation(client, langfuse_creds, alias=alias, started=started)) + assert output.startswith(verdict), ( + f"generation output does not carry the moderation verdict {verdict!r}: {output!r}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 355329585fb..d0b1e8824da 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -60,6 +60,7 @@ class KeyMetadata(BaseModel): priority: str | None = None batch_enqueued_token_limit: int | None = None tag: str | None = None + guardrails: list[str] | None = None class ObjectPermission(BaseModel): @@ -697,6 +698,20 @@ class EmbedResponse(BaseModel): model: str | None = None +# ---------- videos ---------- + + +class VideoCreateBody(BaseModel): + model: str + prompt: str + seconds: str | None = None + + +class VideoCreateResponse(BaseModel): + id: str + status: str | None = None + + # ---------- rerank ---------- @@ -745,6 +760,77 @@ class OcrResponse(BaseModel): pages: list[OcrPage] = [] +# ---------- completions ---------- + + +class CompletionBody(BaseModel): + model: str + prompt: str + max_tokens: int = 8 + n: int = 1 + + +class CompletionChoice(BaseModel): + text: str = "" + + +class CompletionResponse(BaseModel): + choices: list[CompletionChoice] = [] + + +# ---------- images ---------- + + +class ImageGenerationBody(BaseModel): + model: str + prompt: str + n: int = 1 + size: str = "1024x1024" + quality: str = "low" + + +class ImageDatum(BaseModel): + url: str | None = None + b64_json: str | None = None + + +class ImageGenerationResponse(BaseModel): + data: list[ImageDatum] = [] + + +# ---------- audio ---------- + + +class SpeechBody(BaseModel): + model: str + input: str + voice: str = "alloy" + + +class TranscriptionForm(BaseModel): + model: str + + +class TranscriptionResponse(BaseModel): + text: str = "" + + +# ---------- moderations ---------- + + +class ModerationBody(BaseModel): + model: str + input: str + + +class ModerationResult(BaseModel): + flagged: bool + + +class ModerationResponse(BaseModel): + results: list[ModerationResult] = [] + + # ---------- spend logs ---------- @@ -936,6 +1022,23 @@ class RouterSettingsResponse(BaseModel): current_values: RouterCurrentValues +class ConfigListParams(BaseModel): + config_type: Literal["general_settings"] + + +class ConfigField(BaseModel): + """One row of GET /config/list: a general_settings field and the value the + proxy is running with, the two fields a test preconditions on.""" + + model_config = ConfigDict(extra="ignore") + field_name: str + field_value: JsonValue = None + + +class ConfigFieldList(RootModel[tuple[ConfigField, ...]]): + """GET /config/list answers with a bare array of general_settings fields.""" + + class CostMapEntry(BaseModel): model_config = ConfigDict(extra="ignore") litellm_provider: str | None = None @@ -1051,6 +1154,7 @@ class ModelInfoBody(BaseModel): mode: ModelMode | None = None access_groups: list[str] | None = None team_id: str | None = None + allowed_fails: int | None = None allowed_fails_policy: dict[str, int] | None = None diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 2f32361e083..76c57452c69 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -47,6 +47,8 @@ from models import ( AnthropicMessagesResponse, ChatBody, ChatResponse, + ConfigFieldList, + ConfigListParams, CostMap, CostMapEntry, CountTokensBody, @@ -630,6 +632,19 @@ class ProxyClient: provider_live=provider_live, ) + def general_setting_enabled(self, field_name: str) -> bool: + """Whether the proxy is running with the named general_settings flag on, for + a test whose behavior only exists under a config flag the stack has to carry.""" + fields = unwrap( + self.transport.get( + "/config/list", + headers=self.transport.master, + params=ConfigListParams(config_type="general_settings"), + response_type=ConfigFieldList, + ) + ).root + return any(entry.field_name == field_name and entry.field_value is True for entry in fields) + def register_model( self, body: ModelNewBody, listed_for: str | None = None, *, provider_live: bool = False ) -> str: diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index c6acd449884..6f9f57d333e 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -14,3 +14,4 @@ markers = redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set + otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 976c05ffceb..3d5b76f6408 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -42,7 +42,7 @@ REAL_KEY = "os.environ/OPENAI_API_KEY" CACHING_MODEL = "anthropic/claude-haiku-4-5" CACHING_KEY = "os.environ/ANTHROPIC_API_KEY" -CONTENT_FILTERED_MODEL = "azure/gpt-5.4-nano" +AZURE_MODEL = "azure/gpt-5.4-nano" AZURE_KEY = "os.environ/AZURE_API_KEY" AZURE_BASE = "os.environ/AZURE_API_BASE" AZURE_API_VERSION = "2024-10-21" @@ -53,6 +53,7 @@ CONTENT_POLICY_PROMPT = ( ) COOLDOWN_SECONDS = 30.0 +REPLICA_PROPAGATION_SECONDS = 15.0 # The smallest-context chat model OpenAI still serves (16385 tokens). A prompt # past that limit comes back as a real `context_length_exceeded` 400, which is @@ -111,7 +112,7 @@ def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str: return proxy.create_model( name, LiteLLMParamsBody( - model=CONTENT_FILTERED_MODEL, + model=AZURE_MODEL, api_key=AZURE_KEY, api_base=AZURE_BASE, api_version=AZURE_API_VERSION, @@ -120,6 +121,26 @@ def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str: ) +def create_azure_benched_on_first_failure_deployment(proxy: ProxyClient, name: str, cooldown_time: float) -> str: + """The live Azure OpenAI deployment holding all of the group's shuffle weight, + benched on its first failure of any class, with the client's own retries off.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody( + model=AZURE_MODEL, + api_key=AZURE_KEY, + api_base=AZURE_BASE, + api_version=AZURE_API_VERSION, + max_retries=0, + weight=1, + cooldown_time=cooldown_time, + ), + model_info=ModelInfoBody(allowed_fails=0), + ) + ) + + def create_caching_deployment(proxy: ProxyClient, name: str) -> str: """Register the Anthropic deployment whose prompt cache the affinity check pins to.""" return proxy.create_model(name, LiteLLMParamsBody(model=CACHING_MODEL, api_key=CACHING_KEY, weight=1)) diff --git a/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py new file mode 100644 index 00000000000..06174e97d20 --- /dev/null +++ b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py @@ -0,0 +1,139 @@ +"""Live e2e: a client hanging up mid-request under cancel_on_disconnect never +benches the deployment it was talking to. + +The group is the cooldown suite's pair: the live Azure deployment holding all of +the shuffle weight, benched on its first failure of any class with a cooldown that +outlasts the test, plus a healthy backup at weight 0 the shuffle only reaches once +the Azure deployment is benched. A cheap call first proves the Azure deployment +answers the key and warms its auth path. The test then asks for an answer far +longer than CLIENT_HANGS_UP_AFTER_SECONDS of generation, retries off, and hangs up +that many seconds in: late enough that the proxy has handed the call to Azure (a +hang-up before the provider call is in flight cancels nothing the router could +bench, so the cell would pass vacuously). An answer that comes back inside the +window proves nothing and benches nothing either, since a success never counts +against the deployment, so the cell asks again up to HANG_UP_ATTEMPTS times and +fails out loud naming the window only when every ask came back early. After the +cooldown suite's replica propagation window, every one of the next calls has to +come back 200 from the Azure deployment itself, named in x-litellm-model-id; a +single answer from the backup means the hang-up was booked as a failure. + +The test reads `cancel_on_disconnect` back from the proxy first: without the flag +the hang-up cancels nothing and the cell would pass vacuously. +""" + +from __future__ import annotations + +import time + +import pytest +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from e2e_http import AbandonedRequest, StreamingResponse +from lifecycle import ResourceManager +from models import ChatMessage, ReliabilityChatBody, RouterSettingsOverride +from reliability_support import ( + REPLICA_PROPAGATION_SECONDS, + chat_override, + create_azure_benched_on_first_failure_deployment, + create_zero_weight_backup_deployment, + model_id_of, +) + +pytestmark = pytest.mark.e2e + +CLIENT_HANGS_UP_AFTER_SECONDS = 5.0 +HANG_UP_ATTEMPTS = 3 +LONG_ANSWER_MAX_TOKENS = 16384 +BENCH_OUTLASTS_TEST_SECONDS = 300.0 +CALLS_AFTER_HANGUP = 6 + + +def _say_hi(client: ComplexityRouterClient, key: str, group: str) -> StreamingResponse: + return chat_override( + client.proxy, + key, + group, + f"say hi {unique_marker()}", + override=RouterSettingsOverride(num_retries=0), + ) + + +def _ask_for_a_long_answer_then_hang_up( + client: ComplexityRouterClient, key: str, group: str +) -> AbandonedRequest | StreamingResponse: + return client.proxy.transport.abandon( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ReliabilityChatBody( + model=group, + messages=[ + ChatMessage( + role="user", + content=( + "Write an essay on the history of the telegraph with one section per decade from the 1830s " + f"to the 2020s, each section at least 300 words. {unique_marker()}" + ), + ) + ], + max_tokens=LONG_ANSWER_MAX_TOKENS, + router_settings_override=RouterSettingsOverride(num_retries=0), + ), + after=CLIENT_HANGS_UP_AFTER_SECONDS, + ) + + +def _hang_up_mid_answer(client: ComplexityRouterClient, key: str, group: str) -> None: + for attempt in range(1, HANG_UP_ATTEMPTS + 1): + match _ask_for_a_long_answer_then_hang_up(client, key, group): + case AbandonedRequest(): + return + case StreamingResponse(status_code=200): + continue + case StreamingResponse(status_code=status_code, body=body): + pytest.fail( + f"hang-up attempt {attempt} should have found the long answer still in flight after " + f"{CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s, but the proxy answered {status_code}: {body[:300]}" + ) + pytest.fail( + f"the proxy answered all {HANG_UP_ATTEMPTS} long asks within {CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s, so the " + "client never hung up with a call still in flight and the bench this cell guards against could not happen" + ) + + +class TestReliabilityCancelOnDisconnect: + @pytest.mark.covers("reliability.cooldown.client_disconnect.stays_healthy") + def test_client_hanging_up_never_benches_the_deployment( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + assert client.proxy.general_setting_enabled("cancel_on_disconnect"), ( + "this cell needs general_settings.cancel_on_disconnect: true in the proxy config; without it the " + "hang-up cancels nothing and the bench it guards against can never happen" + ) + + group = f"reliability-cooldown-disconnect-{unique_marker()}" + azure = create_azure_benched_on_first_failure_deployment( + client.proxy, group, cooldown_time=BENCH_OUTLASTS_TEST_SECONDS + ) + resources.defer(lambda: client.proxy.delete_model(azure)) + backup = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup)) + + warm_up = _say_hi(client, scoped_key, group) + assert warm_up.status_code == 200 and model_id_of(warm_up) == azure, ( + f"before any hang-up the Azure deployment {azure} should answer the group, got {warm_up.status_code} " + f"from {model_id_of(warm_up)!r}: {warm_up.body[:300]}" + ) + + _hang_up_mid_answer(client, scoped_key, group) + time.sleep(REPLICA_PROPAGATION_SECONDS) + + for call in range(1, CALLS_AFTER_HANGUP + 1): + resp = _say_hi(client, scoped_key, group) + assert resp.status_code == 200, ( + f"call {call} after the hang-up should have been a plain 200 from the group, got " + f"{resp.status_code}: {resp.body[:300]}" + ) + assert model_id_of(resp) == azure, ( + f"call {call} after the hang-up should have been served by the Azure deployment {azure}, the proxy " + f"named {model_id_of(resp)!r}: the cancelled call was booked as a failure and benched it" + ) diff --git a/tests/e2e/router/test_reliability_cooldowns_e2e.py b/tests/e2e/router/test_reliability_cooldowns_e2e.py index 5b5cec09f06..769971e1533 100644 --- a/tests/e2e/router/test_reliability_cooldowns_e2e.py +++ b/tests/e2e/router/test_reliability_cooldowns_e2e.py @@ -43,6 +43,7 @@ from lifecycle import ResourceManager from models import KeyGenerateBody, RouterSettingsOverride from reliability_support import ( COOLDOWN_SECONDS, + REPLICA_PROPAGATION_SECONDS, chat_override, create_always_5xx_deployment, create_always_rate_limited_deployment, @@ -57,7 +58,6 @@ from reliability_support import ( pytestmark = pytest.mark.e2e RECOVERY_GRACE_SECONDS = 10 -REPLICA_PROPAGATION_SECONDS = 15.0 PROPAGATION_POLL_SECONDS = 0.25 BENCH_MARGIN_SECONDS = 4.0 diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 0022c0c4355..a3eec815441 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -13,6 +13,7 @@ from typing import Protocol import e2e_http from e2e_http import ( URL, + AbandonedRequest, AuthHeaders, BinaryStream, NetworkError, @@ -58,6 +59,10 @@ class Transport(Protocol): stream: bool = False, ) -> StreamingResponse: ... + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: ... + def get[R: BaseModel]( self, path: str, @@ -243,6 +248,11 @@ class HttpTransport: timeout=self.request_timeout, ) + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: + return e2e_http.abandon(self._url(path), headers=headers, json=json, after=after) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return e2e_http.probe( self._url(path), @@ -420,6 +430,11 @@ class SplitTransport: ) -> StreamingResponse: return self._route(path).send(path, headers=headers, json=json, params=params, stream=stream) + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: + return self._route(path).abandon(path, headers=headers, json=json, after=after) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return self._route(path).probe(path, params=params, headers=headers) diff --git a/tests/e2e/ui/helpers/mcp.ts b/tests/e2e/ui/helpers/mcp.ts index 554177e11bc..399f746f3a8 100644 --- a/tests/e2e/ui/helpers/mcp.ts +++ b/tests/e2e/ui/helpers/mcp.ts @@ -1,8 +1,21 @@ import { expect, Page as PwPage } from "@playwright/test"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import { navigateToPage } from "./navigation"; import { Page } from "../fixtures/pages"; import { masterKey } from "./traffic"; +export async function listUpstreamToolNames(url: string): Promise { + const client = new Client({ name: "litellm-ui-e2e", version: "0.0.0" }); + await client.connect(new StreamableHTTPClientTransport(new URL(url))); + try { + const { tools } = await client.listTools(); + return tools.map((tool) => tool.name); + } finally { + await client.close(); + } +} + /** Creates an MCP server through the UI's discovery to custom-form flow and returns its name. */ export async function createMcpServer(page: PwPage, url: string): Promise { await navigateToPage(page, Page.McpServers); diff --git a/tests/e2e/ui/package-lock.json b/tests/e2e/ui/package-lock.json index b22673a3535..f56e00506e9 100644 --- a/tests/e2e/ui/package-lock.json +++ b/tests/e2e/ui/package-lock.json @@ -8,11 +8,66 @@ "name": "litellm-ui-e2e", "version": "0.0.0", "devDependencies": { + "@modelcontextprotocol/sdk": "1.30.0", "@playwright/test": "1.58.1", "@types/node": "20.19.37", "typescript": "5.9.3" } }, + "node_modules/@hono/node-server": { + "version": "2.1.1", + "resolved": "https://registry.npmjs.org/@hono/node-server/-/node-server-2.1.1.tgz", + "integrity": "sha512-ELuehkj5VCBdgEw9zs+ivkKwyzzUCSQuE96YmiPvn1ECBoZCczbFXJLeEGMTYjphP6gydh4pHMqEYPVMYUVgQg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=20" + }, + "peerDependencies": { + "hono": "^4" + } + }, + "node_modules/@modelcontextprotocol/sdk": { + "version": "1.30.0", + "resolved": "https://registry.npmjs.org/@modelcontextprotocol/sdk/-/sdk-1.30.0.tgz", + "integrity": "sha512-xKd8OIzlqNzcqcNumGAa6g+PW2kjD5vrpcKOnfldAUPP3j7lnqMPwlTXQm8gF+UwH72z0lqaRbjr9hqGz0eITA==", + "dev": true, + "license": "MIT", + "dependencies": { + "@hono/node-server": "^1.19.9 || ^2.0.5", + "ajv": "^8.17.1", + "ajv-formats": "^3.0.1", + "content-type": "^1.0.5", + "cors": "^2.8.5", + "cross-spawn": "^7.0.5", + "eventsource": "^3.0.2", + "eventsource-parser": "^3.0.0", + "express": "^5.2.1", + "express-rate-limit": "^8.2.1", + "hono": "^4.11.4", + "jose": "^6.1.3", + "json-schema-typed": "^8.0.2", + "pkce-challenge": "^5.0.0", + "raw-body": "^3.0.0", + "zod": "^3.25 || ^4.0", + "zod-to-json-schema": "^3.25.1" + }, + "engines": { + "node": ">=18" + }, + "peerDependencies": { + "@cfworker/json-schema": "^4.1.1", + "zod": "^3.25 || ^4.0" + }, + "peerDependenciesMeta": { + "@cfworker/json-schema": { + "optional": true + }, + "zod": { + "optional": false + } + } + }, "node_modules/@playwright/test": { "version": "1.58.1", "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.1.tgz", @@ -39,6 +94,475 @@ "undici-types": "~6.21.0" } }, + "node_modules/accepts": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/accepts/-/accepts-2.0.0.tgz", + "integrity": "sha512-5cvg6CtKwfgdmVqY1WIiXKc3Q1bkRqGLi+2W/6ao+6Y7gu/RCwRuAhGEzh5B4KlszSuTLgZYuqFqo5bImjNKng==", + "dev": true, + "license": "MIT", + "dependencies": { + "mime-types": "^3.0.0", + "negotiator": "^1.0.0" + }, + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/ajv": { + "version": "8.20.0", + "resolved": "https://registry.npmjs.org/ajv/-/ajv-8.20.0.tgz", + "integrity": "sha512-Thbli+OlOj+iMPYFBVBfJ3OmCAnaSyNn4M1vz9T6Gka5Jt9ba/HIR56joy65tY6kx/FCF5VXNB819Y7/GUrBGA==", + "dev": true, + "license": "MIT", + "dependencies": { + "fast-deep-equal": "^3.1.3", + "fast-uri": "^3.0.1", + "json-schema-traverse": "^1.0.0", + "require-from-string": "^2.0.2" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/epoberezkin" + } + }, + "node_modules/ajv-formats": { + "version": "3.0.1", + "resolved": "https://registry.npmjs.org/ajv-formats/-/ajv-formats-3.0.1.tgz", + "integrity": "sha512-8iUql50EUR+uUcdRQ3HDqa6EVyo3docL8g5WJ3FNcWmu62IbkGUue/pEyLBW8VGKKucTPgqeks4fIU1DA4yowQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "ajv": "^8.0.0" + }, + "peerDependencies": { + "ajv": "^8.0.0" + }, + "peerDependenciesMeta": { + "ajv": { + "optional": true + } + } + }, + "node_modules/body-parser": { + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-2.3.0.tgz", + "integrity": "sha512-2cGmJupaNgg+QUwVLAucDuWuoMZ6EX9iHDRswZ5lsNYEmwPaRknMPCLZz07yTzVq/83p4o/wzbDZbBrTvGGTIw==", + "dev": true, + "license": "MIT", + "dependencies": { + "bytes": "^3.1.2", + "content-type": "^2.0.0", + "debug": "^4.4.3", + "http-errors": "^2.0.1", + "iconv-lite": "^0.7.2", + "on-finished": "^2.4.1", + "qs": "^6.15.2", + "raw-body": "^3.0.2", + "type-is": "^2.1.0" + }, + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/body-parser/node_modules/content-type": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/content-type/-/content-type-2.1.0.tgz", + "integrity": "sha512-mj7UPXE0jaqaOsukNZRUEfEi2AcL7C/vwmwcHV0O97eO1E1pxBZuyjlZrx5seTaNBg1U6+o35wpa35Qfcc+7ag==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/bytes": { + "version": "3.1.2", + "resolved": "https://registry.npmjs.org/bytes/-/bytes-3.1.2.tgz", + "integrity": "sha512-/Nf7TyzTx6S3yRJObOAV7956r8cr2+Oj8AC5dt8wSP3BQAoeX58NoHyCU8P8zGkNXStjTSi6fzO6F0pBdcYbEg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/call-bind-apply-helpers": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/call-bind-apply-helpers/-/call-bind-apply-helpers-1.0.2.tgz", + "integrity": "sha512-Sp1ablJ0ivDkSzjcaJdxEunN5/XvksFJ2sMBFfq6x0ryhQV/2b/KwFe21cMpmHtPOSij8K99/wSfoEuTObmuMQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "es-errors": "^1.3.0", + "function-bind": "^1.1.2" + }, + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/call-bound": { + "version": "1.0.4", + "resolved": "https://registry.npmjs.org/call-bound/-/call-bound-1.0.4.tgz", + "integrity": "sha512-+ys997U96po4Kx/ABpBCqhA9EuxJaQWDQg7295H4hBphv3IZg0boBKuwYpt4YXp6MZ5AmZQnU/tyMTlRpaSejg==", + "dev": true, + "license": "MIT", + "dependencies": { + "call-bind-apply-helpers": "^1.0.2", + "get-intrinsic": "^1.3.0" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/content-disposition": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/content-disposition/-/content-disposition-1.1.0.tgz", + "integrity": "sha512-5jRCH9Z/+DRP7rkvY83B+yGIGX96OYdJmzngqnw2SBSxqCFPd0w2km3s5iawpGX8krnwSGmF0FW5Nhr0Hfai3g==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/content-type": { + "version": "1.0.5", + "resolved": "https://registry.npmjs.org/content-type/-/content-type-1.0.5.tgz", + "integrity": "sha512-nTjqfcBFEipKdXCv4YDQWCfmcLZKm81ldF0pAopTvyrFGVbcR6P/VAAd5G7N+0tTr8QqiU0tFadD6FK4NtJwOA==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/cookie": { + "version": "0.7.2", + "resolved": "https://registry.npmjs.org/cookie/-/cookie-0.7.2.tgz", + "integrity": "sha512-yki5XnKuf750l50uGTllt6kKILY4nQ1eNIQatoXEByZ5dWgnKqbnqmTrBE5B4N7lrMJKQ2ytWMiTO2o0v6Ew/w==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/cookie-signature": { + "version": "1.2.2", + "resolved": "https://registry.npmjs.org/cookie-signature/-/cookie-signature-1.2.2.tgz", + "integrity": "sha512-D76uU73ulSXrD1UXF4KE2TMxVVwhsnCgfAyTg9k8P6KGZjlXKrOLe4dJQKI3Bxi5wjesZoFXJWElNWBjPZMbhg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=6.6.0" + } + }, + "node_modules/cors": { + "version": "2.8.6", + "resolved": "https://registry.npmjs.org/cors/-/cors-2.8.6.tgz", + "integrity": "sha512-tJtZBBHA6vjIAaF6EnIaq6laBBP9aq/Y3ouVJjEfoHbRBcHBAHYcMh/w8LDrk2PvIMMq8gmopa5D4V8RmbrxGw==", + "dev": true, + "license": "MIT", + "dependencies": { + "object-assign": "^4", + "vary": "^1" + }, + "engines": { + "node": ">= 0.10" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/cross-spawn": { + "version": "7.0.6", + "resolved": "https://registry.npmjs.org/cross-spawn/-/cross-spawn-7.0.6.tgz", + "integrity": "sha512-uV2QOWP2nWzsy2aMp8aRibhi9dlzF5Hgh5SHaB9OiTGEyDTiJJyx0uy51QXdyWbtAHNua4XJzUKca3OzKUd3vA==", + "dev": true, + "license": "MIT", + "dependencies": { + "path-key": "^3.1.0", + "shebang-command": "^2.0.0", + "which": "^2.0.1" + }, + "engines": { + "node": ">= 8" + } + }, + "node_modules/debug": { + "version": "4.4.3", + "resolved": "https://registry.npmjs.org/debug/-/debug-4.4.3.tgz", + "integrity": "sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==", + "dev": true, + "license": "MIT", + "dependencies": { + "ms": "^2.1.3" + }, + "engines": { + "node": ">=6.0" + }, + "peerDependenciesMeta": { + "supports-color": { + "optional": true + } + } + }, + "node_modules/depd": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", + "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/dunder-proto": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", + "integrity": "sha512-KIN/nDJBQRcXw0MLVhZE9iQHmG68qAVIBg9CqmUYjmQIhgij9U5MFvrqkUL5FbtyyzZuOeOt0zdeRe4UY7ct+A==", + "dev": true, + "license": "MIT", + "dependencies": { + "call-bind-apply-helpers": "^1.0.1", + "es-errors": "^1.3.0", + "gopd": "^1.2.0" + }, + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/ee-first": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/ee-first/-/ee-first-1.1.1.tgz", + "integrity": "sha512-WMwm9LhRUo+WUaRN+vRuETqG89IgZphVSNkdFgeb6sS/E4OrDIN7t48CAewSHXc6C8lefD8KKfr5vY61brQlow==", + "dev": true, + "license": "MIT" + }, + "node_modules/encodeurl": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/encodeurl/-/encodeurl-2.0.0.tgz", + "integrity": "sha512-Q0n9HRi4m6JuGIV1eFlmvJB7ZEVxu93IrMyiMsGC0lrMJMWzRgx6WGquyfQgZVb31vhGgXnfmPNNXmxnOkRBrg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/es-define-property": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/es-define-property/-/es-define-property-1.0.1.tgz", + "integrity": "sha512-e3nRfgfUZ4rNGL232gUgX06QNyyez04KdjFrF+LTRoOXmrOgFKDg4BCdsjW8EnT69eqdYGmRpJwiPVYNrCaW3g==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/es-errors": { + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/es-errors/-/es-errors-1.3.0.tgz", + "integrity": "sha512-Zf5H2Kxt2xjTvbJvP2ZWLEICxA6j+hAmMzIlypy4xcBg1vKVnx89Wy0GbS+kf5cwCVFFzdCFh2XSCFNULS6csw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/es-object-atoms": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/es-object-atoms/-/es-object-atoms-1.1.2.tgz", + "integrity": "sha512-HWcBoN6NileqtSydK2FqHbS/LoDd2pqrnQHLyJzBj4kOp/ky2MWMN694xOfkK8/SnUsW2DH7EfyVlydKCsm1Zw==", + "dev": true, + "license": "MIT", + "dependencies": { + "es-errors": "^1.3.0" + }, + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/escape-html": { + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/escape-html/-/escape-html-1.0.3.tgz", + "integrity": "sha512-NiSupZ4OeuGwr68lGIeym/ksIZMJodUGOSCZ/FSnTxcrekbvqrgdUxlJOMpijaKZVjAJrWrGs/6Jy8OMuyj9ow==", + "dev": true, + "license": "MIT" + }, + "node_modules/etag": { + "version": "1.8.1", + "resolved": "https://registry.npmjs.org/etag/-/etag-1.8.1.tgz", + "integrity": "sha512-aIL5Fx7mawVa300al2BnEE4iNvo1qETxLrPI/o05L7z6go7fCw1J6EQmbK4FmJ2AS7kgVF/KEZWufBfdClMcPg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/eventsource": { + "version": "3.0.7", + "resolved": "https://registry.npmjs.org/eventsource/-/eventsource-3.0.7.tgz", + "integrity": "sha512-CRT1WTyuQoD771GW56XEZFQ/ZoSfWid1alKGDYMmkt2yl8UXrVR4pspqWNEcqKvVIzg6PAltWjxcSSPrboA4iA==", + "dev": true, + "license": "MIT", + "dependencies": { + "eventsource-parser": "^3.0.1" + }, + "engines": { + "node": ">=18.0.0" + } + }, + "node_modules/eventsource-parser": { + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/eventsource-parser/-/eventsource-parser-3.1.1.tgz", + "integrity": "sha512-EKN1vKAMcZ8MlYMpaNuxN6R9yakzH6uajHcHVTqWJzvu5pWw9DyhbP35HH8MVBQ+dZjAfDxk+A8NiR9KWaXiyQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18.0.0" + } + }, + "node_modules/express": { + "version": "5.2.1", + "resolved": "https://registry.npmjs.org/express/-/express-5.2.1.tgz", + "integrity": "sha512-hIS4idWWai69NezIdRt2xFVofaF4j+6INOpJlVOLDO8zXGpUVEVzIYk12UUi2JzjEzWL3IOAxcTubgz9Po0yXw==", + "dev": true, + "license": "MIT", + "dependencies": { + "accepts": "^2.0.0", + "body-parser": "^2.2.1", + "content-disposition": "^1.0.0", + "content-type": "^1.0.5", + "cookie": "^0.7.1", + "cookie-signature": "^1.2.1", + "debug": "^4.4.0", + "depd": "^2.0.0", + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "etag": "^1.8.1", + "finalhandler": "^2.1.0", + "fresh": "^2.0.0", + "http-errors": "^2.0.0", + "merge-descriptors": "^2.0.0", + "mime-types": "^3.0.0", + "on-finished": "^2.4.1", + "once": "^1.4.0", + "parseurl": "^1.3.3", + "proxy-addr": "^2.0.7", + "qs": "^6.14.0", + "range-parser": "^1.2.1", + "router": "^2.2.0", + "send": "^1.1.0", + "serve-static": "^2.2.0", + "statuses": "^2.0.1", + "type-is": "^2.0.1", + "vary": "^1.1.2" + }, + "engines": { + "node": ">= 18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/express-rate-limit": { + "version": "8.7.0", + "resolved": "https://registry.npmjs.org/express-rate-limit/-/express-rate-limit-8.7.0.tgz", + "integrity": "sha512-hOwV7WOxXfjRpAM1DSJWZDXx3GhplwD8IfwuwvogD8i1Qnkgosw/H45s4ZnFAUHDAhPjlY9hLBvJhKmGMyY26g==", + "dev": true, + "license": "MIT", + "dependencies": { + "debug": "^4.4.3", + "ip-address": "^10.2.0" + }, + "engines": { + "node": ">= 16" + }, + "funding": { + "url": "https://github.com/sponsors/express-rate-limit" + }, + "peerDependencies": { + "express": ">= 4.11" + } + }, + "node_modules/fast-deep-equal": { + "version": "3.1.3", + "resolved": "https://registry.npmjs.org/fast-deep-equal/-/fast-deep-equal-3.1.3.tgz", + "integrity": "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==", + "dev": true, + "license": "MIT" + }, + "node_modules/fast-uri": { + "version": "3.1.8", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.8.tgz", + "integrity": "sha512-GZMtZUTNRpOVIECoXwLNZS5xUGE+mVNbTB8h/7Rwh2TFWcBQiPzTgyZi05BF9UMZKkLJv8XBRJTlU7zg8+ZfMg==", + "dev": true, + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/fastify" + }, + { + "type": "opencollective", + "url": "https://opencollective.com/fastify" + } + ], + "license": "BSD-3-Clause" + }, + "node_modules/finalhandler": { + "version": "2.1.1", + "resolved": "https://registry.npmjs.org/finalhandler/-/finalhandler-2.1.1.tgz", + "integrity": "sha512-S8KoZgRZN+a5rNwqTxlZZePjT/4cnm0ROV70LedRHZ0p8u9fRID0hJUZQpkKLzro8LfmC8sx23bY6tVNxv8pQA==", + "dev": true, + "license": "MIT", + "dependencies": { + "debug": "^4.4.0", + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "on-finished": "^2.4.1", + "parseurl": "^1.3.3", + "statuses": "^2.0.1" + }, + "engines": { + "node": ">= 18.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/forwarded": { + "version": "0.2.0", + "resolved": "https://registry.npmjs.org/forwarded/-/forwarded-0.2.0.tgz", + "integrity": "sha512-buRG0fpBtRHSTCOASe6hD258tEubFoRLb4ZNA6NxMVHNw2gOcwHo9wyablzMzOA5z9xA9L1KNjk/Nt6MT9aYow==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/fresh": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/fresh/-/fresh-2.0.0.tgz", + "integrity": "sha512-Rx/WycZ60HOaqLKAi6cHRKKI7zxWbJ31MhntmtwMoaTeF7XFH9hhBp8vITaMidfljRQ6eYWCKkaTK+ykVJHP2A==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, "node_modules/fsevents": { "version": "2.3.2", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", @@ -54,6 +578,396 @@ "node": "^8.16.0 || ^10.6.0 || >=11.0.0" } }, + "node_modules/function-bind": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz", + "integrity": "sha512-7XHNxH7qX9xG5mIwxkhumTox/MIRNcOgDrxWsMt2pAr23WHp6MrRlN7FBSFpCpr+oVO0F744iUgR82nJMfG2SA==", + "dev": true, + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/get-intrinsic": { + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/get-intrinsic/-/get-intrinsic-1.3.0.tgz", + "integrity": "sha512-9fSjSaos/fRIVIp+xSJlE6lfwhES7LNtKaCBIamHsjr2na1BiABJPo0mOjjz8GJDURarmCPGqaiVg5mfjb98CQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "call-bind-apply-helpers": "^1.0.2", + "es-define-property": "^1.0.1", + "es-errors": "^1.3.0", + "es-object-atoms": "^1.1.1", + "function-bind": "^1.1.2", + "get-proto": "^1.0.1", + "gopd": "^1.2.0", + "has-symbols": "^1.1.0", + "hasown": "^2.0.2", + "math-intrinsics": "^1.1.0" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/get-proto": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/get-proto/-/get-proto-1.0.1.tgz", + "integrity": "sha512-sTSfBjoXBp89JvIKIefqw7U2CCebsc74kiY6awiGogKtoSGbgjYE/G/+l9sF3MWFPNc9IcoOC4ODfKHfxFmp0g==", + "dev": true, + "license": "MIT", + "dependencies": { + "dunder-proto": "^1.0.1", + "es-object-atoms": "^1.0.0" + }, + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/gopd": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/gopd/-/gopd-1.2.0.tgz", + "integrity": "sha512-ZUKRh6/kUFoAiTAtTYPZJ3hw9wNxx+BIBOijnlG9PnrJsCcSjs1wyyD6vJpaYtgnzDrKYRSqf3OO6Rfa93xsRg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/has-symbols": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/has-symbols/-/has-symbols-1.1.0.tgz", + "integrity": "sha512-1cDNdwJ2Jaohmb3sg4OmKaMBwuC48sYni5HUw2DvsC8LjGTLK9h+eb1X6RyuOHe4hT0ULCW68iomhjUoKUqlPQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/hasown": { + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.4.tgz", + "integrity": "sha512-T2UbfbBEF32wiepXIsMlTW9+dDYC6wMh/t/vYA4tuOMKqWz/n3vr1NFSxQiyP+zk2mXsoMA/i/7qV6LKut1t1A==", + "dev": true, + "license": "MIT", + "dependencies": { + "function-bind": "^1.1.2" + }, + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/hono": { + "version": "4.13.8", + "resolved": "https://registry.npmjs.org/hono/-/hono-4.13.8.tgz", + "integrity": "sha512-/Gng7NfoykZl2pjukW5Z6+8Yxm3BPRf86GTbQnt0SbySkvax4fyL4H3HhY1cCpBGmiW9XDRFzRV+CXK2W8QudQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=16.9.0" + } + }, + "node_modules/http-errors": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", + "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "depd": "~2.0.0", + "inherits": "~2.0.4", + "setprototypeof": "~1.2.0", + "statuses": "~2.0.2", + "toidentifier": "~1.0.1" + }, + "engines": { + "node": ">= 0.8" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/iconv-lite": { + "version": "0.7.3", + "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.7.3.tgz", + "integrity": "sha512-IKXpvIzjnC9XTAUbVBcMfGS0EPaIXtW6v+zr+RRp+hqULEpo0owZax6wyRwPOJbWbzjYspQwusTsfVr0ifh4uQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "safer-buffer": ">= 2.1.2 < 3.0.0" + }, + "engines": { + "node": ">=0.10.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/inherits": { + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", + "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", + "dev": true, + "license": "ISC" + }, + "node_modules/ip-address": { + "version": "10.7.2", + "resolved": "https://registry.npmjs.org/ip-address/-/ip-address-10.7.2.tgz", + "integrity": "sha512-7H/2gFSIitxc0hG3nOI1glS8QLo/EHBFFLk8vEUjXY/xu0AdL8jZ9U1IzO2PUm0d2D/ofQcAifb0g6OBkt8U7w==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 12" + } + }, + "node_modules/ipaddr.js": { + "version": "1.9.1", + "resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-1.9.1.tgz", + "integrity": "sha512-0KI/607xoxSToH7GjN1FfSbLoU0+btTicjsQSWQlh/hZykN8KpmMf7uYwPW3R+akZ6R/w18ZlXSHBYXiYUPO3g==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.10" + } + }, + "node_modules/is-promise": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/is-promise/-/is-promise-4.0.0.tgz", + "integrity": "sha512-hvpoI6korhJMnej285dSg6nu1+e6uxs7zG3BYAm5byqDsgJNWwxzM6z6iZiAgQR4TJ30JmBTOwqZUw3WlyH3AQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/isexe": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/isexe/-/isexe-2.0.0.tgz", + "integrity": "sha512-RHxMLp9lnKHGHRng9QFhRCMbYAcVpn69smSGcq3f36xjgVVWThj4qqLbTLlq7Ssj8B+fIQ1EuCEGI2lKsyQeIw==", + "dev": true, + "license": "ISC" + }, + "node_modules/jose": { + "version": "6.2.12", + "resolved": "https://registry.npmjs.org/jose/-/jose-6.2.12.tgz", + "integrity": "sha512-9NiFmJEex0sy2Dk58j2UGBSHgUs2ypF9eZSu4L6vjOX3Dp96Sw1F3uL+H+D1sx02jZZdzUT0HgvCy59CuvXcWw==", + "dev": true, + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/panva" + } + }, + "node_modules/json-schema-traverse": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/json-schema-traverse/-/json-schema-traverse-1.0.0.tgz", + "integrity": "sha512-NM8/P9n3XjXhIZn1lLhkFaACTOURQXjWhV4BA/RnOv8xvgqtqpAX9IO4mRQxSx1Rlo4tqzeqb0sOlruaOy3dug==", + "dev": true, + "license": "MIT" + }, + "node_modules/json-schema-typed": { + "version": "8.0.2", + "resolved": "https://registry.npmjs.org/json-schema-typed/-/json-schema-typed-8.0.2.tgz", + "integrity": "sha512-fQhoXdcvc3V28x7C7BMs4P5+kNlgUURe2jmUT1T//oBRMDrqy1QPelJimwZGo7Hg9VPV3EQV5Bnq4hbFy2vetA==", + "dev": true, + "license": "BSD-2-Clause" + }, + "node_modules/math-intrinsics": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/math-intrinsics/-/math-intrinsics-1.1.0.tgz", + "integrity": "sha512-/IXtbwEk5HTPyEwyKX6hGkYXxM9nbj64B+ilVJnC/R6B0pH5G4V3b0pVbL7DBj4tkhBAppbQUlf6F6Xl9LHu1g==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/media-typer": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/media-typer/-/media-typer-1.1.1.tgz", + "integrity": "sha512-yz3xRaG20c6/BOzvYoDaGtPmGscs7YivItZEEqe6GbwNfHuxu9YNmvnEkMzKldAGY4/80pRcQRZSEnhquk9XuQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/merge-descriptors": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/merge-descriptors/-/merge-descriptors-2.0.0.tgz", + "integrity": "sha512-Snk314V5ayFLhp3fkUREub6WtjBfPdCPY1Ln8/8munuLuiYhsABgBVWsozAG+MWMbVEvcdcpbi9R7ww22l9Q3g==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, + "node_modules/mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", + "dev": true, + "license": "MIT", + "dependencies": { + "mime-db": "^1.54.0" + }, + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/ms": { + "version": "2.1.3", + "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.3.tgz", + "integrity": "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==", + "dev": true, + "license": "MIT" + }, + "node_modules/negotiator": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/negotiator/-/negotiator-1.1.0.tgz", + "integrity": "sha512-NMPBRMJgiQHjbd8phG3Vebdx4kZ1H121rbl5IkMqeOsahptB9BKo/d7oJ3zTXqTgagn2bWlNSXkh0QUGM31RYg==", + "dev": true, + "license": "MIT", + "dependencies": { + "content-type": "^2.1.0" + }, + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/negotiator/node_modules/content-type": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/content-type/-/content-type-2.1.0.tgz", + "integrity": "sha512-mj7UPXE0jaqaOsukNZRUEfEi2AcL7C/vwmwcHV0O97eO1E1pxBZuyjlZrx5seTaNBg1U6+o35wpa35Qfcc+7ag==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/object-assign": { + "version": "4.1.1", + "resolved": "https://registry.npmjs.org/object-assign/-/object-assign-4.1.1.tgz", + "integrity": "sha512-rJgTQnkUnH1sFw8yT6VSU3zD3sWmu6sZhIseY8VX+GRu3P6F7Fu+JNDoXfklElbLJSnc3FUQHVe4cU5hj+BcUg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/object-inspect": { + "version": "1.13.4", + "resolved": "https://registry.npmjs.org/object-inspect/-/object-inspect-1.13.4.tgz", + "integrity": "sha512-W67iLl4J2EXEGTbfeHCffrjDfitvLANg0UlX3wFUUSTx92KXRFegMHUVgSqE+wvhAbi4WqjGg9czysTV2Epbew==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/on-finished": { + "version": "2.4.1", + "resolved": "https://registry.npmjs.org/on-finished/-/on-finished-2.4.1.tgz", + "integrity": "sha512-oVlzkg3ENAhCk2zdv7IJwd/QUD4z2RxRwpkcGY8psCVcCYZNq4wYnVWALHM+brtuJjePWiYF/ClmuDr8Ch5+kg==", + "dev": true, + "license": "MIT", + "dependencies": { + "ee-first": "1.1.1" + }, + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/once": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz", + "integrity": "sha512-lNaJgI+2Q5URQBkccEKHTQOPaXdUxnZZElQTZY0MFUAuaEqe1E+Nyvgdz/aIyNi6Z9MzO5dv1H8n58/GELp3+w==", + "dev": true, + "license": "ISC", + "dependencies": { + "wrappy": "1" + } + }, + "node_modules/parseurl": { + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/parseurl/-/parseurl-1.3.3.tgz", + "integrity": "sha512-CiyeOxFT/JZyN5m0z9PfXw4SCBJ6Sygz1Dpl0wqjlhDEGGBP1GnsUVEL0p63hoG1fcj3fHynXi9NYO4nWOL+qQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/path-key": { + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/path-key/-/path-key-3.1.1.tgz", + "integrity": "sha512-ojmeN0qd+y0jszEtoY48r0Peq5dwMEkIlCOu6Q5f41lfkswXuKtYrhgoTpLnyIcHm24Uhqx+5Tqm2InSwLhE6Q==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=8" + } + }, + "node_modules/path-to-regexp": { + "version": "8.4.2", + "resolved": "https://registry.npmjs.org/path-to-regexp/-/path-to-regexp-8.4.2.tgz", + "integrity": "sha512-qRcuIdP69NPm4qbACK+aDogI5CBDMi1jKe0ry5rSQJz8JVLsC7jV8XpiJjGRLLol3N+R5ihGYcrPLTno6pAdBA==", + "dev": true, + "license": "MIT", + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/pkce-challenge": { + "version": "5.0.1", + "resolved": "https://registry.npmjs.org/pkce-challenge/-/pkce-challenge-5.0.1.tgz", + "integrity": "sha512-wQ0b/W4Fr01qtpHlqSqspcj3EhBvimsdh0KlHhH8HRZnMsEa0ea2fTULOXOS9ccQr3om+GcGRk4e+isrZWV8qQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=16.20.0" + } + }, "node_modules/playwright": { "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.1.tgz", @@ -86,6 +1000,311 @@ "node": ">=18" } }, + "node_modules/proxy-addr": { + "version": "2.0.8", + "resolved": "https://registry.npmjs.org/proxy-addr/-/proxy-addr-2.0.8.tgz", + "integrity": "sha512-5nnx0yGyVUcY6t9RnWcARWtwT9F1D8O9rt08htPvnd49W1IgZtmLkhu9WfMzQj1cFxjHIO6connUNVW5k7AVyQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "forwarded": "0.2.0", + "ipaddr.js": "1.9.1" + }, + "engines": { + "node": ">= 0.10" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/qs": { + "version": "6.16.0", + "resolved": "https://registry.npmjs.org/qs/-/qs-6.16.0.tgz", + "integrity": "sha512-h6fhOIaRrID2CbEY2fqs+7t+UXZo+MLAnU5gRIq85uFtdiUPCdsApMlHhXogKVM4HM2DVbIjGNTTYH2OcmP1vA==", + "dev": true, + "license": "BSD-3-Clause", + "dependencies": { + "es-define-property": "^1.0.1", + "side-channel": "^1.1.1" + }, + "engines": { + "node": ">=0.6" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/range-parser": { + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/range-parser/-/range-parser-1.3.0.tgz", + "integrity": "sha512-hek2mFQpPuI4E1BBKrSto+BU3e3x4xuarsbiwr3+lf7p44juvFMV0XFWQAP3xUyqXA4RrXLIoaSUGbSt056ZMw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.6" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/raw-body": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-3.0.2.tgz", + "integrity": "sha512-K5zQjDllxWkf7Z5xJdV0/B0WTNqx6vxG70zJE4N0kBs4LovmEYWJzQGxC9bS9RAKu3bgM40lrd5zoLJ12MQ5BA==", + "dev": true, + "license": "MIT", + "dependencies": { + "bytes": "~3.1.2", + "http-errors": "~2.0.1", + "iconv-lite": "~0.7.0", + "unpipe": "~1.0.0" + }, + "engines": { + "node": ">= 0.10" + } + }, + "node_modules/require-from-string": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/require-from-string/-/require-from-string-2.0.2.tgz", + "integrity": "sha512-Xf0nWe6RseziFMu+Ap9biiUbmplq6S9/p+7w7YXP/JBHhrUDDUhwa+vANyubuqfZWTveU//DYVGsDG7RKL/vEw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/router": { + "version": "2.2.0", + "resolved": "https://registry.npmjs.org/router/-/router-2.2.0.tgz", + "integrity": "sha512-nLTrUKm2UyiL7rlhapu/Zl45FwNgkZGaCpZbIHajDYgwlJCOzLSk+cIPAnsEqV955GjILJnKbdQC1nVPz+gAYQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "debug": "^4.4.0", + "depd": "^2.0.0", + "is-promise": "^4.0.0", + "parseurl": "^1.3.3", + "path-to-regexp": "^8.0.0" + }, + "engines": { + "node": ">= 18" + } + }, + "node_modules/safer-buffer": { + "version": "2.1.2", + "resolved": "https://registry.npmjs.org/safer-buffer/-/safer-buffer-2.1.2.tgz", + "integrity": "sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==", + "dev": true, + "license": "MIT" + }, + "node_modules/send": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/send/-/send-1.2.1.tgz", + "integrity": "sha512-1gnZf7DFcoIcajTjTwjwuDjzuz4PPcY2StKPlsGAQ1+YH20IRVrBaXSWmdjowTJ6u8Rc01PoYOGHXfP1mYcZNQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "debug": "^4.4.3", + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "etag": "^1.8.1", + "fresh": "^2.0.0", + "http-errors": "^2.0.1", + "mime-types": "^3.0.2", + "ms": "^2.1.3", + "on-finished": "^2.4.1", + "range-parser": "^1.2.1", + "statuses": "^2.0.2" + }, + "engines": { + "node": ">= 18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/serve-static": { + "version": "2.2.1", + "resolved": "https://registry.npmjs.org/serve-static/-/serve-static-2.2.1.tgz", + "integrity": "sha512-xRXBn0pPqQTVQiC8wyQrKs2MOlX24zQ0POGaj0kultvoOCstBQM5yvOhAVSUwOMjQtTvsPWoNCHfPGwaaQJhTw==", + "dev": true, + "license": "MIT", + "dependencies": { + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "parseurl": "^1.3.3", + "send": "^1.2.0" + }, + "engines": { + "node": ">= 18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/setprototypeof": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/setprototypeof/-/setprototypeof-1.2.0.tgz", + "integrity": "sha512-E5LDX7Wrp85Kil5bhZv46j8jOeboKq5JMmYM3gVGdGH8xFpPWXUMsNrlODCrkoxMEeNi/XZIwuRvY4XNwYMJpw==", + "dev": true, + "license": "ISC" + }, + "node_modules/shebang-command": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/shebang-command/-/shebang-command-2.0.0.tgz", + "integrity": "sha512-kHxr2zZpYtdmrN1qDjrrX/Z1rR1kG8Dx+gkpK1G4eXmvXswmcE1hTWBWYUzlraYw1/yZp6YuDY77YtvbN0dmDA==", + "dev": true, + "license": "MIT", + "dependencies": { + "shebang-regex": "^3.0.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/shebang-regex": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/shebang-regex/-/shebang-regex-3.0.0.tgz", + "integrity": "sha512-7++dFhtcx3353uBaq8DDR4NuxBetBzC7ZQOhmTQInHEd6bSrXdiEyzCvG07Z44UYdLShWUyXt5M/yhz8ekcb1A==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=8" + } + }, + "node_modules/side-channel": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/side-channel/-/side-channel-1.1.1.tgz", + "integrity": "sha512-6x6dK6zJdpTzF4sQeNYxwtvBzf6Eg4GtlesS94HOvTudUeyK2WXAaIfmDgsyslYrRBeFIlsi54AYsFGUuhmvrQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "es-errors": "^1.3.0", + "object-inspect": "^1.13.4", + "side-channel-list": "^1.0.1", + "side-channel-map": "^1.0.1", + "side-channel-weakmap": "^1.0.2" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/side-channel-list": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/side-channel-list/-/side-channel-list-1.0.1.tgz", + "integrity": "sha512-mjn/0bi/oUURjc5Xl7IaWi/OJJJumuoJFQJfDDyO46+hBWsfaVM65TBHq2eoZBhzl9EchxOijpkbRC8SVBQU0w==", + "dev": true, + "license": "MIT", + "dependencies": { + "es-errors": "^1.3.0", + "object-inspect": "^1.13.4" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/side-channel-map": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/side-channel-map/-/side-channel-map-1.0.1.tgz", + "integrity": "sha512-VCjCNfgMsby3tTdo02nbjtM/ewra6jPHmpThenkTYh8pG9ucZ/1P8So4u4FGBek/BjpOVsDCMoLA/iuBKIFXRA==", + "dev": true, + "license": "MIT", + "dependencies": { + "call-bound": "^1.0.2", + "es-errors": "^1.3.0", + "get-intrinsic": "^1.2.5", + "object-inspect": "^1.13.3" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/side-channel-weakmap": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/side-channel-weakmap/-/side-channel-weakmap-1.0.2.tgz", + "integrity": "sha512-WPS/HvHQTYnHisLo9McqBHOJk2FkHO/tlpvldyrnem4aeQp4hai3gythswg6p01oSoTl58rcpiFAjF2br2Ak2A==", + "dev": true, + "license": "MIT", + "dependencies": { + "call-bound": "^1.0.2", + "es-errors": "^1.3.0", + "get-intrinsic": "^1.2.5", + "object-inspect": "^1.13.3", + "side-channel-map": "^1.0.1" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/statuses": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", + "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/toidentifier": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/toidentifier/-/toidentifier-1.0.1.tgz", + "integrity": "sha512-o5sSPKEkg/DIQNmH43V0/uerLrpzVedkUh8tGNvaeXpfpuwjKenlSox/2O/BTlZUtEe+JG7s5YhEz608PlAHRA==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.6" + } + }, + "node_modules/type-is": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/type-is/-/type-is-2.1.0.tgz", + "integrity": "sha512-faYHw0anBbc/kWF3zFTEnxSFOAGUX9GFbOBthvDdLsIlEoWOFOtS0zgCiQYwIskL9iGXZL3kAXD8OoZ4GmMATA==", + "dev": true, + "license": "MIT", + "dependencies": { + "content-type": "^2.0.0", + "media-typer": "^1.1.0", + "mime-types": "^3.0.0" + }, + "engines": { + "node": ">= 18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/type-is/node_modules/content-type": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/content-type/-/content-type-2.1.0.tgz", + "integrity": "sha512-mj7UPXE0jaqaOsukNZRUEfEi2AcL7C/vwmwcHV0O97eO1E1pxBZuyjlZrx5seTaNBg1U6+o35wpa35Qfcc+7ag==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, "node_modules/typescript": { "version": "5.9.3", "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", @@ -106,6 +1325,69 @@ "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", "dev": true, "license": "MIT" + }, + "node_modules/unpipe": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/unpipe/-/unpipe-1.0.0.tgz", + "integrity": "sha512-pjy2bYhSsufwWlKwPc+l3cN7+wuJlK6uz0YdJEOlQDbl6jo/YlPi4mb8agUkVC8BF7V8NuzeyPNqRksA3hztKQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/vary": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/vary/-/vary-1.1.2.tgz", + "integrity": "sha512-BNGbWLfd0eUPabhkXUVm0j8uuvREyTh5ovRa/dyow/BqAbZJyC+5fU+IzQOzmAKzYqYRAISoRhdQr3eIZ/PXqg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, + "node_modules/which": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/which/-/which-2.0.2.tgz", + "integrity": "sha512-BLI3Tl1TW3Pvl70l3yq3Y64i+awpwXqsGBYWkkqMtnbXgrMD+yj7rhW0kuEDxzJaYXGjEW5ogapKNMEKNMjibA==", + "dev": true, + "license": "ISC", + "dependencies": { + "isexe": "^2.0.0" + }, + "bin": { + "node-which": "bin/node-which" + }, + "engines": { + "node": ">= 8" + } + }, + "node_modules/wrappy": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/wrappy/-/wrappy-1.0.2.tgz", + "integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==", + "dev": true, + "license": "ISC" + }, + "node_modules/zod": { + "version": "4.6.5", + "resolved": "https://registry.npmjs.org/zod/-/zod-4.6.5.tgz", + "integrity": "sha512-v5l/aFXZQeai4awLbOpSoHecE9UiMrnfx75tEXLjNonXVARxQ5mOeipTjROUchszUNCqnE+hqAMujRsRHsut2Q==", + "dev": true, + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/colinhacks" + } + }, + "node_modules/zod-to-json-schema": { + "version": "3.25.2", + "resolved": "https://registry.npmjs.org/zod-to-json-schema/-/zod-to-json-schema-3.25.2.tgz", + "integrity": "sha512-O/PgfnpT1xKSDeQYSCfRI5Gy3hPf91mKVDuYLUHZJMiDFptvP41MSnWofm8dnCm0256ZNfZIM7DSzuSMAFnjHA==", + "dev": true, + "license": "ISC", + "peerDependencies": { + "zod": "^3.25.28 || ^4" + } } } } diff --git a/tests/e2e/ui/package.json b/tests/e2e/ui/package.json index ede759d97cb..78130412be9 100644 --- a/tests/e2e/ui/package.json +++ b/tests/e2e/ui/package.json @@ -9,6 +9,7 @@ "e2e:migration:root": "playwright test --config migration.serverRootPath.config.ts" }, "devDependencies": { + "@modelcontextprotocol/sdk": "1.30.0", "@playwright/test": "1.58.1", "@types/node": "20.19.37", "typescript": "5.9.3" diff --git a/tests/e2e/ui/tests/mcp/mcpTools.spec.ts b/tests/e2e/ui/tests/mcp/mcpTools.spec.ts index 225ca8b9449..2390a78755e 100644 --- a/tests/e2e/ui/tests/mcp/mcpTools.spec.ts +++ b/tests/e2e/ui/tests/mcp/mcpTools.spec.ts @@ -1,17 +1,20 @@ import { test, expect, Locator } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; -import { createMcpServer, deleteMcpServerByName, openMcpToolsTab } from "../../helpers/mcp"; +import { createMcpServer, deleteMcpServerByName, listUpstreamToolNames, openMcpToolsTab } from "../../helpers/mcp"; // Listing and calling MCP tools, which needs a server that really answers; the create-only spec // points at an unreachable URL on purpose. // -// This spec makes a read-only network call to DeepWiki's public MCP server, from the proxy rather -// than the browser. It needs no credentials, so there is no secret to leak from a public repo. +// This spec makes read-only network calls to DeepWiki's public MCP server: from the proxy, and from +// the test runner to learn which tools the upstream advertises today, so the tool list is never +// pinned here. It needs no credentials, so there is no secret to leak from a public repo. // // A DeepWiki outage turns this red for something that is not a litellm regression. That is left // visible rather than auto-skipped: skipping on connection trouble also skips when the proxy's own // MCP client breaks, which is the regression this exists to catch. E2E_SKIP_EXTERNAL_MCP=1 opts out. const MCP_SERVER_URL = "https://mcp.deepwiki.com/mcp"; +// Read from DeepWiki's tools/list on 2026-09-22. One name has to be pinned so the call-tool test can +// fill a known input (repoName); the listing test checks it is still advertised before the UI checks. const TOOL_NAME = "read_wiki_structure"; const TOOL_ARG_REPO = "BerriAI/litellm"; @@ -36,14 +39,17 @@ test.describe("MCP Tools", () => { }); test("MCP Tools tab lists the tools the upstream server advertises", async ({ page }) => { + const upstreamTools = await listUpstreamToolNames(MCP_SERVER_URL); + expect(upstreamTools).toContain(TOOL_NAME); + // Fetched through the proxy on mount, so allow for a cold upstream connection. const toolList = page.locator(".mcp-tools-scrollable"); await expect(toolList).toBeVisible({ timeout: 30_000 }); - // Non-empty would still pass if the proxy returned some other server's tools. - await expect(toolCard(toolList, TOOL_NAME)).toBeVisible(); - await expect(toolCard(toolList, "ask_question")).toBeVisible(); - await expect(toolCard(toolList, "read_wiki_contents")).toBeVisible(); + for (const name of upstreamTools) { + await expect(toolCard(toolList, name)).toBeVisible(); + } + await expect(toolList.locator("h4.font-mono")).toHaveCount(upstreamTools.length); // No other tool's name or description contains this string, so exactly one card survives. await page.getByPlaceholder("Search tools...").fill(TOOL_NAME); diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 5d9a17acc49..e1b5940935f 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -172,12 +172,39 @@ "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_row": [ "other.provider_wire.fal_ai.gpt_image_generation_quality_size_wire_and_keyed_pricing" ], + "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_generation_prices_non_canonical_size_from_nearest_row": [ + "other.provider_wire.fal_ai.gpt_image_generation_noncanonical_size_uses_nearest_keyed_row" + ], + "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_sdk_response_honors_dump_options": [ + "other.provider_wire.fal_ai.sdk_image_response_dump_options" + ], "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image": [ "other.provider_wire.fal_ai.flux_dev_endpoint_and_per_image_pricing" ], + "tests/integration/providers/test_fal_ai_passthrough_wire.py::test_fal_queue_submit_charges_and_polls_pass_through_free": [ + "other.provider_wire.fal_ai.passthrough_queue_submit_charges_and_polls_do_not" + ], "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_gpt_image_25_edit_inlines_upload_as_data_url_and_charges_keyed_row": [ "other.provider_wire.fal_ai.image_edit_json_data_urls_and_keyed_pricing" ], + "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_flux_lora_depth_edit_sends_single_image_url_and_charges_flat_row": [ + "other.provider_wire.fal_ai.flux_lora_depth_edit_single_image_url_and_flat_pricing" + ], + "tests/integration/providers/test_fal_ai_chat_wire.py::test_fal_moondream3_chat_sends_prompt_image_and_reasoning": [ + "other.provider_wire.fal_ai.moondream3_chat_query_wire_and_token_pricing" + ], + "tests/integration/providers/test_xiaomi_mimo_wire.py::test_xiaomi_mimo_nonstream_surfaces_reasoning_and_charges_registry_price[mimo-v2.6-pro]": [ + "other.provider_wire.xiaomi_mimo.reasoning_content_and_registry_pricing" + ], + "tests/integration/providers/test_xiaomi_mimo_wire.py::test_xiaomi_mimo_nonstream_surfaces_reasoning_and_charges_registry_price[mimo-v2.6-flash]": [ + "other.provider_wire.xiaomi_mimo.reasoning_content_and_registry_pricing" + ], + "tests/integration/providers/test_xiaomi_mimo_wire.py::test_xiaomi_mimo_stream_delivers_reasoning_then_answer_deltas": [ + "other.provider_wire.xiaomi_mimo.reasoning_and_answer_stream_as_deltas" + ], + "tests/integration/providers/test_xiaomi_mimo_wire.py::test_xiaomi_mimo_tool_call_is_forwarded_and_returned": [ + "other.provider_wire.xiaomi_mimo.tool_call_survives_translation" + ], "tests/integration/providers/test_fal_ai_video_wire.py::test_fal_h3_video_create_uses_canonical_body_and_status_path": [ "other.provider_wire.fal_ai.video_queue_create_status_and_content_download" ], diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index ea8bf230d05..ac3fcd33d2e 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json from collections.abc import Mapping from pathlib import Path from types import MappingProxyType @@ -8,6 +9,7 @@ from typing import Annotated, Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator, model_validator CASES_PATH: Final = Path(__file__).resolve().parent / "cost_tracking_cases.json" +PRIOR_RESPONSE_ID_MARKER: Final = "$PRIOR_RESPONSE_ID" class SearchContextCostPerQuery(BaseModel): @@ -312,6 +314,21 @@ class CostTrackingTestCase(BaseModel): usage: Final = self.response.body.get("usage") return isinstance(usage, dict) and isinstance(usage.get("cost"), (int, float)) + @property + def chains_prior_response(self) -> bool: + return self.request.get("previous_response_id") == PRIOR_RESPONSE_ID_MARKER + + @property + def can_chain_prior_response(self) -> bool: + return ( + self.chains_prior_response + and self.endpoint == "/v1/responses" + and isinstance(self.response, JsonResponse) + and isinstance(self.response.body.get("id"), str) + and not isinstance(self.expected, FailureExpected) + and not (isinstance(self.expected, ExactExpected) and self.expected.rollups) + ) + class BatchOutputLine(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") @@ -681,6 +698,11 @@ def data_errors() -> tuple[str, ...]: for marker in ('"id": "call_$REQUEST_ID"', '"id": "toolu_$REQUEST_ID"') ) ) + invalid_prior_response_chains: Final = sorted( + case.name + for case in CASES + if PRIOR_RESPONSE_ID_MARKER in json.dumps(case.request) and not case.can_chain_prior_response + ) return tuple( message for message in ( @@ -700,6 +722,10 @@ def data_errors() -> tuple[str, ...]: f"pinned tool IDs contain $REQUEST_ID: {invalid_pinned_tool_ids}" if invalid_pinned_tool_ids else None, + f"{PRIOR_RESPONSE_ID_MARKER} needs a non-rollup, non-failure /v1/responses JSON response with a string id" + f" as previous_response_id: {invalid_prior_response_chains}" + if invalid_prior_response_chains + else None, ) if message is not None ) diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index 17ebc793fae..09dfa66012f 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -27733,7 +27733,7 @@ "request": { "model": "$MODEL", "input": "continue this text", - "previous_response_id": "resp_scripted_prior" + "previous_response_id": "$PRIOR_RESPONSE_ID" }, "response": { "content_type": "application/json", diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index efab17acba4..82878634677 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -9,13 +9,14 @@ import time import uuid import wave import zlib +from collections.abc import Mapping from hashlib import sha256 from itertools import islice from typing import Final, cast import httpx import pytest -from integration._support.client import JSON_OBJECT, Gateway +from integration._support.client import JSON_OBJECT, Gateway, string_value from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.assertions import assert_exact, assert_recount from integration.cost_calculation.conftest import ( @@ -112,6 +113,19 @@ def _replace_model(value: JsonValue, model_name: str) -> JsonValue: return value +def _prime_prior_response( + gateway: Gateway, request_path: str, request_values: Mapping[str, JsonValue], key: str +) -> str: + primed: Final = gateway.request( + "POST", + request_path, + {field: value for field, value in request_values.items() if field != "previous_response_id"}, + key=key, + ) + assert primed.is_success, f"priming response failed: {primed.status_code}: {primed.text[:400]}" + return string_value(JSON_OBJECT.validate_json(primed.content)["id"]) + + @pytest.mark.parametrize("case", _CASES) def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) -> None: marker: Final = sha256(case.name.encode()).hexdigest()[:12] @@ -175,21 +189,6 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) if isinstance(expected, ExactExpected) and expected.rollups else None ) - request_body: Final = JSON_OBJECT.validate_python( - { - **base_request_values, - **( - {"model": fallback_deployment.model_name, "fallbacks": [model_name]} - if fallback_deployment is not None - else {} - ), - **( - {"user": end_user_id, "cache": {"no-cache": True}} - if end_user_id is not None - else {} - ), - } - ) request_headers: Final = ( { "x-pass-x-scripted-scenario": scenario_id, @@ -207,6 +206,27 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) if passthrough_provider is not None else case.endpoint ) + prior_response_id: Final = ( + _prime_prior_response(gateway, request_path, base_request_values, key) + if case.chains_prior_response + else None + ) + request_body: Final = JSON_OBJECT.validate_python( + { + **base_request_values, + **( + {"model": fallback_deployment.model_name, "fallbacks": [model_name]} + if fallback_deployment is not None + else {} + ), + **( + {"user": end_user_id, "cache": {"no-cache": True}} + if end_user_id is not None + else {} + ), + **({"previous_response_id": prior_response_id} if prior_response_id is not None else {}), + } + ) if case.disconnect_after_frames is not None: with gateway.client.stream( "POST", @@ -250,7 +270,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) - rows: Final = poll_rows(key, len(responses)) + rows: Final = poll_rows(key, len(responses) + (prior_response_id is not None)) if isinstance(expected, RecountExpected): row: Final = rows[0] assert_recount(case.name, expected, row) diff --git a/tests/integration/providers/test_fal_ai_chat_wire.py b/tests/integration/providers/test_fal_ai_chat_wire.py new file mode 100644 index 00000000000..2bb1ac3f168 --- /dev/null +++ b/tests/integration/providers/test_fal_ai_chat_wire.py @@ -0,0 +1,99 @@ +import json +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_MODEL: Final = "fal-ai/moondream3-preview/query" +_PROMPT: Final = "what is in this image?" +_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]]) + + +def _catalog_cost(key: str, field: str) -> float: + cost_map: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes()) + cost_value: Final = cost_map[key][field] + assert isinstance(cost_value, (int, float)) + return float(cost_value) + + +def _approx(value: float) -> object: + return pytest.approx(value, rel=1e-6) # pyright: ignore[reportUnknownMemberType] # pytest lacks typed approx stubs + + +@pytest.mark.covers("other.provider_wire.fal_ai.moondream3_chat_query_wire_and_token_pricing") +def test_fal_moondream3_chat_sends_prompt_image_and_reasoning(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.headers["authorization"] == "Key synthetic-fal-key" + assert request.target == f"/{_MODEL}" + assert request.headers["content-type"] == "application/json" + assert _JSON_OBJECT.validate_json(request.body) == { + "prompt": _PROMPT, + "image_url": "https://example.com/pic.png", + "reasoning": False, + "temperature": 0.2, + } + return Reply( + body=json.dumps( + { + "output": "a red circle on a blue background", + "reasoning": "inspected the shapes", + "finish_reason": "stop", + "usage_info": { + "input_tokens": 11, + "output_tokens": 7, + "prefill_time_ms": 1.0, + "decode_time_ms": 2.0, + "ttft_ms": 1.5, + }, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + wire_url: Final = wire.url + model: Final = scenario.model(model=f"fal_ai/{_MODEL}", api_base=wire.url, api_key="synthetic-fal-key") + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": _PROMPT}, + {"type": "image_url", "image_url": {"url": "https://example.com/pic.png"}}, + ], + } + ], + "reasoning_effort": "none", + "temperature": 0.2, + }, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["choices"] == [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "role": "assistant", + "content": "a red circle on a blue background", + "reasoning_content": "inspected the shapes", + }, + } + ] + assert payload["usage"] == {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18} + cost: Final = float(response.headers["x-litellm-response-cost"]) + assert cost == _approx( + 11 * _catalog_cost(f"fal_ai/{_MODEL}", "input_cost_per_token") + + 7 * _catalog_cost(f"fal_ai/{_MODEL}", "output_cost_per_token") + ) + assert [(request.method, request.target) for request in wire.drain()] == [("POST", f"/{_MODEL}")] diff --git a/tests/integration/providers/test_fal_ai_image_wire.py b/tests/integration/providers/test_fal_ai_image_wire.py index f9ceac0b037..02f24f9e369 100644 --- a/tests/integration/providers/test_fal_ai_image_wire.py +++ b/tests/integration/providers/test_fal_ai_image_wire.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Final import httpx +import litellm import pytest from integration._support.client import Gateway from integration._support.wire import Reply, Request, wire_server @@ -120,6 +121,63 @@ def test_fal_gpt_image_25_generation_sends_quality_and_size_and_charges_keyed_ro ] +@pytest.mark.covers("other.provider_wire.fal_ai.gpt_image_generation_noncanonical_size_uses_nearest_keyed_row") +def test_fal_gpt_image_25_generation_prices_non_canonical_size_from_nearest_row(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.headers["authorization"] == "Key synthetic-fal-key" + assert request.target == "/openai/gpt-image-2.5/flare/text-to-image" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body == {"prompt": _PROMPT, "quality": "low", "image_size": {"width": 1536, "height": 1024}} + return Reply(body=_image_response(((f"{wire_url}/files/noncanonical.png", 1536, 1024),), _PROMPT)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + wire_url: Final = wire.url + model: Final = scenario.model( + model=f"fal_ai/{_GPT_IMAGE_MODEL}", api_base=wire.url, api_key="synthetic-fal-key" + ) + response: Final = gateway.request( + "POST", + "/v1/images/generations", + {"model": model, "prompt": _PROMPT, "quality": "low", "size": "1536x1024"}, + ) + assert response.status_code == 200, response.text + cost: Final = _response_cost(response) + assert cost == _approx(_catalog_cost("fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image")) + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", "/openai/gpt-image-2.5/flare/text-to-image") + ] + + +@pytest.mark.covers("other.provider_wire.fal_ai.sdk_image_response_dump_options") +def test_fal_gpt_image_sdk_response_honors_dump_options() -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.headers["authorization"] == "Key synthetic-fal-key" + assert request.target == "/openai/gpt-image-2.5/flare/text-to-image" + assert _JSON_OBJECT.validate_json(request.body) == {"prompt": _PROMPT, "quality": "low"} + return Reply(body=_image_response((("https://example.com/fal.png", 1024, 1536),), _PROMPT)) + + with wire_server(respond) as wire: + response: Final = litellm.image_generation( + model=_GPT_IMAGE_MODEL, + prompt=_PROMPT, + quality="low", + api_base=wire.url, + api_key="synthetic-fal-key", + custom_llm_provider="fal_ai", + ) + assert response.model_dump(exclude_none=True)["data"] == [ + { + "url": "https://example.com/fal.png", + "provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"}, + } + ] + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", "/openai/gpt-image-2.5/flare/text-to-image") + ] + + @pytest.mark.covers("other.provider_wire.fal_ai.flux_dev_endpoint_and_per_image_pricing") def test_fal_flux_dev_generation_targets_dev_endpoint_and_charges_per_image(gateway: Gateway) -> None: def respond(request: Request) -> Reply: @@ -205,3 +263,44 @@ def test_fal_gpt_image_25_edit_inlines_upload_as_data_url_and_charges_keyed_row( assert [(request.method, request.target) for request in wire.drain()] == [ ("POST", "/openai/gpt-image-2.5/flare/edit") ] + + +@pytest.mark.covers("other.provider_wire.fal_ai.flux_lora_depth_edit_single_image_url_and_flat_pricing") +def test_fal_flux_lora_depth_edit_sends_single_image_url_and_charges_flat_row(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.headers["authorization"] == "Key synthetic-fal-key" + assert request.target == "/fal-ai/flux-lora-depth" + assert request.headers["content-type"] == "application/json" + assert _JSON_OBJECT.validate_json(request.body) == { + "prompt": _PROMPT, + "image_url": "data:image/png;base64," + base64.b64encode(_PNG_BYTES).decode(), + } + return Reply(body=_image_response(((f"{wire_url}/files/depth.png", 1024, 1024),), _PROMPT)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + wire_url: Final = wire.url + model: Final = scenario.model( + model="fal_ai/fal-ai/flux-lora-depth", api_base=wire.url, api_key="synthetic-fal-key" + ) + response: Final = gateway.client.post( + "/v1/images/edits", + data={"model": model, "prompt": _PROMPT}, + files={"image": ("red_circle.png", _PNG_BYTES, "image/png")}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["data"] == [ + { + "url": f"{wire.url}/files/depth.png", + "b64_json": None, + "revised_prompt": None, + "provider_specific_fields": {"width": 1024, "height": 1024, "content_type": "image/png"}, + } + ] + cost: Final = _response_cost(response) + assert cost == _approx(_catalog_cost("fal_ai/fal-ai/flux-lora-depth")) + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", "/fal-ai/flux-lora-depth") + ] diff --git a/tests/integration/providers/test_fal_ai_passthrough_wire.py b/tests/integration/providers/test_fal_ai_passthrough_wire.py new file mode 100644 index 00000000000..f103135124e --- /dev/null +++ b/tests/integration/providers/test_fal_ai_passthrough_wire.py @@ -0,0 +1,86 @@ +import json +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +_MODEL: Final = "fal-ai/trellis-2" +_REQUEST_BODY: Final = {"image_url": "https://example.com/in.png", "resolution": 1536} +_UPSTREAM_BODY: Final = { + "model_glb": { + "url": "https://fal.media/model.glb", + "content_type": "model/gltf-binary", + "file_name": "model.glb", + "file_size": 123, + } +} +_EXPECTED_SPEND: Final = 0.35 + + +@pytest.mark.covers("other.provider_wire.fal_ai.passthrough_queue_submit_charges_and_polls_do_not") +def test_fal_queue_submit_charges_and_polls_pass_through_free(gateway: Gateway, tmp_path) -> None: + def respond(request: Request) -> Reply: + assert request.headers["authorization"] == "Key synthetic-fal-key" + if request.method == "POST": + assert request.target == f"/{_MODEL}" + assert json.loads(request.body) == _REQUEST_BODY + return Reply(body=json.dumps({"request_id": "req-1", "status": "IN_QUEUE"}).encode()) + if request.target == f"/{_MODEL}/requests/req-1/status": + return Reply(body=json.dumps({"status": "COMPLETED"}).encode()) + assert request.target == f"/{_MODEL}/requests/req-1" + return Reply(body=json.dumps(_UPSTREAM_BODY).encode()) + + config: Final = tmp_path / "proxy_config.yaml" + config.write_text( + "model_list: []\n" + "general_settings:\n" + " master_key: os.environ/LITELLM_MASTER_KEY\n" + " database_url: os.environ/DATABASE_URL\n" + " store_model_in_db: true\n" + " disable_spend_logs: false\n" + " proxy_batch_write_at: 1\n" + "router_settings:\n" + " disable_cooldowns: true\n" + ) + with wire_server(respond) as wire: + with owned_proxy( + gateway, + tmp_path, + {"FAL_AI_QUEUE_API_BASE": wire.url, "FAL_AI_API_KEY": "synthetic-fal-key"}, + config=config, + ) as candidate: + submit: Final = candidate.request("POST", f"/fal_ai/{_MODEL}", _REQUEST_BODY) + assert submit.status_code == 200, submit.text + assert json.loads(submit.content) == {"request_id": "req-1", "status": "IN_QUEUE"} + status_response: Final = candidate.request("GET", f"/fal_ai/{_MODEL}/requests/req-1/status") + assert status_response.status_code == 200, status_response.text + assert json.loads(status_response.content) == {"status": "COMPLETED"} + result_response: Final = candidate.request("GET", f"/fal_ai/{_MODEL}/requests/req-1") + assert result_response.status_code == 200, result_response.text + assert json.loads(result_response.content) == _UPSTREAM_BODY + submit_spend: Final = eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (submit.headers["x-litellm-call-id"],), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert float(submit_spend[0]["spend"]) == pytest.approx(_EXPECTED_SPEND) + poll_rows: Final = eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=ANY(%s)', + ([status_response.headers["x-litellm-call-id"], result_response.headers["x-litellm-call-id"]],), + ), + lambda values: len(values) == 2, + seconds=70, + ) + assert sorted(float(row["spend"]) for row in poll_rows) == [0.0, 0.0] + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", f"/{_MODEL}"), + ("GET", f"/{_MODEL}/requests/req-1/status"), + ("GET", f"/{_MODEL}/requests/req-1"), + ] diff --git a/tests/integration/providers/test_fal_ai_video_wire.py b/tests/integration/providers/test_fal_ai_video_wire.py index 827818c6780..276ea2868a7 100644 --- a/tests/integration/providers/test_fal_ai_video_wire.py +++ b/tests/integration/providers/test_fal_ai_video_wire.py @@ -76,6 +76,9 @@ def test_fal_h3_video_create_uses_canonical_body_and_status_path(gateway: Gatewa request_id: Final = "fal-h3-req-" + uuid.uuid4().hex def respond(request: Request) -> Reply: + if request.target == f"/files/{request_id}.mp4": + assert request.method == "GET" + return Reply(body=_MP4, content_type="video/mp4") assert request.headers["authorization"] == "Key synthetic-fal-key" if request.method == "POST": assert request.target == f"/{_H3_MODEL}" diff --git a/tests/integration/providers/test_xiaomi_mimo_wire.py b/tests/integration/providers/test_xiaomi_mimo_wire.py new file mode 100644 index 00000000000..96b9dc17bdf --- /dev/null +++ b/tests/integration/providers/test_xiaomi_mimo_wire.py @@ -0,0 +1,258 @@ +import json +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +_BACKENDS: Final = ("mimo-v2.6-pro", "mimo-v2.6-flash") +_API_KEY: Final = "synthetic-xiaomi-key" +_ARITHMETIC_PROMPT: Final = "What is 17 + 26? Answer with just the number." +_WEATHER_PROMPT: Final = "What is the weather in Paris? Use the tool." +_COUNTING_PROMPT: Final = "Count from 1 to 5, one number per line." +_WEATHER_TOOL: Final[JsonValue] = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} +_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]]) + + +class _Delta(BaseModel): + model_config = ConfigDict(extra="ignore") + content: str | None = None + reasoning_content: str | None = None + + +class _Choice(BaseModel): + model_config = ConfigDict(extra="ignore") + delta: _Delta + finish_reason: str | None = None + + +class _Chunk(BaseModel): + model_config = ConfigDict(extra="ignore") + id: str + choices: tuple[_Choice, ...] + + +def _catalog_cost(backend: str, field: str) -> float: + cost_map: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes()) + cost_value: Final = cost_map[f"xiaomi_mimo/{backend}"][field] + assert isinstance(cost_value, (int, float)) + return float(cost_value) + + +def _approx(value: float) -> object: + return pytest.approx(value, rel=1e-6) # pyright: ignore[reportUnknownMemberType] # pytest lacks typed approx stubs + + +def _completion(identity: str, backend: str, message: Mapping[str, object], finish: str) -> bytes: + return json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": backend, + "choices": [{"index": 0, "message": message, "finish_reason": finish}], + "usage": {"prompt_tokens": 23, "completion_tokens": 41, "total_tokens": 64}, + } + ).encode() + + +def _frame(identity: str, backend: str, delta: Mapping[str, object], finish: str | None = None) -> bytes: + value: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": backend, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + return b"data: " + json.dumps(value).encode() + b"\n\n" + + +def _assert_provider_request(request: Request, backend: str, prompt: str) -> dict[str, JsonValue]: + assert request.method == "POST" + assert request.target == "/chat/completions" + assert request.headers["authorization"] == f"Bearer {_API_KEY}" + assert request.headers["content-type"] == "application/json" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == backend + assert body["messages"] == [{"role": "user", "content": prompt}] + return body + + +@pytest.mark.covers("other.provider_wire.xiaomi_mimo.reasoning_content_and_registry_pricing") +@pytest.mark.parametrize("backend", _BACKENDS) +def test_xiaomi_mimo_nonstream_surfaces_reasoning_and_charges_registry_price(gateway: Gateway, backend: str) -> None: + identity: Final = f"xiaomi-cost-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _assert_provider_request(request, backend, _ARITHMETIC_PROMPT) + assert body["max_tokens"] == 256 + assert "max_completion_tokens" not in body + return Reply( + body=_completion( + identity, + backend, + {"role": "assistant", "content": "43", "reasoning_content": "17 plus 26 is 43."}, + "stop", + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"xiaomi_mimo/{backend}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _ARITHMETIC_PROMPT}], + "max_completion_tokens": 256, + }, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["id"] == identity + assert payload["choices"] == [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "role": "assistant", + "content": "43", + "reasoning_content": "17 plus 26 is 43.", + "provider_specific_fields": {"refusal": None}, + }, + "provider_specific_fields": {}, + } + ] + assert payload["usage"] == {"prompt_tokens": 23, "completion_tokens": 41, "total_tokens": 64} + expected_cost: Final = 23 * _catalog_cost(backend, "input_cost_per_token") + 41 * _catalog_cost( + backend, "output_cost_per_token" + ) + assert float(response.headers["x-litellm-response-cost"]) == _approx(expected_cost) + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (23, 41) + spend: Final = rows[0]["spend"] + assert isinstance(spend, (int, float, str)) + assert float(spend) == _approx(expected_cost) + + +@pytest.mark.covers("other.provider_wire.xiaomi_mimo.reasoning_and_answer_stream_as_deltas") +def test_xiaomi_mimo_stream_delivers_reasoning_then_answer_deltas(gateway: Gateway) -> None: + backend: Final = _BACKENDS[0] + identity: Final = f"xiaomi-stream-{uuid.uuid4().hex}" + frames: Final = ( + _frame(identity, backend, {"role": "assistant", "reasoning_content": "Count "}), + _frame(identity, backend, {"reasoning_content": "up by one."}), + _frame(identity, backend, {"content": "1\n2\n"}), + _frame(identity, backend, {"content": "3\n4\n5"}), + _frame(identity, backend, {}, finish="stop"), + b"data: [DONE]\n\n", + ) + + def respond(request: Request) -> Reply: + body: Final = _assert_provider_request(request, backend, _COUNTING_PROMPT) + assert body["stream"] is True + return Reply(content_type="text/event-stream", chunks=frames) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"xiaomi_mimo/{backend}", api_base=wire.url, api_key=_API_KEY) + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={"model": model, "messages": [{"role": "user", "content": _COUNTING_PROMPT}], "stream": True}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read() + lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: ")) + assert lines[-1] == "data: [DONE]" + chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) + assert {chunk.id for chunk in chunks} == {identity} + choices: Final = tuple(choice for chunk in chunks for choice in chunk.choices) + assert "".join(choice.delta.reasoning_content or "" for choice in choices) == "Count up by one." + assert "".join(choice.delta.content or "" for choice in choices) == "1\n2\n3\n4\n5" + assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ("stop",) + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +@pytest.mark.covers("other.provider_wire.xiaomi_mimo.tool_call_survives_translation") +def test_xiaomi_mimo_tool_call_is_forwarded_and_returned(gateway: Gateway) -> None: + backend: Final = _BACKENDS[1] + identity: Final = f"xiaomi-tool-{uuid.uuid4().hex}" + tool_call: Final = { + "id": "call_paris", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps({"city": "Paris"})}, + } + + def respond(request: Request) -> Reply: + body: Final = _assert_provider_request(request, backend, _WEATHER_PROMPT) + assert body["tools"] == [_WEATHER_TOOL] + assert body["tool_choice"] == "auto" + return Reply( + body=_completion( + identity, + backend, + { + "role": "assistant", + "content": None, + "reasoning_content": "Need the tool.", + "tool_calls": [tool_call], + }, + "tool_calls", + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"xiaomi_mimo/{backend}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _WEATHER_PROMPT}], + "tools": [_WEATHER_TOOL], + "tool_choice": "auto", + }, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["choices"] == [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "content": None, + "reasoning_content": "Need the tool.", + "tool_calls": [tool_call], + "provider_specific_fields": {"refusal": None}, + }, + "provider_specific_fields": {}, + } + ] + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] diff --git a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py index 1d4e13474a7..6c57e59f7e3 100644 --- a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py +++ b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py @@ -413,6 +413,7 @@ PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES = { "/comprehendmedical/{operation}": {"POST"}, "/transcribe": {"POST"}, "/transcribe/{operation}": {"POST"}, + "/tinyfish/{endpoint:path}": {"GET", "POST"}, } diff --git a/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py b/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py index 498f0a734a3..ab5fd7f80ee 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py @@ -53,6 +53,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_enabled(): # Setup logging object with model info litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() @@ -132,6 +133,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_disabled(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() @@ -194,6 +196,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_no_usage_chunk(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() @@ -249,6 +252,7 @@ async def test_vertex_ai_anthropic_streaming_model_extraction(): response.aiter_bytes = mock_aiter_bytes litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + litellm_logging_obj.litellm_params = {} litellm_logging_obj.model_call_details = {} litellm_logging_obj.completion_start_time = None litellm_logging_obj.async_success_handler = AsyncMock() diff --git a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py index c4b1f4f3afd..23b6b302202 100644 --- a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py +++ b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py @@ -22,10 +22,8 @@ from litellm.proxy.proxy_server import token_counter def _fake_hf_tokenizer(num_tokens: int) -> MagicMock: - encoding = MagicMock() - encoding.__len__.return_value = num_tokens tokenizer = MagicMock() - tokenizer.encode_batch_fast.return_value = [encoding] + tokenizer.encode_batch_fast.return_value = [[0] * num_tokens] return tokenizer @@ -58,7 +56,7 @@ async def test_custom_tokenizer_from_model_info_is_used(monkeypatch): ) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - with patch.object(litellm.utils, "Tokenizer") as mock_tokenizer_cls: + with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls: mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(7) response = await token_counter( @@ -92,7 +90,7 @@ async def test_model_without_custom_tokenizer_uses_default(monkeypatch): ) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - with patch.object(litellm.utils, "Tokenizer") as mock_tokenizer_cls: + with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls: response = await token_counter( request=TokenCountRequest( model="gpt-4", diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 3f2c04336a7..e95ed42013b 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -19,6 +19,8 @@ from litellm.proxy._types import ( ) from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( _to_response, + _token_hash_for_create, + _token_hash_for_update, create_jwt_key_mapping, delete_jwt_key_mapping, info_jwt_key_mapping, @@ -1515,6 +1517,132 @@ def test_jwt_client_id_field_does_not_raise_on_duplicate(): assert auth.virtual_key_claim_field == "new_field" +# ────────────────────────────────────────────── +# Tests: identifying the mapped key by hash instead of plaintext +# ────────────────────────────────────────────── + +_TOKEN_HASH = "1923314ae0efc8b2523c7d421bac5a7cf88df291273b139948b526d396974a41" + + +def test_create_stores_a_supplied_token_hash_verbatim(): + """A caller that holds only the hash gets it stored as given, not hashed again.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", token=_TOKEN_HASH + ) + + assert _token_hash_for_create(data) == _TOKEN_HASH + + +def test_create_hashes_a_supplied_plaintext_key(): + """Supplying `key` keeps the original behaviour, so existing configs are unaffected.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest, hash_token + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test-key" + ) + + assert _token_hash_for_create(data) == hash_token("sk-test-key") + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({}, id="neither"), + pytest.param({"key": "sk-test-key", "token": _TOKEN_HASH}, id="both"), + ], +) +def test_create_requires_exactly_one_identifier(kwargs): + """Neither or both is a 400, so a mapping can never be created ambiguously.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", **kwargs + ) + + with pytest.raises(HTTPException) as exc_info: + _token_hash_for_create(data) + + assert exc_info.value.status_code == 400 + assert "exactly one" in exc_info.value.detail.lower() + + +@pytest.mark.parametrize( + "token", + [ + pytest.param("sk-not-a-hash", id="plaintext-key"), + pytest.param("abc123", id="too-short"), + pytest.param(_TOKEN_HASH.upper(), id="uppercase"), + pytest.param(_TOKEN_HASH + "0", id="too-long"), + pytest.param(_TOKEN_HASH[:-1] + "g", id="non-hex-character"), + ], +) +def test_create_rejects_a_token_that_is_not_a_sha256_hash(token): + """hash_token hashes unconditionally, so a bad `token` would be stored as a hash + of a hash and then silently match nothing at auth time.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", token=token + ) + + with pytest.raises(HTTPException) as exc_info: + _token_hash_for_create(data) + + assert exc_info.value.status_code == 400 + assert "SHA-256" in exc_info.value.detail + + +def test_update_leaves_the_mapped_key_alone_when_neither_is_given(): + """Updating only the description must not blank out the mapped key.""" + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + data = UpdateJWTKeyMappingRequest(id="mapping-1", description="new text") + + assert _token_hash_for_update(data) is None + + +def test_update_stores_a_supplied_token_hash_verbatim(): + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + data = UpdateJWTKeyMappingRequest(id="mapping-1", token=_TOKEN_HASH) + + assert _token_hash_for_update(data) == _TOKEN_HASH + + +def test_update_hashes_a_supplied_plaintext_key(): + from litellm.proxy._types import UpdateJWTKeyMappingRequest, hash_token + + data = UpdateJWTKeyMappingRequest(id="mapping-1", key="sk-rotated") + + assert _token_hash_for_update(data) == hash_token("sk-rotated") + + +def test_update_rejects_both_identifiers(): + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + data = UpdateJWTKeyMappingRequest(id="mapping-1", key="sk-abc", token=_TOKEN_HASH) + + with pytest.raises(HTTPException) as exc_info: + _token_hash_for_update(data) + + assert exc_info.value.status_code == 400 + assert "at most one" in exc_info.value.detail.lower() + + +def test_update_rejects_a_token_that_is_not_a_sha256_hash(): + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + data = UpdateJWTKeyMappingRequest(id="mapping-1", token="sk-not-a-hash") + + with pytest.raises(HTTPException) as exc_info: + _token_hash_for_update(data) + + assert exc_info.value.status_code == 400 + assert "SHA-256" in exc_info.value.detail + + # ────────────────────────────────────────────── # Tests: cache eviction must happen AFTER the DB write commits # ────────────────────────────────────────────── diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index f5e8d861d79..9cdac341b1f 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -1540,18 +1540,17 @@ async def test_user_budget_lookup_is_also_unenforced_when_the_database_is_down() """ KNOWN LIMITATION, pinned deliberately rather than discovered later. - `get_user_object` cannot tell "row absent" from "database unreachable": the - absent case raises inside its own try (auth_checks.py:2177) and the handler - at :2213 rewrites every exception into the same - `ValueError("User doesn't exist in db...")`. A connection error, a query - timeout and a malformed row all reach us as that one type and message. + `get_user_object` lets a connection-level outage propagate as-is and rewrites + every other read failure (a query-level Prisma error, a malformed row) into + the same `ValueError("User doesn't exist in db...")` as an absent row, and + `_read_user_model_max_budget` swallows every exception either way. - So tolerating the absent case, which the test above requires, unavoidably - tolerates an outage too, and a user who DOES have a per-model budget goes + So tolerating the absent case, which the test above requires, also + tolerates an outage, and a user who DOES have a per-model budget goes unenforced while the DB is unreachable. This is pre-existing behaviour of - `get_user_object` that the virtual-key path inherits identically; it is not - introduced here. Distinguishing them needs a dedicated exception type for - the absent case and a change to both auth paths. + the virtual-key path; it is not introduced here. Distinguishing them needs + `_read_user_model_max_budget` to let an outage through the way the JWT + path does. """ from litellm.caching.dual_cache import DualCache from litellm.proxy.auth.user_api_key_auth import _read_user_model_max_budget diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index a0a9b71787c..ca7303e4c6d 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -578,6 +578,7 @@ def test_qdrant_semantic_cache_set_cache(): assert ( upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key" ) + assert qdrant_cache.sync_client.put.call_args.kwargs["params"] == {"wait": "true"} @pytest.mark.asyncio @@ -650,6 +651,7 @@ async def test_qdrant_semantic_cache_async_set_cache(): assert ( upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key" ) + assert qdrant_cache.async_client.put.call_args.kwargs["params"] == {"wait": "true"} def test_qdrant_semantic_cache_custom_vector_size(): diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 01f2a13d252..70ab4fe07de 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -4,6 +4,7 @@ and the typed StandardLoggingPayload adapter. These need no OTel SDK.""" import json import logging import re +from collections.abc import Mapping from pathlib import Path from typing import Final @@ -151,9 +152,7 @@ def test_llm_call_span_name(): def _all_constants(cls): return { - getattr(cls, name) - for name in vars(cls) - if not name.startswith("__") and isinstance(getattr(cls, name), str) + getattr(cls, name) for name in vars(cls) if not name.startswith("__") and isinstance(getattr(cls, name), str) } @@ -465,9 +464,7 @@ def test_mcp_tool_call_content_gated_off_by_default(): off = MCPToolCallSpanData.from_standard_logging_payload(_mcp_payload()) assert off.arguments_json is None and off.result_json is None - on = MCPToolCallSpanData.from_standard_logging_payload( - _mcp_payload(), capture_content=True - ) + on = MCPToolCallSpanData.from_standard_logging_payload(_mcp_payload(), capture_content=True) assert on.arguments_json is not None and '"Paris"' in on.arguments_json assert on.result_json is not None and "21" in on.result_json @@ -689,9 +686,7 @@ def test_content_capture_gated_off_by_default(): payload = _sample_payload( messages=[{"role": "user", "content": "secret prompt"}], ) - payload["response"]["choices"] = [ - {"finish_reason": "stop", "message": {"role": "assistant", "content": "secret"}} - ] + payload["response"]["choices"] = [{"finish_reason": "stop", "message": {"role": "assistant", "content": "secret"}}] data = LLMCallSpanData.from_standard_logging_payload(payload) assert data.messages_in == () assert data.choices_out == () @@ -911,6 +906,221 @@ def test_ocr_pages_without_markdown_stay_empty(): assert data.choices_out == () +def _assistant_choice(content: str, finish_reason: str | None = None) -> dict[str, object]: + return { + "message": {"role": "assistant", "content": content, "refusal": None, "tool_calls": None}, + "finish_reason": finish_reason, + } + + +def _route_payload(call_type: str, model: str, response: Mapping[str, object]) -> dict[str, object]: + return _sample_payload(call_type=call_type, model=model, messages=None, response=response) + + +def test_text_completion_choices_become_assistant_messages_in_choice_order() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload( + "atext_completion", + "gpt-3.5-turbo-instruct", + { + "id": "cmpl-1", + "object": "text_completion", + "choices": [ + {"index": 0, "text": " first", "finish_reason": "length", "logprobs": None}, + {"index": 1, "text": " second", "finish_reason": "stop", "logprobs": None}, + ], + }, + ), + capture_content=True, + ) + + assert data.choices_out == (_assistant_choice(" first", "length"), _assistant_choice(" second", "stop")) + assert data.finish_reasons == ("length", "stop") + assert data.response_id == "cmpl-1" + + +def test_text_completion_choices_follow_the_content_capture_gate_but_finish_reasons_do_not() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload( + "atext_completion", "gpt-3.5-turbo-instruct", {"choices": [{"text": "x", "finish_reason": "stop"}]} + ) + ) + + assert data.choices_out == () + assert data.finish_reasons == ("stop",) + + +def test_chat_choices_with_a_message_are_passed_through_untouched_even_beside_a_stray_text_key() -> None: + choice: Final = { + "index": 0, + "finish_reason": "stop", + "text": "no", + "message": {"role": "assistant", "content": "chat"}, + } + data: Final = LLMCallSpanData.from_standard_logging_payload( + _sample_payload(response={"choices": [choice]}), capture_content=True + ) + + assert data.choices_out == (choice,) + + +def test_transcription_text_becomes_one_assistant_choice() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload("atranscription", "gpt-4o-mini-transcribe", {"text": "What is the weather like?", "task": "x"}), + capture_content=True, + ) + + assert data.choices_out == (_assistant_choice("What is the weather like?"),) + assert data.finish_reasons == () + + +def test_empty_transcription_text_stays_empty() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload("atranscription", "gpt-4o-mini-transcribe", {"text": ""}), capture_content=True + ) + + assert data.choices_out == () + + +def test_moderation_results_become_one_verdict_per_input_naming_the_hit_categories() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload( + "amoderation", + "omni-moderation-latest", + { + "id": "modr-1", + "results": [ + { + "flagged": True, + "categories": {"harassment": False, "violence": True, "self-harm": True}, + "category_scores": {"harassment": 0.01, "violence": 0.98, "self-harm": 0.7}, + }, + {"flagged": False, "categories": {"violence": False}}, + {"flagged": True}, + ], + }, + ), + capture_content=True, + ) + + assert data.choices_out == (_assistant_choice("flagged: violence, self-harm\n\nnot flagged\n\nflagged"),) + assert data.response_id == "modr-1" + + +def test_moderation_output_follows_the_content_capture_gate() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload("amoderation", "omni-moderation-latest", {"results": [{"flagged": True}]}) + ) + + assert data.choices_out == () + + +def test_moderation_results_without_a_verdict_produce_no_output() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload("amoderation", "omni-moderation-latest", {"results": [{"categories": {"violence": True}}]}), + capture_content=True, + ) + + assert data.choices_out == () + + +def test_image_data_becomes_a_size_summary_and_never_carries_the_base64_payload() -> None: + encoded: Final = "QUJDRA==" + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload( + "aimage_generation", + "gpt-image-1-mini", + { + "created": 1, + "data": [ + {"b64_json": encoded, "revised_prompt": "a red bicycle"}, + {"url": "https://images.example/cat.png"}, + {"b64_json": "QUJDREVGR0g="}, + ], + }, + ), + capture_content=True, + ) + + assert data.choices_out == ( + _assistant_choice( + "a red bicycle\nb64_json image (4 bytes)\n\nhttps://images.example/cat.png\n\nb64_json image (8 bytes)" + ), + ) + assert encoded not in json.dumps(data.choices_out) + + +def test_image_data_without_a_url_or_payload_stays_empty_and_embeddings_are_not_images() -> None: + images: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload("aimage_generation", "gpt-image-1-mini", {"data": [{"revised_prompt": "x"}]}), + capture_content=True, + ) + embeddings: Final = LLMCallSpanData.from_standard_logging_payload( + _embedding_payload([[0.1, 0.2]]), capture_content=True + ) + + assert images.choices_out == () + assert embeddings.choices_out == () + + +def test_speech_summary_becomes_a_media_type_and_byte_count_choice() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload( + "aspeech", "gpt-4o-mini-tts", {"object": "binary", "content_type": "audio/mpeg", "num_bytes": 48210} + ), + capture_content=True, + ) + + assert data.choices_out == (_assistant_choice("audio/mpeg (48210 bytes)"),) + + +def test_speech_summary_without_a_media_type_is_the_byte_count_and_follows_the_capture_gate() -> None: + response: Final = {"object": "binary", "content_type": None, "num_bytes": 7} + shown: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload("aspeech", "gpt-4o-mini-tts", response), capture_content=True + ) + gated: Final = LLMCallSpanData.from_standard_logging_payload(_route_payload("aspeech", "gpt-4o-mini-tts", response)) + + assert shown.choices_out == (_assistant_choice("7 bytes"),) + assert gated.choices_out == () + + +def test_speech_response_without_a_byte_count_produces_no_output() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload("aspeech", "gpt-4o-mini-tts", {"object": "binary", "content_type": "audio/mpeg"}), + capture_content=True, + ) + + assert data.choices_out == () + + +def test_speech_binary_response_is_logged_as_its_summary_not_dropped() -> None: + import httpx + + from litellm.litellm_core_utils.litellm_logging import _extract_response_obj_and_hidden_params + from litellm.types.llms.openai import HttpxBinaryResponseContent + + raw: Final = httpx.Response(200, headers={"content-type": "audio/mpeg"}, content=b"\x00" * 1234) + response_obj, hidden_params = _extract_response_obj_and_hidden_params(HttpxBinaryResponseContent(raw), None) + + assert response_obj == {"object": "binary", "content_type": "audio/mpeg", "num_bytes": 1234} + assert hidden_params is None + + +def test_speech_binary_response_still_streaming_reports_the_bytes_downloaded_so_far() -> None: + import httpx + + from litellm.types.llms.openai import HttpxBinaryResponseContent + + unread: Final = httpx.Response(200, stream=httpx.ByteStream(b"\x00" * 10)) + + assert HttpxBinaryResponseContent(unread).logging_summary() == { + "object": "binary", + "content_type": None, + "num_bytes": 0, + } + + def test_request_identity_prefers_canonical_team_keys(): from litellm.integrations.otel.model.payloads import RequestIdentity @@ -931,9 +1141,7 @@ def test_request_identity_prefers_canonical_team_keys(): def test_request_identity_falls_back_to_legacy_team_keys(): from litellm.integrations.otel.model.payloads import RequestIdentity - payload = _sample_payload( - metadata={"team_id": "legacy-team", "team_alias": "legacy"} - ) + payload = _sample_payload(metadata={"team_id": "legacy-team", "team_alias": "legacy"}) ident = RequestIdentity.from_payload(payload) assert ident.team_id == "legacy-team" assert ident.team_alias == "legacy" @@ -952,7 +1160,10 @@ def test_request_identity_falls_back_to_legacy_team_keys(): }, "from-header", ), - ({"proxy_server_request": {"headers": {"langfuse_trace_name": ""}}, "metadata": {"trace_name": "body"}}, "body"), + ( + {"proxy_server_request": {"headers": {"langfuse_trace_name": ""}}, "metadata": {"trace_name": "body"}}, + "body", + ), ({"proxy_server_request": {"headers": {}}, "metadata": {"user_api_key_team_id": "t1"}}, None), ({}, None), ], @@ -1002,7 +1213,15 @@ def test_caller_trace_name_prefers_the_langfuse_header_over_body_metadata(reques ), ({}, TraceControls()), ], - ids=["body", "headers-beat-body", "anthropic-body", "non-string-tags-dropped", "scalar-coercion", "mutation-controls-ignored", "empty"], + ids=[ + "body", + "headers-beat-body", + "anthropic-body", + "non-string-tags-dropped", + "scalar-coercion", + "mutation-controls-ignored", + "empty", + ], ) def test_caller_trace_controls_carry_user_session_and_tags(request_data, expected): assert caller_trace_controls({"litellm_params": request_data}) == expected @@ -1168,9 +1387,7 @@ def test_content_capture_opt_in_retains_bodies(): payload = _sample_payload( messages=[{"role": "user", "content": "secret prompt"}], ) - payload["response"]["choices"] = [ - {"finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}} - ] + payload["response"]["choices"] = [{"finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}] data = LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True) assert data.messages_in and data.messages_in[0]["content"] == "secret prompt" assert data.choices_out and data.choices_out[0]["message"]["content"] == "hi" @@ -1187,41 +1404,17 @@ def test_capture_span_content_resolves_modes(): # default (no_content) → off assert OpenTelemetryV2Config().capture_span_content is False + assert OpenTelemetryV2Config(capture_message_content=CaptureMessageContent.SPAN_ONLY).capture_span_content is True assert ( - OpenTelemetryV2Config( - capture_message_content=CaptureMessageContent.SPAN_ONLY - ).capture_span_content - is True - ) - assert ( - OpenTelemetryV2Config( - capture_message_content=CaptureMessageContent.SPAN_AND_EVENT - ).capture_span_content - is True + OpenTelemetryV2Config(capture_message_content=CaptureMessageContent.SPAN_AND_EVENT).capture_span_content is True ) # event-only does not authorize span-attribute content - assert ( - OpenTelemetryV2Config( - capture_message_content=CaptureMessageContent.EVENT_ONLY - ).capture_span_content - is False - ) + assert OpenTelemetryV2Config(capture_message_content=CaptureMessageContent.EVENT_ONLY).capture_span_content is False # V1 accepted UPPER_SNAKE_CASE; the env value is case-insensitive so an # operator carrying ``SPAN_AND_EVENT`` forward still enables capture. - assert ( - OpenTelemetryV2Config( - capture_message_content="SPAN_AND_EVENT" - ).capture_span_content - is True - ) - assert ( - OpenTelemetryV2Config(capture_message_content="SPAN_ONLY").capture_span_content - is True - ) - assert ( - OpenTelemetryV2Config(capture_message_content="NO_CONTENT").capture_span_content - is False - ) + assert OpenTelemetryV2Config(capture_message_content="SPAN_AND_EVENT").capture_span_content is True + assert OpenTelemetryV2Config(capture_message_content="SPAN_ONLY").capture_span_content is True + assert OpenTelemetryV2Config(capture_message_content="NO_CONTENT").capture_span_content is False def test_capture_message_content_normalizer_only_touches_strings(): diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py index 9fa198c4ec5..1e2ae24a329 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_vendor_mappers.py @@ -6,6 +6,8 @@ backends, so one trace lights up every configured destination. """ import json +from collections.abc import Mapping +from typing import Final import pytest @@ -249,6 +251,34 @@ def test_langfuse_mapper_renders_an_ocr_call_with_the_page_markdown_as_output(): assert attrs["langfuse.observation.type"] == "generation" +@pytest.mark.parametrize( + ("call_type", "response", "expected_content"), + [ + ("atext_completion", {"choices": [{"text": "Paris.", "finish_reason": "stop"}]}, "Paris."), + ("atranscription", {"text": "What is the weather like?"}, "What is the weather like?"), + ("amoderation", {"results": [{"flagged": True, "categories": {"violence": True}}]}, "flagged: violence"), + ("aimage_generation", {"data": [{"b64_json": "QUJDRA=="}]}, "b64_json image (4 bytes)"), + ("aspeech", {"object": "binary", "content_type": "audio/mpeg", "num_bytes": 9}, "audio/mpeg (9 bytes)"), + ], +) +def test_langfuse_mapper_renders_every_non_chat_route_output_as_an_assistant_message( + call_type: str, response: Mapping[str, object], expected_content: str +) -> None: + payload: Final[dict[str, object]] = { + "call_type": call_type, + "custom_llm_provider": "openai", + "model": "m", + "messages": None, + "response": response, + } + attrs: Final = LangfuseMapper().map(LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True)) + + assert json.loads(attrs["langfuse.observation.output"]) == [ + {"role": "assistant", "content": expected_content, "refusal": None, "tool_calls": None} + ] + assert attrs["langfuse.observation.type"] == "generation" + + # --------------------------------------------------------------------------- # # Weave # --------------------------------------------------------------------------- # diff --git a/tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py b/tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py new file mode 100644 index 00000000000..989009b309a --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_zero_cost_metric.py @@ -0,0 +1,179 @@ +import datetime +from typing import Final + +import pytest +from prometheus_client import REGISTRY +from prometheus_client.samples import Sample + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.utils import StandardLoggingZeroCostDiagnostic + +METRIC: Final = "litellm_zero_cost_requests_total" +MISSING_KEY_DIAGNOSTIC: Final[StandardLoggingZeroCostDiagnostic] = { + "reason": "missing_pricing_key", + "pricing_model": "dep-1", + "missing_pricing_keys": ("input_cost_per_token", "output_cost_per_token"), +} + + +def _clear_prometheus_registry() -> None: + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +def _samples(metric_name: str) -> list[Sample]: + return [sample for metric in REGISTRY.collect() for sample in metric.samples if sample.name == metric_name] + + +def _payload(zero_cost_diagnostic: StandardLoggingZeroCostDiagnostic | None) -> dict[str, object]: + return { + "id": "t", + "call_type": "completion", + "response_cost": 0.0, + "status": "success", + "total_tokens": 30, + "prompt_tokens": 20, + "completion_tokens": 10, + "startTime": 1.0, + "endTime": 2.0, + "completionStartTime": 1.5, + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "model_group": "per-second-priced-chat", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "request_tags": [], + "end_user": None, + "cache_hit": False, + "stream": False, + "response": {"id": "chatcmpl-1"}, + "model_parameters": {}, + "zero_cost_diagnostic": zero_cost_diagnostic, + "metadata": { + "user_api_key_hash": "h", + "user_api_key_alias": "a", + "user_api_key_team_id": "t", + "user_api_key_team_alias": "ta", + "user_api_key_user_id": "u", + "user_api_key_user_email": "e@x.com", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, + "requester_metadata": None, + "user_api_key_end_user_id": None, + "usage_object": None, + }, + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + } + + +async def _log_success( + logger: PrometheusLogger, zero_cost_diagnostic: StandardLoggingZeroCostDiagnostic | None +) -> None: + now: Final = datetime.datetime.now() + kwargs: Final = { + "model": "openai/gpt-5.4-nano", + "litellm_params": {"metadata": {}}, + "standard_logging_object": _payload(zero_cost_diagnostic), + "stream": False, + "start_time": now - datetime.timedelta(seconds=3), + "api_call_start_time": now - datetime.timedelta(seconds=2), + "completion_start_time": now - datetime.timedelta(seconds=1), + "end_time": now, + } + await logger.async_log_success_event(kwargs, None, now, now) + + +async def _log_failure( + logger: PrometheusLogger, zero_cost_diagnostic: StandardLoggingZeroCostDiagnostic | None +) -> None: + now: Final = datetime.datetime.now() + kwargs: Final = { + "model": "openai/gpt-5.4-nano", + "litellm_params": {"metadata": {}}, + "standard_logging_object": {**_payload(zero_cost_diagnostic), "status": "failure"}, + "exception": Exception("stream cut off after the usage chunk"), + "stream": True, + "start_time": now - datetime.timedelta(seconds=3), + "end_time": now, + } + await logger.async_log_failure_event(kwargs, None, now, now) + + +@pytest.mark.asyncio +async def test_failure_event_counts_a_zero_cost_request_by_model_and_reason() -> None: + _clear_prometheus_registry() + try: + logger: Final = PrometheusLogger() + await _log_failure(logger, None) + assert _samples(METRIC) == [] + + await _log_failure(logger, MISSING_KEY_DIAGNOSTIC) + + samples: Final = _samples(METRIC) + assert len(samples) == 1 + assert samples[0].labels == { + "requested_model": "per-second-priced-chat", + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "api_provider": "openai", + "reason": "missing_pricing_key", + } + assert samples[0].value == 1.0 + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_success_event_counts_a_zero_cost_request_by_model_and_reason() -> None: + _clear_prometheus_registry() + try: + logger: Final = PrometheusLogger() + await _log_success(logger, MISSING_KEY_DIAGNOSTIC) + await _log_success(logger, MISSING_KEY_DIAGNOSTIC) + + samples: Final = _samples(METRIC) + assert len(samples) == 1 + assert samples[0].labels == { + "requested_model": "per-second-priced-chat", + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "api_provider": "openai", + "reason": "missing_pricing_key", + } + assert samples[0].value == 2.0 + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_request_without_a_diagnostic_leaves_the_counter_untouched() -> None: + _clear_prometheus_registry() + try: + await _log_success(PrometheusLogger(), None) + + assert _samples(METRIC) == [] + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_label_filter_that_drops_reason_still_counts_the_request() -> None: + _clear_prometheus_registry() + previous_config: Final = litellm.prometheus_metrics_config + litellm.prometheus_metrics_config = [ + {"group": "zero_cost", "metrics": [METRIC], "include_labels": ["requested_model"]} + ] + try: + await _log_success(PrometheusLogger(), MISSING_KEY_DIAGNOSTIC) + + samples: Final = _samples(METRIC) + assert len(samples) == 1 + assert samples[0].labels == {"requested_model": "per-second-priced-chat"} + assert samples[0].value == 1.0 + finally: + litellm.prometheus_metrics_config = previous_config + _clear_prometheus_registry() diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py new file mode 100644 index 00000000000..0e453e3f5eb --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py @@ -0,0 +1,157 @@ +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import ( + ZERO_COST_COUNTER_NAME, + diagnose_zero_cost, + used_pricing_keys, + zero_cost_warning, +) +from litellm.types.utils import CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, Usage + +PER_SECOND_ENTRY: Final = {"input_cost_per_second": 0.00042, "output_cost_per_second": 0.00042} +FREE_ENTRY: Final = {"input_cost_per_token": 0, "output_cost_per_token": 0, "cache_read_input_token_cost": 2e-08} +PRICED_ENTRY: Final = {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06} +TEXT_USAGE: Final = Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + + +def test_missing_pricing_key_names_every_rate_the_usage_needs() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=False + ) + + assert diagnostic == { + "reason": "missing_pricing_key", + "pricing_model": "dep-1", + "missing_pricing_keys": ("input_cost_per_token", "output_cost_per_token"), + } + + +def test_only_the_absent_rate_is_reported() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry={"input_cost_per_token": 1e-06}, calculation_failed=False + ) + + assert diagnostic is not None + assert diagnostic["missing_pricing_keys"] == ("output_cost_per_token",) + + +def test_free_model_stays_silent() -> None: + assert ( + diagnose_zero_cost(usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=FREE_ENTRY, calculation_failed=False) + is None + ) + + +@pytest.mark.parametrize("calculation_failed", [False, True]) +def test_request_without_usage_stays_silent(calculation_failed: bool) -> None: + usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + + assert ( + diagnose_zero_cost( + usage=usage, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=calculation_failed + ) + is None + ) + + +@pytest.mark.parametrize( + "entry", + [ + {"litellm_provider": "openai", "mode": "chat", "supports_prompt_caching": True}, + {"tiered_pricing": [{"range": [0, 128000], "input_cost_per_token": 0, "output_cost_per_token": 0}]}, + {"tiered_pricing": "not a tier table", "litellm_provider": "openai"}, + ], +) +def test_entry_that_declares_no_rate_stays_silent(entry: Mapping[str, object]) -> None: + assert ( + diagnose_zero_cost(usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=entry, calculation_failed=False) + is None + ) + + +def test_tiered_rate_counts_as_a_declared_rate() -> None: + entry = {"tiered_pricing": [{"range": [0, 128000], "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}]} + + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=entry, calculation_failed=False + ) + + assert diagnostic is not None + assert diagnostic["reason"] == "missing_pricing_key" + + +def test_priced_entry_that_still_prices_to_zero_is_pricing_not_applied() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PRICED_ENTRY, calculation_failed=False + ) + + assert diagnostic == {"reason": "pricing_not_applied", "pricing_model": "dep-1", "missing_pricing_keys": ()} + + +def test_calculator_failure_on_a_priced_entry_is_cost_calculation_error() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PRICED_ENTRY, calculation_failed=True + ) + + assert diagnostic == {"reason": "cost_calculation_error", "pricing_model": "dep-1", "missing_pricing_keys": ()} + + +def test_calculator_failure_on_a_free_entry_stays_silent() -> None: + assert ( + diagnose_zero_cost(usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=FREE_ENTRY, calculation_failed=True) + is None + ) + + +def test_calculator_failure_on_an_entry_that_declares_no_rate_stays_silent() -> None: + entry: Final = {"litellm_provider": "openai", "mode": "chat", "supports_prompt_caching": True} + assert ( + diagnose_zero_cost(usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=entry, calculation_failed=True) + is None + ) + + +def test_audio_tokens_need_the_audio_rates() -> None: + usage = Usage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=10, text_tokens=0), + completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=5, text_tokens=15), + ) + + assert used_pricing_keys(usage) == ( + "input_cost_per_audio_token", + "output_cost_per_token", + "output_cost_per_audio_token", + ) + diagnostic = diagnose_zero_cost( + usage=usage, pricing_model="gemini-audio", pricing_entry=PRICED_ENTRY, calculation_failed=False + ) + assert diagnostic is not None + assert diagnostic["missing_pricing_keys"] == ("input_cost_per_audio_token", "output_cost_per_audio_token") + + +def test_warning_names_the_request_the_entry_the_missing_keys_and_the_counter() -> None: + diagnostic = diagnose_zero_cost( + usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=False + ) + assert diagnostic is not None + + message = zero_cost_warning( + diagnostic, + model_group="per-second-priced-chat", + model="openai/gpt-5.4-nano", + custom_llm_provider="openai", + usage=TEXT_USAGE, + ) + + assert "model_group=per-second-priced-chat" in message + assert "model=openai/gpt-5.4-nano" in message + assert "provider=openai" in message + assert "prompt_tokens=10 completion_tokens=20" in message + assert "pricing entry 'dep-1' has no input_cost_per_token, output_cost_per_token" in message + assert f'{ZERO_COST_COUNTER_NAME}{{reason="missing_pricing_key"}}' in message diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index c67f72680a8..62cf5680266 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_any_messages_to_chat_completion_str_messages_conversion, hoist_images_from_tool_messages, is_encrypted_reasoning_block, + merge_consecutive_system_messages, responses_reasoning_items_from_thinking_blocks, split_concatenated_json_objects, strip_encrypted_reasoning_from_messages, @@ -1846,3 +1847,95 @@ class TestEncryptedReasoningReplay: strip_encrypted_reasoning_from_messages(messages) assert messages == before + + +class TestMergeConsecutiveSystemMessages: + def test_merges_each_run_of_string_system_messages_with_a_blank_line(self): + messages = [ + {"role": "system", "content": "You are terse.", "cache_control": {"type": "ephemeral"}}, + {"role": "system", "content": "Skills: none."}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"}, + {"role": "system", "content": "Reminder A"}, + {"role": "system", "content": "Reminder B"}, + {"role": "user", "content": "Bye"}, + ] + + merged = merge_consecutive_system_messages(messages) + + assert merged == [ + {"role": "system", "content": "You are terse.\n\nSkills: none.", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"}, + {"role": "system", "content": "Reminder A\n\nReminder B"}, + {"role": "user", "content": "Bye"}, + ] + + def test_merges_into_text_parts_when_any_system_content_is_a_list(self): + cached_part = {"type": "text", "text": "Skills: none.", "cache_control": {"type": "ephemeral"}} + messages = [ + {"role": "system", "content": "You are terse."}, + {"role": "system", "content": [cached_part, {"type": "text", "text": "Be brief."}]}, + {"role": "system", "content": "Answer in English."}, + {"role": "user", "content": "Hello"}, + ] + + merged = merge_consecutive_system_messages(messages) + + assert merged == [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are terse."}, + cached_part, + {"type": "text", "text": "Be brief."}, + {"type": "text", "text": "Answer in English."}, + ], + }, + {"role": "user", "content": "Hello"}, + ] + assert merged[0]["content"][1] is cached_part + + @pytest.mark.parametrize( + "messages", + [ + [{"role": "system", "content": "You are terse."}, {"role": "user", "content": "Hello"}], + [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi"}], + [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hello"}, + {"role": "system", "content": "Reminder"}, + ], + [], + ], + ids=["single-system", "no-system", "separated-systems", "empty"], + ) + def test_leaves_messages_without_consecutive_system_messages_untouched(self, messages): + before = copy.deepcopy(messages) + + merged = merge_consecutive_system_messages(messages) + + assert merged == before + assert [message is original for message, original in zip(merged, messages)] == [True] * len(messages) + + @pytest.mark.parametrize( + ("messages", "expected_content"), + [ + ([{"role": "system"}, {"role": "system", "content": "Skills: none."}], "Skills: none."), + ([{"role": "system", "content": "You are terse."}, {"role": "system"}], "You are terse."), + ( + [{"role": "system"}, {"role": "system", "content": [{"type": "text", "text": "Be brief."}]}], + [{"type": "text", "text": "Be brief."}], + ), + ], + ids=["missing-then-str", "str-then-missing", "missing-then-list"], + ) + def test_skips_system_messages_without_content_when_merging(self, messages, expected_content): + merged = merge_consecutive_system_messages([*messages, {"role": "user", "content": "Hello"}]) + + assert merged == [{"role": "system", "content": expected_content}, {"role": "user", "content": "Hello"}] + + def test_keeps_the_first_message_when_no_system_message_in_the_run_has_content(self): + merged = merge_consecutive_system_messages([{"role": "system"}, {"role": "system"}, {"role": "user", "content": "Hi"}]) + + assert merged == [{"role": "system"}, {"role": "user", "content": "Hi"}] diff --git a/tests/test_litellm/litellm_core_utils/test_agentic_followup_kwargs.py b/tests/test_litellm/litellm_core_utils/test_agentic_followup_kwargs.py new file mode 100644 index 00000000000..af0fbdcc35b --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_agentic_followup_kwargs.py @@ -0,0 +1,66 @@ +from collections.abc import Mapping +from typing import Final + +from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs + + +def _build( + *, + request_kwargs: dict[str, object], + patch_kwargs: dict[str, object], + request_params: set[str], + fingerprints: list[str] | None = None, +) -> Mapping[str, object]: + return build_agentic_followup_kwargs( + request_kwargs=request_kwargs, + patch_kwargs=patch_kwargs, + request_params=request_params, + depth=0, + max_loops=3, + fingerprints=fingerprints if fingerprints is not None else [], + fingerprint="fp", + ) + + +def test_followup_kwargs_never_repeat_a_request_param(): + """Neither source may re-add a key the caller already sends as a request param, or the follow-up call raises a duplicate keyword""" + followup: Final = _build( + request_kwargs={"prompt_cache_key": "thread-1", "api_base": "https://a"}, + patch_kwargs={"prompt_cache_key": "thread-1", "metadata": {"user": "u1"}}, + request_params={"prompt_cache_key", "model", "input"}, + ) + + assert followup.keys().isdisjoint({"prompt_cache_key", "model", "input"}) + assert followup["api_base"] == "https://a" + assert followup["metadata"] == {"user": "u1"} + + +def test_followup_kwargs_let_the_plan_override_the_request(): + followup: Final = _build( + request_kwargs={"api_base": "https://request", "timeout": 5}, + patch_kwargs={"api_base": "https://plan"}, + request_params=set(), + ) + + assert followup["api_base"] == "https://plan" + assert followup["timeout"] == 5 + + +def test_followup_kwargs_carry_the_loop_bookkeeping_without_touching_the_inputs(): + fingerprints: Final = ["earlier"] + request_kwargs: Final = {"_agentic_loop_depth": 0, "max_agentic_loops": 9} + patch_kwargs: Final = {"_agentic_loop_fingerprints": ["stale"]} + + followup: Final = _build( + request_kwargs=request_kwargs, + patch_kwargs=patch_kwargs, + request_params=set(), + fingerprints=fingerprints, + ) + + assert followup["_agentic_loop_depth"] == 1 + assert followup["max_agentic_loops"] == 3 + assert followup["_agentic_loop_fingerprints"] == ["earlier", "fp"] + assert fingerprints == ["earlier"] + assert request_kwargs == {"_agentic_loop_depth": 0, "max_agentic_loops": 9} + assert patch_kwargs == {"_agentic_loop_fingerprints": ["stale"]} diff --git a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py index 434daab6ab5..0d7f735e6e5 100644 --- a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -343,6 +343,40 @@ async def test_dispatcher_runs_followup_with_incremented_depth_and_patched_messa assert logger.cleanup_calls == 1 +@pytest.mark.asyncio +async def test_dispatcher_followup_does_not_repeat_a_request_param_found_in_request_kwargs( + restore_callbacks, +): + """Request kwargs that repeat a request param must not crash the follow-up + with a duplicate keyword, whether or not the plan copies them too.""" + followup = _plain_model_response("done") + request_kwargs = {"temperature": 0.2, "api_base": "https://a"} + plan = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=_patched_messages(), kwargs=dict(request_kwargs)), + ) + litellm.callbacks = [_GateOnlyLogger(plan=plan, tool_calls={"tool_calls": [{"id": "call_abc"}]})] + + acompletion_mock = AsyncMock(return_value=followup) + with patch.object(litellm, "acompletion", acompletion_mock): + result = await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={"temperature": 0.2}, + kwargs=dict(request_kwargs), + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + assert result is followup + acompletion_mock.assert_awaited_once() + call_kwargs = acompletion_mock.await_args.kwargs + assert call_kwargs["temperature"] == 0.2 + assert call_kwargs["api_base"] == "https://a" + + @pytest.mark.asyncio async def test_dispatcher_raises_when_depth_reaches_max_agentic_loops( restore_callbacks, diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index b2ad13c205e..bca61a0e76f 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -6,6 +6,8 @@ import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, + bind_budget_reservation_to_callbacks, + budget_reservation_from_metadata, drop_params_env_flag, drop_params_flag, get_or_create_metadata_bucket, @@ -13,7 +15,60 @@ from litellm.litellm_core_utils.core_helpers import ( normalize_drop_params, reconstruct_model_name, redact_nested_match_and_regex_keys, + unbind_budget_reservation_from_callbacks, ) +from litellm.proxy._types import UserAPIKeyAuth + + +class TestBudgetReservationBinding: + """The request-end release skips a reservation a cost callback has claimed, so the claim + must land on the one dict auth stamped, through whichever metadata field or auth object + carries it, and a failed call must be able to hand it back.""" + + @staticmethod + def _reservation() -> dict: + return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + @pytest.mark.parametrize("metadata_variable_name", ["metadata", "litellm_metadata"]) + def test_reservation_stamped_on_the_metadata_is_bound(self, metadata_variable_name: str): + reservation = self._reservation() + + bind_budget_reservation_to_callbacks({metadata_variable_name: {"user_api_key_budget_reservation": reservation}}) + + assert reservation["callback_bound"] is True + + def test_reservation_reachable_only_through_the_auth_object_is_bound(self): + reservation = self._reservation() + user_api_key_auth = UserAPIKeyAuth(token="hashed") + user_api_key_auth.budget_reservation = reservation + + bind_budget_reservation_to_callbacks({"metadata": {"user_api_key_auth": user_api_key_auth}}) + + assert reservation["callback_bound"] is True + + def test_reservation_reachable_only_through_a_dumped_auth_object_is_bound(self): + reservation = self._reservation() + + bind_budget_reservation_to_callbacks({"metadata": {"user_api_key_auth": {"budget_reservation": reservation}}}) + + assert reservation["callback_bound"] is True + + def test_unbind_hands_a_claimed_reservation_back(self): + reservation = self._reservation() + litellm_params = {"litellm_metadata": {"user_api_key_budget_reservation": reservation}} + bind_budget_reservation_to_callbacks(litellm_params) + + unbind_budget_reservation_from_callbacks(litellm_params) + + assert reservation["callback_bound"] is False + + def test_request_without_a_reservation_binds_nothing(self): + metadata = {"user_api_key_auth": UserAPIKeyAuth(token="hashed")} + + bind_budget_reservation_to_callbacks({"metadata": metadata, "litellm_metadata": None}) + + assert budget_reservation_from_metadata(metadata) is None + assert "user_api_key_budget_reservation" not in metadata class TestGetOrCreateMetadataBucket: diff --git a/tests/test_litellm/litellm_core_utils/test_decode_special_tokens.py b/tests/test_litellm/litellm_core_utils/test_decode_special_tokens.py index d1a5f78a859..dd2daef5484 100644 --- a/tests/test_litellm/litellm_core_utils/test_decode_special_tokens.py +++ b/tests/test_litellm/litellm_core_utils/test_decode_special_tokens.py @@ -1,23 +1,11 @@ -from tokenizers import AddedToken, Tokenizer -from tokenizers.models import WordLevel -from tokenizers.pre_tokenizers import Whitespace -from tokenizers.processors import TemplateProcessing - from litellm import decode, encode +from tokenizers import Tokenizer + +TOKENIZER_JSON = """{"version":"1.0","truncation":null,"padding":null,"added_tokens":[{"id":3,"content":"[BOS]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false,"special":true}],"normalizer":null,"pre_tokenizer":{"type":"Whitespace"},"post_processor":{"type":"TemplateProcessing","single":[{"SpecialToken":{"id":"[BOS]","type_id":0}},{"Sequence":{"id":"A","type_id":0}}],"pair":[{"Sequence":{"id":"A","type_id":0}},{"Sequence":{"id":"B","type_id":1}}],"special_tokens":{"[BOS]":{"id":"[BOS]","ids":[3],"tokens":["[BOS]"]}}},"decoder":null,"model":{"type":"WordLevel","vocab":{"[UNK]":0,"Hello":1,"World":2},"unk_token":"[UNK]"}}""" def _create_custom_tokenizer(): - tokenizer = Tokenizer( - WordLevel({"[UNK]": 0, "Hello": 1, "World": 2}, unk_token="[UNK]") - ) - tokenizer.pre_tokenizer = Whitespace() - tokenizer.add_special_tokens([AddedToken("[BOS]", special=True)]) - bos_token_id = tokenizer.token_to_id("[BOS]") - assert bos_token_id is not None - tokenizer.post_processor = TemplateProcessing( - single="[BOS] $A", - special_tokens=[("[BOS]", bos_token_id)], - ) + tokenizer = Tokenizer.from_str(TOKENIZER_JSON) return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index a34bc2af59d..4c963d14ada 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -67,6 +67,15 @@ class TestGetLitellmParamsKwargsExtraction: assert "s3_endpoint_url" not in result_without_s3_kwargs assert "s3_region_name" not in result_without_s3_kwargs + def test_s3_credential_kwargs_are_forwarded_for_s3_signing(self): + result = get_litellm_params(s3_access_key_id="s3-key", s3_secret_access_key="s3-secret") + assert result["s3_access_key_id"] == "s3-key" + assert result["s3_secret_access_key"] == "s3-secret" + + result_without_s3_kwargs = get_litellm_params() + assert "s3_access_key_id" not in result_without_s3_kwargs + assert "s3_secret_access_key" not in result_without_s3_kwargs + def test_subset_of_kwargs_only_includes_provided(self): """Only provided kwargs appear, others remain absent.""" result = get_litellm_params(azure_ad_token="token123") diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 53fee36b3a8..262dabb7c1b 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -448,6 +448,23 @@ async def test_refetch_records_the_blob_id_of_the_bytes_served_and_the_fetch_eta assert get_model_cost_map_provenance() == {"source_revision": git_blob_id(body), "etag": 'W/"abc123"'} +@pytest.mark.asyncio +async def test_loaded_catalog_snapshot_follows_the_fetched_map_and_ignores_later_registrations(monkeypatch): + import litellm + + edited = json.loads(_real_map_bytes()) + edited["gpt-5.4-mini"]["max_input_tokens"] = 777 + client, _ = _mock_client([httpx.Response(200, content=json.dumps(edited).encode())]) + + result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) + + assert isinstance(result, ModelCostMapReloaded) + monkeypatch.setattr(litellm, "model_cost", result.model_cost_map) + litellm.register_model({"gpt-5.4-mini": {"max_input_tokens": 2048}}, persist_across_reloads=False) + assert litellm.model_cost["gpt-5.4-mini"]["max_input_tokens"] == 2048 + assert GetModelCostMap.loaded_model_cost_map()["gpt-5.4-mini"]["max_input_tokens"] == 777 + + @pytest.mark.asyncio async def test_refetch_revision_follows_the_bytes_not_the_url(): edited = json.loads(_real_map_bytes()) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 277ae33a076..021e012f29d 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,9 +1,10 @@ import asyncio import contextlib import datetime +import logging import os import sys -from collections.abc import Callable +from collections.abc import Callable, Iterator, Mapping from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -25,6 +26,7 @@ from litellm.litellm_core_utils.litellm_logging import ( set_callbacks, ) from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, @@ -302,6 +304,417 @@ def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata(): litellm.model_cost.pop(custom_model_id, None) +class TestZeroCostDiagnostic: + DEPLOYMENT_ID: Final = "lit7898-query-only-priced-deployment" + MODEL_GROUP: Final = "query-only-priced-chat" + QUERY_ONLY_PRICING: Final = {"input_cost_per_query": 0.00042} + PER_SECOND_PRICING: Final = {"input_cost_per_second": 0.00042, "output_cost_per_second": 0.00042} + FREE_PRICING: Final = {"input_cost_per_token": 0, "output_cost_per_token": 0} + + @pytest.fixture(params=["query_only", "free"]) + def deployment_pricing(self, request: pytest.FixtureRequest) -> Iterator[Mapping[str, float]]: + pricing: Final = self.QUERY_ONLY_PRICING if request.param == "query_only" else self.FREE_PRICING + litellm.register_model(model_cost={self.DEPLOYMENT_ID: pricing}, persist_across_reloads=False) + try: + yield pricing + finally: + litellm.model_cost.pop(self.DEPLOYMENT_ID, None) + + def _logging_obj( + self, + pricing: Mapping[str, object], + stream: bool = False, + model: str = "openai/gpt-5.4-nano", + call_type: str = "completion", + deployment_id: str | None = DEPLOYMENT_ID, + custom_llm_provider: str = "openai", + ) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model=model, + messages=[{"role": "user", "content": "Hi"}], + stream=stream, + call_type=call_type, + start_time=time.time(), + litellm_call_id="lit7898", + function_id="fn", + ) + self._route_to_deployment( + logging_obj, pricing, model=model, deployment_id=deployment_id, custom_llm_provider=custom_llm_provider + ) + return logging_obj + + def _route_to_deployment( + self, + logging_obj: LitellmLogging, + pricing: Mapping[str, object], + model: str = "openai/gpt-5.4-nano", + deployment_id: str | None = DEPLOYMENT_ID, + custom_llm_provider: str = "openai", + ) -> None: + model_info: Final = pricing if deployment_id is None else {"id": deployment_id, **pricing} + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params={"metadata": {"model_group": self.MODEL_GROUP, "model_info": model_info}}, + custom_llm_provider=custom_llm_provider, + ) + + @staticmethod + def _response( + usage: litellm.Usage | None = None, model: str = "gpt-5.4-nano", **hidden_params: object + ) -> ModelResponse: + response: Final = ModelResponse( + model=model, + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="hello"))], + usage=usage, + ) + response._hidden_params = {"custom_llm_provider": "openai", **hidden_params} + return response + + @staticmethod + def _zero_cost_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.name == "LiteLLM" and record.levelno == logging.WARNING and "priced at $0" in record.getMessage() + ] + + def _assert_flagged(self, logging_obj: LitellmLogging, caplog: pytest.LogCaptureFixture) -> None: + assert logging_obj.model_call_details["zero_cost_diagnostic"] == { + "reason": "missing_pricing_key", + "pricing_model": self.DEPLOYMENT_ID, + "missing_pricing_keys": ("input_cost_per_token", "output_cost_per_token"), + } + warnings: Final = self._zero_cost_warnings(caplog) + assert len(warnings) == 1 + assert f"model_group={self.MODEL_GROUP}" in warnings[0] + assert f"pricing entry '{self.DEPLOYMENT_ID}' has no input_cost_per_token, output_cost_per_token" in warnings[0] + + def test_zero_cost_with_a_missing_rate_warns_once_and_is_recorded( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + first_cost: Final = logging_obj._response_cost_calculator(result=self._response(usage)) + second_cost: Final = logging_obj._response_cost_calculator(result=self._response(usage)) + + assert first_cost == 0.0 + assert second_cost == 0.0 + if deployment_pricing is self.FREE_PRICING: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + + def test_usage_less_stream_chunk_does_not_hide_the_final_response_diagnostic( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=8, completion_tokens=2, total_tokens=10) + logging_obj: Final = self._logging_obj(deployment_pricing, stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._response_cost_calculator(result=self._response(usage=None)) + logging_obj._response_cost_calculator(result=self._response(usage)) + + if deployment_pricing is self.FREE_PRICING: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + + def test_terminal_responses_stream_event_is_judged_by_its_inner_response( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + logging_obj: Final = self._logging_obj(deployment_pricing, stream=True, call_type="aresponses") + event: Final = ResponseCompletedEvent( + type="response.completed", + response=ResponsesAPIResponse( + id="resp-lit7898", + created_at=1, + object="response", + status="completed", + model="gpt-5.4-nano", + output=[], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=20, total_tokens=30), + ), + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cost: Final = logging_obj._response_cost_calculator(result=event) + + assert cost == 0.0 + if deployment_pricing is self.FREE_PRICING: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + + def test_precomputed_zero_hidden_cost_is_flagged_and_lands_in_the_payload( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + response: Final = self._response(usage, response_cost=0.0, model_id=self.DEPLOYMENT_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._process_hidden_params_and_response_cost( + response, start_time=datetime.datetime.now(), end_time=datetime.datetime.now() + ) + + payload: Final = logging_obj.model_call_details["standard_logging_object"] + assert payload["response_cost"] == 0.0 + if deployment_pricing is self.FREE_PRICING: + assert payload["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + assert payload["zero_cost_diagnostic"] == logging_obj.model_call_details["zero_cost_diagnostic"] + + def test_uncomputed_hidden_cost_is_not_a_zero_cost( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + response: Final = self._response(usage, response_cost=None, model_id=self.DEPLOYMENT_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._process_hidden_params_and_response_cost( + response, start_time=datetime.datetime.now(), end_time=datetime.datetime.now() + ) + + assert logging_obj.model_call_details["standard_logging_object"]["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + + def test_unbilled_read_route_with_usage_stays_silent( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing, call_type="aget_responses") + response: Final = self._response(usage, response_cost=0.0, model_id=self.DEPLOYMENT_ID) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._process_hidden_params_and_response_cost( + response, start_time=datetime.datetime.now(), end_time=datetime.datetime.now() + ) + + assert logging_obj.model_call_details["standard_logging_object"]["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + + def test_unmapped_model_that_fails_cost_calculation_stays_silent(self, caplog: pytest.LogCaptureFixture) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj( + {}, model="openai/lit7898-unmapped-model", deployment_id="lit7898-unmapped-deployment" + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cost: Final = logging_obj._response_cost_calculator( + result=self._response(usage, model="lit7898-unmapped-model") + ) + + assert cost is None + assert logging_obj.model_call_details["response_cost_failure_debug_information"] is not None + assert logging_obj.model_call_details.get("zero_cost_diagnostic") is None + assert self._zero_cost_warnings(caplog) == [] + + def test_malformed_usage_never_raises_out_of_the_cost_calculator( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + logging_obj: Final = self._logging_obj(deployment_pricing) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + cost: Final = logging_obj._response_cost_calculator( + result={"model": "gpt-5.4-nano", "usage": {"prompt_tokens": "n/a", "completion_tokens": 3}} + ) + + assert cost is None + assert logging_obj.model_call_details.get("zero_cost_diagnostic") is None + assert self._zero_cost_warnings(caplog) == [] + + def test_usage_less_evaluation_between_two_zero_cost_findings_does_not_warn_twice( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=8, completion_tokens=2, total_tokens=10) + logging_obj: Final = self._logging_obj(deployment_pricing, stream=True, call_type="anthropic_messages") + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._response_cost_calculator(result=self._response(usage=None)) + logging_obj._response_cost_calculator(result=self._response(usage)) + logging_obj._response_cost_calculator(result=self._response(usage=None)) + logging_obj._response_cost_calculator(result=self._response(usage)) + + if deployment_pricing is self.FREE_PRICING: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + return + self._assert_flagged(logging_obj, caplog) + + def test_retry_that_prices_clears_the_diagnostic_and_a_later_zero_cost_is_recorded_silently( + self, caplog: pytest.LogCaptureFixture + ) -> None: + priced_id: Final = "lit7898-priced-deployment" + priced_pricing: Final = {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06} + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + litellm.register_model( + model_cost={self.DEPLOYMENT_ID: self.QUERY_ONLY_PRICING, priced_id: priced_pricing}, + persist_across_reloads=False, + ) + try: + logging_obj: Final = self._logging_obj(self.QUERY_ONLY_PRICING) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logging_obj._response_cost_calculator(result=self._response(usage)) == 0.0 + self._assert_flagged(logging_obj, caplog) + + self._route_to_deployment(logging_obj, priced_pricing, deployment_id=priced_id) + assert logging_obj._response_cost_calculator(result=self._response(usage)) == pytest.approx(5e-05) + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + + self._route_to_deployment(logging_obj, self.QUERY_ONLY_PRICING) + assert logging_obj._response_cost_calculator(result=self._response(usage)) == 0.0 + + assert logging_obj.model_call_details["zero_cost_diagnostic"]["reason"] == "missing_pricing_key" + assert len(self._zero_cost_warnings(caplog)) == 1 + finally: + litellm.model_cost.pop(self.DEPLOYMENT_ID, None) + litellm.model_cost.pop(priced_id, None) + + def test_one_request_evaluated_against_two_cost_map_entries_warns_once( + self, caplog: pytest.LogCaptureFixture + ) -> None: + dated_model: Final = "lit7898-nano-2026-03-17" + requested_model: Final = "lit7898-nano" + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + cost_map_entry: Final = {"litellm_provider": "openai", "mode": "chat", **self.QUERY_ONLY_PRICING} + litellm.register_model( + model_cost={dated_model: cost_map_entry, requested_model: cost_map_entry}, persist_across_reloads=False + ) + try: + logging_obj: Final = self._logging_obj( + {}, model=f"openai/{requested_model}", deployment_id="lit7898-cost-map-deployment" + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logging_obj._response_cost_calculator(result=self._response(usage, model=dated_model)) == 0.0 + assert logging_obj._response_cost_calculator(result=self._response(usage, model=requested_model)) == 0.0 + + assert logging_obj.model_call_details["zero_cost_diagnostic"]["pricing_model"] == requested_model + warnings: Final = self._zero_cost_warnings(caplog) + assert len(warnings) == 1 + assert f"pricing entry '{dated_model}' has no input_cost_per_token, output_cost_per_token" in warnings[0] + finally: + litellm.model_cost.pop(dated_model, None) + litellm.model_cost.pop(requested_model, None) + + def test_free_deployment_without_a_router_id_is_judged_by_its_own_pricing( + self, caplog: pytest.LogCaptureFixture + ) -> None: + global_model: Final = "lit7898-priced-global" + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + litellm.register_model( + model_cost={ + global_model: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = self._logging_obj(self.FREE_PRICING, model=global_model, deployment_id=None) + response: Final = self._response(usage, model=global_model, response_cost=0.0) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logging_obj._process_hidden_params_and_response_cost( + response, start_time=datetime.datetime.now(), end_time=datetime.datetime.now() + ) + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + finally: + litellm.model_cost.pop(global_model, None) + + def test_cache_hit_priced_for_saved_cost_stays_silent( + self, deployment_pricing: Mapping[str, float], caplog: pytest.LogCaptureFixture + ) -> None: + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + logging_obj: Final = self._logging_obj(deployment_pricing) + logging_obj.model_call_details["cache_hit"] = True + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logging_obj._response_cost_calculator(result=self._response(usage), cache_hit=False) == 0.0 + + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + + def test_per_second_priced_deployment_bills_the_call_duration_and_stays_silent( + self, caplog: pytest.LogCaptureFixture + ) -> None: + per_second_id: Final = "lit8315-per-second-priced-deployment" + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + litellm.register_model(model_cost={per_second_id: self.PER_SECOND_PRICING}, persist_across_reloads=False) + try: + logging_obj: Final = self._logging_obj(self.PER_SECOND_PRICING, deployment_id=per_second_id) + response: Final = self._response(usage) + response._response_ms = 1000.0 + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.00084) + + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert self._zero_cost_warnings(caplog) == [] + finally: + litellm.model_cost.pop(per_second_id, None) + + @pytest.mark.parametrize("spilled_over", [True, False]) + def test_ptu_deployment_is_judged_by_the_entry_the_calculator_priced_with( + self, spilled_over: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + router_model_id: Final = "lit7898-ptu-router-model-id" + served_model: Final = "azure/lit7898-ptu-served-model" + ptu_model_info: Final = { + "team_id": "team-1", + "ptu_count": 100, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01", + **self.FREE_PRICING, + } + usage: Final = litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + litellm.register_model( + model_cost={ + router_model_id: {**self.FREE_PRICING, "litellm_provider": "azure", "mode": "chat"}, + served_model: {**self.QUERY_ONLY_PRICING, "litellm_provider": "azure", "mode": "chat"}, + }, + persist_across_reloads=False, + ) + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + try: + logging_obj: Final = self._logging_obj( + ptu_model_info, model=served_model, deployment_id=router_model_id, custom_llm_provider="azure" + ) + spillover_headers: Final = {"llm_provider-x-ms-is-spilled-over": "true"} if spilled_over else {} + response: Final = self._response( + usage, model=served_model, custom_llm_provider="azure", additional_headers=spillover_headers + ) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logging_obj._response_cost_calculator(result=response) == 0.0 + + warnings: Final = self._zero_cost_warnings(caplog) + if not spilled_over: + assert logging_obj.model_call_details["zero_cost_diagnostic"] is None + assert warnings == [] + return + assert logging_obj.model_call_details["zero_cost_diagnostic"] == { + "reason": "missing_pricing_key", + "pricing_model": served_model, + "missing_pricing_keys": ("input_cost_per_token", "output_cost_per_token"), + } + assert len(warnings) == 1 + assert f"pricing entry '{served_model}' has no input_cost_per_token, output_cost_per_token" in warnings[0] + finally: + litellm.model_cost.pop(router_model_id, None) + litellm.model_cost.pop(served_model, None) + + class TestGetRouterModelId: """Tests for the get_router_model_id helper method.""" @@ -407,7 +820,6 @@ class TestGetRouterDeploymentModelInfo: logging_obj.litellm_params = {"api_base": ""} assert logging_obj.get_router_deployment_model_info() is None - def test_a_published_batch_rate_never_displaces_a_declared_standard_rate(self) -> None: """Ownership is per token direction, not per field. @@ -1111,7 +1523,9 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch): @pytest.mark.asyncio -async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(monkeypatch: pytest.MonkeyPatch): +async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log( + monkeypatch: pytest.MonkeyPatch, +): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.responses.main import base_llm_http_handler @@ -3109,6 +3523,20 @@ def _make_logging_obj(stream: bool) -> LitellmLogging: ) +def test_get_response_ms_measures_a_float_start_time_against_a_datetime_end_time(): + """The files paths construct the logging object with ``time.time()`` while the success + handler stamps a datetime end, and the per-second cost path reads this window.""" + logging_obj = _make_logging_obj(stream=False) + logging_obj.update_environment_variables( + model="openai/codex-mini-latest", user="", optional_params={}, litellm_params={} + ) + start_seconds = logging_obj.model_call_details["start_time"] + assert isinstance(start_seconds, float) + logging_obj.model_call_details["end_time"] = datetime.datetime.fromtimestamp(start_seconds + 1.5) + + assert logging_obj.get_response_ms() == pytest.approx(1500) + + def test_get_assembled_streaming_response_returns_none_for_non_streaming(): """Non-streaming requests should return None so the streaming block is skipped.""" import datetime @@ -7066,22 +7494,41 @@ async def test_classifier_audit_matches_provider_transport(provider: str) -> Non return httpx.Response(200, json=mock_responses_api_response(content).model_dump()) if provider == "anthropic": - return httpx.Response(200, json={ - "id": "msg-audit", "type": "message", "role": "assistant", "model": "claude-haiku-4-5", - "content": [{"type": "text", "text": content}], "stop_reason": "end_turn", - "usage": {"input_tokens": 10, "output_tokens": 5}, - }) + return httpx.Response( + 200, + json={ + "id": "msg-audit", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": content}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) if provider == "bedrock": - return httpx.Response(200, json={ - "output": {"message": {"role": "assistant", "content": [{"text": content}]}}, - "stopReason": "end_turn", "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, - "metrics": {"latencyMs": 1}, - }) - return httpx.Response(200, json={ - "id": "chatcmpl-audit", "object": "chat.completion", "created": 0, "model": "gpt-5.6", - "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - }) + return httpx.Response( + 200, + json={ + "output": {"message": {"role": "assistant", "content": [{"text": content}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + "metrics": {"latencyMs": 1}, + }, + ) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-audit", + "object": "chat.completion", + "created": 0, + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) async def capture(kwargs, response_obj, start_time, end_time): logs.put_nowait(kwargs["standard_logging_object"]) @@ -7092,11 +7539,15 @@ async def test_classifier_audit_matches_provider_transport(provider: str) -> Non handler.client = http_client client: Final = ( AsyncAzureOpenAI( - api_key="transport-only", azure_endpoint="https://azure.invalid", - api_version="2025-04-01-preview", http_client=http_client, + api_key="transport-only", + azure_endpoint="https://azure.invalid", + api_version="2025-04-01-preview", + http_client=http_client, ) - if provider == "azure" else AsyncOpenAI(api_key="transport-only", http_client=http_client) - if provider == "openai" else handler + if provider == "azure" + else AsyncOpenAI(api_key="transport-only", http_client=http_client) + if provider == "openai" + else handler ) model: Final = { "openai": "openai/gpt-5.6", @@ -7109,23 +7560,44 @@ async def test_classifier_audit_matches_provider_transport(provider: str) -> Non async def run(marker: str) -> None: if provider == "responses": await litellm.aresponses( - model=model, api_key="transport-only", client=client, max_output_tokens=128, - instructions="classifier-rubric", input=marker, + model=model, + api_key="transport-only", + client=client, + max_output_tokens=128, + instructions="classifier-rubric", + input=marker, metadata={"internal_call_origin": "autorouter_classifier"}, proxy_server_request={"body": {}, "originating_request_masked": {"input": f"source-only-{marker}"}}, - success_callback=[capture], num_retries=0, + success_callback=[capture], + num_retries=0, ) return await litellm.acompletion( - model=model, api_key="transport-only", client=client, max_tokens=128, - aws_access_key_id="transport-only", aws_secret_access_key="transport-only", aws_region_name="us-east-1", + model=model, + api_key="transport-only", + client=client, + max_tokens=128, + aws_access_key_id="transport-only", + aws_secret_access_key="transport-only", + aws_region_name="us-east-1", messages=[{"role": "system", "content": "classifier-rubric"}, {"role": "user", "content": marker}], metadata={"internal_call_origin": "autorouter_classifier"}, proxy_server_request={"body": {}, "originating_request_masked": {"input": f"source-only-{marker}"}}, - success_callback=[capture], num_retries=0, - **({"api_base": "https://azure.invalid", "api_version": "2025-04-01-preview"} if provider == "azure" else {}), - **({"extra_body": {"audit_context": "provider-extra"}, "extra_headers": {"X-Audit": "header-only-secret"}} - if provider in ("openai", "azure") else {}), + success_callback=[capture], + num_retries=0, + **( + {"api_base": "https://azure.invalid", "api_version": "2025-04-01-preview"} + if provider == "azure" + else {} + ), + **( + { + "extra_body": {"audit_context": "provider-extra"}, + "extra_headers": {"X-Audit": "header-only-secret"}, + } + if provider in ("openai", "azure") + else {} + ), ) await asyncio.gather(run("request-one"), run("request-two")) @@ -7149,14 +7621,17 @@ async def test_classifier_audit_matches_provider_transport(provider: str) -> Non @pytest.mark.parametrize("redaction", ["none", "global", "request", "header"]) @pytest.mark.parametrize("status", ["success", "failure"]) @pytest.mark.parametrize("call_type", ["completion", "acompletion", "responses", "aresponses"]) -def test_classifier_audit_obeys_message_logging_before_payload_emission(logging_obj, monkeypatch, redaction, status, call_type): +def test_classifier_audit_obeys_message_logging_before_payload_emission( + logging_obj, monkeypatch, redaction, status, call_type +): from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload monkeypatch.setattr(litellm, "turn_off_message_logging", redaction == "global") params: Final = { - "metadata": {"internal_call_origin": "autorouter_classifier", **( - {"headers": {"x-litellm-enable-message-redaction": "true"}} if redaction == "header" else {} - )}, + "metadata": { + "internal_call_origin": "autorouter_classifier", + **({"headers": {"x-litellm-enable-message-redaction": "true"}} if redaction == "header" else {}), + }, "proxy_server_request": {"body": {}, "originating_request_masked": {"input": "source-only"}}, } logging_obj.call_type = call_type @@ -7169,8 +7644,12 @@ def test_classifier_audit_obeys_message_logging_before_payload_emission(logging_ ) now: Final = datetime.datetime.now() payload: Final = get_standard_logging_object_payload( - kwargs={**logging_obj.model_call_details, "call_type": call_type}, init_response_obj={}, - start_time=now, end_time=now, logging_obj=logging_obj, status=status, + kwargs={**logging_obj.model_call_details, "call_type": call_type}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status=status, ) assert payload is not None if redaction == "none": @@ -7421,7 +7900,13 @@ def _completed_responses_event(usage: ResponseAPIUsage) -> ResponseCompletedEven return ResponseCompletedEvent( type="response.completed", response=ResponsesAPIResponse( - id="resp-1", created_at=1, object="response", status="completed", model="codex-mini-latest", output=[], usage=usage + id="resp-1", + created_at=1, + object="response", + status="completed", + model="codex-mini-latest", + output=[], + usage=usage, ), ) @@ -7441,7 +7926,9 @@ def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost() now = datetime.datetime.now() assembled = logging_obj._get_assembled_streaming_response( - result=_completed_responses_event(ResponseAPIUsage(input_tokens=12, output_tokens=2, total_tokens=14, cost=0.0042)), + result=_completed_responses_event( + ResponseAPIUsage(input_tokens=12, output_tokens=2, total_tokens=14, cost=0.0042) + ), start_time=now, end_time=now, is_async=True, @@ -7467,3 +7954,40 @@ def test_get_assembled_streaming_response_without_usage_cost_leaves_pricing_to_t assert "additional_headers" not in assembled._hidden_params price_map_cost = logging_obj._response_cost_calculator(result=assembled) assert price_map_cost is not None and 0 < price_map_cost != 0.0042 + + +def test_response_cost_calculator_prices_terminal_responses_event_from_its_response(): + logging_obj: Final = _responses_stream_logging_obj() + inner_response: Final = ResponsesAPIResponse( + id="resp-priced", + created_at=1, + object="response", + status="completed", + model="gpt-4o-mini", + output=[], + usage=ResponseAPIUsage(input_tokens=1840, output_tokens=412, total_tokens=2252), + ) + event: Final = ResponseCompletedEvent(type="response.completed", response=inner_response) + + event_cost: Final = logging_obj._response_cost_calculator(result=event) + inner_cost: Final = logging_obj._response_cost_calculator(result=inner_response) + + assert event_cost is not None and event_cost > 0 + assert event_cost == inner_cost + assert logging_obj.cost_breakdown["input_cost"] is not None and logging_obj.cost_breakdown["input_cost"] > 0 + + +class TestBudgetReservationBinding: + """The proxy builds a logging object for every route before calling anything, so a + logging object seeing the reservation is no promise that a cost callback will settle + it: the claim belongs to the call wrapper, and this object must leave it unbound.""" + + def test_update_environment_variables_leaves_the_reservation_unbound(self, logging_obj): + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + logging_obj.update_environment_variables( + litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}, optional_params={} + ) + + assert logging_obj.litellm_params["metadata"]["user_api_key_budget_reservation"] is reservation + assert reservation["callback_bound"] is False diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 276a67e0bd4..22298c00219 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -304,6 +304,26 @@ class TestPerformRedaction: assert delta["thinking_blocks"] is None assert delta["audio"] is None + def test_redacts_text_completion_choices_in_standard_logging_object(self): + details = { + "standard_logging_object": { + "response": { + "object": "text_completion", + "choices": [ + {"text": " Paris.", "finish_reason": "stop", "index": 0}, + {"text": "\n\nBlue", "finish_reason": "length", "index": 1}, + ], + } + } + } + + perform_redaction(details, None) + + assert details["standard_logging_object"]["response"]["choices"] == [ + {"text": "redacted-by-litellm", "finish_reason": "stop", "index": 0}, + {"text": "redacted-by-litellm", "finish_reason": "length", "index": 1}, + ] + def test_redacts_object_choices_inside_model_response_dict(self): result = { "choices": [ diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index f19a8891609..5ce6a4b1ce9 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -786,23 +786,23 @@ def test_token_counter(): import unittest -from litellm.utils import _select_tokenizer_helper, claude_json_str, encoding +from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding # Clear the cache at module load to ensure clean state -_select_tokenizer_helper.cache_clear() +_load_huggingface_tokenizer.cache_clear() class TestTokenizerSelection(unittest.TestCase): def setUp(self): """Clear the LRU cache before each test method. - The _select_tokenizer_helper function is decorated with @lru_cache, - which can cause cache hits from previous tests when running with + The HuggingFace tokenizers behind _select_tokenizer_helper are cached with + @lru_cache, which can cause cache hits from previous tests when running with --dist=loadscope (tests from same file run on same worker). """ - _select_tokenizer_helper.cache_clear() + _load_huggingface_tokenizer.cache_clear() - @patch("litellm.utils.Tokenizer.from_pretrained") + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") def test_llama3_tokenizer_api_failure(self, mock_from_pretrained): # Setup mock to raise an error mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") @@ -817,7 +817,7 @@ class TestTokenizerSelection(unittest.TestCase): self.assertEqual(result["type"], "openai_tokenizer") self.assertEqual(result["tokenizer"], encoding) - @patch("litellm.utils.Tokenizer.from_pretrained") + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") def test_cohere_tokenizer_api_failure(self, mock_from_pretrained): # Setup mock to raise an error mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") @@ -837,10 +837,10 @@ class TestTokenizerSelection(unittest.TestCase): self.assertEqual(result["type"], "openai_tokenizer") self.assertEqual(result["tokenizer"], encoding) - @patch("litellm.utils.Tokenizer.from_str") - def test_claude_tokenizer_api_failure(self, mock_from_str): + @patch("litellm.utils.tokenizer_dispatch.anthropic") + def test_claude_tokenizer_api_failure(self, mock_anthropic): # Setup mock to raise an error - mock_from_str.side_effect = Exception("Failed to load tokenizer") + mock_anthropic.side_effect = Exception("Failed to load tokenizer") # Add Claude model to the list for testing litellm.anthropic_models = ["claude-2"] @@ -849,13 +849,13 @@ class TestTokenizerSelection(unittest.TestCase): result = _select_tokenizer_helper("claude-2") # Verify the attempt to load Claude tokenizer - mock_from_str.assert_called_once_with(claude_json_str) + mock_anthropic.assert_called_once_with() # Verify fallback to OpenAI tokenizer self.assertEqual(result["type"], "openai_tokenizer") self.assertEqual(result["tokenizer"], encoding) - @patch("litellm.utils.Tokenizer.from_pretrained") + @patch("litellm.utils.tokenizer_dispatch.from_pretrained") def test_llama2_tokenizer_api_failure(self, mock_from_pretrained): # Setup mock to raise an error mock_from_pretrained.side_effect = Exception("Failed to load tokenizer") diff --git a/tests/test_litellm/litellm_core_utils/test_tokenizer.py b/tests/test_litellm/litellm_core_utils/test_tokenizer.py new file mode 100644 index 00000000000..aa4a0fc6a1c --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_tokenizer.py @@ -0,0 +1,403 @@ +import copy +import os +import pickle +import subprocess +import sys +from pathlib import Path +from typing import Final, Literal + +import pytest +import tiktoken +from tokenizers import Tokenizer as ReferenceTokenizer + +import litellm +from litellm.caching._embedding_router import truncate_embedding_input +from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding +from litellm.utils import claude_json_str +from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON + + +@pytest.mark.parametrize( + "name", ("cl100k_base", "o200k_base", "p50k_base", "p50k_edit", "r50k_base", "gpt2", "o200k_harmony") +) +@pytest.mark.parametrize( + "text", ("hello world", "café 漢字 🙂", "", "a\ud800b", "\ud83d\ude42", "🙂\ud83d\ude42\udfff", " " * 64) +) +def test_openai_encoding_matches_python_unicode_and_batches(name: str, text: str) -> None: + reference: Final = tiktoken.get_encoding(name) + encoding: Final = OpenAIEncoding.from_tiktoken(name) + expected: Final = reference.encode(text) + + assert encoding.encode(text) == expected + assert encoding.count(text) == len(expected) + assert encoding.encode_batch([text], num_threads=2) == reference.encode_batch([text], num_threads=2) + assert encoding.encode_ordinary_batch([text]) == reference.encode_ordinary_batch([text]) + assert encoding.decode_batch([expected]) == reference.decode_batch([expected]) + assert encoding.decode_bytes_batch([expected]) == reference.decode_bytes_batch([expected]) + + +@pytest.mark.parametrize("allowed", (frozenset(), frozenset({"<|endoftext|>"}), "all")) +@pytest.mark.parametrize("disallowed", (frozenset(), frozenset({"<|fim_prefix|>"}), "all")) +def test_openai_special_token_options_match_python( + allowed: frozenset[str] | Literal["all"], disallowed: frozenset[str] | Literal["all"] +) -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + encoding: Final = OpenAIEncoding.from_tiktoken(reference.name) + text: Final = "hello<|endoftext|><|fim_prefix|>world" + allowed_set: Final = reference.special_tokens_set if allowed == "all" else allowed + disallowed_set: Final = reference.special_tokens_set - allowed_set if disallowed == "all" else disallowed + if any(token in text for token in disallowed_set): + with pytest.raises(ValueError, match="disallowed special token"): + encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed) + return + assert encoding.encode(text, allowed_special=allowed, disallowed_special=disallowed) == reference.encode( + text, allowed_special=allowed, disallowed_special=disallowed + ) + assert encoding.special_tokens_set == reference.special_tokens_set + assert encoding.eot_token == reference.eot_token + + +@pytest.mark.parametrize("errors", ("replace", "ignore", "backslashreplace", "strict")) +def test_openai_partial_token_decoding_preserves_error_policy(errors: str) -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + encoding: Final = OpenAIEncoding.from_tiktoken(reference.name) + tokens: Final = reference.encode("🙂")[:1] + assert encoding.decode_bytes(tokens) == reference.decode_bytes(tokens) + if errors == "strict": + with pytest.raises(UnicodeDecodeError): + encoding.decode(tokens, errors=errors) + return + assert encoding.decode(tokens, errors=errors) == reference.decode(tokens, errors=errors) + assert encoding.decode_tokens_bytes(tokens) == reference.decode_tokens_bytes(tokens) + + +def test_public_encoding_and_semantic_cache_preserve_truncated_unicode() -> None: + reference: Final = tiktoken.get_encoding(litellm.encoding.name) + text: Final = "🙂" + tokens: Final = reference.encode(text) + + assert litellm.encoding.encode(text, disallowed_special=()) == tokens + assert litellm.encoding.encode_batch([text]) == [tokens] + assert litellm.decode(tokens=tokens[:1]) == reference.decode(tokens[:1]) + assert truncate_embedding_input(text, "", 1) == reference.decode(tokens[:1]) + + +@pytest.mark.parametrize("add_special_tokens", (True, False)) +def test_huggingface_encoding_preserves_result_fields_and_serialization(add_special_tokens: bool) -> None: + reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON) + expected: Final = reference.encode("Hello World", add_special_tokens=add_special_tokens) + actual: Final = tokenizer.encode("Hello World", add_special_tokens=add_special_tokens) + + assert (actual.ids, actual.tokens, actual.type_ids, actual.offsets, actual.word_ids, actual.sequence_ids) == ( + expected.ids, + expected.tokens, + expected.type_ids, + expected.offsets, + expected.word_ids, + expected.sequence_ids, + ) + assert (actual.attention_mask, actual.special_tokens_mask, actual.n_sequences, len(actual)) == ( + expected.attention_mask, + expected.special_tokens_mask, + expected.n_sequences, + len(expected), + ) + assert copy.deepcopy(actual).ids == expected.ids + assert pickle.loads(pickle.dumps(actual)).offsets == expected.offsets + assert tokenizer.decode(actual.ids, skip_special_tokens=False) == reference.decode( + expected.ids, skip_special_tokens=False + ) + + +def test_huggingface_character_offsets_and_pretokenized_pairs_match_python() -> None: + reference: Final = ReferenceTokenizer.from_str(claude_json_str) + tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str) + text: Final = "café 漢字 🙂" + actual: Final = tokenizer.encode(text) + expected: Final = reference.encode(text) + + assert actual.offsets == expected.offsets + assert actual.ids == expected.ids + assert ( + tokenizer.encode(["hello", "world"], ["again"], is_pretokenized=True).ids + == reference.encode(["hello", "world"], ["again"], is_pretokenized=True).ids + ) + + +def test_huggingface_batches_apply_padding_across_inputs() -> None: + reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + reference.enable_padding(pad_id=0, pad_token="[UNK]") + tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str()) + inputs: Final = ["Hello", ("Hello World", "World")] + expected: Final = reference.encode_batch(inputs) + actual: Final = tokenizer.encode_batch(inputs) + fast: Final = tokenizer.encode_batch_fast(inputs) + + assert [(item.ids, item.attention_mask, item.offsets) for item in actual] == [ + (item.ids, item.attention_mask, item.offsets) for item in expected + ] + assert [item.ids for item in fast] == [item.ids for item in expected] + assert tokenizer.decode_batch([item.ids for item in actual]) == reference.decode_batch( + [item.ids for item in expected] + ) + + +def test_caller_supplied_huggingface_tokenizer_preserves_public_encode_and_count() -> None: + tokenizer: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + custom: Final = {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + expected: Final = tokenizer.encode("Hello World").ids + + assert litellm.encode(text="Hello World", custom_tokenizer=custom) == expected + assert litellm.token_counter(text="Hello World", custom_tokenizer=custom) == len(expected) + assert litellm.decode(tokens=expected, custom_tokenizer=custom) == "Hello World" + + +def test_caller_supplied_tiktoken_treats_special_spellings_as_text() -> None: + tokenizer: Final = tiktoken.get_encoding("cl100k_base") + custom: Final = {"type": "openai_tokenizer", "tokenizer": tokenizer} + text: Final = "<|endoftext|>" + + assert litellm.encode(text=text, custom_tokenizer=custom) == tokenizer.encode(text, disallowed_special=()) + + +def test_public_tokenizer_objects_survive_pickle_and_deepcopy(tmp_path: Path) -> None: + custom: Final = litellm.create_tokenizer(TOKENIZER_JSON) + tokenizer: Final = custom["tokenizer"] + path: Final = tmp_path / "tokenizer.json" + tokenizer.save(str(path)) + + assert copy.deepcopy(custom)["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids + assert ( + pickle.loads(pickle.dumps(custom))["tokenizer"].encode("Hello World").ids == tokenizer.encode("Hello World").ids + ) + assert HuggingFaceTokenizer.from_file(str(path)).encode("Hello World").ids == tokenizer.encode("Hello World").ids + assert copy.deepcopy(litellm.encoding).encode("hello") == litellm.encoding.encode("hello") + assert pickle.loads(pickle.dumps(litellm.encoding)).encode("hello") == litellm.encoding.encode("hello") + + +@pytest.mark.parametrize("offline", ("0", "1")) +def test_hub_loader_preserves_environment_auth_cache_and_offline(tmp_path: Path, offline: str) -> None: + script: Final = """ +import json +import sys +from pathlib import Path +sys.path.insert(0, sys.argv[1]) +import httpx +import huggingface_hub +from huggingface_hub.errors import LocalEntryNotFoundError +import litellm +payload = sys.argv[2].encode() +offline = sys.argv[3] == "1" +observed = [] +def handle(request): + assert not offline, "offline loading issued a request" + if request.url.path.endswith("/tokenizer.json"): + observed.append(request.headers.get("authorization")) + if request.headers.get("authorization") != "Bearer audit-fixture-token": + return httpx.Response(401) + return httpx.Response(200, headers={"content-length": str(len(payload)), "etag": '"fixture"', "x-repo-commit": "a" * 40}, content=payload if request.method == "GET" else b"") +if not offline: + huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(handle))) +try: + tokenizer = litellm.create_pretrained_tokenizer("test-fixture/tokenizer")["tokenizer"] +except LocalEntryNotFoundError: + assert offline + assert observed == [] +else: + assert not offline + assert "Bearer audit-fixture-token" in observed + assert tokenizer.decode(tokenizer.encode("Hello World").ids) == "Hello World" + assert tuple(Path(sys.argv[4]).rglob("tokenizer.json")) +print("compatible") +""" + result: Final = subprocess.run( + [ + sys.executable, + "-I", + "-c", + script, + str(Path(litellm.__file__).parent.parent), + TOKENIZER_JSON, + offline, + str(tmp_path / "cache"), + ], + capture_output=True, + text=True, + timeout=30, + env={ + **os.environ, + "HF_HOME": str(tmp_path / "home"), + "HF_HUB_CACHE": str(tmp_path / "cache"), + "HF_ENDPOINT": "http://127.0.0.1:9", + "HF_TOKEN": "audit-fixture-token", + "HF_HUB_OFFLINE": offline, + "HF_HUB_DISABLE_IMPLICIT_TOKEN": "0", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + }, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.strip() == "compatible" + + +@pytest.mark.parametrize("rust", (None, "0", "1")) +def test_tokenization_without_native_extension_stays_offline(tmp_path: Path, rust: str | None) -> None: + script: Final = """ +import importlib.abc +import sys +sys.path.insert(0, sys.argv[1]) +def reject_network(event, args): + if event == "socket.connect": + raise AssertionError("tokenizer attempted a network connection") +sys.addaudithook(reject_network) +class Block(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "litellm.rust_bridge._native": + raise ImportError("native extension is unavailable") +sys.meta_path.insert(0, Block()) +import litellm +from litellm.rust_bridge.tokenizer import get_encoding +import tiktoken +from tokenizers import Tokenizer +assert isinstance(litellm.encoding, tiktoken.Encoding) +for name in ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit"): + encoding = get_encoding(name) + text = "offline café 漢字 🙂" + " " * 64 + assert encoding.decode(encoding.encode(text)) == text +ids = litellm.encode(text="hello world") +assert litellm.decode(tokens=ids) == "hello world" +assert litellm.token_counter(model=None, text="hello world") == len(ids) +custom = litellm.create_tokenizer(sys.argv[2]) +assert isinstance(custom["tokenizer"], Tokenizer) +custom["tokenizer"].enable_padding(pad_id=0, pad_token="[UNK]") +assert litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) == "Hello World" +print("compatible") +""" + result: Final = subprocess.run( + [sys.executable, "-I", "-c", script, str(Path(litellm.__file__).parent.parent), TOKENIZER_JSON], + capture_output=True, + text=True, + timeout=30, + cwd=tmp_path, + env={ + **{key: value for key, value in os.environ.items() if key != "LITELLM_RUST"}, + **({"LITELLM_RUST": rust} if rust is not None else {}), + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "TIKTOKEN_CACHE_DIR": str(tmp_path / "unused-tokenizer-cache"), + }, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.strip() == "compatible" + assert not (tmp_path / "unused-tokenizer-cache").exists() + + +@pytest.mark.parametrize("is_pretokenized", (False, True)) +def test_huggingface_batch_sequence_containers_match_python(is_pretokenized: bool) -> None: + reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + tokenizer: Final = HuggingFaceTokenizer.from_str(TOKENIZER_JSON) + inputs: Final = [["Hello", "World"], ("Hello", "World")] + actual: Final = tokenizer.encode_batch(inputs, is_pretokenized=is_pretokenized) + expected: Final = reference.encode_batch(inputs, is_pretokenized=is_pretokenized) + assert [(item.ids, item.type_ids, item.sequence_ids) for item in actual] == [ + (item.ids, item.type_ids, item.sequence_ids) for item in expected + ] + + +@pytest.mark.parametrize("name", ("cl100k_base", "o200k_base", "p50k_edit", "gpt2")) +def test_openai_encoding_exposes_the_tiktoken_vocabulary_surface(name: str) -> None: + reference: Final = tiktoken.get_encoding(name) + encoding: Final = OpenAIEncoding.from_tiktoken(name) + text: Final = "hello fanta" + + assert repr(encoding) == repr(reference) == f"" + assert (encoding.name, encoding.n_vocab, encoding.max_token_value) == ( + reference.name, + reference.n_vocab, + reference.max_token_value, + ) + assert encoding.token_byte_values() == reference.token_byte_values() + assert encoding.encode_single_token("hello") == reference.encode_single_token("hello") + assert encoding.encode_single_token(b"<|endoftext|>") == reference.eot_token + assert [encoding.is_special_token(token) for token in (0, reference.eot_token)] == [False, True] + assert encoding.decode_with_offsets(reference.encode(text)) == reference.decode_with_offsets(reference.encode(text)) + assert encoding.encode_to_numpy(text).tolist() == reference.encode_to_numpy(text).tolist() + stable, completions = encoding.encode_with_unstable(text) + expected_stable, expected_completions = reference.encode_with_unstable(text) + assert (stable, sorted(completions)) == (expected_stable, sorted(expected_completions)) + with pytest.raises(KeyError): + encoding.encode_single_token("<|not-a-token|>") + + +def test_huggingface_tokenizer_exposes_the_tokenizers_vocabulary_surface() -> None: + reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + reference.enable_padding(pad_id=0, pad_token="[UNK]", length=4) + reference.enable_truncation(max_length=3, stride=1, strategy="only_first", direction="left") + tokenizer: Final = HuggingFaceTokenizer.from_str(reference.to_str()) + + assert tokenizer.token_to_id("Hello") == reference.token_to_id("Hello") == 1 + assert tokenizer.id_to_token(3) == reference.id_to_token(3) == "[BOS]" + assert tokenizer.id_to_token(99) is None + assert tokenizer.get_vocab() == reference.get_vocab() + assert tokenizer.get_vocab(with_added_tokens=False) == reference.get_vocab(with_added_tokens=False) + assert tokenizer.get_vocab_size() == reference.get_vocab_size() == 4 + assert tokenizer.get_vocab_size(with_added_tokens=False) == reference.get_vocab_size(with_added_tokens=False) + added: Final = tokenizer.get_added_tokens_decoder() + expected_added: Final = reference.get_added_tokens_decoder() + assert {token_id: str(token) for token_id, token in added.items()} == { + token_id: str(token) for token_id, token in expected_added.items() + } + assert added[3].special == expected_added[3].special + assert tokenizer.num_special_tokens_to_add(False) == reference.num_special_tokens_to_add(False) == 1 + assert tokenizer.num_special_tokens_to_add(True) == reference.num_special_tokens_to_add(True) == 0 + assert tokenizer.padding == reference.padding + assert tokenizer.truncation == reference.truncation + assert tokenizer.encode_special_tokens == reference.encode_special_tokens is False + assert HuggingFaceTokenizer.from_buffer(TOKENIZER_JSON.encode()).encode("Hello").ids == [3, 1] + assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).padding is None + assert HuggingFaceTokenizer.from_str(TOKENIZER_JSON).truncation is None + + +def test_huggingface_encoding_exposes_the_tokenizers_lookup_and_mutation_surface() -> None: + reference: Final = ReferenceTokenizer.from_str(claude_json_str) + tokenizer: Final = HuggingFaceTokenizer.from_str(claude_json_str) + text: Final = "hello wide world" + actual: Final = tokenizer.encode(text, "again") + expected: Final = reference.encode(text, "again") + + lookups: Final = ( + lambda encoding: [encoding.token_to_chars(index) for index in range(len(encoding))], + lambda encoding: [encoding.token_to_word(index) for index in range(len(encoding))], + lambda encoding: [encoding.token_to_sequence(index) for index in range(len(encoding))], + lambda encoding: [encoding.char_to_token(position) for position in range(len(text))], + lambda encoding: [encoding.char_to_word(position) for position in range(len(text))], + lambda encoding: [encoding.char_to_token(position, 1) for position in range(5)], + lambda encoding: [encoding.word_to_tokens(word) for word in range(3)], + lambda encoding: [encoding.word_to_chars(word) for word in range(3)], + lambda encoding: [encoding.word_to_tokens(0, 1), encoding.word_to_chars(0, 1)], + ) + for lookup in lookups: + assert lookup(actual) == lookup(expected) + assert repr(actual) == repr(expected) + + actual.truncate(4, stride=1, direction="left") + expected.truncate(4, stride=1, direction="left") + assert (actual.ids, [item.ids for item in actual.overflowing]) == ( + expected.ids, + [item.ids for item in expected.overflowing], + ) + actual.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="") + expected.pad(6, direction="left", pad_id=7, pad_type_id=1, pad_token="") + assert (actual.ids, actual.attention_mask, actual.type_ids, actual.tokens) == ( + expected.ids, + expected.attention_mask, + expected.type_ids, + expected.tokens, + ) + actual.set_sequence_id(3) + expected.set_sequence_id(3) + assert actual.sequence_ids == expected.sequence_ids + merged: Final = type(actual).merge([actual, tokenizer.encode("more")]) + assert merged.ids == type(expected).merge([expected, reference.encode("more")]).ids + assert merged.offsets == type(expected).merge([expected, reference.encode("more")]).offsets + with pytest.raises(ValueError, match="direction"): + actual.pad(8, direction="sideways") diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py index af7befecc33..895b3b57f7b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py @@ -6,6 +6,8 @@ regression they guard is the one a caller sees: a tier the proxy advertises has leaves the adapter, in the shape the target expects. """ +from typing import Final + import pytest from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( @@ -36,6 +38,126 @@ def _reasoning_effort_sent(model: str, provider: str, reasoning_effort: object) return completion_kwargs.get("reasoning_effort") +def _reasoning_effort_sent_for_thinking( + model: str, + provider: str | None, + thinking: dict[str, object], + *, + tools: list[dict[str, object]] | None = None, + api_base: str | None = None, +) -> object: + extra_kwargs: Final = { + key: value for key, value in (("custom_llm_provider", provider), ("api_base", api_base)) if value is not None + } + completion_kwargs, _ = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + max_tokens=1024, + messages=MESSAGES, + model=model, + metadata=None, + stop_sequences=None, + stream=False, + system=None, + temperature=None, + thinking=thinking, + tool_choice=None, + tools=tools, + top_k=None, + top_p=None, + output_format=None, + extra_kwargs=extra_kwargs, + ) + return completion_kwargs.get("reasoning_effort") + + +SUMMARIZED_THINKING = {"type": "enabled", "budget_tokens": 4096, "summary": "auto"} +PLAIN_THINKING = {"type": "enabled", "budget_tokens": 4096} +MULTIPLY_TOOL = { + "name": "multiply", + "description": "Multiply two integers", + "input_schema": {"type": "object", "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}}}, +} + + +class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: + """Only the Responses API takes ``reasoning_effort`` as a dict. Databricks answered the wrapped + ``{"effort", "summary"}`` with ``field 'reasoning_effort' expects input with json type 'string' + but got 'object'``, so a target that stays on chat completions has to get the plain tier and a + target the bridge picks up has to keep the summary it can honor.""" + + @pytest.mark.parametrize( + "model, provider", + [ + ("databricks/databricks-qwen35-122b-a10b", "databricks"), + ("databricks-qwen35-122b-a10b", "databricks"), + ("fireworks_ai/kimi-k3", "fireworks_ai"), + ], + ) + def test_a_chat_target_gets_the_plain_tier(self, local_model_cost_map: None, model: str, provider: str) -> None: + assert _reasoning_effort_sent_for_thinking(model, provider, SUMMARIZED_THINKING) == "high" + + def test_auto_summary_stays_a_plain_tier_on_a_chat_target( + self, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + + sent = _reasoning_effort_sent_for_thinking("databricks/databricks-qwen35-122b-a10b", "databricks", PLAIN_THINKING) + + assert sent == "high" + + @pytest.mark.parametrize( + "model, provider", + [ + ("azure/responses/gpt-5-mini", "azure"), + ("gpt-5-mini", "openai"), + ("databricks/databricks-gpt-5-5", "databricks"), + ], + ) + def test_a_bridged_target_keeps_the_summary(self, local_model_cost_map: None, model: str, provider: str) -> None: + sent = _reasoning_effort_sent_for_thinking(model, provider, SUMMARIZED_THINKING) + + assert sent == {"effort": "high", "summary": "auto"} + + def test_auto_summary_still_reaches_a_bridged_target( + self, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + + sent = _reasoning_effort_sent_for_thinking("azure/responses/gpt-5-mini", "azure", PLAIN_THINKING) + + assert sent == {"effort": "high", "summary": "detailed"} + + @pytest.mark.parametrize( + "api_base, expected", + [ + ("https://foo.services.ai.azure.com/openai/v1", "high"), + ("https://foo.eastus.models.ai.azure.com", {"effort": "high", "summary": "auto"}), + ], + ) + def test_a_foundry_deployment_is_judged_by_its_api_base( + self, local_model_cost_map: None, api_base: str, expected: object + ) -> None: + """``completion()`` keeps a gpt-5.5 deployment with function tools on Foundry's chat route when + its ``api_base`` is a Foundry OpenAI host, and bridges it to Responses when the base makes it an + Azure OpenAI deployment. The adapter has to read the same ``api_base`` to land on the same call.""" + sent = _reasoning_effort_sent_for_thinking( + "azure_ai/gpt-5.5", "azure_ai", SUMMARIZED_THINKING, tools=[MULTIPLY_TOOL], api_base=api_base + ) + + assert sent == expected + + def test_a_provider_resolved_from_the_api_base_gets_the_plain_tier(self, local_model_cost_map: None) -> None: + sent = _reasoning_effort_sent_for_thinking( + "kimi-k3", None, SUMMARIZED_THINKING, api_base="https://api.together.xyz/v1" + ) + + assert sent == "high" + + def test_a_chained_gateway_keeps_the_dict_for_its_own_bridge(self, local_model_cost_map: None) -> None: + sent = _reasoning_effort_sent_for_thinking("litellm_proxy/gpt-5.4", "litellm_proxy", SUMMARIZED_THINKING) + + assert sent == {"effort": "high", "summary": "auto"} + + class TestTheNormalizedTierIsTheTierSent: """The bug in the caller's terms: a proxy advertising kimi-k3 ``max`` accepted the request and then put ``high`` on the wire. Every spelling of the entry has to survive the adapter, including diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index fc5d807bc23..d835db63d83 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -16,6 +16,7 @@ import json from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import litellm @@ -1507,6 +1508,55 @@ async def test_summary_model_denied_when_team_member_scope_excludes_it(): assert result.applied_edits[0].get("error") == "summary_model_access_denied" +async def test_summary_model_denied_when_team_membership_read_hits_a_db_outage(): + """A member-level scope that cannot be read fails closed: the summary + model is not invoked while the membership row is unreachable.""" + messages = _simple_messages() + mock_call = AsyncMock(return_value=_make_mock_response("x")) + + auth = _fake_user_api_key_auth(key_models=["all-proxy-models"], team_id="team-outage") + auth.user_id = "user-outage" + + class _UnreachableMembershipPrisma: + class db: + class litellm_teammembership: + @staticmethod + async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None: + raise httpx.ConnectError("All connection attempts failed") + + with ( + patch( + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting", + return_value="claude-haiku-4-5", + ), + patch("litellm.token_counter", return_value=200_000), + patch( + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model", + mock_call, + ), + patch( + "litellm.proxy.auth.auth_checks.get_user_object", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.auth.auth_checks.get_project_object", + AsyncMock(return_value=None), + ), + patch("litellm.proxy.proxy_server.prisma_client", _UnreachableMembershipPrisma()), + ): + result = await apply_compact_20260112( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec=_EDIT_SPEC_DEFAULT, + user_api_key_auth=auth, + ) + + mock_call.assert_not_awaited() + assert result.applied_edits[0].get("error") == "summary_model_access_denied" + + async def test_summary_model_denied_when_key_over_model_budget(): """A caller whose per-model budget for the summary model is exhausted cannot trigger the summary call via compaction.""" diff --git a/tests/test_litellm/llms/azure/test_azure.py b/tests/test_litellm/llms/azure/test_azure.py index 6b6832f623c..86065c7adc6 100644 --- a/tests/test_litellm/llms/azure/test_azure.py +++ b/tests/test_litellm/llms/azure/test_azure.py @@ -1,10 +1,13 @@ """Tests for litellm/llms/azure/azure.py AzureChatCompletion handler behaviour.""" +import asyncio import time from typing import Final -from openai import AzureOpenAI +import pytest +from openai import AsyncAzureOpenAI, AzureOpenAI +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.azure import AzureChatCompletion @@ -52,3 +55,25 @@ def test_sync_streaming_stamps_response_headers_on_the_logging_obj() -> None: ) assert logging_obj.model_call_details["response_headers"] == {"x-ms-is-spilled-over": "true"} + + +class _CancelledRawCompletions: + async def create(self, **kwargs): + raise asyncio.CancelledError() + + +@pytest.mark.asyncio +async def test_acompletion_propagates_cancelled_error() -> None: + client = AsyncAzureOpenAI( + api_key="fake-key", + api_version="2024-02-01", + azure_endpoint="https://fake-resource.openai.azure.com", + ) + client.chat.completions.with_raw_response = _CancelledRawCompletions() + + with pytest.raises(asyncio.CancelledError): + await litellm.acompletion( + model="azure/fake-deployment", + messages=[{"role": "user", "content": "hi"}], + client=client, + ) diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 83ec85f1176..caf941ebd19 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -597,7 +597,8 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type): "litellm.files.main.azure_files_instance.initialize_azure_sdk_client" ) elif ( - call_type == CallTypes.avideo_content + call_type == CallTypes.avideo_generation + or call_type == CallTypes.avideo_content or call_type == CallTypes.avideo_list or call_type == CallTypes.avideo_remix or call_type == CallTypes.avideo_create_character diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index df042ce5902..117814a41ff 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -279,6 +279,16 @@ def test_context_window_suffix_stripped_for_cost_lookup(): ) +def test_legacy_mantle_route_prefix_stripped_for_cost_lookup(): + """The mantle/ route token is a routing prefix like openai/, so a bedrock/mantle/ + deployment must resolve the bare Bedrock model for cost lookup while still routing to Mantle.""" + from litellm.llms.bedrock.common_utils import get_bedrock_base_model, strip_bedrock_routing_prefix + + assert strip_bedrock_routing_prefix("mantle/anthropic.claude-sonnet-5") == "anthropic.claude-sonnet-5" + assert get_bedrock_base_model("bedrock/mantle/anthropic.claude-sonnet-5") == "anthropic.claude-sonnet-5" + assert BedrockModelInfo.get_bedrock_route("bedrock/mantle/anthropic.claude-sonnet-5") == "mantle" + + def test_output_config_effort_normalization_uses_model_info_ceiling(monkeypatch): import litellm.llms.bedrock.common_utils as mod @@ -926,3 +936,31 @@ def test_every_bedrock_config_get_error_class_keeps_provider_headers(config): def test_bedrock_get_error_class_audit_covers_every_surface(): assert len(_bedrock_configs_with_get_error_class()) >= 30 + + +def test_s3_static_key_pair_returns_the_pair_when_both_keys_are_set(): + from litellm.llms.bedrock.common_utils import s3_static_key_pair + + assert s3_static_key_pair( + { + "aws_access_key_id": "bedrock-key", + "aws_secret_access_key": "bedrock-secret", + "s3_access_key_id": "s3-key", + "s3_secret_access_key": "s3-secret", + } + ) == ("s3-key", "s3-secret") + + +@pytest.mark.parametrize( + "partial_s3_pair", + [ + {}, + {"s3_access_key_id": "s3-key"}, + {"s3_secret_access_key": "s3-secret"}, + {"s3_access_key_id": "", "s3_secret_access_key": ""}, + ], +) +def test_s3_static_key_pair_is_none_without_a_full_pair(partial_s3_pair): + from litellm.llms.bedrock.common_utils import s3_static_key_pair + + assert s3_static_key_pair({"aws_access_key_id": "bedrock-key", **partial_s3_pair}) is None diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index 09be2118001..37cf49a85ec 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -447,6 +447,63 @@ async def test_mantle_anthropic_messages_sends_workspace_header_and_clean_body() assert "aws_bedrock_project_id" not in requests[0]["body"] +async def _send_anthropic_messages_with_betas(**request_params: object) -> dict: + import litellm + + requests = [] + + async def mock_post(self, url, data=None, headers=None, **kwargs): + requests.append(_capture_request(url=url, headers=headers or {}, data=data)) + return _anthropic_response(url) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + **request_params, + ) + finally: + await litellm.close_litellm_async_clients() + + assert len(requests) == 1 + return requests[0] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_mantle_anthropic_messages_sends_every_beta_in_the_header_not_the_body(): + sent = await _send_anthropic_messages_with_betas( + extra_headers={"anthropic-beta": "context-1m-2025-08-07,interleaved-thinking-2025-05-14"}, + context_management={"edits": [{"type": "clear_tool_uses_20250919"}]}, + ) + + assert ( + sent["headers"]["anthropic-beta"] + == "context-1m-2025-08-07,context-management-2025-06-27,interleaved-thinking-2025-05-14" + ) + assert sent["headers"]["anthropic-version"] == "2023-06-01" + assert sent["body"]["context_management"] == {"edits": [{"type": "clear_tool_uses_20250919"}]} + assert "anthropic_beta" not in sent["body"] + assert "anthropic_version" not in sent["body"] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_mantle_anthropic_messages_drops_the_beta_header_when_mantle_rejects_every_value(): + sent = await _send_anthropic_messages_with_betas(extra_headers={"anthropic-beta": "code-execution-2025-08-25"}) + + assert "anthropic-beta" not in sent["headers"] + assert "anthropic_beta" not in sent["body"] + + def _usageless_anthropic_response(url: str) -> httpx.Response: return httpx.Response( status_code=200, diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index fa4b7439dd8..67a8d045036 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -4051,3 +4051,83 @@ async def test_responses_agentic_followup_does_not_repeat_request_params_from_pl assert followup_calls[0]["prompt_cache_key"] == "thread-1" assert followup_calls[0]["metadata"] == {"user": "u1"} assert followup_calls[0]["_agentic_loop_depth"] == 1 + + +@pytest.mark.asyncio +async def test_responses_agentic_followup_sends_the_plans_request_param_over_a_stale_kwargs_copy(monkeypatch): + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch + + followup_calls: list[dict[str, object]] = [] + + async def fake_aresponses(**kwargs: object) -> str: + followup_calls.append(kwargs) + return "followup-response" + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + + await BaseLLMHTTPHandler()._execute_responses_agentic_plan( + plan=AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + optional_params={"prompt_cache_key": "from-plan-params"}, + kwargs={"prompt_cache_key": "stale-copy"}, + ), + ), + model="gpt-5", + response_api_optional_request_params={"prompt_cache_key": "from-request"}, + logging_obj=Mock(litellm_call_id="call-1"), + kwargs={}, + depth=0, + max_loops=3, + fingerprints=[], + fingerprint="fp", + callback=CustomLogger(), + ) + + assert followup_calls[0]["prompt_cache_key"] == "from-plan-params" + + +@pytest.mark.asyncio +async def test_chat_completion_agentic_followup_does_not_repeat_request_params_from_plan_kwargs(monkeypatch): + """A plan whose kwargs repeat a request param, or the explicitly passed model, must not crash the chat follow-up with a duplicate keyword""" + from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch + + followup_calls: list[dict[str, object]] = [] + + async def fake_acompletion(**kwargs: object) -> str: + followup_calls.append(kwargs) + return "followup-response" + + monkeypatch.setattr(litellm, "acompletion", fake_acompletion) + request_kwargs: Final = {"temperature": 0.2, "api_base": "https://a", "model": "gpt-5"} + plan: Final = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + optional_params={"temperature": 0.2}, + kwargs=dict(request_kwargs), + ), + ) + + response: Final = await BaseLLMHTTPHandler()._execute_chat_completion_agentic_plan( + plan=plan, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + optional_params={"temperature": 0.2}, + kwargs=dict(request_kwargs), + custom_llm_provider="openai", + depth=0, + max_loops=3, + fingerprints=[], + fingerprint="fp", + ) + + assert response == "followup-response" + assert len(followup_calls) == 1 + assert followup_calls[0]["temperature"] == 0.2 + assert followup_calls[0]["api_base"] == "https://a" + assert followup_calls[0]["model"] == "openai/gpt-5" diff --git a/tests/test_litellm/llms/databricks/chat/__init__.py b/tests/test_litellm/llms/databricks/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py new file mode 100644 index 00000000000..a3391a2c585 --- /dev/null +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -0,0 +1,79 @@ +import json +from typing import Final + +import httpx +import respx + +import litellm + + +def test_completion_merges_leading_system_and_developer_messages_for_chat_template_models( + respx_mock: respx.MockRouter, +): + upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + response: Final = litellm.completion( + model="databricks/my-custom-model", + messages=[ + {"role": "system", "content": "You are terse."}, + {"role": "developer", "content": "Skills: none."}, + {"role": "user", "content": "Hello"}, + ], + api_base="https://example.databricks.test/serving-endpoints", + api_key="fake-databricks-api-key", + num_retries=0, + ) + + assert upstream.call_count == 1 + request_body: Final = json.loads(upstream.calls[0].request.read()) + assert request_body["messages"] == [ + {"role": "system", "content": "You are terse.\n\nSkills: none."}, + {"role": "user", "content": "Hello"}, + ] + assert response.choices[0].message.content == "Answer" + + +def test_completion_merges_system_messages_when_one_has_empty_content(respx_mock: respx.MockRouter): + upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + litellm.completion( + model="databricks/my-custom-model", + messages=[ + {"role": "system", "content": "You are terse."}, + {"role": "system", "content": ""}, + {"role": "user", "content": "Hello"}, + ], + api_base="https://example.databricks.test/serving-endpoints", + api_key="fake-databricks-api-key", + num_retries=0, + ) + + request_body: Final = json.loads(upstream.calls[0].request.read()) + assert request_body["messages"] == [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hello"}, + ] diff --git a/tests/test_litellm/llms/fal_ai/chat/test_fal_ai_chat_transformation.py b/tests/test_litellm/llms/fal_ai/chat/test_fal_ai_chat_transformation.py new file mode 100644 index 00000000000..41e8fc0c8c5 --- /dev/null +++ b/tests/test_litellm/llms/fal_ai/chat/test_fal_ai_chat_transformation.py @@ -0,0 +1,358 @@ +import httpx +import pytest + +import litellm +from litellm.llms.fal_ai.chat.transformation import FalAIChatConfig, FalAIError +from litellm.types.utils import LlmProviders, ModelResponse +from litellm.utils import ProviderConfigManager + +MODEL = "fal-ai/moondream3-preview/query" + + +@pytest.fixture(autouse=True) +def _use_local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +def _messages(*content): + return [ + { + "role": "user", + "content": [{"type": "text", "text": text} for text in content[:1]] + + [{"type": "image_url", "image_url": {"url": c}} for c in content[1:]], + } + ] + + +def test_provider_config_manager_resolves_fal_ai_chat_config(): + config = ProviderConfigManager.get_provider_chat_config(model=MODEL, provider=LlmProviders.FAL_AI) + assert isinstance(config, FalAIChatConfig) + + +def test_get_complete_url_targets_fal_endpoint(): + assert ( + FalAIChatConfig().get_complete_url( + api_base=None, api_key=None, model=MODEL, optional_params={}, litellm_params={} + ) + == "https://fal.run/fal-ai/moondream3-preview/query" + ) + + +def test_get_complete_url_strips_fal_ai_model_prefix(): + assert ( + FalAIChatConfig().get_complete_url( + api_base=None, api_key=None, model=f"fal_ai/{MODEL}", optional_params={}, litellm_params={} + ) + == "https://fal.run/fal-ai/moondream3-preview/query" + ) + + +def test_validate_environment_uses_fal_key_scheme(): + headers = FalAIChatConfig().validate_environment( + headers={}, model=MODEL, messages=[], optional_params={}, litellm_params={}, api_key="secret" + ) + assert headers["Authorization"] == "Key secret" + + +def test_transform_request_joins_text_parts_and_extracts_image_url(): + body = FalAIChatConfig().transform_request( + model=MODEL, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is"}, + {"type": "text", "text": "in this image?"}, + {"type": "image_url", "image_url": {"url": "https://example.com/pic.png"}}, + ], + } + ], + optional_params={"temperature": 0.2, "top_p": 0.9, "reasoning": False}, + litellm_params={}, + headers={}, + ) + assert body == { + "prompt": "what is\nin this image?", + "image_url": "https://example.com/pic.png", + "temperature": 0.2, + "top_p": 0.9, + "reasoning": False, + } + + +def test_transform_request_passes_data_url_through(): + body = FalAIChatConfig().transform_request( + model=MODEL, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + ], + } + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body["image_url"] == "data:image/png;base64,AAAA" + + +def test_transform_request_accepts_single_user_message(): + body = FalAIChatConfig().transform_request( + model=MODEL, + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "describe"}, {"type": "image_url", "image_url": "https://a"}], + } + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body["prompt"] == "describe" + assert body["image_url"] == "https://a" + + +def test_transform_request_rejects_system_message(): + with pytest.raises(FalAIError, match="exactly one user message") as exc_info: + FalAIChatConfig().transform_request( + model=MODEL, + messages=[ + {"role": "system", "content": "be terse"}, + { + "role": "user", + "content": [{"type": "text", "text": "describe"}, {"type": "image_url", "image_url": "https://a"}], + }, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert exc_info.value.status_code == 400 + + +def test_transform_request_rejects_multi_turn_history(): + with pytest.raises(FalAIError, match="exactly one user message") as exc_info: + FalAIChatConfig().transform_request( + model=MODEL, + messages=[ + { + "role": "user", + "content": [{"type": "text", "text": "first"}, {"type": "image_url", "image_url": "https://a"}], + }, + {"role": "assistant", "content": "an answer"}, + { + "role": "user", + "content": [{"type": "text", "text": "second"}, {"type": "image_url", "image_url": "https://b"}], + }, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert exc_info.value.status_code == 400 + + +def test_transform_request_rejects_zero_images(): + with pytest.raises(FalAIError, match="exactly one image_url"): + FalAIChatConfig().transform_request( + model=MODEL, + messages=[{"role": "user", "content": [{"type": "text", "text": "describe"}]}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + +def test_transform_request_rejects_two_images(): + with pytest.raises(FalAIError, match="exactly one image_url"): + FalAIChatConfig().transform_request( + model=MODEL, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "compare"}, + {"type": "image_url", "image_url": {"url": "https://a"}}, + {"type": "image_url", "image_url": {"url": "https://b"}}, + ], + } + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + +def test_transform_request_rejects_missing_text(): + with pytest.raises(FalAIError, match="require text"): + FalAIChatConfig().transform_request( + model=MODEL, + messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://a"}}]}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + +def test_transform_request_rejects_streaming(): + with pytest.raises(FalAIError, match="streaming"): + FalAIChatConfig().transform_request( + model=MODEL, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": "https://a"}}, + ], + } + ], + optional_params={"stream": True}, + litellm_params={}, + headers={}, + ) + + +def test_completion_dispatch_rejects_streaming(): + with pytest.raises(litellm.BadRequestError): + litellm.completion( + model=MODEL, + custom_llm_provider="fal_ai", + stream=True, + messages=[{"role": "user", "content": "describe"}], + ) + + +@pytest.mark.parametrize( + "effort,expected", + [("none", False), ("minimal", False), ("low", True), ("medium", True), ("high", True)], +) +def test_map_openai_params_maps_reasoning_effort(effort, expected): + mapped = FalAIChatConfig().map_openai_params( + non_default_params={"reasoning_effort": effort}, optional_params={}, model=MODEL, drop_params=False + ) + assert mapped["reasoning"] is expected + + +def test_map_openai_params_drops_unknown_reasoning_effort_when_dropping(): + mapped = FalAIChatConfig().map_openai_params( + non_default_params={"reasoning_effort": "extreme"}, optional_params={}, model=MODEL, drop_params=True + ) + assert "reasoning" not in mapped + + +def test_map_openai_params_maps_sampling_params(): + mapped = FalAIChatConfig().map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.7, "max_tokens": 10}, + optional_params={}, + model=MODEL, + drop_params=False, + ) + assert mapped == {"temperature": 0.5, "top_p": 0.7} + + +def test_transform_response_maps_output_reasoning_usage_and_finish_reason(): + raw = httpx.Response( + 200, + json={ + "output": "a red circle", + "reasoning": "looked at shapes", + "finish_reason": "stop", + "usage_info": { + "input_tokens": 11, + "output_tokens": 4, + "prefill_time_ms": 1.0, + "decode_time_ms": 2.0, + "ttft_ms": 1.5, + }, + }, + ) + response = FalAIChatConfig().transform_response( + model=MODEL, + raw_response=raw, + model_response=ModelResponse(), + logging_obj=None, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert response.choices[0].message.content == "a red circle" + assert response.choices[0].message.reasoning_content == "looked at shapes" + assert response.choices[0].finish_reason == "stop" + assert response.usage.prompt_tokens == 11 + assert response.usage.completion_tokens == 4 + assert response.usage.total_tokens == 15 + assert response.model == MODEL + + +def test_transform_response_omits_reasoning_when_null(): + raw = httpx.Response( + 200, + json={ + "output": "a red circle", + "reasoning": None, + "finish_reason": "stop", + "usage_info": {"input_tokens": 3, "output_tokens": 2}, + }, + ) + response = FalAIChatConfig().transform_response( + model=MODEL, + raw_response=raw, + model_response=ModelResponse(), + logging_obj=None, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert response.choices[0].message.content == "a red circle" + assert getattr(response.choices[0].message, "reasoning_content", None) is None + assert response.usage.total_tokens == 5 + + +def test_transform_response_rejects_body_missing_output(): + raw = httpx.Response( + 200, + json={"reasoning": "looked", "usage_info": {"input_tokens": 3, "output_tokens": 2}}, + ) + with pytest.raises(FalAIError) as exc_info: + FalAIChatConfig().transform_response( + model=MODEL, + raw_response=raw, + model_response=ModelResponse(), + logging_obj=None, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert exc_info.value.status_code == 422 + + +def test_transform_response_rejects_body_missing_usage_info(): + raw = httpx.Response(200, json={"output": "a red circle"}) + with pytest.raises(FalAIError) as exc_info: + FalAIChatConfig().transform_response( + model=MODEL, + raw_response=raw, + model_response=ModelResponse(), + logging_obj=None, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert exc_info.value.status_code == 422 diff --git a/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py b/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py new file mode 100644 index 00000000000..d99701db9e8 --- /dev/null +++ b/tests/test_litellm/llms/fal_ai/image_edit/test_fal_ai_flux_lora_depth_transformation.py @@ -0,0 +1,116 @@ +import base64 +import io + +import pytest + +import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils +from litellm.llms.fal_ai.image_edit import ( + FalAIFluxLoraDepthEditConfig, + FalAIImageEditConfig, + get_fal_ai_image_edit_config, +) +from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageObject, ImageResponse, LlmProviders +from litellm.utils import ProviderConfigManager + +PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"\x00" * 16 +MODEL = "fal-ai/flux-lora-depth" + + +@pytest.fixture(autouse=True) +def _use_local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", ["fal-ai/flux-lora-depth", "flux-lora-depth", "fal_ai/fal-ai/flux-lora-depth"]) +def test_dispatch_selects_flux_lora_depth_config(model): + assert isinstance(get_fal_ai_image_edit_config(model), FalAIFluxLoraDepthEditConfig) + + +def test_dispatch_keeps_gpt_image_config_for_openai_edit_models(): + config = get_fal_ai_image_edit_config("openai/gpt-image-2.5/flare/edit") + assert type(config) is FalAIImageEditConfig + + +def test_provider_config_manager_resolves_flux_lora_depth(): + config = ProviderConfigManager.get_provider_image_edit_config(model=MODEL, provider=LlmProviders.FAL_AI) + assert isinstance(config, FalAIFluxLoraDepthEditConfig) + + +@pytest.mark.parametrize("model", ["fal-ai/flux-lora-depth", "flux-lora-depth"]) +def test_get_complete_url_targets_endpoint_without_edit_suffix(model): + url = FalAIFluxLoraDepthEditConfig().get_complete_url(model=model, api_base=None, litellm_params={}) + assert url == "https://fal.run/fal-ai/flux-lora-depth" + + +def test_get_supported_openai_params_excludes_quality_mask_background(): + params = FalAIFluxLoraDepthEditConfig().get_supported_openai_params(model=MODEL) + assert "quality" not in params + assert "mask" not in params + assert "background" not in params + + +def test_map_openai_params_translates_n_and_size(): + mapped = FalAIFluxLoraDepthEditConfig().map_openai_params( + image_edit_optional_params=ImageEditOptionalRequestParams(n=2, size="1024x1536", quality="high"), + model=MODEL, + drop_params=False, + ) + assert mapped == {"num_images": 2, "image_size": {"width": 1024, "height": 1536}} + + +def test_transform_request_sends_single_image_url_as_data_url(): + body, files = FalAIFluxLoraDepthEditConfig().transform_image_edit_request( + model=MODEL, + prompt="follow the depth map", + image=io.BytesIO(PNG_BYTES), + image_edit_optional_request_params={"num_images": 1}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert files == () + assert body["prompt"] == "follow the depth map" + assert body["image_url"] == "data:image/png;base64," + base64.b64encode(PNG_BYTES).decode() + assert "image_urls" not in body + assert body["num_images"] == 1 + + +def test_transform_request_passes_remote_url_through_untouched(): + body, _ = FalAIFluxLoraDepthEditConfig().transform_image_edit_request( + model=MODEL, + prompt="follow the depth map", + image="https://example.com/depth.png", + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["image_url"] == "https://example.com/depth.png" + + +def test_transform_request_rejects_two_images(): + with pytest.raises(ValueError, match="exactly one control image"): + FalAIFluxLoraDepthEditConfig().transform_image_edit_request( + model=MODEL, + prompt="follow the depth map", + image=["https://example.com/a.png", "https://example.com/b.png"], + image_edit_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + +def test_image_edit_cost_uses_flat_output_cost_per_image(): + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model=MODEL, + completion_response=ImageResponse(data=[ImageObject(url="https://example.com/out.png")]), + custom_llm_provider="fal_ai", + optional_params={}, + call_type="aimage_edit", + ) + assert cost == litellm.model_cost[f"fal_ai/{MODEL}"]["output_cost_per_image"] > 0 diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py index 6fb34d9f88e..56dcba04b5c 100644 --- a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py +++ b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py @@ -1,11 +1,12 @@ +from typing import Final + import pytest import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils -from litellm.llms.fal_ai.cost_calculator import cost_calculator +from litellm.llms.fal_ai.cost_calculator import cost_calculator, fal_ai_passthrough_cost from litellm.types.utils import ImageObject, ImageResponse - @pytest.fixture(autouse=True) def _use_local_model_cost_map(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") @@ -78,14 +79,27 @@ def test_gpt_image_response_dimensions_override_request_size(): assert cost == expected -def test_gpt_image_response_dimensions_fall_back_to_request_size_when_unpriced(): +def test_gpt_image_response_dimensions_use_nearest_keyed_row_when_unpriced(): model = "fal_ai/openai/gpt-image-2.5/flare/text-to-image" cost = cost_calculator( model=model, image_response=_image_response_with_dimensions(((777, 888),)), optional_params={"quality": "low", "image_size": {"width": 1024, "height": 1536}}, ) - expected = litellm.model_cost[f"fal_ai/low/1024-x-1536/{model.removeprefix('fal_ai/')}"]["output_cost_per_image"] + expected = litellm.model_cost[f"fal_ai/low/1024-x-768/{model.removeprefix('fal_ai/')}"]["output_cost_per_image"] + assert cost == expected + + +def test_gpt_image_25_noncanonical_response_uses_nearest_keyed_row(): + model: Final = "fal_ai/openai/gpt-image-2.5/flare/text-to-image" + cost: Final = cost_calculator( + model=model, + image_response=_image_response_with_dimensions(((1536, 1024),)), + optional_params={"quality": "low", "image_size": {"width": 1536, "height": 1024}}, + ) + expected: Final = litellm.model_cost[ + "fal_ai/low/1024-x-1536/openai/gpt-image-2.5/flare/text-to-image" + ]["output_cost_per_image"] assert cost == expected @@ -160,3 +174,32 @@ def test_image_edit_call_type_routes_to_fal_keyed_pricing(): call_type="aimage_edit", ) assert cost == litellm.model_cost[f"fal_ai/medium/1024-x-1024/{model}"]["output_cost_per_image"] > 0 + + +def test_passthrough_trellis_charges_flat_rate(): + assert ( + fal_ai_passthrough_cost("fal-ai/trellis", {}) + == litellm.model_cost["fal_ai/fal-ai/trellis"]["output_cost_per_image"] + > 0 + ) + + +@pytest.mark.parametrize("resolution", [512, 1024, 1536]) +def test_passthrough_trellis_2_resolution_picks_keyed_tier(resolution): + assert ( + fal_ai_passthrough_cost("fal-ai/trellis-2", {"resolution": resolution}) + == litellm.model_cost["fal_ai/fal-ai/trellis-2"][f"output_cost_per_image_{resolution}"] + > 0 + ) + + +def test_passthrough_trellis_2_without_resolution_falls_back_to_default_rate(): + assert ( + fal_ai_passthrough_cost("fal-ai/trellis-2", {"image_url": "https://a"}) + == litellm.model_cost["fal_ai/fal-ai/trellis-2"]["output_cost_per_image"] + > 0 + ) + + +def test_passthrough_unknown_model_returns_none(): + assert fal_ai_passthrough_cost("fal-ai/no-such-model", {"resolution": 512}) is None diff --git a/tests/test_litellm/llms/xai/test_xai_key_fallback.py b/tests/test_litellm/llms/xai/test_xai_key_fallback.py index 092e4951547..cbc507c5ee3 100644 --- a/tests/test_litellm/llms/xai/test_xai_key_fallback.py +++ b/tests/test_litellm/llms/xai/test_xai_key_fallback.py @@ -12,6 +12,9 @@ from litellm.types.router import GenericLiteLLMParams class FakeLogging: + def __init__(self) -> None: + self.litellm_params: dict = {} + def update_from_kwargs(self, **kwargs): pass diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 087c5a03498..35e7055bbc0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -12,6 +12,7 @@ from starlette.datastructures import Headers from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, UnloadableEntitlementError, + _agent_capped_servers, _is_mcp_admitted_user_subject, ) from litellm.proxy._types import ( @@ -21,6 +22,7 @@ from litellm.proxy._types import ( SpecialMCPServerNames, UserAPIKeyAuth, ) +from litellm.types.agents import AgentCaller @pytest.mark.asyncio @@ -4169,10 +4171,114 @@ async def test_get_allowed_mcp_servers_for_key_prefers_in_memory_permission(): global_mcp_server_manager.registry.pop("direct-server", None) +@pytest.mark.parametrize( + ("agent_servers", "group_ceiling", "expected"), + [ + ([], frozenset({"server_1"}), ("server_1",)), + ([], frozenset({"server_1", "server_2", "server_3"}), ("server_1", "server_2")), + ([], frozenset(), ()), + (["server_2"], frozenset({"server_1", "server_2"}), ("server_2",)), + (["server_1"], frozenset({"server_2"}), ()), + (["server_1"], None, ("server_1",)), + ], +) +def test_agent_capped_servers_intersects_agent_config_and_access_groups(agent_servers, group_ceiling, expected): + """The agent's attached access groups cap the key/team servers alongside its own + object_permission; groups naming no server deny all.""" + assert _agent_capped_servers(["server_1", "server_2"], agent_servers, group_ceiling) == expected + + +def test_agent_capped_servers_without_agent_restrictions_is_uncapped(): + assert _agent_capped_servers(["server_1", "server_2"], [], None) is None + + @pytest.mark.asyncio class TestAgentMCPPermissions: """Test agent-level MCP server and tool permission intersection.""" + @staticmethod + def _agent_key_acting_for(user_id: str, team_id: str | None) -> UserAPIKeyAuth: + agent_key = UserAPIKeyAuth(api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1") + agent_key.agent_caller = AgentCaller(user_id=user_id, team_id=team_id) + return agent_key + + @staticmethod + def _team_servers(grants: dict[str, list[str]]) -> AsyncMock: + async def by_team(user_api_key_auth: UserAPIKeyAuth | None = None) -> list[str]: + assert user_api_key_auth is not None + return grants.get(user_api_key_auth.team_id or "", []) + + return AsyncMock(side_effect=by_team) + + @staticmethod + def _user_servers(grants: dict[str, list[str] | None]) -> AsyncMock: + async def by_user(user_api_key_auth: UserAPIKeyAuth | None = None) -> list[str] | None: + assert user_api_key_auth is not None + return grants.get(user_api_key_auth.user_id or "", []) + + return AsyncMock(side_effect=by_user) + + async def test_agent_key_acting_for_a_user_is_capped_at_the_invoking_teams_servers(self): + """LIT-8014: the agent's own key reaches server_1 and server_2, but the human who invoked it + belongs to a team granted only server_2, so on their behalf the agent reaches only server_2.""" + agent_key = self._agent_key_acting_for(user_id="alice", team_id="callers") + + with ( + patch.object( # test-quality-ok: the level resolvers read proxy_server globals with no injection seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server_1", "server_2"]) + ), + patch.object( # test-quality-ok: same seam, keyed by which team is being asked about + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + self._team_servers({"callers": ["server_2", "server_3"]}), + ), + patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here + MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) + ), + patch.object( # test-quality-ok: neither the agent's owner nor the caller has a personal grant + MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({}) + ), + ): + assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == ["server_2"] + + async def test_agent_key_acting_for_a_teamless_user_is_capped_at_that_users_servers(self): + agent_key = self._agent_key_acting_for(user_id="alice", team_id=None) + + with ( + patch.object( # test-quality-ok: the level resolvers read proxy_server globals with no injection seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server_1", "server_2"]) + ), + patch.object( # test-quality-ok: same seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", self._team_servers({}) + ), + patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here + MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) + ), + patch.object( # test-quality-ok: same seam, keyed by which user is being asked about + MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({"alice": ["server_1"]}) + ), + ): + assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == ["server_1"] + + async def test_agent_key_acting_for_a_caller_whose_entitlement_is_unreadable_reaches_nothing(self): + agent_key = self._agent_key_acting_for(user_id="alice", team_id=None) + + with ( + patch.object( # test-quality-ok: the level resolvers read proxy_server globals with no injection seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server_1"]) + ), + patch.object( # test-quality-ok: same seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", self._team_servers({}) + ), + patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here + MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) + ), + patch.object( # test-quality-ok: None is the resolver's own "entitlement unresolvable" signal + MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({"alice": None}) + ), + ): + assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == [] + async def test_get_allowed_mcp_servers_agent_intersection(self): """Key/team allow [server_1, server_2]; agent allows [server_1]. Result = [server_1].""" user_api_key_auth = UserAPIKeyAuth( @@ -4208,6 +4314,46 @@ class TestAgentMCPPermissions: assert sorted(result) == ["server_1", "server_2"] mock_agent.assert_called_once_with(user_api_key_auth) + async def test_agent_access_group_server_ceiling_expands_group_servers(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + asked: list[str] = [] + + async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None: + asked.append(agent_id) + return AgentAccessGroupCeiling( + access_group_ids=("ag-1",), + models=frozenset(), + mcp_server_ids=frozenset({"aliased-server"}), + agent_ids=frozenset(), + ) + + global_mcp_server_manager.registry["ag-server-id"] = MCPServer( + server_id="ag-server-id", + name="ag-server", + server_name="ag-server", + alias="aliased-server", + url="https://ag-server.example.com", + transport=MCPTransport.http, + ) + try: + result = await MCPRequestHandler._get_agent_access_group_server_ceiling( + UserAPIKeyAuth(api_key="test-key", agent_id="agent-ag"), resolve + ) + finally: + global_mcp_server_manager.registry.pop("ag-server-id", None) + + assert result == frozenset({"ag-server-id"}) + assert asked == ["agent-ag"] + assert ( + await MCPRequestHandler._get_agent_access_group_server_ceiling(UserAPIKeyAuth(api_key="k"), resolve) + is None + ) + assert asked == ["agent-ag"] + async def test_get_allowed_mcp_servers_key_team_agent_intersection(self): """Key allows [1, 2], agent allows [2, 3]. Result = [2].""" user_api_key_auth = UserAPIKeyAuth( @@ -5790,12 +5936,13 @@ class TestMCPDcrBridgeDelegateAdmission: @staticmethod def _wrapped_user_lookup_error(original: BaseException) -> ValueError: - """Reproduce get_user_object's real exception contract (litellm/proxy/auth/auth_checks.py): it - catches every DB failure in a broad ``except`` and re-raises a bare ``ValueError``, so the - original error (a missing-user Exception or a real outage) survives only as ``__context__``. - Injecting a raw ConnectionError/Exception instead would exercise a shape production never - produces and let a chain-blind outage classifier pass. That wrapping fidelity is itself pinned by - test_get_user_object_wraps_db_outage_as_valueerror_preserving_context in test_auth_checks.""" + """Reproduce get_user_object's exception contract (litellm/proxy/auth/auth_checks.py): a read + failure that is not a database outage is re-raised as a bare ``ValueError`` with the original + error only as ``__context__``, while an outage propagates raw (pinned by + test_get_user_object_surfaces_a_db_outage_as_503_not_as_a_missing_user and + test_get_user_object_still_reports_a_non_outage_read_failure_as_a_missing_user in + test_auth_checks). The wrapped shape is the harder one for the outage classifier, so injecting + it here keeps a chain-blind classifier from passing.""" try: raise original except BaseException: diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py new file mode 100644 index 00000000000..e744e84d671 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py @@ -0,0 +1,146 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, + resolve_agent_access_group_ceiling, +) +from litellm.types.agents import AgentResponse + +_CARD: Final = {"name": "agent", "url": "http://localhost:9999", "version": "1.0.0"} + + +def _agent(access_group_ids: list[str] | None) -> AgentResponse: + return AgentResponse( + agent_id="agent-1", agent_name="agent", agent_card_params=_CARD, access_group_ids=access_group_ids + ) + + +def _group( + group_id: str, + models: tuple[str, ...] = (), + mcp_servers: tuple[str, ...] = (), + agents: tuple[str, ...] = (), +) -> LiteLLM_AccessGroupTable: + return LiteLLM_AccessGroupTable( + access_group_id=group_id, + access_group_name=group_id, + access_model_names=list(models), + access_mcp_server_ids=list(mcp_servers), + access_agent_ids=list(agents), + ) + + +def _loaders(agent: AgentResponse | None, groups: dict[str, LiteLLM_AccessGroupTable]): + async def load_agent(agent_id: str) -> tuple[str, ...]: + return tuple(agent.access_group_ids or ()) if agent is not None else () + + async def load_group(group_id: str) -> LiteLLM_AccessGroupTable | None: + return groups.get(group_id) + + return load_agent, load_group + + +@pytest.mark.asyncio +@pytest.mark.parametrize("access_group_ids", [None, []]) +async def test_agent_without_access_groups_has_no_ceiling(access_group_ids: list[str] | None): + load_agent, load_group = _loaders(_agent(access_group_ids), {"g1": _group("g1", models=("gpt-5",))}) + + assert await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) is None + + +@pytest.mark.asyncio +async def test_unknown_agent_has_no_ceiling(): + load_agent, load_group = _loaders(None, {}) + + assert await resolve_agent_access_group_ceiling("missing", load_agent, load_group) is None + + +@pytest.mark.asyncio +async def test_ceiling_is_the_union_of_every_attached_group(): + load_agent, load_group = _loaders( + _agent(["g1", "g2"]), + { + "g1": _group("g1", models=("gpt-5",), mcp_servers=("mcp-a",), agents=("agent-b",)), + "g2": _group("g2", models=("claude-sonnet",), mcp_servers=("mcp-b",), agents=("agent-c",)), + }, + ) + + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) + + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1", "g2"), + models=frozenset({"gpt-5", "claude-sonnet"}), + mcp_server_ids=frozenset({"mcp-a", "mcp-b"}), + agent_ids=frozenset({"agent-b", "agent-c"}), + ) + + +@pytest.mark.asyncio +async def test_unloadable_group_contributes_nothing_but_the_ceiling_still_applies(): + load_agent, load_group = _loaders(_agent(["g1", "gone"]), {"g1": _group("g1", models=("gpt-5",))}) + + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) + + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1", "gone"), + models=frozenset({"gpt-5"}), + mcp_server_ids=frozenset(), + agent_ids=frozenset(), + ) + + +@pytest.mark.asyncio +async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted(): + load_agent, load_group = _loaders(_agent(["gone"]), {}) + + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) + + assert ceiling is not None + assert ceiling.models == frozenset() + assert ceiling.mcp_server_ids == frozenset() + assert ceiling.agent_ids == frozenset() + + +@pytest.mark.asyncio +async def test_default_agent_loader_reads_the_attached_groups_from_the_registry(): + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + _, load_group = _loaders(None, {"g1": _group("g1", models=("gpt-5",))}) + global_agent_registry.register_agent(_agent(["g1"])) + try: + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_access_group=load_group) + finally: + global_agent_registry.deregister_agent("agent") + + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1",), models=frozenset({"gpt-5"}), mcp_server_ids=frozenset(), agent_ids=frozenset() + ) + + +@pytest.mark.asyncio +async def test_default_loader_treats_a_missing_group_as_unreadable(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group + from litellm.proxy.auth import auth_checks + + async def missing_group(**_: object) -> LiteLLM_AccessGroupTable: + raise HTTPException(status_code=404, detail={"error": "Access group doesn't exist in db."}) + + monkeypatch.setattr(proxy_server, "prisma_client", object()) + monkeypatch.setattr(auth_checks, "get_access_object", missing_group) + + assert await _load_access_group("gone") is None + + +@pytest.mark.asyncio +async def test_default_loader_returns_nothing_without_a_db(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group + + monkeypatch.setattr(proxy_server, "prisma_client", None) + + assert await _load_access_group("ag-1") is None diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py new file mode 100644 index 00000000000..b08964503c8 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py @@ -0,0 +1,57 @@ +from typing import Final + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth, agent_caller_from_headers +from litellm.types.agents import AgentCaller + +_AGENT_KEY: Final = UserAPIKeyAuth(api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1") + + +def test_agent_key_echoing_both_ids_acts_for_that_user_and_team() -> None: + headers: Final = {"X-LiteLLM-User-Id": " alice ", "x-litellm-team-id": "callers"} + + assert agent_caller_from_headers(headers, _AGENT_KEY) == AgentCaller(user_id="alice", team_id="callers") + + +def test_agent_key_echoing_only_a_user_id_acts_for_a_teamless_user() -> None: + assert agent_caller_from_headers({"x-litellm-user-id": "alice"}, _AGENT_KEY) == AgentCaller(user_id="alice") + + +@pytest.mark.parametrize("headers", [{}, {"x-litellm-user-id": " ", "x-litellm-team-id": ""}]) +def test_agent_key_echoing_no_caller_acts_for_itself(headers: dict[str, str]) -> None: + assert agent_caller_from_headers(headers, _AGENT_KEY) is None + + +def test_caller_headers_on_a_key_without_an_agent_are_ignored() -> None: + plain_key: Final = UserAPIKeyAuth(api_key="plain-key", user_id="bob") + + assert agent_caller_from_headers({"x-litellm-user-id": "alice", "x-litellm-team-id": "callers"}, plain_key) is None + + +def test_caller_auth_stands_for_the_invoking_user_not_the_agent() -> None: + agent_key: Final = UserAPIKeyAuth( + api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1" + ) + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + + caller_auth: Final = agent_caller_auth(agent_key) + + assert caller_auth is not None + assert (caller_auth.user_id, caller_auth.team_id, caller_auth.agent_id, caller_auth.api_key) == ( + "alice", + "callers", + None, + None, + ) + assert agent_caller_auth(_AGENT_KEY) is None + + +def test_agent_caller_cannot_be_set_from_a_request_payload() -> None: + forged: Final = UserAPIKeyAuth.model_validate( + {"api_key": "agent-key", "agent_id": "agent-1", "agent_caller": {"user_id": "alice", "team_id": "callers"}} + ) + + assert forged.agent_caller is None + assert "agent_caller" not in forged.model_dump() diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 383b72e5c58..a87716375e8 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -9,10 +9,10 @@ from unittest.mock import AsyncMock, patch import pytest - from litellm.constants import UI_SESSION_TOKEN_TEAM_ID -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentAccess, AgentRequestHandler, @@ -20,6 +20,7 @@ from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( UnrestrictedAgentAccess, accessible_agents, ) +from litellm.types.agents import AgentCaller def _registry_with(*agent_names: str) -> AgentRegistry: @@ -157,6 +158,130 @@ class TestAgentRequestHandler: is False ), agent_id + @staticmethod + def _ceiling_resolver(agent_ids: frozenset[str] | None) -> tuple[CeilingResolver, list[str]]: + """A resolver that records the agent ids it was asked about and answers with a fixed + ceiling, or None when the agent has no access groups attached.""" + asked: Final[list[str]] = [] + + async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None: + asked.append(agent_id) + if agent_ids is None: + return None + return AgentAccessGroupCeiling( + access_group_ids=("ag-1",), models=frozenset(), mcp_server_ids=frozenset(), agent_ids=agent_ids + ) + + return resolve, asked + + @staticmethod + def _key_granting(agent_ids: list[str], agent_id: str | None) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id=agent_id, + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="obj-1", agents=agent_ids), + ) + + async def test_agent_access_groups_cap_an_otherwise_unrestricted_key(self): + """A key with no agent grant of its own may still only reach the agents its + agent's attached access groups name.""" + agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent") + resolve, asked = self._ceiling_resolver(frozenset({"agent-beta"})) + + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-beta"}) + ) + assert await AgentRequestHandler.is_agent_allowed("agent-beta", agent_key, resolve) is True + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key, resolve) is False + assert asked == ["caller-agent"] * 3 + + @staticmethod + def _team_grants(grants: dict[str, AgentAccess]) -> AsyncMock: + async def by_team(user_api_key_auth: UserAPIKeyAuth | None = None) -> AgentAccess: + assert user_api_key_auth is not None + return grants.get(user_api_key_auth.team_id or "", UnrestrictedAgentAccess()) + + return AsyncMock(side_effect=by_team) + + async def test_agent_key_acting_for_a_user_is_capped_at_the_invoking_teams_agents(self): + """LIT-8014: the agent's key and access groups reach alpha and beta, but the human who + invoked it belongs to a team granted only beta, so on their behalf the agent reaches only beta.""" + agent_key: Final = self._key_granting(["agent-alpha", "agent-beta"], agent_id="caller-agent") + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + resolve, _ = self._ceiling_resolver(frozenset({"agent-alpha", "agent-beta", "agent-gamma"})) + + with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam + AgentRequestHandler, + "_get_allowed_agents_for_team", + self._team_grants({"callers": RestrictedAgentAccess(frozenset({"agent-beta", "agent-gamma"}))}), + ) as mock_team: + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-beta"}) + ) + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key, resolve) is False + + assert {call.args[0].team_id for call in mock_team.call_args_list} == {None, "callers"} + + async def test_agent_key_acting_for_a_user_whose_team_grants_no_agent_reaches_none(self): + agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent") + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + resolve, _ = self._ceiling_resolver(None) + + with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam + AgentRequestHandler, + "_get_allowed_agents_for_team", + self._team_grants({"callers": RestrictedAgentAccess(frozenset())}), + ): + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset() + ) + + async def test_agent_key_acting_for_an_ungranted_caller_keeps_its_own_agents(self): + agent_key: Final = self._key_granting(["agent-alpha"], agent_id="caller-agent") + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + resolve, _ = self._ceiling_resolver(None) + + with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam + AgentRequestHandler, "_get_allowed_agents_for_team", self._team_grants({}) + ): + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-alpha"}) + ) + + + async def test_agent_access_groups_intersect_with_key_grants(self): + agent_key: Final = self._key_granting(["agent-alpha", "agent-beta"], agent_id="caller-agent") + resolve, _ = self._ceiling_resolver(frozenset({"agent-beta", "agent-gamma"})) + + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-beta"}) + ) + assert await AgentRequestHandler.is_agent_allowed("agent-gamma", agent_key, resolve) is False + + async def test_agent_access_groups_naming_no_agent_deny_every_agent(self): + agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent") + resolve, _ = self._ceiling_resolver(frozenset()) + + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess(frozenset()) + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key, resolve) is False + + async def test_agent_without_access_groups_keeps_key_grants(self): + agent_key: Final = self._key_granting(["agent-alpha"], agent_id="caller-agent") + resolve, asked = self._ceiling_resolver(None) + + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-alpha"}) + ) + assert asked == ["caller-agent"] + + async def test_key_without_agent_never_consults_agent_access_groups(self): + plain_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + resolve, asked = self._ceiling_resolver(frozenset()) + + assert await AgentRequestHandler.resolve_agent_access(plain_key, resolve) == UnrestrictedAgentAccess() + assert asked == [] + async def test_empty_access_group_denies_every_agent(self): """LIT-5143: a key restricted to an access group that resolves to no agents is restricted to nothing, not unrestricted. A failed group lookup still fails open.""" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 441e9640ef9..b9a260f5b14 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -16,6 +16,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.agents import AgentCaller AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]] @@ -511,6 +512,24 @@ async def test_message_methods_forward_caller_identity_headers(method: str): assert forwarded_headers.get("X-LiteLLM-Team-Id") == "team-xyz" +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["message/send", "message/stream"]) +async def test_agent_calling_another_agent_forwards_the_human_who_invoked_it(method: str): + """LIT-8014: an agent acting for alice calls a second agent through the proxy. That hop must + carry alice, not the first agent's owner, so the chain stays capped at what alice may reach.""" + mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS) + agent_key = UserAPIKeyAuth(api_key="sk-agent", user_id="agent-owner", team_id="agent-team", agent_id="agent-1") + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + + captured = await _invoke_message_method(method, mock_request, agent_key) + + forwarded_headers = captured.agent_extra_headers or {} + assert (forwarded_headers.get("X-LiteLLM-User-Id"), forwarded_headers.get("X-LiteLLM-Team-Id")) == ( + "alice", + "callers", + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["message/send", "message/stream"]) async def test_message_methods_send_the_entra_bearer_for_azure_agents(method: str): diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index 231626c7eb5..b036e0dac4d 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -15,6 +15,7 @@ from litellm.proxy.agent_endpoints.agent_registry import ( _restore_redacted_litellm_params, redact_sensitive_agent_litellm_params, ) +from litellm.types.agents import PatchAgentRequest # Obviously-fake stand-ins for a real AWS credential pair (LIT-6736 regression # fixtures) -- never a real key shape, and must never appear in any response. @@ -990,3 +991,138 @@ async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted(): stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY assert stored_params["is_public"] is True + + +def _agent_row_mock(access_group_ids: list[str]) -> MagicMock: + row: Final = MagicMock() + row.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + "access_group_ids": access_group_ids, + } + row.object_permission = None + return row + + +@pytest.mark.asyncio +async def test_add_agent_to_db_persists_deduplicated_access_group_ids(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_create = AsyncMock(return_value=_agent_row_mock(["ag-1", "ag-2"])) + mock_prisma.db.litellm_agentstable.create = mock_create + + result: Final = await registry.add_agent_to_db( + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "access_group_ids": ["ag-1", "ag-2", "ag-1"], + }, + prisma_client=mock_prisma, + created_by="test-user", + ) + + assert tuple(mock_create.call_args.kwargs["data"]["access_group_ids"]) == ("ag-1", "ag-2") + assert result.access_group_ids == ["ag-1", "ag-2"] + + +@pytest.mark.asyncio +async def test_add_agent_to_db_without_access_group_ids_leaves_column_to_its_default(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_create = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.create = mock_create + + await registry.add_agent_to_db( + agent={"agent_name": "Test Agent", "agent_card_params": _sample_agent_card_params()}, + prisma_client=mock_prisma, + created_by="test-user", + ) + + assert "access_group_ids" not in mock_create.call_args.kwargs["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("patch_body", "expected"), + [ + ({"access_group_ids": ["ag-2", "ag-3"]}, ["ag-2", "ag-3"]), + ({"access_group_ids": []}, []), + ({"access_group_ids": None}, []), + ], +) +async def test_patch_agent_in_db_replaces_access_group_ids_when_provided( + patch_body: PatchAgentRequest, expected: list[str] +): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Test Agent", + "litellm_params": {}, + "object_permission_id": None, + "access_group_ids": ["ag-1"], + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock(expected)) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", agent=patch_body, prisma_client=mock_prisma, updated_by="test-user" + ) + + assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected) + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_keeps_access_group_ids_when_omitted(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Old Name", + "litellm_params": {}, + "object_permission_id": None, + "access_group_ids": ["ag-1"], + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock(["ag-1"])) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", agent={"agent_name": "New Name"}, prisma_client=mock_prisma, updated_by="test-user" + ) + + assert "access_group_ids" not in mock_update.call_args.kwargs["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body_access_group_ids", "expected"), + [(["ag-9", "ag-9"], ["ag-9"]), (None, []), ("omitted", [])], +) +async def test_update_agent_in_db_always_writes_access_group_ids(body_access_group_ids, expected: list[str]): + """PUT is a full replacement: omitting the field clears any previously attached groups.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, access_group_ids=["ag-1"]) + ) + mock_update = AsyncMock(return_value=_agent_row_mock(expected)) + mock_prisma.db.litellm_agentstable.update = mock_update + body: Final = { + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"model": "bedrock/agentcore/my-agent"}, + **({} if body_access_group_ids == "omitted" else {"access_group_ids": body_access_group_ids}), + } + + await registry.update_agent_in_db( + agent_id="agent-123", agent=body, prisma_client=mock_prisma, updated_by="test-user" + ) + + assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index 9a9ccd9a213..801f61aa498 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -198,6 +198,62 @@ class TestProxyExceptionAnthropicEnvelope: assert fallback.status_code == 500 assert json.loads(fallback.body)["error"]["type"] == "api_error" + @staticmethod + def _call_id_error_response(general_settings, provider_specific_fields=None): + import litellm.proxy.anthropic_endpoints.endpoints as ep + from litellm.proxy._types import ProxyException + + request = MagicMock() + request.headers = {} + exc = ProxyException( + message="Rate limit exceeded", + type="rate_limit_error", + param=None, + code=429, + headers={"x-litellm-call-id": "call-8302"}, + provider_specific_fields=provider_specific_fields, + ) + with patch("litellm.proxy.proxy_server.general_settings", general_settings): + return ep._anthropic_error_json_response(exc, request) + + def test_anthropic_error_copies_the_call_id_into_the_error_when_opted_in(self): + """With include_call_id_in_error_body on, error.litellm_call_id is byte-identical to + the x-litellm-call-id header and lives inside the error object, which is what the + Anthropic SDK keeps as e.body.""" + response = self._call_id_error_response({"include_call_id_in_error_body": True}) + + assert response.headers["x-litellm-call-id"] == "call-8302" + assert json.loads(response.body) == { + "type": "error", + "error": { + "type": "rate_limit_error", + "message": "Rate limit exceeded", + "litellm_call_id": "call-8302", + }, + } + + def test_anthropic_error_keeps_provider_specific_fields_next_to_the_call_id(self): + response = self._call_id_error_response( + {"include_call_id_in_error_body": True}, + provider_specific_fields={"guardrail": "keyword-block"}, + ) + + assert json.loads(response.body)["error"] == { + "type": "rate_limit_error", + "message": "Rate limit exceeded", + "provider_specific_fields": {"guardrail": "keyword-block"}, + "litellm_call_id": "call-8302", + } + + def test_anthropic_error_leaves_the_envelope_alone_when_opted_out(self): + response = self._call_id_error_response({}) + + assert response.headers["x-litellm-call-id"] == "call-8302" + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "rate_limit_error", "message": "Rate limit exceeded"}, + } + class TestHttpExceptionDictDetail: @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 44c1d6a3c6b..14b739e60ba 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -34,20 +34,26 @@ from litellm.proxy._types import ( UserAPIKeyAuth, WebhookEvent, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver +from litellm.types.agents import AgentCaller from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _cache_management_object, _can_object_call_model, _can_object_call_vector_stores, + _check_agent_access_group_model_access, _check_end_user_budget, _check_team_member_budget, _fetch_key_object_from_db_with_reconnect, _get_fuzzy_user_object, + CallerTeamLoader, + CallerUserLoader, _get_team_db_check, _log_budget_lookup_failure, _tag_max_budget_check, _team_max_budget_check, _virtual_key_max_budget_alert_check, + _check_agent_caller_model_access, _virtual_key_max_budget_check, _virtual_key_soft_budget_check, get_key_object, @@ -69,6 +75,8 @@ from litellm.constants import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from prisma.errors import DataError from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, TAG_REGISTRY_OVERFLOW_SENTINEL, @@ -886,36 +894,80 @@ async def test_get_user_object_upsert_sets_budget_reset_at(monkeypatch, has_budg assert "budget_reset_at" not in creation_args -@pytest.mark.asyncio -async def test_get_user_object_wraps_db_outage_as_valueerror_preserving_context(): - """Pin get_user_object's exception contract: it catches every DB failure in a broad except and - re-raises a bare ValueError, so a real outage survives only as __context__ rather than as the - exception type. The MCP dcr_bridge admission and refresh paths depend on this to tell a transient - outage (retry, 503) from a missing user (fail closed), which is why they classify across the cause - chain instead of the top exception's type. If this wrapping ever changes, that classification must - change with it, so this test guards the contract the callers rely on.""" - from unittest.mock import AsyncMock, MagicMock, patch +def _user_read_raising(error: Exception) -> tuple[MagicMock, MagicMock]: + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=error) + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + cache.async_set_cache = AsyncMock() + return prisma_client, cache - mock_prisma_client = MagicMock() - mock_prisma_client.db = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=ConnectionError("can't reach database server") - ) - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_set_cache = AsyncMock() + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "outage", + [ + httpx.ConnectError("All connection attempts failed"), + httpx.ReadTimeout("timed out"), + DataError( + data={ + "user_facing_error": { + "message": "Can't reach database server at `127.0.0.1:41071`", + "error_code": "P1001", + } + } + ), + ], + ids=["connect_error", "read_timeout", "p1001_as_data_error"], +) +async def test_get_user_object_surfaces_a_db_outage_as_503_not_as_a_missing_user(outage): + from litellm.proxy.auth.auth_exception_handler import _as_proxy_exception + + prisma_client, cache = _user_read_raising(outage) with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): - with pytest.raises(ValueError, match="User doesn't exist in db\\.") as exc_info: + with pytest.raises(type(outage)) as raised: await get_user_object( - user_id="outage-contract-probe-user", - prisma_client=mock_prisma_client, - user_api_key_cache=mock_cache, + user_id="outage-probe-user", + prisma_client=prisma_client, + user_api_key_cache=cache, user_id_upsert=False, proxy_logging_obj=None, ) - assert isinstance(exc_info.value.__context__, ConnectionError) + assert raised.value is outage + assert PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(raised.value) is outage + surfaced = _as_proxy_exception(raised.value) + assert (surfaced.code, surfaced.type) == ("503", ProxyErrorTypes.no_db_connection) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure", + [ + DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}}), + RuntimeError("row validation failed"), + ], + ids=["query_level_data_error", "runtime_error"], +) +async def test_get_user_object_still_reports_a_non_outage_read_failure_as_a_missing_user(failure): + from litellm.proxy.auth.auth_exception_handler import _as_proxy_exception + + prisma_client, cache = _user_read_raising(failure) + + with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): + with pytest.raises(ValueError, match="User doesn't exist in db\\.") as raised: + await get_user_object( + user_id="data-error-probe-user", + prisma_client=prisma_client, + user_api_key_cache=cache, + user_id_upsert=False, + proxy_logging_obj=None, + ) + + assert raised.value.__context__ is failure + surfaced = _as_proxy_exception(raised.value) + assert (surfaced.code, surfaced.type) == ("401", ProxyErrorTypes.auth_error) @pytest.mark.asyncio @@ -8930,6 +8982,69 @@ def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False +def _agent_model_ceiling_resolver( + models: frozenset[str] | None, +) -> tuple[CeilingResolver, list[str]]: + """Resolver that records the agent ids it was asked about and answers with a fixed model + ceiling, or None when the agent has no access groups attached.""" + asked: Final[list[str]] = [] + + async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None: + asked.append(agent_id) + if models is None: + return None + return AgentAccessGroupCeiling( + access_group_ids=("ag-1",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset() + ) + + return resolve, asked + + +@pytest.mark.asyncio +async def test_agent_access_groups_cap_models_even_when_key_allows_them(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + resolve, asked = _agent_model_ceiling_resolver(frozenset({"gpt-5"})) + + assert await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve) is True + + with pytest.raises(ProxyException) as exc_info: + await _check_agent_access_group_model_access("claude-sonnet", agent_key, None, resolve) + + assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied + assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN) + assert asked == ["agent-1", "agent-1"] + + +@pytest.mark.asyncio +async def test_agent_access_groups_naming_no_model_deny_every_model(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=[]) + resolve, _ = _agent_model_ceiling_resolver(frozenset()) + + with pytest.raises(ProxyException) as exc_info: + await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve) + + assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied + + +@pytest.mark.asyncio +async def test_agent_without_access_groups_adds_no_model_ceiling(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + resolve, asked = _agent_model_ceiling_resolver(None) + + assert await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve) is True + assert await _check_agent_access_group_model_access("claude-sonnet", agent_key, None, resolve) is True + assert asked == ["agent-1", "agent-1"] + + +@pytest.mark.asyncio +async def test_key_without_agent_never_consults_agent_access_groups(): + plain_key: Final = UserAPIKeyAuth(token="plain-token", models=["gpt-5"]) + resolve, asked = _agent_model_ceiling_resolver(frozenset()) + + assert await _check_agent_access_group_model_access("gpt-5", plain_key, None, resolve) is True + assert asked == [] + + @pytest.mark.asyncio async def test_team_member_budget_check_temp_budget_increase_extends_cap(): """Spend above max_budget but below max_budget + active temp increase @@ -9086,3 +9201,117 @@ async def test_team_member_budget_check_adds_temp_increase_to_live_team_default( proxy_logging_obj=ProxyLogging(user_api_key_cache=None), ) assert exc_info.value.max_budget == expected_cap + + +def _agent_key_acting_for(user_id: str | None, team_id: str | None) -> UserAPIKeyAuth: + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + agent_key.agent_caller = AgentCaller(user_id=user_id, team_id=team_id) + return agent_key + + +def _caller_loaders( + team: LiteLLM_TeamTable | None, + user: LiteLLM_UserTable | None, +) -> tuple[CallerTeamLoader, CallerUserLoader, list[str]]: + """Loaders that hand back fixed caller rows and record the agent_caller they were asked about.""" + asked: Final[list[str]] = [] + + async def load_team(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + asked.append(f"team:{valid_token.agent_caller.team_id if valid_token.agent_caller else None}") + return team + + async def load_user(valid_token: UserAPIKeyAuth) -> LiteLLM_UserTable | None: + asked.append(f"user:{valid_token.agent_caller.user_id if valid_token.agent_caller else None}") + return user + + return load_team, load_user, asked + + +async def _cache_with_membership(user_id: str, team_id: str, allowed_models: list[str] | None) -> UserApiKeyCache: + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + cache: Final = UserApiKeyCache() + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=allowed_models) if allowed_models else None, + ), + model_type=LiteLLM_TeamMembership, + ) + return cache + + +async def _check_caller_models( + agent_key: UserAPIKeyAuth, + model: str, + load_team: CallerTeamLoader, + load_user: CallerUserLoader, + cache: UserApiKeyCache | None = None, +) -> None: + await _check_agent_caller_model_access( + model=model, + valid_token=agent_key, + llm_router=None, + prisma_client=None, + user_api_key_cache=cache or UserApiKeyCache(), + proxy_logging_obj=MagicMock(), + load_team=load_team, + load_user=load_user, + ) + + +@pytest.mark.asyncio +async def test_agent_key_acting_for_a_team_is_capped_at_that_teams_models(): + """LIT-8014: the invoking team may only call gpt-5, so the agent's own claude grant does not help.""" + agent_key: Final = _agent_key_acting_for(user_id="alice", team_id="team-a") + load_team, load_user, asked = _caller_loaders(LiteLLM_TeamTable(team_id="team-a", models=["gpt-5"]), None) + cache: Final = await _cache_with_membership("alice", "team-a", allowed_models=None) + + await _check_caller_models(agent_key, "gpt-5", load_team, load_user, cache) + with pytest.raises(ProxyException) as exc_info: + await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user, cache) + + assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied + assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN) + assert asked == ["team:team-a", "team:team-a"] + + +@pytest.mark.asyncio +async def test_agent_key_acting_for_a_team_member_is_capped_at_the_members_scope(): + agent_key: Final = _agent_key_acting_for(user_id="alice", team_id="team-a") + load_team, load_user, _ = _caller_loaders( + LiteLLM_TeamTable(team_id="team-a", models=["gpt-5", "claude-sonnet"]), None + ) + cache: Final = await _cache_with_membership("alice", "team-a", allowed_models=["gpt-5"]) + + await _check_caller_models(agent_key, "gpt-5", load_team, load_user, cache) + with pytest.raises(ProxyException) as exc_info: + await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user, cache) + + assert "User=alice, Team=team-a" in exc_info.value.internal_message + + +@pytest.mark.asyncio +async def test_agent_key_acting_for_a_teamless_user_is_capped_at_that_users_models(): + agent_key: Final = _agent_key_acting_for(user_id="alice", team_id=None) + load_team, load_user, asked = _caller_loaders(None, LiteLLM_UserTable(user_id="alice", models=["gpt-5"])) + + await _check_caller_models(agent_key, "gpt-5", load_team, load_user) + with pytest.raises(ProxyException) as exc_info: + await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user) + + assert exc_info.value.type == ProxyErrorTypes.user_model_access_denied + assert asked == ["team:None", "user:alice", "team:None", "user:alice"] + + +@pytest.mark.asyncio +async def test_agent_key_without_an_echoed_caller_keeps_its_own_models(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + load_team, load_user, asked = _caller_loaders(LiteLLM_TeamTable(team_id="team-a", models=[]), None) + + await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user) + + assert asked == [] diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 0969a913605..cd14f630130 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -3350,15 +3350,26 @@ async def test_auth_builder_single_team_db_fallback_when_jwt_has_no_team( mock_get_membership.assert_not_called() -@pytest.mark.asyncio -async def test_auth_builder_single_team_fallback_membership_error_skips_no_raise(): - """ - get_team_object succeeds but get_team_membership raises — do not set team; no exception. - """ - from fastapi import HTTPException +class _UnreachableMembershipPrisma: + class db: + class litellm_teammembership: + @staticmethod + async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None: + raise httpx.ConnectError("All connection attempts failed") - user_id = "u_mem_fail" - team_id_val = "team_mem_fail" + +@pytest.mark.asyncio +async def test_auth_builder_single_team_fallback_membership_outage_raises_instead_of_dropping_the_team(): + """ + get_team_object succeeds but the membership read hits a database outage: the + outage propagates (auth maps it to 503) instead of the team being dropped. + """ + from litellm.proxy.auth.auth_exception_handler import _as_proxy_exception + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + user_id = "u_mem_outage" + team_id_val = "team_mem_outage" user_object = LiteLLM_UserTable( user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER, @@ -3367,6 +3378,7 @@ async def test_auth_builder_single_team_fallback_membership_error_skips_no_raise team_table = LiteLLM_TeamTable(team_id=team_id_val) jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + cache = UserApiKeyCache() with ( patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, @@ -3416,34 +3428,26 @@ async def test_auth_builder_single_team_fallback_membership_error_skips_no_raise "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, ) as mock_get_team, - patch( - "litellm.proxy.auth.handle_jwt.get_team_membership", - new_callable=AsyncMock, - ) as mock_get_membership, ): mock_auth_jwt.return_value = {"sub": user_id, "scope": ""} mock_get_team.return_value = team_table - mock_get_membership.side_effect = HTTPException( - status_code=500, detail="membership lookup failed" - ) - result = await JWTAuthManager.auth_builder( - api_key="test_jwt_token", - jwt_handler=jwt_handler, - request_data={"model": "gpt-4"}, - general_settings={"enforce_rbac": False}, - route="/chat/completions", - prisma_client=None, - user_api_key_cache=None, - parent_otel_span=None, - proxy_logging_obj=None, - ) + with pytest.raises(httpx.ConnectError) as raised: + await JWTAuthManager.auth_builder( + api_key="test_jwt_token", + jwt_handler=jwt_handler, + request_data={"model": "gpt-4"}, + general_settings={"enforce_rbac": False}, + route="/chat/completions", + prisma_client=_UnreachableMembershipPrisma(), + user_api_key_cache=cache, + parent_otel_span=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) - assert result["team_id"] is None - assert result["team_object"] is None - assert result["team_membership"] is None - mock_get_team.assert_called() - mock_get_membership.assert_called_once() + mock_get_team.assert_called() + surfaced = _as_proxy_exception(raised.value) + assert (surfaced.code, surfaced.type) == ("503", ProxyErrorTypes.no_db_connection) # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index da36071a5b4..f03abe8f124 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -59,6 +59,7 @@ from litellm.proxy.auth.user_api_key_auth import ( _user_api_key_auth_builder, get_api_key, user_api_key_auth, + user_api_key_auth_websocket_for_model, ) from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata @@ -9043,3 +9044,90 @@ async def test_router_settings_model_group_alias_authorizes_target_for_team(monk await authorize() assert (await request.json())["model"] == target assert get_client_requested_model(request) == "AgentX-LLM" + + +@pytest.mark.asyncio +async def test_reserve_budget_after_common_checks_hands_the_reservation_to_the_request_state(): + from fastapi import Request + + request = Request(scope={"type": "http"}) + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=reservation), + ): + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/batches/batch_123/cancel", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={}, + request=request, + ) + + assert user_api_key_auth_obj.budget_reservation is reservation + assert request.state.budget_reservation is reservation + assert request.scope["state"]["budget_reservation"] is reservation + + +@pytest.mark.asyncio +async def test_reserve_budget_after_common_checks_clears_the_request_state_when_budget_checks_skip(): + from fastapi import Request + + request = Request(scope={"type": "http", "state": {"budget_reservation": {"reserved_cost": 0.5}}}) + + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=UserAPIKeyAuth(token="test_token"), + request_data={"model": "free-model"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=True, + general_settings={}, + request=request, + ) + + assert request.state.budget_reservation is None + + +@pytest.mark.asyncio +async def test_websocket_auth_hands_the_reservation_to_the_socket_state(): + from fastapi import WebSocket + + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + websocket = WebSocket( + scope={ + "type": "websocket", + "path": "/v1/realtime", + "headers": [(b"authorization", b"Bearer sk-1234")], + "query_string": b"model=gpt-realtime", + }, + receive=AsyncMock(), + send=AsyncMock(), + ) + + async def auth_that_reserves(request, api_key): + request.state.budget_reservation = reservation + return UserAPIKeyAuth(token="hashed", budget_reservation=reservation) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + new=AsyncMock(side_effect=auth_that_reserves), + ): + result = await user_api_key_auth_websocket_for_model(websocket, model="gpt-realtime") + + assert result.budget_reservation == reservation + assert websocket.state.budget_reservation is reservation + assert websocket.scope["state"]["budget_reservation"] is reservation diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index a48c64eb4a0..9353c149d15 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -5,13 +5,16 @@ import shlex import stat import sys import time +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError from pathlib import Path +from threading import Event +from typing import Final from unittest.mock import patch import pytest from click.testing import CliRunner -from litellm.litellm_core_utils.private_json import commit_staged_json +from litellm.litellm_core_utils.private_json import commit_staged_json, write_private_bytes from litellm.proxy.client.cli.commands.claude_settings import ( ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, AUTOROUTE_BACKUP_PATH, @@ -773,7 +776,7 @@ class TestStatusLine: script = tmp_path / "lite" / "statusline.py" command = install_statusline_script(script) - assert script.read_bytes() == pathlib.Path(statusline_script.__file__).read_bytes() + assert script.read_bytes().split(b"\n", 1)[1] == pathlib.Path(statusline_script.__file__).read_bytes() assert shlex.split(command) == [sys.executable, str(script)] assert command == statusline_command(script) assert stat.S_IMODE(script.stat().st_mode) == 0o600 @@ -783,15 +786,13 @@ class TestStatusLine: def test_a_reinstall_replaces_the_script_in_one_step_and_a_refused_one_leaves_the_old_script_whole(self, tmp_path): # Claude Code may be running the script at the moment `lite` reinstalls it; the file it has open # must stay complete, and a reinstall that cannot land must not leave a truncated script behind. - from litellm.proxy.client.cli.commands import statusline_script - script = tmp_path / "lite" / "statusline.py" install_statusline_script(script) - bundled = pathlib.Path(statusline_script.__file__).read_bytes() + bundled = script.read_bytes() with script.open("rb") as running: install_statusline_script(script) assert running.read() == bundled - assert [child.name for child in script.parent.iterdir()] == ["statusline.py"] + assert {child.name for child in script.parent.iterdir()} <= {"statusline.py", "statusline.py.lock"} if os.geteuid() != 0: script.parent.chmod(0o500) @@ -802,6 +803,141 @@ class TestStatusLine: script.parent.chmod(0o700) assert script.read_bytes() == bundled + @pytest.mark.parametrize( + ("installed_version", "older_version"), + (("2.10.0", "2.9.0"), ("2.1.0", "2.1.0rc1"), ("2.1.0rc1", "2.1.0.dev2"), ("2.1.0.post1", "2.1.0")), + ) + def test_an_older_cli_preserves_the_newer_footer( + self, tmp_path: Path, capsys: pytest.CaptureFixture[str], installed_version: str, older_version: str + ) -> None: + script: Final = tmp_path / "statusline.py" + command: Final = install_statusline_script(script, package_version=installed_version) + installed: Final = script.read_bytes() + modified: Final = script.stat().st_mtime_ns + + assert install_statusline_script(script, package_version=older_version) == command + + assert script.read_bytes() == installed + assert script.stat().st_mtime_ns == modified + assert f"Keeping the status line from LiteLLM {installed_version}" in capsys.readouterr().err + + @pytest.mark.parametrize( + "old_header", (b"", b"# litellm-statusline-version: invalid\n", b"# litellm-statusline-version: \xff\n") + ) + def test_a_legacy_or_damaged_version_marker_is_repaired(self, tmp_path: Path, old_header: bytes) -> None: + from litellm.proxy.client.cli.commands import statusline_script + + script: Final = tmp_path / "statusline.py" + script.write_bytes(old_header + b"print('old footer')\n") + + install_statusline_script(script, package_version="2.1.0") + + assert script.read_bytes() == ( + b"# litellm-statusline-version: 2.1.0\n" + Path(statusline_script.__file__).read_bytes() + ) + + @pytest.mark.parametrize("next_version", ("2.1.0", "2.2.0")) + def test_an_equal_or_newer_cli_refreshes_the_footer(self, tmp_path: Path, next_version: str) -> None: + from litellm.proxy.client.cli.commands import statusline_script + + script: Final = tmp_path / "statusline.py" + script.write_bytes(b"# litellm-statusline-version: 2.1.0\nprint('old footer')\n") + + install_statusline_script(script, package_version=next_version) + + assert script.read_bytes() == ( + f"# litellm-statusline-version: {next_version}\n".encode() + Path(statusline_script.__file__).read_bytes() + ) + + def test_configure_keeps_a_newer_footer_while_updating_settings( + self, tmp_path: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + script: Final = tmp_path / "statusline.py" + script.write_bytes(b"# litellm-statusline-version: 999999.0.0\nprint('newer footer')\n") + installed: Final = script.read_bytes() + rig: Final = _Rig(tmp_path, {"theme": "dark"}) + + rig.configure(script_path=script) + + assert script.read_bytes() == installed + assert rig.read()["statusLine"]["command"] == statusline_command(script) + assert rig.read()["env"]["ANTHROPIC_BASE_URL"] == PROXY + assert "Keeping the status line" in capsys.readouterr().err + + @pytest.mark.parametrize("package_version", ("unknown", "", "invalid-version")) + @pytest.mark.parametrize("existing", (None, b"print('legacy footer')\n", b"# litellm-statusline-version: invalid\n")) + def test_an_unknown_cli_version_can_install_and_refresh_an_unversioned_footer( + self, tmp_path: Path, package_version: str, existing: bytes | None + ) -> None: + from litellm.proxy.client.cli.commands import statusline_script + + script: Final = tmp_path / "statusline.py" + if existing is not None: + script.write_bytes(existing) + + assert install_statusline_script(script, package_version=package_version) == statusline_command(script) + assert script.read_bytes() == Path(statusline_script.__file__).read_bytes() + + def test_an_unknown_cli_version_preserves_a_versioned_footer( + self, tmp_path: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + script: Final = tmp_path / "statusline.py" + command: Final = install_statusline_script(script, package_version="2.1.0") + installed: Final = script.read_bytes() + + assert install_statusline_script(script, package_version="unknown") == command + assert script.read_bytes() == installed + assert "Keeping the status line from LiteLLM 2.1.0" in capsys.readouterr().err + + @pytest.mark.parametrize(("first_version", "second_version"), (("2.0", "3.0"), ("3.0", "2.0"))) + def test_overlapping_installs_keep_the_newest_footer( + self, tmp_path: Path, first_version: str, second_version: str + ) -> None: + from litellm.proxy.client.cli.commands import statusline_script + + script: Final = tmp_path / "statusline.py" + first_writing: Final = Event() + release_first: Final = Event() + second_started: Final = Event() + + def paused_write(path: str, data: bytes) -> None: + first_writing.set() + assert release_first.wait(5), "First installer was never released" + write_private_bytes(path, data) + + def second_install() -> str: + second_started.set() + return install_statusline_script(script, package_version=second_version) + + with ThreadPoolExecutor(max_workers=2) as pool: + first: Final = pool.submit(install_statusline_script, script, package_version=first_version, write=paused_write) + try: + assert first_writing.wait(5), "First installer did not reach the write" + second: Final = pool.submit(second_install) + assert second_started.wait(5), "Second installer did not start" + with pytest.raises(FutureTimeoutError): + second.result(timeout=0.5) + finally: + release_first.set() + assert first.result(timeout=5) == statusline_command(script) + assert second.result(timeout=5) == statusline_command(script) + + assert script.read_bytes() == b"# litellm-statusline-version: 3.0\n" + Path(statusline_script.__file__).read_bytes() + + def test_a_failed_install_keeps_the_footer_and_releases_the_lock(self, tmp_path: Path) -> None: + script: Final = tmp_path / "statusline.py" + install_statusline_script(script, package_version="2.0") + installed: Final = script.read_bytes() + + def failed_write(path: str, data: bytes) -> None: + raise OSError("disk full") + + with pytest.raises(ClaudeSettingsError, match="disk full"): + install_statusline_script(script, package_version="3.0", write=failed_write) + assert script.read_bytes() == installed + assert install_statusline_script(script, package_version="3.0") == statusline_command(script) + assert script.read_bytes().startswith(b"# litellm-statusline-version: 3.0\n") + def test_configure_installs_it_and_unconfigure_removes_only_ours(self, tmp_path): rig = _Rig(tmp_path, {"theme": "dark"}) script = tmp_path / "statusline.py" diff --git a/tests/test_litellm/proxy/common_utils/test_error_body_call_id.py b/tests/test_litellm/proxy/common_utils/test_error_body_call_id.py new file mode 100644 index 00000000000..8872b5de397 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_error_body_call_id.py @@ -0,0 +1,35 @@ +import pytest + +from litellm.proxy._types import ConfigGeneralSettings +from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id + + +@pytest.mark.parametrize( + "general_settings, call_id, expected", + [ + ({"include_call_id_in_error_body": True}, "call-1", "call-1"), + ({"include_call_id_in_error_body": True}, None, None), + ({"include_call_id_in_error_body": True}, "", None), + ({"include_call_id_in_error_body": False}, "call-1", None), + ({"include_call_id_in_error_body": "true"}, "call-1", None), + ({}, "call-1", None), + ], +) +def test_only_the_boolean_opt_in_with_a_real_id_yields_a_body_call_id(general_settings, call_id, expected): + """The setting is off by default and only a literal true turns it on; without an id + there is nothing to copy, so the body must never get a fabricated one.""" + assert error_body_call_id(general_settings, call_id) == expected + + +def test_with_call_id_appends_the_key_without_touching_the_input(): + error = {"message": "bad input", "type": "invalid_request_error", "param": None, "code": "400"} + + assert with_call_id(error, "call-1") == {**error, "litellm_call_id": "call-1"} + assert with_call_id(error, None) == error + assert "litellm_call_id" not in error + + +def test_the_setting_name_is_a_config_general_settings_field(): + """The yaml key the docs name and the key the runtime reads must be the same field.""" + assert ConfigGeneralSettings.model_validate({"include_call_id_in_error_body": True}).include_call_id_in_error_body + assert ConfigGeneralSettings.model_validate({}).include_call_id_in_error_body is None diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index 6f7c20166c5..ca2ff8bcce1 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -207,7 +207,9 @@ async def test_get_agent_with_read_through_recovers_agent_by_name(clean_agent_re @pytest.mark.asyncio -async def test_get_agent_with_read_through_returns_none_for_unknown_agent(clean_agent_registry, monkeypatch): +async def test_get_agent_with_read_through_returns_none_for_unknown_agent( + clean_agent_registry, fresh_agent_read_through, monkeypatch +): from unittest.mock import AsyncMock, MagicMock import litellm.proxy.proxy_server as proxy_server diff --git a/tests/test_litellm/proxy/conftest.py b/tests/test_litellm/proxy/conftest.py index 8ef5017a952..49dc8d02bdb 100644 --- a/tests/test_litellm/proxy/conftest.py +++ b/tests/test_litellm/proxy/conftest.py @@ -273,3 +273,12 @@ def create_proxy_test_client( # Initialize proxy asyncio.run(initialize(config=config_fp, debug=init_options.get("debug", False))) return TestClient(app) + + +@pytest.fixture +def fresh_agent_read_through(monkeypatch): + from litellm.proxy.common_utils import registry_read_through + + read_through = registry_read_through.RegistryReadThrough(resync=registry_read_through._resync_agents) + monkeypatch.setattr(registry_read_through, "agent_registry_read_through", read_through) + return read_through diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 33614d2eeca..89e72debbf5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -4,6 +4,7 @@ Tests PII detection and masking for different message formats """ import asyncio +import json from contextlib import asynccontextmanager from unittest.mock import MagicMock, patch @@ -18,7 +19,7 @@ from litellm.proxy.guardrails.guardrail_hooks.presidio import ( ) from litellm.exceptions import GuardrailRaisedException from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType -from litellm.types.utils import Choices, Message, ModelResponse +from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices from litellm.exceptions import BlockedPiiEntityError @@ -2331,47 +2332,320 @@ async def test_apply_guardrail_masks_on_request(): assert "John Smith" not in result["texts"][0] +def _anthropic_sse(event_type: str, payload: dict) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() + + +def _anthropic_text_deltas(chunks: list[bytes]) -> list[tuple[int, str]]: + deltas = [] + for line in b"".join(chunks).decode().split("\n"): + if not line.startswith("data: "): + continue + event = json.loads(line[6:]) + if event.get("type") == "content_block_delta" and event["delta"].get("type") == "text_delta": + deltas.append((event["index"], event["delta"]["text"])) + return deltas + + +def _chat_delta_chunk(text: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-out-mask", + choices=[StreamingChoices(index=0, delta=Delta(content=text, role="assistant"), finish_reason=finish_reason)], + created=1, + model="gpt-4", + object="chat.completion.chunk", + ) + + @pytest.mark.asyncio -async def test_apply_to_output_streaming_bytes_only_logs_warning(): +async def test_apply_to_output_streaming_chat_chunks_are_masked_as_one_response(): """ - Regression test: when apply_to_output=True and the stream contains only - bytes chunks (Anthropic native SSE), output masking is skipped. - A warning must be logged so operators are aware. + Structured chat completion chunks are buffered, assembled and masked as a + whole, so a card number split across deltas cannot reach the caller. """ guardrail = _OPTIONAL_PresidioPIIMasking( mock_testing=True, apply_to_output=True, + mock_redacted_text={"text": "my card is "}, + ) + + async def mock_stream(): + yield _chat_delta_chunk("my card is 4111") + yield _chat_delta_chunk(" 1111 1111 1111") + yield _chat_delta_chunk("", finish_reason="stop") + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={"messages": [{"role": "user", "content": "what is my card"}]}, + ): + collected.append(chunk) + + assert all(isinstance(chunk, ModelResponseStream) for chunk in collected) + joined = "".join(chunk.choices[0].delta.content or "" for chunk in collected) + assert joined == "my card is " + assert collected[-1].choices[0].finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_bytes_after_chat_chunks_are_passed_through_in_order(): + """ + Once structured chunks have been buffered, a trailing bytes frame belongs to + the same stream and must be forwarded rather than treated as a new SSE stream. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": "hello"}, + ) + trailer = b"data: [DONE]\n\n" + + async def mock_stream(): + yield _chat_delta_chunk("hello", finish_reason="stop") + yield trailer + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + assert collected[0] == trailer + assert len(collected) == 2 + assert isinstance(collected[1], ModelResponseStream) + assert collected[1].choices[0].delta.content == "hello" + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_anthropic_sse_bytes_masks_text_split_across_deltas(): + """ + Anthropic native /v1/messages streams reach the post_call hook as raw SSE + bytes. Output masking must run over the whole content block so a card + number split across text_delta events cannot reach the caller. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, ) byte_chunks = [ - b'data: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n', - b'data: {"type":"content_block_delta","delta":{"text":" world"}}\n\n', + _anthropic_sse( + "message_start", + {"type": "message_start", "message": {"id": "msg_1", "model": "claude", "content": [], "usage": {}}}, + ), + _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "4111"}}, + ), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " 1111 1111 1111"}}, + ), + _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {}}), + _anthropic_sse("message_stop", {"type": "message_stop"}), ] async def mock_stream(): for b in byte_chunks: yield b - mock_user_api_key = UserAPIKeyAuth(api_key="test-key") + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + assert all(isinstance(chunk, bytes) for chunk in collected) + joined = b"".join(collected).decode() + assert "4111" not in joined + assert "".join(text for _, text in _anthropic_text_deltas(collected)) == "" + assert joined.count("event: message_start") == 1 + assert joined.count("event: message_stop") == 1 + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_anthropic_sse_bytes_without_pii_are_forwarded_unchanged(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": "Hello world"}, + ) + + byte_chunks = [ + _anthropic_sse( + "message_start", + {"type": "message_start", "message": {"id": "msg_1", "model": "claude", "content": [], "usage": {}}}, + ), + _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}, + ), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " world"}}, + ), + _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {}}), + _anthropic_sse("message_stop", {"type": "message_stop"}), + ] + + async def mock_stream(): + for b in byte_chunks: + yield b collected = [] - with patch("litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger") as mock_logger: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + assert collected == byte_chunks + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_anthropic_sse_bytes_fail_closed_when_presidio_is_unreachable(): + """ + The raw SSE stream is fully drained before masking, so a Presidio outage + must surface as an error to the caller: replaying the unscanned frames + would hand over whatever PII the model generated. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + presidio_analyzer_api_base="http://127.0.0.1:9", + presidio_anonymizer_api_base="http://127.0.0.1:9", + ) + + byte_chunks = [ + _anthropic_sse( + "message_start", + {"type": "message_start", "message": {"id": "msg_1", "model": "claude", "content": [], "usage": {}}}, + ), + _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello world"}}, + ), + _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse("message_stop", {"type": "message_stop"}), + ] + + async def mock_stream(): + for b in byte_chunks: + yield b + + collected = [] + + async def collect_masked_stream(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), response=mock_stream(), request_data={}, ): collected.append(chunk) - # All bytes should be yielded through - assert len(collected) == len(byte_chunks) - for original, received in zip(byte_chunks, collected): - assert original == received + with pytest.raises(Exception, match="Presidio PII analysis failed"): + await collect_masked_stream() - # Warning must be logged about skipped masking - mock_logger.warning.assert_called_once() - warning_msg = mock_logger.warning.call_args[0][0] - assert "Output PII masking was skipped" in warning_msg + assert collected == [] + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_anthropic_sse_bytes_block_action_raises_instead_of_replaying(): + """ + A BLOCK on generated PII must refuse the streaming /v1/messages response the + same way it refuses the non streaming one, not replay the raw frames. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + apply_to_output=True, + mock_testing=False, + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + pii_entities_config={PiiEntityType.CREDIT_CARD: PiiAction.BLOCK}, + ) + + byte_chunks = [ + _anthropic_sse( + "message_start", + {"type": "message_start", "message": {"id": "msg_1", "model": "claude", "content": [], "usage": {}}}, + ), + _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "4111 1111 1111 1111"}}, + ), + _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse("message_stop", {"type": "message_stop"}), + ] + + async def mock_stream(): + for b in byte_chunks: + yield b + + analyzer_hit = [{"entity_type": "CREDIT_CARD", "score": 0.99, "start": 0, "end": 19}] + collected = [] + + async def collect_masked_stream(): + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + with patch.object(guardrail, "_get_session_iterator", _make_mock_session_iterator(analyzer_hit)): + with pytest.raises(BlockedPiiEntityError): + await collect_masked_stream() + + assert collected == [] + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_propagates_upstream_error_when_nothing_was_buffered(): + """ + An upstream guardrail that rejects the stream before the first chunk must + surface as an error to the caller, not as an empty 200 stream. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + + async def failing_stream(): + raise RuntimeError("upstream guardrail rejected the stream") + yield b"" + + with pytest.raises(RuntimeError, match="upstream guardrail rejected the stream"): + async for _ in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=failing_stream(), + request_data={}, + ): + pass @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index d1d22d0d7c2..c90f88ec110 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -2,7 +2,7 @@ import logging from types import SimpleNamespace -from typing import Final +from typing import TYPE_CHECKING, Final, Literal import pytest @@ -13,7 +13,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.llms import load_guardrail_translation_mappings +from litellm.llms import discover_guardrail_translation_mappings, load_guardrail_translation_mappings from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, @@ -41,7 +41,10 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai ) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.utils import CallTypes, Delta, ModelResponseStream, StreamingChoices +from litellm.types.utils import CallTypes, Delta, GenericGuardrailAPIInputs, ModelResponseStream, StreamingChoices + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj class RecordingGuardrail(CustomGuardrail): @@ -61,6 +64,18 @@ class RecordingGuardrail(CustomGuardrail): return {"texts": inputs.get("texts", [])} +class RewritingGuardrail(RecordingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: CustomGuardrail.apply_guardrail contract + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + recorded: Final = await super().apply_guardrail(inputs, request_data, input_type, logging_obj=logging_obj) + return GenericGuardrailAPIInputs(texts=[f"{text} [GUARDRAILED]" for text in recorded["texts"]]) + + class _NoopTranslation(BaseTranslation): """Test translation handler that simply echoes input/output.""" @@ -115,9 +130,7 @@ class TestUnifiedLLMGuardrails: assert msgs[0]["content"] == "sys" def test_effective_skip_respects_per_guardrail_over_global(self, monkeypatch): - monkeypatch.setattr( - litellm, "skip_system_message_in_guardrail", True, raising=False - ) + monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", True, raising=False) class G: skip_system_message_in_guardrail = False @@ -130,21 +143,15 @@ class TestUnifiedLLMGuardrails: assert effective_skip_system_message_for_guardrail(G2()) is True @pytest.mark.asyncio - async def test_openai_handler_skips_system_in_guardrail_inputs( - self, monkeypatch - ): - monkeypatch.setattr( - litellm, "skip_system_message_in_guardrail", True, raising=False - ) + async def test_openai_handler_skips_system_in_guardrail_inputs(self, monkeypatch): + monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", True, raising=False) captured = {} class MockGuardrail: skip_system_message_in_guardrail = None - async def apply_guardrail( - self, inputs, request_data, input_type, logging_obj=None - ): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): captured["inputs"] = inputs return inputs @@ -169,21 +176,15 @@ class TestUnifiedLLMGuardrails: assert data["messages"][0]["content"] == "secret system" @pytest.mark.asyncio - async def test_openai_handler_per_guardrail_skip_false_overrides_global( - self, monkeypatch - ): - monkeypatch.setattr( - litellm, "skip_system_message_in_guardrail", True, raising=False - ) + async def test_openai_handler_per_guardrail_skip_false_overrides_global(self, monkeypatch): + monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", True, raising=False) captured = {} class MockGuardrail: skip_system_message_in_guardrail = False - async def apply_guardrail( - self, inputs, request_data, input_type, logging_obj=None - ): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): captured["inputs"] = inputs return inputs @@ -201,10 +202,7 @@ class TestUnifiedLLMGuardrails: ) assert "sys" in captured["inputs"]["texts"] - roles = { - m.get("role") - for m in (captured["inputs"].get("structured_messages") or []) - } + roles = {m.get("role") for m in (captured["inputs"].get("structured_messages") or [])} assert "system" in roles class TestSkipToolMessageForChatCompletions: @@ -229,12 +227,8 @@ class TestUnifiedLLMGuardrails: assert all(m["role"] != "tool" for m in out) assert msgs[2]["content"] == "tool result" - def test_effective_skip_tool_respects_per_guardrail_over_global( - self, monkeypatch - ): - monkeypatch.setattr( - litellm, "skip_tool_message_in_guardrail", True, raising=False - ) + def test_effective_skip_tool_respects_per_guardrail_over_global(self, monkeypatch): + monkeypatch.setattr(litellm, "skip_tool_message_in_guardrail", True, raising=False) class G: skip_tool_message_in_guardrail = False @@ -248,18 +242,14 @@ class TestUnifiedLLMGuardrails: @pytest.mark.asyncio async def test_openai_handler_skips_tool_in_guardrail_inputs(self, monkeypatch): - monkeypatch.setattr( - litellm, "skip_tool_message_in_guardrail", True, raising=False - ) + monkeypatch.setattr(litellm, "skip_tool_message_in_guardrail", True, raising=False) captured = {} class MockGuardrail: skip_tool_message_in_guardrail = None - async def apply_guardrail( - self, inputs, request_data, input_type, logging_obj=None - ): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): captured["inputs"] = inputs return inputs @@ -299,21 +289,15 @@ class TestUnifiedLLMGuardrails: assert data["messages"][2]["content"] == "secret tool result" @pytest.mark.asyncio - async def test_openai_handler_per_guardrail_skip_tool_false_overrides_global( - self, monkeypatch - ): - monkeypatch.setattr( - litellm, "skip_tool_message_in_guardrail", True, raising=False - ) + async def test_openai_handler_per_guardrail_skip_tool_false_overrides_global(self, monkeypatch): + monkeypatch.setattr(litellm, "skip_tool_message_in_guardrail", True, raising=False) captured = {} class MockGuardrail: skip_tool_message_in_guardrail = False - async def apply_guardrail( - self, inputs, request_data, input_type, logging_obj=None - ): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): captured["inputs"] = inputs return inputs @@ -331,10 +315,7 @@ class TestUnifiedLLMGuardrails: ) assert "tr" in captured["inputs"]["texts"] - roles = { - m.get("role") - for m in (captured["inputs"].get("structured_messages") or []) - } + roles = {m.get("role") for m in (captured["inputs"].get("structured_messages") or [])} assert "tool" in roles class TestAsyncPreCallHook: @@ -360,6 +341,38 @@ class TestUnifiedLLMGuardrails: assert guardrail.event_history == [GuardrailEventHooks.pre_mcp_call] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "call_type", + ["avideo_generation", "acreate_video", "avideo_remix", "avideo_edit", "avideo_extension"], + ) + async def test_video_routes_scan_prompt_and_keep_rewrite(self, monkeypatch, call_type: str) -> None: + """LIT-6685: /v1/videos dispatches call_type="avideo_generation", which the + hook once swallowed as an unknown CallTypes value and returned unscanned. + Runs against the discovered handler map so the video package must really exist.""" + _patch_translation_mappings(monkeypatch, discover_guardrail_translation_mappings()) + handler = UnifiedLLMGuardrails() + guardrail = RewritingGuardrail() + data = { + "guardrail_to_apply": guardrail, + "model": "veo-3.1-fast", + "prompt": "a paper boat on a stream", + "seconds": "4", + } + + result = await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + cache=DualCache(), + data=data, + call_type=call_type, + ) + + assert guardrail.event_history == [GuardrailEventHooks.pre_call] + assert [call["inputs"]["texts"] for call in guardrail.apply_calls] == [["a paper boat on a stream"]] + assert guardrail.apply_calls[0]["inputs"]["model"] == "veo-3.1-fast" + assert result["prompt"] == "a paper boat on a stream [GUARDRAILED]" + assert result["seconds"] == "4" + class TestAsyncModerationHook: @pytest.mark.asyncio async def test_uses_mcp_event_type(self): @@ -424,7 +437,9 @@ class TestUnifiedLLMGuardrails: async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj=None): # type: ignore[override] return data - async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None): # type: ignore[override] + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None + ): # type: ignore[override] return response async def process_output_streaming_response( @@ -493,9 +508,7 @@ class TestUnifiedLLMGuardrails: response=mock_stream(), request_data=request_data, ): - content = ( - item.choices[0].delta.content if item.choices[0].delta else None - ) + content = item.choices[0].delta.content if item.choices[0].delta else None yielded_contents.append(content) # Every chunk should have non-empty content @@ -546,23 +559,18 @@ class TestUnifiedLLMGuardrails: ], ) @pytest.mark.asyncio - async def test_post_call_scans_output_on_every_registered_alias( - self, request_route: str - ) -> None: + async def test_post_call_scans_output_on_every_registered_alias(self, request_route: str) -> None: handler = UnifiedLLMGuardrails() guardrail = RecordingGuardrail() await handler.async_post_call_success_hook( data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, - user_api_key_dict=UserAPIKeyAuth( - api_key="test-key", request_route=request_route - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route=request_route), response=self._responses_api_response(), ) assert guardrail.apply_calls, ( - f"guardrail never ran for request_route={request_route!r}; model " - f"output reached the client unscanned" + f"guardrail never ran for request_route={request_route!r}; model output reached the client unscanned" ) assert guardrail.apply_calls[0]["input_type"] == "response" assert guardrail.apply_calls[0]["inputs"]["texts"] == ["Paris"] @@ -592,18 +600,14 @@ class TestUnifiedLLMGuardrails: assert CallTypes.responses in mappings @pytest.mark.asyncio - async def test_unresolvable_route_skips_scanning_and_says_so( - self, caplog: pytest.LogCaptureFixture - ) -> None: + async def test_unresolvable_route_skips_scanning_and_says_so(self, caplog: pytest.LogCaptureFixture) -> None: handler = UnifiedLLMGuardrails() guardrail = RecordingGuardrail() with caplog.at_level(logging.WARNING): result = await handler.async_post_call_success_hook( data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, - user_api_key_dict=UserAPIKeyAuth( - api_key="test-key", request_route="/cursor/chat/completions" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/cursor/chat/completions"), response=self._responses_api_response(), ) @@ -622,9 +626,7 @@ class TestUnifiedLLMGuardrails: with caplog.at_level(logging.WARNING): await handler.async_post_call_success_hook( data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, - user_api_key_dict=UserAPIKeyAuth( - api_key="test-key", request_route="/v1/chat/completions" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"), response=self._responses_api_response(), ) @@ -734,15 +736,10 @@ class TestUnifiedLLMGuardrails: assert guardrail.event_history == [GuardrailEventHooks.pre_call] assert len(guardrail.apply_calls) == 1 assert guardrail.apply_calls[0]["input_type"] == "request" - assert ( - "https://arxiv.org/pdf/2201.04234" - in guardrail.apply_calls[0]["inputs"]["texts"] - ) + assert "https://arxiv.org/pdf/2201.04234" in guardrail.apply_calls[0]["inputs"]["texts"] # Data should be returned with document intact - assert ( - result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234" - ) + assert result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234" @pytest.mark.asyncio async def test_moderation_hook_invokes_ocr_handler(self): @@ -770,10 +767,7 @@ class TestUnifiedLLMGuardrails: assert guardrail.event_history == [GuardrailEventHooks.during_call] assert len(guardrail.apply_calls) == 1 - assert ( - "https://example.com/scan.png" - in guardrail.apply_calls[0]["inputs"]["texts"] - ) + assert "https://example.com/scan.png" in guardrail.apply_calls[0]["inputs"]["texts"] @pytest.mark.asyncio async def test_post_call_success_hook_guardrails_ocr_output(self): @@ -789,9 +783,7 @@ class TestUnifiedLLMGuardrails: def should_run_guardrail(self, data, event_type): # type: ignore[override] return True - async def apply_guardrail( - self, inputs, request_data, input_type, **kwargs - ): + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): texts = inputs.get("texts", []) return {"texts": [t.replace("SECRET", "[REDACTED]") for t in texts]} @@ -1538,9 +1530,7 @@ class TestStreamingTransform: # And the redacted text ("SECRET") reached the wire on some non-tool # chunk (i.e. the text terminator). transformed = "".join( - item.choices[0].delta.content or "" - for item in out - if item.choices and not item.choices[0].delta.tool_calls + item.choices[0].delta.content or "" for item in out if item.choices and not item.choices[0].delta.tool_calls ) assert "SECRET" in transformed assert "secret" not in transformed @@ -1685,7 +1675,9 @@ class TestStreamingTransform: _stream_chunk("went home."), ModelResponseStream( choices=[ - StreamingChoices(index=0, delta=Delta(content=None, role="assistant", tool_calls=None), finish_reason=None), + StreamingChoices( + index=0, delta=Delta(content=None, role="assistant", tool_calls=None), finish_reason=None + ), StreamingChoices( index=1, delta=Delta( @@ -1985,9 +1977,7 @@ class TestStreamingHttpErrorFrames: guardrail = _EosHttpBlockingGuardrail() chunks = _anthropic_message_chunks(["hello ", "world"]) - out = await _drive_stream( - UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages" - ) + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") raw = b"".join(c for c in out if isinstance(c, bytes)).decode() assert "hello " in raw @@ -2011,9 +2001,7 @@ class TestStreamingHttpErrorFrames: }, ] - out = await _drive_stream( - UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses" - ) + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") assert chunks[0] in out and chunks[1] in out assert chunks[2] not in out @@ -2076,9 +2064,7 @@ class TestStreamingGuardrailInformationBucket: for chunk in chunks: yield chunk - user_api_key_dict = UserAPIKeyAuth( - api_key="test-key", user_id="user-1", request_route="/v1/chat/completions" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="user-1", request_route="/v1/chat/completions") request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}} out = [] @@ -2407,7 +2393,5 @@ class TestTranslationMappingsAreReadLive: assert len(guardrail.apply_calls) == 1 assert not [ - name - for name, value in vars(unified_module).items() - if isinstance(value, dict) and CallTypes.aocr in value + name for name, value in vars(unified_module).items() if isinstance(value, dict) and CallTypes.aocr in value ] diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index d687f8d1c8d..13e57408bd0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -17,7 +17,9 @@ from litellm.proxy._types import ( LitellmUserRoles, UserAPIKeyAuth, ) +from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.proxy_server import app +from litellm.types.agents import AgentResponse def _make_access_group_record( @@ -126,6 +128,7 @@ def client_and_mocks(monkeypatch): mock_agents_table = MagicMock() mock_agents_table.find_many = AsyncMock(return_value=[]) + mock_agents_table.update = AsyncMock(return_value=None) @asynccontextmanager async def mock_tx(): @@ -133,6 +136,7 @@ def client_and_mocks(monkeypatch): litellm_accessgrouptable=mock_access_group_table, litellm_teamtable=mock_team_table, litellm_verificationtoken=mock_key_table, + litellm_agentstable=mock_agents_table, ) yield tx @@ -158,15 +162,9 @@ def client_and_mocks(monkeypatch): mock_proxy_logging = MagicMock() mock_proxy_logging.internal_usage_cache = MagicMock() mock_proxy_logging.internal_usage_cache.dual_cache = MagicMock() - mock_proxy_logging.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock( - return_value=None - ) - mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( - return_value=None - ) - mock_proxy_logging.internal_usage_cache.dual_cache.async_set_cache = AsyncMock( - return_value=None - ) + mock_proxy_logging.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock(return_value=None) + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None) + mock_proxy_logging.internal_usage_cache.dual_cache.async_set_cache = AsyncMock(return_value=None) monkeypatch.setattr(ps, "proxy_logging_obj", mock_proxy_logging) admin_user = UserAPIKeyAuth( @@ -239,9 +237,7 @@ def test_create_access_group_duplicate_name_conflict(client_and_mocks): "unique constraint violation", ], ) -def test_create_access_group_race_condition_returns_409( - client_and_mocks, error_message -): +def test_create_access_group_race_condition_returns_409(client_and_mocks, error_message): """Create race condition: Prisma unique constraint surfaces as 409, not 500.""" client, _, mock_table, *_ = client_and_mocks @@ -288,9 +284,7 @@ def test_create_access_group_500_on_non_constraint_prisma_error(client_and_mocks # Use raise_server_exceptions=False so unhandled exceptions become 500 responses test_client = TestClient(app, raise_server_exceptions=False) - resp = test_client.post( - "/v1/access_group", json={"access_group_name": "test-group"} - ) + resp = test_client.post("/v1/access_group", json={"access_group_name": "test-group"}) assert resp.status_code == 500 @@ -558,9 +552,7 @@ def test_update_access_group_empty_body(client_and_mocks): """Update with empty body succeeds; only updated_by is set.""" client, _, mock_table, *_ = client_and_mocks - existing = _make_access_group_record( - access_group_id="ag-update", access_group_name="unchanged" - ) + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="unchanged") mock_table.find_unique = AsyncMock(return_value=existing) resp = client.put("/v1/access_group/ag-update", json={}) @@ -576,14 +568,10 @@ def test_update_access_group_name_success(client_and_mocks): """Update access_group_name succeeds when new name is unique.""" client, _, mock_table, *_ = client_and_mocks - existing = _make_access_group_record( - access_group_id="ag-update", access_group_name="old-name" - ) + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="old-name") mock_table.find_unique = AsyncMock(return_value=existing) - resp = client.put( - "/v1/access_group/ag-update", json={"access_group_name": "new-name"} - ) + resp = client.put("/v1/access_group/ag-update", json={"access_group_name": "new-name"}) assert resp.status_code == 200 mock_table.update.assert_awaited_once() call_kwargs = mock_table.update.call_args.kwargs @@ -594,19 +582,13 @@ def test_update_access_group_name_duplicate_conflict(client_and_mocks): """Update access_group_name to existing name returns 409 (unique constraint).""" client, _, mock_table, *_ = client_and_mocks - existing = _make_access_group_record( - access_group_id="ag-update", access_group_name="old-name" - ) + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="old-name") mock_table.find_unique = AsyncMock(return_value=existing) mock_table.update = AsyncMock( - side_effect=Exception( - "Unique constraint failed on the fields: (`access_group_name`)" - ) + side_effect=Exception("Unique constraint failed on the fields: (`access_group_name`)") ) - resp = client.put( - "/v1/access_group/ag-update", json={"access_group_name": "taken-name"} - ) + resp = client.put("/v1/access_group/ag-update", json={"access_group_name": "taken-name"}) assert resp.status_code == 409 assert "already exists" in resp.json()["detail"] mock_table.update.assert_awaited_once() @@ -620,21 +602,15 @@ def test_update_access_group_name_duplicate_conflict(client_and_mocks): "unique constraint violation", ], ) -def test_update_access_group_name_unique_constraint_returns_409( - client_and_mocks, error_message -): +def test_update_access_group_name_unique_constraint_returns_409(client_and_mocks, error_message): """Update access_group_name: Prisma unique constraint surfaces as 409.""" client, _, mock_table, *_ = client_and_mocks - existing = _make_access_group_record( - access_group_id="ag-update", access_group_name="old-name" - ) + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="old-name") mock_table.find_unique = AsyncMock(return_value=existing) mock_table.update = AsyncMock(side_effect=Exception(error_message)) - resp = client.put( - "/v1/access_group/ag-update", json={"access_group_name": "race-name"} - ) + resp = client.put("/v1/access_group/ag-update", json={"access_group_name": "race-name"}) assert resp.status_code == 409 assert "already exists" in resp.json()["detail"] @@ -690,9 +666,7 @@ def test_delete_access_group_forbidden_non_admin(client_and_mocks, user_role): def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks): """Delete removes access_group_id from teams and keys before deleting the group.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable mock_key_table = mock_prisma.db.litellm_verificationtoken @@ -722,10 +696,61 @@ def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks): where={"token": "key-token-1"}, data={"access_group_ids": []}, ) - mock_access_group_table.delete.assert_awaited_once_with( - where={"access_group_id": "ag-to-delete"} + mock_access_group_table.delete.assert_awaited_once_with(where={"access_group_id": "ag-to-delete"}) + + +def test_delete_access_group_detaches_group_from_agents(client_and_mocks): + """Delete strips the group from every agent that had it attached, so agents are not left + pointing at a group that no longer exists (which would deny them every model, server and agent).""" + client, mock_prisma, mock_access_group_table, _mock_cache, _mock_proxy_logging = client_and_mocks + mock_agents_table = mock_prisma.db.litellm_agentstable + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + agent_with_group = MagicMock() + agent_with_group.agent_id = "agent-1" + agent_with_group.access_group_ids = ["ag-keep", "ag-to-delete"] + mock_agents_table.find_many = AsyncMock(return_value=[agent_with_group]) + global_agent_registry.register_agent( + AgentResponse( + agent_id="agent-1", + agent_name="detach-test-agent", + agent_card_params={"name": "detach-test-agent", "url": "http://localhost:9", "version": "1"}, + access_group_ids=["ag-keep", "ag-to-delete"], + ) ) + try: + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + mock_agents_table.update.assert_awaited_once_with( + where={"agent_id": "agent-1"}, + data={"access_group_ids": ("ag-keep",)}, + ) + mock_access_group_table.delete.assert_awaited_once_with(where={"access_group_id": "ag-to-delete"}) + registered = global_agent_registry.get_agent_by_id("agent-1") + assert registered is not None + assert tuple(registered.access_group_ids or ()) == ("ag-keep",) + finally: + global_agent_registry.deregister_agent("detach-test-agent") + + +def test_delete_access_group_without_attached_agents_leaves_agents_untouched(client_and_mocks): + client, mock_prisma, mock_access_group_table, _mock_cache, _mock_proxy_logging = client_and_mocks + mock_agents_table = mock_prisma.db.litellm_agentstable + + mock_access_group_table.find_unique = AsyncMock( + return_value=_make_access_group_record(access_group_id="ag-to-delete") + ) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + mock_agents_table.find_many.assert_awaited_once_with(where={"access_group_ids": {"hasSome": ("ag-to-delete",)}}) + mock_agents_table.update.assert_not_awaited() + @pytest.mark.parametrize( "team_cache_group_ids,key_cache_group_ids,expected_team_ids_after,expected_key_ids_after", @@ -792,9 +817,7 @@ def test_delete_access_group_patches_cached_team_and_key( """Delete patches cached team/key objects to remove the deleted access_group_id.""" from litellm.proxy._types import LiteLLM_TeamTableCachedObj - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable mock_key_table = mock_prisma.db.litellm_verificationtoken @@ -820,13 +843,9 @@ def test_delete_access_group_patches_cached_team_and_key( team_id="team-1", access_group_ids=list(team_cache_group_ids), ) - mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( - return_value=cached_team - ) + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=cached_team) else: - mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( - return_value=None - ) + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None) # user_api_key_cache is queried both for teams (fallback after dual_cache) and # hashed keys — return the right stub per ``key``. A single AsyncMock(return_value=key) @@ -834,9 +853,7 @@ def test_delete_access_group_patches_cached_team_and_key( # Use a synchronous side_effect (not async def): AsyncMock awaits coroutine side_effects # inconsistently across Python/unittest versions; sync returns are awaited as immediate results. def user_cache_get_side_effect(*args, **kwargs): - cache_key = ( - kwargs.get("key") if "key" in kwargs else (args[0] if args else None) - ) + cache_key = kwargs.get("key") if "key" in kwargs else (args[0] if args else None) if cache_key == "team_id:team-1": if team_cache_group_ids is None: return None @@ -868,14 +885,11 @@ def test_delete_access_group_patches_cached_team_and_key( team_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "team_id:team-1" - or (len(c.args) >= 1 and c.args[0] == "team_id:team-1") + if c.kwargs.get("key", "") == "team_id:team-1" or (len(c.args) >= 1 and c.args[0] == "team_id:team-1") ] assert len(team_set_calls) >= 1, "Expected team cache to be patched" # The cached team object should have the updated access_group_ids - written_team = ( - team_set_calls[0].kwargs.get("value") or team_set_calls[0].args[1] - ) + written_team = team_set_calls[0].kwargs.get("value") or team_set_calls[0].args[1] if isinstance(written_team, LiteLLM_TeamTableCachedObj): assert written_team.access_group_ids == expected_team_ids_after else: @@ -883,8 +897,7 @@ def test_delete_access_group_patches_cached_team_and_key( team_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "team_id:team-1" - or (len(c.args) >= 1 and c.args[0] == "team_id:team-1") + if c.kwargs.get("key", "") == "team_id:team-1" or (len(c.args) >= 1 and c.args[0] == "team_id:team-1") ] assert len(team_set_calls) == 0, "Should not patch team cache when not cached" @@ -892,8 +905,7 @@ def test_delete_access_group_patches_cached_team_and_key( key_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "hashed-key-1" - or (len(c.args) >= 1 and c.args[0] == "hashed-key-1") + if c.kwargs.get("key", "") == "hashed-key-1" or (len(c.args) >= 1 and c.args[0] == "hashed-key-1") ] assert len(key_set_calls) >= 1, "Expected key cache to be patched" written_key = key_set_calls[0].kwargs.get("value") or key_set_calls[0].args[1] @@ -903,17 +915,14 @@ def test_delete_access_group_patches_cached_team_and_key( key_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "hashed-key-1" - or (len(c.args) >= 1 and c.args[0] == "hashed-key-1") + if c.kwargs.get("key", "") == "hashed-key-1" or (len(c.args) >= 1 and c.args[0] == "hashed-key-1") ] assert len(key_set_calls) == 0, "Should not patch key cache when not cached" def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): """Delete patches key cache — mock returns UserAPIKeyAuth (what UserApiKeyCache emits after deserialize).""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable mock_key_table = mock_prisma.db.litellm_verificationtoken @@ -929,9 +938,7 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): mock_key_table.find_unique = AsyncMock(return_value=key_with_group) # No team in cache - mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( - return_value=None - ) + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None) # Serialized shape from Redis dict; UserApiKeyCache.async_get_cache(model_type=...) yields a model — simulate that. cached_key_payload = { @@ -940,18 +947,14 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): } def user_cache_get_dict_when_key_matches(*args, **kwargs): - cache_key = ( - kwargs.get("key") if "key" in kwargs else (args[0] if args else None) - ) + cache_key = kwargs.get("key") if "key" in kwargs else (args[0] if args else None) if cache_key == "team_id:team-1": return None if cache_key == "hashed-key-dict": return UserAPIKeyAuth.model_validate(cached_key_payload) return None - mock_cache.async_get_cache = AsyncMock( - side_effect=user_cache_get_dict_when_key_matches - ) + mock_cache.async_get_cache = AsyncMock(side_effect=user_cache_get_dict_when_key_matches) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 204 @@ -960,8 +963,7 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): key_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "hashed-key-dict" - or (len(c.args) >= 1 and c.args[0] == "hashed-key-dict") + if c.kwargs.get("key", "") == "hashed-key-dict" or (len(c.args) >= 1 and c.args[0] == "hashed-key-dict") ] assert len(key_set_calls) >= 1, "Expected key cache to be patched" written_key = key_set_calls[0].kwargs.get("value") or key_set_calls[0].args[1] @@ -988,9 +990,7 @@ def test_delete_access_group_404_on_p2025_or_record_not_found(client_and_mocks): existing = _make_access_group_record(access_group_id="ag-to-delete") mock_table.find_unique = AsyncMock(return_value=existing) - mock_table.delete = AsyncMock( - side_effect=Exception("P2025: Record to delete does not exist") - ) + mock_table.delete = AsyncMock(side_effect=Exception("P2025: Record to delete does not exist")) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 404 @@ -1039,9 +1039,7 @@ def test_delete_access_group_500_on_generic_exception(client_and_mocks): ("delete", "/v1/unified_access_group/ag-123", lambda: {}), ], ) -def test_access_group_endpoints_db_not_connected( - client_and_mocks, monkeypatch, method, url, factory -): +def test_access_group_endpoints_db_not_connected(client_and_mocks, monkeypatch, method, url, factory): """All endpoints return 500 when DB is not connected.""" client, *_ = client_and_mocks @@ -1049,9 +1047,7 @@ def test_access_group_endpoints_db_not_connected( resp = getattr(client, method)(url, **factory()) assert resp.status_code == 500 - assert ( - resp.json()["detail"]["error"] == CommonProxyErrors.db_not_connected_error.value - ) + assert resp.json()["detail"]["error"] == CommonProxyErrors.db_not_connected_error.value # --------------------------------------------------------------------------- @@ -1107,9 +1103,7 @@ def test_attached_team_ids_by_group_keeps_column_order_then_appends_unmirrored_t def test_create_access_group_syncs_assigned_teams(client_and_mocks): """Create adds access_group_id to each assigned team's access_group_ids in DB.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable team_record = _make_team_record("team-1") @@ -1132,9 +1126,7 @@ def test_create_access_group_syncs_assigned_teams(client_and_mocks): def test_create_access_group_syncs_assigned_keys(client_and_mocks): """Create adds access_group_id to each assigned key's access_group_ids in DB.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_key_table = mock_prisma.db.litellm_verificationtoken key_record = MagicMock() @@ -1148,9 +1140,7 @@ def test_create_access_group_syncs_assigned_keys(client_and_mocks): ) assert resp.status_code == 201 - mock_key_table.find_unique.assert_awaited_once_with( - where={"token": "hashed-token-1"} - ) + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "hashed-token-1"}) mock_key_table.update.assert_awaited_once() call_kwargs = mock_key_table.update.call_args.kwargs assert call_kwargs["where"] == {"token": "hashed-token-1"} @@ -1200,14 +1190,10 @@ def test_create_access_group_idempotent_team_sync(client_and_mocks): def test_update_access_group_syncs_added_teams(client_and_mocks): """Update adds access_group_id to newly assigned teams.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - existing = _make_access_group_record( - access_group_id="ag-update", assigned_team_ids=["team-existing"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-existing"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) team_record = _make_team_record("team-new") @@ -1248,14 +1234,10 @@ def test_update_access_group_rejects_nonexistent_team(client_and_mocks): def test_update_access_group_syncs_removed_teams(client_and_mocks): """Update removes access_group_id from de-assigned teams.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - existing = _make_access_group_record( - access_group_id="ag-update", assigned_team_ids=["team-keep", "team-remove"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-keep", "team-remove"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) team_to_remove = _make_team_record("team-remove", ["ag-update"]) @@ -1268,9 +1250,7 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks): ) assert resp.status_code == 200 - mock_team_table.find_unique.assert_awaited_once_with( - where={"team_id": "team-remove"} - ) + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-remove"}) mock_team_table.update.assert_awaited_once() call_kwargs = mock_team_table.update.call_args.kwargs assert call_kwargs["where"] == {"team_id": "team-remove"} @@ -1296,19 +1276,15 @@ def test_update_access_group_detaches_team_the_mirror_missed(client_and_mocks): mock_team_table.update.assert_awaited_once() call_kwargs = mock_team_table.update.call_args.kwargs assert call_kwargs["where"] == {"team_id": "team-unmirrored"} - assert call_kwargs["data"]["access_group_ids"] == ["ag-other"] + assert tuple(call_kwargs["data"]["access_group_ids"]) == ("ag-other",) def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_mocks): """Update does not sync teams when assigned_team_ids is absent from the payload.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - existing = _make_access_group_record( - access_group_id="ag-update", assigned_team_ids=["team-1"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-1"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) resp = client.put("/v1/access_group/ag-update", json={"description": "new desc"}) @@ -1320,14 +1296,10 @@ def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_moc def test_update_access_group_syncs_added_keys(client_and_mocks): """Update adds access_group_id to newly assigned keys.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_key_table = mock_prisma.db.litellm_verificationtoken - existing = _make_access_group_record( - access_group_id="ag-update", assigned_key_ids=["old-token"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_key_ids=["old-token"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) key_record = MagicMock() @@ -1350,14 +1322,10 @@ def test_update_access_group_syncs_added_keys(client_and_mocks): def test_update_access_group_syncs_removed_keys(client_and_mocks): """Update removes access_group_id from de-assigned keys.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_key_table = mock_prisma.db.litellm_verificationtoken - existing = _make_access_group_record( - access_group_id="ag-update", assigned_key_ids=["keep-token", "remove-token"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_key_ids=["keep-token", "remove-token"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) key_to_remove = MagicMock() @@ -1385,9 +1353,7 @@ def test_update_access_group_syncs_removed_keys(client_and_mocks): def test_delete_access_group_handles_out_of_sync_assigned_teams(client_and_mocks): """Delete includes teams from assigned_team_ids even when not found by hasSome query.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable # Access group has assigned_team_ids but the team's access_group_ids is not synced @@ -1409,18 +1375,14 @@ def test_delete_access_group_handles_out_of_sync_assigned_teams(client_and_mocks assert resp.status_code == 204 # find_unique is called for the out-of-sync team (included via union with assigned_team_ids) - mock_team_table.find_unique.assert_awaited_once_with( - where={"team_id": "team-out-of-sync"} - ) + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-out-of-sync"}) # No update needed since team's access_group_ids doesn't contain "ag-to-delete" mock_team_table.update.assert_not_awaited() def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks): """Delete includes keys from assigned_key_ids even when not found by hasSome query.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_key_table = mock_prisma.db.litellm_verificationtoken existing = _make_access_group_record( @@ -1439,9 +1401,7 @@ def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 204 - mock_key_table.find_unique.assert_awaited_once_with( - where={"token": "token-out-of-sync"} - ) + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "token-out-of-sync"}) mock_key_table.update.assert_not_awaited() @@ -1536,10 +1496,16 @@ def test_list_access_groups_resolves_names_with_one_query_per_table(client_and_m mock_table.find_many = AsyncMock( return_value=[ _make_access_group_record( - access_group_id="ag-1", access_mcp_server_ids=["mcp-a"], access_agent_ids=["agent-a"], assigned_key_ids=["key-a"] + access_group_id="ag-1", + access_mcp_server_ids=["mcp-a"], + access_agent_ids=["agent-a"], + assigned_key_ids=["key-a"], ), _make_access_group_record( - access_group_id="ag-2", access_mcp_server_ids=["mcp-b"], access_agent_ids=["agent-b"], assigned_key_ids=["key-b"] + access_group_id="ag-2", + access_mcp_server_ids=["mcp-b"], + access_agent_ids=["agent-b"], + assigned_key_ids=["key-b"], ), ] ) @@ -1573,7 +1539,10 @@ def test_list_access_groups_skips_lookups_when_nothing_to_resolve(client_and_moc """Groups with no MCP servers, agents or keys must not trigger an empty IN () query per table.""" client, mock_prisma, mock_table, *_ = client_and_mocks mock_table.find_many = AsyncMock( - return_value=[_make_access_group_record(access_group_id="ag-1"), _make_access_group_record(access_group_id="ag-2")] + return_value=[ + _make_access_group_record(access_group_id="ag-1"), + _make_access_group_record(access_group_id="ag-2"), + ] ) resp = client.get("/v1/access_group") diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index c66d490f122..47540eb5d6d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -7878,3 +7878,25 @@ class TestDeleteMCPGatewaySessions: assert result.terminated_sessions == 2 assert {s.user_id for s in result.sessions} == {"bob"} assert "sk-live-bob" not in result.model_dump_json() + + +class TestGetMcpToolsWireShape: + @pytest.mark.asyncio + async def test_get_mcp_tools_returns_each_tool_in_mcp_wire_spelling(self): + from mcp.types import ListToolsResult, Tool + + add_schema = {"type": "object", "properties": {"a": {"type": "integer"}}, "required": ["a"]} + listed = ListToolsResult( + tools=[Tool(name="add", description="Add", inputSchema=add_schema, outputSchema={"type": "integer"})] + ) + with patch( + "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + AsyncMock(return_value=listed), + ): + result = await mgmt_endpoints.get_mcp_tools(user_api_key_dict=generate_mock_user_api_key_auth()) + + (tool,) = result["tools"] + assert tool["inputSchema"] == add_schema + assert tool["outputSchema"] == {"type": "integer"} + assert "_meta" in tool + assert not {"input_schema", "output_schema", "meta"} & tool.keys() diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 5e6c37c41dd..bd252169131 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -4063,6 +4063,294 @@ class TestModelInfoServerDerivedPricingFilter: assert written["access_groups"] == ["prod"] +class TestModelInfoCostMapEchoFilter: + """LIT-5534. ``/model/info`` fills a deployment's ``model_info`` from the cost map (context + limits, mode, provider, supported params, capability flags), and the Admin UI edit form sends + that whole blob back on any save. Only values that still equal the cost-map entry are the + echo; a value the operator changed is a real override and stays.""" + + def test_echoed_cost_map_metadata_is_not_persisted(self): + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + entry = litellm.get_model_info("openai/gpt-5.6") + echo = {**entry, "id": "dep-echo-0", "db_model": True, "access_groups": ["prod"]} + db_model = Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model="openai/gpt-5.6"), + model_info=ModelInfo(id="dep-echo-0"), + ) + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**echo)), + ) + + info = json.loads(result["model_info"]) + assert info["access_groups"] == ["prod"] + assert set(info).isdisjoint(entry) + assert "max_input_tokens" not in info and "mode" not in info and "supports_vision" not in info, ( + "cost-map metadata must not be persisted from an unchanged /model/info echo" + ) + + def test_an_edited_value_survives_the_echo_filter(self): + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + entry = litellm.get_model_info("openai/gpt-5.6") + echo = { + **entry, + "id": "dep-echo-1", + "db_model": True, + "access_groups": ["prod"], + "max_input_tokens": entry["max_input_tokens"] + 1, + "mode": "completion" if entry["mode"] != "completion" else "chat", + } + db_model = Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model="openai/gpt-5.6"), + model_info=ModelInfo(id="dep-echo-1"), + ) + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**echo)), + ) + + info = json.loads(result["model_info"]) + assert info["max_input_tokens"] == echo["max_input_tokens"] + assert info["mode"] == echo["mode"] + assert "litellm_provider" not in info + assert "supported_openai_params" not in info + + def test_metadata_without_a_cost_map_key_is_persisted(self): + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + from litellm.types.utils import echoed_cost_map_fields + + entry = litellm.get_model_info("openai/gpt-5.6") + assert echoed_cost_map_fields({"max_input_tokens": entry["max_input_tokens"]}, entry) == () + db_model = Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model="openai/gpt-5.6"), + model_info=ModelInfo(id="dep-echo-2"), + ) + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment( + model_info=ModelInfo( + id="dep-echo-2", + max_input_tokens=entry["max_input_tokens"], + mode=entry["mode"], + ) + ), + ) + + info = json.loads(result["model_info"]) + assert info["max_input_tokens"] == entry["max_input_tokens"] + assert info["mode"] == entry["mode"] + + def test_a_stored_mode_survives_an_echoed_save(self): + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + entry = litellm.get_model_info("openai/gpt-5.6") + db_model = Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model="openai/gpt-5.6"), + model_info=ModelInfo(id="dep-echo-3", mode=entry["mode"]), + ) + echo = {**entry, "id": "dep-echo-3", "db_model": True, "access_groups": ["prod"]} + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**echo)), + ) + + info = json.loads(result["model_info"]) + assert info["mode"] == entry["mode"] + assert "max_input_tokens" not in info + + def test_resetting_an_override_to_the_cost_map_value_removes_it(self): + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + entry = litellm.get_model_info("openai/gpt-5.6") + db_model = Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model="openai/gpt-5.6"), + model_info=ModelInfo(id="dep-echo-4", mode="chat", max_input_tokens=2048), + ) + echo = {**entry, "id": "dep-echo-4", "db_model": True, "access_groups": ["staging"]} + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**echo)), + ) + + info = json.loads(result["model_info"]) + assert "max_input_tokens" not in info + assert info["mode"] == "chat" + assert info["access_groups"] == ["staging"] + + def test_reset_is_recognised_after_the_router_registered_the_override(self, monkeypatch: pytest.MonkeyPatch): + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + pristine = litellm.get_model_info("openai/gpt-5.6") + polluted = {**pristine, "max_input_tokens": 2048} + monkeypatch.setattr(litellm, "get_model_info", lambda model, **_: polluted) + db_model = Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model="openai/gpt-5.6"), + model_info=ModelInfo(id="dep-echo-8", max_input_tokens=2048), + ) + echo = {**pristine, "id": "dep-echo-8", "db_model": True} + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**echo)), + ) + + info = json.loads(result["model_info"]) + assert "max_input_tokens" not in info, info + + def test_reset_to_a_remote_catalog_value_that_differs_from_the_bundled_one(self, monkeypatch: pytest.MonkeyPatch): + from types import MappingProxyType + + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + bundled = litellm.get_model_info("openai/gpt-5.6") + remote = {**bundled, "max_input_tokens": bundled["max_input_tokens"] + 1} + remote_catalog = MappingProxyType({remote["key"]: MappingProxyType(remote)}) + monkeypatch.setattr(litellm, "get_model_info", lambda model, **_: {**remote, "max_input_tokens": 2048}) + db_model = Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model="openai/gpt-5.6"), + model_info=ModelInfo(id="dep-echo-9", max_input_tokens=2048), + ) + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**{**remote, "id": "dep-echo-9", "db_model": True})), + loaded_catalog=lambda: remote_catalog, + ) + + info = json.loads(result["model_info"]) + assert "max_input_tokens" not in info, info + + def test_echo_is_compared_against_the_deployments_lookup_not_the_key(self): + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + lookup_pairs: Final = ( + ("openai/gpt-5.6", "gpt-5.6"), + ("openai/gpt-4.1-mini", "gpt-4.1-mini"), + ) + lookup_data: Final = tuple( + (deployment_model, deployment_entry, differing_fields) + for deployment_model, key_model in lookup_pairs + for deployment_entry in (litellm.get_model_info(deployment_model),) + for key_entry in (litellm.get_model_info(key_model),) + for differing_fields in ( + frozenset( + k for k in deployment_entry if k in key_entry and deployment_entry[k] != key_entry[k] + ), + ) + if differing_fields + ) + if not lookup_data: + pytest.skip("No deployment/key cost-map lookup differences are available") + + deployment_model, entry, differing_fields = lookup_data[0] + assert differing_fields + db_model = Deployment( + model_name=deployment_model, + litellm_params=LiteLLM_Params(model=deployment_model), + model_info=ModelInfo(id="dep-echo-5"), + ) + echo = {**entry, "id": "dep-echo-5", "db_model": True, "access_groups": ["prod"]} + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**echo)), + ) + + info = json.loads(result["model_info"]) + assert not frozenset(info).intersection(frozenset(entry) - frozenset(("mode",))) + + def test_base_model_wins_over_litellm_params_model_for_the_lookup(self): + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + entry = litellm.get_model_info("azure/gpt-5.6") + db_model = Deployment( + model_name="azure/my-deploy", + litellm_params=LiteLLM_Params(model="azure/my-deploy"), + model_info=ModelInfo(id="dep-echo-6", base_model="azure/gpt-5.6"), + ) + echo = { + **entry, + "id": "dep-echo-6", + "base_model": "azure/gpt-5.6", + "db_model": True, + } + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**echo)), + ) + + info = json.loads(result["model_info"]) + assert not frozenset(info).intersection(frozenset(entry) - frozenset(("mode",))) + assert info["base_model"] == "azure/gpt-5.6" + + def test_encrypted_stored_model_is_decrypted_for_the_lookup(self, monkeypatch): + import litellm + + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") + entry = litellm.get_model_info("openai/gpt-5.6") + db_model = Deployment( + model_name="gpt-5.6", + litellm_params=LiteLLM_Params(model=encrypt_value_helper(value="openai/gpt-5.6")), + model_info=ModelInfo(id="dep-echo-7", mode="chat"), + ) + echo = {**entry, "id": "dep-echo-7", "db_model": True, "access_groups": ["prod"]} + + result = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**echo)), + ) + + info = json.loads(result["model_info"]) + assert not frozenset(info).intersection(frozenset(entry) - frozenset(("mode",))) + assert info["mode"] == "chat" + assert info["access_groups"] == ["prod"] + + class TestUpdateDBModelClearCacheControlInjectionPoints: def test_explicit_null_removes_stored_injection_points(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( diff --git a/tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py b/tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py new file mode 100644 index 00000000000..f37a20dff8b --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py @@ -0,0 +1,349 @@ +""" +Tests for BudgetReservationReleaseMiddleware. + +Auth reserves budget before the handler runs and hands the reservation to the +request or socket state. A litellm call made through the async client wrapper +claims it for the cost callback that runs after the call; anything still unclaimed +when the response is done or the socket has closed would keep the spend counter +pinned until its TTL, so the middleware releases it. +""" + +import asyncio +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from datetime import datetime +from typing import Final + +import pytest +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, Response, StreamingResponse +from starlette.routing import Route +from starlette.types import ASGIApp, Message, Receive, Scope, Send +from starlette.websockets import WebSocket + +import litellm +from litellm.caching import DualCache +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.middleware.budget_reservation_release_middleware import ( + BudgetReservationReleaseMiddleware, +) +from litellm.proxy.spend_tracking.budget_reservation import ( + reconcile_budget_reservation, + release_unbound_budget_reservation, + reserve_budget_for_request, +) +from litellm.proxy.utils import ProxyLogging +from litellm.utils import Rules, function_setup + +KEY_TOKEN: Final = "hashed-release-middleware-key" +COUNTER_KEY: Final = f"spend:key:{KEY_TOKEN}" +CHAT_BODY: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + + +@pytest.fixture +def spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> DualCache: + cache: Final = DualCache() + monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", None) + return cache + + +@pytest.fixture +def no_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + for callback_list_name in ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, callback_list_name, []) + + +async def _reserve() -> dict: + reservation: Final = await reserve_budget_for_request( + request_body=CHAT_BODY, + route="/v1/chat/completions", + llm_router=None, + valid_token=UserAPIKeyAuth(token=KEY_TOKEN, max_budget=1.0, spend=0.0), + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + ) + assert reservation is not None + assert reservation["reserved_cost"] > 0 + return reservation + + +async def _chat(reservation: dict, **kwargs: object) -> object: + return await litellm.acompletion( + **CHAT_BODY, + metadata={"user_api_key_budget_reservation": reservation}, + **kwargs, + ) + + +def _proxy_pre_call_setup(route_type: str, reservation: dict) -> None: + function_setup( + original_function=route_type, + rules_obj=Rules(), + start_time=datetime.now(), + **CHAT_BODY, + litellm_call_id="proxy-pre-call-setup", + metadata={"user_api_key_budget_reservation": reservation}, + ) + + +def _app( + handler: Callable[[Request], Awaitable[Response]], + release: Callable[[Mapping[str, object]], Awaitable[None]] = release_unbound_budget_reservation, +) -> Starlette: + app: Final = Starlette(routes=[Route("/", handler, methods=["POST"])]) + app.add_middleware(BudgetReservationReleaseMiddleware, release=release) + return app + + +async def _post(app: ASGIApp) -> None: + scope: Final = { + "type": "http", + "method": "POST", + "path": "/", + "raw_path": b"/", + "headers": [], + "query_string": b"", + "scheme": "http", + "server": ("testserver", 80), + "client": ("testclient", 1), + } + + body_delivered: Final = asyncio.Event() + client_never_disconnects: Final = asyncio.Event() + + async def receive() -> Message: + if body_delivered.is_set(): + await client_never_disconnects.wait() + body_delivered.set() + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: Message) -> None: + return None + + await app(scope, receive, send) + + +def _counter(spend_counter_cache: DualCache) -> float | None: + return spend_counter_cache.in_memory_cache.get_cache(key=COUNTER_KEY) + + +@pytest.mark.asyncio +async def test_unbound_reservation_is_released_after_the_response(spend_counter_cache: DualCache): + reservation: Final = await _reserve() + assert _counter(spend_counter_cache) == pytest.approx(reservation["reserved_cost"]) + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + return JSONResponse({"id": "batch_123", "status": "cancelling"}) + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_unbound_reservation_is_released_when_the_handler_raises(spend_counter_cache: DualCache): + reservation: Final = await _reserve() + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + raise RuntimeError("upstream refused the cancel") + + with pytest.raises(RuntimeError): + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reservation_seen_only_by_the_proxy_pre_call_logging_object_is_released( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + + async def cancel_batch_without_a_client_wrapper() -> dict: + return {"id": "batch_123", "status": "cancelling"} + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acancel_batch", reservation) + return JSONResponse(await cancel_batch_without_a_client_wrapper()) + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reservation_of_a_failed_call_is_released_after_the_error_response( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + refused: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o") + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acompletion", reservation) + try: + await _chat(reservation, mock_response=refused) + except litellm.AuthenticationError: + return JSONResponse({"error": {"message": "bad key"}}, status_code=401) + raise AssertionError("the mocked call must fail") + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_reservation_claimed_by_a_completed_call_is_left_for_the_callback( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + reserved_cost: Final = reservation["reserved_cost"] + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acompletion", reservation) + response: Final = await _chat(reservation, mock_response="ok") + return JSONResponse(response.model_dump()) + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(reserved_cost) + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_reservation_claimed_by_a_streaming_call_is_left_for_the_callback_that_finishes_after_the_response( + spend_counter_cache: DualCache, no_callbacks: None +): + reservation: Final = await _reserve() + reserved_cost: Final = reservation["reserved_cost"] + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + _proxy_pre_call_setup("acompletion", reservation) + stream: Final = await _chat(reservation, mock_response="ok", stream=True) + + async def sse() -> AsyncIterator[bytes]: + async for chunk in stream: + yield f"data: {chunk.model_dump_json()}\n\n".encode() + yield b"data: [DONE]\n\n" + + return StreamingResponse(sse(), media_type="text/event-stream") + + await _post(_app(handler)) + + assert _counter(spend_counter_cache) == pytest.approx(reserved_cost) + assert reservation["finalized"] is False + + actual_cost: Final = reserved_cost / 4 + await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=actual_cost) + + assert _counter(spend_counter_cache) == pytest.approx(actual_cost) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_unbound_reservation_of_a_websocket_session_is_released_when_the_socket_closes( + spend_counter_cache: DualCache, +): + reservation: Final = await _reserve() + + async def listen_without_a_provider_key(scope: Scope, receive: Receive, send: Send) -> None: + websocket: Final = WebSocket(scope, receive, send) + websocket.state.budget_reservation = reservation + await websocket.close(code=1011, reason="Required 'DEEPGRAM_API_KEY' in environment") + + async def receive() -> Message: + return {"type": "websocket.connect"} + + async def send(message: Message) -> None: + return None + + middleware: Final = BudgetReservationReleaseMiddleware( + listen_without_a_provider_key, release=release_unbound_budget_reservation + ) + await middleware({"type": "websocket", "path": "/deepgram/v1/listen", "headers": []}, receive, send) + + assert _counter(spend_counter_cache) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_release_runs_once_per_request_with_the_stamped_reservation(): + released: Final = [] + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + async def release(budget_reservation: Mapping[str, object]) -> None: + released.append(budget_reservation) + + async def handler(request: Request) -> Response: + request.state.budget_reservation = reservation + return JSONResponse({}) + + await _post(_app(handler, release=release)) + + assert released == [reservation] + assert released[0] is reservation + + +@pytest.mark.asyncio +async def test_request_without_a_reservation_releases_nothing(): + released: Final = [] + + async def release(budget_reservation: Mapping[str, object]) -> None: + released.append(budget_reservation) + + async def unauthenticated(request: Request) -> Response: + return JSONResponse({}) + + async def budget_checks_skipped(request: Request) -> Response: + request.state.budget_reservation = None + return JSONResponse({}) + + await _post(_app(unauthenticated, release=release)) + await _post(_app(budget_checks_skipped, release=release)) + + assert released == [] + + +@pytest.mark.asyncio +async def test_lifespan_scopes_pass_through(): + released: Final = [] + seen: Final = [] + + async def release(budget_reservation: Mapping[str, object]) -> None: + released.append(budget_reservation) + + async def inner(scope: Scope, receive: Receive, send: Send) -> None: + seen.append(scope["type"]) + + async def receive() -> Message: + return {"type": "lifespan.startup"} + + async def send(message: Message) -> None: + return None + + middleware: Final = BudgetReservationReleaseMiddleware(inner, release=release) + await middleware({"type": "lifespan", "state": {"budget_reservation": {"reserved_cost": 1.0}}}, receive, send) + + assert seen == ["lifespan"] + assert released == [] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_fal_ai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_fal_ai_passthrough_logging_handler.py new file mode 100644 index 00000000000..1c945b9110b --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_fal_ai_passthrough_logging_handler.py @@ -0,0 +1,132 @@ +"""Fal AI pass-through: upstream URL to model extraction and resolution-keyed spend tracking.""" + +from datetime import datetime +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.fal_ai_passthrough_logging_handler import ( + FalAIPassthroughLoggingHandler, +) +from litellm.types.utils import ImageResponse + +pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map") + +UPSTREAM_URL: Final = "https://queue.fal.run/fal-ai/trellis-2" + + +def _logging_obj(call_id: str = "call-fal") -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model="unknown", + messages=[{"role": "user", "content": "passthrough"}], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id=call_id, + function_id="passthrough", + ) + + +def test_is_fal_ai_route_matches_only_the_fal_ai_provider(): + assert FalAIPassthroughLoggingHandler.is_fal_ai_route(UPSTREAM_URL, "fal_ai") is True + assert FalAIPassthroughLoggingHandler.is_fal_ai_route(UPSTREAM_URL, "deepgram") is False + assert FalAIPassthroughLoggingHandler.is_fal_ai_route(UPSTREAM_URL, None) is False + + +def test_handler_extracts_model_urls_and_resolution_keyed_cost(): + upstream_body: Final = { + "model_glb": {"url": "https://fal.media/model.glb", "content_type": "model/gltf-binary"}, + "images": [{"url": "https://fal.media/preview.png"}], + "timings": {"inference": 1.2}, + } + logging_obj: Final = _logging_obj() + expected_cost: Final = litellm.model_cost["fal_ai/fal-ai/trellis-2"]["output_cost_per_image_1536"] + + handler_result = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler( + response_body=upstream_body, + request_body={"image_url": "https://example.com/in.png", "resolution": 1536}, + logging_obj=logging_obj, + url_route=UPSTREAM_URL, + kwargs={"litellm_params": {"metadata": {}}}, + ) + + result = handler_result["result"] + assert isinstance(result, ImageResponse) + assert [image.url for image in result.data or ()] == [ + "https://fal.media/model.glb", + "https://fal.media/preview.png", + ] + assert result._hidden_params["response_cost"] == pytest.approx(expected_cost) + assert handler_result["kwargs"]["model"] == "fal-ai/trellis-2" + assert handler_result["kwargs"]["custom_llm_provider"] == "fal_ai" + assert handler_result["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handler_result["kwargs"]["litellm_params"] == {"metadata": {}} + assert logging_obj.model == "fal-ai/trellis-2" + assert logging_obj.model_call_details["model"] == "fal-ai/trellis-2" + assert logging_obj.model_call_details["custom_llm_provider"] == "fal_ai" + assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected_cost) + + +def test_handler_charges_nothing_and_names_the_model_for_queue_status_and_result_polls(): + for upstream_url in ( + "https://queue.fal.run/fal-ai/trellis-2/requests/req-1/status", + "https://queue.fal.run/fal-ai/trellis-2/requests/req-1", + ): + logging_obj: Final = _logging_obj() + handler_result = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler( + response_body={"status": "COMPLETED"}, + request_body={}, + logging_obj=logging_obj, + url_route=upstream_url, + kwargs={}, + ) + assert handler_result["kwargs"]["model"] == "fal-ai/trellis-2" + assert handler_result["kwargs"]["response_cost"] is None + assert logging_obj.model_call_details["response_cost"] is None + + +def test_handler_charges_for_queue_submit(): + handler_result = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler( + response_body={"request_id": "req-1", "status": "IN_QUEUE"}, + request_body={"image_url": "https://example.com/in.png", "resolution": 1536}, + logging_obj=_logging_obj(), + url_route="https://queue.fal.run/fal-ai/trellis-2", + kwargs={}, + ) + assert handler_result["kwargs"]["model"] == "fal-ai/trellis-2" + assert handler_result["kwargs"]["response_cost"] == pytest.approx( + litellm.model_cost["fal_ai/fal-ai/trellis-2"]["output_cost_per_image_1536"] + ) + + +def test_handler_strips_queue_base_path_prefix(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("FAL_AI_QUEUE_API_BASE", "https://gw.example/fal/queue") + handler_result = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler( + response_body={"request_id": "req-1", "status": "IN_QUEUE"}, + request_body={"image_url": "https://example.com/in.png", "resolution": 1536}, + logging_obj=_logging_obj(), + url_route="https://gw.example/fal/queue/fal-ai/trellis-2", + kwargs={}, + ) + + assert handler_result["kwargs"]["model"] == "fal-ai/trellis-2" + assert handler_result["kwargs"]["response_cost"] == pytest.approx( + litellm.model_cost["fal_ai/fal-ai/trellis-2"]["output_cost_per_image_1536"] + ) + + +def test_handler_without_url_values_returns_empty_image_response_and_no_cost(): + handler_result = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler( + response_body={"status": "COMPLETED"}, + request_body={}, + logging_obj=_logging_obj(), + url_route="https://queue.fal.run/fal-ai/no-such-model", + kwargs={}, + ) + + assert isinstance(handler_result["result"], ImageResponse) + assert not handler_result["result"].data + assert handler_result["kwargs"]["response_cost"] is None + assert handler_result["kwargs"]["model"] == "fal-ai/no-such-model" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py new file mode 100644 index 00000000000..1b83c5140ca --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py @@ -0,0 +1,400 @@ +import asyncio +import json +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + _BACKGROUND_BILLING_TASKS, + TinyFishPassthroughLoggingHandler, + is_tinyfish_agent_url, + resolve_tinyfish_agent_api_base, + resolve_tinyfish_cost_per_step, + run_id_from_sse_frames, + sse_poller_spawned, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.types.passthrough_endpoints.tinyfish import is_allowed_tinyfish_endpoint + +RUN_URL = "https://agent.tinyfish.ai/v1/automation/run" +RUN_ASYNC_URL = "https://agent.tinyfish.ai/v1/automation/run-async" + + +def _make_logging_obj() -> MagicMock: + logging_obj = MagicMock() + logging_obj.litellm_call_id = "test-call-id" + logging_obj.model_call_details = {} + return logging_obj + + +def _make_response(method: str, url: str, body: dict) -> httpx.Response: + request = httpx.Request(method, url) + return httpx.Response(200, request=request, text=json.dumps(body)) + + +class _FakeClient: + """Payload items are dicts served with status_code, or (status, dict) tuples for scripted failures.""" + + def __init__(self, payloads: list, status_code: int = 200): + self.payloads = payloads + self.status_code = status_code + self.requested_urls: list[str] = [] + + async def get(self, url: str, headers: dict) -> httpx.Response: + self.requested_urls.append(url) + item = self.payloads[min(len(self.requested_urls) - 1, len(self.payloads) - 1)] + status, payload = item if isinstance(item, tuple) else (self.status_code, item) + return httpx.Response(status, text=json.dumps(payload), request=httpx.Request("GET", url)) + + +@pytest.fixture +def tinyfish_env(monkeypatch): + monkeypatch.setenv("TINYFISH_API_KEY", "sk-tf-test") + monkeypatch.delenv("TINYFISH_COST_PER_STEP", raising=False) + monkeypatch.delenv("TINYFISH_AGENT_API_BASE", raising=False) + + +class TestCostResolution: + def test_default_rate(self, tinyfish_env): + assert resolve_tinyfish_cost_per_step() == pytest.approx(0.016) + + def test_env_override(self, tinyfish_env, monkeypatch): + monkeypatch.setenv("TINYFISH_COST_PER_STEP", "0.02") + assert resolve_tinyfish_cost_per_step() == pytest.approx(0.02) + + def test_invalid_env_falls_back_to_default(self, tinyfish_env, monkeypatch): + monkeypatch.setenv("TINYFISH_COST_PER_STEP", "free") + assert resolve_tinyfish_cost_per_step() == pytest.approx(0.016) + + +class TestBillingGate: + @pytest.mark.parametrize( + "method,url,expected", + [ + ("POST", RUN_URL, True), + ("POST", RUN_ASYNC_URL, True), + ("POST", "https://agent.tinyfish.ai/v1/automation/run-sse", True), + ("GET", "https://agent.tinyfish.ai/v1/runs", False), + ("GET", "https://agent.tinyfish.ai/v1/runs/run-123?screenshots=none", False), + ("POST", "https://agent.tinyfish.ai/v1/runs/run-123/cancel", False), + ], + ) + def test_only_run_submissions_are_billed(self, method, url, expected): + assert TinyFishPassthroughLoggingHandler.should_log_request(method, url) is expected + + def test_polling_writes_no_spend_row(self, tinyfish_env): + logging_obj = _make_logging_obj() + logging_obj.dispatch_success_handlers = AsyncMock() + poll_url = "https://agent.tinyfish.ai/v1/runs/run-123" + + asyncio.run( + PassThroughEndpointLogging().pass_through_async_success_handler( + httpx_response=_make_response("GET", poll_url, {"run_id": "run-123", "status": "RUNNING"}), + response_body={"run_id": "run-123", "status": "RUNNING"}, + logging_obj=logging_obj, + url_route=poll_url, + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={}, + passthrough_logging_payload={"url": poll_url}, + custom_llm_provider="tinyfish", + ) + ) + + logging_obj.dispatch_success_handlers.assert_not_awaited() + + +class TestBlockingRunBilling: + def _handle(self, response_body: dict, logging_obj: MagicMock): + return TinyFishPassthroughLoggingHandler.tinyfish_passthrough_handler( + httpx_response=_make_response("POST", RUN_URL, response_body), + response_body=response_body, + logging_obj=logging_obj, + url_route=RUN_URL, + result=json.dumps(response_body), + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"url": "https://scrapeme.live/shop", "goal": "extract products"}, + ) + + def test_bills_steps_times_rate(self, tinyfish_env): + logging_obj = _make_logging_obj() + run = {"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 3, "result": {"products": []}} + + handler_result = self._handle(run, logging_obj) + + assert handler_result["kwargs"]["model"] == "tinyfish/automation-run" + assert handler_result["kwargs"]["custom_llm_provider"] == "tinyfish" + assert handler_result["kwargs"]["response_cost"] == pytest.approx(0.048) + assert "standard_logging_object" in handler_result["kwargs"] + assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.048) + + def test_env_rate_override_applies(self, tinyfish_env, monkeypatch): + monkeypatch.setenv("TINYFISH_COST_PER_STEP", "0.5") + run = {"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 2} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] == pytest.approx(1.0) + + def test_failed_run_logs_without_cost(self, tinyfish_env): + run = {"run_id": "run-1", "status": "FAILED", "num_of_steps": 2, "error": {"code": "AGENT_FAILURE"}} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] is None + + def test_cancelled_run_logs_without_cost(self, tinyfish_env): + run = {"run_id": "run-1", "status": "CANCELLED", "num_of_steps": 2} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] is None + + def test_null_steps_logs_without_cost(self, tinyfish_env): + run = {"run_id": "run-1", "status": "RUNNING", "num_of_steps": None} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] is None + + def test_unexpected_error_shape_still_bills(self, tinyfish_env): + run = {"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 2, "error": {"retry_after": "5s"}} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] == pytest.approx(0.032) + + +class TestRunAsyncBilling: + def test_poll_and_log_bills_once_terminal(self, tinyfish_env): + logging_obj = _make_logging_obj() + logging_obj.dispatch_success_handlers = AsyncMock() + fake_client = _FakeClient( + payloads=[{"run_id": "run-9", "status": "COMPLETED", "num_of_steps": 4, "result": "ok"}] + ) + + asyncio.run( + TinyFishPassthroughLoggingHandler._poll_and_log( + run_id="run-9", + logging_obj=logging_obj, + result="", + start_time=datetime.now(), + cache_hit=False, + kwargs={}, + client=fake_client, + ) + ) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + awaited_kwargs = logging_obj.dispatch_success_handlers.await_args.kwargs + assert awaited_kwargs["response_cost"] == pytest.approx(0.064) + assert awaited_kwargs["model"] == "tinyfish/automation-run" + assert fake_client.requested_urls == ["https://agent.tinyfish.ai/v1/runs/run-9?screenshots=none"] + assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.064) + + def test_transient_poll_failure_keeps_polling(self, tinyfish_env): + fake_client = _FakeClient( + payloads=[(500, {}), {"run_id": "run-9", "status": "COMPLETED", "num_of_steps": 3}] + ) + + run = asyncio.run( + TinyFishPassthroughLoggingHandler._poll_until_terminal("run-9", fake_client, poll_interval_seconds=0.0) + ) + + assert run is not None + assert run["num_of_steps"] == 3 + assert len(fake_client.requested_urls) == 2 + + def test_gives_up_after_consecutive_poll_failures(self, tinyfish_env): + fake_client = _FakeClient(payloads=[(500, {})]) + + run = asyncio.run( + TinyFishPassthroughLoggingHandler._poll_until_terminal("run-9", fake_client, poll_interval_seconds=0.0) + ) + + assert run is None + assert len(fake_client.requested_urls) == 12 + + def test_traversal_run_id_is_rejected(self, tinyfish_env): + fake_client = _FakeClient(payloads=[{}]) + + run = asyncio.run(TinyFishPassthroughLoggingHandler._fetch_run("../vault/items", fake_client)) + + assert run is None + assert fake_client.requested_urls == [] + + def test_upstream_error_status_returns_none(self, tinyfish_env): + fake_client = _FakeClient(payloads=[{"error": {"code": "NOT_FOUND"}}], status_code=404) + + run = asyncio.run(TinyFishPassthroughLoggingHandler._fetch_run("run-1", fake_client)) + + assert run is None + + +class TestRunCostStatusGate: + def test_poller_bills_zero_for_terminal_failed_run(self, tinyfish_env): + logging_obj = _make_logging_obj() + logging_obj.dispatch_success_handlers = AsyncMock() + fake_client = _FakeClient(payloads=[{"run_id": "run-9", "status": "FAILED", "num_of_steps": 4}]) + + asyncio.run( + TinyFishPassthroughLoggingHandler._poll_and_log( + run_id="run-9", + logging_obj=logging_obj, + result="", + start_time=datetime.now(), + cache_hit=False, + kwargs={}, + client=fake_client, + ) + ) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + assert logging_obj.dispatch_success_handlers.await_args.kwargs["response_cost"] is None + assert len(fake_client.requested_urls) == 1 + + +class TestRunIdFromSseFrames: + def test_finds_run_id_in_first_frame(self): + frames = b'data: {"run_id": "run-7", "event": "INITIALIZED"}\n\ndata: {"run_id": "run-7", "event": "ACTION"}\n\n' + assert run_id_from_sse_frames(frames) == "run-7" + + def test_skips_frames_without_run_id(self): + frames = b': keepalive\n\ndata: not-json\n\ndata: {"event": "HEARTBEAT"}\n\n' + assert run_id_from_sse_frames(frames) is None + + +class TestStartSseRunBilling: + def test_spawns_detached_poller_that_bills_once(self, tinyfish_env): + logging_obj = _make_logging_obj() + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.model_call_details["litellm_params"] = {"metadata": {"user_api_key_hash": "hash-team-a"}} + fake_client = _FakeClient( + payloads=[{"run_id": "run-7", "status": "COMPLETED", "num_of_steps": 5, "result": "done"}] + ) + + async def _run() -> None: + tasks_before = set(_BACKGROUND_BILLING_TASKS) + TinyFishPassthroughLoggingHandler.start_sse_run_billing( + run_id="run-7", + litellm_logging_obj=logging_obj, + start_time=datetime.now(), + client=fake_client, + ) + assert sse_poller_spawned(logging_obj) + await asyncio.gather(*(_BACKGROUND_BILLING_TASKS - tasks_before)) + + asyncio.run(_run()) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + awaited_kwargs = logging_obj.dispatch_success_handlers.await_args.kwargs + assert awaited_kwargs["response_cost"] == pytest.approx(0.08) + # a missing call id makes every poller row a NULL request_id primary-key collision + assert awaited_kwargs["standard_logging_object"]["id"] == "test-call-id" + # SLO consumers (Prometheus, Langfuse) must see the caller's attribution despite the empty poller kwargs + assert awaited_kwargs["standard_logging_object"]["metadata"]["user_api_key_hash"] == "hash-team-a" + assert fake_client.requested_urls == ["https://agent.tinyfish.ai/v1/runs/run-7?screenshots=none"] + + def test_flag_defaults_to_not_spawned(self): + assert not sse_poller_spawned(_make_logging_obj()) + def test_collected_chunks_price_via_run_fetch(self, tinyfish_env): + logging_obj = _make_logging_obj() + chunks = [ + 'data: {"type": "STARTED", "run_id": "run-7", "status": "RUNNING"}', + 'data: {"type": "PROGRESS", "run_id": "run-7"}', + 'data: {"type": "COMPLETE", "run_id": "run-7", "status": "COMPLETED", "result": "done"}', + ] + fake_client = _FakeClient( + payloads=[{"run_id": "run-7", "status": "COMPLETED", "num_of_steps": 5, "result": "done"}] + ) + + payload = asyncio.run( + TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks( + litellm_logging_obj=logging_obj, + url_route="https://agent.tinyfish.ai/v1/automation/run-sse", + start_time=datetime.now(), + all_chunks=chunks, + end_time=datetime.now(), + client=fake_client, + ) + ) + + assert payload["kwargs"]["response_cost"] == pytest.approx(0.08) + assert payload["kwargs"]["model"] == "tinyfish/automation-run" + assert fake_client.requested_urls == ["https://agent.tinyfish.ai/v1/runs/run-7?screenshots=none"] + + def test_stream_without_run_id_logs_without_cost(self, tinyfish_env): + fake_client = _FakeClient(payloads=[{}]) + + payload = asyncio.run( + TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks( + litellm_logging_obj=_make_logging_obj(), + url_route="https://agent.tinyfish.ai/v1/automation/run-sse", + start_time=datetime.now(), + all_chunks=["data: not-json", ": keepalive"], + end_time=datetime.now(), + client=fake_client, + ) + ) + + assert payload["kwargs"]["response_cost"] is None + assert fake_client.requested_urls == [] + + +class TestRouteDetection: + def test_provider_tag_claims_route(self): + assert PassThroughEndpointLogging().is_tinyfish_route("https://example.com/x", "tinyfish") + + def test_agent_host_claims_route(self): + assert PassThroughEndpointLogging().is_tinyfish_route("https://agent.tinyfish.ai/v1/runs", None) + + def test_other_providers_do_not_claim(self): + assert not PassThroughEndpointLogging().is_tinyfish_route("https://api.openai.com/v1", "openai") + + def test_env_base_override_claims_route(self, monkeypatch): + monkeypatch.setenv("TINYFISH_AGENT_API_BASE", "https://agent.staging.tinyfish.ai") + assert is_tinyfish_agent_url("https://agent.staging.tinyfish.ai/v1/runs/x") + assert not is_tinyfish_agent_url("https://agent.tinyfish.ai/v1/runs/x") + + def test_schemeless_env_base_is_normalized(self, monkeypatch): + monkeypatch.setenv("TINYFISH_AGENT_API_BASE", "agent.staging.tinyfish.ai") + assert resolve_tinyfish_agent_api_base() == "https://agent.staging.tinyfish.ai" + assert is_tinyfish_agent_url("https://agent.staging.tinyfish.ai/v1/runs/x") + + +class TestEndpointAllowlist: + @pytest.mark.parametrize( + "method,path,expected", + [ + ("POST", "/v1/automation/run", True), + ("POST", "/v1/automation/run-async", True), + ("POST", "/v1/automation/run-sse", True), + ("GET", "/v1/runs", False), + ("GET", "/v1/runs/run-abc-123", True), + ("POST", "/v1/runs/run-abc-123/cancel", True), + ("GET", "/v1/vault/items", False), + ("GET", "/v1/wallet", False), + ("POST", "/v1/browser-profiles", False), + ("DELETE", "/v1/runs/run-abc-123", False), + ("GET", "/v1/automation/run", False), + ("POST", "/v1/runs", False), + ("GET", "/v1/runs/..", False), + ("POST", "/v1/runs/../automation/run/cancel", False), + ("POST", "/v1/automation/run/", False), + ("POST", "/v1/automation/run-async/", False), + ("POST", "/v1/automation/run-sse/", False), + ("POST", "/v1//automation/run-async", False), + ("GET", "/v1/runs/run-abc-123/", False), + ("POST", "/v1/runs/run-abc-123/cancel/", False), + ], + ) + def test_allowlist(self, method, path, expected): + assert is_allowed_tinyfish_endpoint(method, path) is expected diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py index 345eeeedc31..e0a5ef063e8 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_typesafe_passthrough_logging_handler.py @@ -42,6 +42,7 @@ def _handler_result(response_body: dict, request_body: dict) -> dict: end_time=datetime.now(), cache_hit=False, request_body=request_body, + custom_llm_provider="typesafe", ) @@ -59,6 +60,7 @@ def test_uses_registry_pricing_and_standard_usage(): end_time=datetime.now(), cache_hit=False, request_body={"model": "jev-latest"}, + custom_llm_provider="typesafe", ) expected_cost = 312 * model_cost["input_cost_per_token"] + 48 * model_cost["output_cost_per_token"] @@ -105,6 +107,7 @@ def test_records_model_provider_and_cost_on_logging_details(): end_time=datetime.now(), cache_hit=False, request_body={"model": "jev-latest"}, + custom_llm_provider="typesafe", ) assert result["kwargs"]["model"] == "typesafe/jev-1.13.0" @@ -132,3 +135,76 @@ def test_success_handler_dispatches_to_typesafe_handler(): assert normalized["kwargs"]["custom_llm_provider"] == "typesafe" assert normalized["kwargs"]["model"] == "typesafe/jev-1.13.0" + + +def test_openrouter_decisions_response_is_priced_from_request_model_registry_row(): + logging_obj = _logging_obj() + model_cost = litellm.model_cost["openrouter/typesafe/jev-1.13"] + response = TypeSafePassthroughLoggingHandler.typesafe_passthrough_handler( + httpx_response=_response(), + response_body={ + "model": "typesafe/jev-1.13-20260917", + "usage": {"input_tokens": 282, "output_tokens": 20}, + }, + logging_obj=logging_obj, + url_route="https://openrouter.ai/api/alpha/decisions", + result='{"answers": {}}', + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"model": "typesafe/jev-1.13"}, + custom_llm_provider="openrouter", + ) + + expected_cost = 282 * model_cost["input_cost_per_token"] + 20 * model_cost["output_cost_per_token"] + assert response["kwargs"]["model"] == "openrouter/typesafe/jev-1.13-20260917" + assert response["kwargs"]["custom_llm_provider"] == "openrouter" + assert response["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert response["kwargs"]["combined_usage_object"].prompt_tokens == 282 + assert response["kwargs"]["combined_usage_object"].completion_tokens == 20 + assert response["kwargs"]["combined_usage_object"].total_tokens == 302 + + +def test_success_handler_dispatches_openrouter_to_the_shared_handler(): + logging_obj = _logging_obj() + normalized = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( + httpx_response=_response(), + response_body={ + "model": "typesafe/jev-1.13-20260917", + "usage": {"input_tokens": 282, "output_tokens": 20}, + }, + request_body={"model": "typesafe/jev-1.13"}, + logging_obj=logging_obj, + url_route="https://openrouter.ai/api/alpha/decisions", + result="{}", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + custom_llm_provider="openrouter", + ) + + assert normalized["kwargs"]["custom_llm_provider"] == "openrouter" + assert normalized["kwargs"]["model"] == "openrouter/typesafe/jev-1.13-20260917" + + +def test_success_handler_skips_typesafe_pricing_for_non_decisions_openrouter_routes(): + logging_obj = _logging_obj() + normalized = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( + httpx_response=_response(), + response_body={ + "model": "typesafe/jev-1.13-20260917", + "usage": {"input_tokens": 282, "output_tokens": 20}, + }, + request_body={"model": "typesafe/jev-1.13"}, + logging_obj=logging_obj, + url_route="https://openrouter.ai/api/v1/chat/completions", + result="{}", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + custom_llm_provider="openrouter", + ) + + assert normalized["standard_logging_response_object"] is None + assert "combined_usage_object" not in normalized["kwargs"] + assert normalized["kwargs"].get("model") != "openrouter/typesafe/jev-1.13-20260917" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 636980eb6e3..fd81fcc8e72 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -29,6 +29,7 @@ from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( BaseOpenAIPassThroughHandler, RouteChecks, + _fal_target, _join_url_paths, _proxy_general_settings, anthropic_proxy_route, @@ -38,6 +39,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( bedrock_proxy_route, create_pass_through_route, cursor_proxy_route, + fal_ai_proxy_route, get_azure_ai_search_index_from_endpoint, get_vertex_base_url, is_azure_ai_search_service_level_index_create, @@ -47,6 +49,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( mistral_proxy_route, relay_nvidia_nim_request, openai_proxy_route, + openrouter_proxy_route, typesafe_proxy_route, vertex_discovery_proxy_route, vertex_proxy_route, @@ -7114,3 +7117,454 @@ class TestTypeSafePassthroughRoute: custom_llm_provider="typesafe", is_streaming_request=False, ) + + +class TestFalAIPassthroughRoute: + @pytest.fixture + def client(self, monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("FAL_AI_API_KEY", "fal-test-key") + monkeypatch.delenv("FAL_AI_QUEUE_API_BASE", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual")) + yield TestClient(app) + + def test_submit_forwards_body_and_key_scheme_to_queue_fal_run(self, client: TestClient) -> None: + body: Final = {"image_url": "https://example.com/in.png", "resolution": 1536} + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post("https://queue.fal.run/fal-ai/trellis-2").mock( + return_value=httpx.Response(200, json={"request_id": "req-1", "status": "IN_QUEUE"}) + ) + response = client.post("/fal_ai/fal-ai/trellis-2", json=body) + + assert response.status_code == 200, response.text + assert response.json() == {"request_id": "req-1", "status": "IN_QUEUE"} + sent = route.calls.last.request + assert sent.headers["authorization"] == "Key fal-test-key" + assert json.loads(sent.content or b"{}") == body + + def test_status_get_forwards_to_queue_fal_run(self, client: TestClient) -> None: + with respx.mock(assert_all_called=True) as upstream: + route = upstream.get("https://queue.fal.run/fal-ai/trellis-2/requests/req-1/status").mock( + return_value=httpx.Response(200, json={"status": "COMPLETED"}) + ) + response = client.get("/fal_ai/fal-ai/trellis-2/requests/req-1/status") + + assert response.status_code == 200, response.text + assert response.json() == {"status": "COMPLETED"} + assert route.calls.last.request.headers["authorization"] == "Key fal-test-key" + + def test_honours_fal_ai_queue_api_base_override(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FAL_AI_API_KEY", "fal-test-key") + monkeypatch.setenv("FAL_AI_QUEUE_API_BASE", "https://queue.example/base") + endpoint_func = AsyncMock(return_value={"ok": True}) + create_route = Mock(return_value=endpoint_func) + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + create_route, + ) + request = MagicMock(spec=Request) + request.method = "POST" + request.query_params = {} + request.json = AsyncMock(return_value={}) + + result = asyncio.run( + fal_ai_proxy_route( + endpoint="fal-ai/trellis", + request=request, + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="virtual-key"), + ) + ) + + assert result == {"ok": True} + create_route.assert_called_once_with( + endpoint="fal-ai/trellis", + target="https://queue.example/base/fal-ai/trellis", + custom_headers={"Authorization": "Key fal-test-key"}, + custom_llm_provider="fal_ai", + is_streaming_request=False, + ) + + def test_submit_to_unpriced_endpoint_returns_400_without_upstream_call(self, client: TestClient) -> None: + with respx.mock(assert_all_called=False) as upstream: + route = upstream.post("https://queue.fal.run/fal-ai/unpriced-model").mock( + return_value=httpx.Response(200, json={"request_id": "req-1"}) + ) + response = client.post("/fal_ai/fal-ai/unpriced-model", json={"image_url": "https://example.com/in.png"}) + + assert response.status_code == 400, response.text + assert "no pricing entry" in response.text + assert not route.calls + + def test_status_get_on_unpriced_endpoint_forwards(self, client: TestClient) -> None: + with respx.mock(assert_all_called=True) as upstream: + upstream.get("https://queue.fal.run/fal-ai/unpriced-model/requests/req-9/status").mock( + return_value=httpx.Response(200, json={"status": "IN_PROGRESS"}) + ) + response = client.get("/fal_ai/fal-ai/unpriced-model/requests/req-9/status") + + assert response.status_code == 200, response.text + assert response.json() == {"status": "IN_PROGRESS"} + + def test_missing_fal_key_returns_401(self, client: TestClient, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("FAL_AI_API_KEY", raising=False) + response = client.post("/fal_ai/fal-ai/trellis", json={}) + assert response.status_code == 401 + + +class TestFalTargetSelection: + def test_endpoint_targets_queue_base(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("FAL_AI_QUEUE_API_BASE", raising=False) + assert str(_fal_target("fal-ai/trellis-2")) == "https://queue.fal.run/fal-ai/trellis-2" + + def test_status_path_targets_queue_base(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("FAL_AI_QUEUE_API_BASE", raising=False) + assert str(_fal_target("fal-ai/trellis-2/requests/req-1/status")) == ( + "https://queue.fal.run/fal-ai/trellis-2/requests/req-1/status" + ) + + def test_queue_base_override(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FAL_AI_QUEUE_API_BASE", "https://queue.example/base") + assert str(_fal_target("fal-ai/trellis-2")) == "https://queue.example/base/fal-ai/trellis-2" + + +class TestOpenRouterPassthroughRoute: + @staticmethod + def _request(body: object, query_params: Mapping[str, str] | None = None) -> MagicMock: + request = MagicMock(spec=Request) + request.method = "POST" + request.query_params = query_params or {} + request.json = AsyncMock(return_value=body) + return request + + @pytest.fixture + def client(self, monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("OPENROUTER_API_KEY", "openrouter-test-key") + monkeypatch.setenv("OPENROUTER_API_BASE", "https://openrouter.example/base") + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual")) + yield TestClient(app) + + @pytest.mark.parametrize( + "method, body", + [ + ("GET", None), + ("POST", {"state": "The sky is blue."}), + ("PUT", {"state": "The sky is blue."}), + ("DELETE", None), + ("PATCH", {"state": "The sky is blue."}), + ], + ) + def test_forwards_every_method_and_body_upstream( + self, client: TestClient, method: str, body: dict[str, str] | None + ) -> None: + with respx.mock(assert_all_called=True) as upstream: + route = upstream.request(method, "https://openrouter.example/base/alpha/decisions").mock( + return_value=httpx.Response(200, json={"id": "upstream_123"}) + ) + response = client.request(method, "/openrouter/alpha/decisions", json=body) + + assert (response.status_code, response.json()) == (200, {"id": "upstream_123"}) + sent: Final = route.calls.last.request + assert sent.headers["authorization"] == "Bearer openrouter-test-key" + assert json.loads(sent.content or b"{}") == (body or {}) + + @pytest.mark.asyncio + async def test_forwards_target_auth_provider_and_query(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENROUTER_API_KEY", "openrouter-test-key") + monkeypatch.setenv("OPENROUTER_API_BASE", "https://openrouter.example/base") + + async def fake_upstream(request, *_args): + target: Final = create_route.call_args.kwargs["target"] + upstream_url: Final = httpx.URL(target).copy_merge_params(request.query_params) + return {"upstream_query": parse_qs(upstream_url.query.decode())} + + endpoint_func = AsyncMock(side_effect=fake_upstream) + create_route = Mock(return_value=endpoint_func) + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + create_route, + ) + + request = self._request({"state": "The sky is blue."}, {"trace": "yes"}) + result = await openrouter_proxy_route( + endpoint="alpha/decisions", + request=request, + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="virtual-key"), + ) + + assert result == {"upstream_query": {"trace": ["yes"]}} + endpoint_func.assert_awaited_once() + create_route.assert_called_once_with( + endpoint="alpha/decisions", + target="https://openrouter.example/base/alpha/decisions", + custom_headers={ + "Authorization": "Bearer openrouter-test-key", + "Content-Type": "application/json", + }, + custom_llm_provider="openrouter", + is_streaming_request=False, + ) + + @pytest.mark.asyncio + async def test_uses_default_target_when_base_is_unset(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENROUTER_API_KEY", "openrouter-test-key") + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + + endpoint_func = AsyncMock(return_value={"ok": True}) + create_route = Mock(return_value=endpoint_func) + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + create_route, + ) + + await openrouter_proxy_route( + endpoint="alpha/decisions", + request=self._request({"state": "The sky is blue."}), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="virtual-key"), + ) + + assert create_route.call_args.kwargs["target"] == "https://openrouter.ai/api/alpha/decisions" + + @pytest.mark.asyncio + @pytest.mark.parametrize("endpoint", ["alpha/decisions", "v1/chat/completions"]) + @pytest.mark.parametrize( + "base_env, expected_root", + [ + (None, "https://openrouter.ai/api"), + ("https://openrouter.ai/api/v1", "https://openrouter.ai/api"), + ("https://openrouter.example/base", "https://openrouter.example/base"), + ("https://openrouter.example/base/v1/", "https://openrouter.example/base"), + ], + ) + async def test_derives_api_root_from_configured_base( + self, monkeypatch: pytest.MonkeyPatch, base_env: str | None, expected_root: str, endpoint: str + ) -> None: + monkeypatch.setenv("OPENROUTER_API_KEY", "openrouter-test-key") + if base_env is None: + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + else: + monkeypatch.setenv("OPENROUTER_API_BASE", base_env) + + endpoint_func = AsyncMock(return_value={"ok": True}) + create_route = Mock(return_value=endpoint_func) + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + create_route, + ) + + await openrouter_proxy_route( + endpoint=endpoint, + request=self._request({"state": "The sky is blue."}), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key="virtual-key"), + ) + + assert create_route.call_args.kwargs["target"] == f"{expected_root}/{endpoint}" + + +class TestTinyFishProxyRoute: + """Tests for the TinyFish Agent pass-through route, faking the upstream HTTP boundary.""" + + RUN_BODY = {"url": "https://scrapeme.live/shop", "goal": "Extract the first 2 product names. Return JSON."} + + @pytest.fixture + def tinyfish_client(self, monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("TINYFISH_API_KEY", "sk-tf-upstream") + monkeypatch.delenv("TINYFISH_AGENT_API_BASE", raising=False) + monkeypatch.delenv("TINYFISH_ALLOW_AUTHENTICATED_RUNS", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual")) + yield TestClient(app) + + def test_forwards_run_with_server_key_not_callers(self, tinyfish_client: TestClient) -> None: + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 2}) + ) + response = tinyfish_client.post( + "/tinyfish/v1/automation/run", json=self.RUN_BODY, headers={"X-API-Key": "sk-callers-virtual-key"} + ) + + assert (response.status_code, response.json()["run_id"]) == (200, "run-1") + assert route.calls.last.request.headers["x-api-key"] == "sk-tf-upstream" + + @pytest.mark.parametrize( + "method,path", + [ + ("GET", "/tinyfish/v1/vault/items"), + ("GET", "/tinyfish/v1/wallet"), + ("POST", "/tinyfish/v1/browser-profiles"), + ("GET", "/tinyfish/v1/automation/run"), + ("GET", "/tinyfish/v1/runs"), + ], + ) + def test_blocks_endpoints_outside_allowlist(self, tinyfish_client: TestClient, method: str, path: str) -> None: + with respx.mock: + response = tinyfish_client.request(method, path) + + assert response.status_code == 403 + assert "not an allowed TinyFish Agent passthrough endpoint" in response.json()["detail"] + + @pytest.mark.parametrize( + "path", + [ + "/tinyfish/v1/automation/run/", + "/tinyfish/v1/automation/run-async/", + "/tinyfish/v1/automation/run-sse/", + "/tinyfish/v1//automation/run-async", + ], + ) + def test_submit_paths_with_extra_slashes_are_rejected_before_forwarding( + self, tinyfish_client: TestClient, path: str + ) -> None: + with respx.mock(assert_all_called=False) as upstream: + upstream.post(url__regex=r"https://agent\.tinyfish\.ai/.*").mock( + return_value=httpx.Response(200, json={"run_id": "run-slash", "status": "PENDING"}) + ) + response = tinyfish_client.post(path, json=self.RUN_BODY) + + assert response.status_code == 403 + assert "not an allowed TinyFish Agent passthrough endpoint" in response.json()["detail"] + assert upstream.calls.call_count == 0 + + def test_rejects_authenticated_run_fields_by_default(self, tinyfish_client: TestClient) -> None: + with respx.mock: + response = tinyfish_client.post("/tinyfish/v1/automation/run", json={**self.RUN_BODY, "use_vault": True}) + + assert response.status_code == 403 + assert "use_vault" in response.json()["detail"] + + def test_env_opt_in_allows_authenticated_run_fields( + self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("TINYFISH_ALLOW_AUTHENTICATED_RUNS", "true") + + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-2", "status": "COMPLETED", "num_of_steps": 1}) + ) + response = tinyfish_client.post("/tinyfish/v1/automation/run", json={**self.RUN_BODY, "use_vault": True}) + + assert response.status_code == 200 + assert json.loads(route.calls.last.request.content)["use_vault"] is True + + def test_returns_401_on_missing_api_key( + self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("TINYFISH_API_KEY") + + with respx.mock: + response = tinyfish_client.get("/tinyfish/v1/runs/run-123") + + assert response.status_code == 401 + assert "TINYFISH_API_KEY" in response.json()["detail"] + + def test_env_base_override_changes_target( + self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("TINYFISH_AGENT_API_BASE", "https://agent.staging.tinyfish.ai") + + with respx.mock(assert_all_called=True) as upstream: + upstream.get("https://agent.staging.tinyfish.ai/v1/runs/run-123").mock( + return_value=httpx.Response(200, json={"run_id": "run-123", "status": "RUNNING"}) + ) + response = tinyfish_client.get("/tinyfish/v1/runs/run-123") + + assert (response.status_code, response.json()["status"]) == (200, "RUNNING") + + @pytest.mark.parametrize( + "body", + [ + {"custom_body": {"url": "https://scrapeme.live/shop", "goal": "g", "use_vault": True}}, + {"url": "https://scrapeme.live/shop", "goal": "g", "stream": True}, + {"url": "https://scrapeme.live/shop", "goal": "g", "query_params": {"x": "1"}}, + ], + ) + def test_rejects_passthrough_envelope_controls(self, tinyfish_client: TestClient, body: dict) -> None: + """custom_body smuggled vault fields past the 403 gate and a stream flag flipped the + billing mode, because the generic passthrough honors both from the caller's body.""" + with respx.mock as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 1}) + ) + response = tinyfish_client.post("/tinyfish/v1/automation/run", json=body) + + assert response.status_code == 400 + assert "envelope" in response.json()["detail"] + assert not route.called + + def test_rejects_envelope_stream_on_cancel(self, tinyfish_client: TestClient) -> None: + with respx.mock as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/runs/run-1/cancel").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "CANCELLED"}) + ) + response = tinyfish_client.post("/tinyfish/v1/runs/run-1/cancel", json={"stream": True}) + + assert response.status_code == 400 + assert not route.called + + @pytest.mark.parametrize( + "content,content_type", + [ + ("url=https%3A%2F%2Fscrapeme.live%2Fshop&goal=g&stream=true", "application/x-www-form-urlencoded"), + ("url=https%3A%2F%2Fscrapeme.live%2Fshop&goal=g&use_vault=true", "application/x-www-form-urlencoded"), + ('{"url": "https://scrapeme.live/shop", "goal": "g", "use_vault": true}', "text/plain"), + ('[{"url": "https://scrapeme.live/shop", "goal": "g", "stream": true}]', "application/json"), + ], + ) + def test_rejects_bodies_that_are_not_json_objects( + self, tinyfish_client: TestClient, content: str, content_type: str + ) -> None: + """A form-encoded body carried stream and use_vault past both field gates, because + the gates only saw fields the body parsed to as JSON.""" + with respx.mock as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 1}) + ) + response = tinyfish_client.post( + "/tinyfish/v1/automation/run", content=content, headers={"Content-Type": content_type} + ) + + assert response.status_code == 400 + assert "JSON object" in response.json()["detail"] + assert not route.called + + def test_cancel_without_body_forwards(self, tinyfish_client: TestClient) -> None: + with respx.mock(assert_all_called=True) as upstream: + upstream.post("https://agent.tinyfish.ai/v1/runs/run-1/cancel").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "CANCELLED"}) + ) + response = tinyfish_client.post("/tinyfish/v1/runs/run-1/cancel") + + assert (response.status_code, response.json()["status"]) == (200, "CANCELLED") + + +class TestTinyFishRouteTimeout: + def test_default_covers_upstream_run_cap(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _tinyfish_route_timeout + + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + assert _tinyfish_route_timeout() == 1500.0 + + def test_operator_configured_timeout_wins(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _tinyfish_route_timeout + + monkeypatch.setattr(proxy_server, "general_settings", {"pass_through_request_timeout": 30}, raising=False) + assert _tinyfish_route_timeout() is None diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index fb89e3a6973..6ad850866b7 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4256,6 +4256,112 @@ async def test_pass_through_request_non_streaming_success_unchanged(): mock_success_handler.assert_called_once() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "upstream_status_code, claimed_by_the_success_handler", + [(200, True), (500, False)], + ids=["success-claims-the-reservation", "upstream-error-leaves-it-for-the-request-end-release"], +) +async def test_pass_through_request_claims_the_budget_reservation_only_when_its_success_handler_runs( + upstream_status_code: int, claimed_by_the_success_handler: bool +): + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed") + user_api_key_dict.budget_reservation = reservation + upstream_response: Final = httpx.Response( + status_code=upstream_status_code, + headers={"content-type": "application/json"}, + content=b'{"status": "upstream"}', + request=httpx.Request("POST", "http://target-api.com/api/generate"), + ) + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") as mock_get_client, + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, + ): + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/mock-upstream/api/generate" + mock_request.body = AsyncMock(return_value=b'{"prompt": "hi"}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + response = await pass_through_request( + request=mock_request, + target="http://target-api.com/api/generate", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) + + assert response.status_code == upstream_status_code + assert reservation["callback_bound"] is claimed_by_the_success_handler + assert mock_worker.ensure_initialized_and_enqueue.call_count == int(claimed_by_the_success_handler) + + +@pytest.mark.asyncio +async def test_pass_through_request_leaves_the_budget_reservation_for_the_request_end_release_when_its_success_handler_cannot_be_enqueued(): + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed") + user_api_key_dict.budget_reservation = reservation + upstream_response: Final = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"status": "upstream"}', + request=httpx.Request("POST", "http://target-api.com/api/generate"), + ) + + def refuse_to_enqueue(async_coroutine): + async_coroutine.close() + raise RuntimeError("logging worker is shutting down") + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") as mock_get_client, + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker, + ): + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=refuse_to_enqueue) + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/mock-upstream/api/generate" + mock_request.body = AsyncMock(return_value=b'{"prompt": "hi"}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + with pytest.raises(ProxyException): + await pass_through_request( + request=mock_request, + target="http://target-api.com/api/generate", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) + + assert reservation["callback_bound"] is False + + @pytest.mark.asyncio async def test_pass_through_request_internal_failure_still_raises_proxy_exception(): """ @@ -4651,6 +4757,90 @@ async def test_pass_through_request_upstream_error_body_stays_buffered(): await fake_client.aclose() +_UPSTREAM_JSON_ERROR: Final = b'{"error": {"message": "bad request", "type": "invalid_request_error"}}' + + +async def _relay_upstream_through_pass_through_request( + general_settings, status_code, content_type, body, callback_headers=None +): + from litellm.proxy._types import UserAPIKeyAuth + + fake_client, cleanup = _inject_fake_passthrough_client( + _FakeUpstreamTransport( + status_code=status_code, + headers={"content-type": content_type}, + stream=_RecordingUpstreamByteStream((body,)), + ), + timeout=313.0, + ) + try: + with ExitStack() as stack: + mock_proxy_logging, _ = _enter_relay_logging_mocks(stack, {}) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=callback_headers) + stack.enter_context(patch("litellm.proxy.proxy_server.general_settings", general_settings)) + return await pass_through_request( + request=_relay_client_request(), + target="http://upstream.test/v1/messages", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-relay-test"), + timeout=313.0, + ) + finally: + cleanup() + await fake_client.aclose() + + +@pytest.mark.asyncio +async def test_pass_through_error_body_carries_the_call_id_when_opted_in(): + """With include_call_id_in_error_body on, a buffered upstream JSON error gets a top-level + litellm_call_id byte-identical to the x-litellm-call-id header, and content-length still + matches the rewritten body.""" + response = await _relay_upstream_through_pass_through_request( + {"include_call_id_in_error_body": True}, 400, "application/json", _UPSTREAM_JSON_ERROR + ) + + call_id = response.headers["x-litellm-call-id"] + assert response.status_code == 400 + assert json.loads(response.body) == {**json.loads(_UPSTREAM_JSON_ERROR), "litellm_call_id": call_id} + assert int(response.headers["content-length"]) == len(response.body) + + +@pytest.mark.asyncio +async def test_pass_through_error_body_call_id_follows_a_restamped_header(): + """A post_call_response_headers_hook that rewrites x-litellm-call-id wins in the header, so the + body copies the emitted header value rather than the id the proxy generated.""" + response = await _relay_upstream_through_pass_through_request( + {"include_call_id_in_error_body": True}, + 400, + "application/json", + _UPSTREAM_JSON_ERROR, + callback_headers={"x-litellm-call-id": "restamped-by-hook"}, + ) + + assert response.headers["x-litellm-call-id"] == "restamped-by-hook" + assert json.loads(response.body)["litellm_call_id"] == "restamped-by-hook" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "general_settings, status_code, content_type, body", + [ + ({}, 400, "application/json", _UPSTREAM_JSON_ERROR), + ({"include_call_id_in_error_body": True}, 502, "text/plain", b"upstream exploded"), + ({"include_call_id_in_error_body": True}, 200, "application/json", b'{"id": "msg_1", "type": "message"}'), + ], +) +async def test_pass_through_body_stays_byte_identical_outside_the_opt_in( + general_settings, status_code, content_type, body +): + """Opted out, a non-JSON error, or a success body: the upstream bytes are relayed as-is.""" + response = await _relay_upstream_through_pass_through_request(general_settings, status_code, content_type, body) + + assert response.status_code == status_code + assert response.body == body + assert "x-litellm-call-id" in response.headers + + _PARTIAL_RELAY_WARNING_MARKER = "ended before upstream body was fully relayed" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index ea6adc35b9a..00a3606bbf5 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -7,8 +7,14 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx import litellm +import litellm.proxy.pass_through_endpoints.llm_provider_handlers.tinyfish_passthrough_logging_handler as tinyfish_handler_module +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + mark_sse_poller_spawned, + sse_poller_spawned, +) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -838,3 +844,286 @@ async def test_chunk_processor_bills_partial_google_usage_on_mid_stream_exceptio assert failure_payload["completion_tokens"] == 12 assert failure_payload["response_cost"] > 12 * 3.75e-06 assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) + + +class TestTinyFishStreamBilling: + """SSE billing is owned by the detached poller spawned on the first run_id frame; it must + survive gen.aclose() (client disconnect) and the stream-end path must not double-bill.""" + + RUNS_URL = "https://agent.tinyfish.ai/v1/runs/run-sse-1?screenshots=none" + SSE_ROUTE = "https://agent.tinyfish.ai/v1/automation/run-sse" + + @pytest.fixture + def tinyfish_env(self, monkeypatch): + monkeypatch.setenv("TINYFISH_API_KEY", "sk-tf-test") + monkeypatch.delenv("TINYFISH_COST_PER_STEP", raising=False) + monkeypatch.delenv("TINYFISH_AGENT_API_BASE", raising=False) + # aiohttp transport bypasses respx; force plain httpx and drop any cached aiohttp client + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + def _tinyfish_logging_obj(self): + obj = _unarmed_logging_obj() + obj.model_call_details = {} + obj.dispatch_success_handlers = AsyncMock() + return obj + + def _spawned_since(self, tasks_before): + return tinyfish_handler_module._BACKGROUND_BILLING_TASKS - tasks_before + + @pytest.mark.asyncio + async def test_run_id_split_across_chunks_spawns_one_poller(self, tinyfish_env): + chunks = [ + b'data: {"run_id": "run-s', + b'se-1", "event": "INITIALIZED"}\n\n', + b'data: {"run_id": "run-sse-1", "event": "COMPLETE"}\n\n', + ] + logging_obj = self._tinyfish_logging_obj() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + with respx.mock(assert_all_called=True) as upstream: + upstream.get(self.RUNS_URL).respond( + json={"run_id": "run-sse-1", "status": "COMPLETED", "num_of_steps": 2, "result": "ok"} + ) + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=_make_streaming_response(chunks), + request_body={"url": "https://scrapeme.live/shop", "goal": "extract"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + route_streaming_logging=AsyncMock(), + ): + received.append(chunk) + + assert received == chunks + assert sse_poller_spawned(logging_obj) + spawned = self._spawned_since(tasks_before) + assert len(spawned) == 1 + await asyncio.gather(*spawned) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + assert logging_obj.dispatch_success_handlers.await_args.kwargs["response_cost"] == pytest.approx(0.032) + + @pytest.mark.asyncio + async def test_poller_survives_client_disconnect_and_bills(self, tinyfish_env): + chunks = [ + b'data: {"run_id": "run-sse-1", "event": "INITIALIZED"}\n\n', + b'data: {"run_id": "run-sse-1", "event": "ACTION"}\n\n', + ] + logging_obj = self._tinyfish_logging_obj() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + with respx.mock(assert_all_called=True) as upstream: + upstream.get(self.RUNS_URL).respond( + json={"run_id": "run-sse-1", "status": "COMPLETED", "num_of_steps": 2, "result": "ok"} + ) + gen = PassThroughStreamingHandler.chunk_processor( + response=_make_streaming_response(chunks), + request_body={"url": "https://scrapeme.live/shop", "goal": "extract"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + route_streaming_logging=AsyncMock(), + ) + await gen.__anext__() + spawned = self._spawned_since(tasks_before) + assert len(spawned) == 1 + await gen.aclose() + + task = next(iter(spawned)) + assert not task.cancelled() + await task + + logging_obj.dispatch_success_handlers.assert_awaited_once() + assert logging_obj.dispatch_success_handlers.await_args.kwargs["response_cost"] == pytest.approx(0.032) + + @pytest.mark.asyncio + async def test_stream_without_run_id_spawns_nothing(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + async for _ in PassThroughStreamingHandler.chunk_processor( + response=_make_streaming_response([b": keepalive\n\n", b'data: {"event": "HEARTBEAT"}\n\n']), + request_body={}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + route_streaming_logging=AsyncMock(), + ): + pass + + assert not sse_poller_spawned(logging_obj) + assert self._spawned_since(tasks_before) == set() + + @pytest.mark.asyncio + async def test_stream_end_skips_dispatch_when_poller_owns_billing(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + mark_sse_poller_spawned(logging_obj) + + await PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + request_body={}, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + raw_bytes=[b'data: {"run_id": "run-sse-1", "event": "COMPLETE"}\n\n'], + end_time=datetime.now(), + ) + + logging_obj.dispatch_success_handlers.assert_not_awaited() + + @pytest.mark.asyncio + async def test_stream_end_fallback_still_logs_when_no_poller_spawned(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + + await PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + request_body={}, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + raw_bytes=[b": keepalive\n\n"], + end_time=datetime.now(), + ) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + assert logging_obj.dispatch_success_handlers.await_args.kwargs["response_cost"] is None + + @pytest.mark.asyncio + async def test_upstream_error_after_spawn_skips_failure_dispatch(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + logging_obj.dispatch_failure_handlers = MagicMock() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + async def _aiter_bytes(): + yield b'data: {"run_id": "run-sse-1", "event": "INITIALIZED"}\n\n' + raise httpx.ReadTimeout("upstream died") + + response = MagicMock(spec=httpx.Response) + response.status_code = 200 + response.aiter_bytes = _aiter_bytes + + async def _consume(): + async for _ in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + route_streaming_logging=AsyncMock(), + ): + pass + + with pytest.raises(httpx.ReadTimeout): + await _consume() + + spawned = self._spawned_since(tasks_before) + assert len(spawned) == 1 + # the poller owns the single row; a failure dispatch would collide on its request_id + logging_obj.dispatch_failure_handlers.assert_not_called() + for task in spawned: + task.cancel() + + @pytest.mark.asyncio + async def test_unterminated_run_id_frame_late_spawns_poller(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + await PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + request_body={}, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + raw_bytes=[b'data: {"run_id": "run-sse-1", "event": "INITIALIZED"}'], + end_time=datetime.now(), + ) + + spawned = self._spawned_since(tasks_before) + assert len(spawned) == 1 + assert sse_poller_spawned(logging_obj) + # the poller polls to terminal instead of the old single fetch that mispriced a RUNNING run at $0 + logging_obj.dispatch_success_handlers.assert_not_awaited() + for task in spawned: + task.cancel() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deferred_dispatch_armed", + [False, True], + ids=["enqueued-at-end-of-stream", "parked-for-deferred-dispatch"], +) +async def test_chunk_processor_claims_the_budget_reservation_before_handing_it_to_the_cost_callback( + deferred_dispatch_armed: bool, +): + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + response = _make_streaming_response([b"event-1", b"event-2"]) + logging_obj = _unarmed_logging_obj() + logging_obj.litellm_params = {"metadata": {"user_api_key_budget_reservation": reservation}} + if deferred_dispatch_armed: + logging_obj._on_deferred_stream_complete = AsyncMock() + claimed_when_the_callback_ran = [] + + async def cost_callback(**kwargs): + claimed_when_the_callback_ran.append(reservation["callback_bound"]) + + async for _ in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/bedrock/model/claude/invoke-with-response-stream", + route_streaming_logging=cost_callback, + ): + pass + + if deferred_dispatch_armed: + (parked_cost_callback,) = logging_obj._deferred_stream_complete_args + await parked_cost_callback + else: + await GLOBAL_LOGGING_WORKER.flush() + + assert reservation["callback_bound"] is True + assert claimed_when_the_callback_ran == [True] + + +@pytest.mark.asyncio +async def test_chunk_processor_leaves_the_budget_reservation_for_the_request_end_release_when_the_cost_callback_cannot_be_enqueued(): + reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + response = _make_streaming_response([b"event-1", b"event-2"]) + logging_obj = _unarmed_logging_obj() + logging_obj.litellm_params = {"metadata": {"user_api_key_budget_reservation": reservation}} + + def refuse_to_enqueue(async_coroutine): + async_coroutine.close() + raise RuntimeError("logging worker is shutting down") + + with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", side_effect=refuse_to_enqueue): + async for _ in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/bedrock/model/claude/invoke-with-response-stream", + route_streaming_logging=AsyncMock(), + ): + pass + + assert reservation["callback_bound"] is False diff --git a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py index 53ea761daa7..b3028ae71dd 100644 --- a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py +++ b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py @@ -95,6 +95,71 @@ async def test_openai_exception_handler_invalid_empty_code_defaults_to_500(): } +def _call_id_exception(headers): + return ProxyException( + message="bad input", + type="invalid_request_error", + param="model", + code=400, + headers=headers, + ) + + +@pytest.mark.asyncio +async def test_openai_exception_handler_copies_the_call_id_into_the_error_when_opted_in(monkeypatch): + """With include_call_id_in_error_body on, error.litellm_call_id is byte-identical to the + x-litellm-call-id header, so a pasted str(e) names the request to look up.""" + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}) + exc = _call_id_exception({"x-litellm-call-id": "call-8302"}) + + response = await openai_exception_handler(request=_make_request(), exc=exc) + body = json.loads(response.body) + + assert response.headers["x-litellm-call-id"] == "call-8302" + assert body == { + "error": { + "message": "bad input", + "type": "invalid_request_error", + "param": "model", + "code": "400", + "litellm_call_id": "call-8302", + } + } + + +@pytest.mark.asyncio +async def test_openai_exception_handler_leaves_the_error_alone_when_opted_out(monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + exc = _call_id_exception({"x-litellm-call-id": "call-8302"}) + + response = await openai_exception_handler(request=_make_request(), exc=exc) + body = json.loads(response.body) + + assert response.headers["x-litellm-call-id"] == "call-8302" + assert body == { + "error": { + "message": "bad input", + "type": "invalid_request_error", + "param": "model", + "code": "400", + } + } + + +@pytest.mark.asyncio +async def test_openai_exception_handler_never_fabricates_a_call_id(monkeypatch): + """An error raised before a call id exists (auth failures, say) carries no header, + and the body must not invent one.""" + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}) + exc = _call_id_exception({}) + + response = await openai_exception_handler(request=_make_request(), exc=exc) + body = json.loads(response.body) + + assert "x-litellm-call-id" not in response.headers + assert "litellm_call_id" not in body["error"] + + # --------------------------------------------------------------------------- # _close_dangling_otel_server_span # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 0e0025f5194..13488106df4 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -4,7 +4,7 @@ import json import math from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import Final +from typing import Final, cast import pytest @@ -26,12 +26,14 @@ from litellm.proxy.spend_tracking.budget_reservation import ( _get_team_member_budget_counter, count_request_input_tokens, estimate_request_max_cost, + release_unbound_budget_reservation, reserve_budget_for_request, ) from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm.rust_bridge import bindings, configuration from litellm.rust_bridge import token_counter as rust_token_counter +from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo TOKEN_COUNTING_ROUTES: Final = ( @@ -238,6 +240,22 @@ class _FakeUpstream(Exception): pass +class _FakeTokenizer: + """Stands in for one shared native `Tokenizer`; only its name identifies it.""" + + def __init__(self, name: str, json: str | None = None) -> None: + self.name = name + self.json = json + + +def _fake_native_tokenizers(monkeypatch: pytest.MonkeyPatch, anthropic_json: str | None = None) -> None: + """Point the counter's tokenizer lookups at fakes; the codec path keeps falling back to Python.""" + fakes: Final = {name: _FakeTokenizer(name) for name in ("cl100k_base", "o200k_base")} + anthropic: Final = _FakeTokenizer("anthropic", anthropic_json) + monkeypatch.setattr(tokenizer_dispatch, "native_encoding", fakes.__getitem__) + monkeypatch.setattr(tokenizer_dispatch, "native_anthropic", lambda: anthropic) + + class _FakeNative: RustBridgeDeclined = _FakeDeclined RustUpstreamError = _FakeUpstream @@ -256,19 +274,13 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`.""" + """Stands in for the native `TokenCounter` class, built over a loaded `Tokenizer`.""" def __init__(self) -> None: self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = [] - def __call__(self, tokenizer_json: str) -> _RecordingCounter: - return _RecordingCounter(self, "anthropic") - - def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: - return _RecordingCounter(self, "cl100k_base") - - def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: - return _RecordingCounter(self, "o200k_base") + def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _RecordingCounter: + return _RecordingCounter(self, cast(rust_token_counter.RustTokenizer, tokenizer.name)) class _DecliningCounter: @@ -277,19 +289,14 @@ class _DecliningCounter: class _DecliningFactory: - def __call__(self, tokenizer_json: str) -> _DecliningCounter: - return _DecliningCounter() - - def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter: - return _DecliningCounter() - - def from_o200k_ranks(self, rank_file: str) -> _DecliningCounter: + def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _DecliningCounter: return _DecliningCounter() @pytest.fixture def rust_counter(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) + _fake_native_tokenizers(monkeypatch) rust_token_counter._counter.cache_clear() configuration.reset_rust_configuration() yield @@ -546,3 +553,41 @@ async def test_team_member_reservation_counter_adds_temp_increase_to_live_team_d assert counter is not None assert counter.max_budget == expected_max_budget assert counter.fallback_spend == 0.5 + + +@pytest.mark.asyncio +async def test_reservation_starts_unbound_to_any_callback(): + reservation: Final = await _reserve("/v1/responses") + + assert reservation is not None + assert reservation["callback_bound"] is False + + +@pytest.mark.asyncio +async def test_release_unbound_budget_reservation_frees_the_counter(spend_counter_cache: DualCache): + counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" + reservation: Final = await _reserve_for_tiny_budget_key( + "/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + ) + assert reservation is not None + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) + + await release_unbound_budget_reservation(reservation) + + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.0) + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_callback(spend_counter_cache: DualCache): + counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" + reservation: Final = await _reserve_for_tiny_budget_key( + "/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + ) + assert reservation is not None + reservation["callback_bound"] = True + + await release_unbound_budget_reservation(reservation) + + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) + assert reservation["finalized"] is False diff --git a/tests/test_litellm/proxy/spend_tracking/test_input_tokens.py b/tests/test_litellm/proxy/spend_tracking/test_input_tokens.py new file mode 100644 index 00000000000..49bbe148386 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_input_tokens.py @@ -0,0 +1,191 @@ +"""Tests for input-token counting shared across the reservation path's models.""" + +from __future__ import annotations + +import json +from types import MappingProxyType +from typing import Final, cast + +import pytest + +import litellm +from litellm.proxy.spend_tracking.input_tokens import ( + TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS, + count_input_tokens, + count_input_tokens_for_model, +) +from litellm.rust_bridge import bindings, configuration, token_counter +from litellm.rust_bridge import tokenizer as tokenizer_dispatch +from litellm.rust_bridge.token_counter import RustTokenizer + +ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929" +CL100K_MODEL: Final = "gpt-4" +O200K_MODEL: Final = "gpt-4o" +PYTHON_ONLY_MODEL: Final = "replicate/meta/llama-2-70b-chat" +MESSAGES: Final = [{"role": "user", "content": "hello"}] +RUST_TOKENS: Final = 777 + + +class _FakeDeclined(Exception): + pass + + +class _FakeUpstream(Exception): + pass + + +class _FakeTokenizer: + """Stands in for one shared native `Tokenizer`; only its name identifies it.""" + + def __init__(self, name: str, json: str | None = None) -> None: + self.name = name + self.json = json + + +def _fake_native_tokenizers(monkeypatch: pytest.MonkeyPatch, anthropic_json: str | None = None) -> None: + """Point the counter's tokenizer lookups at fakes; the codec path keeps falling back to Python.""" + fakes: Final = {name: _FakeTokenizer(name) for name in ("cl100k_base", "o200k_base")} + anthropic: Final = _FakeTokenizer("anthropic", anthropic_json) + monkeypatch.setattr(tokenizer_dispatch, "native_encoding", fakes.__getitem__) + monkeypatch.setattr(tokenizer_dispatch, "native_anthropic", lambda: anthropic) + + +class _FakeNative: + RustBridgeDeclined = _FakeDeclined + RustUpstreamError = _FakeUpstream + + +class _RecordingCounter: + def __init__(self, factory: _RecordingFactory, tokenizer: RustTokenizer) -> None: + self.factory = factory + self.tokenizer = tokenizer + + async def acount_request(self, body: bytes) -> object: + self.factory.calls.append((self.tokenizer, body)) + return {"model": "", "input_tokens": RUST_TOKENS} + + +class _RecordingFactory: + def __init__(self) -> None: + self.calls: list[tuple[RustTokenizer, bytes]] = [] + + def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _RecordingCounter: + return _RecordingCounter(self, cast(RustTokenizer, tokenizer.name)) + + +class _DecliningCounter: + async def acount_request(self, body: bytes) -> object: + raise _FakeDeclined("unsupported request shape") + + +class _DecliningFactory: + def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _DecliningCounter: + return _DecliningCounter() + + +@pytest.fixture(autouse=True) +def _reset_bridge(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) + _fake_native_tokenizers(monkeypatch) + token_counter.TOKEN_COUNTER.reset() + token_counter._counter.cache_clear() + configuration.reset_rust_configuration() + yield + token_counter.TOKEN_COUNTER.reset() + token_counter._counter.cache_clear() + configuration.reset_rust_configuration() + + +def _body(model: object) -> tuple[dict[str, object], bytes]: + body: Final = {"model": model, "messages": MESSAGES} + return body, json.dumps(body).encode() + + +@pytest.mark.asyncio +async def test_models_sharing_a_tokenizer_are_counted_once_and_merged() -> None: + factory: Final = _RecordingFactory() + litellm.rust(True) + token_counter.TOKEN_COUNTER.override(factory) + request_body, raw_body = _body([ANTHROPIC_MODEL, CL100K_MODEL, O200K_MODEL, "gpt-5", PYTHON_ONLY_MODEL]) + + counts: Final = await count_input_tokens( + request_body=request_body, + raw_body=raw_body, + models=(ANTHROPIC_MODEL, CL100K_MODEL, O200K_MODEL, "gpt-5", PYTHON_ONLY_MODEL), + ) + + assert factory.calls == [("anthropic", raw_body), ("cl100k_base", raw_body), ("o200k_base", raw_body)] + assert dict(counts) == { + ANTHROPIC_MODEL: RUST_TOKENS, + CL100K_MODEL: RUST_TOKENS, + O200K_MODEL: RUST_TOKENS, + "gpt-5": RUST_TOKENS, + PYTHON_ONLY_MODEL: count_input_tokens_for_model(request_body=request_body, model=PYTHON_ONLY_MODEL), + } + + +@pytest.mark.asyncio +async def test_rust_disabled_counts_everything_in_python() -> None: + factory: Final = _RecordingFactory() + litellm.rust(False) + token_counter.TOKEN_COUNTER.override(factory) + request_body, raw_body = _body([ANTHROPIC_MODEL, CL100K_MODEL]) + + counts: Final = await count_input_tokens( + request_body=request_body, raw_body=raw_body, models=(ANTHROPIC_MODEL, CL100K_MODEL) + ) + + assert factory.calls == [] + assert dict(counts) == { + model: count_input_tokens_for_model(request_body=request_body, model=model) + for model in (ANTHROPIC_MODEL, CL100K_MODEL) + } + + +@pytest.mark.asyncio +async def test_missing_raw_body_counts_in_python() -> None: + factory: Final = _RecordingFactory() + litellm.rust(True) + token_counter.TOKEN_COUNTER.override(factory) + request_body, _ = _body(ANTHROPIC_MODEL) + + counts: Final = await count_input_tokens(request_body=request_body, raw_body=None, models=(ANTHROPIC_MODEL,)) + + assert factory.calls == [] + assert counts[ANTHROPIC_MODEL] == count_input_tokens_for_model(request_body=request_body, model=ANTHROPIC_MODEL) + + +@pytest.mark.asyncio +async def test_missing_binding_counts_in_python() -> None: + litellm.rust(True) + token_counter.TOKEN_COUNTER.override(None) + request_body, raw_body = _body(ANTHROPIC_MODEL) + + counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw_body, models=(ANTHROPIC_MODEL,)) + + assert counts[ANTHROPIC_MODEL] == count_input_tokens_for_model(request_body=request_body, model=ANTHROPIC_MODEL) + + +@pytest.mark.asyncio +async def test_declined_request_counts_in_python() -> None: + litellm.rust(True) + token_counter.TOKEN_COUNTER.override(_DecliningFactory()) + request_body, raw_body = _body(ANTHROPIC_MODEL) + + counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw_body, models=(ANTHROPIC_MODEL,)) + + assert counts[ANTHROPIC_MODEL] == count_input_tokens_for_model(request_body=request_body, model=ANTHROPIC_MODEL) + assert counts[ANTHROPIC_MODEL] != RUST_TOKENS + + +@pytest.mark.asyncio +async def test_large_input_is_still_counted() -> None: + request_body: Final = { + "model": CL100K_MODEL, + "messages": [{"role": "user", "content": "x" * (TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS + 1)}], + } + + counts: Final = await count_input_tokens(request_body=request_body, raw_body=None, models=(CL100K_MODEL,)) + + assert counts[CL100K_MODEL] == count_input_tokens_for_model(request_body=request_body, model=CL100K_MODEL) + assert isinstance(counts, MappingProxyType) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 86bf896188f..b3913079bb2 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -39,8 +39,6 @@ from litellm.proxy.common_utils.user_api_key_cache import ( model_access_group_spend_counter_key, ) from litellm.proxy.spend_tracking.budget_reservation import ( - TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS, - _approximate_input_size, _get_model_access_group_budget_counters, estimate_request_max_cost, get_budget_window_start, @@ -49,6 +47,10 @@ from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, reserve_budget_for_request, ) +from litellm.proxy.spend_tracking.input_tokens import ( + TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS, + _approximate_input_size, +) from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 0b872400be0..218d8246715 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2343,6 +2343,39 @@ class TestCommonRequestProcessingHelpers: assert isinstance(response, JSONResponse) assert response.headers["x-litellm-model-id"] == "fallback-deployment" + @staticmethod + async def _first_chunk_error_response(**create_response_kwargs): + async def mock_generator(): + yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n' + yield "data: [DONE]\n\n" + + return await create_response( + mock_generator(), + "text/event-stream", + {"x-litellm-call-id": "call-8302"}, + **create_response_kwargs, + ) + + async def test_create_response_first_chunk_error_carries_the_call_id_when_opted_in(self): + """A stream that fails on its first chunk answers as JSON, and with + include_call_id_in_error_body on that JSON names the request like the + non-streaming error path does, byte-identical to the header.""" + response = await self._first_chunk_error_response(general_settings={"include_call_id_in_error_body": True}) + + assert isinstance(response, JSONResponse) + assert response.status_code == 403 + assert response.headers["x-litellm-call-id"] == "call-8302" + assert json.loads(response.body) == { + "error": {"code": 403, "message": "forbidden", "litellm_call_id": "call-8302"} + } + + async def test_create_response_first_chunk_error_body_is_unchanged_by_default(self): + response = await self._first_chunk_error_response() + + assert isinstance(response, JSONResponse) + assert response.headers["x-litellm-call-id"] == "call-8302" + assert json.loads(response.body) == {"error": {"code": 403, "message": "forbidden"}} + async def test_create_streaming_response_disables_proxy_buffering(self): """Regression for #28384: every StreamingResponse create_response returns must carry the headers that stop nginx/ingress/Envoy from buffering the @@ -9122,6 +9155,62 @@ class TestStreamingResponseHeadersFollowFallback: assert result.status_code == 400 assert result.headers["x-litellm-applied-guardrails"] == "stream-blocker" + @pytest.mark.asyncio + async def test_streaming_first_chunk_error_carries_the_call_id_when_opted_in(self, monkeypatch): + """The opt-in reaches the streaming path through base_process_llm_request, so a stream + that fails on its first chunk answers with the call id inside its JSON error body, + byte-identical to the x-litellm-call-id header.""" + + def select_data_generator(**kwargs): + async def generator(): + yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n' + yield "data: [DONE]\n\n" + + return generator() + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "lit-8302-call" + logging_obj._defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + processor = ProxyBaseLLMRequestProcessing( + data={"model": "oa", "stream": True, "litellm_logging_obj": logging_obj} + ) + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock( + side_effect=lambda data, user_api_key_dict, response: response + ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + async def fake_route_request(**kwargs): + async def call(): + return SimpleNamespace(_hidden_params={}, fallback_headers_adopted=False) + + return call() + + monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + + result = await processor.base_process_llm_request( + request=Request(scope={"type": "http", "headers": []}), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings={"include_call_id_in_error_body": True}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + is_streaming_request=True, + skip_pre_call_logic=True, + ) + + assert isinstance(result, JSONResponse) + assert result.status_code == 403 + assert result.headers["x-litellm-call-id"] == "lit-8302-call" + assert json.loads(result.body)["error"]["litellm_call_id"] == "lit-8302-call" + class _MessagesFallbackStream: def __init__(self) -> None: diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index dd330d32ce6..f45715953f9 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -489,6 +489,68 @@ async def test_post_call_stream_masking_guardrail_keeps_own_iterator_on_anthropi assert delivered == chunks +@pytest.mark.asyncio +async def test_post_call_stream_presidio_output_masking_masks_anthropic_messages_stream(monkeypatch): + """Regression: the presidio output-masking callback built by initialize_presidio + was rerouted onto the unified scan-only path on /v1/messages, so a card number + the analyzer flagged still streamed to the caller unmasked.""" + import json + + from litellm.caching.caching import DualCache + from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler + from litellm.types.guardrails import SupportedGuardrailIntegrations + + handler = InMemoryGuardrailHandler() + result = handler.initialize_guardrail( + guardrail={ + "guardrail_name": "presidio-card-mask", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value, + "mode": ["pre_call", "post_call"], + "default_on": True, + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + "pii_entities_config": {"CREDIT_CARD": "MASK"}, + "mock_redacted_text": {"text": "", "items": []}, + }, + } + ) + guardrail_id = result["guardrail_id"] + callbacks = [ + handler.guardrail_id_to_custom_guardrail[guardrail_id], + *handler.guardrail_id_to_sibling_callbacks[guardrail_id], + ] + monkeypatch.setattr(litellm, "callbacks", callbacks) + + chunks = _anthropic_stream_chunks(["4111", " 1111 1111 1111"]) + + async def fake_stream(): + for chunk in chunks: + yield chunk + + delivered = [] + async for chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), + request_data={ + "model": "claude-sonnet-5", + "litellm_logging_obj": _streaming_logging_obj(), + "metadata": {}, + }, + ): + delivered.append(chunk) + + wire = b"".join(delivered).decode() + text_deltas = [ + json.loads(line[6:])["delta"]["text"] + for line in wire.split("\n") + if line.startswith("data: ") and json.loads(line[6:]).get("delta", {}).get("type") == "text_delta" + ] + assert "4111" not in wire, wire + assert "".join(text_deltas) == "", wire + assert wire.count("event: message_stop") == 1, wire + + class _AppliesGuardrail(CustomGuardrail): """Implements the unified interface only, so the proxy routes it to unified_guardrail.""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 2792e176e0e..26bfd5c52bd 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -14877,7 +14877,7 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch): - from tokenizers import Tokenizer + from litellm.rust_bridge._native import Tokenizer from litellm import Router from tests.test_litellm.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags @@ -14890,7 +14890,7 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp time.sleep(0.3) return claude_tokenizer - monkeypatch.setattr(litellm.utils, "Tokenizer", SlowHubTokenizer) + monkeypatch.setattr("litellm.rust_bridge.tokenizer.from_pretrained", SlowHubTokenizer.from_pretrained) monkeypatch.setattr( "litellm.proxy.proxy_server.llm_router", Router( @@ -14914,7 +14914,7 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch): - from tokenizers import Tokenizer + from litellm.rust_bridge._native import Tokenizer from litellm import Router from litellm.types.router import DeploymentTypedDict @@ -14931,7 +14931,7 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi }, } - monkeypatch.setattr(litellm.utils, "Tokenizer", MagicMock(from_pretrained=from_pretrained)) + monkeypatch.setattr("litellm.rust_bridge.tokenizer.from_pretrained", from_pretrained) monkeypatch.setattr( "litellm.proxy.proxy_server.llm_router", Router( diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 6cbbc279748..0b51062dd66 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -1282,7 +1282,7 @@ async def test_route_request_routing_group_name_passes_model_gate(): @pytest.mark.asyncio -async def test_route_request_a2a_agent_miss_does_not_consume_model_read_through(monkeypatch): +async def test_route_request_a2a_agent_miss_does_not_consume_model_read_through(fresh_agent_read_through, monkeypatch): from types import SimpleNamespace from unittest.mock import AsyncMock diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index 5fb4392eec6..7f77f938323 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/tests/test_litellm/proxy/utils/helpers/test_model_access.py @@ -110,8 +110,7 @@ def test_create_model_info_response_happy_path_no_metadata(): "owned_by": result["owned_by"], "created_is_int": isinstance(result["created"], int), "metadata_absent": "metadata" not in result, - "max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int) - and result["max_input_tokens"] > 0, + "max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int) and result["max_input_tokens"] > 0, "max_output_tokens_positive_int": isinstance(result["max_output_tokens"], int) and result["max_output_tokens"] > 0, } @@ -205,9 +204,7 @@ def test_validate_model_access_happy_path_single_model_in_list(): def test_validate_model_access_happy_path_batch_all_accessible(): summary = { - "result": validate_model_access( - "gpt-4o,claude-haiku", ["gpt-4o", "claude-haiku", "gemini"] - ), + "result": validate_model_access("gpt-4o,claude-haiku", ["gpt-4o", "claude-haiku", "gemini"]), "input": "gpt-4o,claude-haiku", "available": ["gpt-4o", "claude-haiku", "gemini"], } @@ -389,9 +386,7 @@ async def test_get_available_models_for_user_error_path_complete_list_raises( def _boom(**_kwargs): raise RuntimeError("downstream failure") - monkeypatch.setattr( - "litellm.proxy.auth.model_checks.get_complete_model_list", _boom - ) + monkeypatch.setattr("litellm.proxy.auth.model_checks.get_complete_model_list", _boom) user_api_key_dict = UserAPIKeyAuth( api_key="sk-test-key", user_id="user-1", @@ -481,6 +476,7 @@ async def test_get_available_models_for_user_without_access_groups_grants_nothin ) assert result == [] + @pytest.mark.asyncio async def test_get_available_models_for_user_resolves_key_access_group_models( monkeypatch, @@ -521,3 +517,78 @@ async def test_get_available_models_for_user_resolves_key_access_group_models( user_api_key_cache=MagicMock(), ) assert result == ["model-b"] + + +def _agent_ceiling(models: frozenset[str] | None): + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + + async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None: + if models is None: + return None + return AgentAccessGroupCeiling( + access_group_ids=("ag-agent",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset() + ) + + return resolve + + +def _agent_key(models: list[str]) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-agent-key", user_id="user-1", agent_id="agent-1", models=models) + + +@pytest.mark.asyncio +async def test_agent_key_listing_is_capped_to_its_access_groups(): + result = await get_available_models_for_user( + user_api_key_dict=_agent_key(["model-a", "model-b", "model-c"]), + llm_router=_router_with_models(["model-a", "model-b", "model-c"]), + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(frozenset({"model-b", "model-d"})), + ) + assert result == ["model-b"] + + +@pytest.mark.asyncio +async def test_agent_key_listing_is_empty_when_its_groups_grant_no_model(): + result = await get_available_models_for_user( + user_api_key_dict=_agent_key(["model-a"]), + llm_router=_router_with_models(["model-a"]), + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(frozenset()), + ) + assert result == [] + + +@pytest.mark.asyncio +async def test_agent_ceiling_expands_a_model_access_group_name_for_listing(): + router = _router_with_models(["model-a", "model-b"]) + router.get_model_access_groups.return_value = {"fast-models": ["model-b"]} + result = await get_available_models_for_user( + user_api_key_dict=_agent_key(["model-a", "model-b"]), + llm_router=router, + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(frozenset({"fast-models"})), + ) + assert result == ["model-b"] + + +@pytest.mark.asyncio +async def test_listing_is_unchanged_without_an_agent_or_without_attached_groups(): + router = _router_with_models(["model-a", "model-b"]) + plain_key = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-plain", user_id="user-1", models=["model-a", "model-b"]), + llm_router=router, + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(frozenset({"model-a"})), + ) + agent_without_groups = await get_available_models_for_user( + user_api_key_dict=_agent_key(["model-a", "model-b"]), + llm_router=router, + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(None), + ) + assert (plain_key, agent_without_groups) == (["model-a", "model-b"], ["model-a", "model-b"]) diff --git a/tests/test_litellm/responses/test_dispatch.py b/tests/test_litellm/responses/test_dispatch.py index 2990360d550..45cb5c4f1ad 100644 --- a/tests/test_litellm/responses/test_dispatch.py +++ b/tests/test_litellm/responses/test_dispatch.py @@ -13,7 +13,7 @@ from litellm.responses.dispatch import ( ) from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, Rule +from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.responses.entrypoints import ( NATIVE_ARESPONSES, @@ -26,7 +26,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse INPUT: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final = () -RUST_RULES: Final = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED),) +RUST_RULES: Final = (RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED),) def _response(model: str = "gpt-4o") -> ResponsesAPIResponse: @@ -102,7 +102,8 @@ async def test_async_python_route_forwards_original_call_shape() -> None: response: Final = _response() async def python( - *call_args: object, **call_kwargs: object # kwargs-ok: records call shape + *call_args: object, + **call_kwargs: object, # kwargs-ok: records call shape ) -> ResponsesAPIResponse: captured.append((call_args, call_kwargs)) return response @@ -143,9 +144,7 @@ def test_native_receives_normalized_request_and_original_call_shape() -> None: "custom_llm_provider": "anthropic", "litellm_metadata": metadata, } - captured: Final[ - list[tuple[LiteLLMResponsesRequest, tuple[object, ...], Mapping[str, object]]] - ] = [] + captured: Final[list[tuple[LiteLLMResponsesRequest, tuple[object, ...], Mapping[str, object]]]] = [] response: Final = _response("anthropic/claude-sonnet-4-5") def python(*call_args: object, **call_kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: rejected fallback @@ -228,9 +227,7 @@ def test_internal_async_marker_bypasses_native() -> None: ((), {}), ), ) -def test_binding_errors_delegate_unchanged_to_python( - args: tuple[object, ...], kwargs: Mapping[str, object] -) -> None: +def test_binding_errors_delegate_unchanged_to_python(args: tuple[object, ...], kwargs: Mapping[str, object]) -> None: captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = [] response: Final = _response() diff --git a/tests/test_litellm/responses/test_responses_api_bridge_flag.py b/tests/test_litellm/responses/test_responses_api_bridge_flag.py index 2f64cc8debc..16135106b41 100644 --- a/tests/test_litellm/responses/test_responses_api_bridge_flag.py +++ b/tests/test_litellm/responses/test_responses_api_bridge_flag.py @@ -262,6 +262,126 @@ class TestUseResponsesApiBridgeFlag: assert request_body["messages"] == [{"role": "user", "content": "Hello"}] assert response.output[0].content[0].text == "Answer" + def test_bridge_drops_client_metadata_even_when_allowed_openai_params_names_it( + self, respx_mock: respx.MockRouter + ): + upstream: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + response: Final = litellm.responses( + model="openai/my-custom-model", + input="Hello", + use_chat_completions_api=True, + allowed_openai_params=["client_metadata"], + client_metadata={"turn_id": "turn-1", "thread_id": "thread-1"}, + api_key="fake-provider-api-key", + num_retries=0, + ) + + assert upstream.call_count == 1 + request_body: Final = json.loads(upstream.calls[0].request.read()) + assert "client_metadata" not in request_body + assert request_body["messages"] == [{"role": "user", "content": "Hello"}] + assert response.output[0].content[0].text == "Answer" + + def test_bridge_merges_instructions_and_developer_input_for_databricks(self, respx_mock: respx.MockRouter): + upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + response: Final = litellm.responses( + model="databricks/my-custom-model", + instructions="You are terse.", + input=[ + {"role": "developer", "content": [{"type": "input_text", "text": "Skills: none."}]}, + {"role": "user", "content": [{"type": "input_text", "text": "Hello"}]}, + ], + client_metadata={"turn_id": "turn-1", "thread_id": "thread-1"}, + chat_template_kwargs={"thinking": True}, + api_base="https://example.databricks.test/serving-endpoints", + api_key="fake-databricks-api-key", + num_retries=0, + ) + + assert upstream.call_count == 1 + request_body: Final = json.loads(upstream.calls[0].request.read()) + assert request_body["messages"] == [ + { + "role": "system", + "content": [{"type": "text", "text": "You are terse."}, {"type": "text", "text": "Skills: none."}], + }, + {"role": "user", "content": [{"type": "text", "text": "Hello"}]}, + ] + assert "client_metadata" not in request_body + assert request_body["chat_template_kwargs"] == {"thinking": True} + assert response.output[0].content[0].text == "Answer" + + def test_bridge_drops_client_metadata_for_provider_without_native_config(self, respx_mock: respx.MockRouter): + upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + response: Final = litellm.responses( + model="databricks/my-custom-model", + input="Hello", + client_metadata={ + "turn_id": "turn-1", + "thread_id": "thread-1", + "session_id": "session-1", + "root_turn_id": "turn-1", + "x-codex-installation-id": "install-1", + "x-codex-turn-metadata": '{"turn_id":"turn-1"}', + }, + chat_template_kwargs={"thinking": True}, + api_base="https://example.databricks.test/serving-endpoints", + api_key="fake-databricks-api-key", + num_retries=0, + ) + + assert upstream.call_count == 1 + request_body: Final = json.loads(upstream.calls[0].request.read()) + assert "client_metadata" not in request_body + assert request_body["chat_template_kwargs"] == {"thinking": True} + assert request_body["messages"] == [{"role": "user", "content": "Hello"}] + assert response.output[0].content[0].text == "Answer" + def test_bridge_keeps_deployment_credentials_while_dropping_unknown_params(self, respx_mock: respx.MockRouter): upstream: Final = respx_mock.post( "https://example-resource.openai.azure.com/openai/deployments/my-deployment/chat/completions", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 97e7a799053..baa15ac1568 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -79,6 +79,7 @@ from litellm.router_strategy.complexity_router.jev_classifier import ( JevSystemOneResponse, JevUsage, ) +from litellm.router_strategy.complexity_router.llm_v2 import LLM_V2_PROMPT_VERSION from litellm.router_strategy.complexity_router.tier_predictor import ( TierGlobalStatistic, TrainedTierArtifact, @@ -3240,6 +3241,41 @@ class TestCapabilityClassifier: ) assert response.model == "capable-model" assert response.routing_decision["cause"] == "capability_classifier_fallback" + assert "classifier_p_solve" not in response.routing_decision + assert "classifier_threshold" not in response.routing_decision + + @pytest.mark.asyncio + @pytest.mark.parametrize("bypass", ("literal_keyword_match", "session_affinity_pin", "housekeeping")) + async def test_bypasses_do_not_reuse_the_previous_capability_forecast( + self, + mock_router_instance: MagicMock, + bypass: Literal["literal_keyword_match", "session_affinity_pin", "housekeeping"], + ) -> None: + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=0.8))) + mock_router_instance.cache = DualCache() + router: Final = self._router( + mock_router_instance, + session_affinity=bypass == "session_affinity_pin", + keyword_tier_rules=[{"keywords": ["quick lookup"], "tier": "SIMPLE"}], + ) + original: Final = await router.async_pre_routing_hook( + model="capability-router", + request_kwargs={"metadata": {"session_id": "forecast-bypass"}}, + messages=[{"role": "user", "content": "Hello!"}], + ) + result: Final = await router.async_pre_routing_hook( + model="capability-router", + request_kwargs={"metadata": {"session_id": "forecast-bypass"}}, + messages=[{"role": "user", "content": TITLE_ASK if bypass == "housekeeping" else "quick lookup"}], + ) + + assert original is not None and original.routing_decision is not None + assert original.routing_decision["classifier_p_solve"] == 0.8 + assert result is not None and result.routing_decision is not None + assert result.routing_decision["cause"] == bypass + assert "classifier_p_solve" not in result.routing_decision + assert "classifier_threshold" not in result.routing_decision + mock_router_instance.acompletion.assert_awaited_once() CUSTOM_TIER_LABELS: Dict[str, str] = { @@ -14686,6 +14722,121 @@ class TestModalityRouting: @pytest.mark.usefixtures("local_model_cost_map") class TestHealthFallbackDispatch: + @pytest.mark.asyncio + @pytest.mark.parametrize("classifier", ("capability", "llm_v2")) + @pytest.mark.parametrize("calibrated", (False, True), ids=("raw", "calibrated")) + @pytest.mark.parametrize("rewrite", ("modality_escalation", "health_failover", "health_default_fallback")) + async def test_classifier_forecasts_survive_placement_rewrites( + self, + classifier: Literal["capability", "llm_v2"], + calibrated: bool, + rewrite: Literal["modality_escalation", "health_failover", "health_default_fallback"], + ) -> None: + calibration: Final = {"slope": 0.8, "intercept": 0.1} + classifier_config: Final = ( + { + "capability_classifier_config": { + "efficient_tier": "SIMPLE", + "capable_tier": "REASONING", + "base_threshold": 0.0, + "threshold_step": 0.1, + **({"calibration": {"version": "test-v1", **calibration}} if calibrated else {}), + } + } + if classifier == "capability" + else { + "llm_v2_config": { + "efficient_profile": "Small coding solver", + "capable_profile": "Large coding solver", + "harness": "Repository tools", + "max_quality_gap": 0.0, + **( + { + "calibration": { + "version": "test-v1", + "prompt_version": LLM_V2_PROMPT_VERSION, + "efficient": calibration, + "capable": calibration, + } + } + if calibrated + else {} + ), + } + } + ) + router: Final = self._router( + config={ + "classifier_type": classifier, + "classifier_llm_config": {"model": "fallback", "timeout_ms": 10000}, + "tiers": {"SIMPLE": "primary", "REASONING": "peer"}, + "tier_labels": {"SIMPLE": "Entry", "REASONING": "Advanced"}, + "modality_routing": True, + **classifier_config, + } + ) + verdict: Final = ( + _capability_reply(p_solve=0.0) + if classifier == "capability" + else json.dumps( + { + "crux": "Preserve existing behavior", + "demands": {"reasoning": "routine", "scope": "localized", "specification": "clear"}, + "verification": "relevant", + "forecasts": { + "efficient": {"likely_failure": "Miss an edge case", "p_solve": 0.0}, + "capable": {"likely_failure": "Miss an edge case", "p_solve": 0.0}, + }, + } + ) + ) + judge_response: Final = litellm.ModelResponse( + choices=[{"message": {"role": "assistant", "content": verdict}}], + usage={"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, + ) + with respx.mock(assert_all_mocked=True) as upstream: + upstream.post(host="fallback.test").respond(json=judge_response.model_dump()) + original: Final = await router.async_pre_routing_hook( + model="health-router", request_kwargs={}, messages=[{"role": "user", "content": "Hello!"}] + ) + for deployment in router.model_list: + deployment["model_info"]["supports_vision"] = ( + rewrite != "modality_escalation" or deployment["model_name"] != "primary" + ) + if rewrite != "modality_escalation": + self._unavailable(router, "primary-id", "cooldown") + if rewrite == "health_default_fallback": + self._unavailable(router, "peer-id", "cooldown") + result: Final = await router.async_pre_routing_hook( + model="health-router", request_kwargs={}, messages=TestModalityRouting.IMAGE_MESSAGE + ) + + assert original is not None and original.routing_decision is not None + assert original.model == "primary" + assert result is not None and result.routing_decision is not None + decision: Final = result.routing_decision + assert decision["cause"] == rewrite + assert result.model == ("fallback" if rewrite == "health_default_fallback" else "peer") + expected: Final = { + field: value for field, value in original.routing_decision.items() if field.startswith("classifier_") + } + assert expected["classifier_p_solve" if classifier == "capability" else "classifier_efficient_p_solve"] == 0.0 + assert ("classifier_calibration_version" in expected) is calibrated + assert {field: value for field, value in decision.items() if field.startswith("classifier_")} == expected + if rewrite == "health_default_fallback": + assert "tier" not in decision and "tier_label" not in decision + 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")) async def test_health_rewrites_preserve_the_original_heuristic_v2_forecast(self, peer: bool) -> None: diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/test_litellm/router_strategy/test_llm_v2.py index 27d31cbe640..6fb6df3265d 100644 --- a/tests/test_litellm/router_strategy/test_llm_v2.py +++ b/tests/test_litellm/router_strategy/test_llm_v2.py @@ -477,6 +477,10 @@ async def test_user_turn_mode_reuses_forecast_until_a_new_user_requirement() -> assert first.model == second.model == "efficient" assert first.routing_decision["cause"] == "llm_v2_classifier" assert first.routing_decision["classifier_cost"] == 0.001 + assert second is not None and second.routing_decision is not None + assert second.routing_decision["cause"] == "user_turn_continuation" + assert "classifier_efficient_p_solve" not in second.routing_decision + assert "classifier_capable_p_solve" not in second.routing_decision client.acompletion.assert_awaited_once() client.acompletion.return_value = _response(_verdict(0.3, 0.9).model_dump_json()) updated: Final = await router.async_pre_routing_hook( diff --git a/tests/test_litellm/rust_bridge/ocr/test_secrets.py b/tests/test_litellm/rust_bridge/ocr/test_secrets.py new file mode 100644 index 00000000000..085a42dd373 --- /dev/null +++ b/tests/test_litellm/rust_bridge/ocr/test_secrets.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.integrations.custom_secret_manager import CustomSecretManager +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge import configuration +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service + + +class _VaultSecrets(CustomSecretManager): + def __init__(self) -> None: + super().__init__(secret_manager_name="rust_bridge_ocr_test") + + async def async_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return "vault-key" if secret_name == "MISTRAL_API_KEY" else None + + def sync_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return "vault-key" if secret_name == "MISTRAL_API_KEY" else None + + +async def _call(asynchronous: bool, api_base: str) -> OCRResponse: + if asynchronous: + return await litellm.aocr( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + api_base=api_base, + ) + return litellm.ocr( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + api_base=api_base, + ) + + +_RESPONSE: Final = { + "pages": [{"index": 0, "markdown": "parsed document", "images": []}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +@pytest.mark.parametrize("rust_enabled", ("0", "1")) +@pytest.mark.parametrize("access_mode", ("read_only", "read_and_write")) +@pytest.mark.parametrize("system", (None, KeyManagementSystem.CUSTOM)) +async def test_readable_secret_managers_keep_python_ocr_fallback( + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + rust_enabled: str, + access_mode: str, + system: KeyManagementSystem | None, +) -> None: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", rust_enabled) + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets()) + monkeypatch.setattr(litellm, "_key_management_system", system) + monkeypatch.setattr( + litellm, + "_key_management_settings", + KeyManagementSettings(access_mode=access_mode, hosted_keys=["MISTRAL_API_KEY"]), + ) + configuration.reset_rust_configuration() + + with recording_service() as server: + server.default_response = ResponseSpec(body=_RESPONSE) + result: Final = await _call(asynchronous, server.base_url) + + assert result.pages[0].markdown == "parsed document" + assert len(server.requests) == 1 + expected_key: Final = "vault-key" if system is KeyManagementSystem.CUSTOM else "environment-key" + assert server.requests[0].headers["authorization"] == f"Bearer {expected_key}" + assert "x-litellm-rust" not in result._hidden_params.get("additional_headers", {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_no_secret_client_leaves_dormant_binding_settings_unread( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1") + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + monkeypatch.setattr(litellm, "secret_manager_client", None) + monkeypatch.setattr(litellm, "_key_management_settings", object()) + configuration.reset_rust_configuration() + + with recording_service() as server: + server.default_response = ResponseSpec(body=_RESPONSE) + result: Final = await _call(asynchronous, server.base_url) + + assert result.pages[0].markdown == "parsed document" + assert len(server.requests) == 1 + assert server.requests[0].headers["authorization"] == "Bearer environment-key" + assert result._hidden_params["additional_headers"]["x-litellm-rust"] == "true" diff --git a/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py b/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py index 1f6a214398a..05f2d13a079 100644 --- a/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py +++ b/tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py @@ -9,9 +9,10 @@ import pytest from pydantic import TypeAdapter import litellm +from litellm._internal_context import is_internal_call from litellm.litellm_core_utils.litellm_logging import Logging from litellm.rust_bridge import callbacks_legacy_python as legacy -from litellm.rust_bridge.callbacks_legacy_python import check_limits, setup +from litellm.rust_bridge.callbacks_legacy_python import check_limits, failure_handler, setup _OCR_KWARGS: Final = MappingProxyType( { @@ -81,6 +82,82 @@ def test_setup_builds_a_logger_when_none_is_supplied(call_type: str, kwargs: Map assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"] +def _budget_reservation() -> dict: + return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + + +def _kwargs_with_a_budget_reservation(reservation: dict) -> dict[str, object]: + return {**_OCR_KWARGS, "metadata": {"user_api_key_budget_reservation": reservation}} + + +def test_setup_claims_the_budget_reservation_for_an_async_call() -> None: + reservation: Final = _budget_reservation() + + setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=True) + + assert reservation["callback_bound"] is True + + +def test_setup_claims_the_budget_reservation_a_supplied_logger_already_saw() -> None: + reservation: Final = _budget_reservation() + supplied: Final = _supplied_logger() + supplied.update_environment_variables( + litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}, optional_params={} + ) + assert reservation["callback_bound"] is False + + setup("aocr", (), {**_OCR_KWARGS, "litellm_logging_obj": supplied}, datetime.datetime.now(), asynchronous=True) + + assert reservation["callback_bound"] is True + + +def test_setup_leaves_the_budget_reservation_alone_for_a_sync_call() -> None: + reservation: Final = _budget_reservation() + + setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=False) + + assert reservation["callback_bound"] is False + + +def test_setup_leaves_the_budget_reservation_alone_for_an_internal_call() -> None: + reservation: Final = _budget_reservation() + token: Final = is_internal_call.set(True) + try: + setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=True) + finally: + is_internal_call.reset(token) + + assert reservation["callback_bound"] is False + + +def test_failure_handler_hands_the_budget_reservation_back_for_an_async_call() -> None: + reservation: Final = _budget_reservation() + now: Final = datetime.datetime.now() + result: Final = setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), now, asynchronous=True) + assert reservation["callback_bound"] is True + + pending: Final = failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=True) + + assert reservation["callback_bound"] is False + assert pending is not None + pending.close() + + +def test_failure_handler_of_an_internal_call_leaves_the_outer_budget_reservation_claim_in_place() -> None: + reservation: Final = _budget_reservation() + now: Final = datetime.datetime.now() + result: Final = setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), now, asynchronous=True) + token: Final = is_internal_call.set(True) + try: + pending: Final = failure_handler(result.logger, RuntimeError("inner step failed"), now, now, asynchronous=True) + finally: + is_internal_call.reset(token) + + assert reservation["callback_bound"] is True + assert pending is not None + pending.close() + + CONTRACT_PATH: Final = ( Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json" ) diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/test_litellm/rust_bridge/test_catalog.py index 147e863baf5..82e3766e8a2 100644 --- a/tests/test_litellm/rust_bridge/test_catalog.py +++ b/tests/test_litellm/rust_bridge/test_catalog.py @@ -6,8 +6,21 @@ from typing import Final import pytest from litellm.rust_bridge import catalog, configuration -from litellm.rust_bridge.catalog import Context, Delivery, Route, Rule +from litellm.rust_bridge.catalog import ( + CacheContext, + CacheRule, + Context, + Delivery, + Route, + RouteContext, + RouteRule, + Rules, + SecretManagerContext, + SecretManagerRule, +) from litellm.rust_bridge.configuration import Decision, Rollout +from litellm.types.caching import LiteLLMCacheType +from litellm.types.secret_managers.main import KeyManagementSystem @pytest.fixture(autouse=True) @@ -34,13 +47,13 @@ def test_shipped_decisions( configuration.rust(process) if environment is not None: monkeypatch.setenv("LITELLM_RUST", environment) - context: Final = Context(route, provider=provider, model="test-model", delivery=delivery) + context: Final = RouteContext(route, provider=provider, model="test-model", delivery=delivery) if route is Route.OCR: enabled: Final = environment == "1" if environment is not None else process is not False assert catalog.rollout(context) is Rollout.RUST_OPT_OUT assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) - elif route is Route.MESSAGES: + elif route in (Route.MESSAGES, Route.TOKEN_COUNTER, Route.TOKENIZER): enabled: Final = environment == "1" if environment is not None else process is True assert catalog.rollout(context) is Rollout.RUST_OPT_IN assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) @@ -57,31 +70,63 @@ def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pyt configuration.rust(True) monkeypatch.setenv("LITELLM_RUST", "1") - assert catalog.rollout(Context(route), rules=()) is Rollout.PYTHON_ONLY - assert catalog.decision(Context(route), rules=()) is Decision.PYTHON + assert catalog.rollout(RouteContext(route), rules=()) is Rollout.PYTHON_ONLY + assert catalog.decision(RouteContext(route), rules=()) is Decision.PYTHON + + +@pytest.mark.parametrize( + "context", + ( + *(CacheContext(backend.value) for backend in LiteLLMCacheType), + *(SecretManagerContext(system.value) for system in KeyManagementSystem), + CacheContext("custom"), + SecretManagerContext("unknown"), + ), +) +def test_backend_rollouts_stay_on_python_when_global_rust_is_enabled( + monkeypatch: pytest.MonkeyPatch, context: Context +) -> None: + configuration.rust(True) + monkeypatch.setenv("LITELLM_RUST", "1") + + assert catalog.rollout(context) is Rollout.PYTHON_ONLY + assert catalog.decision(context) is Decision.PYTHON + + +def test_response_cache_rules_select_the_whole_backend_runtime() -> None: + rules: Final = ( + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + CacheRule(Rollout.PYTHON_ONLY), + ) + + assert catalog.decision(CacheContext(backend="local"), rules) is Decision.RUST_REQUIRED + assert catalog.decision(CacheContext(backend="redis"), rules) is Decision.PYTHON @pytest.mark.parametrize( ("context", "expected"), ( - (Context(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), Decision.RUST_REQUIRED), - (Context(Route.RESPONSES, provider="openai", model="m"), Decision.PYTHON), - (Context(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.STREAMING), Decision.PYTHON), - (Context(Route.RESPONSES, provider="openai", model="other", delivery=Delivery.WEBSOCKET), Decision.PYTHON), - (Context(Route.RESPONSES, provider="anthropic", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON), - (Context(Route.MESSAGES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON), + ( + RouteContext(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), + Decision.RUST_REQUIRED, + ), + (RouteContext(Route.RESPONSES, provider="openai", model="m"), Decision.PYTHON), + (RouteContext(Route.RESPONSES, provider="openai", model="m", delivery=Delivery.STREAMING), Decision.PYTHON), + (RouteContext(Route.RESPONSES, provider="openai", model="other", delivery=Delivery.WEBSOCKET), Decision.PYTHON), + (RouteContext(Route.RESPONSES, provider="anthropic", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON), + (RouteContext(Route.MESSAGES, provider="openai", model="m", delivery=Delivery.WEBSOCKET), Decision.PYTHON), ), ) -def test_first_matching_rule_respects_every_constraint(context: Context, expected: Decision) -> None: +def test_first_matching_rule_respects_every_constraint(context: RouteContext, expected: Decision) -> None: rules: Final = ( - Rule( + RouteRule( Route.RESPONSES, Rollout.RUST_REQUIRED, providers=frozenset({"openai"}), models=frozenset({"m"}), deliveries=frozenset({Delivery.WEBSOCKET}), ), - Rule(Route.RESPONSES, Rollout.PYTHON_ONLY), + RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), ) assert catalog.decision(context, rules) is expected @@ -96,4 +141,78 @@ def test_textract_ocr_has_no_python_path_to_opt_out_to( if environment is not None: monkeypatch.setenv("LITELLM_RUST", environment) - assert catalog.decision(Context(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED + assert catalog.decision(RouteContext(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED + + +@pytest.mark.parametrize( + ("context", "expected"), + ( + (RouteContext(Route.OCR, provider="local"), Decision.RUST_REQUIRED), + (RouteContext(Route.OCR, provider="other"), Decision.PYTHON), + (RouteContext(Route.MESSAGES, provider="local"), Decision.PYTHON), + (CacheContext("local"), Decision.RUST_WITH_FALLBACK), + (CacheContext("other"), Decision.PYTHON), + (SecretManagerContext("local"), Decision.PYTHON), + (SecretManagerContext("other"), Decision.RUST_REQUIRED), + ), +) +def test_mixed_rules_select_only_the_matching_domain(context: Context, expected: Decision) -> None: + rules: Final[Rules] = ( + CacheRule(Rollout.RUST_OPT_OUT, backends=frozenset({"local"})), + CacheRule(Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + SecretManagerRule(Rollout.RUST_REQUIRED), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"local"})), + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + ) + + assert catalog.decision(context, rules) is expected + + +@pytest.mark.parametrize("context", (RouteContext(Route.OCR), CacheContext("local"), SecretManagerContext("local"))) +@pytest.mark.parametrize( + ("rollout", "process", "environment", "expected"), + ( + (Rollout.PYTHON_ONLY, True, "1", Decision.PYTHON), + (Rollout.RUST_REQUIRED, False, "0", Decision.RUST_REQUIRED), + (Rollout.RUST_OPT_IN, None, None, Decision.PYTHON), + (Rollout.RUST_OPT_OUT, None, None, Decision.RUST_WITH_FALLBACK), + (Rollout.RUST_OPT_IN, True, None, Decision.RUST_WITH_FALLBACK), + (Rollout.RUST_OPT_OUT, False, None, Decision.PYTHON), + (Rollout.RUST_OPT_IN, False, "1", Decision.RUST_WITH_FALLBACK), + (Rollout.RUST_OPT_OUT, True, "0", Decision.PYTHON), + ), +) +def test_all_domains_share_rollout_switches_and_first_match( + monkeypatch: pytest.MonkeyPatch, + context: Context, + rollout: Rollout, + process: bool | None, + environment: str | None, + expected: Decision, +) -> None: + configuration.rust(process) + if environment is not None: + monkeypatch.setenv("LITELLM_RUST", environment) + rules: Final[Rules] = ( + RouteRule(Route.OCR, rollout), + CacheRule(rollout), + SecretManagerRule(rollout), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED), + CacheRule(Rollout.RUST_REQUIRED), + SecretManagerRule(Rollout.RUST_REQUIRED), + ) + + assert catalog.decision(context, rules) is expected + assert catalog.decision(context, ()) is Decision.PYTHON + + +@pytest.mark.parametrize("context", (RouteContext(Route.OCR), CacheContext("local"), SecretManagerContext("local"))) +def test_empty_constraints_match_nothing(context: Context) -> None: + rules: Final[Rules] = ( + RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset()), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset()), + SecretManagerRule(Rollout.RUST_REQUIRED, systems=frozenset()), + ) + + assert catalog.decision(context, rules) is Decision.PYTHON diff --git a/tests/test_litellm/rust_bridge/test_dispatch.py b/tests/test_litellm/rust_bridge/test_dispatch.py index 66f8d114f7a..9a3793a772e 100644 --- a/tests/test_litellm/rust_bridge/test_dispatch.py +++ b/tests/test_litellm/rust_bridge/test_dispatch.py @@ -6,7 +6,7 @@ import pytest from litellm.rust_bridge import configuration from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Context, Delivery, Route, Rule, Rules +from litellm.rust_bridge.catalog import CacheRule, Delivery, Route, RouteContext, RouteRule, Rules, SecretManagerRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.dispatch import PublicDispatch @@ -22,14 +22,15 @@ def binding() -> NativeBinding[object]: return bound -def test_route_without_rules_forwards_before_request_projection() -> None: +@pytest.mark.parametrize("rules", ((), (CacheRule(Rollout.RUST_REQUIRED), SecretManagerRule(Rollout.RUST_REQUIRED)))) +def test_route_without_rules_forwards_before_request_projection(rules: Rules) -> None: stream: Final[Iterator[int]] = iter((1, 2)) def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request: pytest.fail("Python-only routes must not project the request") dispatch: Final = PublicDispatch( - route=Route.CHAT_COMPLETIONS, request=reject_request, context=lambda _: Context(Route.CHAT_COMPLETIONS) + route=Route.CHAT_COMPLETIONS, request=reject_request, context=lambda _: RouteContext(Route.CHAT_COMPLETIONS) ) result: Final = dispatch.run( ("model",), @@ -37,15 +38,15 @@ def test_route_without_rules_forwards_before_request_projection() -> None: python=lambda *args, **kwargs: stream, binding=binding(), native=lambda hook, request, args, kwargs: pytest.fail("Python-only routes must not call native"), - rules=(), + rules=rules, ) assert result is stream def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None: rules: Final[Rules] = ( - Rule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), - Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED), + RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY), + RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED), ) def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request: @@ -54,7 +55,7 @@ def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None dispatch: Final = PublicDispatch( route=Route.CHAT_COMPLETIONS, request=reject_request, - context=lambda _: Context(Route.CHAT_COMPLETIONS), + context=lambda _: RouteContext(Route.CHAT_COMPLETIONS), ) expected: Final = object() result: Final = dispatch.run( @@ -69,12 +70,12 @@ def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None def test_disabled_optional_rust_rule_forwards_before_projection() -> None: - rules: Final[Rules] = (Rule(Route.OCR, Rollout.RUST_OPT_OUT),) + rules: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),) def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request: pytest.fail("Disabled optional Rust must not project the request") - dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: Context(Route.OCR)) + dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: RouteContext(Route.OCR)) expected: Final = object() configuration.rust(False) try: @@ -95,12 +96,14 @@ def test_native_stream_result_is_not_consumed_or_wrapped() -> None: request: Final = Request(model="streaming-model") stream: Final[Iterator[int]] = iter((1, 2)) rules: Final[Rules] = ( - Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.STREAMING})), + CacheRule(Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY), + RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.STREAMING})), ) dispatch: Final = PublicDispatch( route=Route.CHAT_COMPLETIONS, request=lambda args, kwargs: request, - context=lambda value: Context(Route.CHAT_COMPLETIONS, model=value.model, delivery=Delivery.STREAMING), + context=lambda value: RouteContext(Route.CHAT_COMPLETIONS, model=value.model, delivery=Delivery.STREAMING), ) def native(request: Request, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Iterator[int]: @@ -122,7 +125,8 @@ def test_native_stream_result_is_not_consumed_or_wrapped() -> None: @pytest.mark.asyncio -async def test_async_route_without_rules_preserves_async_iterator_result() -> None: +@pytest.mark.parametrize("rules", ((), (CacheRule(Rollout.RUST_REQUIRED), SecretManagerRule(Rollout.RUST_REQUIRED)))) +async def test_async_route_without_rules_preserves_async_iterator_result(rules: Rules) -> None: async def chunks() -> AsyncGenerator[int, None]: yield 1 @@ -135,7 +139,7 @@ async def test_async_route_without_rules_preserves_async_iterator_result() -> No return stream dispatch: Final = PublicDispatch( - route=Route.RESPONSES, request=reject_request, context=lambda _: Context(Route.RESPONSES) + route=Route.RESPONSES, request=reject_request, context=lambda _: RouteContext(Route.RESPONSES) ) result: Final = await dispatch.arun( ("model",), @@ -143,7 +147,7 @@ async def test_async_route_without_rules_preserves_async_iterator_result() -> No python=python, binding=binding(), native=lambda hook, request, args, kwargs: pytest.fail("Python-only routes must not call native"), - rules=(), + rules=rules, ) assert result is stream await stream.aclose() @@ -152,11 +156,13 @@ async def test_async_route_without_rules_preserves_async_iterator_result() -> No @pytest.mark.asyncio async def test_async_dispatch_accepts_websocket_style_none_result() -> None: request: Final = Request(model="realtime-model") - rules: Final[Rules] = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.WEBSOCKET})),) + rules: Final[Rules] = ( + RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED, deliveries=frozenset({Delivery.WEBSOCKET})), + ) dispatch: Final = PublicDispatch( route=Route.RESPONSES, request=lambda args, kwargs: request, - context=lambda value: Context(Route.RESPONSES, model=value.model, delivery=Delivery.WEBSOCKET), + context=lambda value: RouteContext(Route.RESPONSES, model=value.model, delivery=Delivery.WEBSOCKET), ) async def python(*args: object, **kwargs: object) -> None: # kwargs-ok: public pass-through shape @@ -183,14 +189,14 @@ async def test_async_dispatch_accepts_websocket_style_none_result() -> None: def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() -> None: rules: Final[Rules] = ( - Rule(Route.MESSAGES, Rollout.RUST_REQUIRED), - Rule(Route.OCR, Rollout.PYTHON_ONLY, providers=frozenset({"mistral"})), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + RouteRule(Route.OCR, Rollout.PYTHON_ONLY, providers=frozenset({"mistral"})), ) def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request: pytest.fail("Rules that cannot select Rust must not project the request") - dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: Context(Route.OCR)) + dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: RouteContext(Route.OCR)) expected: Final = object() result: Final = dispatch.run( ("model",), @@ -206,11 +212,11 @@ def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() - @pytest.mark.asyncio async def test_async_bypass_forwards_to_python_without_native() -> None: request: Final = Request(model="bypassed-model") - rules: Final[Rules] = (Rule(Route.RESPONSES, Rollout.RUST_REQUIRED),) + rules: Final[Rules] = (RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED),) dispatch: Final = PublicDispatch( route=Route.RESPONSES, request=lambda args, kwargs: request, - context=lambda value: Context(Route.RESPONSES, model=value.model), + context=lambda value: RouteContext(Route.RESPONSES, model=value.model), bypass=lambda value: value.model == "bypassed-model", ) expected: Final = object() diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index fa6c0b30413..bff7ded3114 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -10,7 +10,7 @@ from litellm.exceptions import APIError from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import bindings, configuration, runtime -from litellm.rust_bridge.catalog import Context, Delivery, Route, Rule +from litellm.rust_bridge.catalog import Delivery, Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout @@ -39,7 +39,7 @@ class NativeFn(Protocol): def __call__(self) -> str: ... -CONTEXT: Final = Context(Route.MESSAGES, provider="anthropic", model="model") +CONTEXT: Final = RouteContext(Route.MESSAGES, provider="anthropic", model="model") RUST: Final = "rust" PYTHON: Final = "python" @@ -50,8 +50,8 @@ def binding(native: NativeFn | None) -> bindings.NativeBinding[NativeFn]: return bound -def rules(rollout: Rollout) -> tuple[Rule, ...]: - return (Rule(Route.MESSAGES, rollout, providers=frozenset({"anthropic"})),) +def rules(rollout: Rollout) -> tuple[RouteRule, ...]: + return (RouteRule(Route.MESSAGES, rollout, providers=frozenset({"anthropic"})),) class Recorder: @@ -74,7 +74,7 @@ def recorder(native_effect: BaseException | None = None) -> Recorder: return Recorder(native_effect) -def run(rollout: Rollout, calls: Recorder, *, native_missing: bool = False, context: Context = CONTEXT) -> str: +def run(rollout: Rollout, calls: Recorder, *, native_missing: bool = False, context: RouteContext = CONTEXT) -> str: return runtime.run( context, binding=binding(None if native_missing else calls.rust), @@ -146,8 +146,8 @@ def test_context_outside_rule_stays_on_python() -> None: calls: Final = recorder() configuration.rust(True) - assert run(Rollout.RUST_REQUIRED, calls, context=Context(Route.MESSAGES, provider="openai")) == "python" - assert run(Rollout.RUST_REQUIRED, calls, context=Context(Route.RESPONSES, provider="anthropic")) == "python" + assert run(Rollout.RUST_REQUIRED, calls, context=RouteContext(Route.MESSAGES, provider="openai")) == "python" + assert run(Rollout.RUST_REQUIRED, calls, context=RouteContext(Route.RESPONSES, provider="anthropic")) == "python" assert calls.calls == (PYTHON, PYTHON) @@ -155,20 +155,20 @@ def test_context_outside_rule_stays_on_python() -> None: @pytest.mark.parametrize( "context", ( - Context(Route.CHAT_COMPLETIONS, provider="anthropic"), - Context(Route.CHAT_COMPLETIONS, provider="bedrock"), - Context(Route.RESPONSES, provider="openai"), - Context(Route.TRANSCRIPTION, provider="openai"), + RouteContext(Route.CHAT_COMPLETIONS, provider="anthropic"), + RouteContext(Route.CHAT_COMPLETIONS, provider="bedrock"), + RouteContext(Route.RESPONSES, provider="openai"), + RouteContext(Route.TRANSCRIPTION, provider="openai"), ), ) @pytest.mark.parametrize("delivery", tuple(Delivery)) async def test_shipped_python_routes_never_load_native( - monkeypatch: pytest.MonkeyPatch, context: Context, delivery: Delivery + monkeypatch: pytest.MonkeyPatch, context: RouteContext, delivery: Delivery ) -> None: monkeypatch.setenv("LITELLM_RUST", "1") configuration.rust(True) calls: Final = recorder() - request: Final = Context(context.route, provider=context.provider, delivery=delivery) + request: Final = RouteContext(context.route, provider=context.provider, delivery=delivery) def reject_load(value: object) -> NativeFn | None: pytest.fail("Python-only dispatch must not load a native binding") diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index 023f02cffbb..3a86a69ed8b 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -1,12 +1,8 @@ -import dataclasses import logging -from pathlib import Path from typing import Final import httpx import pytest -from pydantic import TypeAdapter -from typing_extensions import ReadOnly, TypedDict import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager @@ -15,34 +11,6 @@ from litellm.rust_bridge import settings from litellm.secret_managers.main import get_secret_str from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem -CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" - - -class SettingSpec(TypedDict): - adapter: ReadOnly[str] - required: ReadOnly[bool] - precedence: ReadOnly[str] - sensitive: ReadOnly[bool] - shapes: ReadOnly[list[str]] - unsupported_live: ReadOnly[str | None] - - -class SettingsGroup(TypedDict): - version: ReadOnly[int] - fields: ReadOnly[dict[str, SettingSpec]] - - -def test_the_rust_contract_matches_the_returned_fields() -> None: - contract: Final = TypeAdapter(dict[str, SettingsGroup]).validate_json(CONTRACT_PATH.read_text()) - - assert {name: tuple(group["fields"]) for name, group in contract.items()} == { - "http_settings": tuple(field.name for field in dataclasses.fields(settings.http_settings())), - "url_policy": tuple(field.name for field in dataclasses.fields(settings.url_policy())), - "provider_defaults": tuple(field.name for field in dataclasses.fields(settings.provider_defaults())), - "secret_manager": tuple(field.name for field in dataclasses.fields(settings.secret_manager())), - } - - def test_url_policy_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "user_url_validation", False) monkeypatch.setattr(litellm, "user_url_allowed_hosts", ["docs.internal:8443"]) @@ -140,6 +108,75 @@ def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.Mon assert settings.secret_manager() == settings.SecretManager(readable=False) +def test_secret_manager_projects_custom_settings(monkeypatch: pytest.MonkeyPatch) -> None: + manager_settings: Final = KeyManagementSettings( + access_mode="read_and_write", + hosted_keys=["MISTRAL_API_KEY"], + primary_secret_name="primary", + aws_region_name="us-east-1", + ) + client: Final = _VaultSecrets({"MISTRAL_API_KEY": "vault-key"}) + monkeypatch.setattr(litellm, "secret_manager_client", client) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", manager_settings) + + assert settings.secret_manager_binding() == settings.SecretManagerBinding( + system="custom", + access_mode="read_and_write", + hosted_keys=["MISTRAL_API_KEY"], + primary_secret_name="primary", + store_virtual_keys=manager_settings.store_virtual_keys, + prefix_for_stored_virtual_keys=manager_settings.prefix_for_stored_virtual_keys, + kms_key_id=manager_settings.kms_key_id, + custom_secret_manager=manager_settings.custom_secret_manager, + aws_region_name="us-east-1", + aws_role_name=manager_settings.aws_role_name, + aws_session_name=manager_settings.aws_session_name, + aws_external_id=manager_settings.aws_external_id, + aws_profile_name=manager_settings.aws_profile_name, + aws_web_identity_token=manager_settings.aws_web_identity_token, + aws_sts_endpoint=manager_settings.aws_sts_endpoint, + replica_regions=manager_settings.replica_regions, + client=client, + settings_object=manager_settings, + ) + + +def test_secret_manager_without_a_client_has_no_system(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "secret_manager_client", None) + + assert settings.secret_manager_binding().system is None + + +def test_secret_manager_uses_key_management_defaults(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "secret_manager_client", None) + monkeypatch.setattr(litellm, "_key_management_settings", None) + + defaults: Final = KeyManagementSettings() + result: Final = settings.secret_manager_binding() + + assert result == settings.SecretManagerBinding( + system=None, + access_mode=defaults.access_mode, + hosted_keys=defaults.hosted_keys, + primary_secret_name=defaults.primary_secret_name, + store_virtual_keys=defaults.store_virtual_keys, + prefix_for_stored_virtual_keys=defaults.prefix_for_stored_virtual_keys, + kms_key_id=defaults.kms_key_id, + custom_secret_manager=defaults.custom_secret_manager, + aws_region_name=defaults.aws_region_name, + aws_role_name=defaults.aws_role_name, + aws_session_name=defaults.aws_session_name, + aws_external_id=defaults.aws_external_id, + aws_profile_name=defaults.aws_profile_name, + aws_web_identity_token=defaults.aws_web_identity_token, + aws_sts_endpoint=defaults.aws_sts_endpoint, + replica_regions=defaults.replica_regions, + client=None, + settings_object=None, + ) + + def test_provider_defaults_read_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "vertex_project", "configured-project") monkeypatch.setattr(litellm, "vertex_location", "europe-west4") diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py index 71aa79cc4bb..3da291c898d 100644 --- a/tests/test_litellm/rust_bridge/test_token_counter.py +++ b/tests/test_litellm/rust_bridge/test_token_counter.py @@ -12,25 +12,32 @@ from types import MappingProxyType from typing import Final import pytest -import tiktoken -from tokenizers import Tokenizer import litellm from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding -from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens +from litellm.proxy.spend_tracking.input_tokens import count_input_tokens, count_input_tokens_for_model from litellm.rust_bridge import bindings, configuration from litellm.rust_bridge import token_counter as bridge +from litellm.rust_bridge import tokenizer as tokenizer_dispatch +from litellm.rust_bridge._native import Tokenizer from litellm.utils import claude_json_str MODEL: Final = "claude-sonnet-4-5-20250929" CL100K_MODEL: Final = "gpt-4" O200K_MODEL: Final = "gpt-4o" +MODEL_BY_TOKENIZER: Final[MappingProxyType[bridge.RustTokenizer, str]] = MappingProxyType( + {"anthropic": MODEL, "cl100k_base": CL100K_MODEL, "o200k_base": O200K_MODEL} +) TOKENIZERS: Final[tuple[bridge.RustTokenizer, ...]] = ("anthropic", "cl100k_base", "o200k_base") -RANK_FILE_LINES: Final = MappingProxyType({"cl100k_base": 100_256, "o200k_base": 199_998}) BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode() +def _counted(body: dict[str, object], model: str) -> tuple[bytes, dict[str, object]]: + raw: Final = json.dumps({**body, "model": model}).encode() + return raw, json.loads(raw) + + class _FakeDeclined(Exception): pass @@ -39,14 +46,41 @@ class _FakeUpstream(Exception): pass +class _FakeTokenizer: + """Stands in for one shared native `Tokenizer`; only its name identifies it.""" + + def __init__(self, name: str, json: str | None = None) -> None: + self.name = name + self.json = json + + +def _fake_native_tokenizers(monkeypatch: pytest.MonkeyPatch, anthropic_json: str | None = None) -> None: + """Point the counter's tokenizer lookups at fakes while the bridge is faked; the codec path + keeps falling back to Python. Parity tests that restore the real extension get the real + lookups back.""" + fakes: Final = {name: _FakeTokenizer(name) for name in ("cl100k_base", "o200k_base")} + anthropic: Final = _FakeTokenizer("anthropic", anthropic_json) + real_encoding: Final = tokenizer_dispatch.native_encoding + real_anthropic: Final = tokenizer_dispatch.native_anthropic + + def faked() -> bool: + return isinstance(bindings.get_native_bridge(), _FakeNative) + + monkeypatch.setattr( + tokenizer_dispatch, "native_encoding", lambda name: fakes[name] if faked() else real_encoding(name) + ) + monkeypatch.setattr(tokenizer_dispatch, "native_anthropic", lambda: anthropic if faked() else real_anthropic()) + + class _FakeNative: RustBridgeDeclined = _FakeDeclined RustUpstreamError = _FakeUpstream class _RecordingCounter: - def __init__(self, tokenizer_json: str) -> None: - self.tokenizer_json = tokenizer_json + def __init__(self, tokenizer: _FakeTokenizer, fast: bool) -> None: + self.tokenizer = tokenizer + self.fast = fast self.bodies: list[bytes] = [] async def acount_request(self, body: bytes) -> object: @@ -55,25 +89,16 @@ class _RecordingCounter: class _RecordingFactory: - """Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files.""" + """Stands in for the native `TokenCounter` class, built over a loaded `Tokenizer`.""" def __init__(self) -> None: self.counters: list[_RecordingCounter] = [] - self.rank_files: list[str] = [] - def __call__(self, tokenizer_json: str) -> _RecordingCounter: - counter = _RecordingCounter(tokenizer_json) + def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _RecordingCounter: + counter = _RecordingCounter(tokenizer, fast) self.counters.append(counter) return counter - def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter: - self.rank_files.append(rank_file) - return self("cl100k_base") - - def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter: - self.rank_files.append(rank_file) - return self("o200k_base") - class _RaisingCounter: def __init__(self, error: Exception) -> None: @@ -89,13 +114,7 @@ class _RaisingFactory: def __init__(self, error: Exception) -> None: self.error = error - def __call__(self, tokenizer_json: str) -> _RaisingCounter: - return _RaisingCounter(self.error) - - def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter: - return _RaisingCounter(self.error) - - def from_o200k_ranks(self, rank_file: str) -> _RaisingCounter: + def from_tokenizer(self, tokenizer: _FakeTokenizer, fast: bool = False) -> _RaisingCounter: return _RaisingCounter(self.error) @@ -105,6 +124,7 @@ def _reset_bridge(monkeypatch: pytest.MonkeyPatch): bridge._counter.cache_clear() configuration.reset_rust_configuration() monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative()) + _fake_native_tokenizers(monkeypatch, anthropic_json=claude_json_str) yield bridge.TOKEN_COUNTER.reset() bridge._counter.cache_clear() @@ -117,8 +137,12 @@ async def test_disabled_bridge_never_constructs_a_counter(tokenizer: bridge.Rust factory: Final = _RecordingFactory() litellm.rust(False) bridge.TOKEN_COUNTER.override(factory) + model: Final = MODEL_BY_TOKENIZER[tokenizer] + raw, request_body = _counted({"messages": [{"role": "user", "content": "hello"}]}, model) - assert await bridge.count_input_tokens(BODY, tokenizer) is None + counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw, models=(model,)) + + assert counts[model] == count_input_tokens_for_model(request_body=request_body, model=model) assert factory.counters == [] @@ -128,31 +152,33 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No litellm.rust(True) bridge.TOKEN_COUNTER.override(factory) - first: Final = await bridge.count_input_tokens(BODY, "anthropic") - second: Final = await bridge.count_input_tokens(BODY, "anthropic") + first: Final = await bridge.native_count(factory, "anthropic", BODY) + second: Final = await bridge.native_count(factory, "anthropic", BODY) assert first == bridge.InputTokenCount(model=MODEL, input_tokens=42) assert second == first assert len(factory.counters) == 1 assert factory.counters[0].bodies == [BODY, BODY] - assert json.loads(factory.counters[0].tokenizer_json)["model"]["type"] == "BPE" + assert factory.counters[0].fast is False + assert factory.counters[0].tokenizer is tokenizer_dispatch.native_anthropic() + assert json.loads(factory.counters[0].tokenizer.json or "")["model"]["type"] == "BPE" @pytest.mark.asyncio @pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base")) -async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokenizer: bridge.RustTokenizer) -> None: +async def test_tiktoken_counter_is_built_over_the_shared_encoding_once(tokenizer: bridge.RustTokenizer) -> None: factory: Final = _RecordingFactory() litellm.rust(True) bridge.TOKEN_COUNTER.override(factory) - first: Final = await bridge.count_input_tokens(BODY, tokenizer) - second: Final = await bridge.count_input_tokens(BODY, tokenizer) + first: Final = await bridge.native_count(factory, tokenizer, BODY) + second: Final = await bridge.native_count(factory, tokenizer, BODY) assert first == second == bridge.InputTokenCount(model=MODEL, input_tokens=42) - assert len(factory.rank_files) == 1 - assert factory.rank_files[0].startswith("IQ== 0\n") - assert factory.rank_files[0].count("\n") == RANK_FILE_LINES[tokenizer] - assert factory.counters[0].tokenizer_json == tokenizer + assert len(factory.counters) == 1 + assert factory.counters[0].tokenizer.name == tokenizer + assert factory.counters[0].tokenizer is tokenizer_dispatch.native_encoding(tokenizer) + assert factory.counters[0].fast is False assert factory.counters[0].bodies == [BODY, BODY] @@ -162,13 +188,13 @@ async def test_each_tokenizer_gets_its_own_cached_counter() -> None: litellm.rust(True) bridge.TOKEN_COUNTER.override(factory) - await bridge.count_input_tokens(BODY, "anthropic") - await bridge.count_input_tokens(BODY, "cl100k_base") - await bridge.count_input_tokens(BODY, "o200k_base") - await bridge.count_input_tokens(BODY, "anthropic") - await bridge.count_input_tokens(BODY, "o200k_base") + await bridge.native_count(factory, "anthropic", BODY) + await bridge.native_count(factory, "cl100k_base", BODY) + await bridge.native_count(factory, "o200k_base", BODY) + await bridge.native_count(factory, "anthropic", BODY) + await bridge.native_count(factory, "o200k_base", BODY) - assert [counter.tokenizer_json for counter in factory.counters][1:] == ["cl100k_base", "o200k_base"] + assert [counter.tokenizer.name for counter in factory.counters] == ["anthropic", "cl100k_base", "o200k_base"] assert [len(counter.bodies) for counter in factory.counters] == [2, 1, 2] @@ -176,8 +202,11 @@ async def test_each_tokenizer_gets_its_own_cached_counter() -> None: async def test_missing_native_module_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: litellm.rust(True) monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) + raw, request_body = _counted({"messages": [{"role": "user", "content": "hello"}]}, MODEL) - assert [await bridge.count_input_tokens(BODY, tokenizer) for tokenizer in TOKENIZERS] == [None, None, None] + counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw, models=(MODEL,)) + + assert counts[MODEL] == count_input_tokens_for_model(request_body=request_body, model=MODEL) @pytest.mark.asyncio @@ -185,8 +214,12 @@ async def test_missing_native_module_falls_back(monkeypatch: pytest.MonkeyPatch) async def test_declined_request_falls_back(tokenizer: bridge.RustTokenizer) -> None: litellm.rust(True) bridge.TOKEN_COUNTER.override(_RaisingFactory(_FakeDeclined("request has no messages"))) + model: Final = MODEL_BY_TOKENIZER[tokenizer] + raw, request_body = _counted({"messages": [{"role": "user", "content": "hello"}]}, model) - assert await bridge.count_input_tokens(BODY, tokenizer) is None + counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw, models=(model,)) + + assert counts[model] == count_input_tokens_for_model(request_body=request_body, model=model) @pytest.mark.asyncio @@ -194,8 +227,12 @@ async def test_declined_request_falls_back(tokenizer: bridge.RustTokenizer) -> N async def test_runtime_failure_falls_back(tokenizer: bridge.RustTokenizer) -> None: litellm.rust(True) bridge.TOKEN_COUNTER.override(_RaisingFactory(RuntimeError("encode failed"))) + model: Final = MODEL_BY_TOKENIZER[tokenizer] + raw, request_body = _counted({"messages": [{"role": "user", "content": "hello"}]}, model) - assert await bridge.count_input_tokens(BODY, tokenizer) is None + counts: Final = await count_input_tokens(request_body=request_body, raw_body=raw, models=(model,)) + + assert counts[model] == count_input_tokens_for_model(request_body=request_body, model=model) @pytest.mark.parametrize( @@ -273,8 +310,8 @@ def test_rust_tokenizer_names_the_encoding_python_actually_counts_with(model: st "Hello, world! camelCase ABCdef \u00e9\u00e8 12345 \u3053\u3093\u306b\u3061\u306f <|endoftext|>\r\n" * 9 ) python_count: Final = litellm.token_counter(model=model, text=text) - cl100k_count: Final = len(tiktoken.get_encoding("cl100k_base").encode(text, disallowed_special=())) - o200k_count: Final = len(tiktoken.get_encoding("o200k_base").encode(text, disallowed_special=())) + cl100k_count: Final = Tokenizer.from_tiktoken("cl100k_base").count(text) + o200k_count: Final = Tokenizer.from_tiktoken("o200k_base").count(text) assert cl100k_count != o200k_count match bridge.rust_tokenizer(model): case "cl100k_base": @@ -282,7 +319,7 @@ def test_rust_tokenizer_names_the_encoding_python_actually_counts_with(model: st case "o200k_base": assert python_count == o200k_count case "anthropic": - assert python_count == len(Tokenizer.from_str(claude_json_str).encode(text).ids) + assert python_count == Tokenizer.from_json(claude_json_str).count(text) assert python_count not in {cl100k_count, o200k_count} case None: pytest.fail(f"{model} must have a Rust tokenizer") @@ -386,12 +423,11 @@ async def test_native_count_matches_python_budget_counter( litellm.rust(True) body: Final = json.dumps(request_body).replace(MODEL, model) - rust_count: Final = await bridge.count_input_tokens(body.encode(), tokenizer) - python_count: Final = _count_input_tokens(request_body=json.loads(body), model=model) + request_body_parsed: Final = json.loads(body) + counts: Final = await count_input_tokens(request_body=request_body_parsed, raw_body=body.encode(), models=(model,)) + python_count: Final = count_input_tokens_for_model(request_body=request_body_parsed, model=model) - assert rust_count is not None - assert rust_count.model == json.loads(body).get("model") - assert rust_count.input_tokens == python_count + assert counts[model] == python_count @pytest.mark.asyncio @@ -405,15 +441,14 @@ async def test_tiktoken_counts_long_text_exactly_where_python_chunks( litellm.rust(True) text: Final = "x " * 20_000 body: Final = {"model": model, "messages": [{"role": "user", "content": text}]} - encoding: Final = tiktoken.get_encoding(tokenizer) - exact: Final = 3 + len(encoding.encode("user")) + len(encoding.encode(text)) + 3 + encoding: Final = Tokenizer.from_tiktoken(tokenizer) + exact: Final = 3 + encoding.count("user") + encoding.count(text) + 3 chunks: Final = -(-len(text) // TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS) - rust_count: Final = await bridge.count_input_tokens(json.dumps(body).encode(), tokenizer) - python_count: Final = _count_input_tokens(request_body=body, model=model) + counts: Final = await count_input_tokens(request_body=body, raw_body=json.dumps(body).encode(), models=(model,)) + python_count: Final = count_input_tokens_for_model(request_body=body, model=model) - assert rust_count is not None - assert rust_count.input_tokens == exact + assert counts[model] == exact assert python_count is not None assert exact < python_count <= exact + chunks @@ -440,5 +475,9 @@ async def test_native_declines_shapes_python_prices_differently( native: Final = pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) litellm.rust(True) + model: Final = MODEL_BY_TOKENIZER[tokenizer] + raw, parsed = _counted(request_body, model) - assert await bridge.count_input_tokens(json.dumps(request_body).encode(), tokenizer) is None + counts: Final = await count_input_tokens(request_body=parsed, raw_body=raw, models=(model,)) + + assert counts.get(model) == count_input_tokens_for_model(request_body=parsed, model=model) diff --git a/tests/test_litellm/rust_bridge/test_tokenizer.py b/tests/test_litellm/rust_bridge/test_tokenizer.py new file mode 100644 index 00000000000..0de7ad50b1e --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_tokenizer.py @@ -0,0 +1,134 @@ +from collections.abc import Generator +from typing import Final + +import pytest +import tiktoken +from tokenizers import Tokenizer + +import litellm +from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding +from litellm.rust_bridge import configuration, tokenizer +from litellm.utils import _select_tokenizer +from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON + + +@pytest.fixture(autouse=True) +def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: + monkeypatch.delenv("LITELLM_RUST", raising=False) + configuration.reset_rust_configuration() + yield + tokenizer.TOKENIZER.reset() + configuration.reset_rust_configuration() + + +@pytest.mark.parametrize("environment", (None, "0", "1")) +@pytest.mark.parametrize("process", (None, False, True)) +def test_tokenizer_factories_follow_rollout( + monkeypatch: pytest.MonkeyPatch, environment: str | None, process: bool | None +) -> None: + configuration.rust(process) + if environment is not None: + monkeypatch.setenv("LITELLM_RUST", environment) + enabled: Final = environment == "1" if environment is not None else process is True + encoding: Final = tokenizer.get_encoding("cl100k_base") + custom: Final = litellm.create_tokenizer(TOKENIZER_JSON) + reference: Final = Tokenizer.from_str(TOKENIZER_JSON) + + assert isinstance(encoding, OpenAIEncoding if enabled else tiktoken.Encoding) + assert isinstance(custom["tokenizer"], HuggingFaceTokenizer if enabled else Tokenizer) + assert encoding.encode("café 漢字 🙂") == tiktoken.get_encoding(encoding.name).encode("café 漢字 🙂") + assert litellm.encode(text="Hello World", custom_tokenizer=custom) == reference.encode("Hello World").ids + assert litellm.token_counter(text="Hello World", custom_tokenizer=custom) == len(reference.encode("Hello World")) + + +def test_missing_native_binding_keeps_python_tokenizer_api() -> None: + configuration.rust(True) + tokenizer.TOKENIZER.override(None) + encoding: Final = tokenizer.get_encoding("cl100k_base") + custom: Final = litellm.create_tokenizer(TOKENIZER_JSON)["tokenizer"] + + assert isinstance(encoding, tiktoken.Encoding) + assert isinstance(custom, Tokenizer) + custom.enable_padding(pad_id=0, pad_token="[UNK]") + assert [item.ids for item in custom.encode_batch(["Hello", "Hello World"])] == [[3, 1, 0], [3, 1, 2]] + + +def test_cached_selection_follows_backend_changes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True) + configuration.rust(True) + native: Final = _select_tokenizer("dispatch-fixture")["tokenizer"] + configuration.rust(False) + python: Final = _select_tokenizer("dispatch-fixture")["tokenizer"] + + assert isinstance(native, OpenAIEncoding) + assert isinstance(python, tiktoken.Encoding) + assert native.encode("hello") == python.encode("hello") + + +def test_declined_native_factory_falls_back_before_tokenizing() -> None: + from litellm.rust_bridge._native import RustBridgeDeclined + + class UnavailableTokenizer: + @staticmethod + def from_json(json: str) -> None: + raise RustBridgeDeclined("huggingface feature is disabled") + + configuration.rust(True) + binding: Final = tokenizer._as_factory(UnavailableTokenizer) + tokenizer.TOKENIZER.override(binding) + custom: Final = litellm.create_tokenizer(TOKENIZER_JSON) + + assert isinstance(custom["tokenizer"], Tokenizer) + assert ( + litellm.decode(tokens=litellm.encode(text="Hello World", custom_tokenizer=custom), custom_tokenizer=custom) + == "Hello World" + ) + + +@pytest.mark.parametrize( + ("model", "text"), + ( + ("gpt-4o", "hello <|endoftext|> world"), + ("gpt-3.5-turbo", "café 漢字 🙂"), + ("text-davinci-003", " def f():\n return 1\n"), + ("tokenizer-parity-fixture", "hello again"), + ), +) +def test_public_token_api_is_identical_across_backends(monkeypatch: pytest.MonkeyPatch, model: str, text: str) -> None: + """`litellm.token_counter`, `encode` and `decode` return the same values whichever backend + the catalog picks; only the object types differ.""" + monkeypatch.setattr(litellm, "anthropic_models", {*litellm.anthropic_models, "tokenizer-parity-fixture"}) + messages: Final = [{"role": "user", "content": text}, {"role": "assistant", "content": "ok"}] + + def observe() -> tuple[int, int, list[int], str]: + ids: Final = litellm.encode(model=model, text=text) + return ( + litellm.token_counter(model=model, text=text), + litellm.token_counter(model=model, messages=messages), + ids, + litellm.decode(model=model, tokens=ids), + ) + + configuration.rust(False) + python: Final = observe() + configuration.rust(True) + rust: Final = observe() + + assert rust == python + + +def test_cached_huggingface_tokenizers_follow_backend_changes(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer as RustHuggingFaceTokenizer + from litellm.utils import _load_huggingface_tokenizer + + monkeypatch.setattr(litellm, "anthropic_models", {*litellm.anthropic_models, "tokenizer-cache-fixture"}) + _load_huggingface_tokenizer.cache_clear() + configuration.rust(True) + native: Final = _select_tokenizer("tokenizer-cache-fixture")["tokenizer"] + configuration.rust(False) + python: Final = _select_tokenizer("tokenizer-cache-fixture")["tokenizer"] + configuration.rust(True) + + assert isinstance(native, RustHuggingFaceTokenizer) + assert isinstance(python, Tokenizer) + assert _select_tokenizer("tokenizer-cache-fixture")["tokenizer"] is native diff --git a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py index e449d4392d8..0eae041f535 100644 --- a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py +++ b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py @@ -56,10 +56,11 @@ def _write_wheel( metadata_tags: tuple[str, ...] | None = (_EXPECTED_TAG,), dist_info: str = _DIST_INFO, duplicate_wheel: bool = False, + native_bytes: bytes = b"synthetic native extension", ) -> Path: wheel: Final = tmp_path / f"litellm-1.100.0-{filename_tag}.whl" with zipfile.ZipFile(wheel, "w", compression=zipfile.ZIP_DEFLATED) as archive: - archive.writestr(_NATIVE_MEMBER, b"synthetic native extension") + archive.writestr(_NATIVE_MEMBER, native_bytes) archive.writestr( f"{dist_info}/METADATA", "Metadata-Version: 2.1\nName: litellm\nVersion: 1.100.0\n", @@ -195,3 +196,17 @@ def test_rejects_production_module_exposing_panic_hook(tmp_path: Path) -> None: wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) assert _run_verifier(wheel, exposes_panic=True) == 1 + + +@pytest.mark.parametrize("embedded", (False, True)) +def test_vocabulary_is_packaged_once(tmp_path: Path, embedded: bool) -> None: + ranks: Final = b"AA== 0\nAQ== 1\nAg== 2\n" + wheel: Final = _write_wheel( + tmp_path, + filename_tag=_EXPECTED_TAG, + native_bytes=b"native engine" + (ranks if embedded else b""), + ) + with zipfile.ZipFile(wheel, "a") as archive: + archive.writestr("litellm/litellm_core_utils/tokenizers/" + "a" * 40, ranks) + + assert _run_verifier(wheel) == (1 if embedded else 0) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 1d6c229f9ce..bdeedb5f38c 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1,3 +1,4 @@ +import datetime import time from typing import Final @@ -29,6 +30,7 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, Usage, ) +from litellm.types.videos.main import VideoObject @pytest.fixture @@ -3038,21 +3040,23 @@ def test_completion_cost_logs_cache_and_reasoning_breakdown_for_custom_pricing() assert total == pytest.approx(100 * 1e-6 + 800 * 1e-7 + 100 * 1.25e-6 + 500 * 2e-6) -def test_cost_per_token_per_second_pricing(monkeypatch): +@pytest.mark.parametrize("custom_llm_provider", ["together_ai", "openai", "anthropic", "bedrock", "azure"]) +def test_cost_per_token_per_second_pricing(monkeypatch, custom_llm_provider: str): """ Models priced by duration (input/output_cost_per_second) with no per-token rates - must be billed as cost_per_second * response_time_ms / 1000 in cost_per_token. + must be billed as cost_per_second * response_time_ms / 1000 in cost_per_token, + whether or not the provider has its own cost calculator. """ monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - model = "test-per-second-pricing-model" + model = f"test-per-second-pricing-{custom_llm_provider}" litellm.register_model( model_cost={ model: { "input_cost_per_second": 0.02, "output_cost_per_second": 0.04, - "litellm_provider": "together_ai", + "litellm_provider": custom_llm_provider, "mode": "chat", } } @@ -3060,7 +3064,7 @@ def test_cost_per_token_per_second_pricing(monkeypatch): prompt_cost, completion_cost_value = cost_per_token( model=model, - custom_llm_provider="together_ai", + custom_llm_provider=custom_llm_provider, prompt_tokens=10, completion_tokens=20, response_time_ms=1500.0, @@ -3070,6 +3074,143 @@ def test_cost_per_token_per_second_pricing(monkeypatch): assert completion_cost_value == pytest.approx(0.04 * 1.5) +def test_cost_per_token_keeps_token_pricing_when_per_second_rates_are_also_set(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + model = "test-token-and-per-second-pricing-model" + litellm.register_model( + model_cost={ + model: { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "input_cost_per_second": 0.02, + "output_cost_per_second": 0.04, + "litellm_provider": "openai", + "mode": "chat", + } + } + ) + + prompt_cost, completion_cost_value = cost_per_token( + model=model, + custom_llm_provider="openai", + prompt_tokens=10, + completion_tokens=20, + response_time_ms=1500.0, + ) + + assert prompt_cost == pytest.approx(10 * 1e-6) + assert completion_cost_value == pytest.approx(20 * 2e-6) + + +def _logging_obj_with_call_window(duration_ms: float) -> Logging: + start_time: Final = datetime.datetime(2026, 9, 21, 12, 0, 0) + logging_obj: Final = Logging( + model="gpt-5.4-nano", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=start_time, + litellm_call_id="per-second-call-window", + function_id="f", + ) + logging_obj.model_call_details["start_time"] = start_time + logging_obj.model_call_details["end_time"] = start_time + datetime.timedelta(milliseconds=duration_ms) + return logging_obj + + +@pytest.mark.parametrize( + ("stamped_response_ms", "total_time", "logged_duration_ms", "expected_seconds"), + [(None, 0.0, 1500.0, 1.5), (3000.0, 0.0, 1500.0, 3.0), (None, 2500.0, 1500.0, 2.5), (3000.0, 2500.0, 1500.0, 3.0)], +) +def test_completion_cost_per_second_deployment_bills_the_call_duration( + monkeypatch, + stamped_response_ms: float | None, + total_time: float, + logged_duration_ms: float, + expected_seconds: float, +): + """ + A deployment priced only per second bills the stamped ``_response_ms`` when there is one, + then the caller's explicit ``total_time``, and the logging object's start/end window otherwise + (a streamed response is never stamped). + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + deployment_id = "per-second-openai-deployment" + litellm.register_model( + model_cost={ + deployment_id: { + "input_cost_per_second": 0.02, + "output_cost_per_second": 0.04, + "litellm_provider": "openai", + "mode": "chat", + } + } + ) + response = ModelResponse( + model="gpt-5.4-nano", + usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18), + ) + response._response_ms = stamped_response_ms + + cost = completion_cost( + completion_response=response, + model="openai/gpt-5.4-nano", + custom_llm_provider="openai", + custom_pricing=True, + router_model_id=deployment_id, + total_time=total_time, + litellm_logging_obj=_logging_obj_with_call_window(logged_duration_ms), + ) + + assert cost == pytest.approx((0.02 + 0.04) * expected_seconds) + + +@pytest.mark.parametrize("mode", ["audio_transcription", "audio_speech", "video_generation", "realtime"]) +def test_cost_per_token_leaves_media_second_rates_to_their_dedicated_paths(monkeypatch, mode: str): + """ + A media-mode entry's per-second rates price audio or video seconds, which the dedicated + transcription, speech, video, and realtime paths bill from the media itself, so a call that + reaches the generic path with only a wall-clock duration must not bill them. + """ + model = f"test-media-per-second-{mode}" + monkeypatch.setitem( + litellm.model_cost, + model, + {"input_cost_per_second": 0.02, "output_cost_per_second": 0.4, "litellm_provider": "openai", "mode": mode}, + ) + + assert cost_per_token(model=model, custom_llm_provider="openai", response_time_ms=2000.0) == (0.0, 0.0) + + +def test_completion_cost_video_status_poll_bills_nothing_on_a_per_second_video_model(monkeypatch): + """ + Polling a video job returns a ``VideoObject`` with no stamped duration, so the cost path falls + back to the logging object's call window; on a video model priced per output second that + window must not be billed, or every status poll would charge for the seconds it took to answer. + """ + model = "test-veo-per-second-poll" + monkeypatch.setitem( + litellm.model_cost, + model, + {"output_cost_per_second": 0.4, "litellm_provider": "vertex_ai", "mode": "video_generation"}, + ) + video = VideoObject(id="video_1", object="video", status="completed", model=model, progress=100) + + cost = completion_cost( + completion_response=video, + model=model, + custom_llm_provider="vertex_ai", + call_type=CallTypes.video_retrieve.value, + litellm_logging_obj=_logging_obj_with_call_window(2000.0), + ) + + assert cost == 0.0 + + def _batch_cache_usage() -> Usage: return Usage( prompt_tokens=11000, @@ -3574,6 +3715,45 @@ def test_completion_cost_mantle_native_messages_prices_haiku_from_the_mantle_row ) == pytest.approx(expected), model +def test_completion_cost_legacy_mantle_route_prices_after_router_registration(local_model_cost_map): + """The proxy registers every deployment under its provider-prefixed key at boot. A + bedrock/mantle/ deployment must resolve to the bare Bedrock row there, otherwise the boot + entry is a cost-less capability rule that shadows the priced row and every call on the deployment, + /v1/chat/completions and /v1/messages alike, bills $0.""" + from litellm import Router + + Router( + model_list=[ + { + "model_name": "claude-sonnet-5", + "litellm_params": { + "model": "bedrock/mantle/anthropic.claude-sonnet-5", + "aws_region_name": "us-east-1", + }, + } + ] + ) + assert "bedrock/mantle/anthropic.claude-sonnet-5" not in litellm.model_cost + + response = litellm.ModelResponse( + id="msg_x", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="claude-sonnet-5", + usage={"prompt_tokens": 16, "completion_tokens": 4, "total_tokens": 20}, + ) + row = litellm.model_cost["anthropic.claude-sonnet-5"] + expected = 16 * row["input_cost_per_token"] + 4 * row["output_cost_per_token"] + assert expected > 0 + + for call_type in ("completion", "anthropic_messages"): + assert litellm.completion_cost( + completion_response=response, + model="mantle/anthropic.claude-sonnet-5", + custom_llm_provider="bedrock", + call_type=call_type, + ) == pytest.approx(expected), call_type + + def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map): """An explicit base_model keeps pricing on that model's own key even when the request carries a region with different regional rates, so the private provider model never widens region pricing.""" diff --git a/tests/test_litellm/test_cost_map_guard.py b/tests/test_litellm/test_cost_map_guard.py index 1b4330ed62c..1595bd9864f 100644 --- a/tests/test_litellm/test_cost_map_guard.py +++ b/tests/test_litellm/test_cost_map_guard.py @@ -114,6 +114,40 @@ def test_unclassified_entry_key_is_reported() -> None: assert "Unclassified keys" in failure and "weird_thing" in failure +STALE_HEAD: Final = _snapshot(BASE_MAP, backup=_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)}), schema="{}") +CODE_ONLY: Final = ( + "litellm/utils.py", + "tests/test_litellm/test_utils.py", + "docs/model_prices_and_context_window.json", +) + + +def test_human_pr_that_leaves_the_cost_map_alone_skips_the_file_checks() -> None: + unparseable: Final = guard.Snapshot(cost_map="{not json", backup="", schema="") + assert _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=False) == () + assert _failures(STALE_HEAD, changed_files=(), bot=False) == () + assert _failures(unparseable, changed_files=CODE_ONLY, bot=False) == () + + +@pytest.mark.parametrize("guarded_path", guard.GUARDED_PATHS) +def test_touching_any_cost_map_file_keeps_the_file_checks(guarded_path: str) -> None: + failures: Final = _failures(STALE_HEAD, changed_files=(*CODE_ONLY, guarded_path), bot=False) + assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)] + assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)] + + +def test_bot_pr_always_gets_the_file_checks() -> None: + failures: Final = _failures(STALE_HEAD, changed_files=CODE_ONLY, bot=True) + assert [failure for failure in failures if failure.startswith(guard.BACKUP_PATH)] + assert [failure for failure in failures if failure.startswith(guard.SCHEMA_PATH)] + + +def test_contract_names_the_skip() -> None: + assert guard.contract_for(False, CODE_ONLY) == "human PR, cost map untouched" + assert guard.contract_for(False, (*CODE_ONLY, guard.SCHEMA_PATH)) == "human PR, file checks only" + assert guard.contract_for(True, CODE_ONLY) == "bot contract enforced" + + def test_bot_may_only_touch_the_cost_map_files() -> None: changed = (*guard.GUARDED_PATHS, "litellm/utils.py", ".github/workflows/cost-map-guard.yml") assert _failures(BASE, changed_files=changed, bot=False) == () @@ -147,6 +181,16 @@ def _commit(repo: Path, cost_map: dict[str, object], message: str) -> str: (repo / guard.BACKUP_PATH).parent.mkdir(exist_ok=True) (repo / guard.BACKUP_PATH).write_text(text) (repo / guard.SCHEMA_PATH).write_text(schema_module.render(schema_module.build_schema(cost_map))) + return _git_commit(repo, message) + + +def _commit_code_only(repo: Path, message: str) -> str: + (repo / "litellm").mkdir(exist_ok=True) + (repo / "litellm" / "utils.py").write_text(f"print('{message}')\n") + return _git_commit(repo, message) + + +def _git_commit(repo: Path, message: str) -> str: subprocess.run(("git", "add", "-A"), cwd=repo, check=True) subprocess.run( ("git", "-c", "user.name=t", "-c", "user.email=t@example.com", "commit", "-q", "-m", message), @@ -186,6 +230,40 @@ def test_main_reads_both_revisions_from_git( assert expected_line in result.stdout.splitlines() +def test_main_skips_the_file_checks_on_a_stale_base_the_pr_never_touched(tmp_path: Path) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + _commit(tmp_path, BASE_MAP, "base") + (tmp_path / guard.BACKUP_PATH).write_text(_serialize({**BASE_MAP, "openrouter/b": _entry(3e-06)})) + (tmp_path / guard.SCHEMA_PATH).write_text("{}") + stale_base: Final = _commit_code_only(tmp_path, "stale base with drifted backup and schema") + head: Final = _commit_code_only(tmp_path, "code change on the stale base") + human: Final = _run_guard(tmp_path, stale_base, head, "litellm_fix_pricing") + assert human.returncode == 0, human.stdout + human.stderr + assert "cost map guard passed (human PR, cost map untouched)" in human.stdout.splitlines() + bot: Final = _run_guard(tmp_path, stale_base, head, BOT_REF) + assert bot.returncode == 1 + backup_failure: Final = f"- {guard.BACKUP_PATH} differs from {guard.COST_MAP_PATH}; copy the root file over it" + assert backup_failure in bot.stdout.splitlines() + + +def test_main_keeps_the_file_checks_when_a_cost_map_file_is_renamed(tmp_path: Path) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + base: Final = _commit(tmp_path, BASE_MAP, "base") + subprocess.run(("git", "mv", guard.COST_MAP_PATH, "renamed.json"), cwd=tmp_path, check=True) + head: Final = _git_commit(tmp_path, "rename the cost map") + result: Final = _run_guard(tmp_path, base, head, "litellm_fix_pricing") + assert result.returncode == 1, result.stdout + result.stderr + assert "cost map guard failed (human PR, file checks only):" in result.stdout.splitlines() + + +def test_main_fails_when_the_changed_files_cannot_be_read(tmp_path: Path) -> None: + subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) + head: Final = _commit(tmp_path, BASE_MAP, "head") + result: Final = _run_guard(tmp_path, "0" * 40, head, "litellm_fix_pricing") + assert result.returncode == 1, result.stdout + result.stderr + assert result.stdout.startswith("cost map guard failed: git diff ") + + def test_main_rejects_a_bot_pr_that_edits_code(tmp_path: Path) -> None: subprocess.run(("git", "init", "-q", str(tmp_path)), check=True) base = _commit(tmp_path, BASE_MAP, "base") diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 1e59f4d878e..cf61a6d9f65 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4984,6 +4984,100 @@ async def test_wrapper_async_fires_post_call_failure_deployment_hook_on_internal assert isinstance(recorder.calls[0][1], litellm.AuthenticationError) +def _budget_reservation(callback_bound: bool = False) -> dict: + return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": callback_bound} + + +_BUDGET_RESERVATION_CALL_KWARGS: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} +_BUDGET_RESERVATION_REFUSAL: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o") + + +@pytest.mark.asyncio +async def test_wrapper_async_claims_the_budget_reservation_for_the_cost_callback() -> None: + reservation = _budget_reservation() + + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response="ok", + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is True + + +@pytest.mark.asyncio +async def test_wrapper_async_claims_the_budget_reservation_before_the_stream_is_consumed() -> None: + reservation = _budget_reservation() + + stream = await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response="ok", + stream=True, + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is True + async for _ in stream: + pass + + +@pytest.mark.asyncio +async def test_wrapper_async_claims_the_budget_reservation_a_supplied_logging_object_already_saw() -> None: + reservation = _budget_reservation() + logging_obj, kwargs = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + **_BUDGET_RESERVATION_CALL_KWARGS, + litellm_call_id="proxy-pre-call-setup", + metadata={"user_api_key_budget_reservation": reservation}, + ) + assert reservation["callback_bound"] is False + + await litellm.acompletion(**kwargs, litellm_logging_obj=logging_obj, mock_response="ok") + + assert reservation["callback_bound"] is True + + +@pytest.mark.asyncio +async def test_wrapper_async_hands_the_budget_reservation_back_when_the_call_fails() -> None: + reservation = _budget_reservation() + + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response=_BUDGET_RESERVATION_REFUSAL, + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is False + + +@pytest.mark.asyncio +async def test_wrapper_async_leaves_the_budget_reservation_alone_on_internal_calls() -> None: + claimed_by_the_outer_call = _budget_reservation(callback_bound=True) + never_claimed = _budget_reservation() + + token = is_internal_call.set(True) + try: + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response="ok", + metadata={"user_api_key_budget_reservation": never_claimed}, + ) + with pytest.raises(litellm.AuthenticationError): + await litellm.acompletion( + **_BUDGET_RESERVATION_CALL_KWARGS, + mock_response=_BUDGET_RESERVATION_REFUSAL, + metadata={"user_api_key_budget_reservation": claimed_by_the_outer_call}, + ) + finally: + is_internal_call.reset(token) + + assert never_claimed["callback_bound"] is False + assert claimed_by_the_outer_call["callback_bound"] is True + + @pytest.mark.asyncio async def test_wrapper_async_does_not_fire_failure_hook_for_pre_call_budget_error( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py b/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py index 72e98711f0c..2d5936fabdb 100644 --- a/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py +++ b/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py @@ -13,16 +13,18 @@ GROK_KEY_PREFIXES: Final = ("vertex_ai/xai/grok-", "azure_ai/grok-", "xai/grok-" @pytest.mark.usefixtures("local_model_cost_map") def test_grok_models_with_cache_read_price_advertise_prompt_caching() -> None: - cached_grok_models = tuple( + cached_grok_models: Final = tuple( key for key, entry in litellm.model_cost.items() if key.startswith(GROK_KEY_PREFIXES) and entry.get("cache_read_input_token_cost") ) assert cached_grok_models, "expected at least one grok model with a cache read price" - missing_flag = tuple(key for key in cached_grok_models if supports_prompt_caching(model=key) is not True) + missing_flag: Final = tuple( + key for key in cached_grok_models if get_model_info(model=key).get("supports_prompt_caching") is not True + ) assert missing_flag == (), ( - f"grok models with cache_read_input_token_cost fail supports_prompt_caching: {missing_flag}" + f"grok models with cache_read_input_token_cost fail get_model_info supports_prompt_caching: {missing_flag}" ) @@ -31,8 +33,14 @@ def test_vertex_ai_grok_4_6_supports_prompt_caching_via_get_model_info() -> None routed_model, provider, _, _ = get_llm_provider(model=MODEL) assert (routed_model, provider) == ("xai/grok-4.6", "vertex_ai") - info = get_model_info(model=routed_model, custom_llm_provider=provider) - assert info["litellm_provider"] == "vertex_ai" - assert info.get("supports_prompt_caching") is True + routed_info: Final = get_model_info(model=routed_model, custom_llm_provider=provider) + assert routed_info["litellm_provider"] == "vertex_ai" + assert routed_info.get("supports_prompt_caching") is True + assert routed_info.get("cache_read_input_token_cost") + + catalog_info: Final = get_model_info(model=MODEL) + assert catalog_info["key"] == MODEL + assert catalog_info.get("supports_prompt_caching") is True + assert catalog_info.get("cache_read_input_token_cost") assert supports_prompt_caching(model=MODEL) is True diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index 5f44ba1773e..0a8c9414a0d 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -1,9 +1,16 @@ +import json from typing import Final import pytest -from litellm.types.utils import HiddenParams, all_litellm_params, text_tokens_without_nested_reasoning +from litellm.types.utils import ( + HiddenParams, + ImageObject, + ImageResponse, + all_litellm_params, + text_tokens_without_nested_reasoning, +) def test_rust_is_a_known_litellm_param(): @@ -763,13 +770,70 @@ def test_delta_function_tool_call_unchanged_by_custom_support(): def test_image_response_keeps_background(): """https://github.com/BerriAI/litellm/issues/38649""" - from litellm.types.utils import ImageResponse - response = ImageResponse(created=1, data=[{"b64_json": "aGk="}], background="transparent", output_format="png") assert response.background == "transparent" assert response.model_dump()["background"] == "transparent" +def test_image_response_serialization_honors_dump_options(): + response: Final = ImageResponse( + data=[ + ImageObject( + url="https://example.com/image.png", + provider_specific_fields={"width": 1024, "height": 1536, "content_type": "image/png"}, + ) + ] + ) + expected: Final = [ + { + "url": "https://example.com/image.png", + "provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"}, + } + ] + assert response.model_dump(exclude_none=True)["data"] == expected + assert json.loads(response.model_dump_json(exclude_none=True))["data"] == expected + assert response.model_dump()["data"][0]["provider_specific_fields"] == expected[0]["provider_specific_fields"] + assert "url" not in response.model_dump(exclude={"data": {0: {"url"}}})["data"][0] + assert response.model_dump(include={"data": {"__all__": {"url"}}})["data"] == [ + {"url": "https://example.com/image.png"} + ] + assert response.model_dump(include={"data": {0: True}})["data"] == [ + { + "b64_json": None, + "revised_prompt": None, + "url": "https://example.com/image.png", + "provider_specific_fields": {"width": 1024, "height": 1536, "content_type": "image/png"}, + } + ] + assert response.model_dump(exclude={"data": {0: True}})["data"] == [] + + two_image_response: Final = ImageResponse( + data=[ + ImageObject(url="https://example.com/image.png"), + ImageObject(url="https://example.com/second-image.png"), + ] + ) + assert two_image_response.model_dump(exclude={"data": {1}})["data"] == [ + { + "b64_json": None, + "revised_prompt": None, + "url": "https://example.com/image.png", + "provider_specific_fields": None, + } + ] + assert two_image_response.model_dump(exclude={"data": {-1}})["data"] == [ + { + "b64_json": None, + "revised_prompt": None, + "url": "https://example.com/image.png", + "provider_specific_fields": None, + } + ] + assert two_image_response.model_dump(include={"data": {-1: {"url"}}})["data"] == [ + {"url": "https://example.com/second-image.png"} + ] + + @pytest.mark.parametrize( ("completion_tokens", "text_tokens", "reasoning_tokens", "other_modality_tokens", "expected_text_tokens"), ( diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 51815651eb4..0d3b8ba472d 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -593,7 +593,7 @@ def test_native_projection_errors_never_select_python( import ssl from litellm.rust_bridge import runtime, settings - from litellm.rust_bridge.catalog import Context, Route, Rule + from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR, LiteLLMOcrRequest @@ -621,11 +621,11 @@ def test_native_projection_errors_never_select_python( with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"): runtime.run( - Context(Route.OCR, provider="mistral"), + RouteContext(Route.OCR, provider="mistral"), binding=NATIVE_OCR, native=lambda native: native(request, (), {}), python=python_fallback, - rules=(Rule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),), + rules=(RouteRule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),), ) assert ocr_server.requests == [] diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index e0d92a2e957..0f389edaa27 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -1,19 +1,24 @@ import asyncio import contextvars import gc +import hashlib +import http.server import json +import math import os import threading import time import uuid import weakref -from collections.abc import Generator +from collections.abc import Callable, Generator +from contextlib import ExitStack from datetime import datetime from pathlib import Path from types import SimpleNamespace from typing import Final, Protocol, cast from unittest.mock import Mock from urllib.parse import urlparse +from uuid import uuid4 import boto3 import botocore.config @@ -30,13 +35,22 @@ from litellm.caching.disk_cache import DiskCache from litellm.caching.gcs_cache import GCSCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.caching.redis_semantic_cache import RedisSemanticCache from litellm.caching.s3_cache import S3Cache from litellm.rust_bridge import _native +from litellm.rust_bridge.catalog import CacheRule, Route, RouteRule, SecretManagerRule +from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.response_cache import ResponseCacheRuntime, resolve_response_cache from litellm.types.caching import LiteLLMCacheType +from litellm.types.llms.custom_llm import CustomLLMItem +from litellm.types.utils import EmbeddingResponse from tests.test_litellm_rust.support.fake_gcs import FakeGcs from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.s3_stub import S3Stub +_CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name +_CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name + pytestmark: Final = pytest.mark.requires_rust_extension @@ -49,6 +63,71 @@ def request(key: str = "key") -> dict[str, object]: return {"key": {"preset": key}} +def qdrant_request( + key: str, + messages: list[dict[str, object]], + **kwargs: object, +) -> dict[str, object]: + return {**request(key), "messages": messages, **kwargs} + + +def embedding_vector(text: str) -> list[float]: + raw: Final = hashlib.sha256(text.encode()).digest()[:8] + values: Final = [byte / 127.5 - 1 for byte in raw] + norm: Final = math.sqrt(sum(value * value for value in values)) + return [value / norm for value in values] + + +@pytest.fixture +def qdrant_url() -> str: + value: Final[str | None] = os.environ.get("QDRANT_URL") + if not value: + pytest.skip("QDRANT_URL is required for Qdrant semantic cache tests") + return value.rstrip("/") + + +@pytest.fixture +def fake_embedding_endpoint(monkeypatch: pytest.MonkeyPatch) -> Generator[str]: + class EmbeddingHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self) -> None: + length: Final = int(self.headers["Content-Length"]) + body: Final = json.loads(self.rfile.read(length)) + text: Final = body["input"] + response: Final = { + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": embedding_vector(text), + } + ], + "model": body["model"], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + encoded: Final = json.dumps(response).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, *_args: object) -> None: + return + + server: Final = http.server.ThreadingHTTPServer(("127.0.0.1", 0), EmbeddingHandler) + worker: Final = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + monkeypatch.setenv("OPENAI_API_BASE", f"http://127.0.0.1:{server.server_address[1]}") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + worker.join(timeout=5) + + @pytest.fixture def redis_url() -> Generator[str]: server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis") @@ -113,15 +192,52 @@ def test_existing_constructor_and_global_are_unchanged() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) assert type(facade.cache) is InMemoryCache assert "_native_cache_handle" not in vars(facade) + assert resolve_response_cache(facade) is None with rebound(litellm, "cache", facade): - resolver: Final = _native._CacheTestResolver(litellm) + resolver: Final = _CacheTestResolver(litellm) assert resolver.resolve().kind == "python_callback" resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} +async def test_catalog_constructs_native_runtime_from_public_cache_configuration() -> None: + rules: Final = ( + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + ) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + runtime: Final = resolve_response_cache(facade, rules) + assert isinstance(runtime, ResponseCacheRuntime) + assert runtime.kind == "native" + + sync_request: Final = runtime.request(facade, {"cache_key": "sync"}) + assert sync_request is not None + runtime.store(sync_request, {"answer": 1}) + assert runtime.lookup(sync_request) == {"answer": 1} + assert facade.cache.get_cache("sync") is None + + async_request: Final = runtime.request(facade, {"cache_key": "async"}) + assert async_request is not None + await runtime.async_store(async_request, {"answer": 2}) + assert await runtime.async_lookup(async_request) == {"answer": 2} + assert await facade.cache.async_get_cache("async") is None + + requests: Final = (sync_request, async_request) + expected: Final = { + "values": [{"answer": 1}, {"answer": 2}], + "missing_indices": [], + } + assert runtime.lookup_batch(requests) == expected + assert await runtime.async_lookup_batch(requests) == expected + + await runtime.async_flush() + assert runtime.lookup(sync_request) is None + assert await runtime.async_lookup(async_request) is None + + def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: - resolver: Final = _native._CacheTestResolver(litellm) + resolver: Final = _CacheTestResolver(litellm) enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) enabled: Final = litellm.cache @@ -144,13 +260,13 @@ def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> Non async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: - namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.memory()) - resolver: Final = _native._CacheTestResolver(namespace) + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.memory()) + resolver: Final = _CacheTestResolver(namespace) selected: Final = resolver.resolve() assert selected.kind == "native" selected.store(request(), {"answer": 1}) assert await selected.async_lookup(request()) == {"answer": 1} - with rebound(namespace, "cache", _native._CacheTestHandle.memory()): + with rebound(namespace, "cache", _CacheTestHandle.memory()): replacement: Final = resolver.resolve() await selected.async_store(request(), {"answer": 2}) assert replacement.lookup(request()) is None @@ -183,7 +299,7 @@ async def test_python_callback_preserves_identity_caller_task_context_and_errors raise failure namespace: Final = SimpleNamespace(cache=CustomCache()) - binding: Final = _native._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() assert binding.kind == "python_callback" assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel assert context.get() == "callback" @@ -204,7 +320,7 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: finally: finished.set() - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() async def lookup() -> object: return await binding.async_lookup(None, callback_kwargs={}) @@ -219,9 +335,9 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle: Final = _native._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.memory() handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) native: Final = resolver.resolve() assert native.kind == "native" native.store(request(), {"source": "native"}) @@ -252,12 +368,12 @@ def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not class CustomCache(Cache): pass - handle: Final = _native._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.memory() with pytest.raises(TypeError): handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) facade: Final = Cache(type=LiteLLMCacheType.LOCAL) handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) with rebound(facade, "cache", InMemoryCache()): assert resolver.resolve().kind == "python_callback" with rebound(facade, "ttl", 12): @@ -282,7 +398,7 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: def cyclic_reference() -> weakref.ReferenceType[CustomCache]: callback: Final = CustomCache() namespace: Final = SimpleNamespace(cache=callback) - binding: Final = _native._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() setattr(callback, "binding", binding) return weakref.ref(callback) @@ -293,8 +409,8 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) - namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url, namespace="team")) - binding: Final = _native._CacheTestResolver(namespace).resolve() + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.redis(redis_url, namespace="team")) + binding: Final = _CacheTestResolver(namespace).resolve() response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} client.set("team:sync", str(envelope)) @@ -316,33 +432,31 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd def test_invalid_duration_and_request_shape_fail_before_storage() -> None: - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() for seconds in (-1.0, float("nan"), float("inf")): with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) assert binding.lookup(request()) is None with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - _native._CacheTestHandle.memory(ttl_seconds=-1) + _CacheTestHandle.memory(ttl_seconds=-1) async def test_memory_size_policy_is_applied_by_the_native_host() -> None: - handle: Final = _native._CacheTestHandle.memory(capacity=2, max_entry_bytes=128) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + handle: Final = _CacheTestHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=handle)).resolve() small: Final = {"answer": "ok"} binding.store(request("small"), small) assert await binding.async_lookup(request("small")) == small await binding.async_store(request("large"), {"answer": "x" * 256}) assert binding.lookup(request("large")) is None assert binding.lookup(request("small")) == small - disabled: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.memory(capacity=0)) - ).resolve() + disabled: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory(capacity=0))).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None async def test_native_batch_lookup_and_store_report_partial_hits() -> None: - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() requests: Final = [request("hit"), request("miss"), request("disabled")] requests[2]["controls"] = { "supported_call_type": True, @@ -380,9 +494,7 @@ async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: ) -> object: return result, kwargs - binding: Final = _native._CacheTestResolver( - SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL)) - ).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL))).resolve() assert binding.kind == "python_callback" requests: Final = [request("first"), request("second")] kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}] @@ -410,7 +522,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) cache.cache.set_cache("key", "value") - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=cache)).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=cache)).resolve() assert binding.kind == "python_callback" setattr(cache.cache, "ping", ping) @@ -422,7 +534,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: def test_facade_registration_rejects_mismatched_capacity() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) with pytest.raises(TypeError, match="capacities must match"): - _native._CacheTestHandle.memory(capacity=7)._bind_facade(facade) + _CacheTestHandle.memory(capacity=7)._bind_facade(facade) def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None: @@ -525,19 +637,19 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: redis_flush_size=2, ) with pytest.raises(TypeError, match="default TTLs must match"): - _native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) + _CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) with pytest.raises(TypeError, match="namespaces must match"): - _native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) - _native._CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + _CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) + _CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(redis_url) with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" pool: Final = facade.cache.redis_client.connection_pool with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None @@ -1112,3 +1224,741 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n client.delete("unscoped") client.close() facade.cache.redis_client.close() + + +PARAPHRASE_MARKER: Final = " (paraphrase)" +SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" +SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" +SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset") + + +def _normalized(vector: list[float]) -> list[float]: + norm: Final = math.sqrt(sum(component * component for component in vector)) + return [component / norm for component in vector] + + +def _base_embedding(prompt: str) -> list[float]: + digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest() + return _normalized([float(digest[index] + 1) for index in range(8)]) + + +def _semantic_embedding(prompt: str) -> list[float]: + if PARAPHRASE_MARKER not in prompt: + return _base_embedding(prompt) + base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) + pivot: Final = min(range(8), key=lambda index: abs(base[index])) + direction: Final = _normalized( + [(1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] for index in range(8)] + ) + # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance + return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) + + +class DeterministicEmbedding(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + self.async_calls: list[dict[str, object]] = [] + self.entered = asyncio.Event() + self.gate: asyncio.Event | None = None + + def _respond( + self, + model: str, + input: object, + model_response: EmbeddingResponse, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.calls.append({"model": model, "input": texts}) + model_response.model = model + model_response.data = [ + {"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))} + for index, text in enumerate(texts) + ] + return model_response + + def embedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + return self._respond(model, input, model_response) + + async def aembedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.async_calls.append( + { + "model": model, + "input": texts, + "task": asyncio.current_task(), + "context": SEMANTIC_CONTEXT.get(), + } + ) + SEMANTIC_CONTEXT.set("written-in-aembedding") + self.entered.set() + if self.gate is not None: + await self.gate.wait() + return self._respond(model, input, model_response) + + +@pytest.fixture +def semantic_embedding() -> Generator[DeterministicEmbedding]: + handler: Final = DeterministicEmbedding() + with ExitStack() as stack: + stack.enter_context( + rebound( + litellm, + "custom_provider_map", + [ + *litellm.custom_provider_map, + cast( + CustomLLMItem, + {"provider": "semantic-test", "custom_handler": handler}, + ), + ], + ) + ) + stack.enter_context( + rebound( + litellm, + "_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + ) + ) + stack.enter_context(rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"])) + yield handler + + +@pytest.fixture +def redis_stack() -> Generator[tuple[str, str]]: + url: Final = os.environ.get("LITELLM_REDIS_STACK_URL") + if url is None: + pytest.skip("LITELLM_REDIS_STACK_URL is not set") + index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}" + yield url, index + client: Final = redis.Redis.from_url(url) + try: + client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown + except redis.RedisError: + pass + client.close() + + +def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]: + return { + "key": {"preset": key}, + "messages": [{"role": "user", "content": prompt}], + **extra, + } + + +def semantic_messages(prompt: str) -> list[dict[str, object]]: + return [{"role": "user", "content": prompt}] + + +def semantic_entry_id(prompt: str, tag: str) -> str: + return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest() + + +def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache: + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=similarity_threshold, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade) + return facade + + +def test_redis_semantic_constructor_identity_and_provenance( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + backend: Final = cast(RedisSemanticCache, facade.cache) + assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache" + assert type(backend) is RedisSemanticCache + assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend.similarity_threshold == 0.8 + assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL + handle: Final = cast(object, getattr(facade, "_native_cache_handle")) + assert isinstance(handle, _CacheTestHandle) + assert handle.backend == "redis_semantic" + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + + +def test_redis_semantic_native_and_python_sync_entries_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}} + + binding.store(semantic_request("geo", "what is the capital of france"), response) + + native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}" + stored: Final = client.hgetall(native_hash_key) + assert set(stored) == { + b"entry_id", + b"prompt", + b"response", + b"prompt_vector", + b"inserted_at", + b"updated_at", + b"litellm_cache_key", + }, stored + assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1] + assert stored[b"prompt"] == b"what is the capital of france" + assert stored[b"litellm_cache_key"] == b"geo" + assert len(stored[b"prompt_vector"]) == 32 + decoded: Final = cast(dict[str, object], json.loads(stored[b"response"])) + assert decoded["response"] == response + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "geo", messages=semantic_messages("what is the capital of france") + ) + == decoded + ) + assert semantic_embedding.calls == [ + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["dimension test"]}, + ] + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "math", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}), + messages=semantic_messages("what is 6 times 7"), + ) + python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}" + assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == { + "timestamp": 1700000000.0, + "response": {"answer": 42}, + } + assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42} + client.close() + + +async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + await binding.async_store(semantic_request("async", "name a primary color"), {"answer": "blue"}) + hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" + decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) + python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async", messages=semantic_messages("name a primary color") + ) + assert python_read == decoded + + await binding.async_store_batch( + [ + semantic_request("batch-one", "first batch prompt"), + semantic_request("batch-two", "second batch prompt"), + ], + [{"answer": 1}, {"answer": 2}], + ) + expected: Final = { + key: json.loads(cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response"))) + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ) + } + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ): + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + key, messages=semantic_messages(prompt) + ) + == expected[key] + ), key + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async-python", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), + messages=semantic_messages("python written prompt"), + ) + assert await binding.async_lookup(semantic_request("async-python", "python written prompt")) == {"answer": "python"} + client.close() + + +async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + caller: Final = asyncio.current_task() + SEMANTIC_CONTEXT.set("caller-sentinel") + response: Final = {"choices": [{"text": "paris"}]} + + await binding.async_store(semantic_request("inline", "what is the capital of france"), response) + assert ( + await binding.async_lookup(semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}")) + == response + ) + assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None + assert SEMANTIC_CONTEXT.get() == "written-in-aembedding" + assert semantic_embedding.async_calls == [ + { + "model": "deterministic", + "input": ["what is the capital of france"], + "task": caller, + "context": "caller-sentinel", + }, + { + "model": "deterministic", + "input": [f"what is the capital of france{PARAPHRASE_MARKER}"], + "task": caller, + "context": "written-in-aembedding", + }, + { + "model": "deterministic", + "input": ["python written prompt"], + "task": caller, + "context": "written-in-aembedding", + }, + ], semantic_embedding.async_calls + + +async def test_native_semantic_cancellation_during_embedding_skips_the_backend( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + semantic_embedding.gate = asyncio.Event() + + async def lookup() -> object: + return await binding.async_lookup(semantic_request("cancel", "cancelled prompt")) + + task: Final = asyncio.create_task(lookup()) + await semantic_embedding.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + semantic_embedding.gate.set() + + assert len(semantic_embedding.async_calls) == 1 + assert ( + await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "cancel", messages=semantic_messages("cancelled prompt") + ) + is None + ) + + +def test_redis_semantic_similarity_tag_and_threshold_boundaries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"}) + paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}" + assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"} + assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None + assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None + + strict: Final = semantic_facade(url, index, similarity_threshold=0.99) + strict_binding: Final = _CacheTestResolver(SimpleNamespace(cache=strict)).resolve() + assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None + assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"} + + +def test_redis_semantic_ttl_is_written_only_when_requested( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store({**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1}) + expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" + assert 0 < client.ttl(expiring) <= 12 + + binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2}) + persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}" + assert client.ttl(persistent) == -1 + + binding.store( + {**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5}, + {"answer": 3}, + ) + fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}" + assert client.ttl(fractional) == 2 + client.close() + + +def test_redis_semantic_malformed_response_is_a_miss_for_both_readers( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(semantic_request("bad", "corrupt me"), {"answer": 1}) + hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}" + client.hset(hash_key, "response", b"{not json") + assert binding.lookup(semantic_request("bad", "corrupt me")) is None + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "bad", messages=semantic_messages("corrupt me") + ) + is None + ) + client.close() + + +async def test_redis_semantic_unsupported_operations_raise_not_implemented( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + with pytest.raises(NotImplementedError): + binding.lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_flush() + with pytest.raises(NotImplementedError): + await binding.ping() + + +def test_redis_semantic_requests_without_prompt_are_noops( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(request("plain"), {"answer": 1}) + assert binding.lookup(request("plain")) is None + assert semantic_embedding.calls == [] + assert client.keys(f"{index}:*") == [] + client.close() + + +def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"} + binding.store(scoped, {"answer": "kept"}) + hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}" + assert client.hget(hash_key, "litellm_cache_key") == b"team-a" + assert binding.lookup(scoped) == {"answer": "kept"} + assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None + assert binding.lookup({**scoped, "scope": "team-b"}) is None + client.close() + + +def test_redis_semantic_configuration_drift_falls_back_to_python( + redis_stack: tuple[str, str], + semantic_embedding: DeterministicEmbedding, + monkeypatch: pytest.MonkeyPatch, +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + with rebound(facade.cache, "similarity_threshold", 0.5): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "embedding_model", "other-model"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "_index_name", "other-index"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"): + assert resolver.resolve().kind == "python_callback" + + def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: + return _semantic_embedding(prompt) + + monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding) + assert resolver.resolve().kind == "python_callback" + + +def test_redis_semantic_handle_rejects_wrong_backends( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + + class CustomSemanticCache(RedisSemanticCache): + pass + + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic(object()) + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic( + CustomSemanticCache( + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=f"{index}_subclass", + ) + ) + + facade: Final = semantic_facade(url, index) + with pytest.raises(TypeError, match="backend types must match"): + _CacheTestHandle.redis(url)._bind_facade(facade) + + subclassed_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=index, + ) + with pytest.raises(TypeError): + _CacheTestHandle.redis_semantic(subclassed_facade.cache)._bind_facade(subclassed_facade) + + replacement_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + with pytest.raises(TypeError, match="must be the native embedder"): + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade) + + +def qdrant_facade(qdrant_url: str, collection_name: str) -> Cache: + return Cache( + type=LiteLLMCacheType.QDRANT_SEMANTIC, + qdrant_api_base=qdrant_url, + qdrant_collection_name=collection_name, + similarity_threshold=0.99, + qdrant_semantic_cache_embedding_model="text-embedding-3-small", + qdrant_semantic_cache_vector_size=8, + ) + + +def test_qdrant_semantic_facade_binds_native_and_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "shared prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + facade.cache.set_cache( + "python-key", + {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, + messages=messages, + ) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + assert binding.lookup(qdrant_request("python-key", messages)) == {"id": "py"} + binding.store(qdrant_request("native-key", messages), {"id": "native"}) + python_value: Final = facade.cache.get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} + unrelated: Final = [{"role": "user", "content": "unrelated prompt"}] + assert binding.lookup(qdrant_request("native-key", unrelated)) is None + assert facade.cache.get_cache("native-key", messages=unrelated) is None + assert binding.lookup(qdrant_request("different-key", messages)) is None + assert facade.cache.get_cache("different-key", messages=messages) is None + + +async def test_qdrant_semantic_async_parity(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "async prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + await facade.cache.async_set_cache( + "python-key", + {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, + messages=messages, + ) + assert await binding.async_lookup(qdrant_request("python-key", messages)) == {"id": "py"} + await binding.async_store(qdrant_request("native-key", messages), {"id": "native"}) + python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} + + +async def test_qdrant_semantic_async_store_batch_shares_entries( + qdrant_url: str, fake_embedding_endpoint: str +) -> None: + del fake_embedding_endpoint + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + entries: Final = [ + qdrant_request("batch-one", [{"role": "user", "content": "first batch prompt"}]), + qdrant_request("batch-two", [{"role": "user", "content": "second batch prompt"}]), + ] + await binding.async_store_batch(entries, [{"id": "one"}, {"id": "two"}]) + + assert binding.lookup(entries[0]) == {"id": "one"} + assert binding.lookup(entries[1]) == {"id": "two"} + assert ( + (await facade.cache.async_get_cache("batch-one", messages=entries[0]["messages"]))["response"] + == {"id": "one"} + ) + assert ( + (await facade.cache.async_get_cache("batch-two", messages=entries[1]["messages"]))["response"] + == {"id": "two"} + ) + + +async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( + qdrant_url: str, fake_embedding_endpoint: str +) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "malformed prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + key: Final = "malformed-key" + response: Final = { + "points": [ + { + "id": str(uuid4()), + "vector": embedding_vector("malformed prompt"), + "payload": { + "litellm_cache_key": key, + "text": "malformed prompt", + "response": "not json", + }, + } + ] + } + facade.cache.sync_client.put( + url=f"{qdrant_url}/collections/{collection}/points", + headers=facade.cache.headers, + json=response, + ) + assert binding.lookup(qdrant_request(key, messages)) is None + with pytest.raises(RuntimeError, match="operation is not supported"): + binding.lookup_batch([qdrant_request(key, messages)]) + with pytest.raises(RuntimeError, match="operation is not supported"): + await binding.async_flush() + with pytest.raises(RuntimeError, match="operation is not supported"): + await binding.ping() + + +def test_qdrant_semantic_ignores_request_expiry(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "persistent prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + binding.store(qdrant_request("persistent-key", messages, ttl_seconds=1.0), {"id": "persistent"}) + time.sleep(1.2) + assert binding.lookup(qdrant_request("persistent-key", messages)) == {"id": "persistent"} + python_value: Final = facade.cache.get_cache("persistent-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "persistent"} + + +def test_qdrant_semantic_mutation_and_projection_fallback(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + facade.cache.qdrant_api_key = "rotated" + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + facade.cache.similarity_threshold = 0.5 + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") + unsupported.cache.embedding_max_input_tokens = 100 + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(unsupported) + unsupported.cache.embedding_max_input_tokens = None + unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777" + with pytest.raises(TypeError, match="gRPC"): + handle._bind_facade(unsupported) diff --git a/tests/test_litellm_rust/test_fork_guard.py b/tests/test_litellm_rust/test_fork_guard.py index 086397bab5c..2a8fb6f9fca 100644 --- a/tests/test_litellm_rust/test_fork_guard.py +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -1,5 +1,6 @@ import os import textwrap +from typing import Final import pytest @@ -144,3 +145,76 @@ def test_sdk_call_in_a_child_forked_after_native_use_raises_instead_of_hanging() result = run_child_interpreter(_SDK_CONTRACT, env=env, timeout=120) assert result.returncode == 0, result.stderr + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork only") +@pytest.mark.parametrize("warm_fast_counter", (False, True)) +def test_tokenizers_share_the_native_process_guard(warm_fast_counter: bool) -> None: + script: Final = """ +import asyncio +import os +import litellm +from litellm.proxy.spend_tracking.input_tokens import count_input_tokens +from litellm.rust_bridge import _native +from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer +from litellm.utils import claude_json_str + +litellm.anthropic_models = {*litellm.anthropic_models, "tokenizer-fork-fixture"} +_native.reserve_process_for_forking() +for create in ( + lambda: _native.Tokenizer.from_tiktoken("cl100k_base"), + lambda: _native.Tokenizer.from_json(claude_json_str), + lambda: litellm.token_counter(model="tokenizer-fork-fixture", text="hello"), +): + try: + create() + except _native.ProcessReservedForForking: + pass + else: + raise AssertionError("reserved parent ran a native tokenizer") +assert not _native.process_state_started() + +pid = os.fork() +if pid == 0: + tokenizer = HuggingFaceTokenizer.from_str(claude_json_str) + encoding = _native.Tokenizer.from_tiktoken("cl100k_base") + if os.environ["WARM_FAST_COUNTER"] == "True": + _native.TokenCounter.from_tokenizer(encoding, fast=True) + expected = [item.ids for item in tokenizer.encode_batch(["hello", "world"])] + assert _native.process_state_started() + grandchild = os.fork() + if grandchild == 0: + for call in ( + lambda: tokenizer.encode_batch(["hello", "world"]), + lambda: tokenizer.encode("hello"), + lambda: encoding.count("hello"), + lambda: encoding.count("hello", fast=True), + lambda: _native.TokenCounter.from_tokenizer(encoding), + lambda: _native.TokenCounter.from_tokenizer(encoding, fast=True), + lambda: _native.Tokenizer.from_tiktoken("cl100k_base"), + lambda: asyncio.run(count_input_tokens({"prompt": "hello"}, b'{"prompt": "hello"}', ("counter-fork-fixture",))), + ): + try: + call() + except _native.ForkedAfterNativeRuntimeStarted: + pass + else: + os._exit(1) + os._exit(0) + assert os.waitpid(grandchild, 0)[1] == 0 + assert [item.ids for item in tokenizer.encode_batch(["hello", "world"])] == expected + os._exit(0) +assert os.waitpid(pid, 0)[1] == 0 +""" + result: Final = run_child_interpreter( + script, + env={ + **os.environ, + "OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "WARM_FAST_COUNTER": str(warm_fast_counter), + }, + timeout=30, + ) + + assert result.returncode == 0, result.stderr diff --git a/tests/test_litellm_rust/test_tokenizer.py b/tests/test_litellm_rust/test_tokenizer.py new file mode 100644 index 00000000000..98d5259b652 --- /dev/null +++ b/tests/test_litellm_rust/test_tokenizer.py @@ -0,0 +1,130 @@ +import json +from typing import Final + +import pytest +import tiktoken +from tokenizers import Tokenizer as ReferenceTokenizer + +from litellm.rust_bridge import _native +from litellm.utils import claude_json_str +from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOKENIZER_JSON + +pytestmark = pytest.mark.requires_rust_extension + + +def test_tiktoken_codec_round_trips_and_counts() -> None: + tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") + encoded: Final = tokenizer.encode("hello world") + + assert tokenizer.name == "cl100k_base" + assert tokenizer.count("hello world") == len(encoded) + assert tokenizer.decode(encoded) == "hello world" + + +def test_huggingface_codec_skips_special_tokens() -> None: + tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) + encoded: Final = tokenizer.encode("hello") + + assert "" in tokenizer.decode(encoded, skip_special_tokens=False) + assert tokenizer.decode(encoded, skip_special_tokens=True) == "hello" + + +def test_tiktoken_codec_keeps_the_requested_encoding_name() -> None: + assert _native.Tokenizer.from_tiktoken("gpt2").name == "gpt2" + assert _native.Tokenizer.from_tiktoken("r50k_base").name == "r50k_base" + assert _native.Tokenizer.from_tiktoken("gpt2").encode("hi") == _native.Tokenizer.from_tiktoken("r50k_base").encode( + "hi" + ) + + +def test_tiktoken_codec_exposes_its_vocabulary() -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") + + assert tokenizer.special_tokens() == reference._special_tokens + assert tokenizer.max_token_value() == reference.max_token_value + assert tokenizer.token_byte_values() == reference.token_byte_values() + assert tokenizer.encode_single_token(b"hello") == reference.encode_single_token("hello") + assert tokenizer.is_special_token(reference.eot_token) and not tokenizer.is_special_token(0) + with pytest.raises(KeyError): + tokenizer.encode_single_token(b"<|not-a-token|>") + + +def test_huggingface_codec_rejects_tiktoken_only_calls() -> None: + tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) + with pytest.raises(ValueError, match="requires a tiktoken encoding"): + tokenizer.token_byte_values() + with pytest.raises(ValueError, match="requires a Hugging Face tokenizer"): + _native.Tokenizer.from_tiktoken("cl100k_base").get_vocab() + + +def test_unknown_tiktoken_encoding_raises_value_error() -> None: + with pytest.raises(ValueError, match="unsupported tokenizer"): + _native.Tokenizer.from_tiktoken("unknown-encoding") + + +def test_tiktoken_codec_decodes_truncated_unicode_like_python() -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + tokenizer: Final = _native.Tokenizer.from_tiktoken(reference.name) + encoded: Final = reference.encode("🙂漢字") + + assert tuple(tokenizer.decode(encoded[:end]) for end in range(1, len(encoded) + 1)) == tuple( + reference.decode(encoded[:end]) for end in range(1, len(encoded) + 1) + ) + + +FAST_TEXTS: Final = ( + "", + "hello world <|endoftext|>", + "café 漢字 ع 🙂 line\r\n indented 123456789", + "x a\u0301 fi", +) + + +def test_fast_counting_is_an_opt_in_over_the_same_loaded_tokenizer() -> None: + for tokenizer in ( + _native.Tokenizer.from_tiktoken("cl100k_base"), + _native.Tokenizer.from_tiktoken("o200k_base"), + _native.Tokenizer.from_json(claude_json_str), + ): + assert [tokenizer.count(text, fast=True) for text in FAST_TEXTS] == [ + tokenizer.count(text) for text in FAST_TEXTS + ] + + +@pytest.mark.parametrize( + "name", ("cl100k_base", "o200k_base", "o200k_harmony", "p50k_base", "p50k_edit", "r50k_base", "gpt2") +) +@pytest.mark.asyncio +async def test_token_counter_counts_over_a_shared_tokenizer(name: str) -> None: + messages: Final = [{"role": "user", "content": "hello wide world"}, {"role": "assistant", "content": "ok"}] + body: Final = json.dumps({"model": "gpt-4", "messages": messages}).encode() + tokenizer: Final = _native.Tokenizer.from_tiktoken(name) + reference: Final = tiktoken.get_encoding(name) + for text in FAST_TEXTS: + assert tokenizer.count(text, fast=True) == tokenizer.count(text) == len(reference.encode_ordinary(text)) + + exact: Final = await _native.TokenCounter.from_tokenizer(tokenizer).acount_request(body) + fast: Final = await _native.TokenCounter.from_tokenizer(tokenizer, fast=True).acount_request(body) + + assert exact == fast + assert exact["input_tokens"] == 3 + sum( + 3 + len(reference.encode_ordinary(message["role"])) + len(reference.encode_ordinary(message["content"])) + for message in messages + ) + + +@pytest.mark.parametrize("configured", (False, True)) +@pytest.mark.asyncio +async def test_fast_count_preserves_huggingface_configuration(configured: bool) -> None: + reference: Final = ReferenceTokenizer.from_str(TOKENIZER_JSON) + if configured: + reference.enable_truncation(max_length=3) + reference.enable_padding(pad_id=0, pad_token="[UNK]", length=5) + tokenizer: Final = _native.Tokenizer.from_json(reference.to_str()) + counter: Final = _native.TokenCounter.from_tokenizer(tokenizer, fast=True) + for text in ("", "Hello", "Hello World Hello World", "[BOS] Hello"): + expected: Final = len(reference.encode(text)) + assert tokenizer.count(text, fast=True) == tokenizer.count(text) == expected + result: Final = await counter.acount_request(json.dumps({"prompt": text}).encode()) + assert result["input_tokens"] == expected diff --git a/tests/unit/chat_completions/test_dispatch.py b/tests/unit/chat_completions/test_dispatch.py index 63821c74208..2807ed7f8f7 100644 --- a/tests/unit/chat_completions/test_dispatch.py +++ b/tests/unit/chat_completions/test_dispatch.py @@ -10,7 +10,7 @@ from litellm.chat_completions.dispatch import ( ) from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, Rule +from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.chat_completions.entrypoints import ( NATIVE_ACOMPLETION, NATIVE_COMPLETION, @@ -23,7 +23,7 @@ from litellm.types.utils import ModelResponse MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final = () -RUST_RULES: Final = (Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) +RUST_RULES: Final = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),) def completion_binding(native: NativeCompletion | None) -> NativeBinding[NativeCompletion]: @@ -117,9 +117,7 @@ def test_native_receives_bound_request_and_original_call_shape() -> None: "custom_llm_provider": "anthropic", "metadata": metadata, } - captured: Final[ - list[tuple[LiteLLMChatCompletionsRequest, tuple[object, ...], Mapping[str, object]]] - ] = [] + captured: Final[list[tuple[LiteLLMChatCompletionsRequest, tuple[object, ...], Mapping[str, object]]]] = [] def python(*call_args: object, **call_kwargs: object) -> ModelResponse: # kwargs-ok: rejected Rust fallback pytest.fail("Required Rust dispatch must not call Python") diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py index 50409b2ea2c..37201e8155b 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -7,10 +7,12 @@ through _hidden_params to the x-litellm-callback-duration-ms response header. import asyncio import datetime +from typing import Final from unittest.mock import MagicMock import pytest +import litellm import litellm.litellm_core_utils.llm_response_utils.response_metadata as response_metadata_mod import litellm.proxy.common_request_processing as common_request_processing_mod from litellm.litellm_core_utils.litellm_logging import Logging @@ -22,7 +24,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from litellm.types.utils import ModelResponse +from litellm.types.utils import ModelResponse, Usage class TestCallbackDurationMs: @@ -583,3 +585,57 @@ class TestLoggingInitCallbackDuration: # Should still be set (deep copy of None is essentially a no-op) assert hasattr(obj, "callback_duration_ms") assert obj.callback_duration_ms >= 0 + + +def test_update_response_metadata_prices_per_second_deployment_from_its_stamped_duration(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + deployment_id: Final = "per-second-deployment-response-metadata" + litellm.register_model( + model_cost={ + deployment_id: { + "input_cost_per_second": 0.02, + "output_cost_per_second": 0.04, + "litellm_provider": "openai", + "mode": "chat", + } + } + ) + start_time: Final = datetime.datetime(2026, 9, 21, 12, 0, 0) + logging_obj: Final = Logging( + model="gpt-5.4-nano", + messages=[{"role": "user", "content": "Hello"}], + stream=False, + call_type="completion", + start_time=start_time, + litellm_call_id="per-second-response-metadata", + function_id="f", + ) + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + litellm_params={ + "input_cost_per_second": 0.02, + "output_cost_per_second": 0.04, + "metadata": {"model_info": {"id": deployment_id}}, + }, + optional_params={}, + custom_llm_provider="openai", + ) + logging_obj.model_call_details["end_time"] = start_time + datetime.timedelta(seconds=10) + result: Final = ModelResponse( + model="gpt-5.4-nano", + usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18), + ) + + update_response_metadata( + result=result, + logging_obj=logging_obj, + model="gpt-5.4-nano", + kwargs={"model_info": {"id": deployment_id}}, + start_time=start_time, + end_time=start_time + datetime.timedelta(seconds=2), + ) + + assert result._response_ms == pytest.approx(2000) + assert result._hidden_params["response_cost"] == pytest.approx((0.02 + 0.04) * 2) diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_handler.py b/tests/unit/llms/bedrock/files/test_bedrock_files_handler.py index 639be272351..5c078affffc 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_handler.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_handler.py @@ -270,3 +270,40 @@ async def test_afile_content_assumes_role_with_external_id(monkeypatch): assert s3_client_kwargs["aws_access_key_id"] == "ASIAFILESDOWNLOADROLE" assert s3_client_kwargs["aws_session_token"] == "assumed-session-token" assert response.content == b'{"custom_id": "req-1"}' + + +@pytest.mark.asyncio +async def test_afile_content_builds_the_s3_client_with_the_s3_pair_when_it_differs_from_the_aws_identity(): + import boto3 + + class FakeS3Body: + def read(self): + return b'{"custom_id": "req-1"}' + + class FakeS3Client: + def get_object(self, Bucket, Key): + return {"Body": FakeS3Body()} + + optional_params = { + "_litellm_internal_model_credentials": MappingProxyType({"s3_bucket_name": "safe-bucket"}), + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIABEDROCKONLY", + "aws_secret_access_key": "bedrock-only-secret", + "aws_session_token": "bedrock-only-token", + "s3_access_key_id": "AKIAS3ONLY", + "s3_secret_access_key": "s3-only-secret", + } + + with patch.object(boto3, "client", return_value=FakeS3Client()) as mock_boto3_client: + response = await BedrockFilesHandler().afile_content( + file_content_request={"file_id": "s3://safe-bucket/litellm-bedrock-files-model-id-abc.jsonl"}, + optional_params=optional_params, + timeout=10.0, + max_retries=None, + ) + + s3_client_kwargs = mock_boto3_client.call_args.kwargs + assert s3_client_kwargs["aws_access_key_id"] == "AKIAS3ONLY" + assert s3_client_kwargs["aws_secret_access_key"] == "s3-only-secret" + assert s3_client_kwargs["aws_session_token"] is None, "the aws_* session token belongs to the Bedrock identity" + assert response.content == b'{"custom_id": "req-1"}' diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 2d2de77269b..d0921e68424 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -6,6 +6,7 @@ import json import os from collections.abc import Mapping from contextlib import AsyncExitStack, closing +from types import MappingProxyType from typing import Final from unittest.mock import MagicMock from urllib.parse import unquote, urlparse @@ -3789,3 +3790,82 @@ class TestBedrockFileListTransformation: assert denied.value.status_code == 403 assert "AccessDenied" in denied.value.message + + +_SPLIT_IDENTITY_PARAMS: Final = { + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIABEDROCKONLY", + "aws_secret_access_key": "bedrock-only-secret", + "s3_access_key_id": "AKIAS3ONLY", + "s3_secret_access_key": "s3-only-secret", + "s3_bucket_name": "safe-bucket", +} + + +def _authorization(headers: Mapping[str, str]) -> str: + return {key.lower(): value for key, value in headers.items()}["authorization"] + + +def test_sign_s3_request_uses_the_s3_pair_when_it_differs_from_the_aws_identity(): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + signed_headers, _signed_body = BedrockFilesConfig()._sign_s3_request( + content='{"custom_id": "req-1"}', + api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", + optional_params=dict(_SPLIT_IDENTITY_PARAMS), + ) + + assert _authorization(signed_headers).startswith("AWS4-HMAC-SHA256 Credential=AKIAS3ONLY/"), ( + "the S3 PutObject must be signed by s3_access_key_id, not the Bedrock aws_access_key_id" + ) + + +def test_sign_s3_request_with_the_s3_pair_ignores_ambient_aws_session_token_role_and_profile(monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_SESSION_TOKEN", "pod-token") + monkeypatch.setenv("AWS_ROLE_NAME", "arn:aws:iam::123456789012:role/pod") + monkeypatch.setenv("AWS_PROFILE_NAME", "pod-profile") + signed_headers, _signed_body = BedrockFilesConfig()._sign_s3_request( + content='{"custom_id": "req-1"}', + api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", + optional_params=dict(_SPLIT_IDENTITY_PARAMS), + ) + + lowered: Final = {key.lower(): value for key, value in signed_headers.items()} + assert lowered["authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIAS3ONLY/") + assert "x-amz-security-token" not in lowered, "an ambient AWS_SESSION_TOKEN must not be mixed into the s3_* pair" + + +@pytest.mark.parametrize("method", ["GET", "DELETE"]) +def test_sign_s3_request_without_body_uses_the_s3_pair_when_it_differs_from_the_aws_identity(method): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig, _BedrockS3RequestParams + + signed_headers = BedrockFilesConfig()._sign_s3_request_without_body( + method=method, + api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", + aws_region_name="us-east-1", + request_params=_BedrockS3RequestParams.model_validate(_SPLIT_IDENTITY_PARAMS), + ) + + assert _authorization(signed_headers).startswith("AWS4-HMAC-SHA256 Credential=AKIAS3ONLY/"), ( + f"the S3 {method} must be signed by s3_access_key_id, not the Bedrock aws_access_key_id" + ) + + +def test_transform_file_content_request_signs_with_the_s3_pair_from_litellm_params(): + from litellm.llms.bedrock.files.transformation import S3_SIGNED_REQUEST_HEADERS_PARAM, BedrockFilesConfig + + litellm_params = { + **_SPLIT_IDENTITY_PARAMS, + "_litellm_internal_model_credentials": MappingProxyType({"s3_bucket_name": "safe-bucket"}), + } + BedrockFilesConfig().transform_file_content_request( + file_content_request={"file_id": "s3://safe-bucket/litellm-bedrock-files-model-id-abc.jsonl"}, + optional_params={}, + litellm_params=litellm_params, + ) + + assert _authorization(litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]).startswith( + "AWS4-HMAC-SHA256 Credential=AKIAS3ONLY/" + ) diff --git a/tests/unit/messages/test_dispatch.py b/tests/unit/messages/test_dispatch.py index 586b77d9a25..48eb1adbf51 100644 --- a/tests/unit/messages/test_dispatch.py +++ b/tests/unit/messages/test_dispatch.py @@ -12,7 +12,7 @@ from litellm.messages.dispatch import ( ) from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, Rule, Rules +from litellm.rust_bridge.catalog import Route, RouteRule, Rules from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.messages.entrypoints import ( NATIVE_AMESSAGES, @@ -25,7 +25,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMe MESSAGES: Final = [{"role": "user", "content": "hi"}] PYTHON_RULES: Final[Rules] = () -RUST_RULES: Final[Rules] = (Rule(Route.MESSAGES, Rollout.RUST_REQUIRED),) +RUST_RULES: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),) def messages_binding(native: NativeMessages | None) -> NativeBinding[NativeMessages]: diff --git a/tests/unit/ocr/test_dispatch.py b/tests/unit/ocr/test_dispatch.py index e54d4070ba8..2727ff23449 100644 --- a/tests/unit/ocr/test_dispatch.py +++ b/tests/unit/ocr/test_dispatch.py @@ -12,7 +12,7 @@ from litellm.ocr.dispatch import ( ) from litellm.rust_bridge import catalog from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.catalog import Route, Rule, Rules +from litellm.rust_bridge.catalog import Route, RouteRule, Rules from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.ocr.entrypoints import ( NATIVE_AOCR, @@ -22,8 +22,8 @@ from litellm.rust_bridge.ocr.entrypoints import ( NativeOcr, ) -PYTHON_RULES: Final[Rules] = (Rule(Route.OCR, Rollout.PYTHON_ONLY),) -RUST_RULES: Final[Rules] = (Rule(Route.OCR, Rollout.RUST_REQUIRED),) +PYTHON_RULES: Final[Rules] = (RouteRule(Route.OCR, Rollout.PYTHON_ONLY),) +RUST_RULES: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),) def ocr_binding(native: NativeOcr | None) -> NativeBinding[NativeOcr]: @@ -403,8 +403,8 @@ def test_provider_scoped_rule_sees_the_provider_named_by_the_model_prefix( model: str, custom_llm_provider: str | None, expected: str ) -> None: rules: Final[Rules] = ( - Rule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), - Rule(Route.OCR, Rollout.PYTHON_ONLY), + RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), ) document: Final[Mapping[str, object]] = {"type": "image_url", "image_url": "data:image/png;base64,YQ=="} kwargs: Final[Mapping[str, object]] = ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx index 15b85001cdc..8e100d0c3ed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx @@ -90,6 +90,7 @@ export interface AgentFormValues { guardrails?: string[]; entitlement_models?: string[]; entitlement_agents?: string[]; + access_group_ids?: string[]; allowed_mcp_servers_and_groups?: McpServerSelection; mcp_tool_permissions?: Record; defaultInputModes?: string[]; @@ -121,6 +122,7 @@ export interface AgentRequestPayload { agent_card_params?: Record; litellm_params?: Record; object_permission?: Record; + access_group_ids?: string[]; } interface AgentFormFieldProps { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx index ebc97891744..457ee656415 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx @@ -21,6 +21,9 @@ vi.mock("./agent_card_discovery", () => ({ default: () =>
({ default: () =>
})); vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({ default: () =>
})); vi.mock("@/components/guardrails/GuardrailSelector", () => ({ default: () =>
})); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: () => ({ data: [], isLoading: false, isError: false }), +})); vi.mock("@/components/common_components/team_dropdown", () => ({ default: () =>
})); const a2aInfo: AgentCreateInfo = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx index ccac244f019..fbf5cf8c1fb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx @@ -44,6 +44,14 @@ vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({ default: () => null, })); +vi.mock("@/components/common_components/AccessGroupSelector", () => ({ + default: ({ onChange }: { onChange: (value: string[]) => void }) => ( + + ), +})); + vi.mock("@/components/common_components/team_dropdown", () => ({ default: () => null, })); @@ -141,5 +149,29 @@ describe("AddAgentForm logos", () => { await vi.waitFor(() => expect(networking.createAgentCall).toHaveBeenCalled()); const [, payload] = vi.mocked(networking.createAgentCall).mock.calls[0]; expect(payload.object_permission).toEqual({ mcp_toolsets: ["ts-1"] }); + expect(payload).not.toHaveProperty("access_group_ids"); + }); + + it("includes selected access groups in the create payload", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + vi.mocked(networking.createAgentCall) + .mockReset() + .mockResolvedValue({ + agent_id: "agent-1", + agent_name: "Test Agent", + } as never); + vi.mocked(networking.keyListCall).mockResolvedValue({ keys: [] }); + + renderForm(); + await user.click(screen.getByRole("button", { name: "Next →" })); + await user.click(screen.getByTestId("select-access-group")); + await user.click(screen.getByRole("button", { name: "Next →" })); + await user.click(screen.getByRole("button", { name: "Next →" })); + await user.click(screen.getByText(/Skip for now/)); + await user.click(screen.getByRole("button", { name: "Create Agent →" })); + + await vi.waitFor(() => expect(networking.createAgentCall).toHaveBeenCalled()); + const [, payload] = vi.mocked(networking.createAgentCall).mock.calls[0]; + expect(payload.access_group_ids).toEqual(["ag-1", "ag-2"]); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx index e71fed40209..5bd6ea9b83a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx @@ -50,6 +50,7 @@ import { } from "./AgentFormKit"; import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; +import AccessGroupSelector from "@/components/common_components/AccessGroupSelector"; import GuardrailSelector from "@/components/guardrails/GuardrailSelector"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; @@ -113,6 +114,7 @@ const SHARED_INITIAL_VALUES: AgentFormValues = { mcp_tool_permissions: {}, entitlement_models: [], entitlement_agents: [], + access_group_ids: [], guardrails: [], }; @@ -374,6 +376,9 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok if (Object.keys(objectPermission).length > 0) { agentData.object_permission = objectPermission; } + if (values.access_group_ids?.length) { + agentData.access_group_ids = values.access_group_ids; + } // Wire trace-id flags and budget controls into agent litellm_params (before create call) if (requireTraceIdInbound || requireTraceIdOutbound) { @@ -494,6 +499,22 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok )} + + {({ value, onChange }) => ( + + )} + + { return agentData; }; +export const parseAccessGroupIdsForForm = (agent: { access_group_ids?: string[] | null }) => ({ + access_group_ids: agent.access_group_ids ?? [], +}); + export const parseMcpPermissionsForForm = (agent: any) => ({ allowed_mcp_servers_and_groups: { servers: agent.object_permission?.mcp_servers ?? [], @@ -377,5 +381,6 @@ export const parseAgentForForm = (agent: any) => { // extra_headers: already an array of strings extra_headers: agent.extra_headers ?? [], ...parseMcpPermissionsForForm(agent), + ...parseAccessGroupIdsForForm(agent), }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx index 357f924cbe7..37e00766a75 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx @@ -25,6 +25,10 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ vi.mock("./agent_card_discovery", () => ({ default: () =>
})); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: () => ({ data: [], isLoading: false, isError: false }), +})); + const A2A_AGENT = { agent_id: "agent-1", agent_name: "my-agent", @@ -176,6 +180,7 @@ describe("AgentInfoView update payload", () => { session_tpm_limit: 333, session_rpm_limit: 444, object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, + access_group_ids: [], }); }); @@ -217,6 +222,7 @@ describe("AgentInfoView update payload", () => { session_tpm_limit: 333, session_rpm_limit: 444, object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, + access_group_ids: [], }); }); @@ -295,6 +301,7 @@ describe("AgentInfoView update payload", () => { model: "langgraph/asst_1", }, object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, + access_group_ids: [], }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx index 29d97a20afe..7e6c7c0e05c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx @@ -28,6 +28,28 @@ vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ useMCPServers: () => ({ data: [{ server_id: "srv-1", server_name: "github" }] }), })); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: () => ({ + data: [{ access_group_id: "ag-1", access_group_name: "support-tools" }], + isLoading: false, + isError: false, + }), +})); + +vi.mock("@/components/common_components/AccessGroupSelector", () => ({ + default: ({ value, onChange }: { value?: string[]; onChange: (value: string[]) => void }) => ( +
+ {(value ?? []).join(",")} + + +
+ ), +})); + vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ default: () =>
, })); @@ -76,6 +98,38 @@ describe("AgentInfoView settings", () => { expect(payload.tpm_limit).toBe(42); const clearedMcpGrants = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }; expect(payload.object_permission).toEqual(clearedMcpGrants); + expect(payload.access_group_ids).toEqual([]); + }); + + it("sends the newly attached access group in the update payload", async () => { + render(); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(screen.getByRole("button", { name: "Edit Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Attach ag-1" })); + expect(screen.getByTestId("selected-access-groups")).toHaveTextContent("ag-1"); + + fireEvent.click(screen.getByRole("button", { name: /Save Changes/ })); + + await waitFor(() => expect(networking.patchAgentCall).toHaveBeenCalledTimes(1)); + const [, , payload] = vi.mocked(networking.patchAgentCall).mock.calls[0]; + expect(payload.access_group_ids).toEqual(["ag-1"]); + }); + + it("loads the attached access groups into the editor and sends an empty list once detached", async () => { + vi.mocked(networking.getAgentInfo).mockResolvedValue({ ...agent, access_group_ids: ["ag-1"] }); + render(); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(screen.getByRole("button", { name: "Edit Settings" })); + expect(await screen.findByTestId("selected-access-groups")).toHaveTextContent("ag-1"); + + fireEvent.click(screen.getByRole("button", { name: "Detach all access groups" })); + fireEvent.click(screen.getByRole("button", { name: /Save Changes/ })); + + await waitFor(() => expect(networking.patchAgentCall).toHaveBeenCalledTimes(1)); + const [, , payload] = vi.mocked(networking.patchAgentCall).mock.calls[0]; + expect(payload.access_group_ids).toEqual([]); }); it("shows MCP grants with server names on the overview tab", async () => { @@ -88,4 +142,20 @@ describe("AgentInfoView settings", () => { expect(await screen.findByText("github (srv-1)")).toBeInTheDocument(); }); + + it("shows attached access groups with their names on the overview tab", async () => { + vi.mocked(networking.getAgentInfo).mockResolvedValue({ ...agent, access_group_ids: ["ag-1", "ag-unknown"] }); + + render(); + + expect(await screen.findByText("support-tools (ag-1)")).toBeInTheDocument(); + expect(screen.getByText("ag-unknown")).toBeInTheDocument(); + }); + + it("shows None when the agent has no access groups attached", async () => { + render(); + + expect(await screen.findByText("Access Groups")).toBeInTheDocument(); + expect(screen.getByText("None")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx index 6cb99e9692f..adac456f232 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx @@ -16,6 +16,8 @@ import { Agent } from "@/components/agents/types"; import { KeyResponse } from "@/components/key_team_helpers/key_list"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; +import AccessGroupSelector from "@/components/common_components/AccessGroupSelector"; import KeyInfoView from "@/components/templates/key_info_view"; import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; @@ -26,6 +28,7 @@ import { AGENT_FORM_CONFIG, buildAgentDataFromForm, buildMcpObjectPermission, + parseAccessGroupIdsForForm, parseAgentForForm, parseMcpPermissionsForForm, } from "./agent_config"; @@ -122,7 +125,11 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT } else { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset({ ...parseDynamicAgentForForm(data, typeInfo), ...parseMcpPermissionsForForm(data) }); + form.reset({ + ...parseDynamicAgentForForm(data, typeInfo), + ...parseMcpPermissionsForForm(data), + ...parseAccessGroupIdsForForm(data), + }); } else { form.reset(parseAgentForForm(data)); } @@ -142,7 +149,11 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT if (agentType !== "a2a") { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset({ ...parseDynamicAgentForForm(agent, typeInfo), ...parseMcpPermissionsForForm(agent) }); + form.reset({ + ...parseDynamicAgentForForm(agent, typeInfo), + ...parseMcpPermissionsForForm(agent), + ...parseAccessGroupIdsForForm(agent), + }); } } } @@ -153,12 +164,18 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT const mcpSelection = useWatch({ control: form.control, name: "allowed_mcp_servers_and_groups" }); const mcpToolPermissions = useWatch({ control: form.control, name: "mcp_tool_permissions" }); const { data: mcpServers = [] } = useMCPServers(); + const { data: accessGroups = [] } = useAccessGroups(); const mcpServerLabel = (serverId: string) => { const server = mcpServers.find((s) => s.server_id === serverId); return server?.server_name ? `${server.server_name} (${serverId})` : serverId; }; + const accessGroupLabel = (accessGroupId: string) => { + const group = accessGroups.find((g) => g.access_group_id === accessGroupId); + return group ? `${group.access_group_name} (${accessGroupId})` : accessGroupId; + }; + const discoveryRequest = useMemo( () => buildDiscoveryRequest(detectedAgentType, watchedFormValues || {}, selectedAgentTypeInfo), [watchedFormValues, selectedAgentTypeInfo, detectedAgentType], @@ -221,6 +238,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT await patchAgentCall(accessToken, agentId, { ...updateData, object_permission: buildMcpObjectPermission(values), + access_group_ids: values.access_group_ids ?? [], }); toast.success("Agent updated successfully"); setIsEditing(false); @@ -350,6 +368,17 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {agent.rpm_limit ?? "Unlimited"} {agent.session_tpm_limit ?? "Unlimited"} {agent.session_rpm_limit ?? "Unlimited"} + + {agent.access_group_ids?.length ? ( +
+ {agent.access_group_ids.map((accessGroupId) => ( +
{accessGroupLabel(accessGroupId)}
+ ))} +
+ ) : ( + "None" + )} +
{formatDate(agent.created_at)} {formatDate(agent.updated_at)} @@ -489,6 +518,26 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {rateLimitField("session_rpm_limit", "Session RPM Limit")}
+ +

Access Groups

+ + + {({ value, onChange }) => ( + + )} + + +

MCP Servers

diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx index 6c039b106bf..6743167bff6 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.test.tsx @@ -33,4 +33,33 @@ describe("ExportTypeSelector", () => { renderWithProviders(); expect(screen.getByRole("radio", { name: /Day-by-day by team and model/i })).toBeChecked(); }); + + it("should offer the per-user scope for teams and call onChange with daily_with_users", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders(); + + const option = screen.getByRole("radio", { name: /Day-by-day breakdown by team and user/i }); + await user.click(option); + + expect(onChange).toHaveBeenCalledWith("daily_with_users"); + expect(screen.getByText("Daily metrics for each team, split by key owner")).toBeInTheDocument(); + }); + + it("should hide the per-user scope for user exports while keeping the other scopes", () => { + renderWithProviders(); + + expect(screen.queryByRole("radio", { name: /and user/i })).not.toBeInTheDocument(); + expect( + screen.getByRole("radio", { name: /Day-by-day breakdown by user Daily metrics for each user$/i }), + ).toBeInTheDocument(); + expect(screen.getByRole("radio", { name: /Day-by-day breakdown by user and key/i })).toBeInTheDocument(); + expect(screen.getByRole("radio", { name: /Day-by-day by user and model/i })).toBeInTheDocument(); + }); + + it("should offer the per-user scope for tags", () => { + renderWithProviders(); + + expect(screen.getByRole("radio", { name: /Day-by-day breakdown by tag and user/i })).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx index edd055ba7a7..f6fdece83a8 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx @@ -9,7 +9,7 @@ interface ExportTypeSelectorProps { } const ExportTypeSelector: React.FC = ({ value, onChange, entityType }) => { - const scopes: { value: ExportScope; title: string; description: string }[] = [ + const allScopes: { value: ExportScope; title: string; description: string }[] = [ { value: "daily", title: `Day-by-day breakdown by ${entityType}`, @@ -25,7 +25,13 @@ const ExportTypeSelector: React.FC = ({ value, onChange title: `Day-by-day by ${entityType} and model`, description: "Daily metrics split by model", }, + { + value: "daily_with_users", + title: `Day-by-day breakdown by ${entityType} and user`, + description: `Daily metrics for each ${entityType}, split by key owner`, + }, ]; + const scopes = allScopes.filter((scope) => scope.value !== "daily_with_users" || entityType !== "user"); return (
diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index 30714ad632d..15f193ecc3f 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -2,7 +2,7 @@ import type { DateRangePickerValue } from "@/components/shared/date_picker_types import type { Team } from "@/components/key_team_helpers/key_list"; export type ExportFormat = "csv" | "json"; -export type ExportScope = "daily" | "daily_with_keys" | "daily_with_models"; +export type ExportScope = "daily" | "daily_with_keys" | "daily_with_models" | "daily_with_users"; export type EntityType = "tag" | "team" | "organization" | "customer" | "agent" | "user"; export interface EntitySpendData { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index 7ed014d43b2..3f9cf58ec20 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -8,6 +8,7 @@ import { generateDailyData, generateDailyWithKeysData, generateDailyWithModelsData, + generateDailyWithUsersData, generateExportData, generateMetadata, getEntityBreakdown, @@ -151,6 +152,204 @@ describe("EntityUsageExport utils", () => { "team-2": "Team Two", }; + const usersFixture: EntitySpendData = { + results: [ + { + date: "2025-03-01", + breakdown: { + entities: { + "team-1": { + metrics: { + spend: 16.5, + api_requests: 165, + successful_requests: 156, + failed_requests: 9, + total_tokens: 1650, + prompt_tokens: 940, + completion_tokens: 710, + cache_read_input_tokens: 90, + cache_creation_input_tokens: 60, + }, + api_key_breakdown: { + kA: { + metrics: { + spend: 1.1, + api_requests: 11, + successful_requests: 10, + failed_requests: 1, + total_tokens: 110, + prompt_tokens: 60, + completion_tokens: 50, + cache_read_input_tokens: 6, + cache_creation_input_tokens: 4, + }, + metadata: { + team_id: "team-1", + key_alias: "alice-key", + user_id: "u1", + user_email: "a@x", + }, + }, + kB: { + metrics: { + spend: 2.2, + api_requests: 22, + successful_requests: 20, + failed_requests: 2, + total_tokens: 220, + prompt_tokens: 130, + completion_tokens: 90, + cache_read_input_tokens: 12, + cache_creation_input_tokens: 8, + }, + metadata: { + team_id: "team-1", + user_id: "u1", + user_email: "a@x", + }, + }, + kC: { + metrics: { + spend: 3.3, + api_requests: 33, + successful_requests: 31, + failed_requests: 2, + total_tokens: 330, + prompt_tokens: 190, + completion_tokens: 140, + cache_read_input_tokens: 18, + cache_creation_input_tokens: 12, + }, + metadata: { + team_id: "team-1", + user_id: "u2", + user_email: null, + }, + }, + kD: { + metrics: { + spend: 4.4, + api_requests: 44, + successful_requests: 42, + failed_requests: 2, + total_tokens: 440, + prompt_tokens: 250, + completion_tokens: 190, + cache_read_input_tokens: 24, + cache_creation_input_tokens: 16, + }, + metadata: { + team_id: "team-1", + user_id: null, + }, + }, + kE: { + metrics: { + spend: 5.5, + api_requests: 55, + successful_requests: 53, + failed_requests: 2, + total_tokens: 550, + prompt_tokens: 310, + completion_tokens: 240, + cache_read_input_tokens: 30, + cache_creation_input_tokens: 20, + }, + metadata: { + team_id: "team-1", + user_id: "u3", + key_exists: false, + }, + }, + }, + }, + "team-2": { + metrics: { + spend: 6.6, + api_requests: 66, + successful_requests: 64, + failed_requests: 2, + total_tokens: 660, + prompt_tokens: 370, + completion_tokens: 290, + cache_read_input_tokens: 36, + cache_creation_input_tokens: 24, + }, + api_key_breakdown: { + kF: { + metrics: { + spend: 6.6, + api_requests: 66, + successful_requests: 64, + failed_requests: 2, + total_tokens: 660, + prompt_tokens: 370, + completion_tokens: 290, + cache_read_input_tokens: 36, + cache_creation_input_tokens: 24, + }, + metadata: { + team_id: "team-2", + user_id: "u1", + user_email: "a@x", + }, + }, + }, + }, + }, + }, + }, + { + date: "2025-03-02", + breakdown: { + entities: { + "team-1": { + metrics: { + spend: 7.7, + api_requests: 77, + successful_requests: 75, + failed_requests: 2, + total_tokens: 770, + prompt_tokens: 430, + completion_tokens: 340, + cache_read_input_tokens: 42, + cache_creation_input_tokens: 28, + }, + api_key_breakdown: { + kA: { + metrics: { + spend: 7.7, + api_requests: 77, + successful_requests: 75, + failed_requests: 2, + total_tokens: 770, + prompt_tokens: 430, + completion_tokens: 340, + cache_read_input_tokens: 42, + cache_creation_input_tokens: 28, + }, + metadata: { + team_id: "team-1", + key_alias: "alice-key", + user_id: "u1", + user_email: "a@x", + }, + }, + }, + }, + }, + }, + }, + ], + metadata: { + total_spend: 30.8, + total_api_requests: 308, + total_successful_requests: 295, + total_failed_requests: 13, + total_tokens: 3080, + }, + }; + beforeEach(() => { vi.clearAllMocks(); }); @@ -1056,6 +1255,27 @@ describe("EntityUsageExport utils", () => { expect(keyIds).toContain("key1"); expect(keyIds).toContain("key2"); }); + + it("should emit key owner columns right after Key ID", () => { + const result = generateDailyWithKeysData(usersFixture, "Team"); + + const columnNames = Object.keys(result[0]); + expect(columnNames[columnNames.indexOf("Key ID") + 1]).toBe("User ID"); + expect(columnNames[columnNames.indexOf("User ID") + 1]).toBe("User Email"); + + const kARow = result.find((r) => r["Key ID"] === "kA" && r.Date === "2025-03-01"); + expect(kARow?.["User ID"]).toBe("u1"); + expect(kARow?.["User Email"]).toBe("a@x"); + expect(kARow?.["Key Alias"]).toBe("alice-key"); + + const kCRow = result.find((r) => r["Key ID"] === "kC"); + expect(kCRow?.["User ID"]).toBe("u2"); + expect(kCRow?.["User Email"]).toBe("-"); + + const kDRow = result.find((r) => r["Key ID"] === "kD"); + expect(kDRow?.["User ID"]).toBe("-"); + expect(kDRow?.["User Email"]).toBe("-"); + }); }); describe("generateDailyWithModelsData", () => { @@ -2010,6 +2230,21 @@ describe("EntityUsageExport utils", () => { window.Blob = originalBlob; }); + + it("should generate the daily_with_users filename and include User ID in the rows", () => { + const anchorElement = document.createElement("a"); + vi.spyOn(document, "createElement").mockReturnValue(anchorElement); + vi.useFakeTimers(); + vi.setSystemTime(new Date("2025-03-01T12:00:00Z")); + + handleExportCSV(usersFixture, "daily_with_users", "Team", "team", mockTeamAliasMap); + vi.useRealTimers(); + + expect(anchorElement.download).toBe("team_usage_daily_with_users_2025-03-01.csv"); + + const unparsedRows = vi.mocked(Papa.unparse).mock.calls[0][0] as Record[]; + expect(unparsedRows[0]).toHaveProperty("User ID"); + }); }); describe("handleExportJSON", () => { @@ -2462,4 +2697,312 @@ describe("EntityUsageExport utils", () => { expect(result.find((r) => r["User ID"] === "user-b")?.["User"]).toBe("Grace"); }); }); + + describe("generateDailyWithUsersData", () => { + it("should reconcile spend with daily_with_keys and daily per date and team", () => { + const byUser = generateDailyWithUsersData(usersFixture, "Team"); + const byKey = generateDailyWithKeysData(usersFixture, "Team"); + const daily = generateDailyData(usersFixture, "Team"); + + expect(byUser.length).toBeGreaterThan(0); + expect(byKey.length).toBeGreaterThan(0); + expect(daily.length).toBeGreaterThan(0); + + const sumSpend = (rows: any[]): Record => { + const totals: Record = {}; + rows.forEach((r) => { + const bucket = `${r.Date}|${r["Team ID"]}`; + totals[bucket] = (totals[bucket] || 0) + Number(r["Spend ($)"]); + }); + return totals; + }; + + const userTotals = sumSpend(byUser); + const keyTotals = sumSpend(byKey); + + daily.forEach((row) => { + const bucket = `${row.Date}|${row["Team ID"]}`; + expect(userTotals[bucket]).toBeCloseTo(Number(row["Spend ($)"]), 4); + expect(keyTotals[bucket]).toBeCloseTo(Number(row["Spend ($)"]), 4); + }); + }); + + it("should roll multiple keys owned by one user in a team into a single row", () => { + const rows = generateDailyWithUsersData(usersFixture, "Team"); + const matches = rows.filter((r) => r.Date === "2025-03-01" && r["Team ID"] === "team-1" && r["User ID"] === "u1"); + + expect(matches).toHaveLength(1); + const row = matches[0]; + expect(row.Keys).toBe(2); + expect(row["User Email"]).toBe("a@x"); + expect(row["Spend ($)"]).toBe("3.3000"); + expect(row.Requests).toBe(33); + expect(row["Successful Requests"]).toBe(30); + expect(row["Failed Requests"]).toBe(3); + expect(row["Total Tokens"]).toBe(330); + expect(row["Prompt Tokens"]).toBe(190); + expect(row["Completion Tokens"]).toBe(140); + expect(row["Cache Read Input Tokens"]).toBe(18); + expect(row["Cache Creation Input Tokens"]).toBe(12); + }); + + it("should bucket keys with no owner into an Unassigned row without dropping spend", () => { + const rows = generateDailyWithUsersData(usersFixture, "Team"); + const row = rows.find( + (r) => r.Date === "2025-03-01" && r["Team ID"] === "team-1" && r["User ID"] === "Unassigned", + ); + + expect(row).toBeDefined(); + expect(row?.["User Email"]).toBe("-"); + expect(Number(row?.["Spend ($)"])).toBeCloseTo(4.4, 4); + }); + + it("should keep different users in the same team as separate rows", () => { + const rows = generateDailyWithUsersData(usersFixture, "Team"); + const teamRows = rows.filter((r) => r.Date === "2025-03-01" && r["Team ID"] === "team-1"); + + const u1Rows = teamRows.filter((r) => r["User ID"] === "u1"); + const u2Rows = teamRows.filter((r) => r["User ID"] === "u2"); + expect(u1Rows).toHaveLength(1); + expect(u2Rows).toHaveLength(1); + expect(Number(u2Rows[0]["Spend ($)"])).toBeCloseTo(3.3, 4); + }); + + it("should keep the same user in different teams as separate rows", () => { + const rows = generateDailyWithUsersData(usersFixture, "Team"); + const u1Rows = rows.filter((r) => r.Date === "2025-03-01" && r["User ID"] === "u1"); + + expect(u1Rows).toHaveLength(2); + const team1Row = u1Rows.find((r) => r["Team ID"] === "team-1"); + const team2Row = u1Rows.find((r) => r["Team ID"] === "team-2"); + expect(team1Row?.Keys).toBe(2); + expect(team2Row?.Keys).toBe(1); + expect(Number(team2Row?.["Spend ($)"])).toBeCloseTo(6.6, 4); + }); + + it("should show a dash email when the user has none", () => { + const rows = generateDailyWithUsersData(usersFixture, "Team"); + const row = rows.find((r) => r.Date === "2025-03-01" && r["Team ID"] === "team-1" && r["User ID"] === "u3"); + + expect(row).toBeDefined(); + expect(row?.["User Email"]).toBe("-"); + }); + + it("should still attribute a deleted key to its user", () => { + const rows = generateDailyWithUsersData(usersFixture, "Team"); + const row = rows.find((r) => r.Date === "2025-03-01" && r["Team ID"] === "team-1" && r["User ID"] === "u3"); + + expect(row).toBeDefined(); + expect(Number(row?.["Spend ($)"])).toBeCloseTo(5.5, 4); + expect(row?.Requests).toBe(55); + }); + + it("should group rows under team_id on the aggregated endpoint shape", () => { + const aggregatedFixture: EntitySpendData = { + results: [ + { + date: "2025-03-01", + breakdown: { + entities: {}, + api_keys: { + kA: { + metrics: { + spend: 1.1, + api_requests: 11, + successful_requests: 10, + failed_requests: 1, + total_tokens: 110, + prompt_tokens: 60, + completion_tokens: 50, + cache_read_input_tokens: 6, + cache_creation_input_tokens: 4, + }, + metadata: { team_id: "team-1", user_id: "u1", user_email: "a@x" }, + }, + kF: { + metrics: { + spend: 6.6, + api_requests: 66, + successful_requests: 64, + failed_requests: 2, + total_tokens: 660, + prompt_tokens: 370, + completion_tokens: 290, + cache_read_input_tokens: 36, + cache_creation_input_tokens: 24, + }, + metadata: { team_id: "team-2", user_id: "u2" }, + }, + }, + }, + }, + ], + metadata: usersFixture.metadata, + }; + + const rows = generateDailyWithUsersData(aggregatedFixture, "Team"); + + expect(rows).toHaveLength(2); + const team1Row = rows.find((r) => r["Team ID"] === "team-1"); + const team2Row = rows.find((r) => r["Team ID"] === "team-2"); + expect(team1Row?.["User ID"]).toBe("u1"); + expect(team2Row?.["User ID"]).toBe("u2"); + expect(Number(team1Row?.["Spend ($)"])).toBeCloseTo(1.1, 4); + expect(Number(team2Row?.["Spend ($)"])).toBeCloseTo(6.6, 4); + }); + + it("should emit the exact column order and sort by date ascending", () => { + const rows = generateDailyWithUsersData(usersFixture, "Team"); + + expect(Object.keys(rows[0])).toEqual([ + "Date", + "Team", + "Team ID", + "User ID", + "User Email", + "Keys", + "Spend ($)", + "Requests", + "Successful Requests", + "Failed Requests", + "Total Tokens", + "Prompt Tokens", + "Completion Tokens", + "Cache Read Input Tokens", + "Cache Creation Input Tokens", + ]); + + const dates = rows.map((r) => new Date(r.Date).getTime()); + for (let i = 0; i < dates.length - 1; i++) { + expect(dates[i]).toBeLessThanOrEqual(dates[i + 1]); + } + }); + + it("should dispatch daily_with_users through generateExportData", () => { + expect(generateExportData(usersFixture, "daily_with_users", "Team")).toEqual( + generateDailyWithUsersData(usersFixture, "Team"), + ); + }); + + it("should keep owners separate when entity and user ids contain underscores", () => { + const collisionFixture: EntitySpendData = { + results: [ + { + date: "2025-03-01", + breakdown: { + entities: { + team_1: { + metrics: { spend: 1, api_requests: 1, total_tokens: 10 }, + api_key_breakdown: { + kX: { + metrics: { spend: 1, api_requests: 1, total_tokens: 10 }, + metadata: { team_id: "team_1", user_id: "u1" }, + }, + }, + }, + team: { + metrics: { spend: 2, api_requests: 2, total_tokens: 20 }, + api_key_breakdown: { + kY: { + metrics: { spend: 2, api_requests: 2, total_tokens: 20 }, + metadata: { team_id: "team", user_id: "1_u1" }, + }, + }, + }, + }, + }, + }, + ], + metadata: usersFixture.metadata, + }; + + const rows = generateDailyWithUsersData(collisionFixture, "Team"); + + expect(rows).toHaveLength(2); + const team1Row = rows.find((r) => r["Team ID"] === "team_1"); + expect(team1Row?.["User ID"]).toBe("u1"); + expect(team1Row?.Keys).toBe(1); + expect(team1Row?.["Spend ($)"]).toBe("1.0000"); + const teamRow = rows.find((r) => r["Team ID"] === "team"); + expect(teamRow?.["User ID"]).toBe("1_u1"); + expect(teamRow?.Keys).toBe(1); + expect(teamRow?.["Spend ($)"]).toBe("2.0000"); + }); + + it("should leave daily and daily_with_models output without user columns", () => { + const daily = generateDailyData(usersFixture, "Team"); + expect(daily[0]).not.toHaveProperty("User ID"); + + const modelsFixture: EntitySpendData = { + results: [ + { + date: "2025-03-01", + breakdown: { + entities: { + "team-1": { + metrics: { + spend: 1.1, + api_requests: 11, + successful_requests: 10, + failed_requests: 1, + total_tokens: 110, + prompt_tokens: 60, + completion_tokens: 50, + }, + api_key_breakdown: { + kA: { + metrics: { + spend: 1.1, + api_requests: 11, + successful_requests: 10, + failed_requests: 1, + total_tokens: 110, + }, + metadata: { team_id: "team-1", user_id: "u1", user_email: "a@x" }, + }, + }, + }, + }, + models: { + "gpt-4o": { + metrics: { spend: 1.1, api_requests: 11, total_tokens: 110 }, + api_key_breakdown: { + kA: { + metrics: { + spend: 1.1, + api_requests: 11, + successful_requests: 10, + failed_requests: 1, + total_tokens: 110, + }, + metadata: {}, + }, + }, + }, + }, + }, + }, + ], + metadata: usersFixture.metadata, + }; + + const modelRows = generateDailyWithModelsData(modelsFixture, "Team"); + expect(modelRows).toHaveLength(1); + expect(Object.keys(modelRows[0])).toEqual([ + "Date", + "Team", + "Team ID", + "Model", + "Spend ($)", + "Requests", + "Successful", + "Failed", + "Total Tokens", + "Prompt Tokens", + "Completion Tokens", + "Cache Read Input Tokens", + "Cache Creation Input Tokens", + ]); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 8fd75134bcc..95ce584cc89 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -166,6 +166,8 @@ export const generateDailyWithKeysData = ( entityAlias: string; keyId: string; keyAlias: string | null; + userId: string | null; + userEmail: string | null; metrics: { spend: number; api_requests: number; @@ -200,6 +202,8 @@ export const generateDailyWithKeysData = ( entityAlias, keyId, keyAlias, + userId: keyData?.metadata?.user_id || null, + userEmail: keyData?.metadata?.user_email || null, metrics: { spend: keyData.metrics?.spend || 0, api_requests: keyData.metrics?.api_requests || 0, @@ -236,6 +240,7 @@ export const generateDailyWithKeysData = ( [`${entityLabel} ID`]: item.entityId, "Key Alias": item.keyAlias || "-", "Key ID": item.keyId, + ...(entityLabel === "User" ? {} : { "User ID": item.userId || "-", "User Email": item.userEmail || "-" }), "Spend ($)": formatNumberWithCommas(item.metrics.spend, 4), Requests: item.metrics.api_requests, "Successful Requests": item.metrics.successful_requests, @@ -250,6 +255,71 @@ export const generateDailyWithKeysData = ( return dailyKeyBreakdown.sort((a, b) => new Date(a.Date).getTime() - new Date(b.Date).getTime()); }; +export const generateDailyWithUsersData = ( + spendData: EntitySpendData, + entityLabel: string, + teamAliasMap: Record = {}, +): any[] => { + const aggregatedData: { + [key: string]: { + Date: string; + entityId: string; + entityAlias: string; + userId: string; + userEmail: string | null; + keyIds: Set; + metrics: Record<(typeof METRIC_KEYS)[number], number>; + }; + } = {}; + + spendData.results.forEach((day) => { + Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { + const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); + Object.entries(data.api_key_breakdown || {}).forEach(([keyId, keyData]: [string, any]) => { + const userId = keyData?.metadata?.user_id || "Unassigned"; + const uniqueKey = JSON.stringify([day.date, entityId, userId]); + if (!aggregatedData[uniqueKey]) { + aggregatedData[uniqueKey] = { + Date: day.date, + entityId, + entityAlias, + userId, + userEmail: null, + keyIds: new Set(), + metrics: Object.fromEntries(METRIC_KEYS.map((k) => [k, 0])) as Record<(typeof METRIC_KEYS)[number], number>, + }; + } + const bucket = aggregatedData[uniqueKey]; + bucket.userEmail = bucket.userEmail || keyData?.metadata?.user_email || null; + bucket.keyIds.add(keyId); + for (const k of METRIC_KEYS) { + bucket.metrics[k] += keyData?.metrics?.[k] || 0; + } + }); + }); + }); + + return Object.values(aggregatedData) + .map((item) => ({ + Date: item.Date, + [entityLabel]: item.entityAlias, + [`${entityLabel} ID`]: item.entityId, + "User ID": item.userId, + "User Email": item.userEmail || "-", + Keys: item.keyIds.size, + "Spend ($)": formatNumberWithCommas(item.metrics.spend, 4), + Requests: item.metrics.api_requests, + "Successful Requests": item.metrics.successful_requests, + "Failed Requests": item.metrics.failed_requests, + "Total Tokens": item.metrics.total_tokens, + "Prompt Tokens": item.metrics.prompt_tokens, + "Completion Tokens": item.metrics.completion_tokens, + "Cache Read Input Tokens": item.metrics.cache_read_input_tokens, + "Cache Creation Input Tokens": item.metrics.cache_creation_input_tokens, + })) + .sort((a, b) => new Date(a.Date).getTime() - new Date(b.Date).getTime()); +}; + export const generateDailyWithModelsData = ( spendData: EntitySpendData, entityLabel: string, @@ -340,6 +410,8 @@ export const generateExportData = ( return generateDailyWithKeysData(spendData, entityLabel, teamAliasMap); case "daily_with_models": return generateDailyWithModelsData(spendData, entityLabel, teamAliasMap); + case "daily_with_users": + return generateDailyWithUsersData(spendData, entityLabel, teamAliasMap); default: return generateDailyData(spendData, entityLabel, teamAliasMap); } diff --git a/ui/litellm-dashboard/src/components/agents/types.ts b/ui/litellm-dashboard/src/components/agents/types.ts index 24ff0c0e12c..6adb3fa9dda 100644 --- a/ui/litellm-dashboard/src/components/agents/types.ts +++ b/ui/litellm-dashboard/src/components/agents/types.ts @@ -21,6 +21,7 @@ export interface Agent { [key: string]: any; }; object_permission?: AgentObjectPermission; + access_group_ids?: string[] | null; keys?: AgentAttachedKey[] | null; spend?: number; tpm_limit?: number | null; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index f4011a653ab..82399792674 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6211,6 +6211,7 @@ export const patchAgentCall = async ( rpm_limit?: number | null; session_tpm_limit?: number | null; session_rpm_limit?: number | null; + access_group_ids?: string[]; }, ) => { try { diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx index 17770c54482..5492ceffefc 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx @@ -113,6 +113,148 @@ describe("RoutingDecisionCard", () => { }, ); + it.each(["capability_classifier", "modality_escalation"])( + "shows the recorded Capability forecast for %s", + (cause) => { + render( + , + ); + + expect(screen.getByText("Capability estimates")).toBeInTheDocument(); + expect(screen.getByText("Efficient model solve chance")).toBeInTheDocument(); + expect(screen.getAllByText(/^\d+\.\d%$/).map((value) => value.textContent)).toEqual(["0.0%", "86.4%", "82.0%"]); + expect(screen.getByText("Raw")).toBeInTheDocument(); + expect(screen.getByText("Calibrated")).toBeInTheDocument(); + expect(screen.getByText("Threshold")).toBeInTheDocument(); + expect(screen.getByText("uncertain")).toBeInTheDocument(); + expect(screen.getByText("UNC-2")).toBeInTheDocument(); + expect(screen.getByText("calibration-1")).toBeInTheDocument(); + expect(screen.getByText("Deep")).toBeInTheDocument(); + expect(screen.queryByText("FUSE v2 estimates")).not.toBeInTheDocument(); + }, + ); + + it("omits absent Capability fields while preserving a recorded zero threshold", () => { + render( + , + ); + + expect(screen.getAllByText(/^\d+\.\d%$/).map((value) => value.textContent)).toEqual(["25.0%", "0.0%"]); + for (const label of ["Calibrated", "Calibration", "Boundary", "Rule"]) { + expect(screen.queryByText(label)).not.toBeInTheDocument(); + } + }); + + it.each(["llm_v2_classifier", "default_fallback"])( + "shows the original calibrated FUSE v2 forecast for %s", + (cause) => { + render( + , + ); + + expect(screen.getByText("FUSE v2 estimates")).toBeInTheDocument(); + expect(screen.getAllByText(/^\d+\.\d%$/).map((value) => value.textContent)).toEqual([ + "25.0%", + "91.0%", + "75.0%", + "80.0%", + ]); + expect(screen.getByText("Efficient (raw)")).toBeInTheDocument(); + expect(screen.getByText("Capable (raw)")).toBeInTheDocument(); + expect(screen.getByText("Efficient (calibrated)")).toBeInTheDocument(); + expect(screen.getByText("Capable (calibrated)")).toBeInTheDocument(); + expect(screen.getAllByText(/percentage points$/).map((value) => value.textContent)).toEqual([ + "5.0 percentage points", + "10.0 percentage points", + ]); + expect(screen.getByText("Applied gap")).toBeInTheDocument(); + expect(screen.getByText("Allowed gap")).toBeInTheDocument(); + expect(screen.getByText("calibration-2")).toBeInTheDocument(); + expect(screen.getByText("llm-v2:verification=tests")).toBeInTheDocument(); + expect(screen.getByText("fallback-model")).toBeInTheDocument(); + expect(screen.queryByText("Capability estimates")).not.toBeInTheDocument(); + }, + ); + + it("uses raw FUSE v2 probabilities without calibration and preserves negative and zero gaps", () => { + render( + , + ); + + expect(screen.getAllByText(/^\d+\.\d%$/).map((value) => value.textContent)).toEqual(["50.0%", "0.0%"]); + expect(screen.getAllByText(/percentage points$/).map((value) => value.textContent)).toEqual([ + "-50.0 percentage points", + "0.0 percentage points", + ]); + expect(screen.queryByText(/calibrated|Calibration/)).not.toBeInTheDocument(); + }); + + it.each([ + { classifier_efficient_p_solve: 0, classifier_max_quality_gap: 0.2 }, + { + classifier_efficient_p_solve: 0.4, + classifier_capable_p_solve: 0.9, + classifier_calibrated_efficient_p_solve: 0, + classifier_calibration_version: "partial-calibration", + classifier_max_quality_gap: 0.2, + }, + ])("shows partial FUSE v2 estimates without inventing an applied gap: %j", (fields) => { + render(); + + expect(screen.getByText("FUSE v2 estimates")).toBeInTheDocument(); + expect(screen.getByText("0.0%")).toBeInTheDocument(); + expect(screen.getByText("20.0 percentage points")).toBeInTheDocument(); + expect(screen.queryByText("Applied gap")).not.toBeInTheDocument(); + expect(screen.queryByText("Capable (calibrated)")).not.toBeInTheDocument(); + }); + + it.each([ + ["capability_classifier", "Capability"], + ["llm_v2_classifier", "FUSE v2"], + ["capability_classifier_fallback", "Capable tier, Capability classifier failed"], + ["llm_v2_fallback", "Capable tier, FUSE v2 classifier failed"], + ["session_affinity_pin", "Pinned to session"], + ])("labels %s without inventing a missing forecast", (cause, label) => { + render(); + + expect(screen.getByText(label)).toBeInTheDocument(); + expect(screen.getByText("MEDIUM")).toBeInTheDocument(); + expect(screen.queryByText(/estimates/)).not.toBeInTheDocument(); + }); + it("uses the persisted boundary snapshot, not today's defaults", () => { // Same score, boundaries the operator had configured lower: it lands in a // different band, and the card must say so. diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx index c497b2a94a7..6a8cedd3746 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -24,6 +24,17 @@ export interface RoutingDecision { matched_keyword?: string; escalation_keyword?: string; classifier_model?: string; + classifier_p_solve?: number; + classifier_calibrated_p_solve?: number; + classifier_threshold?: number; + classifier_capability_boundary?: string; + classifier_primary_rule?: string; + classifier_calibration_version?: string; + classifier_efficient_p_solve?: number; + classifier_capable_p_solve?: number; + classifier_calibrated_efficient_p_solve?: number; + classifier_calibrated_capable_p_solve?: number; + classifier_max_quality_gap?: number; classifier_confidence?: number; classifier_probabilities?: Record; classifier_cost?: number; @@ -94,6 +105,10 @@ function describeReasoningOverride(tierLabel: string | undefined, floor: number const CONSTANT_CAUSE_LABELS: Record = { heuristic_scorer: "Heuristic scorer", heuristic_v2: "Heuristic v2", + capability_classifier: "Capability", + capability_classifier_fallback: "Capable tier, Capability classifier failed", + llm_v2_classifier: "FUSE v2", + llm_v2_fallback: "Capable tier, FUSE v2 classifier failed", heuristic_first_short_circuit: "Heuristic scorer, classifier skipped", hybrid_short_circuit: "Heuristic scorer, score clear of every boundary", classifier_plugin: "Custom classifier plugin", @@ -161,6 +176,71 @@ function Row({ label, children }: { label: string; children: React.ReactNode }) ); } +function PercentageRow({ label, value, unit = "%" }: { label: string; value?: number; unit?: string }) { + if (value === undefined) return null; + return ( + + {`${(value * 100).toFixed(1)}${unit}`} + + ); +} + +function CapabilityForecast({ decision }: { decision: RoutingDecision }) { + const { + classifier_p_solve: raw, + classifier_calibrated_p_solve: calibrated, + classifier_threshold: threshold, + classifier_capability_boundary: boundary, + classifier_primary_rule: rule, + classifier_calibration_version: version, + } = decision; + if ([raw, calibrated, threshold, boundary, rule].every((value) => value === undefined)) return null; + + return ( +
+
Capability estimates
+
Efficient model solve chance
+ + + + {boundary && {boundary}} + {rule && {rule}} + {version && {version}} +
+ ); +} + +function FuseV2Forecast({ decision }: { decision: RoutingDecision }) { + const { + classifier_efficient_p_solve: rawEfficient, + classifier_capable_p_solve: rawCapable, + classifier_calibrated_efficient_p_solve: calibratedEfficient, + classifier_calibrated_capable_p_solve: calibratedCapable, + classifier_max_quality_gap: allowedGap, + classifier_calibration_version: version, + } = decision; + if ([rawEfficient, rawCapable, calibratedEfficient, calibratedCapable, allowedGap].every((v) => v === undefined)) { + return null; + } + const isCalibrated = [calibratedEfficient, calibratedCapable, version].some((value) => value !== undefined); + const efficient = isCalibrated ? calibratedEfficient : rawEfficient; + const capable = isCalibrated ? calibratedCapable : rawCapable; + const gap = efficient !== undefined && capable !== undefined ? capable - efficient : undefined; + + return ( +
+
FUSE v2 estimates
+ + + + + + + {version && {version}} +
+ ); +} + export function RoutingDecisionCard({ decision, className, @@ -266,6 +346,9 @@ export function RoutingDecisionCard({
)} + + + {signals && signals.length > 0 && ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 58b5cedc2cf..6a184d0765d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -4872,6 +4872,27 @@ export interface paths { patch: operations["assemblyai_proxy_route_eu_assemblyai__endpoint__patch"]; trace?: never; }; + "/fal_ai/{endpoint}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Fal Ai Proxy Route */ + get: operations["fal_ai_proxy_route_fal_ai__endpoint__get"]; + /** Fal Ai Proxy Route */ + put: operations["fal_ai_proxy_route_fal_ai__endpoint__put"]; + /** Fal Ai Proxy Route */ + post: operations["fal_ai_proxy_route_fal_ai__endpoint__post"]; + /** Fal Ai Proxy Route */ + delete: operations["fal_ai_proxy_route_fal_ai__endpoint__delete"]; + options?: never; + head?: never; + /** Fal Ai Proxy Route */ + patch: operations["fal_ai_proxy_route_fal_ai__endpoint__patch"]; + trace?: never; + }; "/fallback": { parameters: { query?: never; @@ -10646,6 +10667,27 @@ export interface paths { patch: operations["openai_passthrough_route_openai_passthrough__endpoint__patch"]; trace?: never; }; + "/openrouter/{endpoint}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Openrouter Proxy Route */ + get: operations["openrouter_proxy_route_openrouter__endpoint__get"]; + /** Openrouter Proxy Route */ + put: operations["openrouter_proxy_route_openrouter__endpoint__put"]; + /** Openrouter Proxy Route */ + post: operations["openrouter_proxy_route_openrouter__endpoint__post"]; + /** Openrouter Proxy Route */ + delete: operations["openrouter_proxy_route_openrouter__endpoint__delete"]; + options?: never; + head?: never; + /** Openrouter Proxy Route */ + patch: operations["openrouter_proxy_route_openrouter__endpoint__patch"]; + trace?: never; + }; "/organization/daily/activity": { parameters: { query?: never; @@ -16558,6 +16600,64 @@ export interface paths { patch?: never; trace?: never; }; + "/tinyfish/{endpoint}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Tinyfish Proxy Route + * @description Pass-through for the TinyFish Agent API (goal-based web automation). + * + * Forwarded endpoints: + * - POST /v1/automation/run — run to completion (blocking) + * - POST /v1/automation/run-async — submit a run, poll GET /v1/runs/{id} for the result + * - POST /v1/automation/run-sse — run with SSE progress events + * - GET /v1/runs/{id} — run status / result + * - POST /v1/runs/{id}/cancel — cancel a run + * + * Every other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs + * listing, which would let any caller discover other callers' run ids) returns 403: all + * proxy callers share one upstream key. + * + * Credential lookup order: + * 1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through) + * 2. TINYFISH_API_KEY environment variable + * + * [Docs](https://docs.litellm.ai/docs/pass_through/tinyfish) + */ + get: operations["tinyfish_proxy_route_tinyfish__endpoint__get"]; + put?: never; + /** + * Tinyfish Proxy Route + * @description Pass-through for the TinyFish Agent API (goal-based web automation). + * + * Forwarded endpoints: + * - POST /v1/automation/run — run to completion (blocking) + * - POST /v1/automation/run-async — submit a run, poll GET /v1/runs/{id} for the result + * - POST /v1/automation/run-sse — run with SSE progress events + * - GET /v1/runs/{id} — run status / result + * - POST /v1/runs/{id}/cancel — cancel a run + * + * Every other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs + * listing, which would let any caller discover other callers' run ids) returns 403: all + * proxy callers share one upstream key. + * + * Credential lookup order: + * 1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through) + * 2. TINYFISH_API_KEY environment variable + * + * [Docs](https://docs.litellm.ai/docs/pass_through/tinyfish) + */ + post: operations["tinyfish_proxy_route_tinyfish__endpoint__post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/token": { parameters: { query?: never; @@ -23640,6 +23740,8 @@ export interface components { }; /** AgentConfig */ AgentConfig: { + /** Access Group Ids */ + access_group_ids?: string[] | null; agent_card_params: components["schemas"]["AgentCard"]; /** Agent Name */ agent_name: string; @@ -23787,6 +23889,8 @@ export interface components { }; /** AgentResponse */ AgentResponse: { + /** Access Group Ids */ + access_group_ids?: string[] | null; /** Agent Card Params */ agent_card_params: { [key: string]: unknown; @@ -25831,7 +25935,7 @@ export interface components { * CallTypes * @enum {string} */ - CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; + CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; /** CallbackDelete */ CallbackDelete: { /** Callback Name */ @@ -27115,6 +27219,11 @@ export interface components { * @default false */ health_check_skip_disabled_background_models: boolean; + /** + * Include Call Id In Error Body + * @description opt-in to copy the x-litellm-call-id response header's value into JSON error bodies, as error.litellm_call_id on the OpenAI-shaped and /v1/messages routes and as a top-level litellm_call_id on pass-through routes, so an error a client prints names the request to look up. Off by default + */ + include_call_id_in_error_body?: boolean | null; /** * Infer Model From Keys * @description for `/models` endpoint, infers available model based on environment keys (e.g. OPENAI_API_KEY) @@ -27979,7 +28088,9 @@ export interface components { /** Jwt Issuer */ jwt_issuer?: string | null; /** Key */ - key: string; + key?: string | null; + /** Token */ + token?: string | null; }; /** CreateSearchToolRequest */ CreateSearchToolRequest: { @@ -31226,6 +31337,12 @@ export interface components { output_cost_per_character_above_128k_tokens?: number | null; /** Output Cost Per Image */ output_cost_per_image?: number | null; + /** Output Cost Per Image 1024 */ + output_cost_per_image_1024?: number | null; + /** Output Cost Per Image 1536 */ + output_cost_per_image_1536?: number | null; + /** Output Cost Per Image 512 */ + output_cost_per_image_512?: number | null; /** Output Cost Per Image Token */ output_cost_per_image_token?: number | null; /** Output Cost Per Pixel */ @@ -34983,6 +35100,8 @@ export interface components { }; /** PatchAgentRequest */ PatchAgentRequest: { + /** Access Group Ids */ + access_group_ids?: string[] | null; agent_card_params?: components["schemas"]["AgentCard"]; /** Agent Name */ agent_name?: string; @@ -39850,6 +39969,8 @@ export interface components { jwt_issuer?: string | null; /** Key */ key?: string | null; + /** Token */ + token?: string | null; }; /** UpdateKeyRequest */ UpdateKeyRequest: { @@ -42051,6 +42172,12 @@ export interface components { output_cost_per_character_above_128k_tokens?: number | null; /** Output Cost Per Image */ output_cost_per_image?: number | null; + /** Output Cost Per Image 1024 */ + output_cost_per_image_1024?: number | null; + /** Output Cost Per Image 1536 */ + output_cost_per_image_1536?: number | null; + /** Output Cost Per Image 512 */ + output_cost_per_image_512?: number | null; /** Output Cost Per Image Token */ output_cost_per_image_token?: number | null; /** Output Cost Per Pixel */ @@ -49341,6 +49468,161 @@ export interface operations { }; }; }; + fal_ai_proxy_route_fal_ai__endpoint__get: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + fal_ai_proxy_route_fal_ai__endpoint__put: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + fal_ai_proxy_route_fal_ai__endpoint__post: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + fal_ai_proxy_route_fal_ai__endpoint__delete: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + fal_ai_proxy_route_fal_ai__endpoint__patch: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; create_fallback_fallback_post: { parameters: { query?: never; @@ -56330,6 +56612,161 @@ export interface operations { }; }; }; + openrouter_proxy_route_openrouter__endpoint__get: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + openrouter_proxy_route_openrouter__endpoint__put: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + openrouter_proxy_route_openrouter__endpoint__post: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + openrouter_proxy_route_openrouter__endpoint__delete: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + openrouter_proxy_route_openrouter__endpoint__patch: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_organization_daily_activity_organization_daily_activity_get: { parameters: { query?: { @@ -62643,6 +63080,68 @@ export interface operations { }; }; }; + tinyfish_proxy_route_tinyfish__endpoint__get: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + tinyfish_proxy_route_tinyfish__endpoint__post: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; token_endpoint_token_post: { parameters: { query?: { diff --git a/uv.lock b/uv.lock index 543581cbc23..32d7580f61b 100644 --- a/uv.lock +++ b/uv.lock @@ -1162,14 +1162,11 @@ wheels = [ [[package]] name = "click" -version = "8.3.3" +version = "8.5.0" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/bb/63/f9e1ea081ce35720d8b92acde70daaedace594dc93b693c869e0d5910718/click-8.3.3.tar.gz", hash = "sha256:398329ad4837b2ff7cbe1dd166a4c0f8900c3ca3a218de04466f38f6497f18a2", size = 328061, upload-time = "2026-04-22T15:11:27.506Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/0e/7fa0ef50764b67090eca4114772a2abf8b6148198475e54c660b97caeee6/click-8.5.0.tar.gz", hash = "sha256:ba0d2089de75ea0310e2dde03160e6ca10009947fb95a182f9b54021bb272e34", size = 382235, upload-time = "2026-08-26T13:33:14.56Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ae/44/c1221527f6a71a01ec6fbad7fa78f1d50dfa02217385cf0fa3eec7087d59/click-8.3.3-py3-none-any.whl", hash = "sha256:a2bf429bb3033c89fa4936ffb35d5cb471e3719e1f3c8a7c3fff0b8314305613", size = 110502, upload-time = "2026-04-22T15:11:25.044Z" }, + { url = "https://files.pythonhosted.org/packages/58/50/6c0d534c5f134586a8e1ba4e330569e32f057e33372ae556463212fb4cd3/click-8.5.0-py3-none-any.whl", hash = "sha256:255bc9599cf7748b4b1a446ccc735421bd08a2ae529a8b88597d3de5664ee360", size = 125251, upload-time = "2026-08-26T13:33:12.928Z" }, ] [[package]] @@ -1989,11 +1986,11 @@ wheels = [ [[package]] name = "filelock" -version = "3.29.0" +version = "3.32.7" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b5/fe/997687a931ab51049acce6fa1f23e8f01216374ea81374ddee763c493db5/filelock-3.29.0.tar.gz", hash = "sha256:69974355e960702e789734cb4871f884ea6fe50bd8404051a3530bc07809cf90", size = 57571, upload-time = "2026-04-19T15:39:10.068Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/59/e19834834cb01a32febfbb0f8a23a9088088f5d45991824ff2bc3b5e8acb/filelock-3.32.7.tar.gz", hash = "sha256:37b8a3d9811b0f9aef7e5ec5c71bb320de52df51e6ca9bcd6f5ad81187660da7", size = 225154, upload-time = "2026-09-16T00:24:20.907Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/81/47/dd9a212ef6e343a6857485ffe25bba537304f1913bdbed446a23f7f592e1/filelock-3.29.0-py3-none-any.whl", hash = "sha256:96f5f6344709aa1572bbf631c640e4ebeeb519e08da902c39a001882f30ac258", size = 39812, upload-time = "2026-04-19T15:39:08.752Z" }, + { url = "https://files.pythonhosted.org/packages/15/df/31098c5aeb4d966b553641472bd55fcf5fdfac953549894b8a765ba44e91/filelock-3.32.7-py3-none-any.whl", hash = "sha256:65ff0d0190ea42038b32bda4b77834fb05be2cad4c5b9b01aa4dfb3614536e52", size = 100157, upload-time = "2026-09-16T00:24:19.543Z" }, ] [[package]] @@ -2225,11 +2222,11 @@ wheels = [ [[package]] name = "fsspec" -version = "2026.4.0" +version = "2026.7.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d5/8d/1c51c094345df128ca4a990d633fe1a0ff28726c9e6b3c41ba65087bba1d/fsspec-2026.4.0.tar.gz", hash = "sha256:301d8ac70ae90ef3ad05dcf94d6c3754a097f9b5fe4667d2787aa359ec7df7e4", size = 312760, upload-time = "2026-04-29T20:42:38.635Z" } +sdist = { url = "https://files.pythonhosted.org/packages/00/78/f34251dadb8f3921264a1d9b8946f5e542014ee2614b285261b4e40e6775/fsspec-2026.7.0.tar.gz", hash = "sha256:c803c40f4cf860b49dea58ee3e1c33cb9c790520e233537e1340049f89b82a88", size = 317040, upload-time = "2026-07-28T16:34:51.052Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d5/0c/043d5e551459da400957a1395e0febbf771446ff34291afcbe3d8be2a279/fsspec-2026.4.0-py3-none-any.whl", hash = "sha256:11ef7bb35dab8a394fde6e608221d5cf3e8499401c249bebaeaad760a1a8dec2", size = 203402, upload-time = "2026-04-29T20:42:36.842Z" }, + { url = "https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl", hash = "sha256:b57ddbafedfaef7018c1ecab32aa200a9d7ca26b77965f64e48b70061249d279", size = 206583, upload-time = "2026-07-28T16:34:49.538Z" }, ] [[package]] @@ -3146,34 +3143,26 @@ wheels = [ [[package]] name = "hf-xet" -version = "1.5.0" +version = "1.6.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/d8/5c06fc76461418326a7decf8367480c35be11a41fd938633929c60a9ec6b/hf_xet-1.5.0.tar.gz", hash = "sha256:e0fb0a34d9f406eed88233e829a67ec016bec5af19e480eac65a233ea289a948", size = 837196, upload-time = "2026-05-06T06:18:15.583Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1b/ab/522a2ab67f27971a9d48ca666d4fca85ef7d5282d142e31fd087e27b1bbe/hf_xet-1.6.0.tar.gz", hash = "sha256:2e58454a340b3556dfa4972d5451aff4fba8dd42a236600ba1a1d2b1514f0fef", size = 920527, upload-time = "2026-08-03T22:33:13.243Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/68/9b/6912c99070915a4f28119e3c5b52a9abd1eec0ad5cb293b8c967a0c6f5a2/hf_xet-1.5.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:7d70fe2ce97b9db73b9c9b9c81fe3693640aec83416a966c446afea54acfae3c", size = 4023383, upload-time = "2026-05-06T06:17:53.947Z" }, - { url = "https://files.pythonhosted.org/packages/0f/6d/9563cfde59b5d8128a9c7ec972a087f4c782e4f7bac5a85234edfd5d5e49/hf_xet-1.5.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:73a0dae8c71de3b0633a45c73f4a4a5ed09e94b43441d82981a781d4f12baa42", size = 3792751, upload-time = "2026-05-06T06:17:51.791Z" }, - { url = "https://files.pythonhosted.org/packages/07/a5/ed5a0cf35b49a0571af5a8f53416dad1877a718c021c9937c3a53cb45781/hf_xet-1.5.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a60290ec57e9b71767fba7c3645ddafdd0759974b540441510c629c6db6db24a", size = 4456058, upload-time = "2026-05-06T06:17:40.735Z" }, - { url = "https://files.pythonhosted.org/packages/60/fb/3ae8bf2a7a37a4197d0195d7247fd25b3952e15cb8a599e285dfaa6f52b3/hf_xet-1.5.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e5de0f6deada0dada870bb376a11bcd1f08abf3a968a6d118f33e72d1b1eb480", size = 4250783, upload-time = "2026-05-06T06:17:38.412Z" }, - { url = "https://files.pythonhosted.org/packages/a2/9b/8bae40d4d91525085137196e84eb0ed49cf65b5e96e5c3ecdadd8bd0fac2/hf_xet-1.5.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c799d49f1a5544a0ef7591c0ee75e0d6b93d6f56dc7a4979f59f7518d2872216", size = 4445594, upload-time = "2026-05-06T06:18:04.219Z" }, - { url = "https://files.pythonhosted.org/packages/13/59/c74efbbd4e8728172b2cc72a2bc014d2947a4b7bdced932fbd3f5da1a4e5/hf_xet-1.5.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:2baea1b0b989e5c152fe81425f7745ddc8901280ba3d97c98d8cdece7b706c60", size = 4663995, upload-time = "2026-05-06T06:18:06.1Z" }, - { url = "https://files.pythonhosted.org/packages/73/32/8e1e0410af64cda9b139d1dcebdc993a8ff9c8c7c0e2696ae356d75ccc0d/hf_xet-1.5.0-cp313-cp313t-win_amd64.whl", hash = "sha256:526345b3ed45f374f6317349df489167606736c876241ba984105afe7fd4839d", size = 3966608, upload-time = "2026-05-06T06:18:19.74Z" }, - { url = "https://files.pythonhosted.org/packages/fc/34/a8febc8f4edbea8b3e21b02ebc8b628679b84ba7e45cde624a7736b51500/hf_xet-1.5.0-cp313-cp313t-win_arm64.whl", hash = "sha256:786d28e2eb8315d5035544b9d137b4a842d600c434bb91bf7d0d953cce906ad4", size = 3796946, upload-time = "2026-05-06T06:18:17.568Z" }, - { url = "https://files.pythonhosted.org/packages/2a/20/8fc8996afe5815fa1a6be8e9e5c02f24500f409d599e905800d498a4e14d/hf_xet-1.5.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:872d5601e6deea30d15865ede55d29eac6daf5a534ab417b99b6ef6b076dd96c", size = 4023495, upload-time = "2026-05-06T06:18:01.94Z" }, - { url = "https://files.pythonhosted.org/packages/32/6a/93d84463c00cecb561a7508aa6303e35ee2894294eac14245526924415fe/hf_xet-1.5.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:9929561f5abf4581c8ea79587881dfef6b8abb2a0d8a51915936fc2a614f4e73", size = 3792731, upload-time = "2026-05-06T06:18:00.021Z" }, - { url = "https://files.pythonhosted.org/packages/9d/5a/8ec8e0c863b382d00b3c2e2af6ded6b06371be617144a625903a6d562f4b/hf_xet-1.5.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f7b7bbae318e583a86fb21e5a4a175d6721d628a2874f4bd022d0e660c32a682", size = 4456738, upload-time = "2026-05-06T06:17:49.574Z" }, - { url = "https://files.pythonhosted.org/packages/c5/ca/f7effa1a67717da2bcc6b6c28f71c6ca648c77acaec4e2c32f40cbe16d85/hf_xet-1.5.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:cf7b2dc6f31a4ea754bb50f74cde482dcf5d366d184076d8530b9872787f3761", size = 4251622, upload-time = "2026-05-06T06:17:47.096Z" }, - { url = "https://files.pythonhosted.org/packages/65/f2/19247dba3e231cf77dec59ddfb878f00057635ff773d099c9b59d37812c3/hf_xet-1.5.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8dbcbab554c9ef158ef2c991545c3e970ddd8cc7acdcd0a78c5a41095dab4ded", size = 4445667, upload-time = "2026-05-06T06:18:11.983Z" }, - { url = "https://files.pythonhosted.org/packages/7f/64/6f116801a3bcfb6f59f5c251f48cadc47ea54026441c4a385079286a94fa/hf_xet-1.5.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5906bf7718d3636dc13402914736abe723492cb730f744834f5f5b67d3a12702", size = 4664619, upload-time = "2026-05-06T06:18:13.771Z" }, - { url = "https://files.pythonhosted.org/packages/5c/e8/069542d37946ed08669b127e1496fa99e78196d71de8d41eda5e9f1b7a58/hf_xet-1.5.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5f3dc2248fc01cc0a00cd392ab497f1ca373fcbc7e3f2da1f452480b384e839e", size = 3966802, upload-time = "2026-05-06T06:18:28.162Z" }, - { url = "https://files.pythonhosted.org/packages/f9/91/fc6fdec27b14d04e88c386ac0a0129732b53fa23f7c4a78f4b83a039c567/hf_xet-1.5.0-cp314-cp314t-win_arm64.whl", hash = "sha256:b285cea1b5bab46b758772716ba8d6854a1a0310fed1c249d678a8b38601e5a0", size = 3797168, upload-time = "2026-05-06T06:18:26.287Z" }, - { url = "https://files.pythonhosted.org/packages/3d/fb/69ff198a82cae7eb1a69fb84d93b3a3e4816564d76817fe541ddc96874eb/hf_xet-1.5.0-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:dad0dc84e941b8ba3c860659fe1fdc35c049d47cce293f003287757e971a8f56", size = 4030814, upload-time = "2026-05-06T06:17:57.933Z" }, - { url = "https://files.pythonhosted.org/packages/9b/ff/edcc2b40162bef3ff78e14ab637e5f3b89243d6aee72f5949d3bb6a5af83/hf_xet-1.5.0-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:fd6e5a9b0fdac4ed03ed45ef79254a655b1aaab514a02202617fbf643f5fdf7a", size = 3798444, upload-time = "2026-05-06T06:17:55.79Z" }, - { url = "https://files.pythonhosted.org/packages/49/4d/103f76b04310e5e57656696cc184690d20c466af0bca3ca88f8c8ea5d4f3/hf_xet-1.5.0-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3531b1823a0e6d77d80f9ed15ca0e00f0d115094f8ac033d5cae88f4564cc949", size = 4465986, upload-time = "2026-05-06T06:17:44.886Z" }, - { url = "https://files.pythonhosted.org/packages/c4/a2/546f47f464737b3edbab6f8ddb57f2599b93d2cbb66f06abb475ccb48651/hf_xet-1.5.0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:9a0ee58cd18d5ea799f7ed11290bbccbe56bdd8b1d97ca74b9cc49a3945d7a3b", size = 4259865, upload-time = "2026-05-06T06:17:42.639Z" }, - { url = "https://files.pythonhosted.org/packages/95/7f/1be593c1f28613be2e196473481cd81bfc5910795e30a34e8f744f6cac4f/hf_xet-1.5.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1e60df5a42e9bed8628b6416af2cba4cba57ae9f02de226a06b020d98e1aab18", size = 4459835, upload-time = "2026-05-06T06:18:08.026Z" }, - { url = "https://files.pythonhosted.org/packages/aa/b2/703569fc881f3284487e68cda7b42179978480da3c438042a6bbbb4a671c/hf_xet-1.5.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:4b35549ce62601b84da4ff9b24d970032ace3d4430f52d91bcbb26c901d6c690", size = 4672414, upload-time = "2026-05-06T06:18:09.864Z" }, - { url = "https://files.pythonhosted.org/packages/af/37/1b6def445c567286b50aa3b33828158e135b1be44938dde59f11382a500c/hf_xet-1.5.0-cp37-abi3-win_amd64.whl", hash = "sha256:2806c7c17b4d23f8d88f7c4814f838c3b6150773fe339c20af23e1cfaf2797e4", size = 3977238, upload-time = "2026-05-06T06:18:23.621Z" }, - { url = "https://files.pythonhosted.org/packages/62/94/3b66b148778ee100dcfd69c2ca22b57b41b44d3063ceec934f209e9184ce/hf_xet-1.5.0-cp37-abi3-win_arm64.whl", hash = "sha256:b6c9df403040248c76d808d3e047d64db2d923bae593eb244c41e425cf6cd7be", size = 3806916, upload-time = "2026-05-06T06:18:21.7Z" }, + { url = "https://files.pythonhosted.org/packages/41/62/3c062f593bd92ef4e77a0ef39541e3d82a0a1d3947c8a777a02a13a27828/hf_xet-1.6.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:70cbb9c896901600128cb9b6f06e132954fbede1db30f31f7c6c63f84cb7c31d", size = 4074584, upload-time = "2026-08-03T22:32:47.364Z" }, + { url = "https://files.pythonhosted.org/packages/bb/1e/c0ad437dd267a8e435bef594acf781bbc3874ff0b6435b4962d03ecf7cc4/hf_xet-1.6.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:23379c2f9ec8696d952b16414a2bae72cad86a52df869b050698ba60f538c675", size = 3867381, upload-time = "2026-08-03T22:32:49.049Z" }, + { url = "https://files.pythonhosted.org/packages/d5/ee/7c0d7b6ab336167531b1c30af2af003f054af4c749becbd7209ae33a77c3/hf_xet-1.6.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f2f7278c05c22fd60cb436cda1269649b3e81db65ecdc8496e5e164aa4143e7b", size = 4453982, upload-time = "2026-08-03T22:32:50.568Z" }, + { url = "https://files.pythonhosted.org/packages/63/06/ad8eab1c9525246650cbaa821caa3cdbaca734ab1a5b8c91bea09cbd8d69/hf_xet-1.6.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:948f15d3a9545cfe5932f6bd8b440f6ae630aee108f14b7bd6c561f7c2dcc522", size = 4249445, upload-time = "2026-08-03T22:32:52.391Z" }, + { url = "https://files.pythonhosted.org/packages/d8/26/1eee8aedb0dafc1ab9717dc9ac602cde33361b232dc06803f1f6ed18b58c/hf_xet-1.6.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5153e6bb103ad49d6ea9f1b2e230db5a2ea32551ad09a706d2f61d7c7c80d80e", size = 4451099, upload-time = "2026-08-03T22:32:54.114Z" }, + { url = "https://files.pythonhosted.org/packages/67/57/0b88af1f194ab6c9c650547d9cc06bfeaab836ae4dcdb331676bfb8be95a/hf_xet-1.6.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:35cec30d75c6f9eb9c16a77cef68e85a103b72e24d4b473714ec9ff06428bab9", size = 4664712, upload-time = "2026-08-03T22:32:55.547Z" }, + { url = "https://files.pythonhosted.org/packages/53/a0/26b717a9d1840e8abf48dcec64b5ed8fbe472671d38ad28d30e147132b33/hf_xet-1.6.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5789835d7c6bc9436962853192082374297fb72d7eff7e7762ec25ceb7e25338", size = 4025906, upload-time = "2026-08-03T22:32:57.391Z" }, + { url = "https://files.pythonhosted.org/packages/49/f6/4a9966633c6fef83af997e2cff68ec1963676d412bdfd096df2a93b8e185/hf_xet-1.6.0-cp314-cp314t-win_arm64.whl", hash = "sha256:75765820ce4700db3750c94acc8fe27c5fae4c9ec000a0dbac3ca082acf97765", size = 3849221, upload-time = "2026-08-03T22:32:59.123Z" }, + { url = "https://files.pythonhosted.org/packages/a2/50/7afa2c9c787405864fc47a0d1bbc02c62e9101947ed43c1f43899fc7d91d/hf_xet-1.6.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:633dc0cd71d32da58ab8c03ad38e2fac452c15c2b0a2866ebf6ededfe0a5061d", size = 4071729, upload-time = "2026-08-03T22:33:00.721Z" }, + { url = "https://files.pythonhosted.org/packages/4b/69/55b8dcf636142ae660fec1869fcac14c4da2e8412e14d6eee1523be77e9f/hf_xet-1.6.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:f0906082d9932ae0c0057fa194041c22b4e2cdb46b2592ef3b91f020d62a081a", size = 3876287, upload-time = "2026-08-03T22:33:02.251Z" }, + { url = "https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d62671bb130879cef0ee4c9ebe47a14af6c66ec53e6d84dc15936e5ffdfac82f", size = 4464663, upload-time = "2026-08-03T22:33:03.802Z" }, + { url = "https://files.pythonhosted.org/packages/9a/69/1f0cbc2fb22ae6082d094f743d1b8945a3f36f6089cb95f42b7ee348cda7/hf_xet-1.6.0-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0e6e21fa3cdfcdcd76748564bf593870a5e013f47d97cf10aed63aa222cff5b7", size = 4262538, upload-time = "2026-08-03T22:33:05.287Z" }, + { url = "https://files.pythonhosted.org/packages/d1/3a/4f4f2301ade26e404462d3336fa11f7958d914cabbabdd6e03c3c5d5658c/hf_xet-1.6.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:4fc74352a17015bd0ee90038bc9efe38db894cde45f268b6712b04fce8cd0acb", size = 4460520, upload-time = "2026-08-03T22:33:06.81Z" }, + { url = "https://files.pythonhosted.org/packages/ab/5f/311725e2a905534dfee2dcb5b08414f249147f1f12252bfc2bd24caa075c/hf_xet-1.6.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8fb4f71cba6129110c3374a33f919001ff130488fc23553698e34cc1c2a1198c", size = 4675937, upload-time = "2026-08-03T22:33:08.616Z" }, + { url = "https://files.pythonhosted.org/packages/98/b7/8c59a66d15205024662f1d66968136f13893f96df1ddc5087e2e281fc95f/hf_xet-1.6.0-cp38-abi3-win_amd64.whl", hash = "sha256:fb4fadde1b2b70bf4c0c14a6dccbe7194b1c28947fefd5bbe3fed9d940676c3b", size = 4033128, upload-time = "2026-08-03T22:33:10.171Z" }, + { url = "https://files.pythonhosted.org/packages/73/63/ca511b6f802f28cf3489b280fe77475bcca8de85e81a6299d7916b5b5555/hf_xet-1.6.0-cp38-abi3-win_arm64.whl", hash = "sha256:3dc3e35441ba395006af5aaacc40ef2e603c51ef46c3530b9156185f00935ea3", size = 3859359, upload-time = "2026-08-03T22:33:11.725Z" }, ] [[package]] @@ -3381,22 +3370,23 @@ wheels = [ [[package]] name = "huggingface-hub" -version = "1.14.0" +version = "1.32.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "click" }, { name = "filelock" }, { name = "fsspec" }, { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, { name = "httpx" }, { name = "packaging" }, { name = "pyyaml" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, { name = "tqdm" }, - { name = "typer" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/39/40/43109e943fd718b0ccd0cd61eb4f1c347df22bf81f5874c6f22adf44bcff/huggingface_hub-1.14.0.tar.gz", hash = "sha256:d6d2c9cd6be1d02ae9ec6672d5587d10a427f377db688e82528f426a041622c2", size = 782365, upload-time = "2026-05-06T14:14:34.278Z" } +sdist = { url = "https://files.pythonhosted.org/packages/fe/0f/e83fdd856da8fca26bf78d71709ebd120432a0ce535e72b9597cab1eb5bf/huggingface_hub-1.32.0.tar.gz", hash = "sha256:ed70a45498abe86039df7c2f4e5f7575de524be908d3840e8f828d5525eafd6a", size = 1038662, upload-time = "2026-09-17T10:27:48.049Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/89/a5/33b49ba7bea7c41bb37f74ec0f8beea0831e052330196633fe2c77516ea6/huggingface_hub-1.14.0-py3-none-any.whl", hash = "sha256:efe075535c62e130b30e836b138e13785f6f043d1f0539e0a39aa411a99e90b8", size = 661479, upload-time = "2026-05-06T14:14:32.029Z" }, + { url = "https://files.pythonhosted.org/packages/1b/cf/d98dd561d6d0d7b7d7a64d1563f8aaaa7c235daee41c1c9bcc3da62420ed/huggingface_hub-1.32.0-py3-none-any.whl", hash = "sha256:b0c7c80561969d9cdacdd55fce67ba9584cca0b9d4ea80957a3a5c1445fac5c8", size = 842906, upload-time = "2026-09-17T10:27:46.102Z" }, ] [[package]] @@ -4516,14 +4506,18 @@ dependencies = [ { name = "boto3" }, { name = "click" }, { name = "fastuuid" }, + { name = "filelock" }, { name = "httpx", extra = ["http2"] }, + { name = "huggingface-hub" }, { name = "importlib-metadata" }, { name = "jinja2" }, { name = "jsonschema" }, { name = "openai" }, + { name = "packaging" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "python-dotenv" }, + { name = "pyyaml" }, { name = "tiktoken" }, { name = "tokenizers" }, ] @@ -4648,6 +4642,12 @@ utils = [ ] [package.dev-dependencies] +benchmarks = [ + { name = "a2a-sdk" }, + { name = "mcp" }, + { name = "pytest" }, + { name = "pytest-codspeed" }, +] ci = [ { name = "aiodynamo" }, { name = "anthropic" }, @@ -4687,6 +4687,8 @@ dev = [ { name = "keyring" }, { name = "langfuse" }, { name = "mypy" }, + { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "openapi-core" }, { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, @@ -4769,6 +4771,7 @@ requires-dist = [ { name = "fastapi", marker = "extra == 'proxy'", specifier = ">=0.136.3,<1.0" }, { name = "fastapi-sso", marker = "extra == 'proxy'", specifier = ">=0.19.0,<1.0" }, { name = "fastuuid", specifier = ">=0.14.0,<1.0" }, + { name = "filelock", specifier = ">=3.16.1,<4.0" }, { name = "google-cloud-aiplatform", marker = "extra == 'google'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-aiplatform", marker = "extra == 'proxy-runtime'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-iam", marker = "extra == 'extra-proxy'", specifier = ">=2.19.1,<3.0" }, @@ -4784,6 +4787,7 @@ requires-dist = [ { name = "httpx", extras = ["http2"], specifier = ">=0.28.0,<1.0" }, { name = "httpx2", marker = "extra == 'mcp'", specifier = ">=2.5.0,<3" }, { name = "httpx2", marker = "extra == 'proxy'", specifier = ">=2.5.0,<3" }, + { name = "huggingface-hub", specifier = ">=0.34.0,<2.0" }, { name = "importlib-metadata", specifier = ">=8.0.0,<9.0" }, { name = "inquirerpy", marker = "extra == 'cli'", specifier = ">=0.3.4,<1.0" }, { name = "inquirerpy", marker = "extra == 'proxy'", specifier = ">=0.3.4,<1.0" }, @@ -4807,6 +4811,7 @@ requires-dist = [ { name = "opentelemetry-instrumentation-fastapi", marker = "extra == 'proxy-runtime'", specifier = "==0.49b0" }, { name = "opentelemetry-sdk", marker = "extra == 'proxy-runtime'", specifier = "==1.28.0" }, { name = "orjson", marker = "extra == 'proxy'", specifier = ">=3.11.6,<4.0" }, + { name = "packaging", specifier = ">=24.0" }, { name = "polars", marker = "extra == 'proxy'", specifier = ">=1.38.1,<2.0" }, { name = "prisma", marker = "extra == 'extra-proxy'", specifier = ">=0.11.0,<1.0" }, { name = "prometheus-client", marker = "extra == 'proxy-runtime'", specifier = ">=0.20.0,<1.0" }, @@ -4824,6 +4829,7 @@ requires-dist = [ { name = "python-dotenv", specifier = ">=1.0.0,<2.0" }, { name = "python-multipart", marker = "extra == 'proxy'", specifier = ">=0.0.27,<1.0" }, { name = "python3-saml", marker = "extra == 'saml'", specifier = ">=1.16.0,<2.0" }, + { name = "pyyaml", specifier = ">=6.0.3,<7.0" }, { name = "pyyaml", marker = "extra == 'cli'", specifier = ">=6.0.3,<7.0" }, { name = "pyyaml", marker = "extra == 'proxy'", specifier = ">=6.0.3,<7.0" }, { name = "redisvl", marker = "extra == 'extra-proxy'", specifier = ">=0.4.1,<1.0" }, @@ -4850,6 +4856,12 @@ requires-dist = [ provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "saml", "semantic-router", "mlflow", "grpc", "stt-vertex-chirp", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] [package.metadata.requires-dev] +benchmarks = [ + { name = "a2a-sdk", specifier = "==1.1.0" }, + { name = "mcp", specifier = ">=2.2.0,<3" }, + { name = "pytest", specifier = "==9.0.3" }, + { name = "pytest-codspeed", specifier = "==4.3.0" }, +] ci = [ { name = "aiodynamo", specifier = "==24.7" }, { name = "anthropic", specifier = "==0.84.0" }, @@ -4889,6 +4901,7 @@ dev = [ { name = "keyring", specifier = "==25.7.0" }, { name = "langfuse", specifier = "==2.59.7" }, { name = "mypy", specifier = "==1.20.1" }, + { name = "numpy", specifier = ">=1.26.0,<3.0" }, { name = "openapi-core", specifier = "==0.22.0" }, { name = "opentelemetry-api", specifier = "==1.28.0" }, { name = "opentelemetry-exporter-otlp", specifier = "==1.28.0" }, @@ -9167,15 +9180,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" }, ] -[[package]] -name = "shellingham" -version = "1.5.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, -] - [[package]] name = "simple-websocket" version = "1.1.0" @@ -9916,21 +9920,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/19/97/56608b2249fe206a67cd573bc93cd9896e1efb9e98bce9c163bcdc704b88/truststore-0.10.4-py3-none-any.whl", hash = "sha256:adaeaecf1cbb5f4de3b1959b42d41f6fab57b2b1666adb59e89cb0b53361d981", size = 18660, upload-time = "2025-08-12T18:49:01.46Z" }, ] -[[package]] -name = "typer" -version = "0.25.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "annotated-doc" }, - { name = "click" }, - { name = "rich" }, - { name = "shellingham" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e4/51/9aed62104cea109b820bbd6c14245af756112017d309da813ef107d42e7e/typer-0.25.1.tar.gz", hash = "sha256:9616eb8853a09ffeabab1698952f33c6f29ffdbceb4eaeecf571880e8d7664cc", size = 122276, upload-time = "2026-04-30T19:32:16.964Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/f9/2b3ff4e56e5fa7debfaf9eb135d0da96f3e9a1d5b27222223c7296336e5f/typer-0.25.1-py3-none-any.whl", hash = "sha256:75caa44ed46a03fb2dab8808753ffacdbfea88495e74c85a28c5eefcf5f39c89", size = 58409, upload-time = "2026-04-30T19:32:18.271Z" }, -] - [[package]] name = "types-awscrt" version = "0.34.1"