diff --git a/.github/workflows/osv-scan.yml b/.github/workflows/osv-scan.yml index 9dd321f88db..0cd94fdd9e2 100644 --- a/.github/workflows/osv-scan.yml +++ b/.github/workflows/osv-scan.yml @@ -7,11 +7,6 @@ on: - litellm_internal_staging - litellm_oss_branch - "litellm_**" - paths: - - uv.lock - - ui/litellm-dashboard/package-lock.json - - osv-scanner.toml - - .github/workflows/osv-scan.yml schedule: - cron: "23 6 * * *" workflow_dispatch: diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index de7e1b68346..950d6ca31a6 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -14,7 +14,7 @@ permissions: jobs: lint: runs-on: ubuntu-latest - timeout-minutes: 10 + timeout-minutes: 15 steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 @@ -87,9 +87,11 @@ jobs: run: | uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')" - - name: Run basedpyright type checking + - name: Check basedpyright budget (delta vs base) + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} run: | - (uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py + (uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA" - name: Check for circular imports run: | diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml new file mode 100644 index 00000000000..3d0a159cdc7 --- /dev/null +++ b/.github/workflows/test-rust.yml @@ -0,0 +1,65 @@ +name: LiteLLM Rust + +on: + push: + paths: + - "litellm-rust/**" + - ".github/workflows/test-rust.yml" + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_branch + - "litellm_**" + paths: + - "litellm-rust/**" + - ".github/workflows/test-rust.yml" + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + rust-checks: + name: rustfmt, clippy, test + runs-on: ubuntu-latest + timeout-minutes: 10 + defaults: + run: + working-directory: litellm-rust + env: + CARGO_TERM_COLOR: always + + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Rust + run: | + rustup toolchain install stable --profile minimal --component clippy,rustfmt + rustup default stable + + - name: Cache Cargo registry and target + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + litellm-rust/target + key: ${{ runner.os }}-cargo-${{ hashFiles('litellm-rust/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo- + + - name: Check Rust formatting + run: cargo fmt --check + + - name: Run Clippy + run: cargo clippy --workspace --all-targets --locked -- -D warnings + + - name: Run Rust tests + run: cargo test --workspace --locked diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index a7363ac3b43..2226d519331 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -32,7 +32,9 @@ jobs: tests/test_litellm/repositories tests/test_litellm/images tests/test_litellm/interactions + tests/test_litellm/ocr tests/test_litellm/passthrough + tests/test_litellm/sandbox tests/test_litellm/vector_stores tests/test_litellm/test_*.py workers: 2 diff --git a/Makefile b/Makefile index 27150aec938..076eac0f4a7 100644 --- a/Makefile +++ b/Makefile @@ -125,7 +125,8 @@ lint-ruff-FULL-dev: install-dev else echo "No changed .py files to check."; fi lint-basedpyright: install-dev - ($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py + git fetch origin litellm_internal_staging + ($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging lint-basedpyright-budget-update: install-dev ($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --update diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index d1a576aeb33..2f65f99c292 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -120,6 +120,9 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/robots.txt", # Health (k8s probes) "/health", + # Plugin system + "/api/plugins", + "/plugin-proxy/", ) BACKEND_EXACT_PATHS: frozenset[str] = frozenset( diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 7ba7656e407..f5b0a9aaf81 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -121,7 +121,7 @@ }, "reportReturnType": { "baseline": 126, - "slack": 13 + "slack": 100 }, "reportTypedDictNotRequiredAccess": { "baseline": 20, @@ -157,7 +157,7 @@ }, "reportUnnecessaryComparison": { "baseline": 683, - "slack": 10 + "slack": 100 }, "reportUnnecessaryContains": { "baseline": 4, diff --git a/docs/plugin_architecture.md b/docs/plugin_architecture.md new file mode 100644 index 00000000000..8801761531d --- /dev/null +++ b/docs/plugin_architecture.md @@ -0,0 +1,141 @@ +# LiteLLM Plugin Architecture + +Plugins let external services appear as selectable modes in the litellm UI sidebar alongside the AI Gateway. + +--- + +## Quick start + +### 1. Configure the plugin + +Add a `plugins` block to your litellm `config.yaml`: + +```yaml +general_settings: + master_key: sk-... + plugins: + - name: my-plugin # unique identifier (no spaces) + display_name: My Plugin # shown in the UI dropdown + url: "https://my-plugin.example.com" + plugin_key: "sk-..." # plugin's own auth credential +``` + +`plugin_key` is injected as `Authorization: Bearer ` on every +request proxied through `/plugin-proxy/my-plugin/*`. The caller's litellm +credential is stripped before forwarding so the plugin never receives a live +litellm API key. + +### 2. Implement two endpoints on your service + +| Endpoint | Method | Purpose | +|---|---|---| +| `GET /api/plugin-manifest` | public | Returns plugin metadata for the UI | +| `POST /api/plugin-auth` | public | Decrypts the identity claim for seamless sign-in | + +#### `GET /api/plugin-manifest` + +```json +{ + "name": "my-plugin", + "display_name": "My Plugin", + "version": "1.0.0", + "nav_items": [ + { "key": "home", "label": "Home", "icon": "HomeOutlined", "path": "/" }, + { "key": "reports", "label": "Reports", "icon": "BarChartOutlined", "path": "/reports" } + ], + "capabilities": ["reports", "data"] +} +``` + +#### `POST /api/plugin-auth` + +Receives `{ "session_claim": "" }`. + +The proxy never shares `LITELLM_SALT_KEY` with your plugin. Each plugin is +provisioned with its own dedicated key, derived as +`HMAC-SHA256(LITELLM_SALT_KEY, plugin_name)`. Compute it once on the proxy +host and hand the result to your plugin as a secret (e.g. `PLUGIN_AUTH_KEY`): + +```bash +python -c 'import base64,hmac,hashlib,os; \ +print(base64.urlsafe_b64encode(hmac.new(os.environ["LITELLM_SALT_KEY"].encode(), b"my-plugin", hashlib.sha256).digest()).decode())' +``` + +A compromised plugin holding only this scoped key cannot recover +`LITELLM_SALT_KEY` or decrypt any other litellm secret. + +Decrypt and validate the claim with that key: + +```python +import json, os, time +from cryptography.fernet import Fernet + +_CLAIM_TTL_SECONDS = 30 + +def plugin_auth(session_claim: str) -> dict: + cipher = Fernet(os.environ["PLUGIN_AUTH_KEY"].encode()) + claim = json.loads(cipher.decrypt(session_claim.encode(), ttl=_CLAIM_TTL_SECONDS)) + if claim.get("plugin") != "my-plugin": + raise ValueError("claim audience mismatch") + if int(claim.get("exp", 0)) < int(time.time()): + raise ValueError("claim expired") + return claim +``` + +The claim is `{ "plugin", "user_id", "user_role", "exp" }`; it carries no +litellm bearer token. Establish the plugin's own session from `user_id` / +`user_role` and authenticate API calls back to litellm through the +`/plugin-proxy/my-plugin/*` reverse proxy, which injects `plugin_key` for you. + +--- + +## How iframe auth works + +``` +litellm UI + ├─ GET /api/plugins/auth-token -> { session_claim } + └─ postMessage({ type:"litellm-auth", session_claim }, pluginOrigin) + │ + ▼ +Plugin iframe browser + └─ POST /api/plugin-auth { session_claim } + │ + ▼ +Plugin server + ├─ decrypt(session_claim, PLUGIN_AUTH_KEY) -> { user_id, user_role, exp } + └─ establish plugin session -> stored in sessionStorage +``` + +No litellm bearer token ever leaves the proxy; the claim only conveys the +caller's identity and expires after 30 seconds. A postMessage intercept +yields ciphertext that is useless without the plugin's scoped key. + +--- + +## Proxy routes + +- `GET /api/plugins` — list registered plugins (`name`, `display_name`, `url`). `plugin_key` is **never** returned; it stays server-side. Requires an authenticated caller. +- `GET /api/plugins/auth-token?plugin_name=` — short-lived encrypted identity claim for the named plugin. Requires `LITELLM_SALT_KEY` to be set (503 otherwise) and the plugin to be registered (404 otherwise). +- `ANY /plugin-proxy/{name}/{path}` — authenticated reverse proxy to the plugin backend. Restricted to `proxy_admin`. + +--- + +## Reverse proxy behaviour + +When an admin (or server-to-server caller) hits `/plugin-proxy//`, the proxy authenticates the caller locally, then rewrites the request before forwarding it to the plugin's `url`: + +- **Every litellm credential header is stripped** — `Authorization`, `x-api-key`, `API-Key`, `x-goog-api-key`, `Ocp-Apim-Subscription-Key`, `x-litellm-api-key`, any configured `litellm_key_header_name`, plus `Cookie`. The plugin can never be handed the caller's live litellm key. +- **`plugin_key` is injected** as `Authorization: Bearer ` — the only credential the plugin receives. +- **Caller identity is forwarded** as `x-litellm-user-id` and `x-litellm-user-role` so the plugin can run its own authorization. These are informational, not credentials. +- **Responses are sandboxed** — `Content-Security-Policy: sandbox` and `X-Content-Type-Options: nosniff` are set so plugin-controlled bytes served from the litellm origin cannot execute against the dashboard. + +--- + +## Security checklist + +- [ ] `LITELLM_SALT_KEY` is set on the proxy and never shared with the plugin +- [ ] The plugin holds only its derived `HMAC(LITELLM_SALT_KEY, plugin_name)` key, provisioned as a dedicated secret +- [ ] `plugin_key` is a dedicated credential scoped to the plugin (not your litellm master key) +- [ ] Plugin's `POST /api/plugin-auth` enforces the claim's `plugin` audience and `exp` (30s TTL) +- [ ] Plugin treats `x-litellm-user-id` / `x-litellm-user-role` as identity hints, not as proof of authentication +- [ ] Plugin service URL uses HTTPS in production diff --git a/litellm-rust/.gitignore b/litellm-rust/.gitignore new file mode 100644 index 00000000000..b83d22266ac --- /dev/null +++ b/litellm-rust/.gitignore @@ -0,0 +1 @@ +/target/ diff --git a/litellm-rust/ADDING_A_PROVIDER.md b/litellm-rust/ADDING_A_PROVIDER.md new file mode 100644 index 00000000000..2fa81798605 --- /dev/null +++ b/litellm-rust/ADDING_A_PROVIDER.md @@ -0,0 +1,9 @@ +# Adding a provider / route to litellm-rust + +Three layers, same for every route (see `ocr` and `realtime` as references): + +1. **Transform contract (pure)** — `crates/core/src//transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) + types in `types.rs`. No network, env, or auth. +2. **Provider config (pure)** — `crates/providers/src///transformation.rs`: implement that trait as a `const __CONFIG`, mirroring the Python provider tree. Add parity unit tests. +3. **HTTP / transport (the host)** — `crates/providers/src/.rs` (e.g. `ocr.rs`, `realtime.rs`): the callable fn (`run_ocr`, `realtime`). It resolves the key, builds the auth header, builds URL + transforms via the config, then does the network call. This is the only layer allowed to do I/O. + +**Calling:** the host invokes the route fn — the Python bridge calls `run_ocr`; the `ai-gateway` server calls `realtime`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`. diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md new file mode 100644 index 00000000000..1d2987e0a1a --- /dev/null +++ b/litellm-rust/CLAUDE.md @@ -0,0 +1,88 @@ +# CLAUDE.md + +This file defines the rules for Rust work in LiteLLM. + +## Core Boundary + +The `core` and `providers` crates describe work; hosts execute work. + +Route-level Rust structure mirrors LiteLLM's Python responsibilities: +- `core/src//` owns the route contract, shared types, and provider + template traits. For OCR, this means `core/src/ocr`. +- `providers/src///transformation.rs` owns the + provider-specific transform. For Mistral OCR, this means + `providers/src/mistral/ocr/transformation.rs`. +- Future network execution belongs in a host/transport layer such as + `llm_http_handler`, not inside `core` or `providers`. + +Allowed in `core` and `providers`: +- Pure request transforms +- Pure response transforms +- Pure stream chunk normalization +- Shared data types and validation errors +- Deterministic token/cost helper logic + +Not allowed in `core` or `providers`: +- Network calls +- Environment variable or secret reads +- Filesystem access +- Database or cache access +- Provider SDK signing or auth flows +- Logging callbacks, spend writes, or custom callbacks +- Global mutable runtime state + +Python owns rollout state and fallback while Rust is being introduced. Rust +paths must be off by default until parity tests prove equivalence with Python. + +## Production Bar + +Rust code in this workspace is held to a strict parity and robustness bar from +the first PR: + +- Correctness parity is proven with tests. Do not rely on README claims or + manual inspection for a port that mirrors Python behavior. +- Every provider transform must have unit tests for supported-parameter + filtering, request body shape, response normalization, missing/null fields, + and bad-input errors. +- When Rust is exposed through Python, add Python tests that prove disabled, + enabled, and unavailable-bridge fallback behavior. +- Avoid panics on user/provider input. Return typed errors and let the host map + them to Python exceptions or HTTP responses. +- OCR handles documents that often contain personal data. Do not log document + contents, base64 payloads, provider response bodies, or secrets. +- Error messages must be useful but data-minimized. Truncate or sanitize any + upstream body before it crosses a host boundary. +- Treat empty or whitespace-only credentials, URLs, and config values as absent + at the host/config resolution layer. +- Preserve Python output shape intentionally. If a field is always serialized as + `null` for Python parity, leave a short comment explaining that parity choice. + +## Host I/O Rules + +These rules apply when adding future crates or modules that execute network I/O, +such as `ai-gateway`, router hosts, or standalone servers: + +- Set connect and full-request timeouts. No unbounded waits. +- Reuse HTTP clients; do not construct clients per request. +- Prefer rustls TLS for portable Python wheels and Linux images unless there is + a documented reason not to. +- Add request IDs and structured tracing at the host layer, without logging OCR + document contents or secrets. +- Do not echo raw upstream response bodies to callers. Sanitize and bound them. +- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is + impossible by construction and documented. + +## Checks + +Run these before pushing Rust changes. The same checks run in GitHub Actions +for changes under `litellm-rust/`. + +```bash +cd litellm-rust +cargo fmt --check +cargo clippy --workspace --all-targets -- -D warnings +cargo test --workspace +``` + +When a Rust path is exposed through Python, add Python parity tests that compare +the existing Python output with the Rust-backed output. diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock new file mode 100644 index 00000000000..2ec86fc00e7 --- /dev/null +++ b/litellm-rust/Cargo.lock @@ -0,0 +1,1754 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" + +[[package]] +name = "cc" +version = "1.2.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-io", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi", + "wasip2", + "wasm-bindgen", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "hyper" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "tokio", + "tokio-rustls", + "tower-service", + "webpki-roots", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03d04c30968dffe80775bd4d7fb676131cd04a1fb46d2686dbffbaec2d9dfd31" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "litellm-core" +version = "0.1.0" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.18", +] + +[[package]] +name = "litellm-providers" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-core", + "reqwest", + "serde_json", + "tokio", + "tokio-tungstenite", +] + +[[package]] +name = "litellm-python-bridge" +version = "0.1.0" +dependencies = [ + "litellm-core", + "litellm-providers", + "pyo3", + "serde_json", +] + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "mio" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.23.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls", + "socket2", + "thiserror 2.0.18", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +dependencies = [ + "bytes", + "getrandom 0.3.4", + "lru-slab", + "rand 0.9.4", + "ring", + "rustc-hash", + "rustls", + "rustls-pki-types", + "slab", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.60.2", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rand" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "webpki-roots", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + +[[package]] +name = "rustls" +version = "0.23.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" +dependencies = [ + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pki-types" +version = "1.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" +dependencies = [ + "futures-util", + "log", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tungstenite", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18e5b8366ee7a95b16d32197d0b2604b43a0be89dc5fac9f8e96ccafbaedda8a" +dependencies = [ + "byteorder", + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.8.6", + "rustls", + "rustls-pki-types", + "sha1", + "thiserror 1.0.69", + "utf-8", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.125" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ddb3f79143bced6de84270411622a2699cee572fc0875aeaf1e7867cf9fca1a" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.75" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "503b14d284f2c8dac03b819967e155ea753f573586193b2b2c95990cb5d69280" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.125" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e21a184b13fb19e157296e2c46056aec9092264fab83e4ba59e68c61b323c3d" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.125" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fecefd9c35bd935a20fc3fc344b5f29138961e4f47fb03297d88f2587afb5ebd" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.125" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23939e44bb9a5d7576fa2b563dc2e136628f1224e88a8deed09e04858b77871f" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-sys" +version = "0.3.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6430a72df5eb332242960fe84b3002a241163998241eb596d4f739b9757061d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webpki-roots" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm 0.52.6", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml new file mode 100644 index 00000000000..0b1528672a7 --- /dev/null +++ b/litellm-rust/Cargo.toml @@ -0,0 +1,24 @@ +[workspace] +members = [ + "crates/core", + "crates/providers", + "crates/python-bridge", +] +resolver = "2" + +[workspace.package] +edition = "2021" +license = "MIT" +repository = "https://github.com/BerriAI/litellm" + +[workspace.dependencies] +litellm-core = { path = "crates/core" } +litellm-providers = { path = "crates/providers" } +pyo3 = "0.23.5" +reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +thiserror = "2.0" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } +tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } +futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } diff --git a/litellm-rust/README.md b/litellm-rust/README.md new file mode 100644 index 00000000000..15ad1855420 --- /dev/null +++ b/litellm-rust/README.md @@ -0,0 +1,34 @@ +# LiteLLM Rust + +This workspace contains the staged Rust implementation for LiteLLM. + +Rust starts as a pure transform core used by the existing Python host. Python +continues to own auth, configuration, network I/O, retries, routing, logging, +callbacks, spend tracking, and customer plugins until each Rust path has parity +coverage and production evidence. + +## Layout + +```text +crates/ + core/ Route contracts, shared pure types, errors, and templates. + src/ocr/ + providers/ Provider-specific pure transforms. + src/mistral/ocr/transformation.rs + python-bridge/ PyO3 bridge for Python LiteLLM. +``` + +The folder shape should follow the Python provider tree: +`providers/src///transformation.rs`. The bridge should expose +one function per top-level route, starting with `ocr(payload)`. + +## Checks + +Run these before pushing Rust changes. GitHub Actions runs the same checks for +changes under `litellm-rust/`. + +```bash +cargo fmt --check +cargo clippy --workspace --all-targets -- -D warnings +cargo test --workspace +``` diff --git a/litellm-rust/crates/core/CLAUDE.md b/litellm-rust/crates/core/CLAUDE.md new file mode 100644 index 00000000000..20873878967 --- /dev/null +++ b/litellm-rust/crates/core/CLAUDE.md @@ -0,0 +1,47 @@ +# CLAUDE.md + +Rules for `litellm-rust/crates/core`. + +## Responsibility + +`core` owns shared data types, typed errors, and deterministic helper contracts. +It must stay pure and host-independent. + +Allowed: +- Shared request/response structs. +- Typed errors with stable, non-sensitive messages. +- Deterministic validation helpers. +- Serialization helpers that intentionally mirror Python output shape. +- Route templates that match Python base config responsibilities, such as + `ocr::transformation::OcrProviderConfig`. + +Not allowed: +- Network, filesystem, database, cache, or environment access. +- Secret reads or auth/header construction. +- Logging callbacks, tracing spans, spend writes, or customer callbacks. +- Provider-specific branching that belongs in `providers`. +- Panics for user/provider-controlled input. + +## Typed Contracts (core rule) + +Trait and function boundaries MUST be strongly typed. No stringly-typed JSON +(`&str` / `String` / `Vec` / bare `serde_json::Value`) as a transform +input or output. Parse wire bytes into typed structs/enums at the host edge; +`core` and `providers` operate only on those types (e.g. `RealtimeEvent`, +`RealtimeTransformResult`, `OcrRequestData`). A `type`-style discriminator is a +typed field on a struct, not a raw string threaded through the API. + +## Structure + +Use route names directly under `src/`: `ocr`, future `messages`, +`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not +invent broad names like `engine` for route contracts. + +## Parity Rules + +- Every shared type used by a provider transform needs unit tests for + serialization shape. +- If Python parity requires always emitting a `null` field instead of omitting + it, document that in code and pin it with a test. +- Error enums should preserve enough detail for Python/HTTP hosts to map errors + consistently without exposing document contents or upstream bodies. diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml new file mode 100644 index 00000000000..e54002fe5e8 --- /dev/null +++ b/litellm-rust/crates/core/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "litellm-core" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs new file mode 100644 index 00000000000..645e261f76d --- /dev/null +++ b/litellm-rust/crates/core/src/error.rs @@ -0,0 +1,33 @@ +use thiserror::Error; + +pub type CoreResult = Result; + +#[derive(Debug, Error, PartialEq, Eq)] +pub enum CoreError { + #[error("expected {expected}, got {actual}")] + InvalidType { + expected: &'static str, + actual: &'static str, + }, + #[error("missing required field: {0}")] + MissingField(&'static str), + #[error("invalid response: {0}")] + InvalidResponse(String), + #[error("{0}")] + Auth(String), + #[error("OCR request failed with status {status}: {body}")] + Http { status: u16, body: String }, + #[error("OCR network error: {0}")] + Network(String), +} + +pub fn json_type_name(value: &serde_json::Value) -> &'static str { + match value { + serde_json::Value::Null => "null", + serde_json::Value::Bool(_) => "bool", + serde_json::Value::Number(_) => "number", + serde_json::Value::String(_) => "string", + serde_json::Value::Array(_) => "array", + serde_json::Value::Object(_) => "object", + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs new file mode 100644 index 00000000000..5d4d5bfd142 --- /dev/null +++ b/litellm-rust/crates/core/src/lib.rs @@ -0,0 +1,5 @@ +pub mod error; +pub mod ocr; +pub mod realtime; + +pub use error::{CoreError, CoreResult}; diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs new file mode 100644 index 00000000000..ec2fbb969a6 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -0,0 +1,2 @@ +pub mod transformation; +pub mod types; diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs new file mode 100644 index 00000000000..7353d9d22c4 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -0,0 +1,32 @@ +use serde_json::{Map, Value}; + +use crate::CoreResult; + +use super::types::{OcrRequestData, OcrResponseData}; + +pub trait OcrProviderConfig { + fn supported_ocr_params(&self) -> &'static [&'static str]; + + fn map_ocr_params(&self, non_default_params: &Map) -> Map { + let mut mapped_params = Map::new(); + for (param, value) in non_default_params { + if self.supported_ocr_params().contains(¶m.as_str()) { + mapped_params.insert(param.clone(), value.clone()); + } + } + mapped_params + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult; + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult; +} diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs new file mode 100644 index 00000000000..1a72b8f1d66 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -0,0 +1,29 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct OcrRequestData { + pub data: Value, + pub files: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct OcrResponseData { + pub pages: Vec, + pub model: String, + pub document_annotation: Option, + pub usage_info: Option, + pub object: String, +} + +impl OcrResponseData { + pub fn into_json(self) -> Value { + serde_json::json!({ + "pages": self.pages, + "model": self.model, + "document_annotation": self.document_annotation, + "usage_info": self.usage_info, + "object": self.object, + }) + } +} diff --git a/litellm-rust/crates/core/src/realtime/mod.rs b/litellm-rust/crates/core/src/realtime/mod.rs new file mode 100644 index 00000000000..ec2fbb969a6 --- /dev/null +++ b/litellm-rust/crates/core/src/realtime/mod.rs @@ -0,0 +1,2 @@ +pub mod transformation; +pub mod types; diff --git a/litellm-rust/crates/core/src/realtime/transformation.rs b/litellm-rust/crates/core/src/realtime/transformation.rs new file mode 100644 index 00000000000..a4baa27a6c2 --- /dev/null +++ b/litellm-rust/crates/core/src/realtime/transformation.rs @@ -0,0 +1,22 @@ +use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; +use crate::CoreResult; + +pub trait RealtimeProviderConfig { + /// Build the upstream WebSocket URL (e.g. `wss://api.openai.com/v1/realtime?model=…`). + /// Pure string construction only — no network, no env. + fn complete_url(&self, api_base: Option<&str>, model: &str) -> String; + + /// Transform a client → backend event before it is forwarded upstream. + fn transform_realtime_request( + &self, + event: &RealtimeEvent, + model: &str, + ) -> CoreResult; + + /// Transform a backend → client event before it is forwarded downstream. + fn transform_realtime_response( + &self, + event: &RealtimeEvent, + model: &str, + ) -> CoreResult; +} diff --git a/litellm-rust/crates/core/src/realtime/types.rs b/litellm-rust/crates/core/src/realtime/types.rs new file mode 100644 index 00000000000..3b59224b6e9 --- /dev/null +++ b/litellm-rust/crates/core/src/realtime/types.rs @@ -0,0 +1,60 @@ +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +/// A single realtime event exchanged over the WebSocket. +/// +/// The `type` discriminator is a typed field; the remaining fields are +/// preserved losslessly in `data` so a transform can pass an event through, or +/// inspect/modify specific fields, without enumerating every event variant. +/// Wire (de)serialization happens at the host edge — `core`/`providers` operate +/// only on this typed form. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct RealtimeEvent { + #[serde(rename = "type")] + pub event_type: String, + #[serde(flatten)] + pub data: Map, +} + +/// One or more typed events produced by a realtime transform. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct RealtimeTransformResult { + pub events: Vec, +} + +impl RealtimeTransformResult { + /// Forward a single event unchanged (the OpenAI baseline). + pub fn passthrough(event: RealtimeEvent) -> Self { + Self { + events: vec![event], + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn event(raw: &str) -> RealtimeEvent { + serde_json::from_str(raw).expect("valid event json") + } + + #[test] + fn realtime_event_round_trips_type_and_extra_fields() { + let raw = r#"{"type":"response.output_text.delta","delta":"hi","response_id":"r1"}"#; + let parsed = event(raw); + assert_eq!(parsed.event_type, "response.output_text.delta"); + assert_eq!(parsed.data.get("delta"), Some(&Value::String("hi".into()))); + // Re-serializing yields a semantically-equal event (key order may differ). + let reparsed: RealtimeEvent = + serde_json::from_str(&serde_json::to_string(&parsed).unwrap()).unwrap(); + assert_eq!(parsed, reparsed); + } + + #[test] + fn passthrough_produces_single_element_vec() { + let parsed = event(r#"{"type":"session.update"}"#); + let result = RealtimeTransformResult::passthrough(parsed.clone()); + assert_eq!(result.events, vec![parsed]); + } +} diff --git a/litellm-rust/crates/providers/CLAUDE.md b/litellm-rust/crates/providers/CLAUDE.md new file mode 100644 index 00000000000..0f7fdcda2aa --- /dev/null +++ b/litellm-rust/crates/providers/CLAUDE.md @@ -0,0 +1,53 @@ +# CLAUDE.md + +Rules for `litellm-rust/crates/providers`. + +## Responsibility + +`providers` owns provider-specific pure transforms. It mirrors the existing +Python provider modules closely enough that parity review is mechanical. + +Provider files should map to the Python provider tree: + +```text +providers/src///transformation.rs +``` + +For example, Mistral OCR lives at +`providers/src/mistral/ocr/transformation.rs`, matching +`litellm/llms/mistral/ocr/transformation.py`. + +Allowed: +- Provider request transforms. +- Provider response normalization. +- Supported-parameter filtering. +- Provider-specific validation that does not require I/O or secrets. + +Not allowed: +- HTTP clients or provider SDK calls. +- Environment variable reads. +- API key resolution or auth header construction. +- Logging, callbacks, spend tracking, retries, routing, cooldowns, or fallbacks. +- Panics on bad user/provider input. + +## Required Tests + +Every provider transform must include focused unit tests for: +- Supported params matching the Python provider config. +- Unknown params being dropped or transformed the same way as Python. +- Request body shape matching Python output. +- Response normalization with complete, missing, null, and extra fields. +- Bad input returning typed errors. + +For OCR specifically, assume documents can contain personal data. Tests should +prove transforms do not copy document contents into error messages. + +## Implementation Rules + +- Prefer static supported-parameter lists over allocating strings on every call. +- Keep transforms deterministic and allocation-conscious, but choose clarity over + premature micro-optimization for tiny parameter lists. +- Use typed errors from `core`; avoid stringly-typed error plumbing. +- Add comments only when they explain Python-parity decisions or provider quirks. +- Put route-level provider dispatch in a route file such as `providers/src/ocr.rs`. + Do not move provider-specific transform logic into the Python bridge. diff --git a/litellm-rust/crates/providers/Cargo.toml b/litellm-rust/crates/providers/Cargo.toml new file mode 100644 index 00000000000..cb8a91aa5c2 --- /dev/null +++ b/litellm-rust/crates/providers/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "litellm-providers" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-core.workspace = true +reqwest.workspace = true +serde_json.workspace = true +tokio.workspace = true +tokio-tungstenite.workspace = true +futures-util.workspace = true + +[dev-dependencies] +serde_json.workspace = true diff --git a/litellm-rust/crates/providers/src/lib.rs b/litellm-rust/crates/providers/src/lib.rs new file mode 100644 index 00000000000..34b69d88cbf --- /dev/null +++ b/litellm-rust/crates/providers/src/lib.rs @@ -0,0 +1,4 @@ +pub mod mistral; +pub mod ocr; +pub mod openai; +pub mod realtime; diff --git a/litellm-rust/crates/providers/src/mistral/mod.rs b/litellm-rust/crates/providers/src/mistral/mod.rs new file mode 100644 index 00000000000..3621ff6a2fd --- /dev/null +++ b/litellm-rust/crates/providers/src/mistral/mod.rs @@ -0,0 +1 @@ +pub mod ocr; diff --git a/litellm-rust/crates/providers/src/mistral/ocr/mod.rs b/litellm-rust/crates/providers/src/mistral/ocr/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/providers/src/mistral/ocr/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/providers/src/mistral/ocr/transformation.rs b/litellm-rust/crates/providers/src/mistral/ocr/transformation.rs new file mode 100644 index 00000000000..fd691177783 --- /dev/null +++ b/litellm-rust/crates/providers/src/mistral/ocr/transformation.rs @@ -0,0 +1,292 @@ +use litellm_core::error::{json_type_name, CoreError, CoreResult}; +use litellm_core::ocr::transformation::OcrProviderConfig; +use litellm_core::ocr::types::{OcrRequestData, OcrResponseData}; +use serde_json::{Map, Value}; + +const SUPPORTED_OCR_PARAMS: &[&str] = &[ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + "document_annotation_prompt", + "extract_header", + "extract_footer", + "table_format", + "confidence_scores_granularity", + "id", +]; + +/// Default Mistral API base, used when the caller does not override `api_base`. +pub const MISTRAL_DEFAULT_API_BASE: &str = "https://api.mistral.ai/v1"; + +/// Environment variable holding the Mistral API key. +pub const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY"; + +/// Error message raised when no Mistral API key can be resolved. +pub const MISSING_KEY_MESSAGE: &str = "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params"; + +/// Build the complete OCR endpoint URL, de-duplicating a trailing `/v1`. +/// +/// Blank/whitespace `api_base` is treated as absent (guard at resolution time). +pub fn complete_url(api_base: Option<&str>) -> String { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(MISTRAL_DEFAULT_API_BASE) + .trim_end_matches('/'); + + if base.ends_with("/v1") { + format!("{base}/ocr") + } else { + format!("{base}/v1/ocr") + } +} + +/// Resolve the Mistral API key from the explicit param or the environment. +/// +/// Blank/whitespace values are treated as absent. Returns `CoreError::Auth` +/// when no usable key is available. +/// +/// Note: the env fallback only reads the process environment. Secret-manager +/// backends (AWS/Azure/GCP/Vault) are resolved on the Python side and passed in +/// via `api_key`; this fallback is a last resort for direct/standalone use. +pub fn resolve_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + api_key + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) +} + +pub struct MistralOcrConfig; + +pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig; + +impl OcrProviderConfig for MistralOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + SUPPORTED_OCR_PARAMS + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult { + if !document.is_object() { + return Err(CoreError::InvalidType { + expected: "object", + actual: json_type_name(&document), + }); + } + + let mut data = Map::new(); + data.insert("model".to_string(), Value::String(model.to_string())); + data.insert("document".to_string(), document); + for (param, value) in optional_params { + data.insert(param, value); + } + + Ok(OcrRequestData { + data: Value::Object(data), + files: None, + }) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + let response_object = response_json + .as_object() + .ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + })?; + + let pages = response_object + .get("pages") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let model = response_object + .get("model") + .and_then(Value::as_str) + .unwrap_or(model) + .to_string(); + let document_annotation = response_object.get("document_annotation").cloned(); + let usage_info = response_object.get("usage_info").cloned(); + + Ok(OcrResponseData { + pages, + model, + document_annotation, + usage_info, + object: "ocr".to_string(), + }) + } +} + +pub fn supported_ocr_params() -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.supported_ocr_params() +} + +pub fn map_ocr_params(non_default_params: &Map) -> Map { + MISTRAL_OCR_CONFIG.map_ocr_params(non_default_params) +} + +pub fn transform_ocr_request( + model: &str, + document: Value, + optional_params: Map, +) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) +} + +pub fn transform_ocr_response(model: &str, response_json: Value) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn supported_params_match_python_mistral_ocr_config() { + assert_eq!( + supported_ocr_params(), + &[ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + "document_annotation_prompt", + "extract_header", + "extract_footer", + "table_format", + "confidence_scores_granularity", + "id", + ] + ); + } + + #[test] + fn map_ocr_params_drops_unknown_params() { + let params = json!({ + "extract_header": true, + "unsupported_param": "value", + "pages": [0, 1] + }); + let mapped = map_ocr_params(params.as_object().unwrap()); + + assert_eq!(mapped.get("extract_header"), Some(&json!(true))); + assert_eq!(mapped.get("pages"), Some(&json!([0, 1]))); + assert!(!mapped.contains_key("unsupported_param")); + } + + #[test] + fn transform_ocr_request_builds_mistral_body() { + let document = json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }); + let optional_params = json!({ + "include_image_base64": true, + "table_format": "html" + }) + .as_object() + .unwrap() + .clone(); + + let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params) + .expect("request should transform"); + + assert_eq!( + result.data, + json!({ + "model": "mistral-ocr-latest", + "document": document, + "include_image_base64": true, + "table_format": "html" + }) + ); + assert_eq!(result.files, None); + } + + #[test] + fn transform_ocr_request_rejects_non_object_document() { + let err = transform_ocr_request("mistral-ocr-latest", json!("bad"), Map::new()) + .expect_err("string document should be rejected"); + + assert_eq!( + err, + CoreError::InvalidType { + expected: "object", + actual: "string", + } + ); + } + + #[test] + fn transform_ocr_response_normalizes_mistral_json() { + let response = json!({ + "pages": [{"index": 0, "markdown": "hello"}], + "model": "mistral-ocr-2505-completion", + "document_annotation": null, + "usage_info": {"pages_processed": 1} + }); + + let result = transform_ocr_response("mistral-ocr-latest", response) + .expect("response should transform"); + + assert_eq!(result.pages, vec![json!({"index": 0, "markdown": "hello"})]); + assert_eq!(result.model, "mistral-ocr-2505-completion"); + assert_eq!(result.document_annotation, Some(Value::Null)); + assert_eq!(result.usage_info, Some(json!({"pages_processed": 1}))); + assert_eq!(result.object, "ocr"); + } + + #[test] + fn complete_url_defaults_and_dedupes_v1() { + assert_eq!(complete_url(None), "https://api.mistral.ai/v1/ocr"); + assert_eq!(complete_url(Some(" ")), "https://api.mistral.ai/v1/ocr"); + assert_eq!( + complete_url(Some("https://proxy.internal")), + "https://proxy.internal/v1/ocr" + ); + assert_eq!( + complete_url(Some("https://proxy.internal/v1/")), + "https://proxy.internal/v1/ocr" + ); + } + + #[test] + fn resolve_api_key_prefers_param_then_env() { + let no_env = |_: &str| None; + assert_eq!( + resolve_api_key(Some("sk-param"), &no_env).unwrap(), + "sk-param" + ); + + let with_env = |key: &str| (key == MISTRAL_API_KEY_ENV).then(|| "sk-env".to_string()); + assert_eq!(resolve_api_key(None, &with_env).unwrap(), "sk-env"); + // Blank param falls through to the environment. + assert_eq!(resolve_api_key(Some(" "), &with_env).unwrap(), "sk-env"); + } + + #[test] + fn resolve_api_key_errors_when_absent() { + let err = resolve_api_key(None, &|_| None).expect_err("missing key should error"); + assert_eq!(err, CoreError::Auth(MISSING_KEY_MESSAGE.to_string())); + } +} diff --git a/litellm-rust/crates/providers/src/ocr.rs b/litellm-rust/crates/providers/src/ocr.rs new file mode 100644 index 00000000000..dcd56a5f0b4 --- /dev/null +++ b/litellm-rust/crates/providers/src/ocr.rs @@ -0,0 +1,127 @@ +//! End-to-end OCR orchestration. +//! +//! Owns the whole Mistral OCR call so the Python side stays a thin bridge: +//! resolve the API key, build the URL + body via the pure transforms, POST it, +//! and normalize the response. The HTTP client is built once and reused. + +use std::sync::OnceLock; +use std::time::Duration; + +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrProviderConfig; +use litellm_core::CoreResult; +use serde_json::{Map, Value}; + +use crate::mistral::ocr::transformation as mistral; +use crate::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; + +/// OCR over large documents can take a while; bound it generously rather than +/// hanging forever on an unresponsive upstream. The client-level limit is the +/// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``. +const OCR_TIMEOUT_SECS: u64 = 600; + +/// Maximum upstream body characters retained in error messages. OCR responses +/// can echo document contents and prompts; keep enough for debugging without +/// forwarding sensitive payloads across the host boundary. +const ERROR_BODY_MAX_CHARS: usize = 256; + +/// Process-wide blocking HTTP client (connection pool + TLS reused across calls). +fn http_client() -> &'static reqwest::blocking::Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + reqwest::blocking::Client::builder() + .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) + .build() + .expect("failed to build reqwest client") + }) +} + +fn truncate_error_body(body: &str) -> String { + if body.chars().count() <= ERROR_BODY_MAX_CHARS { + return body.to_string(); + } + let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect(); + format!("{truncated}... (truncated)") +} + +/// Perform a Mistral OCR call end to end and return the normalized response as +/// JSON (the shape the Python `OCRResponse` model expects). +/// +/// Blocking: intended to be called with the GIL released from the Python bridge. +pub fn run_ocr( + model: &str, + document: Value, + api_key: Option<&str>, + api_base: Option<&str>, + optional_params: Map, + timeout: Option, +) -> CoreResult { + let config = &MISTRAL_OCR_CONFIG; + + let api_key = mistral::resolve_api_key(api_key, &|key| std::env::var(key).ok())?; + let url = mistral::complete_url(api_base); + let filtered_params = config.map_ocr_params(&optional_params); + let body = config + .transform_ocr_request(model, document, filtered_params)? + .data; + + let mut request = http_client().post(&url).bearer_auth(&api_key).json(&body); + if let Some(duration) = timeout { + request = request.timeout(duration); + } + + let response = request + .send() + .map_err(|err| CoreError::Network(err.to_string()))?; + + let status = response.status(); + let text = response + .text() + .map_err(|err| CoreError::Network(err.to_string()))?; + + if !status.is_success() { + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + }); + } + + let response_json: Value = serde_json::from_str(&text) + .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; + + Ok(config + .transform_ocr_response(model, response_json)? + .into_json()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn truncate_error_body_passes_short_strings_through() { + let body = "Unauthorized"; + assert_eq!(truncate_error_body(body), "Unauthorized"); + } + + #[test] + fn truncate_error_body_caps_long_payloads() { + let body = "x".repeat(ERROR_BODY_MAX_CHARS + 50); + let truncated = truncate_error_body(&body); + + assert!(truncated.ends_with("... (truncated)")); + let prefix_chars = truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(); + assert_eq!(prefix_chars, ERROR_BODY_MAX_CHARS); + } + + #[test] + fn truncate_error_body_does_not_split_multibyte_chars() { + let body = "é".repeat(ERROR_BODY_MAX_CHARS + 10); + let truncated = truncate_error_body(&body); + assert!(truncated.is_char_boundary(truncated.len())); + } +} diff --git a/litellm-rust/crates/providers/src/openai/mod.rs b/litellm-rust/crates/providers/src/openai/mod.rs new file mode 100644 index 00000000000..403e32975cf --- /dev/null +++ b/litellm-rust/crates/providers/src/openai/mod.rs @@ -0,0 +1 @@ +pub mod realtime; diff --git a/litellm-rust/crates/providers/src/openai/realtime/mod.rs b/litellm-rust/crates/providers/src/openai/realtime/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/providers/src/openai/realtime/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/providers/src/openai/realtime/transformation.rs b/litellm-rust/crates/providers/src/openai/realtime/transformation.rs new file mode 100644 index 00000000000..2e127c699e0 --- /dev/null +++ b/litellm-rust/crates/providers/src/openai/realtime/transformation.rs @@ -0,0 +1,189 @@ +use litellm_core::realtime::transformation::RealtimeProviderConfig; +use litellm_core::realtime::types::{RealtimeEvent, RealtimeTransformResult}; +use litellm_core::CoreResult; + +/// Default OpenAI API base, used when the caller does not override `api_base`. +pub const OPENAI_REALTIME_DEFAULT_API_BASE: &str = "https://api.openai.com"; + +/// Path appended to the resolved host base to reach the realtime endpoint. +pub const OPENAI_REALTIME_PATH: &str = "/v1/realtime"; + +/// Percent-encode a query value, escaping any char outside the RFC 3986 +/// unreserved set (`A-Za-z0-9-._~`). Keeps us dependency-free; common realtime +/// model slugs have no special chars, but this stays correct for the rest. +fn percent_encode(value: &str) -> String { + let mut encoded = String::with_capacity(value.len()); + for byte in value.bytes() { + let unreserved = byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~'); + if unreserved { + encoded.push(byte as char); + } else { + encoded.push('%'); + encoded.push_str(&format!("{byte:02X}")); + } + } + encoded +} + +/// Build the realtime WebSocket URL, porting Python's `OpenAIRealtime._construct_url`. +/// +/// Blank/whitespace `api_base` is treated as absent (guard at resolution time), +/// falling back to the default. The scheme is swapped to its WebSocket +/// equivalent (`https://`→`wss://`, `http://`→`ws://`); bases already using +/// `ws`/`wss` are left untouched. A bare host or unrecognized scheme defaults to +/// secure `wss://` so we never hand a scheme-less URL to the connector (this is +/// a deliberate hardening over Python's `_construct_url`, which would emit a +/// scheme-less URL here). A trailing `/` is trimmed before the path and +/// `?model=` are appended. +pub fn complete_url(api_base: Option<&str>, model: &str) -> String { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(OPENAI_REALTIME_DEFAULT_API_BASE); + + let base = if let Some(rest) = base.strip_prefix("https://") { + format!("wss://{rest}") + } else if let Some(rest) = base.strip_prefix("http://") { + format!("ws://{rest}") + } else if base.starts_with("wss://") || base.starts_with("ws://") { + base.to_string() + } else { + format!("wss://{base}") + }; + + let base = base.trim_end_matches('/'); + + format!( + "{base}{OPENAI_REALTIME_PATH}?model={}", + percent_encode(model) + ) +} + +pub struct OpenAiRealtimeConfig; + +pub const OPENAI_REALTIME_CONFIG: OpenAiRealtimeConfig = OpenAiRealtimeConfig; + +impl RealtimeProviderConfig for OpenAiRealtimeConfig { + fn complete_url(&self, api_base: Option<&str>, model: &str) -> String { + complete_url(api_base, model) + } + + fn transform_realtime_request( + &self, + event: &RealtimeEvent, + _model: &str, + ) -> CoreResult { + Ok(RealtimeTransformResult::passthrough(event.clone())) + } + + fn transform_realtime_response( + &self, + event: &RealtimeEvent, + _model: &str, + ) -> CoreResult { + Ok(RealtimeTransformResult::passthrough(event.clone())) + } +} + +pub fn transform_realtime_request( + event: &RealtimeEvent, + model: &str, +) -> CoreResult { + OPENAI_REALTIME_CONFIG.transform_realtime_request(event, model) +} + +pub fn transform_realtime_response( + event: &RealtimeEvent, + model: &str, +) -> CoreResult { + OPENAI_REALTIME_CONFIG.transform_realtime_response(event, model) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn complete_url_defaults_to_openai_wss() { + assert_eq!( + complete_url(None, "gpt-4o-realtime-preview"), + "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_blank_base_uses_default() { + assert_eq!( + complete_url(Some(" "), "gpt-4o-realtime-preview"), + "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_swaps_http_to_ws() { + assert_eq!( + complete_url(Some("http://localhost:8080"), "gpt-4o-realtime-preview"), + "ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_dedupes_trailing_slash() { + assert_eq!( + complete_url(Some("https://api.openai.com/"), "gpt-4o-realtime-preview"), + "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_custom_base() { + assert_eq!( + complete_url(Some("https://oai.azure.example"), "gpt-4o-realtime-preview"), + "wss://oai.azure.example/v1/realtime?model=gpt-4o-realtime-preview" + ); + } + + #[test] + fn complete_url_preserves_existing_wss_scheme() { + assert_eq!( + complete_url(Some("wss://api.openai.com"), "gpt-realtime"), + "wss://api.openai.com/v1/realtime?model=gpt-realtime" + ); + } + + #[test] + fn complete_url_bare_host_defaults_to_wss() { + assert_eq!( + complete_url(Some("api.openai.com"), "gpt-realtime"), + "wss://api.openai.com/v1/realtime?model=gpt-realtime" + ); + } + + #[test] + fn complete_url_percent_encodes_model_space() { + assert_eq!( + complete_url(None, "gpt 4o"), + "wss://api.openai.com/v1/realtime?model=gpt%204o" + ); + } + + #[test] + fn transform_realtime_request_passthrough_preserves_event() { + let event: RealtimeEvent = + serde_json::from_str(r#"{"type":"session.update","session":{"voice":"alloy"}}"#) + .expect("valid event"); + let result = + transform_realtime_request(&event, "gpt-realtime").expect("passthrough is infallible"); + assert_eq!(result.events, vec![event]); + } + + #[test] + fn transform_realtime_response_passthrough_preserves_event() { + let event: RealtimeEvent = + serde_json::from_str(r#"{"type":"response.output_audio.delta","delta":"abc=="}"#) + .expect("valid event"); + let result = + transform_realtime_response(&event, "gpt-realtime").expect("passthrough is infallible"); + assert_eq!(result.events, vec![event]); + } +} diff --git a/litellm-rust/crates/providers/src/realtime.rs b/litellm-rust/crates/providers/src/realtime.rs new file mode 100644 index 00000000000..f2946db4137 --- /dev/null +++ b/litellm-rust/crates/providers/src/realtime.rs @@ -0,0 +1,193 @@ +//! End-to-end OpenAI realtime invocation. +//! +//! The host-facing entry point, mirroring `providers::ocr::run_ocr`: open the +//! WebSocket to OpenAI, drive typed events through the pure +//! `OPENAI_REALTIME_CONFIG` transforms, and collect the response events. +//! Network, auth header, key resolution, and wire (de)serialization live here so +//! the `transformation` module stays pure and typed. + +use std::time::Duration; + +use futures_util::{SinkExt, StreamExt}; +use litellm_core::error::CoreError; +use litellm_core::realtime::transformation::RealtimeProviderConfig; +use litellm_core::realtime::types::RealtimeEvent; +use litellm_core::CoreResult; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::Message; + +use crate::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; + +/// Environment variable holding the OpenAI API key (last-resort fallback). +const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; + +/// Default overall ceiling for a single realtime invocation. +const DEFAULT_TIMEOUT_SECS: u64 = 60; + +const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; + +/// Resolve the OpenAI API key from the explicit param or the environment. +/// +/// Blank/whitespace values are treated as absent (guard at resolution time). +fn resolve_api_key(api_key: Option<&str>) -> CoreResult { + api_key + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| { + std::env::var(OPENAI_API_KEY_ENV) + .ok() + .filter(|key| !key.trim().is_empty()) + }) + .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) +} + +/// True for events that end a realtime turn: a completed response or an error. +fn is_terminal_event(event: &RealtimeEvent) -> bool { + event.event_type == "response.done" || event.event_type == "error" +} + +/// Invoke the OpenAI realtime API end to end over a WebSocket. +/// +/// Sends each `input_events` entry after passing it through +/// `transform_realtime_request`, then collects backend events — each passed +/// through `transform_realtime_response` — until a terminal event +/// (`response.done` / `error`) arrives, the socket closes, or the `timeout` +/// elapses. Returns the transformed backend events in arrival order. +/// +/// Mirrors `run_ocr`: pure transforms come from `core`/`providers`; the network, +/// auth header, key resolution, and JSON (de)serialization are owned here. +pub async fn realtime( + model: &str, + input_events: Vec, + api_key: Option<&str>, + api_base: Option<&str>, + timeout: Option, +) -> CoreResult> { + let config = &OPENAI_REALTIME_CONFIG; + let api_key = resolve_api_key(api_key)?; + let url = config.complete_url(api_base, model); + + let mut request = url + .as_str() + .into_client_request() + .map_err(|err| CoreError::Network(err.to_string()))?; + // GA realtime API: only Authorization is needed. The legacy + // `OpenAI-Beta: realtime=v1` header opts into the now-removed beta request + // shape and triggers `beta_api_shape_disabled`, so we do not send it. + request.headers_mut().insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {api_key}")) + .map_err(|err| CoreError::Auth(err.to_string()))?, + ); + + let (mut ws, _response) = connect_async(request) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + + for event in &input_events { + for outbound in config.transform_realtime_request(event, model)?.events { + let payload = serde_json::to_string(&outbound) + .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; + ws.send(Message::Text(payload)) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + } + } + + let deadline = timeout.unwrap_or_else(|| Duration::from_secs(DEFAULT_TIMEOUT_SECS)); + let mut received: Vec = Vec::new(); + + let collect = async { + while let Some(message) = ws.next().await { + match message.map_err(|err| CoreError::Network(err.to_string()))? { + Message::Text(text) => { + let event: RealtimeEvent = serde_json::from_str(text.as_str()) + .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; + for outbound in config.transform_realtime_response(&event, model)?.events { + let terminal = is_terminal_event(&outbound); + received.push(outbound); + if terminal { + return Ok::<(), CoreError>(()); + } + } + } + Message::Close(_) => return Ok(()), + _ => {} + } + } + Ok(()) + }; + + tokio::time::timeout(deadline, collect) + .await + .map_err(|_| CoreError::Network("realtime call timed out".to_string()))??; + + Ok(received) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn event(raw: &str) -> RealtimeEvent { + serde_json::from_str(raw).expect("valid event json") + } + + #[test] + fn resolve_api_key_prefers_param_then_blank_falls_through() { + assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test"); + // A blank param with no env set should error. + if std::env::var(OPENAI_API_KEY_ENV).is_err() { + assert!(resolve_api_key(Some(" ")).is_err()); + } + } + + #[test] + fn is_terminal_event_matches_done_and_error_only() { + assert!(is_terminal_event(&event(r#"{"type":"response.done"}"#))); + assert!(is_terminal_event(&event(r#"{"type":"error","error":{}}"#))); + assert!(!is_terminal_event(&event( + r#"{"type":"response.output_text.delta"}"# + ))); + } + + /// Live end-to-end check against OpenAI. Ignored by default (CI never runs + /// it); run explicitly with `OPENAI_API_KEY` set: + /// `cargo test -p litellm-providers realtime_invokes_openai -- --ignored --nocapture` + #[tokio::test] + #[ignore = "hits the live OpenAI realtime API; needs OPENAI_API_KEY"] + async fn realtime_invokes_openai_and_responds() { + let key = + std::env::var(OPENAI_API_KEY_ENV).expect("set OPENAI_API_KEY to run this ignored test"); + + let response_create = event( + r#"{"type":"response.create","response":{"output_modalities":["text"],"instructions":"Respond with exactly: hello world"}}"#, + ); + + let events = realtime( + "gpt-realtime", + vec![response_create], + Some(&key), + None, + Some(Duration::from_secs(30)), + ) + .await + .expect("realtime call should succeed"); + + let types: Vec<&str> = events.iter().map(|e| e.event_type.as_str()).collect(); + eprintln!("received {} events: {:?}", events.len(), types); + + assert!( + types.contains(&"response.done"), + "expected a response.done event, got: {types:?}" + ); + assert!( + types.contains(&"response.output_text.delta"), + "expected streamed text output, got: {types:?}" + ); + } +} diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md new file mode 100644 index 00000000000..efa1a554c9c --- /dev/null +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -0,0 +1,36 @@ +# CLAUDE.md + +Rules for `litellm-rust/crates/python-bridge`. + +## Responsibility + +`python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms. +Keep this crate thin. It adapts Python objects to Rust payloads and returns +Python-compatible dictionaries. + +## Bridge Shape + +- Prefer one stable method per top-level LiteLLM route, for example + `ocr(payload)`. +- Do not add one exported PyO3 function per provider helper unless there is a + measured reason. +- Provider dispatch belongs in Rust route modules such as + `litellm_providers::ocr`, not in this PyO3 crate. +- Python owns rollout state and fallback. Rust should return errors; Python + decides whether to raise or fall back. + +## Data Handling + +- OCR payloads can contain personal data and large base64 images. Do not log + payloads or provider responses. +- Avoid copying large payloads more than needed. The current JSON round-trip is + acceptable for the first scaffold, but future performance work should evaluate + direct PyO3 conversion before expanding Rust coverage to image-heavy paths. +- Do not expose raw Rust errors that include document contents or upstream + bodies. + +## Tests + +- `cargo test --workspace` must compile this crate. +- Python tests must cover bridge disabled, bridge enabled, and module-missing + fallback behavior for every exposed route. diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml new file mode 100644 index 00000000000..80b6478daac --- /dev/null +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "litellm-python-bridge" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[lib] +name = "litellm_python_bridge" +crate-type = ["cdylib"] + +[dependencies] +litellm-core.workspace = true +litellm-providers.workspace = true +pyo3 = { workspace = true, features = ["extension-module"] } +serde_json.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/gil.rs b/litellm-rust/crates/python-bridge/src/gil.rs new file mode 100644 index 00000000000..dc1b591735c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/gil.rs @@ -0,0 +1,32 @@ +//! GIL accounting. +//! +//! A single chokepoint for releasing the GIL around blocking work. Every +//! blocking call in the bridge goes through [`release_gil`] instead of calling +//! `Python::allow_threads` directly, so the release count stays accurate and we +//! have one place to extend later (timing histograms, per-call labels, etc.). + +use std::sync::atomic::{AtomicU64, Ordering}; + +use pyo3::prelude::*; + +/// Number of times the bridge has released the GIL since process start. +static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); + +/// Release the GIL around `f`, recording the release. +/// +/// `f` must not touch any Python state — that is what makes releasing the GIL +/// safe. Returning the value back to Python re-acquires the GIL at the call +/// site, after `f` has finished. +pub fn release_gil(py: Python<'_>, f: F) -> T +where + F: FnOnce() -> T + Send, + T: Send, +{ + GIL_RELEASES.fetch_add(1, Ordering::Relaxed); + py.allow_threads(f) +} + +/// Total GIL releases performed by the bridge so far. +pub fn release_count() -> u64 { + GIL_RELEASES.load(Ordering::Relaxed) +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs new file mode 100644 index 00000000000..15e93f7b00c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -0,0 +1,100 @@ +use std::time::Duration; + +use litellm_core::error::CoreError; +use litellm_providers::ocr::run_ocr; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyAny, PyDict}; +use serde_json::{Map, Value}; + +mod gil; + +fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { + let json = py.import("json")?; + let encoded: String = json.call_method1("dumps", (value,))?.extract()?; + serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string())) +} + +fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { + let json = py.import("json")?; + let encoded = + serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?; + Ok(json.call_method1("loads", (encoded,))?.unbind()) +} + +/// Map a core error to the closest Python exception. Caller-input problems +/// (auth, bad types, missing fields) -> `ValueError`; everything else +/// (network, upstream status, parse failures) -> `RuntimeError`. +fn core_error_to_pyerr(err: CoreError) -> PyErr { + match err { + CoreError::Auth(message) => PyValueError::new_err(message), + CoreError::InvalidType { .. } | CoreError::MissingField(_) => { + PyValueError::new_err(err.to_string()) + } + other => PyRuntimeError::new_err(other.to_string()), + } +} + +/// Perform a Mistral OCR call end to end and return the response as a dict. +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, optional_params=None, timeout_seconds=None))] +fn ocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let document = py_to_json(py, document.bind(py))?; + + let optional_params = match optional_params { + Some(params) => match py_to_json(py, params.bind(py))? { + Value::Object(map) => map, + _ => return Err(PyValueError::new_err("optional_params must be a dict")), + }, + None => Map::new(), + }; + + let timeout = timeout_seconds.and_then(|secs| { + if secs.is_finite() && secs > 0.0 { + Some(Duration::from_secs_f64(secs)) + } else { + None + } + }); + + // Release the GIL during the blocking HTTP call (counted for observability). + let result = gil::release_gil(py, || { + run_ocr( + &model, + document, + api_key.as_deref(), + api_base.as_deref(), + optional_params, + timeout, + ) + }); + + match result { + Ok(value) => json_to_py(py, value), + Err(err) => Err(core_error_to_pyerr(err)), + } +} + +/// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe +/// how often the bridge has dropped the GIL for blocking work. +#[pyfunction] +fn gil_stats(py: Python<'_>) -> PyResult> { + let stats = PyDict::new(py); + stats.set_item("releases", gil::release_count())?; + Ok(stats.into_any().unbind()) +} + +#[pymodule] +fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(ocr, module)?)?; + module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + Ok(()) +} diff --git a/litellm/__init__.py b/litellm/__init__.py index d21234d2a81..d0513f77b35 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -80,6 +80,7 @@ from litellm.constants import ( WANDB_MODELS, REPEATED_STREAMING_CHUNK_LIMIT, request_timeout, + request_timeout_explicitly_set as request_timeout_explicitly_set, open_ai_embedding_models, cohere_embedding_models, bedrock_embedding_models, @@ -673,6 +674,7 @@ elevenlabs_models: Set = set() dashscope_models: Set = set() moonshot_models: Set = set() publicai_models: Set = set() +darkbloom_models: Set = set() v0_models: Set = set() morph_models: Set = set() lambda_ai_models: Set = set() @@ -927,6 +929,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None): moonshot_models.add(key) elif value.get("litellm_provider") == "publicai": publicai_models.add(key) + elif value.get("litellm_provider") == "darkbloom": + darkbloom_models.add(key) elif value.get("litellm_provider") == "v0": v0_models.add(key) elif value.get("litellm_provider") == "morph": @@ -1075,6 +1079,7 @@ model_list = list( | dashscope_models | moonshot_models | publicai_models + | darkbloom_models | v0_models | morph_models | lambda_ai_models @@ -1179,6 +1184,7 @@ models_by_provider: dict = { "modelscope": modelscope_models, "moonshot": moonshot_models, "publicai": publicai_models, + "darkbloom": darkbloom_models, "v0": v0_models, "morph": morph_models, "lambda_ai": lambda_ai_models, @@ -1400,6 +1406,7 @@ from .skills.main import ( ) from .containers.main import * from .ocr.main import * +from .ocr.rust_bridge import use_litellm_rust from .rag.main import * from .sandbox.main import * from .search.main import * @@ -1922,9 +1929,6 @@ if TYPE_CHECKING: from .llms.fireworks_ai.completion.transformation import ( FireworksAITextCompletionConfig as FireworksAITextCompletionConfig, ) - from .llms.fireworks_ai.audio_transcription.transformation import ( - FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig, - ) from .llms.fireworks_ai.embed.fireworks_ai_transformation import ( FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index e653b40fd04..4f131354d2e 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -260,7 +260,6 @@ LLM_CONFIG_NAMES = ( "SambaNovaEmbeddingConfig", "FireworksAIConfig", "FireworksAITextCompletionConfig", - "FireworksAIAudioTranscriptionConfig", "FireworksAIEmbeddingConfig", "FriendliaiChatConfig", "JinaAIEmbeddingConfig", @@ -1027,10 +1026,6 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.fireworks_ai.completion.transformation", "FireworksAITextCompletionConfig", ), - "FireworksAIAudioTranscriptionConfig": ( - ".llms.fireworks_ai.audio_transcription.transformation", - "FireworksAIAudioTranscriptionConfig", - ), "FireworksAIEmbeddingConfig": ( ".llms.fireworks_ai.embed.fireworks_ai_transformation", "FireworksAIEmbeddingConfig", diff --git a/litellm/constants.py b/litellm/constants.py index c0e265c0e4a..212d34357f8 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -201,6 +201,18 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int( # Provider-specific API base URLs XAI_API_BASE = "https://api.x.ai/v1" +OPEN_SANDBOX_API_BASE_ENV_VAR = "OPEN_SANDBOX_API_BASE" +OPEN_SANDBOX_API_KEY_ENV_VAR = "OPEN_SANDBOX_API_KEY" +OPEN_SANDBOX_DEFAULT_TEMPLATE = "opensandbox/code-interpreter:v1.1.0" +_OPEN_SANDBOX_FALLBACK_ENTRYPOINT = "/opt/code-interpreter/code-interpreter.sh" +OPEN_SANDBOX_DEFAULT_ENTRYPOINT = (_OPEN_SANDBOX_FALLBACK_ENTRYPOINT,) +OPEN_SANDBOX_DEFAULT_LANGUAGE = "python" +OPEN_SANDBOX_DEFAULT_CPU_LIMIT = "1" +OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT = "2Gi" +OPEN_SANDBOX_EXECD_PORT = 44772 +OPEN_SANDBOX_DEFAULT_TIMEOUT = 300 +OPEN_SANDBOX_READY_TIMEOUT = 30.0 +OPEN_SANDBOX_POLL_INTERVAL = 0.2 DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = int( os.getenv("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024) @@ -456,6 +468,7 @@ HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: float = 5.0 request_timeout: float = float( os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS))) ) +request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ DEFAULT_A2A_AGENT_TIMEOUT: float = float( os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000) ) # 10 minutes @@ -867,6 +880,7 @@ openai_compatible_providers: List = [ "docker_model_runner", "ragflow", "pinstripes", # Pinstripes - JSON-configured provider + "darkbloom", ] openai_text_completion_compatible_providers: List = ( [ # providers that support `/v1/completions` diff --git a/litellm/integrations/code_interpreter_interception/__init__.py b/litellm/integrations/code_interpreter_interception/__init__.py new file mode 100644 index 00000000000..2256356b6f5 --- /dev/null +++ b/litellm/integrations/code_interpreter_interception/__init__.py @@ -0,0 +1,15 @@ +""" +Code Interpreter Interception Module + +Converts the native OpenAI Responses ``code_interpreter`` tool into a function +tool, runs the model-emitted code in a sandbox, and feeds the result back into +the agentic loop. +""" + +from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, +) + +__all__ = [ + "CodeInterpreterInterceptionLogger", +] diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py new file mode 100644 index 00000000000..362581937d7 --- /dev/null +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -0,0 +1,839 @@ +""" +Code Interpreter Interception Handler + +CustomLogger that swaps the native OpenAI Responses ``code_interpreter`` tool for +a function tool, executes the code the model emits inside a sandbox, and feeds the +captured stdout back through the typed agentic loop plan. +""" + +import json +import time +import uuid +from typing import Any, Literal, TypedDict, cast + +import litellm +from pydantic import ValidationError + +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.code_interpreter_interception import ( + CodeInterpreterInterceptionConfig, +) +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, + CHAT_COMPLETION_AGENTIC_SURFACE, + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + is_interception_internal_key, +) +from litellm.types.llms.openai import ( + ChatCompletionAssistantMessage, + ChatCompletionAssistantToolCall, + ChatCompletionToolMessage, +) +from litellm.types.utils import ( + CallTypes, + ChatCompletionMessageToolCall, + ModelResponse, +) + +LITELLM_CODE_EXECUTION_TOOL_NAME = "litellm_code_execution" +_INTERCEPTION_ACTIVE_KEY = "_code_interpreter_interception_active" +_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +_CONVERTED_STREAM_KEY = "_code_interpreter_interception_converted_stream" +_LITELLM_METADATA_KEY = "litellm_metadata" +_CACHE_TTL_SECONDS = 15 * 60 + + +class CodeExecutionToolCall(TypedDict, total=False): + id: str | None + call_id: str | None + type: Literal["function"] + name: str + arguments: str + + +class CodeInterpreterLogOutput(TypedDict): + type: Literal["logs"] + logs: str + + +class CodeInterpreterCall(TypedDict): + id: str + type: Literal["code_interpreter_call"] + status: Literal["completed"] + code: str + container_id: str | None + outputs: list[CodeInterpreterLogOutput] + + +class CodeExecutionFunctionParameters(TypedDict): + type: Literal["object"] + properties: dict[str, dict[str, str]] + required: list[str] + + +class ResponsesFunctionTool(TypedDict): + type: Literal["function"] + name: str + description: str + parameters: CodeExecutionFunctionParameters + + +class ChatCompletionFunctionDefinition(TypedDict): + name: str + description: str + parameters: CodeExecutionFunctionParameters + + +class ChatCompletionFunctionTool(TypedDict): + type: Literal["function"] + function: ChatCompletionFunctionDefinition + + +CodeExecutionFunctionTool = ResponsesFunctionTool | ChatCompletionFunctionTool + + +class ResponsesFunctionToolChoice(TypedDict): + type: Literal["function"] + name: str + + +class ChatCompletionFunctionToolChoice(TypedDict): + type: Literal["function"] + function: dict[str, str] + + +CodeExecutionFunctionToolChoice = ( + ResponsesFunctionToolChoice | ChatCompletionFunctionToolChoice +) + + +def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> dict[str, Any] | None: + try: + from litellm.sandbox.sandbox_tools import resolve_sandbox_tool + except ImportError: + return None + return resolve_sandbox_tool(sandbox_tool_name) + + +class CodeInterpreterInterceptionLogger(CustomLogger): + """ + CustomLogger that implements transparent code-interpreter execution loops. + + Flow: + 1. Replace the native ``code_interpreter`` tool with a function tool in the + pre-call hook so the model emits code as function-call arguments. + 2. Detect ``litellm_code_execution`` function calls in the model response. + 3. Run the emitted code in a sandbox (reused per request via a server-minted + sandbox key) and build a typed rerun plan that appends the + function_call_output. + """ + + def __init__( + self, + enabled: bool = True, + enabled_providers: list[str] | None = None, + sandbox_tool_name: str | None = None, + sandbox_config: Any | None = None, + ): + super().__init__() + self.enabled = enabled + self.enabled_providers = enabled_providers + self.sandbox_tool_name = sandbox_tool_name + self.sandbox_config = sandbox_config + self._container_cache: dict[str, tuple[Any, dict[str, Any] | None, float]] = {} + + @classmethod + def from_config_yaml( + cls, config: CodeInterpreterInterceptionConfig + ) -> "CodeInterpreterInterceptionLogger": + return cls( + enabled=bool(config.get("enabled", True)), + enabled_providers=config.get("enabled_providers"), + sandbox_tool_name=config.get("sandbox_tool_name"), + ) + + @staticmethod + def initialize_from_proxy_config( + litellm_settings: dict[str, Any], + callback_specific_params: dict[str, Any], + ) -> "CodeInterpreterInterceptionLogger": + params: CodeInterpreterInterceptionConfig = {} + if "code_interpreter_interception_params" in litellm_settings: + params = litellm_settings["code_interpreter_interception_params"] + elif "code_interpreter_interception" in callback_specific_params and isinstance( + callback_specific_params["code_interpreter_interception"], dict + ): + params = cast( + CodeInterpreterInterceptionConfig, + callback_specific_params["code_interpreter_interception"], + ) + return CodeInterpreterInterceptionLogger.from_config_yaml(params) + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, Any], call_type: CallTypes | None + ) -> dict | None: + if not kwargs.get("_agentic_loop_depth"): + kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None) + kwargs.pop(_SANDBOX_KEY, None) + self._strip_interception_metadata(kwargs) + if not self.enabled: + return None + if call_type not in ( + CallTypes.responses, + CallTypes.aresponses, + CallTypes.completion, + CallTypes.acompletion, + ): + return None + if ( + self.enabled_providers is not None + and self._resolve_provider(kwargs) not in self.enabled_providers + ): + return None + + tools = kwargs.get("tools") + if not isinstance(tools, list): + return None + if not any( + isinstance(tool, dict) and tool.get("type") == "code_interpreter" + for tool in tools + ): + return None + + kwargs[_INTERCEPTION_ACTIVE_KEY] = True + kwargs[_SANDBOX_KEY] = uuid.uuid4().hex + if kwargs.get("stream"): + kwargs["stream"] = False + kwargs[_CONVERTED_STREAM_KEY] = True + self._write_interception_metadata(kwargs) + + function_tool = self._get_function_tool(call_type=call_type) + kwargs["tools"] = [ + ( + function_tool + if isinstance(tool, dict) and tool.get("type") == "code_interpreter" + else tool + ) + for tool in tools + ] + if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")): + kwargs["tool_choice"] = self._get_function_tool_choice(call_type=call_type) + return kwargs + + @staticmethod + def _strip_interception_metadata(kwargs: dict[str, Any]) -> None: + metadata = kwargs.get(_LITELLM_METADATA_KEY) + if not isinstance(metadata, dict): + return + filtered_metadata = { + key: value + for key, value in metadata.items() + if not is_interception_internal_key(key) + and not key.startswith("_agentic_loop") + and key != "max_agentic_loops" + } + if filtered_metadata: + kwargs[_LITELLM_METADATA_KEY] = filtered_metadata + else: + kwargs.pop(_LITELLM_METADATA_KEY, None) + + @staticmethod + def _write_interception_metadata(kwargs: dict[str, Any]) -> None: + metadata = kwargs.get(_LITELLM_METADATA_KEY) + metadata = dict(metadata) if isinstance(metadata, dict) else {} + for key in (_INTERCEPTION_ACTIVE_KEY, _SANDBOX_KEY, _CONVERTED_STREAM_KEY): + if key in kwargs: + metadata[key] = kwargs[key] + kwargs[_LITELLM_METADATA_KEY] = metadata + + @staticmethod + def _get_function_parameters() -> CodeExecutionFunctionParameters: + return { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + } + + def _get_function_tool( + self, call_type: CallTypes | None + ) -> CodeExecutionFunctionTool: + description = "Execute python code in a sandbox and return stdout." + if call_type in (CallTypes.completion, CallTypes.acompletion): + return { + "type": "function", + "function": { + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": description, + "parameters": self._get_function_parameters(), + }, + } + return { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": description, + "parameters": self._get_function_parameters(), + } + + @staticmethod + def _get_function_tool_choice( + call_type: CallTypes | None, + ) -> CodeExecutionFunctionToolChoice: + if call_type in (CallTypes.completion, CallTypes.acompletion): + return { + "type": "function", + "function": {"name": LITELLM_CODE_EXECUTION_TOOL_NAME}, + } + return { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + } + + @staticmethod + def _tool_choice_targets_code_interpreter(tool_choice: Any) -> bool: + if not isinstance(tool_choice, dict): + return False + function = tool_choice.get("function") + return ( + tool_choice.get("type") == "code_interpreter" + or tool_choice.get("name") == "code_interpreter" + or tool_choice.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + or ( + isinstance(function, dict) + and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + ) + ) + + def _resolve_provider(self, kwargs: dict[str, Any]) -> str | None: + provider = kwargs.get("custom_llm_provider") + if provider: + return provider + model = kwargs.get("model") + if not isinstance(model, str): + return None + try: + return litellm.get_llm_provider(model=model)[1] + except Exception: + return None + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + if not self.enabled: + return False, {} + if not kwargs.get(_INTERCEPTION_ACTIVE_KEY): + return False, {} + if ( + self.enabled_providers is not None + and custom_llm_provider not in self.enabled_providers + ): + return False, {} + + tool_calls = ( + self._extract_chat_completion_code_execution_tool_calls(response=response) + if kwargs.get("_agentic_loop_api_surface") + == CHAT_COMPLETION_AGENTIC_SURFACE + else self._extract_code_execution_tool_calls(response=response) + ) + if not tool_calls: + return False, {} + + return True, {"tool_calls": tool_calls} + + async def async_build_agentic_loop_plan( + self, + tools: dict, + model: str, + messages: list[dict], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: dict, + logging_obj: Any, + stream: bool, + kwargs: dict, + ) -> AgenticLoopPlan: + if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE: + return await self._build_chat_completion_agentic_loop_plan( + tools=tools, + model=model, + messages=messages, + optional_params=anthropic_messages_optional_request_params, + kwargs=kwargs, + ) + + await self._prune_expired_cache() + tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) + sandbox_key = kwargs.get(_SANDBOX_KEY) + container, params = await self._get_or_create_container(cache_key=sandbox_key) + + try: + container_id = cast(str | None, getattr(container, "id", None)) + input_list = self._normalize_messages(messages) + code_interpreter_calls: list[CodeInterpreterCall] = [] + for tool_call in tool_calls: + arguments = tool_call.get("arguments", "") + code = self._parse_code(arguments) + stdout = await self._run_tool_call( + container=container, params=params, arguments=arguments + ) + input_list.append( + { + "type": "function_call", + "call_id": tool_call.get("call_id"), + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": arguments, + } + ) + input_list.append( + { + "type": "function_call_output", + "call_id": tool_call.get("call_id"), + "output": stdout, + } + ) + code_interpreter_calls.append( + { + "id": f"ci_{uuid.uuid4().hex}", + "type": "code_interpreter_call", + "status": "completed", + "code": code, + "container_id": container_id, + "outputs": ( + [{"type": "logs", "logs": stdout}] if stdout else [] + ), + } + ) + except Exception: + await self._delete_container_for_cache_key(sandbox_key) + raise + + optional_params = anthropic_messages_optional_request_params + request_patch = AgenticLoopRequestPatch( + model=model, + messages=input_list, + tools=self._get_followup_tools( + tools=optional_params.get("tools"), + call_type=CallTypes.responses, + ), + optional_params=self._get_followup_optional_params(optional_params), + kwargs=self._filter_agentic_loop_kwargs(kwargs), + ) + + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={ + "tool_type": "code_interpreter", + "sandbox_key": sandbox_key or "", + "code_interpreter_calls": code_interpreter_calls, + }, + ) + + async def _build_chat_completion_agentic_loop_plan( + self, + tools: dict[str, object], + model: str, + messages: list[dict], + optional_params: dict[str, object], + kwargs: dict[str, object], + ) -> AgenticLoopPlan: + await self._prune_expired_cache() + tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) + sandbox_key = cast(str | None, kwargs.get(_SANDBOX_KEY)) + container, params = await self._get_or_create_container(cache_key=sandbox_key) + + try: + container_id = cast(str | None, getattr(container, "id", None)) + tool_results = [ + await self._build_chat_completion_tool_result( + container=container, + params=params, + tool_call=tool_call, + container_id=container_id, + ) + for tool_call in tool_calls + ] + except Exception: + await self._delete_container_for_cache_key(sandbox_key) + raise + tool_messages = [result[0] for result in tool_results] + code_interpreter_calls = [result[1] for result in tool_results] + + request_patch = AgenticLoopRequestPatch( + model=model, + messages=list(messages) + + [self._build_chat_completion_assistant_message(tool_calls)] + + tool_messages, + tools=self._get_followup_tools( + tools=optional_params.get("tools"), + call_type=CallTypes.completion, + ), + optional_params=self._get_followup_optional_params(optional_params), + kwargs=self._filter_agentic_loop_kwargs(kwargs), + ) + + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=request_patch, + metadata={ + "tool_type": "code_interpreter", + "sandbox_key": sandbox_key or "", + "code_interpreter_calls": code_interpreter_calls, + "response_format": "openai", + }, + ) + + async def _build_chat_completion_tool_result( + self, + container: object, + params: dict[str, Any] | None, + tool_call: CodeExecutionToolCall, + container_id: str | None, + ) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]: + arguments = tool_call.get("arguments", "") + code = self._parse_code(arguments) + stdout = await self._run_tool_call( + container=container, params=params, arguments=arguments + ) + tool_call_id = ( + tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex + ) + return ( + { + "role": "tool", + "tool_call_id": tool_call_id, + "content": stdout, + }, + { + "id": f"ci_{uuid.uuid4().hex}", + "type": "code_interpreter_call", + "status": "completed", + "code": code, + "container_id": container_id, + "outputs": [{"type": "logs", "logs": stdout}] if stdout else [], + }, + ) + + async def async_agentic_loop_cleanup_hook( + self, plan: AgenticLoopPlan, kwargs: dict + ) -> None: + metadata = plan.metadata or {} if plan else {} + await self._delete_container_for_cache_key(metadata.get("sandbox_key")) + + @staticmethod + def _filter_agentic_loop_kwargs(kwargs: dict[str, object]) -> dict[str, object]: + return { + k: v + for k, v in kwargs.items() + if k not in {"litellm_logging_obj", "acompletion"} + and not is_interception_internal_key( + k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES + ) + } + + def _get_followup_tools( + self, tools: object, call_type: CallTypes | None + ) -> list[dict[str, Any]] | None: + if not isinstance(tools, list): + return None + return [ + ( + self._get_function_tool(call_type=call_type) + if isinstance(tool, dict) and tool.get("type") == "code_interpreter" + else tool + ) + for tool in tools + ] + + def _get_followup_optional_params( + self, optional_params: dict[str, object] + ) -> dict[str, object]: + drop_tool_choice = self._tool_choice_targets_code_interpreter( + optional_params.get("tool_choice") + ) + return { + k: v + for k, v in optional_params.items() + if k != "tools" and not (k == "tool_choice" and drop_tool_choice) + } + + async def async_post_agentic_loop_response_hook( + self, response: Any, plan: AgenticLoopPlan, kwargs: dict + ) -> Any: + metadata = plan.metadata or {} if plan else {} + await self._delete_container_for_cache_key(metadata.get("sandbox_key")) + + calls = metadata.get("code_interpreter_calls") + if not calls: + return response + + is_dict = isinstance(response, dict) + output = ( + response.get("output") if is_dict else getattr(response, "output", None) + ) + if not isinstance(output, list): + return response + + def _item_type(item: Any) -> Any: + return ( + item.get("type") + if isinstance(item, dict) + else getattr(item, "type", None) + ) + + insert_at = next( + (i for i, item in enumerate(output) if _item_type(item) == "message"), + len(output), + ) + new_output = output[:insert_at] + list(calls) + output[insert_at:] + if is_dict: + response["output"] = new_output + else: + response.output = new_output + return response + + @staticmethod + def _parse_code(arguments: str) -> str: + try: + return json.loads(arguments).get("code", "") if arguments else "" + except (json.JSONDecodeError, TypeError, AttributeError): + return "" + + async def _run_tool_call( + self, container: Any, params: dict[str, Any] | None, arguments: str + ) -> str: + try: + code = json.loads(arguments).get("code", "") if arguments else "" + except (json.JSONDecodeError, TypeError): + return "[invalid tool arguments: could not parse code]" + + result = await self._run_code(container=container, params=params, code=code) + if getattr(result, "error", None): + error = result.error + message = ( + error.get("value") or error.get("name") + if isinstance(error, dict) + else str(error) + ) + return f"[execution error] {message}" + return getattr(result, "stdout", "") or "" + + async def _get_or_create_container( + self, cache_key: str | None + ) -> tuple[Any, dict[str, Any] | None]: + if cache_key: + cached = self._container_cache.get(cache_key) + if cached is not None: + return cached[0], cached[1] + + container, params = await self._create_container() + if cache_key: + self._container_cache[cache_key] = (container, params, time.time()) + return container, params + + async def _create_container(self) -> tuple[Any, dict[str, Any] | None]: + if self.sandbox_config is not None: + return await self.sandbox_config.acreate_sandbox(), None + + params = _resolve_sandbox_tool(self.sandbox_tool_name) + if params is None: + raise ValueError( + "CodeInterpreterInterception: no sandbox available. Provide a " + "sandbox_config or configure a sandbox tool resolvable via " + "sandbox_tool_name." + ) + container = await litellm.acreate_sandbox( + provider=params["sandbox_provider"], + api_key=params.get("api_key"), + api_base=params.get("api_base"), + ) + return container, params + + async def _run_code( + self, container: Any, params: dict[str, Any] | None, code: str + ) -> Any: + if self.sandbox_config is not None: + return await self.sandbox_config.arun_code(container=container, code=code) + if params is None: + raise ValueError( + "CodeInterpreterInterception: no sandbox available to run code." + ) + return await litellm.arun_code( + provider=params["sandbox_provider"], + container=container, + code=code, + api_key=params.get("api_key"), + ) + + async def _delete_container( + self, container: Any, params: dict[str, Any] | None + ) -> None: + try: + if self.sandbox_config is not None: + await self.sandbox_config.adelete_sandbox(container=container) + return + if params is None: + return + await litellm.adelete_sandbox( + provider=params["sandbox_provider"], + container=container, + api_key=params.get("api_key"), + api_base=params.get("api_base"), + ) + except Exception: + verbose_logger.exception( + "CodeInterpreterInterception: failed to delete sandbox container" + ) + + async def _delete_container_for_cache_key(self, cache_key: str | None) -> None: + if not cache_key: + return + cached = self._container_cache.pop(cache_key, None) + if cached is None: + return + await self._delete_container(container=cached[0], params=cached[1]) + + def _normalize_messages(self, messages: Any) -> list[dict[str, Any]]: + if isinstance(messages, str): + return [{"role": "user", "content": messages}] + if isinstance(messages, list): + return list(messages) + return [] + + def _extract_code_execution_tool_calls( + self, response: object + ) -> list[CodeExecutionToolCall]: + if isinstance(response, dict): + output = response.get("output", []) + else: + output = getattr(response, "output", []) or [] + if not isinstance(output, list): + return [] + + return [ + { + "call_id": ( + item.get("call_id") + if isinstance(item, dict) + else getattr(item, "call_id", None) + ), + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": ( + item.get("arguments") + if isinstance(item, dict) + else getattr(item, "arguments", "") + ), + } + for item in output + if self._is_code_execution_call(item) + ] + + def _extract_chat_completion_code_execution_tool_calls( + self, response: ModelResponse | dict[str, Any] + ) -> list[CodeExecutionToolCall]: + model_response = self._to_model_response(response) + if model_response is None: + return [] + choices = model_response.choices or [] + if not choices: + return [] + message = choices[0].message + tool_calls = message.tool_calls or [] + + return [ + normalized + for tool_call in tool_calls + if (normalized := self._normalize_chat_completion_tool_call(tool_call)) + is not None + ] + + @staticmethod + def _normalize_chat_completion_tool_call( + tool_call: ChatCompletionMessageToolCall, + ) -> CodeExecutionToolCall | None: + if ( + tool_call.type != "function" + or tool_call.function.name != LITELLM_CODE_EXECUTION_TOOL_NAME + ): + return None + + arguments = tool_call.function.arguments + if isinstance(arguments, dict): + arguments = json.dumps(arguments) + elif not isinstance(arguments, str): + arguments = "" if arguments is None else str(arguments) + + return { + "id": tool_call.id, + "call_id": tool_call.id, + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": arguments, + } + + @staticmethod + def _build_chat_completion_assistant_message( + tool_calls: list[CodeExecutionToolCall], + ) -> ChatCompletionAssistantMessage: + return { + "role": "assistant", + "tool_calls": [ + cast( + ChatCompletionAssistantToolCall, + { + "id": tool_call.get("id"), + "type": "function", + "function": { + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": tool_call.get("arguments", ""), + }, + }, + ) + for tool_call in tool_calls + ], + } + + @staticmethod + def _to_model_response( + response: ModelResponse | dict[str, Any], + ) -> ModelResponse | None: + if isinstance(response, ModelResponse): + return response + try: + return ModelResponse(**response) + except (TypeError, ValidationError): + return None + + def _is_code_execution_call(self, item: Any) -> bool: + if isinstance(item, dict): + return ( + item.get("type") == "function_call" + and item.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + ) + return ( + getattr(item, "type", None) == "function_call" + and getattr(item, "name", None) == LITELLM_CODE_EXECUTION_TOOL_NAME + ) + + async def _prune_expired_cache(self) -> None: + now = time.time() + expired = [ + (cache_key, container, params) + for cache_key, ( + container, + params, + created_at, + ) in self._container_cache.items() + if now - created_at > _CACHE_TTL_SECONDS + ] + for cache_key, container, params in expired: + self._container_cache.pop(cache_key, None) + await self._delete_container(container=container, params=params) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 481cf7fce8e..94fb97dff53 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -718,6 +718,24 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ return response + async def async_agentic_loop_cleanup_hook( + self, + plan: AgenticLoopPlan, + kwargs: dict, + ) -> None: + """ + Release resources held for an agentic-loop iteration. + + Runs in a ``finally`` around the follow-up provider call, so it fires + whether the rerun returns normally, hits a loop safety abort, or raises + an upstream error. Implementations must be idempotent because the + post-response hook may already have released the same resource on the + success path. Use ``plan.metadata`` to locate what to clean up. + + Default does nothing. + """ + return None + async def async_should_run_chat_completion_agentic_loop( self, response: Any, diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index a109ba898ff..991b156ae64 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -1,6 +1,7 @@ """Typed configuration for the OpenTelemetry instrumentation.""" from enum import Enum +from functools import lru_cache from typing import Any, List from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator @@ -47,7 +48,12 @@ class _OTelV2Flag(BaseSettings): enabled: bool = Field(default=False, validation_alias=AliasChoices(OTEL_V2_ENV)) +@lru_cache(maxsize=1) def is_otel_v2_enabled() -> bool: + # Resolved once at startup and cached: constructing the pydantic-settings + # model re-scans the environment and cost ~28us, which on the proxy hot path + # (auth, logging-callback setup) compounded into a measurable throughput + # regression. Tests that toggle the env must call ``is_otel_v2_enabled.cache_clear()``. return _OTelV2Flag().enabled diff --git a/litellm/interactions/litellm_responses_transformation/transformation.py b/litellm/interactions/litellm_responses_transformation/transformation.py index 173d4ca8764..0ff1a97cd0b 100644 --- a/litellm/interactions/litellm_responses_transformation/transformation.py +++ b/litellm/interactions/litellm_responses_transformation/transformation.py @@ -300,9 +300,6 @@ class LiteLLMResponsesInteractionsConfig: "total_output_tokens": getattr(usage, "output_tokens", 0), } - # Add role - interactions_response_dict["role"] = "model" - # Add updated (same as created for now) interactions_response_dict["updated"] = created diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py new file mode 100644 index 00000000000..938e892bd50 --- /dev/null +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -0,0 +1,332 @@ +# this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api + +import json +from typing import cast + +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.custom_logger import ( + CHAT_COMPLETION_AGENTIC_SURFACE, + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + AgenticLoopPlan, + AgenticLoopRequestPatch, + is_interception_internal_key, +) +from litellm.types.utils import ModelResponse +from litellm.utils import CustomStreamWrapper + +_FOLLOWUP_INTERNAL_PARAMS = frozenset( + ( + "acompletion", + "litellm_logging_obj", + "custom_llm_provider", + "model_alias_map", + "stream_response", + "custom_prompt_dict", + "_agentic_loop_api_surface", + ) +) + + +def _gate_overridden(callback: CustomLogger) -> bool: + base = CustomLogger.async_should_run_agentic_loop + func = type(callback).async_should_run_agentic_loop + return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + + +def _build_plan_overridden(callback: CustomLogger) -> bool: + base = CustomLogger.async_build_agentic_loop_plan + func = type(callback).async_build_agentic_loop_plan + return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + + +def _post_hook_overridden(callback: CustomLogger) -> bool: + base = CustomLogger.async_post_agentic_loop_response_hook + func = type(callback).async_post_agentic_loop_response_hook + return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + + +def _coerce_int(value: object, default: int) -> int: + return int(value) if isinstance(value, (int, str)) else default + + +def _agentic_loop_settings(kwargs: dict[str, object]) -> tuple[int, int, list[str]]: + depth = _coerce_int(kwargs.get("_agentic_loop_depth"), 0) + max_loops = max(_coerce_int(kwargs.get("max_agentic_loops"), 3), 1) + raw_fingerprints = kwargs.get("_agentic_loop_fingerprints") + fingerprints = ( + [str(fp) for fp in raw_fingerprints] + if isinstance(raw_fingerprints, list) + else [] + ) + return depth, max_loops, fingerprints + + +def _fingerprint_tools(tool_calls: object) -> str: + try: + return json.dumps(tool_calls, sort_keys=True, default=str) + except Exception: + return str(tool_calls) + + +def _check_agentic_loop_safety( + tool_calls: object, + fingerprints: list[str], + depth: int, + max_loops: int, + model: str, +) -> str: + fingerprint = _fingerprint_tools(tool_calls) + if fingerprint in fingerprints: + raise ValueError( + "Agentic loop detected repeated tool-call fingerprint; aborting rerun" + ) + if depth >= max_loops: + raise ValueError(f"Exceeded max_agentic_loops={max_loops} for model={model}") + return fingerprint + + +def _wrap_response_as_fake_stream(response: object) -> object: + if getattr(response, "object", None) == "chat.completion.chunk": + return response + if not hasattr(response, "choices"): + return response + from litellm.llms.base_llm.base_model_iterator import ( + convert_model_response_to_streaming, + ) + + return convert_model_response_to_streaming(cast(ModelResponse, response)) + + +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 _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]: + return { + k: v + for k, v in source.items() + if not is_interception_internal_key( + k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES + ) + and k not in _FOLLOWUP_INTERNAL_PARAMS + } + + +async def _execute_chat_completion_agentic_plan( + *, + plan: AgenticLoopPlan, + callback: CustomLogger, + model: str, + optional_params: dict[str, object], + kwargs: dict[str, object], + logging_obj: object, + custom_llm_provider: str, + depth: int, + max_loops: int, + fingerprints: list[str], + fingerprint: str, +) -> object: + import litellm + + patch = plan.request_patch or AgenticLoopRequestPatch() + if patch.messages is None: + raise ValueError("Agentic loop plan missing patched messages") + + full_model_name = patch.model or model + if "/" not in full_model_name: + full_model_name = f"{custom_llm_provider}/{full_model_name}" + + optional_params_for_followup = {**optional_params, **patch.optional_params} + if patch.tools is not None: + optional_params_for_followup["tools"] = patch.tools + if "tool_choice" not in patch.optional_params: + optional_params_for_followup.pop("tool_choice", None) + + kwargs_for_followup = _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["_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( + model=full_model_name, + messages=patch.messages, + **optional_params_for_followup, + **kwargs_for_followup, + ) + if _post_hook_overridden(callback): + try: + response_followup = ( + await callback.async_post_agentic_loop_response_hook( + response=response_followup, plan=plan, kwargs=kwargs + ) + ) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_post_agentic_loop_response_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + if kwargs.get("_code_interpreter_interception_converted_stream") and not depth: + return _wrap_response_as_fake_stream(response_followup) + return response_followup + finally: + try: + await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + + +async def maybe_run_chat_completion_agentic_loop( + *, + response: ModelResponse, + model: str, + messages: list, + optional_params: dict, + kwargs: dict, + logging_obj: object, + custom_llm_provider: str, + stream: bool, +) -> ModelResponse | CustomStreamWrapper | None: + import litellm + + callbacks = litellm.callbacks + ( + getattr(logging_obj, "dynamic_success_callbacks", None) or [] + ) + depth, max_loops, fingerprints = _agentic_loop_settings(kwargs) + tools = optional_params.get("tools", []) + + for callback in callbacks: + if not isinstance(callback, CustomLogger): + continue + if not _gate_overridden(callback): + continue + + gate_kwargs = { + **kwargs, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + "custom_llm_provider": custom_llm_provider, + } + try: + should_run, tool_calls = await callback.async_should_run_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=gate_kwargs, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in chat completion agentic gate: %s", + str(e), + ) + continue + + if not should_run: + continue + + fingerprint = _check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + + try: + plan_kwargs = { + **kwargs, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + "custom_llm_provider": custom_llm_provider, + } + if not _build_plan_overridden(callback): + return await callback.async_run_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=plan_kwargs, + ) + + plan = await callback.async_build_agentic_loop_plan( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=plan_kwargs, + ) + + if plan.response_override is not None: + return plan.response_override + if plan.terminate: + return response + if not plan.run_agentic_loop: + continue + + return await _execute_chat_completion_agentic_plan( + plan=plan, + callback=callback, + model=model, + optional_params=optional_params, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: %s", + str(e), + ) + + if ( + kwargs.get("_code_interpreter_interception_converted_stream") + and not depth + and hasattr(response, "choices") + ): + return cast( + "ModelResponse | CustomStreamWrapper", + _wrap_response_as_fake_stream(response), + ) + return None diff --git a/litellm/litellm_core_utils/completion_timeout.py b/litellm/litellm_core_utils/completion_timeout.py index 5350d88e593..70c6896a323 100644 --- a/litellm/litellm_core_utils/completion_timeout.py +++ b/litellm/litellm_core_utils/completion_timeout.py @@ -6,10 +6,7 @@ from typing import Callable, Optional, Union import httpx -from litellm.constants import ( - COMPLETION_HTTP_FALLBACK_SECONDS, - DEFAULT_REQUEST_TIMEOUT_SECONDS, -) +from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS class CompletionTimeout: @@ -22,17 +19,13 @@ class CompletionTimeout: """ Used when ``model_timeout`` and kwargs timeouts are all unset. - ``global_timeout`` is :attr:`litellm.request_timeout` (numeric / string), not - :class:`httpx.Timeout`. - - If it equals :data:`~litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS` (6000), - return :data:`~litellm.constants.COMPLETION_HTTP_FALLBACK_SECONDS`. Same if - ``None``. Otherwise return ``float(global_timeout)``. + ``global_timeout`` is the explicitly-configured ``litellm.request_timeout`` + (numeric / string) or ``None`` when it was never set. ``None`` falls back to + :data:`~litellm.constants.COMPLETION_HTTP_FALLBACK_SECONDS`; any explicit value + (including ``6000``) is honored. """ if global_timeout is None: return COMPLETION_HTTP_FALLBACK_SECONDS - if float(global_timeout) == float(DEFAULT_REQUEST_TIMEOUT_SECONDS): - return COMPLETION_HTTP_FALLBACK_SECONDS return float(global_timeout) @staticmethod @@ -50,11 +43,10 @@ class CompletionTimeout: 1. ``model_timeout`` (call argument / merged ``litellm_params``) 2. ``kwargs["timeout"]`` 3. ``kwargs["request_timeout"]`` - 4. Fallback from ``global_timeout`` (:attr:`litellm.request_timeout`) — if it is - the package default (6000), use 600 instead. + 4. ``global_timeout`` (the explicitly-configured ``litellm.request_timeout``), + or 600 when nothing was configured. Coerce :class:`httpx.Timeout` when the provider does not support it. - Explicit ``6000`` on the model or in kwargs is kept as ``6000``. """ resolved: Union[float, str, httpx.Timeout] if model_timeout is not None: diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index edb97b310d7..9b2a9af4126 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1,7 +1,7 @@ import json import re import traceback -from typing import Any, Optional +from typing import Any, Optional, Protocol, cast import httpx @@ -244,6 +244,2011 @@ def extract_and_raise_litellm_exception( ) +class _ProviderHTTPException(Protocol): + status_code: int + message: str + response: httpx.Response + request: httpx.Request + body: object + code: str + llm_provider: str + + +def _map_openai_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + # custom_llm_provider is openai, make it OpenAI + message = get_error_message(error_obj=original_exception) + if message is None: + if hasattr(original_exception, "message"): + message = original_exception.message + else: + message = str(original_exception) + + if message is not None and isinstance( + message, str + ): # done to prevent user-confusion. Relevant issue - https://github.com/BerriAI/litellm/issues/1414 + message = message.replace("OPENAI", custom_llm_provider.upper()) + message = message.replace( + "openai.OpenAIError", + "{}.{}Error".format(custom_llm_provider, custom_llm_provider), + ) + if custom_llm_provider == "openai": + exception_provider = "OpenAI" + "Exception" + else: + exception_provider = ( + custom_llm_provider[0].upper() + custom_llm_provider[1:] + "Exception" + ) + + if ExceptionCheckers.is_error_str_rate_limit(error_str): + raise RateLimitError( + message=f"RateLimitError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + ) + elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): + raise ContextWindowExceededError( + message=f"ContextWindowExceededError: {exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif "invalid_request_error" in error_str and "model_not_found" in error_str: + raise NotFoundError( + message=f"{exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif "A timeout occurred" in error_str: + raise Timeout( + message=f"{exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif ( + ( + "invalid_request_error" in error_str + and "content_policy_violation" in error_str + ) + or ("Invalid prompt" in error_str and "violating our usage policy" in error_str) + or ( + "request was rejected as a result of the safety system" in error_str.lower() + ) + ): + raise ContentPolicyViolationError( + message=f"ContentPolicyViolationError: {exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif ( + "invalid_encrypted_content" in error_str or "could not be verified" in error_str + ): + helpful_message = ( + f"{exception_provider} - {message}\n\n" + " This error occurs when load balancing Responses API across deployments with different API keys.\n" + " Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n" + " Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n" + " router_settings:\n" + " enable_pre_call_checks: true\n" + " optional_pre_call_checks:\n" + " - encrypted_content_affinity\n\n" + " Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing" + ) + raise BadRequestError( + message=helpful_message, + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + ) + elif ( + "invalid_request_error" in error_str + and "Incorrect API key provided" not in error_str + ): + raise BadRequestError( + message=f"{exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + ) + elif ( + "Web server is returning an unknown error" in error_str + or "The server had an error processing your request." in error_str + ): + raise litellm.InternalServerError( + message=f"{exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + ) + elif "Request too large" in error_str: + raise RateLimitError( + message=f"RateLimitError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif ( + "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable" + in error_str + ): + raise AuthenticationError( + message=f"AuthenticationError: {exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif "Mistral API raised a streaming error" in error_str: + _request = httpx.Request(method="POST", url="https://api.openai.com/v1") + raise APIError( + status_code=500, + message=f"{exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + request=_request, + litellm_debug_info=extra_information, + ) + elif hasattr(original_exception, "status_code"): + if original_exception.status_code == 400: + raise BadRequestError( + message=f"{exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 401: + raise AuthenticationError( + message=f"AuthenticationError: {exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 404: + raise NotFoundError( + message=f"NotFoundError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"Timeout Error: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 422: + raise BadRequestError( + message=f"{exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"RateLimitError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 500: + raise InternalServerError( + message=f"InternalServerError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 502: + raise BadGatewayError( + message=f"BadGatewayError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 503: + raise ServiceUnavailableError( + message=f"ServiceUnavailableError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 504: # gateway timeout error + raise Timeout( + message=f"Timeout Error: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + exception_status_code=original_exception.status_code, + ) + else: + raise APIError( + status_code=original_exception.status_code, + message=f"APIError: {exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + request=getattr(original_exception, "request", None), + litellm_debug_info=extra_information, + ) + else: + # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors + # exception_mapping_worked = True + raise APIConnectionError( + message=f"APIConnectionError: {exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + litellm_debug_info=extra_information, + request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), + ) + + +def _map_anthropic_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if ( + "prompt is too long" in error_str + or "prompt: length" in error_str + or ExceptionCheckers.is_error_str_context_window_exceeded(error_str) + ): + raise ContextWindowExceededError( + message="AnthropicError - {}".format(error_str), + model=model, + llm_provider="anthropic", + ) + elif "overloaded_error" in error_str or "Overloaded" in error_str: + raise InternalServerError( + message="AnthropicError - {}".format(error_str), + model=model, + llm_provider="anthropic", + ) + if "Invalid API Key" in error_str: + raise AuthenticationError( + message="AnthropicError - {}".format(error_str), + model=model, + llm_provider="anthropic", + ) + if "content filtering policy" in error_str: + raise ContentPolicyViolationError( + message="AnthropicError - {}".format(error_str), + model=model, + llm_provider="anthropic", + ) + if "Client error '400 Bad Request'" in error_str: + raise BadRequestError( + message="AnthropicError - {}".format(error_str), + model=model, + llm_provider="anthropic", + ) + if hasattr(original_exception, "status_code"): + verbose_logger.debug(f"status_code: {original_exception.status_code}") + if original_exception.status_code == 401: + raise AuthenticationError( + message=f"AnthropicException - {error_str}", + llm_provider="anthropic", + model=model, + ) + elif ( + original_exception.status_code == 400 + or original_exception.status_code == 413 + ): + raise BadRequestError( + message=f"AnthropicException - {error_str}", + model=model, + llm_provider="anthropic", + ) + elif original_exception.status_code == 404: + raise NotFoundError( + message=f"AnthropicException - {error_str}", + model=model, + llm_provider="anthropic", + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"AnthropicException - {error_str}", + model=model, + llm_provider="anthropic", + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"AnthropicException - {error_str}", + llm_provider="anthropic", + model=model, + ) + elif ( + original_exception.status_code == 500 + or original_exception.status_code == 529 + ): + raise litellm.InternalServerError( + message=f"AnthropicException - {error_str}. Handle with `litellm.InternalServerError`.", + llm_provider="anthropic", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 502: + raise BadGatewayError( + message=f"AnthropicException BadGatewayError - {error_str}", + llm_provider="anthropic", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 503: + raise litellm.ServiceUnavailableError( + message=f"AnthropicException - {error_str}. Handle with `litellm.ServiceUnavailableError`.", + llm_provider="anthropic", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 504: # gateway timeout error + raise Timeout( + message=f"AnthropicException Timeout - {error_str}", + model=model, + llm_provider="anthropic", + exception_status_code=original_exception.status_code, + ) + + +def _map_replicate_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if "Incorrect authentication token" in error_str: + raise AuthenticationError( + message=f"ReplicateException - {error_str}", + llm_provider="replicate", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "input is too long" in error_str: + raise ContextWindowExceededError( + message=f"ReplicateException - {error_str}", + model=model, + llm_provider="replicate", + response=getattr(original_exception, "response", None), + ) + elif exception_type == "ModelError": + raise BadRequestError( + message=f"ReplicateException - {error_str}", + model=model, + llm_provider="replicate", + response=getattr(original_exception, "response", None), + ) + elif "Request was throttled" in error_str: + raise RateLimitError( + message=f"ReplicateException - {error_str}", + llm_provider="replicate", + model=model, + response=getattr(original_exception, "response", None), + ) + elif hasattr(original_exception, "status_code"): + if original_exception.status_code == 401: + raise AuthenticationError( + message=f"ReplicateException - {original_exception.message}", + llm_provider="replicate", + model=model, + response=getattr(original_exception, "response", None), + ) + elif ( + original_exception.status_code == 400 + or original_exception.status_code == 413 + ): + raise BadRequestError( + message=f"ReplicateException - {original_exception.message}", + model=model, + llm_provider="replicate", + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 422: + raise UnprocessableEntityError( + message=f"ReplicateException - {original_exception.message}", + model=model, + llm_provider="replicate", + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"ReplicateException - {original_exception.message}", + model=model, + llm_provider="replicate", + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"ReplicateException - {original_exception.message}", + llm_provider="replicate", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 500: + raise ServiceUnavailableError( + message=f"ReplicateException - {original_exception.message}", + llm_provider="replicate", + model=model, + response=getattr(original_exception, "response", None), + ) + raise APIError( + status_code=500, + message=f"ReplicateException - {str(original_exception)}", + llm_provider="replicate", + model=model, + request=httpx.Request( + method="POST", + url="https://api.replicate.com/v1/deployments", + ), + ) + + +def _map_openai_like_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if "authorization denied for" in error_str: + + # Predibase returns the raw API Key in the response - this block ensures it's not returned in the exception + if ( + error_str is not None + and isinstance(error_str, str) + and "bearer" in error_str.lower() + ): + # only keep the first 10 chars after the occurnence of "bearer" + _bearer_token_start_index = error_str.lower().find("bearer") + error_str = error_str[: _bearer_token_start_index + 14] + error_str += "XXXXXXX" + '"' + + raise AuthenticationError( + message=f"{custom_llm_provider.capitalize()}Exception: Authentication Error - {error_str}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): + raise ContextWindowExceededError( + message=f"{custom_llm_provider.capitalize()}Exception: Context Window Error - {error_str}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif "token_quota_reached" in error_str: + raise RateLimitError( + message=f"{custom_llm_provider.capitalize()}Exception: Rate Limit Errror - {error_str}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + ) + elif ( + "The server received an invalid response from an upstream server." in error_str + ): + raise litellm.InternalServerError( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + ) + elif "model_no_support_for_function" in error_str: + raise BadRequestError( + message=f"{custom_llm_provider.capitalize()}Exception - Use 'watsonx_text' route instead. IBM WatsonX does not support `/text/chat` endpoint. - {error_str}", + llm_provider=custom_llm_provider, + model=model, + ) + elif hasattr(original_exception, "status_code"): + if original_exception.status_code == 500: + raise litellm.InternalServerError( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + ) + elif ( + original_exception.status_code == 401 + or original_exception.status_code == 403 + ): + raise AuthenticationError( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + ) + elif original_exception.status_code == 400: + raise BadRequestError( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + ) + elif original_exception.status_code == 404: + raise NotFoundError( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif ( + original_exception.status_code == 422 + or original_exception.status_code == 424 + ): + raise BadRequestError( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 503: + raise ServiceUnavailableError( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 504: # gateway timeout error + raise Timeout( + message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + exception_status_code=original_exception.status_code, + ) + + +def _map_bedrock_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if ( + "too many tokens" in error_str + or "expected maxLength:" in error_str + or "Input is too long" in error_str + or "prompt is too long" in error_str + or "prompt: length: 1.." in error_str + or "Too many input tokens" in error_str + ): + raise ContextWindowExceededError( + message=f"BedrockException: Context Window Error - {error_str}", + model=model, + llm_provider="bedrock", + ) + elif ( + "Conversation blocks and tool result blocks cannot be provided in the same turn." + in error_str + ): + raise BadRequestError( + message=f"BedrockException - {error_str}\n. Enable 'litellm.modify_params=True' (for PROXY do: `litellm_settings::modify_params: True`) to insert a dummy assistant message and fix this error.", + model=model, + llm_provider="bedrock", + response=getattr(original_exception, "response", None), + ) + elif "Malformed input request" in error_str: + raise BadRequestError( + message=f"BedrockException - {error_str}", + model=model, + llm_provider="bedrock", + response=getattr(original_exception, "response", None), + ) + elif "A conversation must start with a user message." in error_str: + raise BadRequestError( + message=f"BedrockException - {error_str}\n. Pass in default user message via `completion(..,user_continue_message=)` or enable `litellm.modify_params=True`.\nFor Proxy: do via `litellm_settings::modify_params: True` or user_continue_message under `litellm_params`", + model=model, + llm_provider="bedrock", + response=getattr(original_exception, "response", None), + ) + elif ( + "Unable to locate credentials" in error_str + or "The security token included in the request is invalid" in error_str + ): + raise AuthenticationError( + message=f"BedrockException Invalid Authentication - {error_str}", + model=model, + llm_provider="bedrock", + response=getattr(original_exception, "response", None), + ) + elif "AccessDeniedException" in error_str: + raise PermissionDeniedError( + message=f"BedrockException PermissionDeniedError - {error_str}", + model=model, + llm_provider="bedrock", + response=getattr(original_exception, "response", None), + ) + elif "throttlingException" in error_str or "ThrottlingException" in error_str: + raise RateLimitError( + message=f"BedrockException: Rate Limit Error - {error_str}", + model=model, + llm_provider="bedrock", + response=getattr(original_exception, "response", None), + ) + elif "Connect timeout on endpoint URL" in error_str or "timed out" in error_str: + raise Timeout( + message=f"BedrockException: Timeout Error - {error_str}", + model=model, + llm_provider="bedrock", + ) + elif "Could not process image" in error_str: + raise litellm.InternalServerError( + message=f"BedrockException - {error_str}", + model=model, + llm_provider="bedrock", + ) + elif hasattr(original_exception, "status_code"): + if original_exception.status_code == 500: + raise ServiceUnavailableError( + message=f"BedrockException - {original_exception.message}", + llm_provider="bedrock", + model=model, + response=httpx.Response( + status_code=500, + request=httpx.Request( + method="POST", url="https://api.openai.com/v1/" + ), + ), + ) + elif original_exception.status_code == 401: + raise AuthenticationError( + message=f"BedrockException - {original_exception.message}", + llm_provider="bedrock", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 400: + raise BadRequestError( + message=f"BedrockException - {original_exception.message}", + llm_provider="bedrock", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 404: + raise NotFoundError( + message=f"BedrockException - {original_exception.message}", + llm_provider="bedrock", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"BedrockException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 422: + raise BadRequestError( + message=f"BedrockException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"BedrockException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 503: + raise ServiceUnavailableError( + message=f"BedrockException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 504: # gateway timeout error + raise Timeout( + message=f"BedrockException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + exception_status_code=original_exception.status_code, + ) + + +def _map_sagemaker_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if "Unable to locate credentials" in error_str: + raise BadRequestError( + message=f"litellm.BadRequestError: SagemakerException - {error_str}", + model=model, + llm_provider="sagemaker", + response=getattr(original_exception, "response", None), + ) + elif "Input validation error: `best_of` must be > 0 and <= 2" in error_str: + raise BadRequestError( + message="SagemakerException - the value of 'n' must be > 0 and <= 2 for sagemaker endpoints", + model=model, + llm_provider="sagemaker", + response=getattr(original_exception, "response", None), + ) + elif ( + "`inputs` tokens + `max_new_tokens` must be <=" in error_str + or "instance type with more CPU capacity or memory" in error_str + ): + raise ContextWindowExceededError( + message=f"SagemakerException - {error_str}", + model=model, + llm_provider="sagemaker", + response=getattr(original_exception, "response", None), + ) + elif hasattr(original_exception, "status_code"): + if original_exception.status_code == 500: + raise ServiceUnavailableError( + message=f"SagemakerException - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + response=httpx.Response( + status_code=500, + request=httpx.Request( + method="POST", url="https://api.openai.com/v1/" + ), + ), + ) + elif original_exception.status_code == 401: + raise AuthenticationError( + message=f"SagemakerException - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 400: + raise BadRequestError( + message=f"SagemakerException - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 404: + raise NotFoundError( + message=f"SagemakerException - {original_exception.message}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"SagemakerException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif ( + original_exception.status_code == 422 + or original_exception.status_code == 424 + ): + raise BadRequestError( + message=f"SagemakerException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"SagemakerException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 503: + raise ServiceUnavailableError( + message=f"SagemakerException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 504: # gateway timeout error + raise Timeout( + message=f"SagemakerException - {original_exception.message}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + exception_status_code=original_exception.status_code, + ) + + +def _map_vertex_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if ( + "Vertex AI API has not been used in project" in error_str + or "Unable to find your project" in error_str + ): + raise BadRequestError( + message=f"litellm.BadRequestError: {custom_llm_provider}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + response=httpx.Response( + status_code=400, + request=httpx.Request( + method="POST", + url=" https://cloud.google.com/vertex-ai/", + ), + ), + litellm_debug_info=extra_information, + ) + if "400 Request payload size exceeds" in error_str: + raise ContextWindowExceededError( + message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + ) + elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): + raise ContextWindowExceededError( + message=f"ContextWindowExceededError: {custom_llm_provider.capitalize()}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif "None Unknown Error." in error_str or "Content has no parts." in error_str: + raise litellm.InternalServerError( + message=f"litellm.InternalServerError: {custom_llm_provider}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + response=httpx.Response( + status_code=500, + content=str(original_exception), + request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore + ), + litellm_debug_info=extra_information, + ) + elif "API key not valid." in error_str: + raise AuthenticationError( + message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif "403" in error_str: + raise BadRequestError( + message=f"{custom_llm_provider.capitalize()}Exception BadRequestError - {error_str}", + model=model, + llm_provider=custom_llm_provider, + response=httpx.Response( + status_code=403, + request=httpx.Request( + method="POST", + url=" https://cloud.google.com/vertex-ai/", + ), + ), + litellm_debug_info=extra_information, + ) + elif ( + "The response was blocked." in error_str + or "Output blocked by content filtering policy" + in error_str # anthropic on vertex ai + ): + raise ContentPolicyViolationError( + message=f"{custom_llm_provider.capitalize()}Exception ContentPolicyViolationError - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + response=httpx.Response( + status_code=400, + request=httpx.Request( + method="POST", + url=" https://cloud.google.com/vertex-ai/", + ), + ), + ) + elif ( + "429 Quota exceeded" in error_str + or "Quota exceeded for" in error_str + or "Resource exhausted" in error_str + or "IndexError: list index out of range" in error_str + or "429 Unable to submit request because the service is temporarily out of capacity." + in error_str + ): + raise RateLimitError( + message=f"litellm.RateLimitError: {custom_llm_provider}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + response=httpx.Response( + status_code=429, + request=httpx.Request( + method="POST", + url=" https://cloud.google.com/vertex-ai/", + ), + ), + ) + elif ( + isinstance(getattr(original_exception, "status_code", None), int) + and 500 <= original_exception.status_code < 600 + and _get_body_error_code(error_str) == 429 + ): + # upstream gateway wraps a 429 inside a 5xx envelope + # e.g. HTTP 500/503 with {"error":{"code":429,...}}. + # Scoped to 5xx so HTTP 400/401 with body code:429 + # still maps to BadRequestError / AuthenticationError. + raise RateLimitError( + message=f"litellm.RateLimitError: {custom_llm_provider}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + response=httpx.Response( + status_code=429, + request=httpx.Request( + method="POST", + url=" https://cloud.google.com/vertex-ai/", + ), + ), + ) + elif ( + "500 Internal Server Error" in error_str + or "The model is overloaded." in error_str + ): + raise litellm.InternalServerError( + message=f"litellm.InternalServerError: {custom_llm_provider}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + if hasattr(original_exception, "status_code"): + if original_exception.status_code == 400: + raise BadRequestError( + message=f"{custom_llm_provider.capitalize()}Exception BadRequestError - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + response=httpx.Response( + status_code=400, + request=httpx.Request( + method="POST", + url="https://cloud.google.com/vertex-ai/", + ), + ), + ) + if original_exception.status_code == 401: + raise AuthenticationError( + message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", + llm_provider=custom_llm_provider, + model=model, + ) + if original_exception.status_code == 403: + raise PermissionDeniedError( + message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", + llm_provider=custom_llm_provider, + model=model, + response=httpx.Response( + status_code=403, + request=httpx.Request( + method="POST", + url="https://cloud.google.com/vertex-ai/", + ), + ), + ) + if original_exception.status_code == 404: + raise NotFoundError( + message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", + llm_provider=custom_llm_provider, + model=model, + ) + if original_exception.status_code == 408: + raise Timeout( + message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", + llm_provider=custom_llm_provider, + model=model, + ) + + if original_exception.status_code == 429: + raise RateLimitError( + message=f"litellm.RateLimitError: {custom_llm_provider.capitalize()}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + response=httpx.Response( + status_code=429, + request=httpx.Request( + method="POST", + url=" https://cloud.google.com/vertex-ai/", + ), + ), + ) + if original_exception.status_code == 500: + raise litellm.InternalServerError( + message=f"{custom_llm_provider.capitalize()}Exception InternalServerError - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + response=httpx.Response( + status_code=500, + content=str(original_exception), + request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore + ), + ) + if original_exception.status_code == 502: + raise APIConnectionError( + message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", + llm_provider=custom_llm_provider, + model=model, + ) + if original_exception.status_code == 503: + raise ServiceUnavailableError( + message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", + llm_provider=custom_llm_provider, + model=model, + ) + + +def _map_cloudflare_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if "Authentication error" in error_str: + raise AuthenticationError( + message=f"Cloudflare Exception - {original_exception.message}", + llm_provider="cloudflare", + model=model, + response=getattr(original_exception, "response", None), + ) + if "must have required property" in error_str: + raise BadRequestError( + message=f"Cloudflare Exception - {original_exception.message}", + llm_provider="cloudflare", + model=model, + response=getattr(original_exception, "response", None), + ) + + +def _map_cohere_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if "invalid api token" in error_str or "No API key provided." in error_str: + raise AuthenticationError( + message=f"CohereException - {original_exception.message}", + llm_provider="cohere", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "invalid type: parameter" in error_str: + raise BadRequestError( + message=f"CohereException - {original_exception.message}", + llm_provider="cohere", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "too many tokens" in error_str: + raise ContextWindowExceededError( + message=f"CohereException - {original_exception.message}", + model=model, + llm_provider="cohere", + response=getattr(original_exception, "response", None), + ) + elif "internal server error" in error_str.lower(): + raise InternalServerError( + message=f"CohereException - {error_str}", + model=model, + llm_provider="cohere", + response=getattr(original_exception, "response", None), + ) + elif hasattr(original_exception, "status_code"): + if ( + original_exception.status_code == 400 + or original_exception.status_code == 498 + ): + raise BadRequestError( + message=f"CohereException - {original_exception.message}", + llm_provider="cohere", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"CohereException - {original_exception.message}", + llm_provider="cohere", + model=model, + ) + elif original_exception.status_code == 500: + raise InternalServerError( + message=f"CohereException - {original_exception.message}", + llm_provider="cohere", + model=model, + response=getattr(original_exception, "response", None), + ) + elif ( + "CohereConnectionError" in exception_type + ): # cohere seems to fire these errors when we load test it (1k+ messages / min) + raise RateLimitError( + message=f"CohereException - {original_exception.message}", + llm_provider="cohere", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "invalid type:" in error_str: + raise BadRequestError( + message=f"CohereException - {original_exception.message}", + llm_provider="cohere", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "Unexpected server error" in error_str: + raise InternalServerError( + message=f"CohereException - {original_exception.message}", + llm_provider="cohere", + model=model, + response=getattr(original_exception, "response", None), + ) + else: + if hasattr(original_exception, "status_code"): + raise APIError( + status_code=original_exception.status_code, + message=f"CohereException - {original_exception.message}", + llm_provider="cohere", + model=model, + request=getattr(original_exception, "request", None), + ) + raise cast(Exception, original_exception) + + +def _map_huggingface_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if "length limit exceeded" in error_str: + raise ContextWindowExceededError( + message=error_str, + model=model, + llm_provider="huggingface", + response=getattr(original_exception, "response", None), + ) + elif "A valid user token is required" in error_str: + raise BadRequestError( + message=error_str, + llm_provider="huggingface", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "Rate limit reached" in error_str: + raise RateLimitError( + message=error_str, + llm_provider="huggingface", + model=model, + response=getattr(original_exception, "response", None), + ) + if hasattr(original_exception, "status_code"): + if original_exception.status_code == 401: + raise AuthenticationError( + message=f"HuggingfaceException - {original_exception.message}", + llm_provider="huggingface", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 400: + raise BadRequestError( + message=f"HuggingfaceException - {original_exception.message}", + model=model, + llm_provider="huggingface", + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"HuggingfaceException - {original_exception.message}", + model=model, + llm_provider="huggingface", + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"HuggingfaceException - {original_exception.message}", + llm_provider="huggingface", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 503: + raise ServiceUnavailableError( + message=f"HuggingfaceException - {original_exception.message}", + llm_provider="huggingface", + model=model, + response=getattr(original_exception, "response", None), + ) + else: + raise APIError( + status_code=original_exception.status_code, + message=f"HuggingfaceException - {original_exception.message}", + llm_provider="huggingface", + model=model, + request=getattr(original_exception, "request", None), + ) + + +def _map_ai21_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if hasattr(original_exception, "message"): + if "Prompt has too many tokens" in original_exception.message: + raise ContextWindowExceededError( + message=f"AI21Exception - {original_exception.message}", + model=model, + llm_provider="ai21", + response=getattr(original_exception, "response", None), + ) + if "Bad or missing API token." in original_exception.message: + raise BadRequestError( + message=f"AI21Exception - {original_exception.message}", + model=model, + llm_provider="ai21", + response=getattr(original_exception, "response", None), + ) + if hasattr(original_exception, "status_code"): + if original_exception.status_code == 401: + raise AuthenticationError( + message=f"AI21Exception - {original_exception.message}", + llm_provider="ai21", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"AI21Exception - {original_exception.message}", + model=model, + llm_provider="ai21", + ) + if original_exception.status_code == 422: + raise BadRequestError( + message=f"AI21Exception - {original_exception.message}", + model=model, + llm_provider="ai21", + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"AI21Exception - {original_exception.message}", + llm_provider="ai21", + model=model, + response=getattr(original_exception, "response", None), + ) + else: + raise APIError( + status_code=original_exception.status_code, + message=f"AI21Exception - {original_exception.message}", + llm_provider="ai21", + model=model, + request=getattr(original_exception, "request", None), + ) + + +def _map_nlp_cloud_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if "detail" in error_str: + if "Input text length should not exceed" in error_str: + raise ContextWindowExceededError( + message=f"NLPCloudException - {error_str}", + model=model, + llm_provider="nlp_cloud", + response=getattr(original_exception, "response", None), + ) + elif "value is not a valid" in error_str: + raise BadRequestError( + message=f"NLPCloudException - {error_str}", + model=model, + llm_provider="nlp_cloud", + response=getattr(original_exception, "response", None), + ) + else: + raise APIError( + status_code=500, + message=f"NLPCloudException - {error_str}", + model=model, + llm_provider="nlp_cloud", + request=getattr(original_exception, "request", None), + ) + if hasattr( + original_exception, "status_code" + ): # https://docs.nlpcloud.com/?shell#errors + if ( + original_exception.status_code == 400 + or original_exception.status_code == 406 + or original_exception.status_code == 413 + or original_exception.status_code == 422 + ): + raise BadRequestError( + message=f"NLPCloudException - {original_exception.message}", + llm_provider="nlp_cloud", + model=model, + response=getattr(original_exception, "response", None), + ) + elif ( + original_exception.status_code == 401 + or original_exception.status_code == 403 + ): + raise AuthenticationError( + message=f"NLPCloudException - {original_exception.message}", + llm_provider="nlp_cloud", + model=model, + response=getattr(original_exception, "response", None), + ) + elif ( + original_exception.status_code == 522 + or original_exception.status_code == 524 + ): + raise Timeout( + message=f"NLPCloudException - {original_exception.message}", + model=model, + llm_provider="nlp_cloud", + ) + elif ( + original_exception.status_code == 429 + or original_exception.status_code == 402 + ): + raise RateLimitError( + message=f"NLPCloudException - {original_exception.message}", + llm_provider="nlp_cloud", + model=model, + response=getattr(original_exception, "response", None), + ) + elif ( + original_exception.status_code == 500 + or original_exception.status_code == 503 + ): + raise APIError( + status_code=original_exception.status_code, + message=f"NLPCloudException - {original_exception.message}", + llm_provider="nlp_cloud", + model=model, + request=getattr(original_exception, "request", None), + ) + elif ( + original_exception.status_code == 504 + or original_exception.status_code == 520 + ): + raise ServiceUnavailableError( + message=f"NLPCloudException - {original_exception.message}", + model=model, + llm_provider="nlp_cloud", + response=getattr(original_exception, "response", None), + ) + else: + raise APIError( + status_code=original_exception.status_code, + message=f"NLPCloudException - {original_exception.message}", + llm_provider="nlp_cloud", + model=model, + request=getattr(original_exception, "request", None), + ) + + +def _map_together_ai_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + try: + error_response = json.loads(error_str) + except Exception: + error_response = {"error": error_str} + if ( + "error" in error_response + and "`inputs` tokens + `max_new_tokens` must be <=" in error_response["error"] + ): + raise ContextWindowExceededError( + message=f"TogetherAIException - {error_response['error']}", + model=model, + llm_provider="together_ai", + response=getattr(original_exception, "response", None), + ) + elif "error" in error_response and "invalid private key" in error_response["error"]: + raise AuthenticationError( + message=f"TogetherAIException - {error_response['error']}", + llm_provider="together_ai", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "error" in error_response and "INVALID_ARGUMENT" in error_response["error"]: + raise BadRequestError( + message=f"TogetherAIException - {error_response['error']}", + model=model, + llm_provider="together_ai", + response=getattr(original_exception, "response", None), + ) + elif "A timeout occurred" in error_str: + raise Timeout( + message=f"TogetherAIException - {error_str}", + model=model, + llm_provider="together_ai", + ) + elif ( + "error" in error_response + and "API key doesn't match expected format." in error_response["error"] + ): + raise BadRequestError( + message=f"TogetherAIException - {error_response['error']}", + model=model, + llm_provider="together_ai", + response=getattr(original_exception, "response", None), + ) + elif ( + "error_type" in error_response and error_response["error_type"] == "validation" + ): + raise BadRequestError( + message=f"TogetherAIException - {error_response['error']}", + model=model, + llm_provider="together_ai", + response=getattr(original_exception, "response", None), + ) + if hasattr(original_exception, "status_code"): + if original_exception.status_code == 408: + raise Timeout( + message=f"TogetherAIException - {original_exception.message}", + model=model, + llm_provider="together_ai", + ) + elif original_exception.status_code == 422: + raise BadRequestError( + message=f"TogetherAIException - {error_response['error']}", + model=model, + llm_provider="together_ai", + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"TogetherAIException - {original_exception.message}", + llm_provider="together_ai", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 524: + raise Timeout( + message=f"TogetherAIException - {original_exception.message}", + llm_provider="together_ai", + model=model, + ) + else: + raise APIError( + status_code=original_exception.status_code, + message=f"TogetherAIException - {original_exception.message}", + llm_provider="together_ai", + model=model, + request=getattr(original_exception, "request", None), + ) + + +def _map_aleph_alpha_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if "This is longer than the model's maximum context length" in error_str: + raise ContextWindowExceededError( + message=f"AlephAlphaException - {original_exception.message}", + llm_provider="aleph_alpha", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "InvalidToken" in error_str or "No token provided" in error_str: + raise BadRequestError( + message=f"AlephAlphaException - {original_exception.message}", + llm_provider="aleph_alpha", + model=model, + response=getattr(original_exception, "response", None), + ) + elif hasattr(original_exception, "status_code"): + verbose_logger.debug(f"status code: {original_exception.status_code}") + if original_exception.status_code == 401: + raise AuthenticationError( + message=f"AlephAlphaException - {original_exception.message}", + llm_provider="aleph_alpha", + model=model, + ) + elif original_exception.status_code == 400: + raise BadRequestError( + message=f"AlephAlphaException - {original_exception.message}", + llm_provider="aleph_alpha", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"AlephAlphaException - {original_exception.message}", + llm_provider="aleph_alpha", + model=model, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 500: + raise ServiceUnavailableError( + message=f"AlephAlphaException - {original_exception.message}", + llm_provider="aleph_alpha", + model=model, + response=getattr(original_exception, "response", None), + ) + raise cast(Exception, original_exception) + raise cast(Exception, original_exception) + + +def _map_ollama_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if isinstance(original_exception, dict): + error_str = original_exception.get("error", "") + else: + error_str = str(original_exception) + if "no such file or directory" in error_str: + raise BadRequestError( + message=f"OllamaException: Invalid Model/Model not loaded - {original_exception}", + model=model, + llm_provider="ollama", + response=getattr(original_exception, "response", None), + ) + elif "Failed to establish a new connection" in error_str: + raise ServiceUnavailableError( + message=f"OllamaException: {original_exception}", + llm_provider="ollama", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "Invalid response object from API" in error_str: + raise BadRequestError( + message=f"OllamaException: {original_exception}", + llm_provider="ollama", + model=model, + response=getattr(original_exception, "response", None), + ) + elif "Read timed out" in error_str: + raise Timeout( + message=f"OllamaException: {original_exception}", + llm_provider="ollama", + model=model, + ) + + +def _map_vllm_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if hasattr(original_exception, "status_code"): + if original_exception.status_code == 0: + raise APIConnectionError( + message=f"VLLMException - {original_exception.message}", + llm_provider="vllm", + model=model, + request=getattr(original_exception, "request", None), + ) + + +def _map_azure_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + message = get_error_message(error_obj=original_exception) + if message is None: + if hasattr(original_exception, "message"): + message = original_exception.message + else: + message = str(original_exception) + + # Azure OpenAI (especially Images) often nests error details under + # body["error"]. Detect content policy violations using the structured + # payload in addition to string matching. + azure_error_code: Optional[str] = None + try: + body_dict = getattr(original_exception, "body", None) or {} + if isinstance(body_dict, dict): + if isinstance(body_dict.get("error"), dict): + azure_error_code = body_dict["error"].get("code") # type: ignore[index] + # Also check inner_error for + # ResponsibleAIPolicyViolation which indicates a + # content policy violation even when the top-level + # code is generic (e.g. "invalid_request_error"). + if azure_error_code != "content_policy_violation": + _inner = body_dict["error"].get( + "inner_error" + ) or body_dict[ # type: ignore[index] + "error" + ].get( + "innererror" + ) # type: ignore[index] + if ( + isinstance(_inner, dict) + and _inner.get("code") == "ResponsibleAIPolicyViolation" + ): + azure_error_code = "content_policy_violation" + else: + azure_error_code = body_dict.get("code") + except Exception: + azure_error_code = None + + if "Internal server error" in error_str: + raise litellm.InternalServerError( + message=f"AzureException Internal server error - {message}", + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + ) + elif "This model's maximum context length is" in error_str: + raise ContextWindowExceededError( + message=f"AzureException ContextWindowExceededError - {message}", + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + ) + elif "DeploymentNotFound" in error_str: + raise NotFoundError( + message=f"AzureException NotFoundError - {message}", + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + ) + elif ( + azure_error_code == "content_policy_violation" + or ExceptionCheckers.is_azure_content_policy_violation_error(error_str) + ): + from litellm.llms.azure.exception_mapping import ( + AzureOpenAIExceptionMapping, + ) + + raise AzureOpenAIExceptionMapping.create_content_policy_violation_error( + message=message, + model=model, + extra_information=extra_information, + original_exception=original_exception, + ) + elif ( + azure_error_code == "invalid_encrypted_content" + or "could not be verified" in error_str + ): + helpful_message = ( + f"AzureException - {message}\n\n" + "This error occurs when load balancing Responses API across deployments with different API keys.\n" + " Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n" + " Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n" + " router_settings:\n" + " enable_pre_call_checks: true\n" + " optional_pre_call_checks:\n" + " - encrypted_content_affinity\n\n" + " Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing" + ) + raise BadRequestError( + message=helpful_message, + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + body=getattr(original_exception, "body", None), + ) + elif "invalid_request_error" in error_str: + raise BadRequestError( + message=f"AzureException BadRequestError - {message}", + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + body=getattr(original_exception, "body", None), + ) + elif ( + "The api_key client option must be set either by passing api_key to the client or by setting" + in error_str + ): + raise AuthenticationError( + message=f"{exception_provider} AuthenticationError - {message}", + llm_provider=custom_llm_provider, + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + ) + elif "Connection error" in error_str: + raise APIConnectionError( + message=f"{exception_provider} APIConnectionError - {message}", + llm_provider=custom_llm_provider, + model=model, + litellm_debug_info=extra_information, + ) + elif hasattr(original_exception, "status_code"): + if original_exception.status_code == 400: + raise BadRequestError( + message=f"AzureException - {message}", + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + body=getattr(original_exception, "body", None), + ) + elif original_exception.status_code == 401: + raise AuthenticationError( + message=f"AzureException AuthenticationError - {message}", + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"AzureException Timeout - {message}", + model=model, + litellm_debug_info=extra_information, + llm_provider="azure", + ) + elif original_exception.status_code == 422: + raise BadRequestError( + message=f"AzureException BadRequestError - {message}", + model=model, + llm_provider="azure", + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"AzureException RateLimitError - {message}", + model=model, + llm_provider="azure", + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 502: + raise BadGatewayError( + message=f"AzureException BadGatewayError - {message}", + model=model, + llm_provider="azure", + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 503: + raise ServiceUnavailableError( + message=f"AzureException ServiceUnavailableError - {message}", + model=model, + llm_provider="azure", + litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), + ) + elif original_exception.status_code == 504: # gateway timeout error + raise Timeout( + message=f"AzureException Timeout - {message}", + model=model, + litellm_debug_info=extra_information, + llm_provider="azure", + exception_status_code=original_exception.status_code, + ) + else: + raise APIError( + status_code=original_exception.status_code, + message=f"AzureException APIError - {message}", + llm_provider="azure", + litellm_debug_info=extra_information, + model=model, + request=httpx.Request(method="POST", url="https://openai.com/"), + ) + else: + # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors + raise APIConnectionError( + message=f"{exception_provider} APIConnectionError - {message}\n{_redact_string(traceback.format_exc())}", + llm_provider="azure", + model=model, + litellm_debug_info=extra_information, + request=httpx.Request(method="POST", url="https://openai.com/"), + ) + + +def _map_openrouter_exception( + *, + model: str, + original_exception: _ProviderHTTPException, + custom_llm_provider: str, + error_str: str, + exception_type: str, + exception_provider: str, + extra_information: str, +) -> None: + if hasattr(original_exception, "status_code"): + if original_exception.status_code == 400: + raise BadRequestError( + message=f"{exception_provider} - {error_str}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 401: + raise AuthenticationError( + message=f"AuthenticationError: {exception_provider} - {error_str}", + llm_provider=custom_llm_provider, + model=model, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 404: + raise NotFoundError( + message=f"NotFoundError: {exception_provider} - {error_str}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 408: + raise Timeout( + message=f"Timeout Error: {exception_provider} - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 422: + raise BadRequestError( + message=f"BadRequestError: {exception_provider} - {error_str}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 429: + raise RateLimitError( + message=f"RateLimitError: {exception_provider} - {error_str}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 503: + raise ServiceUnavailableError( + message=f"ServiceUnavailableError: {exception_provider} - {error_str}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + elif original_exception.status_code == 504: # gateway timeout error + raise Timeout( + message=f"Timeout Error: {exception_provider} - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + exception_status_code=original_exception.status_code, + ) + else: + raise APIError( + status_code=original_exception.status_code, + message=f"APIError: {exception_provider} - {error_str}", + llm_provider=custom_llm_provider, + model=model, + request=getattr(original_exception, "request", None), + litellm_debug_info=extra_information, + ) + else: + # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors + raise APIConnectionError( + message=f"APIConnectionError: {exception_provider} - {error_str}", + llm_provider=custom_llm_provider, + model=model, + litellm_debug_info=extra_information, + request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), + ) + + def exception_type( # type: ignore model, original_exception, @@ -259,6 +2264,9 @@ def exception_type( # type: ignore return original_exception exception_mapping_worked = False exception_provider = custom_llm_provider + mappable_exception: _ProviderHTTPException = cast( + "_ProviderHTTPException", original_exception + ) if litellm.suppress_debug_info is False: print() # noqa: T201 print( # noqa: T201 @@ -382,2078 +2390,199 @@ def exception_type( # type: ignore or custom_llm_provider in litellm.openai_compatible_providers or custom_llm_provider == "mistral" ): - # custom_llm_provider is openai, make it OpenAI - message = get_error_message(error_obj=original_exception) - if message is None: - if hasattr(original_exception, "message"): - message = original_exception.message - else: - message = str(original_exception) - - if message is not None and isinstance( - message, str - ): # done to prevent user-confusion. Relevant issue - https://github.com/BerriAI/litellm/issues/1414 - message = message.replace("OPENAI", custom_llm_provider.upper()) - message = message.replace( - "openai.OpenAIError", - "{}.{}Error".format(custom_llm_provider, custom_llm_provider), - ) - if custom_llm_provider == "openai": - exception_provider = "OpenAI" + "Exception" - else: - exception_provider = ( - custom_llm_provider[0].upper() - + custom_llm_provider[1:] - + "Exception" - ) - - if ExceptionCheckers.is_error_str_rate_limit(error_str): - exception_mapping_worked = True - raise RateLimitError( - message=f"RateLimitError: {exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - ) - elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"ContextWindowExceededError: {exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif ( - "invalid_request_error" in error_str - and "model_not_found" in error_str - ): - exception_mapping_worked = True - raise NotFoundError( - message=f"{exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif "A timeout occurred" in error_str: - exception_mapping_worked = True - raise Timeout( - message=f"{exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif ( - ( - "invalid_request_error" in error_str - and "content_policy_violation" in error_str - ) - or ( - "Invalid prompt" in error_str - and "violating our usage policy" in error_str - ) - or ( - "request was rejected as a result of the safety system" - in error_str.lower() - ) - ): - exception_mapping_worked = True - raise ContentPolicyViolationError( - message=f"ContentPolicyViolationError: {exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif ( - "invalid_encrypted_content" in error_str - or "could not be verified" in error_str - ): - exception_mapping_worked = True - helpful_message = ( - f"{exception_provider} - {message}\n\n" - " This error occurs when load balancing Responses API across deployments with different API keys.\n" - " Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n" - " Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n" - " router_settings:\n" - " enable_pre_call_checks: true\n" - " optional_pre_call_checks:\n" - " - encrypted_content_affinity\n\n" - " Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing" - ) - raise BadRequestError( - message=helpful_message, - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - body=getattr(original_exception, "body", None), - ) - elif ( - "invalid_request_error" in error_str - and "Incorrect API key provided" not in error_str - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"{exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - body=getattr(original_exception, "body", None), - ) - elif ( - "Web server is returning an unknown error" in error_str - or "The server had an error processing your request." in error_str - ): - exception_mapping_worked = True - raise litellm.InternalServerError( - message=f"{exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - ) - elif "Request too large" in error_str: - exception_mapping_worked = True - raise RateLimitError( - message=f"RateLimitError: {exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif ( - "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable" - in error_str - ): - exception_mapping_worked = True - raise AuthenticationError( - message=f"AuthenticationError: {exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif "Mistral API raised a streaming error" in error_str: - exception_mapping_worked = True - _request = httpx.Request( - method="POST", url="https://api.openai.com/v1" - ) - raise APIError( - status_code=500, - message=f"{exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - request=_request, - litellm_debug_info=extra_information, - ) - elif hasattr(original_exception, "status_code"): - exception_mapping_worked = True - if original_exception.status_code == 400: - exception_mapping_worked = True - raise BadRequestError( - message=f"{exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"AuthenticationError: {exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 404: - exception_mapping_worked = True - raise NotFoundError( - message=f"NotFoundError: {exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"Timeout Error: {exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 422: - exception_mapping_worked = True - raise BadRequestError( - message=f"{exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - body=getattr(original_exception, "body", None), - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"RateLimitError: {exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 500: - exception_mapping_worked = True - raise InternalServerError( - message=f"InternalServerError: {exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 502: - exception_mapping_worked = True - raise BadGatewayError( - message=f"BadGatewayError: {exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 503: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"ServiceUnavailableError: {exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 504: # gateway timeout error - exception_mapping_worked = True - raise Timeout( - message=f"Timeout Error: {exception_provider} - {message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - exception_status_code=original_exception.status_code, - ) - else: - exception_mapping_worked = True - raise APIError( - status_code=original_exception.status_code, - message=f"APIError: {exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - request=getattr(original_exception, "request", None), - litellm_debug_info=extra_information, - ) - else: - # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors - # exception_mapping_worked = True - raise APIConnectionError( - message=f"APIConnectionError: {exception_provider} - {message}", - llm_provider=custom_llm_provider, - model=model, - litellm_debug_info=extra_information, - request=httpx.Request( - method="POST", url="https://api.openai.com/v1/" - ), - ) + _map_openai_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif ( custom_llm_provider == "anthropic" or custom_llm_provider == "anthropic_text" ): # one of the anthropics - if ( - "prompt is too long" in error_str - or "prompt: length" in error_str - or ExceptionCheckers.is_error_str_context_window_exceeded(error_str) - ): - exception_mapping_worked = True - raise ContextWindowExceededError( - message="AnthropicError - {}".format(error_str), - model=model, - llm_provider="anthropic", - ) - elif "overloaded_error" in error_str or "Overloaded" in error_str: - exception_mapping_worked = True - raise InternalServerError( - message="AnthropicError - {}".format(error_str), - model=model, - llm_provider="anthropic", - ) - if "Invalid API Key" in error_str: - exception_mapping_worked = True - raise AuthenticationError( - message="AnthropicError - {}".format(error_str), - model=model, - llm_provider="anthropic", - ) - if "content filtering policy" in error_str: - exception_mapping_worked = True - raise ContentPolicyViolationError( - message="AnthropicError - {}".format(error_str), - model=model, - llm_provider="anthropic", - ) - if "Client error '400 Bad Request'" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message="AnthropicError - {}".format(error_str), - model=model, - llm_provider="anthropic", - ) - if hasattr(original_exception, "status_code"): - verbose_logger.debug( - f"status_code: {original_exception.status_code}" - ) - if original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"AnthropicException - {error_str}", - llm_provider="anthropic", - model=model, - ) - elif ( - original_exception.status_code == 400 - or original_exception.status_code == 413 - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"AnthropicException - {error_str}", - model=model, - llm_provider="anthropic", - ) - elif original_exception.status_code == 404: - exception_mapping_worked = True - raise NotFoundError( - message=f"AnthropicException - {error_str}", - model=model, - llm_provider="anthropic", - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"AnthropicException - {error_str}", - model=model, - llm_provider="anthropic", - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"AnthropicException - {error_str}", - llm_provider="anthropic", - model=model, - ) - elif ( - original_exception.status_code == 500 - or original_exception.status_code == 529 - ): - exception_mapping_worked = True - raise litellm.InternalServerError( - message=f"AnthropicException - {error_str}. Handle with `litellm.InternalServerError`.", - llm_provider="anthropic", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 502: - exception_mapping_worked = True - raise BadGatewayError( - message=f"AnthropicException BadGatewayError - {error_str}", - llm_provider="anthropic", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 503: - exception_mapping_worked = True - raise litellm.ServiceUnavailableError( - message=f"AnthropicException - {error_str}. Handle with `litellm.ServiceUnavailableError`.", - llm_provider="anthropic", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 504: # gateway timeout error - exception_mapping_worked = True - raise Timeout( - message=f"AnthropicException Timeout - {error_str}", - model=model, - llm_provider="anthropic", - exception_status_code=original_exception.status_code, - ) - elif custom_llm_provider == "replicate": - if "Incorrect authentication token" in error_str: - exception_mapping_worked = True - raise AuthenticationError( - message=f"ReplicateException - {error_str}", - llm_provider="replicate", - model=model, - response=getattr(original_exception, "response", None), - ) - elif "input is too long" in error_str: - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"ReplicateException - {error_str}", - model=model, - llm_provider="replicate", - response=getattr(original_exception, "response", None), - ) - elif exception_type == "ModelError": - exception_mapping_worked = True - raise BadRequestError( - message=f"ReplicateException - {error_str}", - model=model, - llm_provider="replicate", - response=getattr(original_exception, "response", None), - ) - elif "Request was throttled" in error_str: - exception_mapping_worked = True - raise RateLimitError( - message=f"ReplicateException - {error_str}", - llm_provider="replicate", - model=model, - response=getattr(original_exception, "response", None), - ) - elif hasattr(original_exception, "status_code"): - if original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"ReplicateException - {original_exception.message}", - llm_provider="replicate", - model=model, - response=getattr(original_exception, "response", None), - ) - elif ( - original_exception.status_code == 400 - or original_exception.status_code == 413 - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"ReplicateException - {original_exception.message}", - model=model, - llm_provider="replicate", - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 422: - exception_mapping_worked = True - raise UnprocessableEntityError( - message=f"ReplicateException - {original_exception.message}", - model=model, - llm_provider="replicate", - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"ReplicateException - {original_exception.message}", - model=model, - llm_provider="replicate", - ) - elif original_exception.status_code == 422: - exception_mapping_worked = True - raise UnprocessableEntityError( - message=f"ReplicateException - {original_exception.message}", - llm_provider="replicate", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"ReplicateException - {original_exception.message}", - llm_provider="replicate", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 500: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"ReplicateException - {original_exception.message}", - llm_provider="replicate", - model=model, - response=getattr(original_exception, "response", None), - ) - exception_mapping_worked = True - raise APIError( - status_code=500, - message=f"ReplicateException - {str(original_exception)}", - llm_provider="replicate", + _map_anthropic_exception( model=model, - request=httpx.Request( - method="POST", - url="https://api.replicate.com/v1/deployments", - ), + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) + elif custom_llm_provider == "replicate": + _map_replicate_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, ) elif custom_llm_provider in litellm._openai_like_providers: - if "authorization denied for" in error_str: - exception_mapping_worked = True - - # Predibase returns the raw API Key in the response - this block ensures it's not returned in the exception - if ( - error_str is not None - and isinstance(error_str, str) - and "bearer" in error_str.lower() - ): - # only keep the first 10 chars after the occurnence of "bearer" - _bearer_token_start_index = error_str.lower().find("bearer") - error_str = error_str[: _bearer_token_start_index + 14] - error_str += "XXXXXXX" + '"' - - raise AuthenticationError( - message=f"{custom_llm_provider.capitalize()}Exception: Authentication Error - {error_str}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"{custom_llm_provider.capitalize()}Exception: Context Window Error - {error_str}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif "token_quota_reached" in error_str: - exception_mapping_worked = True - raise RateLimitError( - message=f"{custom_llm_provider.capitalize()}Exception: Rate Limit Errror - {error_str}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - ) - elif ( - "The server received an invalid response from an upstream server." - in error_str - ): - exception_mapping_worked = True - raise litellm.InternalServerError( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - llm_provider=custom_llm_provider, - model=model, - ) - elif "model_no_support_for_function" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"{custom_llm_provider.capitalize()}Exception - Use 'watsonx_text' route instead. IBM WatsonX does not support `/text/chat` endpoint. - {error_str}", - llm_provider=custom_llm_provider, - model=model, - ) - elif hasattr(original_exception, "status_code"): - if original_exception.status_code == 500: - exception_mapping_worked = True - raise litellm.InternalServerError( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - llm_provider=custom_llm_provider, - model=model, - ) - elif ( - original_exception.status_code == 401 - or original_exception.status_code == 403 - ): - exception_mapping_worked = True - raise AuthenticationError( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - llm_provider=custom_llm_provider, - model=model, - ) - elif original_exception.status_code == 400: - exception_mapping_worked = True - raise BadRequestError( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - llm_provider=custom_llm_provider, - model=model, - ) - elif original_exception.status_code == 404: - exception_mapping_worked = True - raise NotFoundError( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - llm_provider=custom_llm_provider, - model=model, - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif ( - original_exception.status_code == 422 - or original_exception.status_code == 424 - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 503: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 504: # gateway timeout error - exception_mapping_worked = True - raise Timeout( - message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - exception_status_code=original_exception.status_code, - ) + _map_openai_like_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif custom_llm_provider == "bedrock": - if ( - "too many tokens" in error_str - or "expected maxLength:" in error_str - or "Input is too long" in error_str - or "prompt is too long" in error_str - or "prompt: length: 1.." in error_str - or "Too many input tokens" in error_str - ): - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"BedrockException: Context Window Error - {error_str}", - model=model, - llm_provider="bedrock", - ) - elif ( - "Conversation blocks and tool result blocks cannot be provided in the same turn." - in error_str - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"BedrockException - {error_str}\n. Enable 'litellm.modify_params=True' (for PROXY do: `litellm_settings::modify_params: True`) to insert a dummy assistant message and fix this error.", - model=model, - llm_provider="bedrock", - response=getattr(original_exception, "response", None), - ) - elif "Malformed input request" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"BedrockException - {error_str}", - model=model, - llm_provider="bedrock", - response=getattr(original_exception, "response", None), - ) - elif "A conversation must start with a user message." in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"BedrockException - {error_str}\n. Pass in default user message via `completion(..,user_continue_message=)` or enable `litellm.modify_params=True`.\nFor Proxy: do via `litellm_settings::modify_params: True` or user_continue_message under `litellm_params`", - model=model, - llm_provider="bedrock", - response=getattr(original_exception, "response", None), - ) - elif ( - "Unable to locate credentials" in error_str - or "The security token included in the request is invalid" - in error_str - ): - exception_mapping_worked = True - raise AuthenticationError( - message=f"BedrockException Invalid Authentication - {error_str}", - model=model, - llm_provider="bedrock", - response=getattr(original_exception, "response", None), - ) - elif "AccessDeniedException" in error_str: - exception_mapping_worked = True - raise PermissionDeniedError( - message=f"BedrockException PermissionDeniedError - {error_str}", - model=model, - llm_provider="bedrock", - response=getattr(original_exception, "response", None), - ) - elif ( - "throttlingException" in error_str - or "ThrottlingException" in error_str - ): - exception_mapping_worked = True - raise RateLimitError( - message=f"BedrockException: Rate Limit Error - {error_str}", - model=model, - llm_provider="bedrock", - response=getattr(original_exception, "response", None), - ) - elif ( - "Connect timeout on endpoint URL" in error_str - or "timed out" in error_str - ): - exception_mapping_worked = True - raise Timeout( - message=f"BedrockException: Timeout Error - {error_str}", - model=model, - llm_provider="bedrock", - ) - elif "Could not process image" in error_str: - exception_mapping_worked = True - raise litellm.InternalServerError( - message=f"BedrockException - {error_str}", - model=model, - llm_provider="bedrock", - ) - elif hasattr(original_exception, "status_code"): - if original_exception.status_code == 500: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"BedrockException - {original_exception.message}", - llm_provider="bedrock", - model=model, - response=httpx.Response( - status_code=500, - request=httpx.Request( - method="POST", url="https://api.openai.com/v1/" - ), - ), - ) - elif original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"BedrockException - {original_exception.message}", - llm_provider="bedrock", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 400: - exception_mapping_worked = True - raise BadRequestError( - message=f"BedrockException - {original_exception.message}", - llm_provider="bedrock", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 404: - exception_mapping_worked = True - raise NotFoundError( - message=f"BedrockException - {original_exception.message}", - llm_provider="bedrock", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"BedrockException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 422: - exception_mapping_worked = True - raise BadRequestError( - message=f"BedrockException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"BedrockException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 503: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"BedrockException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 504: # gateway timeout error - exception_mapping_worked = True - raise Timeout( - message=f"BedrockException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - exception_status_code=original_exception.status_code, - ) + _map_bedrock_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif ( custom_llm_provider == "sagemaker" or custom_llm_provider == "sagemaker_chat" ): - if "Unable to locate credentials" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"litellm.BadRequestError: SagemakerException - {error_str}", - model=model, - llm_provider="sagemaker", - response=getattr(original_exception, "response", None), - ) - elif ( - "Input validation error: `best_of` must be > 0 and <= 2" - in error_str - ): - exception_mapping_worked = True - raise BadRequestError( - message="SagemakerException - the value of 'n' must be > 0 and <= 2 for sagemaker endpoints", - model=model, - llm_provider="sagemaker", - response=getattr(original_exception, "response", None), - ) - elif ( - "`inputs` tokens + `max_new_tokens` must be <=" in error_str - or "instance type with more CPU capacity or memory" in error_str - ): - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"SagemakerException - {error_str}", - model=model, - llm_provider="sagemaker", - response=getattr(original_exception, "response", None), - ) - elif hasattr(original_exception, "status_code"): - if original_exception.status_code == 500: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"SagemakerException - {original_exception.message}", - llm_provider=custom_llm_provider, - model=model, - response=httpx.Response( - status_code=500, - request=httpx.Request( - method="POST", url="https://api.openai.com/v1/" - ), - ), - ) - elif original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"SagemakerException - {original_exception.message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 400: - exception_mapping_worked = True - raise BadRequestError( - message=f"SagemakerException - {original_exception.message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 404: - exception_mapping_worked = True - raise NotFoundError( - message=f"SagemakerException - {original_exception.message}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"SagemakerException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif ( - original_exception.status_code == 422 - or original_exception.status_code == 424 - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"SagemakerException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"SagemakerException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 503: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"SagemakerException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 504: # gateway timeout error - exception_mapping_worked = True - raise Timeout( - message=f"SagemakerException - {original_exception.message}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - exception_status_code=original_exception.status_code, - ) + _map_sagemaker_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif ( custom_llm_provider == LlmProviders.VERTEX_AI or custom_llm_provider == LlmProviders.VERTEX_AI_BETA or custom_llm_provider == LlmProviders.GEMINI ): - if ( - "Vertex AI API has not been used in project" in error_str - or "Unable to find your project" in error_str - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"litellm.BadRequestError: {custom_llm_provider}Exception - {error_str}", - model=model, - llm_provider=custom_llm_provider, - response=httpx.Response( - status_code=400, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ), - litellm_debug_info=extra_information, - ) - if "400 Request payload size exceeds" in error_str: - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", - model=model, - llm_provider=custom_llm_provider, - ) - elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"ContextWindowExceededError: {custom_llm_provider.capitalize()}Exception - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif ( - "None Unknown Error." in error_str - or "Content has no parts." in error_str - ): - exception_mapping_worked = True - raise litellm.InternalServerError( - message=f"litellm.InternalServerError: {custom_llm_provider}Exception - {error_str}", - model=model, - llm_provider=custom_llm_provider, - response=httpx.Response( - status_code=500, - content=str(original_exception), - request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore - ), - litellm_debug_info=extra_information, - ) - elif "API key not valid." in error_str: - exception_mapping_worked = True - raise AuthenticationError( - message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif "403" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"{custom_llm_provider.capitalize()}Exception BadRequestError - {error_str}", - model=model, - llm_provider=custom_llm_provider, - response=httpx.Response( - status_code=403, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ), - litellm_debug_info=extra_information, - ) - elif ( - "The response was blocked." in error_str - or "Output blocked by content filtering policy" - in error_str # anthropic on vertex ai - ): - exception_mapping_worked = True - raise ContentPolicyViolationError( - message=f"{custom_llm_provider.capitalize()}Exception ContentPolicyViolationError - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - response=httpx.Response( - status_code=400, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ), - ) - elif ( - "429 Quota exceeded" in error_str - or "Quota exceeded for" in error_str - or "Resource exhausted" in error_str - or "IndexError: list index out of range" in error_str - or "429 Unable to submit request because the service is temporarily out of capacity." - in error_str - ): - exception_mapping_worked = True - raise RateLimitError( - message=f"litellm.RateLimitError: {custom_llm_provider}Exception - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - response=httpx.Response( - status_code=429, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ), - ) - elif ( - isinstance(getattr(original_exception, "status_code", None), int) - and 500 <= original_exception.status_code < 600 - and _get_body_error_code(error_str) == 429 - ): - # upstream gateway wraps a 429 inside a 5xx envelope - # e.g. HTTP 500/503 with {"error":{"code":429,...}}. - # Scoped to 5xx so HTTP 400/401 with body code:429 - # still maps to BadRequestError / AuthenticationError. - exception_mapping_worked = True - raise RateLimitError( - message=f"litellm.RateLimitError: {custom_llm_provider}Exception - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - response=httpx.Response( - status_code=429, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ), - ) - elif ( - "500 Internal Server Error" in error_str - or "The model is overloaded." in error_str - ): - exception_mapping_worked = True - raise litellm.InternalServerError( - message=f"litellm.InternalServerError: {custom_llm_provider}Exception - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - if hasattr(original_exception, "status_code"): - if original_exception.status_code == 400: - exception_mapping_worked = True - raise BadRequestError( - message=f"{custom_llm_provider.capitalize()}Exception BadRequestError - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - response=httpx.Response( - status_code=400, - request=httpx.Request( - method="POST", - url="https://cloud.google.com/vertex-ai/", - ), - ), - ) - if original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", - llm_provider=custom_llm_provider, - model=model, - ) - if original_exception.status_code == 403: - exception_mapping_worked = True - raise PermissionDeniedError( - message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", - llm_provider=custom_llm_provider, - model=model, - response=httpx.Response( - status_code=403, - request=httpx.Request( - method="POST", - url="https://cloud.google.com/vertex-ai/", - ), - ), - ) - if original_exception.status_code == 404: - exception_mapping_worked = True - raise NotFoundError( - message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", - llm_provider=custom_llm_provider, - model=model, - ) - if original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", - llm_provider=custom_llm_provider, - model=model, - ) - - if original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"litellm.RateLimitError: {custom_llm_provider.capitalize()}Exception - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - response=httpx.Response( - status_code=429, - request=httpx.Request( - method="POST", - url=" https://cloud.google.com/vertex-ai/", - ), - ), - ) - if original_exception.status_code == 500: - exception_mapping_worked = True - raise litellm.InternalServerError( - message=f"{custom_llm_provider.capitalize()}Exception InternalServerError - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - response=httpx.Response( - status_code=500, - content=str(original_exception), - request=httpx.Request(method="completion", url="https://github.com/BerriAI/litellm"), # type: ignore - ), - ) - if original_exception.status_code == 502: - exception_mapping_worked = True - raise APIConnectionError( - message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", - llm_provider=custom_llm_provider, - model=model, - ) - if original_exception.status_code == 503: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"{custom_llm_provider.capitalize()}Exception - {error_str}", - llm_provider=custom_llm_provider, - model=model, - ) + _map_vertex_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif custom_llm_provider == "cloudflare": - if "Authentication error" in error_str: - exception_mapping_worked = True - raise AuthenticationError( - message=f"Cloudflare Exception - {original_exception.message}", - llm_provider="cloudflare", - model=model, - response=getattr(original_exception, "response", None), - ) - if "must have required property" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"Cloudflare Exception - {original_exception.message}", - llm_provider="cloudflare", - model=model, - response=getattr(original_exception, "response", None), - ) + _map_cloudflare_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif ( custom_llm_provider == "cohere" or custom_llm_provider == "cohere_chat" ): # Cohere - if ( - "invalid api token" in error_str - or "No API key provided." in error_str - ): - exception_mapping_worked = True - raise AuthenticationError( - message=f"CohereException - {original_exception.message}", - llm_provider="cohere", - model=model, - response=getattr(original_exception, "response", None), - ) - elif "invalid type: parameter" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"CohereException - {original_exception.message}", - llm_provider="cohere", - model=model, - response=getattr(original_exception, "response", None), - ) - elif "too many tokens" in error_str: - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"CohereException - {original_exception.message}", - model=model, - llm_provider="cohere", - response=getattr(original_exception, "response", None), - ) - elif "internal server error" in error_str.lower(): - exception_mapping_worked = True - raise InternalServerError( - message=f"CohereException - {error_str}", - model=model, - llm_provider="cohere", - response=getattr(original_exception, "response", None), - ) - elif hasattr(original_exception, "status_code"): - if ( - original_exception.status_code == 400 - or original_exception.status_code == 498 - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"CohereException - {original_exception.message}", - llm_provider="cohere", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"CohereException - {original_exception.message}", - llm_provider="cohere", - model=model, - ) - elif original_exception.status_code == 500: - exception_mapping_worked = True - raise InternalServerError( - message=f"CohereException - {original_exception.message}", - llm_provider="cohere", - model=model, - response=getattr(original_exception, "response", None), - ) - elif ( - "CohereConnectionError" in exception_type - ): # cohere seems to fire these errors when we load test it (1k+ messages / min) - exception_mapping_worked = True - raise RateLimitError( - message=f"CohereException - {original_exception.message}", - llm_provider="cohere", - model=model, - response=getattr(original_exception, "response", None), - ) - elif "invalid type:" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"CohereException - {original_exception.message}", - llm_provider="cohere", - model=model, - response=getattr(original_exception, "response", None), - ) - elif "Unexpected server error" in error_str: - exception_mapping_worked = True - raise InternalServerError( - message=f"CohereException - {original_exception.message}", - llm_provider="cohere", - model=model, - response=getattr(original_exception, "response", None), - ) - else: - if hasattr(original_exception, "status_code"): - exception_mapping_worked = True - raise APIError( - status_code=original_exception.status_code, - message=f"CohereException - {original_exception.message}", - llm_provider="cohere", - model=model, - request=getattr(original_exception, "request", None), - ) - raise original_exception + _map_cohere_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif custom_llm_provider == "huggingface": - if "length limit exceeded" in error_str: - exception_mapping_worked = True - raise ContextWindowExceededError( - message=error_str, - model=model, - llm_provider="huggingface", - response=getattr(original_exception, "response", None), - ) - elif "A valid user token is required" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=error_str, - llm_provider="huggingface", - model=model, - response=getattr(original_exception, "response", None), - ) - elif "Rate limit reached" in error_str: - exception_mapping_worked = True - raise RateLimitError( - message=error_str, - llm_provider="huggingface", - model=model, - response=getattr(original_exception, "response", None), - ) - if hasattr(original_exception, "status_code"): - if original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"HuggingfaceException - {original_exception.message}", - llm_provider="huggingface", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 400: - exception_mapping_worked = True - raise BadRequestError( - message=f"HuggingfaceException - {original_exception.message}", - model=model, - llm_provider="huggingface", - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"HuggingfaceException - {original_exception.message}", - model=model, - llm_provider="huggingface", - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"HuggingfaceException - {original_exception.message}", - llm_provider="huggingface", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 503: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"HuggingfaceException - {original_exception.message}", - llm_provider="huggingface", - model=model, - response=getattr(original_exception, "response", None), - ) - else: - exception_mapping_worked = True - raise APIError( - status_code=original_exception.status_code, - message=f"HuggingfaceException - {original_exception.message}", - llm_provider="huggingface", - model=model, - request=getattr(original_exception, "request", None), - ) + _map_huggingface_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif custom_llm_provider == "ai21": - if hasattr(original_exception, "message"): - if "Prompt has too many tokens" in original_exception.message: - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"AI21Exception - {original_exception.message}", - model=model, - llm_provider="ai21", - response=getattr(original_exception, "response", None), - ) - if "Bad or missing API token." in original_exception.message: - exception_mapping_worked = True - raise BadRequestError( - message=f"AI21Exception - {original_exception.message}", - model=model, - llm_provider="ai21", - response=getattr(original_exception, "response", None), - ) - if hasattr(original_exception, "status_code"): - if original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"AI21Exception - {original_exception.message}", - llm_provider="ai21", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"AI21Exception - {original_exception.message}", - model=model, - llm_provider="ai21", - ) - if original_exception.status_code == 422: - exception_mapping_worked = True - raise BadRequestError( - message=f"AI21Exception - {original_exception.message}", - model=model, - llm_provider="ai21", - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"AI21Exception - {original_exception.message}", - llm_provider="ai21", - model=model, - response=getattr(original_exception, "response", None), - ) - else: - exception_mapping_worked = True - raise APIError( - status_code=original_exception.status_code, - message=f"AI21Exception - {original_exception.message}", - llm_provider="ai21", - model=model, - request=getattr(original_exception, "request", None), - ) + _map_ai21_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif custom_llm_provider == "nlp_cloud": - if "detail" in error_str: - if "Input text length should not exceed" in error_str: - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"NLPCloudException - {error_str}", - model=model, - llm_provider="nlp_cloud", - response=getattr(original_exception, "response", None), - ) - elif "value is not a valid" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"NLPCloudException - {error_str}", - model=model, - llm_provider="nlp_cloud", - response=getattr(original_exception, "response", None), - ) - else: - exception_mapping_worked = True - raise APIError( - status_code=500, - message=f"NLPCloudException - {error_str}", - model=model, - llm_provider="nlp_cloud", - request=getattr(original_exception, "request", None), - ) - if hasattr( - original_exception, "status_code" - ): # https://docs.nlpcloud.com/?shell#errors - if ( - original_exception.status_code == 400 - or original_exception.status_code == 406 - or original_exception.status_code == 413 - or original_exception.status_code == 422 - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"NLPCloudException - {original_exception.message}", - llm_provider="nlp_cloud", - model=model, - response=getattr(original_exception, "response", None), - ) - elif ( - original_exception.status_code == 401 - or original_exception.status_code == 403 - ): - exception_mapping_worked = True - raise AuthenticationError( - message=f"NLPCloudException - {original_exception.message}", - llm_provider="nlp_cloud", - model=model, - response=getattr(original_exception, "response", None), - ) - elif ( - original_exception.status_code == 522 - or original_exception.status_code == 524 - ): - exception_mapping_worked = True - raise Timeout( - message=f"NLPCloudException - {original_exception.message}", - model=model, - llm_provider="nlp_cloud", - ) - elif ( - original_exception.status_code == 429 - or original_exception.status_code == 402 - ): - exception_mapping_worked = True - raise RateLimitError( - message=f"NLPCloudException - {original_exception.message}", - llm_provider="nlp_cloud", - model=model, - response=getattr(original_exception, "response", None), - ) - elif ( - original_exception.status_code == 500 - or original_exception.status_code == 503 - ): - exception_mapping_worked = True - raise APIError( - status_code=original_exception.status_code, - message=f"NLPCloudException - {original_exception.message}", - llm_provider="nlp_cloud", - model=model, - request=getattr(original_exception, "request", None), - ) - elif ( - original_exception.status_code == 504 - or original_exception.status_code == 520 - ): - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"NLPCloudException - {original_exception.message}", - model=model, - llm_provider="nlp_cloud", - response=getattr(original_exception, "response", None), - ) - else: - exception_mapping_worked = True - raise APIError( - status_code=original_exception.status_code, - message=f"NLPCloudException - {original_exception.message}", - llm_provider="nlp_cloud", - model=model, - request=getattr(original_exception, "request", None), - ) + _map_nlp_cloud_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif custom_llm_provider == "together_ai": - try: - error_response = json.loads(error_str) - except Exception: - error_response = {"error": error_str} - if ( - "error" in error_response - and "`inputs` tokens + `max_new_tokens` must be <=" - in error_response["error"] - ): - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"TogetherAIException - {error_response['error']}", - model=model, - llm_provider="together_ai", - response=getattr(original_exception, "response", None), - ) - elif ( - "error" in error_response - and "invalid private key" in error_response["error"] - ): - exception_mapping_worked = True - raise AuthenticationError( - message=f"TogetherAIException - {error_response['error']}", - llm_provider="together_ai", - model=model, - response=getattr(original_exception, "response", None), - ) - elif ( - "error" in error_response - and "INVALID_ARGUMENT" in error_response["error"] - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"TogetherAIException - {error_response['error']}", - model=model, - llm_provider="together_ai", - response=getattr(original_exception, "response", None), - ) - elif "A timeout occurred" in error_str: - exception_mapping_worked = True - raise Timeout( - message=f"TogetherAIException - {error_str}", - model=model, - llm_provider="together_ai", - ) - elif ( - "error" in error_response - and "API key doesn't match expected format." - in error_response["error"] - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"TogetherAIException - {error_response['error']}", - model=model, - llm_provider="together_ai", - response=getattr(original_exception, "response", None), - ) - elif ( - "error_type" in error_response - and error_response["error_type"] == "validation" - ): - exception_mapping_worked = True - raise BadRequestError( - message=f"TogetherAIException - {error_response['error']}", - model=model, - llm_provider="together_ai", - response=getattr(original_exception, "response", None), - ) - if hasattr(original_exception, "status_code"): - if original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"TogetherAIException - {original_exception.message}", - model=model, - llm_provider="together_ai", - ) - elif original_exception.status_code == 422: - exception_mapping_worked = True - raise BadRequestError( - message=f"TogetherAIException - {error_response['error']}", - model=model, - llm_provider="together_ai", - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"TogetherAIException - {original_exception.message}", - llm_provider="together_ai", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 524: - exception_mapping_worked = True - raise Timeout( - message=f"TogetherAIException - {original_exception.message}", - llm_provider="together_ai", - model=model, - ) - else: - exception_mapping_worked = True - raise APIError( - status_code=original_exception.status_code, - message=f"TogetherAIException - {original_exception.message}", - llm_provider="together_ai", - model=model, - request=getattr(original_exception, "request", None), - ) + _map_together_ai_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif custom_llm_provider == "aleph_alpha": - if ( - "This is longer than the model's maximum context length" - in error_str - ): - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"AlephAlphaException - {original_exception.message}", - llm_provider="aleph_alpha", - model=model, - response=getattr(original_exception, "response", None), - ) - elif "InvalidToken" in error_str or "No token provided" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"AlephAlphaException - {original_exception.message}", - llm_provider="aleph_alpha", - model=model, - response=getattr(original_exception, "response", None), - ) - elif hasattr(original_exception, "status_code"): - verbose_logger.debug( - f"status code: {original_exception.status_code}" - ) - if original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"AlephAlphaException - {original_exception.message}", - llm_provider="aleph_alpha", - model=model, - ) - elif original_exception.status_code == 400: - exception_mapping_worked = True - raise BadRequestError( - message=f"AlephAlphaException - {original_exception.message}", - llm_provider="aleph_alpha", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"AlephAlphaException - {original_exception.message}", - llm_provider="aleph_alpha", - model=model, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 500: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"AlephAlphaException - {original_exception.message}", - llm_provider="aleph_alpha", - model=model, - response=getattr(original_exception, "response", None), - ) - raise original_exception - raise original_exception + _map_aleph_alpha_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif ( custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat" ): - if isinstance(original_exception, dict): - error_str = original_exception.get("error", "") - else: - error_str = str(original_exception) - if "no such file or directory" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"OllamaException: Invalid Model/Model not loaded - {original_exception}", - model=model, - llm_provider="ollama", - response=getattr(original_exception, "response", None), - ) - elif "Failed to establish a new connection" in error_str: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"OllamaException: {original_exception}", - llm_provider="ollama", - model=model, - response=getattr(original_exception, "response", None), - ) - elif "Invalid response object from API" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"OllamaException: {original_exception}", - llm_provider="ollama", - model=model, - response=getattr(original_exception, "response", None), - ) - elif "Read timed out" in error_str: - exception_mapping_worked = True - raise Timeout( - message=f"OllamaException: {original_exception}", - llm_provider="ollama", - model=model, - ) + _map_ollama_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif custom_llm_provider == "vllm": - if hasattr(original_exception, "status_code"): - if original_exception.status_code == 0: - exception_mapping_worked = True - raise APIConnectionError( - message=f"VLLMException - {original_exception.message}", - llm_provider="vllm", - model=model, - request=getattr(original_exception, "request", None), - ) + _map_vllm_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) elif custom_llm_provider == "azure" or custom_llm_provider == "azure_text": - message = get_error_message(error_obj=original_exception) - if message is None: - if hasattr(original_exception, "message"): - message = original_exception.message - else: - message = str(original_exception) - - # Azure OpenAI (especially Images) often nests error details under - # body["error"]. Detect content policy violations using the structured - # payload in addition to string matching. - azure_error_code: Optional[str] = None - try: - body_dict = getattr(original_exception, "body", None) or {} - if isinstance(body_dict, dict): - if isinstance(body_dict.get("error"), dict): - azure_error_code = body_dict["error"].get("code") # type: ignore[index] - # Also check inner_error for - # ResponsibleAIPolicyViolation which indicates a - # content policy violation even when the top-level - # code is generic (e.g. "invalid_request_error"). - if azure_error_code != "content_policy_violation": - _inner = body_dict["error"].get( - "inner_error" - ) or body_dict[ # type: ignore[index] - "error" - ].get( - "innererror" - ) # type: ignore[index] - if ( - isinstance(_inner, dict) - and _inner.get("code") - == "ResponsibleAIPolicyViolation" - ): - azure_error_code = "content_policy_violation" - else: - azure_error_code = body_dict.get("code") - except Exception: - azure_error_code = None - - if "Internal server error" in error_str: - exception_mapping_worked = True - raise litellm.InternalServerError( - message=f"AzureException Internal server error - {message}", - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) - elif "This model's maximum context length is" in error_str: - exception_mapping_worked = True - raise ContextWindowExceededError( - message=f"AzureException ContextWindowExceededError - {message}", - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) - elif "DeploymentNotFound" in error_str: - exception_mapping_worked = True - raise NotFoundError( - message=f"AzureException NotFoundError - {message}", - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) - elif ( - azure_error_code == "content_policy_violation" - or ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) - ): - exception_mapping_worked = True - from litellm.llms.azure.exception_mapping import ( - AzureOpenAIExceptionMapping, - ) - - raise AzureOpenAIExceptionMapping.create_content_policy_violation_error( - message=message, - model=model, - extra_information=extra_information, - original_exception=original_exception, - ) - elif ( - azure_error_code == "invalid_encrypted_content" - or "could not be verified" in error_str - ): - exception_mapping_worked = True - helpful_message = ( - f"AzureException - {message}\n\n" - "This error occurs when load balancing Responses API across deployments with different API keys.\n" - " Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n" - " Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n" - " router_settings:\n" - " enable_pre_call_checks: true\n" - " optional_pre_call_checks:\n" - " - encrypted_content_affinity\n\n" - " Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing" - ) - raise BadRequestError( - message=helpful_message, - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - body=getattr(original_exception, "body", None), - ) - elif "invalid_request_error" in error_str: - exception_mapping_worked = True - raise BadRequestError( - message=f"AzureException BadRequestError - {message}", - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - body=getattr(original_exception, "body", None), - ) - elif ( - "The api_key client option must be set either by passing api_key to the client or by setting" - in error_str - ): - exception_mapping_worked = True - raise AuthenticationError( - message=f"{exception_provider} AuthenticationError - {message}", - llm_provider=custom_llm_provider, - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) - elif "Connection error" in error_str: - exception_mapping_worked = True - raise APIConnectionError( - message=f"{exception_provider} APIConnectionError - {message}", - llm_provider=custom_llm_provider, - model=model, - litellm_debug_info=extra_information, - ) - elif hasattr(original_exception, "status_code"): - exception_mapping_worked = True - if original_exception.status_code == 400: - exception_mapping_worked = True - raise BadRequestError( - message=f"AzureException - {message}", - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - body=getattr(original_exception, "body", None), - ) - elif original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"AzureException AuthenticationError - {message}", - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"AzureException Timeout - {message}", - model=model, - litellm_debug_info=extra_information, - llm_provider="azure", - ) - elif original_exception.status_code == 422: - exception_mapping_worked = True - raise BadRequestError( - message=f"AzureException BadRequestError - {message}", - model=model, - llm_provider="azure", - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"AzureException RateLimitError - {message}", - model=model, - llm_provider="azure", - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 502: - exception_mapping_worked = True - raise BadGatewayError( - message=f"AzureException BadGatewayError - {message}", - model=model, - llm_provider="azure", - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 503: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"AzureException ServiceUnavailableError - {message}", - model=model, - llm_provider="azure", - litellm_debug_info=extra_information, - response=getattr(original_exception, "response", None), - ) - elif original_exception.status_code == 504: # gateway timeout error - exception_mapping_worked = True - raise Timeout( - message=f"AzureException Timeout - {message}", - model=model, - litellm_debug_info=extra_information, - llm_provider="azure", - exception_status_code=original_exception.status_code, - ) - else: - exception_mapping_worked = True - raise APIError( - status_code=original_exception.status_code, - message=f"AzureException APIError - {message}", - llm_provider="azure", - litellm_debug_info=extra_information, - model=model, - request=httpx.Request( - method="POST", url="https://openai.com/" - ), - ) - else: - # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors - raise APIConnectionError( - message=f"{exception_provider} APIConnectionError - {message}\n{_redact_string(traceback.format_exc())}", - llm_provider="azure", - model=model, - litellm_debug_info=extra_information, - request=httpx.Request(method="POST", url="https://openai.com/"), - ) + _map_azure_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) if custom_llm_provider == "openrouter": - if hasattr(original_exception, "status_code"): - exception_mapping_worked = True - if original_exception.status_code == 400: - exception_mapping_worked = True - raise BadRequestError( - message=f"{exception_provider} - {error_str}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 401: - exception_mapping_worked = True - raise AuthenticationError( - message=f"AuthenticationError: {exception_provider} - {error_str}", - llm_provider=custom_llm_provider, - model=model, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 404: - exception_mapping_worked = True - raise NotFoundError( - message=f"NotFoundError: {exception_provider} - {error_str}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 408: - exception_mapping_worked = True - raise Timeout( - message=f"Timeout Error: {exception_provider} - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 422: - exception_mapping_worked = True - raise BadRequestError( - message=f"BadRequestError: {exception_provider} - {error_str}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 429: - exception_mapping_worked = True - raise RateLimitError( - message=f"RateLimitError: {exception_provider} - {error_str}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 503: - exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"ServiceUnavailableError: {exception_provider} - {error_str}", - model=model, - llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), - litellm_debug_info=extra_information, - ) - elif original_exception.status_code == 504: # gateway timeout error - exception_mapping_worked = True - raise Timeout( - message=f"Timeout Error: {exception_provider} - {error_str}", - model=model, - llm_provider=custom_llm_provider, - litellm_debug_info=extra_information, - exception_status_code=original_exception.status_code, - ) - else: - exception_mapping_worked = True - raise APIError( - status_code=original_exception.status_code, - message=f"APIError: {exception_provider} - {error_str}", - llm_provider=custom_llm_provider, - model=model, - request=getattr(original_exception, "request", None), - litellm_debug_info=extra_information, - ) - else: - # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors - raise APIConnectionError( - message=f"APIConnectionError: {exception_provider} - {error_str}", - llm_provider=custom_llm_provider, - model=model, - litellm_debug_info=extra_information, - request=httpx.Request( - method="POST", url="https://api.openai.com/v1/" - ), - ) + _map_openrouter_exception( + model=model, + original_exception=mappable_exception, + custom_llm_provider=custom_llm_provider, + error_str=error_str, + exception_type=exception_type, + exception_provider=exception_provider, + extra_information=extra_information, + ) if ( "BadRequestError.__init__() missing 1 required positional argument: 'param'" in str(original_exception) diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index e87042b9101..c22d3b99705 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -86,9 +86,7 @@ def get_supported_openai_params( model=model ) elif request_type == "transcription": - return litellm.FireworksAIAudioTranscriptionConfig().get_supported_openai_params( - model=model - ) + return None else: return litellm.FireworksAIConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "nvidia_nim": @@ -191,7 +189,9 @@ def get_supported_openai_params( ) elif custom_llm_provider == "sambanova": if request_type == "embeddings": - litellm.SambaNovaEmbeddingConfig().get_supported_openai_params(model=model) + return litellm.SambaNovaEmbeddingConfig().get_supported_openai_params( + model=model + ) else: return litellm.SambanovaConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "nebius": diff --git a/litellm/litellm_core_utils/request_timeout_resolver.py b/litellm/litellm_core_utils/request_timeout_resolver.py new file mode 100644 index 00000000000..146c39ce9f3 --- /dev/null +++ b/litellm/litellm_core_utils/request_timeout_resolver.py @@ -0,0 +1,29 @@ +"""Single source of truth for whether ``litellm.request_timeout`` was configured. + +``litellm.request_timeout`` always holds a value (the package default, +:data:`~litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS`), so a bare read can't +tell "user asked for this" from "nobody set it". This resolver answers that: + +* ``request_timeout_explicitly_set`` is the authoritative signal, set when the + value comes from the ``REQUEST_TIMEOUT`` env var or ``litellm_settings``. +* A runtime value that differs from the package default (e.g. ``litellm.request_timeout + = 300`` in SDK code) is also treated as explicit, for backwards compatibility. +""" + +from __future__ import annotations + +from typing import Optional + +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS + + +def get_configured_request_timeout() -> Optional[float]: + """Return the explicitly-configured ``litellm.request_timeout``, else ``None``.""" + import litellm + + timeout = float(litellm.request_timeout) + if litellm.request_timeout_explicitly_set: + return timeout + if timeout != float(DEFAULT_REQUEST_TIMEOUT_SECONDS): + return timeout + return None diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 4928dd08386..b14e12de7cd 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -12,6 +12,7 @@ class SensitiveDataMasker: visible_prefix: int = 4, visible_suffix: int = 4, mask_char: str = "*", + mask_short_values: bool = True, ): self.sensitive_patterns = sensitive_patterns or { "password", @@ -38,12 +39,17 @@ class SensitiveDataMasker: self.visible_prefix = visible_prefix self.visible_suffix = visible_suffix self.mask_char = mask_char + self.mask_short_values = mask_short_values def _mask_value(self, value: str) -> str: - if not value or len(str(value)) < (self.visible_prefix + self.visible_suffix): - return value - value_str = str(value) + if not value_str: + return value + if len(value_str) <= (self.visible_prefix + self.visible_suffix): + return ( + self.mask_char * len(value_str) if self.mask_short_values else value_str + ) + masked_length = len(value_str) - (self.visible_prefix + self.visible_suffix) # Handle the case where visible_suffix is 0 to avoid showing the entire string diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index d3330c3dcec..e278483d689 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -6,6 +6,7 @@ import logging import threading import time import traceback +from dataclasses import dataclass from typing import ( Any, AsyncIterator, @@ -97,6 +98,19 @@ def print_verbose(print_statement): pass +@dataclass(frozen=True, slots=True) +class _ProviderChunkParsed: + response_obj: dict[str, Any] + + +@dataclass(frozen=True, slots=True) +class _ProviderChunkEarlyReturn: + value: Any + + +_ProviderChunkResult = Union[_ProviderChunkParsed, _ProviderChunkEarlyReturn] + + class CustomStreamWrapper: def __init__( self, @@ -1145,381 +1159,392 @@ class CustomStreamWrapper: del model_response.choices[0].delta.reasoning_content return + def _dispatch_provider_chunk( + self, + chunk: Any, + model_response: ModelResponseStream, + completion_obj: dict[str, Any], + ) -> _ProviderChunkResult: + response_obj: dict[str, Any] = {} + if ( + isinstance(chunk, ModelResponseStream) + and self.custom_llm_provider is not None + and self.custom_llm_provider in litellm._custom_providers + ): + _has_content = bool( + chunk.choices + and chunk.choices[0].delta is not None + and ( + chunk.choices[0].delta.content or chunk.choices[0].delta.tool_calls + ) + ) + if self.received_finish_reason is not None: + if not _has_content: + raise StopIteration + if chunk.choices and chunk.choices[0].finish_reason: + self.received_finish_reason = chunk.choices[0].finish_reason + if not _has_content: + return _ProviderChunkEarlyReturn(None) + # Strip finish_reason from the content chunk so it appears + # only on the trailing empty-delta chunk (OpenAI spec). + # finish_reason_handler() will emit the proper terminal chunk. + chunk.choices[0].finish_reason = None # type: ignore[assignment] + return _ProviderChunkEarlyReturn(chunk) + + if ( + isinstance(chunk, dict) + and generic_chunk_has_all_required_fields( + chunk=chunk + ) # check if chunk is a generic streaming chunk + ) or ( + self.custom_llm_provider + and self.custom_llm_provider in litellm._custom_providers + ): + if self.received_finish_reason is not None: + _chunk_has_content = isinstance(chunk, dict) and ( + bool(chunk.get("text", "")) + or chunk.get("tool_use") is not None + # Usage-only final chunks are valid and needed to surface + # finish_reason/usage to downstream translators. + or chunk.get("usage") is not None + ) + if not _chunk_has_content and ( + not isinstance(chunk, dict) + or "provider_specific_fields" not in chunk + ): + raise StopIteration + anthropic_response_obj: GChunk = cast(GChunk, chunk) + completion_obj["content"] = anthropic_response_obj["text"] + if anthropic_response_obj["is_finished"]: + self.received_finish_reason = anthropic_response_obj["finish_reason"] + + if anthropic_response_obj["finish_reason"]: + self.intermittent_finish_reason = anthropic_response_obj[ + "finish_reason" + ] + + if anthropic_response_obj["usage"] is not None: + setattr( + model_response, + "usage", + litellm.Usage(**anthropic_response_obj["usage"]), + ) + + if ( + "tool_use" in anthropic_response_obj + and anthropic_response_obj["tool_use"] is not None + ): + completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]] + + if ( + "provider_specific_fields" in anthropic_response_obj + and anthropic_response_obj["provider_specific_fields"] is not None + ): + for key, value in anthropic_response_obj[ + "provider_specific_fields" + ].items(): + setattr(model_response, key, value) + + response_obj = cast(dict[str, Any], anthropic_response_obj) + elif self.model == "replicate" or self.custom_llm_provider == "replicate": + response_obj = self.handle_replicate_chunk(chunk) + completion_obj["content"] = response_obj["text"] + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + elif self.custom_llm_provider and self.custom_llm_provider == "predibase": + response_obj = self.handle_predibase_chunk(chunk) + completion_obj["content"] = response_obj["text"] + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + elif ( + self.custom_llm_provider and self.custom_llm_provider == "baseten" + ): # baseten doesn't provide streaming + completion_obj["content"] = self.handle_baseten_chunk(chunk) + elif ( + self.custom_llm_provider and self.custom_llm_provider == "ai21" + ): # ai21 doesn't provide streaming + response_obj = self.handle_ai21_chunk(chunk) + completion_obj["content"] = response_obj["text"] + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + elif self.custom_llm_provider and self.custom_llm_provider == "maritalk": + response_obj = self.handle_maritalk_chunk(chunk) + completion_obj["content"] = response_obj["text"] + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + elif self.custom_llm_provider and self.custom_llm_provider == "vllm": + completion_obj["content"] = chunk[0].outputs[0].text + elif ( + self.custom_llm_provider and self.custom_llm_provider == "aleph_alpha" + ): # aleph alpha doesn't provide streaming + response_obj = self.handle_aleph_alpha_chunk(chunk) + completion_obj["content"] = response_obj["text"] + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + elif self.custom_llm_provider == "nlp_cloud": + try: + response_obj = self.handle_nlp_cloud_chunk(chunk) + completion_obj["content"] = response_obj["text"] + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + except Exception as e: + if self.received_finish_reason: + raise e + else: + if self.sent_first_chunk is False: + raise Exception("An unknown error occurred with the stream") + self.received_finish_reason = "stop" + elif self.custom_llm_provider == "vertex_ai" and not isinstance( + chunk, ModelResponseStream + ): + chunk = cast(Any, chunk) + import proto # type: ignore + + if hasattr(chunk, "candidates") is True: + try: + try: + completion_obj["content"] = chunk.text # type: ignore + except Exception as e: + original_exception = e + if "Part has no text." in str(e): + ## check for function calling + function_call = ( + chunk.candidates[0].content.parts[0].function_call # type: ignore + ) + + args_dict = {} + + # Check if it's a RepeatedComposite instance + for key, val in function_call.args.items(): + if isinstance( + val, + proto.marshal.collections.repeated.RepeatedComposite, # type: ignore + ): + # If so, convert to list + args_dict[key] = [v for v in val] + else: + args_dict[key] = val + + try: + args_str = json.dumps(args_dict) + except Exception as e: + raise e + _delta_obj = litellm.utils.Delta( + content=None, + tool_calls=[ + { + "id": f"call_{str(uuid.uuid4())}", + "function": { + "arguments": args_str, + "name": function_call.name, + }, + "type": "function", + } + ], + ) + _streaming_response = StreamingChoices(delta=_delta_obj) + _model_response = ModelResponseStream() + _model_response.choices = [_streaming_response] + response_obj = {"original_chunk": _model_response} + else: + raise original_exception + if ( + hasattr(chunk.candidates[0], "finish_reason") # type: ignore + and chunk.candidates[0].finish_reason.name # type: ignore + != "FINISH_REASON_UNSPECIFIED" + ): # every non-final chunk in vertex ai has this + self.received_finish_reason = map_finish_reason( # type: ignore + chunk.candidates[0].finish_reason.name + ) + except Exception: + if chunk.candidates[0].finish_reason.name == "SAFETY": # type: ignore + raise Exception( + f"The response was blocked by VertexAI. {str(chunk)}" + ) + else: + completion_obj["content"] = str(chunk) + elif self.custom_llm_provider == "petals": + if self.completion_stream is None or len(self.completion_stream) == 0: + if self.received_finish_reason is not None: + raise StopIteration + else: + self.received_finish_reason = "stop" + chunk_size = 30 + stream = cast(Any, self.completion_stream) + new_chunk = stream[:chunk_size] + completion_obj["content"] = new_chunk + self.completion_stream = stream[chunk_size:] + elif self.custom_llm_provider == "palm": + # fake streaming + response_obj = {} + if self.completion_stream is None or len(self.completion_stream) == 0: + if self.received_finish_reason is not None: + raise StopIteration + else: + self.received_finish_reason = "stop" + chunk_size = 30 + stream = cast(Any, self.completion_stream) + new_chunk = stream[:chunk_size] + completion_obj["content"] = new_chunk + self.completion_stream = stream[chunk_size:] + elif self.custom_llm_provider == "triton": + response_obj = self.handle_triton_stream(chunk) + completion_obj["content"] = response_obj["text"] + print_verbose(f"completion obj content: {completion_obj['content']}") + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + elif self.custom_llm_provider == "text-completion-openai": + response_obj = self.handle_openai_text_completion_chunk(chunk) + completion_obj["content"] = response_obj["text"] + print_verbose(f"completion obj content: {completion_obj['content']}") + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + if response_obj["usage"] is not None: + setattr( + model_response, + "usage", + litellm.Usage( + prompt_tokens=response_obj["usage"].prompt_tokens, + completion_tokens=response_obj["usage"].completion_tokens, + total_tokens=response_obj["usage"].total_tokens, + ), + ) + elif self.custom_llm_provider == "text-completion-codestral": + if not isinstance(chunk, str): + raise ValueError(f"chunk is not a string: {chunk}") + response_obj = cast( + dict[str, Any], + litellm.CodestralTextCompletionConfig()._chunk_parser(chunk), + ) + completion_obj["content"] = response_obj["text"] + print_verbose(f"completion obj content: {completion_obj['content']}") + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + if "usage" in response_obj is not None: + setattr( + model_response, + "usage", + litellm.Usage( + prompt_tokens=response_obj["usage"].prompt_tokens, + completion_tokens=response_obj["usage"].completion_tokens, + total_tokens=response_obj["usage"].total_tokens, + ), + ) + elif self.custom_llm_provider == "azure_text": + response_obj = self.handle_azure_text_completion_chunk(chunk) + completion_obj["content"] = response_obj["text"] + print_verbose(f"completion obj content: {completion_obj['content']}") + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + elif self.custom_llm_provider == "cached_response": + chunk = cast(ModelResponseStream, chunk) + response_obj = { + "text": chunk.choices[0].delta.content, + "is_finished": True, + "finish_reason": chunk.choices[0].finish_reason, + "original_chunk": chunk, + "tool_calls": ( + chunk.choices[0].delta.tool_calls + if hasattr(chunk.choices[0].delta, "tool_calls") + else None + ), + } + + completion_obj["content"] = response_obj["text"] + if response_obj["tool_calls"] is not None: + completion_obj["tool_calls"] = response_obj["tool_calls"] + print_verbose(f"completion obj content: {completion_obj['content']}") + if hasattr(chunk, "id"): + model_response.id = chunk.id + self.response_id = chunk.id + if hasattr(chunk, "system_fingerprint"): + self.system_fingerprint = chunk.system_fingerprint + if response_obj["is_finished"]: + self.received_finish_reason = response_obj["finish_reason"] + else: # openai / azure chat model + if self.custom_llm_provider in [ + LlmProviders.AZURE.value, + LlmProviders.AZURE_AI.value, + ]: + if isinstance(chunk, BaseModel) and hasattr(chunk, "model"): + # for azure, we need to pass the model from the original chunk + self.model = getattr(chunk, "model", self.model) + response_obj = self.handle_openai_chat_completion_chunk(chunk) + if response_obj is None: + return _ProviderChunkEarlyReturn(None) + completion_obj["content"] = response_obj["text"] + self.intermittent_finish_reason = response_obj.get("finish_reason", None) + if response_obj["is_finished"]: + if response_obj["finish_reason"] == "error": + raise Exception( + "{} raised a streaming error - finish_reason: error, no content string given. Received Chunk={}".format( + self.custom_llm_provider, response_obj + ) + ) + self.received_finish_reason = response_obj["finish_reason"] + if response_obj.get("original_chunk", None) is not None: + if hasattr(response_obj["original_chunk"], "id"): + model_response = self.set_model_id( + response_obj["original_chunk"].id, model_response + ) + if hasattr(response_obj["original_chunk"], "system_fingerprint"): + model_response.system_fingerprint = response_obj[ + "original_chunk" + ].system_fingerprint + self.system_fingerprint = response_obj[ + "original_chunk" + ].system_fingerprint + if response_obj["logprobs"] is not None: + model_response.choices[0].logprobs = response_obj["logprobs"] + + if response_obj["usage"] is not None: + if isinstance(response_obj["usage"], dict): + setattr( + model_response, + "usage", + litellm.Usage( + prompt_tokens=response_obj["usage"].get( + "prompt_tokens", None + ) + or None, + completion_tokens=response_obj["usage"].get( + "completion_tokens", None + ) + or None, + total_tokens=response_obj["usage"].get("total_tokens", None) + or None, + ), + ) + elif isinstance(response_obj["usage"], Usage): + setattr( + model_response, + "usage", + response_obj["usage"], + ) + elif isinstance(response_obj["usage"], BaseModel): + setattr( + model_response, + "usage", + litellm.Usage(**response_obj["usage"].model_dump()), + ) + return _ProviderChunkParsed(response_obj) + def chunk_creator(self, chunk: Any): # type: ignore if hasattr(chunk, "id"): self.response_id = chunk.id model_response = self.model_response_creator() - response_obj: Dict[str, Any] = {} + response_obj: dict[str, Any] = {} try: # return this for all models - completion_obj: Dict[str, Any] = {"content": ""} - from litellm.types.utils import GenericStreamingChunk as GChunk - - if ( - isinstance(chunk, ModelResponseStream) - and self.custom_llm_provider is not None - and self.custom_llm_provider in litellm._custom_providers - ): - _has_content = bool( - chunk.choices - and chunk.choices[0].delta is not None - and ( - chunk.choices[0].delta.content - or chunk.choices[0].delta.tool_calls - ) - ) - if self.received_finish_reason is not None: - if not _has_content: - raise StopIteration - if chunk.choices and chunk.choices[0].finish_reason: - self.received_finish_reason = chunk.choices[0].finish_reason - if not _has_content: - return None - # Strip finish_reason from the content chunk so it appears - # only on the trailing empty-delta chunk (OpenAI spec). - # finish_reason_handler() will emit the proper terminal chunk. - chunk.choices[0].finish_reason = None # type: ignore[assignment] - return chunk - - if ( - isinstance(chunk, dict) - and generic_chunk_has_all_required_fields( - chunk=chunk - ) # check if chunk is a generic streaming chunk - ) or ( - self.custom_llm_provider - and self.custom_llm_provider in litellm._custom_providers - ): - if self.received_finish_reason is not None: - _chunk_has_content = isinstance(chunk, dict) and ( - bool(chunk.get("text", "")) - or chunk.get("tool_use") is not None - # Usage-only final chunks are valid and needed to surface - # finish_reason/usage to downstream translators. - or chunk.get("usage") is not None - ) - if not _chunk_has_content and ( - not isinstance(chunk, dict) - or "provider_specific_fields" not in chunk - ): - raise StopIteration - anthropic_response_obj: GChunk = cast(GChunk, chunk) - completion_obj["content"] = anthropic_response_obj["text"] - if anthropic_response_obj["is_finished"]: - self.received_finish_reason = anthropic_response_obj[ - "finish_reason" - ] - - if anthropic_response_obj["finish_reason"]: - self.intermittent_finish_reason = anthropic_response_obj[ - "finish_reason" - ] - - if anthropic_response_obj["usage"] is not None: - setattr( - model_response, - "usage", - litellm.Usage(**anthropic_response_obj["usage"]), - ) - - if ( - "tool_use" in anthropic_response_obj - and anthropic_response_obj["tool_use"] is not None - ): - completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]] - - if ( - "provider_specific_fields" in anthropic_response_obj - and anthropic_response_obj["provider_specific_fields"] is not None - ): - for key, value in anthropic_response_obj[ - "provider_specific_fields" - ].items(): - setattr(model_response, key, value) - - response_obj = cast(Dict[str, Any], anthropic_response_obj) - elif self.model == "replicate" or self.custom_llm_provider == "replicate": - response_obj = self.handle_replicate_chunk(chunk) - completion_obj["content"] = response_obj["text"] - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - elif self.custom_llm_provider and self.custom_llm_provider == "predibase": - response_obj = self.handle_predibase_chunk(chunk) - completion_obj["content"] = response_obj["text"] - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - elif ( - self.custom_llm_provider and self.custom_llm_provider == "baseten" - ): # baseten doesn't provide streaming - completion_obj["content"] = self.handle_baseten_chunk(chunk) - elif ( - self.custom_llm_provider and self.custom_llm_provider == "ai21" - ): # ai21 doesn't provide streaming - response_obj = self.handle_ai21_chunk(chunk) - completion_obj["content"] = response_obj["text"] - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - elif self.custom_llm_provider and self.custom_llm_provider == "maritalk": - response_obj = self.handle_maritalk_chunk(chunk) - completion_obj["content"] = response_obj["text"] - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - elif self.custom_llm_provider and self.custom_llm_provider == "vllm": - completion_obj["content"] = chunk[0].outputs[0].text - elif ( - self.custom_llm_provider and self.custom_llm_provider == "aleph_alpha" - ): # aleph alpha doesn't provide streaming - response_obj = self.handle_aleph_alpha_chunk(chunk) - completion_obj["content"] = response_obj["text"] - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - elif self.custom_llm_provider == "nlp_cloud": - try: - response_obj = self.handle_nlp_cloud_chunk(chunk) - completion_obj["content"] = response_obj["text"] - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - except Exception as e: - if self.received_finish_reason: - raise e - else: - if self.sent_first_chunk is False: - raise Exception("An unknown error occurred with the stream") - self.received_finish_reason = "stop" - elif self.custom_llm_provider == "vertex_ai" and not isinstance( - chunk, ModelResponseStream - ): - import proto # type: ignore - - if hasattr(chunk, "candidates") is True: - try: - try: - completion_obj["content"] = chunk.text # type: ignore - except Exception as e: - original_exception = e - if "Part has no text." in str(e): - ## check for function calling - function_call = ( - chunk.candidates[0].content.parts[0].function_call # type: ignore - ) - - args_dict = {} - - # Check if it's a RepeatedComposite instance - for key, val in function_call.args.items(): - if isinstance( - val, - proto.marshal.collections.repeated.RepeatedComposite, # type: ignore - ): - # If so, convert to list - args_dict[key] = [v for v in val] - else: - args_dict[key] = val - - try: - args_str = json.dumps(args_dict) - except Exception as e: - raise e - _delta_obj = litellm.utils.Delta( - content=None, - tool_calls=[ - { - "id": f"call_{str(uuid.uuid4())}", - "function": { - "arguments": args_str, - "name": function_call.name, - }, - "type": "function", - } - ], - ) - _streaming_response = StreamingChoices(delta=_delta_obj) - _model_response = ModelResponseStream() - _model_response.choices = [_streaming_response] - response_obj = {"original_chunk": _model_response} - else: - raise original_exception - if ( - hasattr(chunk.candidates[0], "finish_reason") # type: ignore - and chunk.candidates[0].finish_reason.name # type: ignore - != "FINISH_REASON_UNSPECIFIED" - ): # every non-final chunk in vertex ai has this - self.received_finish_reason = map_finish_reason( # type: ignore - chunk.candidates[0].finish_reason.name - ) - except Exception: - if chunk.candidates[0].finish_reason.name == "SAFETY": # type: ignore - raise Exception( - f"The response was blocked by VertexAI. {str(chunk)}" - ) - else: - completion_obj["content"] = str(chunk) - elif self.custom_llm_provider == "petals": - if self.completion_stream is None or len(self.completion_stream) == 0: - if self.received_finish_reason is not None: - raise StopIteration - else: - self.received_finish_reason = "stop" - chunk_size = 30 - new_chunk = self.completion_stream[:chunk_size] # type: ignore[index] - completion_obj["content"] = new_chunk - self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index] - elif self.custom_llm_provider == "palm": - # fake streaming - response_obj = {} - if self.completion_stream is None or len(self.completion_stream) == 0: - if self.received_finish_reason is not None: - raise StopIteration - else: - self.received_finish_reason = "stop" - chunk_size = 30 - new_chunk = self.completion_stream[:chunk_size] # type: ignore[index] - completion_obj["content"] = new_chunk - self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index] - elif self.custom_llm_provider == "triton": - response_obj = self.handle_triton_stream(chunk) - completion_obj["content"] = response_obj["text"] - print_verbose(f"completion obj content: {completion_obj['content']}") - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - elif self.custom_llm_provider == "text-completion-openai": - response_obj = self.handle_openai_text_completion_chunk(chunk) - completion_obj["content"] = response_obj["text"] - print_verbose(f"completion obj content: {completion_obj['content']}") - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - if response_obj["usage"] is not None: - setattr( - model_response, - "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].prompt_tokens, - completion_tokens=response_obj["usage"].completion_tokens, - total_tokens=response_obj["usage"].total_tokens, - ), - ) - elif self.custom_llm_provider == "text-completion-codestral": - if not isinstance(chunk, str): - raise ValueError(f"chunk is not a string: {chunk}") - response_obj = cast( - Dict[str, Any], - litellm.CodestralTextCompletionConfig()._chunk_parser(chunk), - ) - completion_obj["content"] = response_obj["text"] - print_verbose(f"completion obj content: {completion_obj['content']}") - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - if "usage" in response_obj is not None: - setattr( - model_response, - "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].prompt_tokens, - completion_tokens=response_obj["usage"].completion_tokens, - total_tokens=response_obj["usage"].total_tokens, - ), - ) - elif self.custom_llm_provider == "azure_text": - response_obj = self.handle_azure_text_completion_chunk(chunk) - completion_obj["content"] = response_obj["text"] - print_verbose(f"completion obj content: {completion_obj['content']}") - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - elif self.custom_llm_provider == "cached_response": - chunk = cast(ModelResponseStream, chunk) - response_obj = { - "text": chunk.choices[0].delta.content, - "is_finished": True, - "finish_reason": chunk.choices[0].finish_reason, - "original_chunk": chunk, - "tool_calls": ( - chunk.choices[0].delta.tool_calls - if hasattr(chunk.choices[0].delta, "tool_calls") - else None - ), - } - - completion_obj["content"] = response_obj["text"] - if response_obj["tool_calls"] is not None: - completion_obj["tool_calls"] = response_obj["tool_calls"] - print_verbose(f"completion obj content: {completion_obj['content']}") - if hasattr(chunk, "id"): - model_response.id = chunk.id - self.response_id = chunk.id - if hasattr(chunk, "system_fingerprint"): - self.system_fingerprint = chunk.system_fingerprint - if response_obj["is_finished"]: - self.received_finish_reason = response_obj["finish_reason"] - else: # openai / azure chat model - if self.custom_llm_provider in [ - LlmProviders.AZURE.value, - LlmProviders.AZURE_AI.value, - ]: - if isinstance(chunk, BaseModel) and hasattr(chunk, "model"): - # for azure, we need to pass the model from the original chunk - self.model = getattr(chunk, "model", self.model) - response_obj = self.handle_openai_chat_completion_chunk(chunk) - if response_obj is None: - return - completion_obj["content"] = response_obj["text"] - self.intermittent_finish_reason = response_obj.get( - "finish_reason", None - ) - if response_obj["is_finished"]: - if response_obj["finish_reason"] == "error": - raise Exception( - "{} raised a streaming error - finish_reason: error, no content string given. Received Chunk={}".format( - self.custom_llm_provider, response_obj - ) - ) - self.received_finish_reason = response_obj["finish_reason"] - if response_obj.get("original_chunk", None) is not None: - if hasattr(response_obj["original_chunk"], "id"): - model_response = self.set_model_id( - response_obj["original_chunk"].id, model_response - ) - if hasattr(response_obj["original_chunk"], "system_fingerprint"): - model_response.system_fingerprint = response_obj[ - "original_chunk" - ].system_fingerprint - self.system_fingerprint = response_obj[ - "original_chunk" - ].system_fingerprint - if response_obj["logprobs"] is not None: - model_response.choices[0].logprobs = response_obj["logprobs"] - - if response_obj["usage"] is not None: - if isinstance(response_obj["usage"], dict): - setattr( - model_response, - "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].get( - "prompt_tokens", None - ) - or None, - completion_tokens=response_obj["usage"].get( - "completion_tokens", None - ) - or None, - total_tokens=response_obj["usage"].get( - "total_tokens", None - ) - or None, - ), - ) - elif isinstance(response_obj["usage"], Usage): - setattr( - model_response, - "usage", - response_obj["usage"], - ) - elif isinstance(response_obj["usage"], BaseModel): - setattr( - model_response, - "usage", - litellm.Usage(**response_obj["usage"].model_dump()), - ) + completion_obj: dict[str, Any] = {"content": ""} + dispatch_result = self._dispatch_provider_chunk( + chunk=chunk, + model_response=model_response, + completion_obj=completion_obj, + ) + if isinstance(dispatch_result, _ProviderChunkEarlyReturn): + return dispatch_result.value + response_obj = dispatch_result.response_obj model_response.model = self.model ## FUNCTION CALL PARSING @@ -1980,11 +2005,29 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages, - logging_obj=self.logging_obj, - ) + try: + complete_streaming_response = litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, + ) + except Exception as e: + # stream_chunk_builder can re-raise (as APIError) on large agentic + # streams. The raise originates inside this except-StopIteration block, + # so the sibling `except Exception` below does not catch it; it would + # escape __next__ and drop the request from SpendLogs. Recover + # best-effort usage from the raw chunks so cost is still tracked + verbose_logger.warning( + "stream_chunk_builder raised at end-of-stream (%s); logging " + "best-effort usage from chunks.", + str(e), + ) + try: + complete_streaming_response = self.model_response_creator( + chunk={"usage": calculate_total_usage(chunks=self.chunks)} + ) + except Exception: + complete_streaming_response = None response = self.model_response_creator() if complete_streaming_response is not None: @@ -2209,11 +2252,27 @@ class CustomStreamWrapper: except (StopAsyncIteration, StopIteration): if self.sent_last_chunk is True: # log the final chunk with accurate streaming values - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages, - logging_obj=self.logging_obj, - ) + try: + complete_streaming_response = litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, + ) + except Exception as e: + # see sync __next__: a raise from stream_chunk_builder inside this + # except handler escapes __anext__ and drops the request from SpendLogs. + # Recover best-effort usage from the raw chunks so cost is still tracked + verbose_logger.warning( + "stream_chunk_builder raised at end-of-stream (%s); logging " + "best-effort usage from chunks.", + str(e), + ) + try: + complete_streaming_response = self.model_response_creator( + chunk={"usage": calculate_total_usage(chunks=self.chunks)} + ) + except Exception: + complete_streaming_response = None response = self.model_response_creator() if complete_streaming_response is not None: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py index c7c110ff3e3..8714939f025 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/interceptors/advisor.py @@ -84,8 +84,14 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): ) # Optional routing overrides for the advisor sub-call (e.g. proxy routing). # If not set in the tool definition, litellm resolves from env vars. - advisor_api_key: Optional[str] = advisor_tool.get("api_key") - advisor_api_base: Optional[str] = advisor_tool.get("api_base") + # The advisor tool is caller-controlled; only honor a client-supplied + # api_base/api_key when the proxy has enabled clientside credentials, + # otherwise let litellm resolve from server config. + advisor_api_key: Optional[str] = None + advisor_api_base: Optional[str] = None + if _allow_client_side_advisor_credentials(): + advisor_api_key = advisor_tool.get("api_key") + advisor_api_base = advisor_tool.get("api_base") # Build the synthetic tool definition the provider will receive. synthetic_advisor_tool = _make_synthetic_advisor_tool() @@ -181,6 +187,20 @@ class AdvisorOrchestrationHandler(MessagesInterceptor): # --------------------------------------------------------------------------- +def _allow_client_side_advisor_credentials() -> bool: + """Whether a caller-supplied advisor api_base/api_key may be honored. + + Gated on the proxy's ``allow_client_side_credentials`` opt-in. When the + interceptor runs outside the proxy (SDK use), there is no admin boundary + to protect, so client-supplied routing is allowed. + """ + try: + from litellm.proxy.proxy_server import general_settings + except (ImportError, ModuleNotFoundError): + return True + return general_settings.get("allow_client_side_credentials") is True + + def _make_synthetic_advisor_tool() -> Dict: """Build a regular tool definition the executor provider can understand.""" return { diff --git a/litellm/llms/apiserpent/search/transformation.py b/litellm/llms/apiserpent/search/transformation.py index 1eb7d34c875..bc11875ba12 100644 --- a/litellm/llms/apiserpent/search/transformation.py +++ b/litellm/llms/apiserpent/search/transformation.py @@ -53,7 +53,13 @@ class APISerpentSearchConfig(BaseSearchConfig): api_base: Optional[str] = None, **kwargs, ) -> Dict: - api_key = api_key or get_secret_str("APISERPENT_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("APISERPENT_API_KEY",), + base_env_var="APISERPENT_API_BASE", + default_api_base=APISERPENT_BASE, + ) if not api_key: raise ValueError( "APISERPENT_API_KEY is not set. Set `APISERPENT_API_KEY` environment variable." diff --git a/litellm/llms/base_llm/sandbox/transformation.py b/litellm/llms/base_llm/sandbox/transformation.py index 6ad945f47a3..1c012a15fdb 100644 --- a/litellm/llms/base_llm/sandbox/transformation.py +++ b/litellm/llms/base_llm/sandbox/transformation.py @@ -8,10 +8,14 @@ run code -> delete container; `code_interpreter_tool` combines all three. from typing import Any, Union +import httpx + from pydantic import Field, PrivateAttr from litellm.types.llms.base import LiteLLMPydanticObjectBase +SANDBOX_MAX_OUTPUT_BYTES = 10 * 1024 * 1024 + class ContainerHandle(LiteLLMPydanticObjectBase): """A live sandbox container. Carries everything needed to reach it again.""" @@ -53,7 +57,7 @@ class BaseSandboxConfig: *, template: str | None = None, timeout: int | None = None, - allow_internet_access: bool = True, + allow_internet_access: bool | None = None, api_key: str | None = None, **kwargs, ) -> ContainerHandle: @@ -77,3 +81,16 @@ class BaseSandboxConfig: **kwargs, ) -> bool: raise NotImplementedError("adelete_sandbox must be implemented by provider") + + async def _read_capped_lines(self, response: httpx.Response) -> list[str]: + lines: list[str] = [] + total = 0 + async for line in response.aiter_lines(): + total += len(line.encode("utf-8")) + if total > SANDBOX_MAX_OUTPUT_BYTES: + raise ValueError( + f"Sandbox output exceeded {SANDBOX_MAX_OUTPUT_BYTES} bytes; aborting " + "to avoid unbounded memory use." + ) + lines.append(line) + return lines diff --git a/litellm/llms/base_llm/search/transformation.py b/litellm/llms/base_llm/search/transformation.py index 4dfe86685fb..1581d8bb064 100644 --- a/litellm/llms/base_llm/search/transformation.py +++ b/litellm/llms/base_llm/search/transformation.py @@ -3,11 +3,13 @@ Base Search transformation configuration. """ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from urllib.parse import urlsplit import httpx from pydantic import PrivateAttr from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str from litellm.types.llms.base import LiteLLMPydanticObjectBase if TYPE_CHECKING: @@ -16,6 +18,29 @@ else: LiteLLMLoggingObj = Any +def _search_host(url: str) -> str: + return urlsplit(url).netloc.lower() + + +def _is_trusted_search_api_base( + caller_api_base: str, + default_api_base: str | None, + base_env_var: str | None, +) -> bool: + candidate = _search_host(caller_api_base) + if not candidate: + return False + trusted = { + _search_host(base) + for base in ( + default_api_base, + get_secret_str(base_env_var) if base_env_var else None, + ) + if base + } + return candidate in trusted + + class SearchResult(LiteLLMPydanticObjectBase): """Single search result.""" @@ -86,6 +111,60 @@ class BaseSearchConfig: "max_tokens_per_page", } + def _assert_trusted_api_base_for_server_credential( + self, + caller_api_base: str | None, + default_api_base: str | None, + base_env_var: str | None, + credential_name: str, + ) -> None: + """ + Block sending a server-managed credential to a caller-chosen host. + + A caller-supplied api_base is honored when constructing the request URL, so + falling back to a server-configured secret while the caller controls the host + leaks that secret. The provider default and the operator's own api_base + override are the only trusted destinations for a server-managed credential. + """ + if not caller_api_base: + return + if _is_trusted_search_api_base(caller_api_base, default_api_base, base_env_var): + return + raise ValueError( + f"Refusing to send the server-configured {credential_name} to the " + f"caller-supplied api_base '{caller_api_base}'. Pass an explicit api_key " + f"when overriding api_base for this search provider." + ) + + def resolve_server_api_key( + self, + *, + caller_api_key: str | None, + caller_api_base: str | None, + key_env_vars: tuple[str, ...], + base_env_var: str | None, + default_api_base: str | None, + ) -> str | None: + """ + Resolve a single-secret search API key, falling back to a server-managed + secret only when the request targets a trusted host. + + Returns the caller's key when provided, otherwise the first set + server-managed secret (or None when none is set, for keyless providers). + """ + if caller_api_key: + return caller_api_key + server_key = next( + (key for key in (get_secret_str(var) for var in key_env_vars) if key), + None, + ) + if server_key is None: + return None + self._assert_trusted_api_base_for_server_credential( + caller_api_base, default_api_base, base_env_var, key_env_vars[0] + ) + return server_key + def validate_environment( self, headers: Dict, diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 2c9ea187912..c31462a735b 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -10,7 +10,6 @@ from typing import ( Callable, ClassVar, Dict, - List, Literal, Optional, Tuple, @@ -210,32 +209,11 @@ class BaseAWSLLM: """ Return a boto3.Credentials object """ - ## CHECK IS 'os.environ/' passed in - params_to_check: List[Optional[str]] = [ - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - aws_region_name, - aws_session_name, - aws_profile_name, - aws_role_name, - aws_web_identity_token, - aws_sts_endpoint, - aws_external_id, - ] - - # Iterate over parameters and update if needed - for i, param in enumerate(params_to_check): - if param and param.startswith("os.environ/"): - _v = get_secret(param) - if _v is not None and isinstance(_v, str): - params_to_check[i] = _v - elif param is None: # check if uppercase value in env - key = self.aws_authentication_params[i] - if key.upper() in os.environ: - params_to_check[i] = os.getenv(key.upper()) - - # Assign updated values back to parameters + # Only config-sourced credentials are expanded against the environment. + # os.environ/ references in the model config are resolved at load time, + # so any reference still present at this point is caller-supplied input and is + # left as-is rather than expanded into a process environment variable. Each + # unset param falls back to its matching fixed AWS_* ambient env var. ( aws_access_key_id, aws_secret_access_key, @@ -247,7 +225,21 @@ class BaseAWSLLM: aws_web_identity_token, aws_sts_endpoint, aws_external_id, - ) = params_to_check + ) = tuple( + value if value is not None else os.getenv(env_var) + for value, env_var in ( + (aws_access_key_id, "AWS_ACCESS_KEY_ID"), + (aws_secret_access_key, "AWS_SECRET_ACCESS_KEY"), + (aws_session_token, "AWS_SESSION_TOKEN"), + (aws_region_name, "AWS_REGION_NAME"), + (aws_session_name, "AWS_SESSION_NAME"), + (aws_profile_name, "AWS_PROFILE_NAME"), + (aws_role_name, "AWS_ROLE_NAME"), + (aws_web_identity_token, "AWS_WEB_IDENTITY_TOKEN"), + (aws_sts_endpoint, "AWS_STS_ENDPOINT"), + (aws_external_id, "AWS_EXTERNAL_ID"), + ) + ) verbose_logger.debug( "in get credentials\n" @@ -845,6 +837,20 @@ class BaseAWSLLM: f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}" ) + # get_secret() expands environment-variable references (an os.environ/ + # prefix, or a bare name matching an environment variable). Config-sourced + # references are expanded at load time, so such a reference reaching here is + # caller-supplied input; reject it rather than expanding a process-environment + # value for use as the token. + if ( + aws_web_identity_token.startswith("os.environ/") + or aws_web_identity_token in os.environ + ): + raise AwsAuthError( + message="Invalid web identity token reference.", + status_code=400, + ) + oidc_token = get_secret(aws_web_identity_token) if oidc_token is None: diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 75b560b4d6d..9fca7bc61af 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -70,6 +70,7 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( BedrockError, ModelResponseIterator, + build_bedrock_stream_error, get_bedrock_response_stream_shape, get_bedrock_tool_name, ) @@ -1841,23 +1842,7 @@ class AWSEventStreamDecoder: parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: - decoded_body = response_dict["body"].decode() - if isinstance(decoded_body, dict): - error_message = decoded_body.get("message") - elif isinstance(decoded_body, str): - error_message = decoded_body - else: - error_message = "" - exception_status = response_dict["headers"].get(":exception-type") - error_message = exception_status + " " + error_message - raise BedrockError( - status_code=response_dict["status_code"], - message=( - json.dumps(error_message) - if isinstance(error_message, dict) - else error_message - ), - ) + raise build_bedrock_stream_error(response_dict, response_stream_shape) if "chunk" in parsed_response: chunk = parsed_response.get("chunk") if not chunk: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index bdc5da321c6..9f58e5c0f1c 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -7,9 +7,21 @@ Common utilities used across bedrock chat/embedding/image generation import functools import json import os -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Mapping, + Optional, + TypedDict, + Union, +) if TYPE_CHECKING: + from botocore.model import Shape + from litellm.types.llms.bedrock import BedrockCreateBatchRequest import httpx @@ -1132,6 +1144,39 @@ def get_bedrock_response_stream_shape(): return _load_bedrock_response_stream_shape() +class BedrockEventStreamResponseDict(TypedDict): + status_code: int + headers: Mapping[str, str] + body: bytes + + +def build_bedrock_stream_error( + response_dict: BedrockEventStreamResponseDict, + response_stream_shape: Shape | None, +) -> BedrockError: + """Build a BedrockError for a non-200 event-stream error event. + + botocore hard-codes HTTP 400 on every mid-stream error event, so the modeled + ResponseStream member's httpStatusCode is the real status. Resolve it from the + shape and fall back to the raw status when the type is not modeled. + """ + exception_type = response_dict["headers"].get(":exception-type") + decoded_body = response_dict["body"].decode() + message = f"{exception_type} {decoded_body}" if exception_type else decoded_body + + status_code = response_dict["status_code"] + if exception_type is not None and response_stream_shape is not None: + member = response_stream_shape.members.get(exception_type) + if member is not None: + modeled_status = ( + (member.metadata or {}).get("error", {}).get("httpStatusCode") + ) + if modeled_status is not None: + status_code = int(modeled_status) + + return BedrockError(status_code=status_code, message=message) + + class BedrockEventStreamDecoderBase: """ Base class for event stream decoding for Bedrock @@ -1156,23 +1201,7 @@ class BedrockEventStreamDecoderBase: parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: - decoded_body = response_dict["body"].decode() - if isinstance(decoded_body, dict): - error_message = decoded_body.get("message") - elif isinstance(decoded_body, str): - error_message = decoded_body - else: - error_message = "" - exception_status = response_dict["headers"].get(":exception-type") - error_message = exception_status + " " + error_message - raise BedrockError( - status_code=response_dict["status_code"], - message=( - json.dumps(error_message) - if isinstance(error_message, dict) - else error_message - ), - ) + raise build_bedrock_stream_error(response_dict, response_stream_shape) if "chunk" in parsed_response: chunk = parsed_response.get("chunk") if not chunk: diff --git a/litellm/llms/brave/search/transformation.py b/litellm/llms/brave/search/transformation.py index 9dfcd6bc75a..8ffe7dcb126 100644 --- a/litellm/llms/brave/search/transformation.py +++ b/litellm/llms/brave/search/transformation.py @@ -115,7 +115,13 @@ class BraveSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("BRAVE_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("BRAVE_API_KEY",), + base_env_var="BRAVE_API_BASE", + default_api_base=self.BRAVE_API_BASE, + ) if not api_key: raise ValueError( diff --git a/litellm/llms/cloudflare/chat/transformation.py b/litellm/llms/cloudflare/chat/transformation.py index 66e253f304d..68f08741cc5 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -1,26 +1,15 @@ -import json -import time -from typing import AsyncIterator, Iterator, List, Optional, Union +from typing import List, Optional, Union import httpx -import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segments -from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator -from litellm.llms.base_llm.chat.transformation import ( - BaseConfig, - BaseLLMException, - LiteLLMLoggingObj, +from litellm._logging import verbose_logger +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig +from litellm.secret_managers.main import ( + get_secret_str, + normalize_nonempty_secret_str, ) -from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ( - ChatCompletionToolCallChunk, - ChatCompletionUsageBlock, - GenericStreamingChunk, - ModelResponse, - Usage, -) class CloudflareError(BaseLLMException): @@ -34,26 +23,46 @@ class CloudflareError(BaseLLMException): message=message, request=self.request, response=self.response, - ) # Call the base class constructor with the parameters it needs + ) -class CloudflareChatConfig(BaseConfig): - max_tokens: Optional[int] = None - stream: Optional[bool] = None - - def __init__( +class CloudflareChatConfig(OpenAIGPTConfig): + def get_complete_url( self, - max_tokens: Optional[int] = None, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, stream: Optional[bool] = None, - ) -> None: - locals_ = locals().copy() - for key, value in locals_.items(): - if key != "self" and value is not None: - setattr(self.__class__, key, value) + ) -> str: + return super().get_complete_url( + api_base=self._resolve_api_base(api_base), + api_key=api_key, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + stream=stream, + ) - @classmethod - def get_config(cls): - return super().get_config() + @staticmethod + def _resolve_api_base(api_base: Optional[str]) -> str: + if not api_base: + account_id = normalize_nonempty_secret_str( + get_secret_str("CLOUDFLARE_ACCOUNT_ID") + ) + if account_id is None: + raise ValueError( + "Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID in the environment or pass api_base explicitly" + ) + return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1" + trimmed = api_base.rstrip("/") + if trimmed.endswith("/ai/run"): + verbose_logger.warning( + "Cloudflare api_base ending in '/ai/run' is the legacy Workers AI path and no longer serves OpenAI-compatible requests; rewriting to the '/ai/v1' endpoint" + ) + return f"{trimmed[: -len('/ai/run')]}/ai/v1" + return api_base def validate_environment( self, @@ -67,107 +76,18 @@ class CloudflareChatConfig(BaseConfig): ) -> dict: if api_key is None: raise ValueError( - "Missing CloudflareError API Key - A call is being made to cloudflare but no key is set either in the environment variables or via params" + "Missing Cloudflare API Key - A call is being made to cloudflare but no key is set either in the environment variables or via params" ) - headers = { - "accept": "application/json", - "content-type": "apbplication/json", - "Authorization": "Bearer " + api_key, - } - return headers - - def get_complete_url( - self, - api_base: Optional[str], - api_key: Optional[str], - model: str, - optional_params: dict, - litellm_params: dict, - stream: Optional[bool] = None, - ) -> str: - if api_base is None: - account_id = get_secret_str("CLOUDFLARE_ACCOUNT_ID") - api_base = ( - f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/" - ) - encoded_model = encode_url_path_segments(model, field_name="model") - return api_base + encoded_model - - def get_supported_openai_params(self, model: str) -> List[str]: - return [ - "stream", - "max_tokens", - ] - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - supported_openai_params = self.get_supported_openai_params(model=model) - for param, value in non_default_params.items(): - if param == "max_completion_tokens": - optional_params["max_tokens"] = value - elif param in supported_openai_params: - optional_params[param] = value - return optional_params - - def transform_request( - self, - model: str, - messages: List[AllMessageValues], - optional_params: dict, - litellm_params: dict, - headers: dict, - ) -> dict: - config = litellm.CloudflareChatConfig.get_config() - for k, v in config.items(): - if k not in optional_params: - optional_params[k] = v - - data = { - "messages": messages, - **optional_params, - } - return data - - def transform_response( - self, - model: str, - raw_response: httpx.Response, - model_response: ModelResponse, - logging_obj: LiteLLMLoggingObj, - request_data: dict, - messages: List[AllMessageValues], - optional_params: dict, - litellm_params: dict, - encoding: str, - api_key: Optional[str] = None, - json_mode: Optional[bool] = None, - ) -> ModelResponse: - completion_response = raw_response.json() - - # Support both "response" and "response_text" keys (newer models like Nemotron use "response_text") - result = completion_response["result"] - model_response.choices[0].message.content = result.get("response") if result.get("response") is not None else result.get("response_text", "") # type: ignore - - prompt_tokens = litellm.utils.get_token_count(messages=messages, model=model) - completion_tokens = len( - encoding.encode(model_response["choices"][0]["message"].get("content", "")) + return super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, ) - model_response.created = int(time.time()) - model_response.model = "cloudflare/" + model - usage = Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ) - setattr(model_response, "usage", usage) - return model_response - def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: @@ -175,48 +95,3 @@ class CloudflareChatConfig(BaseConfig): status_code=status_code, message=error_message, ) - - def get_model_response_iterator( - self, - streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], - sync_stream: bool, - json_mode: Optional[bool] = False, - ): - return CloudflareChatResponseIterator( - streaming_response=streaming_response, - sync_stream=sync_stream, - json_mode=json_mode, - ) - - -class CloudflareChatResponseIterator(BaseModelResponseIterator): - def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: - try: - text = "" - tool_use: Optional[ChatCompletionToolCallChunk] = None - is_finished = False - finish_reason = "" - usage: Optional[ChatCompletionUsageBlock] = None - provider_specific_fields = None - - index = int(chunk.get("index", 0)) - - if "response" in chunk and chunk["response"] is not None: - text = chunk["response"] - elif "response_text" in chunk and chunk["response_text"] is not None: - text = chunk["response_text"] - - returned_chunk = GenericStreamingChunk( - text=text, - tool_use=tool_use, - is_finished=is_finished, - finish_reason=finish_reason, - usage=usage, - index=index, - provider_specific_fields=provider_specific_fields, - ) - - return returned_chunk - - except json.JSONDecodeError: - raise ValueError(f"Failed to decode JSON from chunk: {chunk}") diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 01c94476431..1000ab12803 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -42,6 +42,9 @@ from litellm.constants import ( HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS, ) from litellm.litellm_core_utils.logging_utils import track_llm_api_timing +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) from litellm.types.llms.custom_http import * if TYPE_CHECKING: @@ -134,6 +137,18 @@ _DEFAULT_TIMEOUT = httpx.Timeout( timeout=COMPLETION_HTTP_FALLBACK_SECONDS, connect=HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS, ) + + +def _default_cached_client_timeout() -> httpx.Timeout: + """Timeout for cached default httpx clients; honors an explicit litellm.request_timeout.""" + configured = get_configured_request_timeout() + if configured is None: + return _DEFAULT_TIMEOUT + return httpx.Timeout( + timeout=configured, connect=HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS + ) + + _STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS = 5.0 _STREAMING_ERROR_BODY_READ_EXECUTOR = concurrent.futures.ThreadPoolExecutor( max_workers=50, @@ -1379,7 +1394,7 @@ def get_async_httpx_client( _new_client = AsyncHTTPHandler(**handler_params) else: _new_client = AsyncHTTPHandler( - timeout=_DEFAULT_TIMEOUT, + timeout=_default_cached_client_timeout(), shared_session=shared_session, ) @@ -1428,7 +1443,7 @@ def _get_httpx_client(params: Optional[dict] = None) -> HTTPHandler: } _new_client = HTTPHandler(**handler_params) else: - _new_client = HTTPHandler(timeout=_DEFAULT_TIMEOUT) + _new_client = HTTPHandler(timeout=_default_cached_client_timeout()) cache.set_cache( key=_cache_key_name, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8ac5b47c6e7..948c90f9f99 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,5 +1,6 @@ import json import ssl +from functools import lru_cache from urllib.parse import parse_qs, urlencode, urlparse, urlunparse from typing import ( TYPE_CHECKING, @@ -13,6 +14,7 @@ from typing import ( Tuple, Union, cast, + get_type_hints, ) import httpx # type: ignore @@ -26,6 +28,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 REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming +from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, @@ -101,6 +104,7 @@ from litellm.types.llms.openai import ( HttpxBinaryResponseContent, OpenAIFileObject, ResponseInputParam, + ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, ) from litellm.types.rerank import RerankResponse @@ -135,6 +139,7 @@ from litellm.utils import ( ImageResponse, ModelResponse, ProviderConfigManager, + async_pre_call_deployment_hook, ) from .http_handler import get_shared_realtime_ssl_context @@ -184,6 +189,47 @@ def _google_genai_streaming_hidden_params( } +@lru_cache(maxsize=None) +def _responses_api_optional_request_param_names() -> frozenset[str]: + return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys()) + + +def _custom_logger_callbacks(logging_obj: Any) -> list[Any]: + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import ( + get_custom_logger_compatible_class, + ) + + dynamic_success_callbacks = getattr(logging_obj, "dynamic_success_callbacks", None) + callbacks = list(litellm.callbacks) + if isinstance(dynamic_success_callbacks, (list, tuple)): + callbacks.extend(dynamic_success_callbacks) + + custom_loggers: list[Any] = [] + for cb in callbacks: + if isinstance(cb, str): + resolved = get_custom_logger_compatible_class(cb) # type: ignore[arg-type] + if resolved is None: + continue + cb = resolved + if isinstance(cb, CustomLogger): + custom_loggers.append(cb) + return custom_loggers + + +def _has_pre_call_deployment_hook(logging_obj: Any) -> bool: + from litellm.integrations.custom_logger import CustomLogger + + base_func = CustomLogger.async_pre_call_deployment_hook + for cb in _custom_logger_callbacks(logging_obj): + cb_func = getattr(type(cb), "async_pre_call_deployment_hook", base_func) + if getattr(cb_func, "__func__", cb_func) is not getattr( + base_func, "__func__", base_func + ): + return True + return False + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -1833,6 +1879,9 @@ class BaseLLMHTTPHandler: data = provider_config.transform_search_request( query=query, optional_params=optional_params, + api_key=api_key, + api_base=api_base, + headers=headers or {}, ) # Get complete URL (pass data for providers that need request body for URL construction) @@ -2224,12 +2273,92 @@ class BaseLLMHTTPHandler: ) raise ValueError("anthropic_messages_handler is not implemented for sync calls") + def _run_sync_responses_pre_call_deployment_hook( + self, + *, + model: str, + input: Union[str, ResponseInputParam], + custom_llm_provider: str, + response_api_optional_request_params: dict[str, Any], + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + ) -> tuple[ + str, + Union[str, ResponseInputParam], + str, + dict[str, Any], + GenericLiteLLMParams, + ]: + if not _has_pre_call_deployment_hook(logging_obj): + return ( + model, + input, + custom_llm_provider, + response_api_optional_request_params, + litellm_params, + ) + + modified_kwargs = run_async_function( + async_pre_call_deployment_hook, + { + **dict(litellm_params), + **response_api_optional_request_params, + "model": model, + "input": input, + "custom_llm_provider": custom_llm_provider, + }, + CallTypes.responses.value, + ) + if modified_kwargs is None: + return ( + model, + input, + custom_llm_provider, + response_api_optional_request_params, + litellm_params, + ) + + optional_param_names = _responses_api_optional_request_param_names() + updated_response_params = { + **response_api_optional_request_params, + **{ + key: value + for key, value in modified_kwargs.items() + if key in optional_param_names + }, + } + updated_litellm_params = GenericLiteLLMParams( + **{ + **dict(litellm_params), + **{ + key: value + for key, value in modified_kwargs.items() + if key not in optional_param_names + and key not in {"model", "input", "custom_llm_provider"} + }, + } + ) + return ( + str(modified_kwargs["model"]) if "model" in modified_kwargs else model, + cast( + Union[str, ResponseInputParam], + modified_kwargs["input"] if "input" in modified_kwargs else input, + ), + ( + str(modified_kwargs["custom_llm_provider"]) + if "custom_llm_provider" in modified_kwargs + else custom_llm_provider + ), + updated_response_params, + updated_litellm_params, + ) + def response_api_handler( self, model: str, input: Union[str, ResponseInputParam], responses_api_provider_config: BaseResponsesAPIConfig, - response_api_optional_request_params: Dict, + response_api_optional_request_params: dict[str, Any], custom_llm_provider: str, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, @@ -2276,6 +2405,21 @@ class BaseLLMHTTPHandler: shared_session=shared_session, ) + ( + model, + input, + custom_llm_provider, + response_api_optional_request_params, + litellm_params, + ) = self._run_sync_responses_pre_call_deployment_hook( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + response_api_optional_request_params=response_api_optional_request_params, + litellm_params=litellm_params, + logging_obj=logging_obj, + ) + if client is None or not isinstance(client, HTTPHandler): sync_httpx_client = _get_httpx_client( params={"ssl_verify": litellm_params.get("ssl_verify", None)} @@ -2407,12 +2551,36 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) - return responses_api_provider_config.transform_response_api_response( - model=model, - raw_response=response, - logging_obj=logging_obj, + initial_response = ( + responses_api_provider_config.transform_response_api_response( + model=model, + raw_response=response, + logging_obj=logging_obj, + ) ) + if self._has_agentic_completion_hook(logging_obj): + final_response = run_async_function( + self._call_agentic_completion_hooks, + response=initial_response, + model=model, + messages=( + input + if isinstance(input, list) + else [{"role": "user", "content": input}] + ), + anthropic_messages_provider_config=responses_api_provider_config, + anthropic_messages_optional_request_params=response_api_optional_request_params, + logging_obj=logging_obj, + stream=False, + custom_llm_provider=custom_llm_provider, + kwargs=dict(litellm_params), + api_surface="responses", + ) + return final_response if final_response is not None else initial_response + + return initial_response + async def async_response_api_handler( self, model: str, @@ -2570,12 +2738,44 @@ class BaseLLMHTTPHandler: provider_config=responses_api_provider_config, ) - return responses_api_provider_config.transform_response_api_response( - model=model, - raw_response=response, - logging_obj=logging_obj, + initial_response = ( + responses_api_provider_config.transform_response_api_response( + model=model, + raw_response=response, + logging_obj=logging_obj, + ) ) + final_response = await self._call_agentic_completion_hooks( + response=initial_response, + model=model, + messages=( + input + if isinstance(input, list) + else [{"role": "user", "content": input}] + ), + anthropic_messages_provider_config=responses_api_provider_config, + anthropic_messages_optional_request_params=response_api_optional_request_params, + logging_obj=logging_obj, + stream=False, + custom_llm_provider=custom_llm_provider, + kwargs=dict(litellm_params), + api_surface="responses", + ) + + result = final_response if final_response is not None else initial_response + if litellm_params.get( + "_code_interpreter_interception_converted_stream" + ) and not litellm_params.get("_agentic_loop_depth"): + return self._wrap_responses_response_as_fake_stream( + result=result, + model=model, + responses_api_provider_config=responses_api_provider_config, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + return result + async def async_delete_response_api_handler( self, response_id: str, @@ -4734,22 +4934,9 @@ class BaseLLMHTTPHandler: agentic callback is detected too. """ from litellm.integrations.custom_logger import CustomLogger - from litellm.litellm_core_utils.litellm_logging import ( - get_custom_logger_compatible_class, - ) base_func = CustomLogger.async_should_run_agentic_loop - callbacks = litellm.callbacks + ( - getattr(logging_obj, "dynamic_success_callbacks", None) or [] - ) - for cb in callbacks: - if isinstance(cb, str): - resolved = get_custom_logger_compatible_class(cb) # type: ignore[arg-type] - if resolved is None: - continue - cb = resolved - if not isinstance(cb, CustomLogger): - continue + for cb in _custom_logger_callbacks(logging_obj): cb_func = getattr(type(cb), "async_should_run_agentic_loop", base_func) if getattr(cb_func, "__func__", cb_func) is not getattr( base_func, "__func__", base_func @@ -4875,6 +5062,132 @@ class BaseLLMHTTPHandler: return response + async def _execute_responses_agentic_plan( + self, + plan: AgenticLoopPlan, + model: str, + response_api_optional_request_params: dict, + logging_obj: "LiteLLMLoggingObj", + kwargs: dict, + depth: int, + max_loops: int, + fingerprints: list[str], + fingerprint: str, + callback: Any | None = None, + ) -> Any: + patch = plan.request_patch or AgenticLoopRequestPatch() + if patch.messages is None: + raise ValueError("Agentic loop plan missing patched responses input") + + optional_params = dict(response_api_optional_request_params) + optional_params.update(patch.optional_params) + if patch.tools is not None: + optional_params["tools"] = patch.tools + optional_params = { + k: v + for k, v in optional_params.items() + if k != "stream" and k != "_code_interpreter_interception_converted_stream" + } + + internal_keys = {"litellm_logging_obj"} + kwargs_for_followup = { + k: v + for k, v in kwargs.items() + if not k.startswith("_websearch_interception") + and not k.startswith("_compression_interception") + and k != "_code_interpreter_interception_converted_stream" + and k not in internal_keys + and k not in optional_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] + + try: + response = await litellm.aresponses( + model=patch.model or model, + input=patch.messages, + **optional_params, + **kwargs_for_followup, + ) + + if callback is not None: + try: + response = await callback.async_post_agentic_loop_response_hook( + response=response, plan=plan, kwargs=kwargs + ) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_post_agentic_loop_response_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + + return response + finally: + if callback is not None: + await self._run_agentic_loop_cleanup( + callback=callback, + plan=plan, + kwargs=kwargs, + logging_obj=logging_obj, + model=model, + ) + + @staticmethod + async def _run_agentic_loop_cleanup( + callback: Any, + plan: AgenticLoopPlan, + kwargs: dict, + logging_obj: "LiteLLMLoggingObj", + model: str, + ) -> None: + try: + await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + + def _wrap_responses_response_as_fake_stream( + self, + result: Any, + model: str, + responses_api_provider_config: Any, + logging_obj: "LiteLLMLoggingObj", + custom_llm_provider: str, + ) -> Any: + """ + Wrap a completed responses result as a synthetic stream. + + Used when an interceptor forced stream=False to run the agentic loop on + the non-streaming path, but the caller originally asked for streaming. + """ + import httpx + + from litellm.responses.streaming_iterator import ( + MockResponsesAPIStreamingIterator, + ) + + payload = result.model_dump() if hasattr(result, "model_dump") else result + raw_response = httpx.Response(status_code=200, json=payload) + return MockResponsesAPIStreamingIterator( + response=raw_response, + model=model, + responses_api_provider_config=responses_api_provider_config, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + async def _execute_chat_completion_agentic_plan( self, plan: AgenticLoopPlan, @@ -4940,6 +5253,7 @@ class BaseLLMHTTPHandler: stream: bool, custom_llm_provider: str, kwargs: Dict, + api_surface: str = "anthropic_messages", ) -> Optional[Any]: """ Call agentic completion hooks for all custom loggers (Anthropic Messages API). @@ -5046,6 +5360,20 @@ class BaseLLMHTTPHandler: if not plan.run_agentic_loop: continue + if api_surface == "responses": + return await self._execute_responses_agentic_plan( + plan=plan, + model=model, + response_api_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + kwargs=kwargs_with_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + callback=callback, + ) + return await self._execute_anthropic_agentic_plan( plan=plan, model=model, @@ -5083,7 +5411,7 @@ class BaseLLMHTTPHandler: else False ) - if websearch_converted_stream: + if api_surface == "anthropic_messages" and websearch_converted_stream: from typing import cast from litellm._logging import verbose_logger @@ -5358,9 +5686,7 @@ class BaseLLMHTTPHandler: import websockets from websockets.asyncio.client import ClientConnection - url = self._append_query_params( - provider_config.get_complete_url(api_base, model, api_key), query_params - ) + url = provider_config.get_complete_url(api_base, model, api_key) headers = provider_config.validate_environment( headers=headers, model=model, diff --git a/litellm/llms/dataforseo/search/transformation.py b/litellm/llms/dataforseo/search/transformation.py index 27c10d740b5..701db586b72 100644 --- a/litellm/llms/dataforseo/search/transformation.py +++ b/litellm/llms/dataforseo/search/transformation.py @@ -61,9 +61,18 @@ class DataForSEOSearchConfig(BaseSearchConfig): password = get_secret_str("DATAFORSEO_PASSWORD") # If api_key is provided in "login:password" format, use it + caller_supplied_credentials = bool(api_key and ":" in api_key) if api_key and ":" in api_key: login, password = api_key.split(":", 1) + if not caller_supplied_credentials and login and password: + self._assert_trusted_api_base_for_server_credential( + api_base, + self.DATAFORSEO_API_BASE, + "DATAFORSEO_API_BASE", + "DATAFORSEO_LOGIN", + ) + if not login: raise ValueError( "DATAFORSEO_LOGIN is not set. Set `DATAFORSEO_LOGIN` environment variable or pass credentials in api_key parameter." diff --git a/litellm/llms/e2b/sandbox/transformation.py b/litellm/llms/e2b/sandbox/transformation.py index 1ce28bc55fb..ecfc1642c97 100644 --- a/litellm/llms/e2b/sandbox/transformation.py +++ b/litellm/llms/e2b/sandbox/transformation.py @@ -16,6 +16,7 @@ from litellm.llms.base_llm.sandbox.transformation import ( BaseSandboxConfig, CodeExecutionResult, ContainerHandle, + SANDBOX_MAX_OUTPUT_BYTES, ) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -29,7 +30,7 @@ E2B_DEFAULT_TEMPLATE = "code-interpreter-v1" E2B_DEFAULT_DOMAIN = "e2b.app" JUPYTER_PORT = 49999 DEFAULT_SANDBOX_TIMEOUT = 300 -MAX_OUTPUT_BYTES = 10 * 1024 * 1024 +MAX_OUTPUT_BYTES = SANDBOX_MAX_OUTPUT_BYTES class E2BSandboxConfig(BaseSandboxConfig): @@ -49,18 +50,22 @@ class E2BSandboxConfig(BaseSandboxConfig): *, template: str | None = None, timeout: int | None = None, - allow_internet_access: bool = True, + allow_internet_access: bool | None = None, api_key: str | None = None, + api_base: str | None = None, metadata: dict | None = None, client: AsyncHTTPHandler | None = None, **kwargs, ) -> ContainerHandle: key = self.validate_environment(api_key=api_key) + base = api_base or E2B_API_BASE body = { "templateID": template or E2B_DEFAULT_TEMPLATE, "timeout": timeout if timeout is not None else DEFAULT_SANDBOX_TIMEOUT, "secure": True, - "allow_internet_access": allow_internet_access, + "allow_internet_access": ( + True if allow_internet_access is None else allow_internet_access + ), } if metadata: body["metadata"] = metadata @@ -68,7 +73,7 @@ class E2BSandboxConfig(BaseSandboxConfig): response = cast( httpx.Response, await self._http(client).post( - url=f"{E2B_API_BASE}/sandboxes", + url=f"{base}/sandboxes", headers={"X-API-Key": key, "Content-Type": "application/json"}, json=body, ), @@ -84,6 +89,7 @@ class E2BSandboxConfig(BaseSandboxConfig): "envd_access_token": data.get("envdAccessToken"), "traffic_access_token": data.get("trafficAccessToken"), "api_key": key, + "api_base": base, } return handle @@ -130,6 +136,7 @@ class E2BSandboxConfig(BaseSandboxConfig): *, container: Union[ContainerHandle, str], api_key: str | None = None, + api_base: str | None = None, client: AsyncHTTPHandler | None = None, **kwargs, ) -> bool: @@ -139,11 +146,12 @@ class E2BSandboxConfig(BaseSandboxConfig): or handle._hidden_params.get("api_key") or self.validate_environment() ) + base = api_base or handle._hidden_params.get("api_base") or E2B_API_BASE try: response = cast( httpx.Response, await self._http(client).delete( - url=f"{E2B_API_BASE}/sandboxes/{handle.id}", + url=f"{base}/sandboxes/{handle.id}", headers={"X-API-Key": key}, ), ) @@ -163,20 +171,6 @@ class E2BSandboxConfig(BaseSandboxConfig): handle._hidden_params = {} return handle - @staticmethod - async def _read_capped_lines(response: httpx.Response) -> list[str]: - lines: list[str] = [] - total = 0 - async for line in response.aiter_lines(): - total += len(line.encode("utf-8")) - if total > MAX_OUTPUT_BYTES: - raise ValueError( - f"Sandbox output exceeded {MAX_OUTPUT_BYTES} bytes; aborting to " - "avoid unbounded memory use." - ) - lines.append(line) - return lines - @staticmethod def _parse_lines(lines: list[str]) -> CodeExecutionResult: def _try_parse(stripped: str): @@ -187,10 +181,9 @@ class E2BSandboxConfig(BaseSandboxConfig): messages = tuple( parsed - for stripped in (line.strip() for line in lines) - if stripped - for parsed in (_try_parse(stripped),) - if parsed is not None + for line in lines + if (stripped := line.strip()) + if (parsed := _try_parse(stripped)) is not None ) def of_type(message_type: str): diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 7a34ededa6b..5cfd14aeaa9 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -65,7 +65,13 @@ class ExaAISearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("EXA_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("EXA_API_KEY",), + base_env_var="EXA_API_BASE", + default_api_base=self.EXA_AI_API_BASE, + ) if not api_key: raise ValueError( "EXA_API_KEY is not set. Set `EXA_API_KEY` environment variable." diff --git a/litellm/llms/fastcrw/search/transformation.py b/litellm/llms/fastcrw/search/transformation.py index ce702266e7b..b571a659cac 100644 --- a/litellm/llms/fastcrw/search/transformation.py +++ b/litellm/llms/fastcrw/search/transformation.py @@ -57,7 +57,13 @@ class FastCRWSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("CRW_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("CRW_API_KEY",), + base_env_var="CRW_API_BASE", + default_api_base=self.FASTCRW_API_BASE, + ) if not api_key: raise ValueError( "CRW_API_KEY is not set. Set `CRW_API_KEY` environment variable." diff --git a/litellm/llms/firecrawl/search/transformation.py b/litellm/llms/firecrawl/search/transformation.py index 18cf1d28c4d..7e01ba58706 100644 --- a/litellm/llms/firecrawl/search/transformation.py +++ b/litellm/llms/firecrawl/search/transformation.py @@ -61,7 +61,13 @@ class FirecrawlSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("FIRECRAWL_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("FIRECRAWL_API_KEY",), + base_env_var="FIRECRAWL_API_BASE", + default_api_base=self.FIRECRAWL_API_BASE, + ) if not api_key: raise ValueError( "FIRECRAWL_API_KEY is not set. Set `FIRECRAWL_API_KEY` environment variable." diff --git a/litellm/llms/fireworks_ai/audio_transcription/transformation.py b/litellm/llms/fireworks_ai/audio_transcription/transformation.py deleted file mode 100644 index 00bb5f26797..00000000000 --- a/litellm/llms/fireworks_ai/audio_transcription/transformation.py +++ /dev/null @@ -1,17 +0,0 @@ -from typing import List - -from litellm.types.llms.openai import OpenAIAudioTranscriptionOptionalParams - -from ...openai.transcriptions.whisper_transformation import ( - OpenAIWhisperAudioTranscriptionConfig, -) -from ..common_utils import FireworksAIMixin - - -class FireworksAIAudioTranscriptionConfig( - FireworksAIMixin, OpenAIWhisperAudioTranscriptionConfig -): - def get_supported_openai_params( - self, model: str - ) -> List[OpenAIAudioTranscriptionOptionalParams]: - return ["language", "prompt", "response_format", "timestamp_granularities"] diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 341c2fc7350..7e4395959b9 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -1,5 +1,15 @@ import json -from typing import Any, List, Literal, Optional, Tuple, Union, cast +from typing import ( + Any, + AsyncIterator, + Iterator, + List, + Literal, + Optional, + Tuple, + Union, + cast, +) import httpx @@ -15,7 +25,6 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, - ChatCompletionImageObject, ChatCompletionToolParam, OpenAIChatCompletionToolParam, ) @@ -25,6 +34,7 @@ from litellm.types.utils import ( Function, Message, ModelResponse, + ModelResponseStream, ProviderSpecificModelInfo, ) from litellm.utils import ( @@ -34,10 +44,34 @@ from litellm.utils import ( supports_tool_choice, ) -from ...openai.chat.gpt_transformation import OpenAIGPTConfig +from ...openai.chat.gpt_transformation import ( + OpenAIChatCompletionStreamingHandler, + OpenAIGPTConfig, +) from ..common_utils import FireworksAIException +def _extract_fireworks_hidden_params(payload: dict) -> dict: + """ + Collect Fireworks-specific response fields (perf_metrics, prompt_token_ids, + per-choice raw_output and token_ids) from a non-streaming completion payload + or a single streaming chunk, so the same data lands in ``_hidden_params`` on + both response paths. + """ + choices = [c for c in (payload.get("choices") or []) if isinstance(c, dict)] + top_level = { + f"fireworks_{field}": payload[field] + for field in ("perf_metrics", "prompt_token_ids") + if field in payload + } + per_choice = { + f"fireworks_{dest}": [c[field] for c in choices if field in c] + for field, dest in (("raw_output", "raw_outputs"), ("token_ids", "token_ids")) + if any(field in c for c in choices) + } + return {**top_level, **per_choice} + + class FireworksAIConfig(OpenAIGPTConfig): """ Reference: https://docs.fireworks.ai/api-reference/post-chatcompletions @@ -60,8 +94,7 @@ class FireworksAIConfig(OpenAIGPTConfig): logprobs: Optional[int] = None reasoning_effort: Optional[str] = None - # Non OpenAI parameters - Fireworks AI only params - prompt_truncate_length: Optional[int] = None + prompt_truncate_len: Optional[int] = None context_length_exceeded_behavior: Optional[Literal["error", "truncate"]] = None def __init__( @@ -80,7 +113,7 @@ class FireworksAIConfig(OpenAIGPTConfig): user: Optional[str] = None, logprobs: Optional[int] = None, reasoning_effort: Optional[str] = None, - prompt_truncate_length: Optional[int] = None, + prompt_truncate_len: Optional[int] = None, context_length_exceeded_behavior: Optional[Literal["error", "truncate"]] = None, ) -> None: locals_ = locals().copy() @@ -108,8 +141,30 @@ class FireworksAIConfig(OpenAIGPTConfig): "response_format", "user", "logprobs", - "prompt_truncate_length", + "prompt_truncate_len", "context_length_exceeded_behavior", + "seed", + "top_logprobs", + "min_p", + "typical_p", + "repetition_penalty", + "mirostat_target", + "mirostat_lr", + "logit_bias", + "echo", + "echo_last", + "ignore_eos", + "prompt_cache_key", + "prompt_cache_isolation_key", + "raw_output", + "perf_metrics_in_response", + "return_token_ids", + "safe_tokenization", + "service_tier", + "speculation", + "prediction", + "stream_options", + "sampling_mask", ] # Only add tools for models that support function calling @@ -133,9 +188,11 @@ class FireworksAIConfig(OpenAIGPTConfig): if supports_tool_choice(model=model, custom_llm_provider="fireworks_ai"): supported_params.append("tool_choice") - # Only add reasoning_effort for models that support it + # Only add reasoning params for models that support it if supports_reasoning(model=model, custom_llm_provider="fireworks_ai"): supported_params.append("reasoning_effort") + supported_params.append("reasoning_history") + supported_params.append("thinking") return supported_params @@ -151,6 +208,18 @@ class FireworksAIConfig(OpenAIGPTConfig): param == "tools" and value is not None for param, value in non_default_params.items() ) + if ( + non_default_params.get("thinking") is not None + and non_default_params.get("reasoning_effort") is not None + ): + raise litellm.BadRequestError( + message=( + "Fireworks AI chat completions does not support specifying both " + "`thinking` and `reasoning_effort` in the same request." + ), + model=model, + llm_provider="fireworks_ai", + ) for param, value in non_default_params.items(): if param == "tool_choice": @@ -174,40 +243,19 @@ class FireworksAIConfig(OpenAIGPTConfig): optional_params["response_format"] = value elif param == "max_completion_tokens": optional_params["max_tokens"] = value + elif param == "reasoning_effort": + if value is True: + optional_params["reasoning_effort"] = "medium" + elif value is False: + optional_params["reasoning_effort"] = "none" + else: + optional_params["reasoning_effort"] = value elif param in supported_openai_params: if value is not None: optional_params[param] = value return optional_params - def _add_transform_inline_image_block( - self, - content: ChatCompletionImageObject, - model: str, - disable_add_transform_inline_image_block: Optional[bool], - ) -> ChatCompletionImageObject: - """ - Add transform_inline to the image_url (allows non-vision models to parse documents/images/etc.) - - ignore if model is a vision model - - ignore if user has disabled this feature - """ - if ( - "vision" in model or disable_add_transform_inline_image_block - ): # allow user to toggle this feature. - return content - if isinstance(content["image_url"], str): - # Skip base64 data URLs — appending #transform=inline corrupts the - # base64 payload and causes an "Incorrect padding" decode error on - # the Fireworks side. Data URLs are already inlined by definition. - # Lower-case before checking: URI schemes are case-insensitive (RFC 3986). - if not content["image_url"].lower().startswith("data:"): - content["image_url"] = f"{content['image_url']}#transform=inline" - elif isinstance(content["image_url"], dict): - url = content["image_url"]["url"] - if not url.lower().startswith("data:"): - content["image_url"]["url"] = f"{url}#transform=inline" - return content - def _transform_tools( self, tools: List[OpenAIChatCompletionToolParam] ) -> List[OpenAIChatCompletionToolParam]: @@ -225,36 +273,46 @@ class FireworksAIConfig(OpenAIGPTConfig): self, messages: List[AllMessageValues], model: str, litellm_params: dict ) -> List[AllMessageValues]: """ - Add 'transform=inline' to the url of the image_url + Strip fields not permitted by FireworksAI from messages. """ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, - migrate_file_to_image_url, ) - disable_add_transform_inline_image_block = cast( - Optional[bool], - litellm_params.get("disable_add_transform_inline_image_block") - or litellm.disable_add_transform_inline_image_block, + supports_vision_value = self._get_model_cost_capability_exact( + model=model, capability="supports_vision" ) - ## For any 'file' message type with pdf content, move to 'image_url' message type - for message in messages: - if message["role"] == "user": - _message_content = message.get("content") - if _message_content is not None and isinstance(_message_content, list): - for idx, content in enumerate(_message_content): - if content["type"] == "file": - _message_content[idx] = migrate_file_to_image_url(content) for message in messages: if message["role"] == "user": _message_content = message.get("content") if _message_content is not None and isinstance(_message_content, list): for content in _message_content: - if content["type"] == "image_url": - content = self._add_transform_inline_image_block( - content=content, + if not isinstance(content, dict): + continue + if content.get("type") == "file": + raise litellm.BadRequestError( + message=( + "Fireworks AI chat completions does not support " + "file content blocks. For PDFs, convert pages to " + "images and send image_url blocks to a Fireworks " + "vision model, or extract text before calling a " + "text-only model." + ), model=model, - disable_add_transform_inline_image_block=disable_add_transform_inline_image_block, + llm_provider="fireworks_ai", + ) + if ( + content.get("type") == "image_url" + and supports_vision_value is False + ): + raise litellm.BadRequestError( + message=( + f"Fireworks AI model {model} does not support " + "image inputs. Use a Fireworks vision model or " + "remove image_url content blocks." + ), + model=model, + llm_provider="fireworks_ai", ) filter_value_from_dict(cast(dict, message), "cache_control") # Remove fields not permitted by FireworksAI (additionalProperties: false @@ -317,43 +375,55 @@ class FireworksAIConfig(OpenAIGPTConfig): return True return ("-" + key_short + "-") in short_name - def _get_model_cost_capability(self, model: str, capability: str) -> Optional[bool]: + @staticmethod + def _short_model_name(model: str) -> str: short_name = model if short_name.startswith("fireworks_ai/"): short_name = short_name[len("fireworks_ai/") :] if short_name.startswith("accounts/fireworks/models/"): short_name = short_name[len("accounts/fireworks/models/") :] + return short_name - candidate_keys = [ + def _get_model_cost_capability_exact( + self, model: str, capability: str + ) -> Optional[bool]: + short_name = self._short_model_name(model) + candidate_keys = ( model, f"fireworks_ai/{short_name}", f"fireworks_ai/accounts/fireworks/models/{short_name}", - ] - + ) for candidate_key in candidate_keys: model_info = litellm.model_cost.get(candidate_key) if model_info is not None and model_info.get(capability) is not None: return cast(Optional[bool], model_info.get(capability)) + return None - # Fallback: preserve historical substring matching for model name - # variants (e.g. fine-tuned or regionally-suffixed versions of a - # known model). Pick the *longest* matching entry so a more specific - # known model (e.g. "qwen3-8b-instruct") wins over a less specific - # one (e.g. "qwen3-8b") when the query model is more specific still. - # Use hyphen-aligned matching to avoid false positives where a short - # known model name is an unrelated substring of a longer one. - best_match_short: Optional[str] = None - best_match_value: Optional[bool] = None - for key_short, model_info in self._get_fireworks_index(): - if model_info.get(capability) is None: - continue - if not self._matches_on_hyphen_boundary(short_name, key_short): - continue - if best_match_short is None or len(key_short) > len(best_match_short): - best_match_short = key_short - best_match_value = cast(Optional[bool], model_info.get(capability)) + def _get_model_cost_capability(self, model: str, capability: str) -> Optional[bool]: + exact = self._get_model_cost_capability_exact( + model=model, capability=capability + ) + if exact is not None: + return exact - return best_match_value + # Fallback: substring matching for model name variants (e.g. fine-tuned + # or regionally-suffixed versions of a known model). Pick the *longest* + # matching entry so a more specific known model (e.g. "qwen3-8b-instruct") + # wins over a less specific one (e.g. "qwen3-8b"). Hyphen-aligned matching + # avoids false positives where a short known name is an unrelated + # substring of a longer one. This stays a soft signal: capability-gated + # hard rejections use the exact lookup so a fuzzy match never blocks a + # custom deployment. + short_name = self._short_model_name(model) + matches = [ + (key_short, cast(Optional[bool], model_info.get(capability))) + for key_short, model_info in self._get_fireworks_index() + if model_info.get(capability) is not None + and self._matches_on_hyphen_boundary(short_name, key_short) + ] + if not matches: + return None + return max(matches, key=lambda match: len(match[0]))[1] def get_provider_info(self, model: str) -> ProviderSpecificModelInfo: supports_function_calling_value = self._get_model_cost_capability( @@ -362,12 +432,16 @@ class FireworksAIConfig(OpenAIGPTConfig): supports_reasoning_value = self._get_model_cost_capability( model=model, capability="supports_reasoning" ) + supports_vision_value = self._get_model_cost_capability( + model=model, capability="supports_vision" + ) + supports_pdf_input_value = self._get_model_cost_capability( + model=model, capability="supports_pdf_input" + ) provider_specific_model_info: ProviderSpecificModelInfo = { "supports_function_calling": True, "supports_prompt_caching": True, # https://docs.fireworks.ai/guides/prompt-caching - "supports_pdf_input": True, # via document inlining - "supports_vision": True, # via document inlining } if supports_function_calling_value is not None: @@ -381,6 +455,14 @@ class FireworksAIConfig(OpenAIGPTConfig): supports_reasoning_value ) + if supports_vision_value is not None: + provider_specific_model_info["supports_vision"] = supports_vision_value + + if supports_pdf_input_value is not None: + provider_specific_model_info["supports_pdf_input"] = ( + supports_pdf_input_value + ) + return provider_specific_model_info def transform_request( @@ -402,6 +484,15 @@ class FireworksAIConfig(OpenAIGPTConfig): if "tools" in optional_params and optional_params["tools"] is not None: tools = self._transform_tools(tools=optional_params["tools"]) optional_params["tools"] = tools + if optional_params.get("stream"): + stream_options = optional_params.get("stream_options") + if stream_options is None: + optional_params["stream_options"] = {"include_usage": True} + elif stream_options.get("include_usage") is not False: + optional_params["stream_options"] = { + **stream_options, + "include_usage": True, + } return super().transform_request( model=model, messages=messages, @@ -494,10 +585,25 @@ class FireworksAIConfig(OpenAIGPTConfig): ) ) - response._hidden_params = {"additional_headers": additional_headers} + response._hidden_params = { + "additional_headers": additional_headers, + **_extract_fireworks_hidden_params(completion_response), + } return response + def get_model_response_iterator( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], + sync_stream: bool, + json_mode: Optional[bool] = False, + ) -> Any: + return FireworksAIChatCompletionStreamingHandler( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + def _get_openai_compatible_provider_info( self, api_base: Optional[str], api_key: Optional[str] ) -> Tuple[Optional[str], Optional[str]]: @@ -554,3 +660,15 @@ class FireworksAIConfig(OpenAIGPTConfig): or get_secret_str("FIREWORKSAI_API_KEY") or get_secret_str("FIREWORKS_AI_TOKEN") ) + + +class FireworksAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): + def chunk_parser(self, chunk: dict) -> ModelResponseStream: + parsed = super().chunk_parser(chunk) + fireworks_fields = _extract_fireworks_hidden_params(chunk) + if fireworks_fields: + parsed.provider_specific_fields = { + **(getattr(parsed, "provider_specific_fields", None) or {}), + **fireworks_fields, + } + return parsed diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 74f6cd4d831..e153d00e6ab 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -103,6 +103,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # bypassing spend and budget accounting. self._pending_usage_metadata: Optional[dict] = None + def _include_function_response_id(self) -> bool: + """Google AI Studio Gemini 3.5+ accepts ``id`` on functionResponses; Vertex AI rejects it.""" + return True + @staticmethod def _usage_detail_alias(details: Any, defaults: Dict[str, int]) -> Dict[str, Any]: if not isinstance(details, dict): @@ -604,10 +608,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) # Build Gemini toolResponse format - function_response = { - "id": call_id, - "response": output_dict, - } + function_response: dict[str, Any] = {"response": output_dict} + if self._include_function_response_id() and call_id: + function_response["id"] = call_id if function_name: function_response["name"] = function_name diff --git a/litellm/llms/google_pse/search/transformation.py b/litellm/llms/google_pse/search/transformation.py index a8aa109cbf0..5cd3f2085a8 100644 --- a/litellm/llms/google_pse/search/transformation.py +++ b/litellm/llms/google_pse/search/transformation.py @@ -85,7 +85,13 @@ class GooglePSESearchConfig(BaseSearchConfig): Google PSE uses API key as a query parameter, not in headers. This method is called but headers are not used for authentication. """ - api_key = api_key or get_secret_str("GOOGLE_PSE_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("GOOGLE_PSE_API_KEY",), + base_env_var="GOOGLE_PSE_API_BASE", + default_api_base=self.GOOGLE_PSE_API_BASE, + ) if not api_key: raise ValueError( "GOOGLE_PSE_API_KEY is not set. Set `GOOGLE_PSE_API_KEY` environment variable." @@ -137,6 +143,7 @@ class GooglePSESearchConfig(BaseSearchConfig): query: Union[str, List[str]], optional_params: dict, api_key: Optional[str] = None, + api_base: str | None = None, search_engine_id: Optional[str] = None, **kwargs, ) -> Dict: @@ -165,8 +172,16 @@ class GooglePSESearchConfig(BaseSearchConfig): # Google PSE only supports single string queries query = " ".join(query) - # Get API credentials - api_key = api_key or get_secret_str("GOOGLE_PSE_API_KEY") + # Get API credentials. The key is sent as a query param to api_base, so + # resolve it host-aware to avoid leaking a server-managed key to a + # caller-supplied host. + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("GOOGLE_PSE_API_KEY",), + base_env_var="GOOGLE_PSE_API_BASE", + default_api_base=self.GOOGLE_PSE_API_BASE, + ) search_engine_id = search_engine_id or get_secret_str("GOOGLE_PSE_ENGINE_ID") if not api_key: diff --git a/litellm/llms/linkup/search/transformation.py b/litellm/llms/linkup/search/transformation.py index 2b17d5642ac..d27ae038f9e 100644 --- a/litellm/llms/linkup/search/transformation.py +++ b/litellm/llms/linkup/search/transformation.py @@ -61,7 +61,13 @@ class LinkupSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("LINKUP_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("LINKUP_API_KEY",), + base_env_var="LINKUP_API_BASE", + default_api_base=self.LINKUP_API_BASE, + ) if not api_key: raise ValueError( "LINKUP_API_KEY is not set. Set `LINKUP_API_KEY` environment variable." diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index f1ad3708236..8d0cf993814 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -247,6 +247,8 @@ class MistralConfig(OpenAIGPTConfig): The above statement is not valid now. Need to plan to remove all the #1,2,3 Mistral API supports content as a list. """ + messages = [self._strip_output_only_fields(m) for m in messages] + ## 1. If 'image_url' or 'file' in content, then transform with base class and mistral-specific handling for m in messages: _content_block = m.get("content") @@ -409,6 +411,25 @@ class MistralConfig(OpenAIGPTConfig): return cleaned_tools + @classmethod + def _strip_output_only_fields(cls, message: AllMessageValues) -> AllMessageValues: + """ + ``reasoning_content`` and ``thinking_blocks`` are output-only fields that + LiteLLM attaches to assistant responses. Mistral's input schema forbids + unknown fields, so replaying them verbatim in a follow-up turn triggers a + 422 ``extra_forbidden``. Drop them before the request is sent. + """ + if message["role"] != "assistant": + return message + return cast( + AllMessageValues, + { + k: v + for k, v in message.items() + if k not in ("reasoning_content", "thinking_blocks") + }, + ) + @classmethod def _handle_name_in_message(cls, message: AllMessageValues) -> AllMessageValues: """ diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 24943563937..d87346fea70 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -115,6 +115,14 @@ "max_completion_tokens": "max_tokens" } }, + "darkbloom": { + "base_url": "https://api.darkbloom.dev/v1", + "api_key_env": "DARKBLOOM_API_KEY", + "api_base_env": "DARKBLOOM_API_BASE", + "param_mappings": { + "max_completion_tokens": "max_tokens" + } + }, "neosantara": { "base_url": "https://api.neosantara.xyz/v1", "api_key_env": "NEOSANTARA_API_KEY", diff --git a/litellm/llms/opensandbox/__init__.py b/litellm/llms/opensandbox/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/litellm/llms/opensandbox/__init__.py @@ -0,0 +1 @@ + diff --git a/litellm/llms/opensandbox/sandbox/__init__.py b/litellm/llms/opensandbox/sandbox/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/litellm/llms/opensandbox/sandbox/__init__.py @@ -0,0 +1 @@ + diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py new file mode 100644 index 00000000000..dc9f8440d30 --- /dev/null +++ b/litellm/llms/opensandbox/sandbox/transformation.py @@ -0,0 +1,598 @@ +import asyncio +import json +import time +from typing import Union, cast + +import httpx + +from litellm.constants import ( + OPEN_SANDBOX_API_BASE_ENV_VAR, + OPEN_SANDBOX_API_KEY_ENV_VAR, + OPEN_SANDBOX_DEFAULT_CPU_LIMIT, + OPEN_SANDBOX_DEFAULT_ENTRYPOINT, + OPEN_SANDBOX_DEFAULT_LANGUAGE, + OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT, + OPEN_SANDBOX_DEFAULT_TEMPLATE, + OPEN_SANDBOX_DEFAULT_TIMEOUT, + OPEN_SANDBOX_EXECD_PORT, + OPEN_SANDBOX_POLL_INTERVAL, + OPEN_SANDBOX_READY_TIMEOUT, +) +from litellm.llms.base_llm.sandbox.transformation import ( + BaseSandboxConfig, + CodeExecutionResult, + ContainerHandle, + SANDBOX_MAX_OUTPUT_BYTES, +) +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.custom_http import httpxSpecialProvider + +DEFAULT_SANDBOX_TIMEOUT = OPEN_SANDBOX_DEFAULT_TIMEOUT +DEFAULT_READY_TIMEOUT = OPEN_SANDBOX_READY_TIMEOUT +DEFAULT_POLL_INTERVAL = OPEN_SANDBOX_POLL_INTERVAL +MAX_OUTPUT_BYTES = SANDBOX_MAX_OUTPUT_BYTES + + +class OpenSandboxSandboxConfig(BaseSandboxConfig): + def _http(self, client: AsyncHTTPHandler | None) -> AsyncHTTPHandler: + if client is not None: + return client + return get_async_httpx_client(llm_provider=httpxSpecialProvider.Sandbox) + + def validate_environment(self, api_key: str | None = None, **kwargs) -> str: + if api_key is not None: + return api_key + return get_secret_str(OPEN_SANDBOX_API_KEY_ENV_VAR) or "" + + async def acreate_sandbox( + self, + *, + template: str | None = None, + timeout: int | None = None, + allow_internet_access: bool | None = None, + api_key: str | None = None, + api_base: str | None = None, + metadata: dict[str, str] | None = None, + env_vars: dict[str, str] | None = None, + resource_limits: dict[str, str] | None = None, + resource_requests: dict[str, str] | None = None, + entrypoint: list[str] | tuple[str, ...] | None = None, + network_policy: dict[str, object] | None = None, + secure_access: bool = False, + use_server_proxy: bool = False, + ready_timeout: float | None = None, + poll_interval: float | None = None, + client: AsyncHTTPHandler | None = None, + **kwargs, + ) -> ContainerHandle: + key = self.validate_environment(api_key=api_key) + base = self._api_base(api_base) + ready_timeout_seconds = ( + float(ready_timeout) if ready_timeout is not None else DEFAULT_READY_TIMEOUT + ) + poll_interval_seconds = ( + float(poll_interval) if poll_interval is not None else DEFAULT_POLL_INTERVAL + ) + body = self._create_body( + template=template, + timeout=timeout, + allow_internet_access=allow_internet_access, + metadata=metadata, + env_vars=env_vars, + resource_limits=resource_limits, + resource_requests=resource_requests, + entrypoint=entrypoint, + network_policy=network_policy, + secure_access=secure_access, + ) + + response = cast( + httpx.Response, + await self._http(client).post( + url=f"{base}/sandboxes", + headers=self._lifecycle_headers(key), + json=body, + ), + ) + data = response.json() + sandbox_id = str(data["id"]) + + if self._sandbox_state(data) != "Running": + await self._wait_until_running( + sandbox_id=sandbox_id, + api_base=base, + headers=self._lifecycle_headers(key), + client=client, + ready_timeout=ready_timeout_seconds, + poll_interval=poll_interval_seconds, + ) + + endpoint, endpoint_headers = await self._wait_for_execd_endpoint( + sandbox_id=sandbox_id, + api_base=base, + headers=self._lifecycle_headers(key), + use_server_proxy=use_server_proxy, + client=client, + ready_timeout=ready_timeout_seconds, + poll_interval=poll_interval_seconds, + ) + + handle = ContainerHandle(id=sandbox_id, provider="opensandbox", domain=base) + handle._hidden_params = { + "api_base": base, + "api_key": key, + "execd_endpoint": endpoint, + "execd_headers": endpoint_headers, + "use_server_proxy": use_server_proxy, + } + return handle + + async def arun_code( + self, + *, + container: Union[ContainerHandle, str], + code: str, + api_key: str | None = None, + api_base: str | None = None, + language: str = OPEN_SANDBOX_DEFAULT_LANGUAGE, + use_server_proxy: bool = False, + ready_timeout: float | None = None, + poll_interval: float | None = None, + client: AsyncHTTPHandler | None = None, + **kwargs, + ) -> CodeExecutionResult: + handle = await self._ensure_handle( + container=container, + api_key=api_key, + api_base=api_base, + use_server_proxy=use_server_proxy, + ready_timeout=( + float(ready_timeout) + if ready_timeout is not None + else DEFAULT_READY_TIMEOUT + ), + poll_interval=( + float(poll_interval) + if poll_interval is not None + else DEFAULT_POLL_INTERVAL + ), + client=client, + ) + endpoint = str(handle._hidden_params["execd_endpoint"]) + endpoint_headers = self._as_str_dict(handle._hidden_params.get("execd_headers")) + base = str( + handle._hidden_params.get("api_base") + or handle.domain + or self._api_base(api_base) + ) + lines = await self._post_code( + url=f"{self._endpoint_base_url(endpoint, base)}/code", + headers={ + "Content-Type": "application/json", + "Accept": "text/event-stream", + "Cache-Control": "no-cache", + **endpoint_headers, + }, + body={ + "code": code, + "context": {"language": language}, + }, + client=client, + ) + return self._parse_lines(lines) + + async def adelete_sandbox( + self, + *, + container: Union[ContainerHandle, str], + api_key: str | None = None, + api_base: str | None = None, + client: AsyncHTTPHandler | None = None, + **kwargs, + ) -> bool: + handle = self._as_handle(container, api_base=api_base) + base = str(handle._hidden_params.get("api_base") or self._api_base(api_base)) + key = self._api_key(api_key=api_key, handle=handle) + try: + response = cast( + httpx.Response, + await self._http(client).delete( + url=f"{base}/sandboxes/{handle.id}", + headers=self._lifecycle_headers(key), + ), + ) + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + return False + raise + return 200 <= response.status_code < 300 + + async def _ensure_handle( + self, + *, + container: Union[ContainerHandle, str], + api_key: str | None, + api_base: str | None, + use_server_proxy: bool, + ready_timeout: float, + poll_interval: float, + client: AsyncHTTPHandler | None, + ) -> ContainerHandle: + handle = self._as_handle(container, api_base=api_base) + if handle._hidden_params.get("execd_endpoint"): + return handle + + base = str(handle._hidden_params.get("api_base") or self._api_base(api_base)) + key = self._api_key(api_key=api_key, handle=handle) + resolved_use_server_proxy = bool( + handle._hidden_params.get("use_server_proxy", use_server_proxy) + ) + endpoint, endpoint_headers = await self._wait_for_execd_endpoint( + sandbox_id=handle.id, + api_base=base, + headers=self._lifecycle_headers(key), + use_server_proxy=resolved_use_server_proxy, + client=client, + ready_timeout=ready_timeout, + poll_interval=poll_interval, + ) + handle.domain = base + handle._hidden_params = { + **handle._hidden_params, + "api_base": base, + "api_key": key, + "execd_endpoint": endpoint, + "execd_headers": endpoint_headers, + "use_server_proxy": resolved_use_server_proxy, + } + return handle + + async def _wait_until_running( + self, + *, + sandbox_id: str, + api_base: str, + headers: dict[str, str], + client: AsyncHTTPHandler | None, + ready_timeout: float, + poll_interval: float, + ) -> None: + deadline = time.monotonic() + ready_timeout + while True: + response = cast( + httpx.Response, + await self._http(client).get( + url=f"{api_base}/sandboxes/{sandbox_id}", + headers=headers, + ), + ) + data = response.json() + state = self._sandbox_state(data) + if state == "Running": + return + if state in {"Failed", "Stopping", "Terminated"}: + raise ValueError(f"OpenSandbox sandbox {sandbox_id} entered {state}") + if time.monotonic() >= deadline: + raise TimeoutError( + f"OpenSandbox sandbox {sandbox_id} was not Running within " + f"{ready_timeout} seconds" + ) + await asyncio.sleep(poll_interval) + + async def _wait_for_execd_endpoint( + self, + *, + sandbox_id: str, + api_base: str, + headers: dict[str, str], + use_server_proxy: bool, + client: AsyncHTTPHandler | None, + ready_timeout: float, + poll_interval: float, + ) -> tuple[str, dict[str, str]]: + deadline = time.monotonic() + ready_timeout + last_error: Exception | None = None + while True: + try: + return await self._get_execd_endpoint( + sandbox_id=sandbox_id, + api_base=api_base, + headers=headers, + use_server_proxy=use_server_proxy, + client=client, + ) + except httpx.HTTPStatusError as e: + if e.response.status_code != 404: + raise + last_error = e + except ValueError as e: + last_error = e + + if time.monotonic() >= deadline: + raise TimeoutError( + f"OpenSandbox execd endpoint for {sandbox_id} was not ready within " + f"{ready_timeout} seconds" + ) from last_error + await asyncio.sleep(poll_interval) + + async def _get_execd_endpoint( + self, + *, + sandbox_id: str, + api_base: str, + headers: dict[str, str], + use_server_proxy: bool, + client: AsyncHTTPHandler | None, + ) -> tuple[str, dict[str, str]]: + response = cast( + httpx.Response, + await self._http(client).get( + url=f"{api_base}/sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}", + headers=headers, + params={"use_server_proxy": use_server_proxy}, + ), + ) + data = response.json() + endpoint = data.get("endpoint") + if not endpoint: + raise ValueError( + f"OpenSandbox did not return an execd endpoint for {sandbox_id}" + ) + return str(endpoint), self._as_str_dict(data.get("headers")) + + async def _post_code( + self, + *, + url: str, + headers: dict[str, str], + body: dict[str, object], + client: AsyncHTTPHandler | None, + ) -> list[str]: + timeout = httpx.Timeout(connect=30.0, read=None, write=30.0, pool=None) + response = cast( + httpx.Response, + await self._http(client).post( + url=url, + headers=headers, + timeout=timeout, + json=body, + stream=True, + ), + ) + return await self._read_capped_lines(response) + + def _api_key(self, *, api_key: str | None, handle: ContainerHandle) -> str: + if api_key is not None: + return api_key + if "api_key" in handle._hidden_params: + return str(handle._hidden_params["api_key"]) + return self.validate_environment() + + @staticmethod + def _create_body( + *, + template: str | None, + timeout: int | None, + allow_internet_access: bool | None, + metadata: dict[str, str] | None, + env_vars: dict[str, str] | None, + resource_limits: dict[str, str] | None, + resource_requests: dict[str, str] | None, + entrypoint: list[str] | tuple[str, ...] | None, + network_policy: dict[str, object] | None, + secure_access: bool, + ) -> dict[str, object]: + body: dict[str, object] = { + "image": {"uri": template or OPEN_SANDBOX_DEFAULT_TEMPLATE}, + "entrypoint": list(entrypoint or OPEN_SANDBOX_DEFAULT_ENTRYPOINT), + "timeout": timeout if timeout is not None else DEFAULT_SANDBOX_TIMEOUT, + "resourceLimits": resource_limits + or OpenSandboxSandboxConfig._default_resource_limits(), + } + if metadata: + body["metadata"] = metadata + if env_vars: + body["env"] = env_vars + if resource_requests: + body["resourceRequests"] = resource_requests + if network_policy is not None: + body["networkPolicy"] = network_policy + elif allow_internet_access is not True: + body["networkPolicy"] = {"defaultAction": "deny", "egress": []} + if secure_access: + body["secureAccess"] = True + return body + + @staticmethod + def _default_resource_limits() -> dict[str, str]: + return { + "cpu": OPEN_SANDBOX_DEFAULT_CPU_LIMIT, + "memory": OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT, + } + + @staticmethod + def _sandbox_state(data: object) -> str | None: + if not isinstance(data, dict): + return None + status = data.get("status") + if not isinstance(status, dict): + return None + state = status.get("state") + return str(state) if state is not None else None + + @staticmethod + def _as_str_dict(value: object) -> dict[str, str]: + if not isinstance(value, dict): + return {} + return {str(k): str(v) for k, v in value.items()} + + @staticmethod + def _api_base(api_base: str | None) -> str: + base = api_base or get_secret_str(OPEN_SANDBOX_API_BASE_ENV_VAR) + if not base: + raise ValueError( + "OpenSandbox api_base is required. Pass api_base or set " + f"{OPEN_SANDBOX_API_BASE_ENV_VAR}." + ) + return str(base).rstrip("/") + + @staticmethod + def _lifecycle_headers(api_key: str) -> dict[str, str]: + headers = {"Content-Type": "application/json"} + if api_key: + headers["OPEN-SANDBOX-API-KEY"] = api_key + return headers + + @staticmethod + def _endpoint_base_url(endpoint: str, api_base: str) -> str: + normalized_endpoint = endpoint.rstrip("/") + if normalized_endpoint.startswith(("http://", "https://")): + return normalized_endpoint + protocol = api_base.split("://", 1)[0] if "://" in api_base else "http" + return f"{protocol}://{normalized_endpoint}" + + @staticmethod + def _as_handle( + container: Union[ContainerHandle, str], *, api_base: str | None + ) -> ContainerHandle: + if isinstance(container, ContainerHandle): + return container + handle = ContainerHandle( + id=str(container), + provider="opensandbox", + domain=OpenSandboxSandboxConfig._api_base(api_base), + ) + handle._hidden_params = {} + return handle + + @staticmethod + def _parse_lines(lines: list[str]) -> CodeExecutionResult: + messages = tuple( + event + for line in lines + if (event := OpenSandboxSandboxConfig._parse_sse_line(line)) is not None + ) + + def of_type(message_type: str): + return (m for m in messages if m.get("type") == message_type) + + error = next( + (OpenSandboxSandboxConfig._normalize_error(m) for m in of_type("error")), + None, + ) + execution_count = next( + ( + OpenSandboxSandboxConfig._as_int(m.get("execution_count")) + for m in of_type("execution_count") + if OpenSandboxSandboxConfig._as_int(m.get("execution_count")) + is not None + ), + None, + ) + + return CodeExecutionResult( + stdout="".join(str(m.get("text", "")) for m in of_type("stdout")), + stderr="".join(str(m.get("text", "")) for m in of_type("stderr")), + results=[ + OpenSandboxSandboxConfig._normalize_result(m) for m in of_type("result") + ], + error=error, + execution_count=execution_count, + ) + + @staticmethod + def _parse_sse_line(line: str) -> dict[str, object] | None: + stripped = line.strip() + if not stripped or stripped.startswith( + ( + ":", + "event:", + "id:", + "retry:", + ) + ): + return None + data = stripped[5:].strip() if stripped.startswith("data:") else stripped + if not data: + return None + try: + parsed = json.loads(data) + except json.JSONDecodeError: + return None + if not isinstance(parsed, dict): + return None + if "type" not in parsed and "code" in parsed and "message" in parsed: + return { + "type": "error", + "error": { + "ename": str(parsed["code"]), + "evalue": str(parsed["message"]), + "traceback": [], + }, + } + return parsed + + @staticmethod + def _normalize_result(message: dict[str, object]) -> dict[str, object]: + results = message.get("results") + if isinstance(results, dict): + return {str(k): v for k, v in results.items()} + return { + str(k): v + for k, v in message.items() + if k not in {"type", "timestamp", "execution_count"} + } + + @staticmethod + def _normalize_error(message: dict[str, object]) -> dict[str, object]: + raw_error = message.get("error") + if isinstance(raw_error, dict): + name = OpenSandboxSandboxConfig._first_non_none_value( + raw_error, "ename", "name", default="" + ) + value = OpenSandboxSandboxConfig._first_non_none_value( + raw_error, "evalue", "value", default="" + ) + traceback = OpenSandboxSandboxConfig._first_non_none_value( + raw_error, "traceback", default=[] + ) + return { + "name": name, + "value": value, + "traceback": traceback, + } + return { + "name": OpenSandboxSandboxConfig._first_non_none_value( + message, "name", default="" + ), + "value": OpenSandboxSandboxConfig._first_non_none_value( + message, "value", "text", default="" + ), + "traceback": OpenSandboxSandboxConfig._first_non_none_value( + message, "traceback", default=[] + ), + } + + @staticmethod + def _as_int(value: object) -> int | None: + if isinstance(value, int): + return value + if isinstance(value, str): + try: + return int(value) + except ValueError: + return None + return None + + @staticmethod + def _first_non_none_value( + values: dict[str, object], *keys: str, default: object + ) -> object: + return next( + (values[key] for key in keys if key in values and values[key] is not None), + default, + ) diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index 85602bf1d86..35a0d84df40 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -67,10 +67,12 @@ class ParallelAISearchConfig(BaseSearchConfig): api_base: Optional[str] = None, **kwargs, ) -> Dict: - api_key = ( - api_key - or get_secret_str("PARALLEL_AI_API_KEY") - or get_secret_str("PARALLEL_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"), + base_env_var="PARALLEL_AI_API_BASE", + default_api_base=self.PARALLEL_AI_API_BASE, ) if not api_key: raise ValueError( diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index bf055f91aa0..ec7ec397ea6 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -98,10 +98,11 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: if num_search_queries > 0 and search_cost_value is not None: # Handle both dict and float formats if isinstance(search_cost_value, dict): - # Use the "low" size as default - tests expect 0.005 / 1000 - search_cost_per_query = ( - _safe_float_cast(search_cost_value.get("search_context_size_low", 0)) - / 1000 + # search_context_cost_per_query stores the per-request price in USD + # (e.g. sonar low = $0.005/request). Use it directly, matching the + # gemini cost calculator which reads the same field per request. + search_cost_per_query = _safe_float_cast( + search_cost_value.get("search_context_size_low", 0) ) else: search_cost_per_query = _safe_float_cast(search_cost_value) diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py index ea96f87957c..55de52c5384 100644 --- a/litellm/llms/perplexity/search/transformation.py +++ b/litellm/llms/perplexity/search/transformation.py @@ -50,7 +50,13 @@ class PerplexitySearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("PERPLEXITYAI_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("PERPLEXITYAI_API_KEY",), + base_env_var="PERPLEXITY_API_BASE", + default_api_base=self.PERPLEXITY_API_BASE, + ) if not api_key: raise ValueError( "PERPLEXITYAI_API_KEY is not set. Set `PERPLEXITYAI_API_KEY` environment variable." diff --git a/litellm/llms/searchapi/search/transformation.py b/litellm/llms/searchapi/search/transformation.py index c04e1377f9c..ae8413684cc 100644 --- a/litellm/llms/searchapi/search/transformation.py +++ b/litellm/llms/searchapi/search/transformation.py @@ -74,7 +74,13 @@ class SearchAPIConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("SEARCHAPI_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("SEARCHAPI_API_KEY",), + base_env_var="SEARCHAPI_API_BASE", + default_api_base=self.SEARCHAPI_API_BASE, + ) if not api_key: raise ValueError( @@ -114,6 +120,7 @@ class SearchAPIConfig(BaseSearchConfig): query: Union[str, List[str]], optional_params: dict, api_key: Optional[str] = None, + api_base: str | None = None, search_engine_id: Optional[str] = None, **kwargs, ) -> Dict: @@ -137,8 +144,16 @@ class SearchAPIConfig(BaseSearchConfig): if isinstance(query, list): query = " ".join(query) - # Get API key from parameter or environment - api_key = api_key or get_secret_str("SEARCHAPI_API_KEY") + # Get API key from parameter or environment. The key is sent as a query + # param to api_base, so resolve it host-aware to avoid leaking a + # server-managed key to a caller-supplied host. + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("SEARCHAPI_API_KEY",), + base_env_var="SEARCHAPI_API_BASE", + default_api_base=self.SEARCHAPI_API_BASE, + ) if not api_key: raise ValueError( "SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable." diff --git a/litellm/llms/searxng/search/transformation.py b/litellm/llms/searxng/search/transformation.py index ee6f3895721..ff68be5709e 100644 --- a/litellm/llms/searxng/search/transformation.py +++ b/litellm/llms/searxng/search/transformation.py @@ -61,7 +61,13 @@ class SearXNGSearchConfig(BaseSearchConfig): Some instances may require authentication via headers. """ # SearXNG typically doesn't require API keys, but support optional auth - api_key = api_key or get_secret_str("SEARXNG_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("SEARXNG_API_KEY",), + base_env_var="SEARXNG_API_BASE", + default_api_base=None, + ) if api_key: headers["Authorization"] = f"Bearer {api_key}" headers["Content-Type"] = "application/json" diff --git a/litellm/llms/serper/search/transformation.py b/litellm/llms/serper/search/transformation.py index 0daccbe652b..dd43f2d2dc9 100644 --- a/litellm/llms/serper/search/transformation.py +++ b/litellm/llms/serper/search/transformation.py @@ -55,7 +55,13 @@ class SerperSearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("SERPER_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("SERPER_API_KEY",), + base_env_var="SERPER_API_BASE", + default_api_base=self.SERPER_API_BASE, + ) if not api_key: raise ValueError( "SERPER_API_KEY is not set. Set `SERPER_API_KEY` environment variable." diff --git a/litellm/llms/tavily/search/transformation.py b/litellm/llms/tavily/search/transformation.py index ec96db96f36..647cfb5fa84 100644 --- a/litellm/llms/tavily/search/transformation.py +++ b/litellm/llms/tavily/search/transformation.py @@ -64,7 +64,13 @@ class TavilySearchConfig(BaseSearchConfig): """ Validate environment and return headers. """ - api_key = api_key or get_secret_str("TAVILY_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("TAVILY_API_KEY",), + base_env_var="TAVILY_API_BASE", + default_api_base=self.TAVILY_API_BASE, + ) if not api_key: raise ValueError( "TAVILY_API_KEY is not set. Set `TAVILY_API_KEY` environment variable." diff --git a/litellm/llms/tinyfish/search/transformation.py b/litellm/llms/tinyfish/search/transformation.py index c4949380e3a..b92f7ca1aff 100644 --- a/litellm/llms/tinyfish/search/transformation.py +++ b/litellm/llms/tinyfish/search/transformation.py @@ -67,7 +67,13 @@ class TinyfishSearchConfig(BaseSearchConfig): api_base: str | None = None, **kwargs: object, ) -> dict[str, str]: - resolved_key = api_key or get_secret_str("TINYFISH_API_KEY") + resolved_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("TINYFISH_API_KEY",), + base_env_var="TINYFISH_API_BASE", + default_api_base=self.TINYFISH_API_BASE, + ) if not resolved_key: raise ValueError( "TINYFISH_API_KEY is not set. Set `TINYFISH_API_KEY` environment variable." diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index d6441db7856..1fe9f15c9f0 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -32,6 +32,9 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): self._project = project self._location = location + def _include_function_response_id(self) -> bool: + return False + # ------------------------------------------------------------------ # URL # ------------------------------------------------------------------ diff --git a/litellm/llms/you_com/search/transformation.py b/litellm/llms/you_com/search/transformation.py index 3c94b991735..0c7916e4c05 100644 --- a/litellm/llms/you_com/search/transformation.py +++ b/litellm/llms/you_com/search/transformation.py @@ -64,7 +64,13 @@ class YouComSearchConfig(BaseSearchConfig): endpoint with the `X-API-Key` header. Otherwise fall through to the keyless free tier; no auth header is required. """ - api_key = api_key or get_secret_str("YOUCOM_API_KEY") + api_key = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("YOUCOM_API_KEY",), + base_env_var="YOUCOM_API_BASE", + default_api_base=self.YOU_COM_API_BASE, + ) headers["Content-Type"] = "application/json" # Pin Accept-Encoding to identity: the keyless `api.you.com/v1/agents/search` # endpoint advertises gzip content-encoding but returns body bytes the diff --git a/litellm/main.py b/litellm/main.py index 63c5798e70a..c3d7ca28c49 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -81,11 +81,17 @@ from litellm.constants import ( from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + maybe_run_chat_completion_agentic_loop, +) from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, get_audio_file_for_health_check, ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) from litellm.litellm_core_utils.get_litellm_params import OPTIONAL_KWARGS_KEYS from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_provider_specific_headers import ( @@ -118,6 +124,10 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.realtime_api.main import _realtime_health_check from litellm.secret_managers.main import get_secret_bool, get_secret_str +from litellm.types.completion import ( + _CompletionDispatchContext, + _CompletionDispatchResult, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CustomPricingLiteLLMParams, @@ -650,6 +660,39 @@ async def acompletion( response_object=response, model_response_object=litellm.ModelResponse(), ) + # Provider-agnostic dispatch point for the chat-completions agentic loop + # (code-interpreter interception, etc). Chat routing forks per provider + # before this (OpenAI goes through the OpenAI SDK in openai.py, others + # through the shared httpx handler), so a dispatch inside any single + # provider handler would miss the others. Here is where every fork + # reconverges, so the loop runs once for all providers. Responses needs + # no equivalent: every provider already funnels through one shared + # handler where the loop is dispatched. + if isinstance(response, litellm.ModelResponse): + looped = await maybe_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + optional_params={ + k: v + for k, v in completion_kwargs.items() + if v is not None + and k + not in ( + "model", + "messages", + "stream", + "acompletion", + "deployment_id", + ) + }, + kwargs=kwargs, + logging_obj=kwargs.get("litellm_logging_obj"), + custom_llm_provider=custom_llm_provider, + stream=bool(stream), + ) + if looped is not None: + response = looped if isinstance(response, CustomStreamWrapper): response.set_logging_event_loop( loop=loop @@ -1084,6 +1127,3825 @@ def _build_custom_pricing_entry( return entry +def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + _azure_detection_model = ctx._azure_detection_model + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + api_version = ctx.api_version + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + max_retries = ctx.max_retries + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + dynamic_params = False + if client is not None and ( + isinstance(client, openai.AzureOpenAI) + or isinstance(client, openai.AsyncAzureOpenAI) + ): + dynamic_params = _check_dynamic_azure_params( + azure_client_params={"api_version": api_version}, + azure_client=client, + ) + + api_type = get_secret("AZURE_API_TYPE") or "azure" + + api_base = api_base or litellm.api_base or get_secret("AZURE_API_BASE") + + api_version = ( + api_version + or litellm.api_version + or get_secret_str("AZURE_API_VERSION") + or litellm.AZURE_DEFAULT_API_VERSION + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.azure_key + or get_secret_str("AZURE_OPENAI_API_KEY") + or get_secret_str("AZURE_API_KEY") + ) + + azure_ad_token = optional_params.get("extra_body", {}).pop( + "azure_ad_token", None + ) or get_secret_str("AZURE_AD_TOKEN") + + azure_ad_token_provider = litellm_params.get("azure_ad_token_provider", None) + + headers = headers or litellm.headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + if max_retries is not None: + optional_params["max_retries"] = max_retries + + if litellm.AzureOpenAIO1Config().is_o_series_model(model=_azure_detection_model): + ## LOAD CONFIG - if set + config = litellm.AzureOpenAIO1Config.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + response = azure_o1_chat_completions.completion( + model=model, + messages=messages, + headers=headers, + api_key=api_key, + api_base=api_base, + api_version=api_version, + dynamic_params=dynamic_params, + azure_ad_token=azure_ad_token, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, # type: ignore + client=client, # pass AsyncAzureOpenAI, AzureOpenAI client + custom_llm_provider=custom_llm_provider, + ) + else: + ## LOAD CONFIG - if set + config = litellm.AzureOpenAIConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + ## COMPLETION CALL + response = azure_chat_completions.completion( + model=model, + messages=messages, + headers=headers, + api_key=api_key, + api_base=api_base, + api_version=api_version, + api_type=api_type, + dynamic_params=dynamic_params, + azure_ad_token=azure_ad_token, + azure_ad_token_provider=azure_ad_token_provider, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, # type: ignore + client=client, # pass AsyncAzureOpenAI, AzureOpenAI client + ) + + if optional_params.get("stream", False): + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={ + "headers": headers, + "api_version": api_version, + "api_base": api_base, + }, + ) + + return response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_azure_text(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + api_version = ctx.api_version + client = ctx.client + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + api_type = get_secret_str("AZURE_API_TYPE") or "azure" + + api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") + + if api_base is None: + raise ValueError( + "api_base is required for Azure OpenAI LLM provider. Either set it dynamically or set the AZURE_API_BASE environment variable." + ) + + api_version = ( + api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.azure_key + or get_secret_str("AZURE_OPENAI_API_KEY") + or get_secret_str("AZURE_API_KEY") + ) + + azure_ad_token = optional_params.get("extra_body", {}).pop( + "azure_ad_token", None + ) or get_secret_str("AZURE_AD_TOKEN") + + azure_ad_token_provider = litellm_params.get("azure_ad_token_provider", None) + + headers = headers or litellm.headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + + ## LOAD CONFIG - if set + config = litellm.AzureOpenAIConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + ## COMPLETION CALL + response = azure_text_completions.completion( + model=model, + messages=messages, + headers=headers, + api_key=api_key, + api_base=api_base, + api_version=cast(str, api_version), + api_type=api_type, + azure_ad_token=azure_ad_token, + azure_ad_token_provider=azure_ad_token_provider, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, + client=client, # pass AsyncAzureOpenAI, AzureOpenAI client + ) + + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={ + "headers": headers, + "api_version": api_version, + "api_base": api_base, + }, + ) + + return response + + +def _complete_deepseek(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_azure_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + + azure_ai_route = AzureFoundryModelInfo.get_azure_ai_route(model) + + # Check if this is an agents route - model format: azure_ai/agents/ + if azure_ai_route == "agents": + from litellm.llms.azure_ai.agents import AzureAIAgentsConfig + + api_base = AzureFoundryModelInfo.get_api_base(api_base) + if api_base is None: + raise ValueError( + "Azure AI Agents requests require an api_base. " + "Set `api_base` or the AZURE_AI_API_BASE env var." + ) + api_key = AzureFoundryModelInfo.get_api_key(api_key) + + response = AzureAIAgentsConfig.completion( + model=model, + messages=messages, + api_base=api_base, + api_key=api_key, + model_response=model_response, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, + acompletion=acompletion, + stream=stream, + headers=headers or litellm.headers, + ) + + # Check if this is a Claude model - route to Azure Anthropic handler + elif "claude" in model.lower(): + # Use Azure Anthropic handler for Claude models + api_base = AzureFoundryModelInfo.get_api_base(api_base) + if api_base is None: + raise ValueError( + "Azure Anthropic requests require an api_base. " + "Set `api_base` or the AZURE_AI_API_BASE env var." + ) + api_key = AzureFoundryModelInfo.get_api_key(api_key) + + # Ensure the URL ends with /v1/messages for Anthropic + if api_base: + api_base = api_base.rstrip("/") + if not api_base.endswith("/v1/messages"): + if "/anthropic" in api_base: + parts = api_base.split("/anthropic", 1) + api_base = parts[0] + "/anthropic" + else: + api_base = api_base + "/anthropic" + api_base = api_base + "/v1/messages" + + response = azure_anthropic_chat_completions.completion( + model=model, + messages=messages, + api_base=api_base, + acompletion=acompletion, + custom_prompt_dict=litellm.custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + headers=headers, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + ) + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + ) + response = response + else: + # Non-Claude models use standard Azure AI flow + api_base = AzureFoundryModelInfo.get_api_base(api_base) + # set API KEY + api_key = AzureFoundryModelInfo.get_api_key(api_key) + + headers = headers or litellm.headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + + ## FOR COHERE + if "command-r" in model: # make sure tool call in messages are str + messages = stringify_json_tool_call_content(messages=messages) + + ## COMPLETION CALL + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, # pass AsyncOpenAI, OpenAI client + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + if optional_params.get("stream", False): + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={"headers": headers}, + ) + + return response + + +def _complete_text_completion_openai( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + text_completion = ctx.text_completion + timeout = ctx.timeout + + openai.api_type = "openai" + + api_base = ( + api_base + or litellm.api_base + or get_secret("OPENAI_BASE_URL") + or get_secret("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + + openai.api_version = None + # set API KEY + + api_key = ( + api_key or litellm.api_key or litellm.openai_key or get_secret("OPENAI_API_KEY") + ) + + headers = headers or litellm.headers + + ## LOAD CONFIG - if set + config = litellm.OpenAITextCompletionConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > openai_text_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + if litellm.organization: + openai.organization = litellm.organization + + ## COMPLETION CALL + _response = openai_text_completions.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + acompletion=acompletion, + client=client, # pass AsyncOpenAI, OpenAI client + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + ) + + if ( + optional_params.get("stream", False) is False + and acompletion is False + and text_completion is False + ): + # convert to chat completion response + _response = ( + litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( + response_object=_response, model_response_object=model_response + ) + ) + + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=_response, + additional_args={"headers": headers}, + ) + return _response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_fireworks_ai( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_heroku(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_ragflow(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_xai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_groq(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there + or litellm.api_base + or get_secret("GROQ_API_BASE") + or "https://api.groq.com/openai/v1" + ) + + # set API KEY + api_key = ( + api_key + or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there + or litellm.groq_key + or get_secret("GROQ_API_KEY") + ) + + headers = headers or litellm.headers + + ## LOAD CONFIG - if set + config = litellm.GroqChatConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_bedrock_mantle( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = api_base or litellm.api_base or get_secret("BEDROCK_MANTLE_API_BASE") + api_key = api_key or litellm.api_key or get_secret("BEDROCK_MANTLE_API_KEY") + headers = headers or litellm.headers + config = litellm.BedrockMantleChatConfig.get_config() + for k, v in config.items(): + if k not in optional_params: + optional_params[k] = v + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + + +def _complete_a2a(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + ( + api_base, + api_key, + headers, + ) = litellm.A2AConfig.resolve_agent_config_from_registry( + model=model, + api_base=api_base, + api_key=api_key, + headers=headers, + optional_params=optional_params, + ) + + # Fall back to environment variables and defaults + api_base = api_base or litellm.api_base or get_secret_str("A2A_API_BASE") + + if api_base is None: + raise Exception( + "api_base is required for A2A provider. " + "Either provide api_base parameter, set A2A_API_BASE environment variable, " + "or register the agent in the proxy with model='a2a/'." + ) + + headers = headers or litellm.headers + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + provider_config=provider_config, + ) + + +def _complete_gigachat(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.api_key + or litellm.gigachat_key + or get_secret("GIGACHAT_API_KEY") + or get_secret("GIGACHAT_CREDENTIALS") + ) + + headers = headers or litellm.headers or {} + + ## COMPLETION CALL + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_sap(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + headers = headers or litellm.headers + ## LOAD CONFIG - if set + config = litellm.GenAIHubOrchestrationConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + return sap_gen_ai_hub_chat_completions.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + shared_session=shared_session, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + api_key=api_key, + api_base=api_base, + stream=stream, + ) + + +def _complete_aiohttp_openai( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there + or litellm.api_base + or get_secret("OPENAI_BASE_URL") + or get_secret("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + # set API KEY + api_key = ( + api_key + or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there + or litellm.openai_key + or get_secret("OPENAI_API_KEY") + ) + + headers = headers or litellm.headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + return base_llm_aiohttp_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + + +def _complete_cometapi(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.cometapi_key + or get_secret_str("COMETAPI_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("COMETAPI_API_BASE") + or "https://api.cometapi.com/v1" + ) + + ## COMPLETION CALL + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + + ## LOGGING + logging.post_call(input=messages, api_key=api_key, original_response=response) + + return response + + +def _complete_minimax(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("MINIMAX_API_BASE") + or "https://api.minimax.io/v1" + ) + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + model_response=model_response, + encoding=_get_encoding(), + logging_obj=logging, + optional_params=optional_params, + timeout=timeout, + litellm_params=litellm_params, + shared_session=shared_session, + acompletion=acompletion, + stream=stream, + api_key=api_key, + headers=headers, + client=client, + provider_config=provider_config, + ) + logging.post_call(input=messages, api_key=api_key, original_response=response) + + return response + + +def _complete_hosted_vllm(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = api_base or litellm.api_base or get_secret_str("HOSTED_VLLM_API_BASE") + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + model_response=model_response, + encoding=_get_encoding(), + logging_obj=logging, + optional_params=optional_params, + timeout=timeout, + litellm_params=litellm_params, + shared_session=shared_session, + acompletion=acompletion, + stream=stream, + api_key=api_key, + headers=headers, + client=client, + provider_config=provider_config, + ) + logging.post_call(input=messages, api_key=api_key, original_response=response) + + return response + + +def _complete_custom_openai( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + metadata = ctx.metadata + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + organization = ctx.organization + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there + or litellm.api_base + or get_secret("OPENAI_BASE_URL") + or get_secret("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + organization = ( + organization + or litellm.organization + or get_secret("OPENAI_ORGANIZATION") + or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 + ) + openai.organization = organization + # set API KEY + api_key = ( + api_key + or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there + or litellm.openai_key + or get_secret("OPENAI_API_KEY") + ) + + headers = headers or litellm.headers + + # Add GitHub Copilot headers (same as /responses endpoint does) + if custom_llm_provider == "github_copilot": + from litellm.llms.github_copilot.authenticator import Authenticator + from litellm.llms.github_copilot.common_utils import ( + get_copilot_default_headers, + ) + + copilot_auth = Authenticator() + copilot_api_key = copilot_auth.get_api_key() + copilot_headers = get_copilot_default_headers(copilot_api_key) + if extra_headers: + copilot_headers.update(extra_headers) + extra_headers = copilot_headers + + if extra_headers is not None: + optional_params["extra_headers"] = extra_headers + + if ( + litellm.enable_preview_features and metadata is not None + ): # [PREVIEW] allow metadata to be passed to OPENAI + openai_metadata = get_requester_metadata(metadata) + if openai_metadata is not None: + optional_params["metadata"] = openai_metadata + + ## LOAD CONFIG - if set + config = litellm.OpenAIConfig.get_config() + for k, v in config.items(): + if ( + k not in optional_params + ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + ## COMPLETION CALL + use_base_llm_http_handler = get_secret_bool( + "EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER" + ) + + try: + if use_base_llm_http_handler: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + model_response=model_response, + encoding=_get_encoding(), + logging_obj=logging, + optional_params=optional_params, + timeout=timeout, + litellm_params=litellm_params, + shared_session=shared_session, + acompletion=acompletion, + stream=stream, + api_key=api_key, + headers=headers, + client=client, + provider_config=provider_config, + ) + else: + response = openai_chat_completions.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + custom_prompt_dict=custom_prompt_dict, + client=client, # pass AsyncOpenAI, OpenAI client + organization=organization, + custom_llm_provider=custom_llm_provider, + shared_session=shared_session, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + if optional_params.get("stream", False): + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={"headers": headers}, + ) + + return response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_mistral(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = api_key or litellm.api_key or get_secret("MISTRAL_API_KEY") + api_base = ( + api_base + or litellm.api_base + or get_secret("MISTRAL_API_BASE") + or "https://api.mistral.ai/v1" + ) + + return base_llm_http_handler.completion( + model=model, + messages=messages, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + model_response=model_response, + encoding=_get_encoding(), + logging_obj=logging, + optional_params=optional_params, + timeout=timeout, + litellm_params=litellm_params, + shared_session=shared_session, + acompletion=acompletion, + stream=stream, + api_key=api_key, + headers=headers, + client=client, + provider_config=provider_config, + ) + + +def _complete_replicate(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + replicate_key = ( + api_key + or litellm.replicate_key + or litellm.api_key + or get_secret("REPLICATE_API_KEY") + or get_secret("REPLICATE_API_TOKEN") + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret("REPLICATE_API_BASE") + or "https://api.replicate.com/v1" + ) + + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + + model_response = replicate_chat_completion( # type: ignore + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), # for calculating input/output tokens + api_key=replicate_key, + logging_obj=logging, + custom_prompt_dict=custom_prompt_dict, + acompletion=acompletion, + headers=headers, + ) + + if optional_params.get("stream", False) is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=replicate_key, + original_response=model_response, + ) + + return model_response + + +def _complete_anthropic_text( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.anthropic_key + or litellm.api_key + or os.environ.get("ANTHROPIC_API_KEY") + ) + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + api_base = cast( + Optional[str], + api_base + or litellm.api_base + or get_secret("ANTHROPIC_API_BASE") + or get_secret("ANTHROPIC_BASE_URL") + or "https://api.anthropic.com/v1/complete", + ) + + # Check if we should disable automatic URL suffix appending + disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") + if ( + api_base is not None + and not disable_url_suffix + and not api_base.endswith("/v1/complete") + ): + api_base += "/v1/complete" + elif disable_url_suffix: + verbose_logger.debug( + "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/complete suffix" + ) + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="anthropic_text", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + ) + + +def _complete_anthropic(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.anthropic_key + or litellm.api_key + or os.environ.get("ANTHROPIC_API_KEY") + ) + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + # call /messages + # default route for all anthropic models + api_base = cast( + Optional[str], + api_base + or litellm.api_base + or get_secret("ANTHROPIC_API_BASE") + or get_secret("ANTHROPIC_BASE_URL") + or "https://api.anthropic.com/v1/messages", + ) + + # Check if we should disable automatic URL suffix appending + disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") + if ( + api_base is not None + and not disable_url_suffix + and not api_base.endswith("/v1/messages") + ): + api_base += "/v1/messages" + elif disable_url_suffix: + verbose_logger.debug( + "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/messages suffix" + ) + + response = anthropic_chat_completions.completion( + model=model, + messages=messages, + api_base=api_base, + acompletion=acompletion, + custom_prompt_dict=litellm.custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), # for calculating input/output tokens + api_key=api_key, + logging_obj=logging, + headers=headers, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + ) + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + ) + return response + + +def _complete_nlp_cloud(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + nlp_cloud_key = ( + api_key + or litellm.nlp_cloud_key + or get_secret("NLP_CLOUD_API_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret("NLP_CLOUD_API_BASE") + or "https://api.nlpcloud.io/v1/gpu/" + ) + + response = nlp_cloud_chat_completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + api_key=nlp_cloud_key, + logging_obj=logging, + ) + + if "stream" in optional_params and optional_params["stream"] is True: + # don't try to access stream object, + response = CustomStreamWrapper( + response, + model, + custom_llm_provider="nlp_cloud", + logging_obj=logging, + ) + + if optional_params.get("stream", False) or acompletion is True: + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + ) + + return response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_aleph_alpha(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + api_key = ctx.api_key + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + aleph_alpha_key = ( + api_key + or litellm.aleph_alpha_key + or get_secret("ALEPH_ALPHA_API_KEY") + or get_secret("ALEPHALPHA_API_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret("ALEPH_ALPHA_API_BASE") + or "https://api.aleph-alpha.com/complete" + ) + + model_response = aleph_alpha.completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + default_max_tokens_to_sample=litellm.max_tokens, + api_key=aleph_alpha_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + ) + + if "stream" in optional_params and optional_params["stream"] is True: + # don't try to access stream object, + return CustomStreamWrapper( + model_response, + model, + custom_llm_provider="aleph_alpha", + logging_obj=logging, + ) + return model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_cohere_chat(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + extra_headers = ctx.extra_headers + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + cohere_key = ( + api_key + or litellm.cohere_key + or get_secret_str("COHERE_API_KEY") + or get_secret_str("CO_API_KEY") + or litellm.api_key + ) + + cohere_route = CohereModelInfo.get_cohere_route(model) + verbose_logger.debug(f"Cohere route: {cohere_route}") + # Set API base based on route + if cohere_route == "v2": + api_base = ( + api_base + or litellm.api_base + or get_secret_str("COHERE_API_BASE") + or "https://api.cohere.com/v2/chat" + ) + # Remove v2/ prefix from model name for the actual API call + if "v2/" in model: + model = model.replace("v2/", "") + else: + api_base = ( + api_base + or litellm.api_base + or get_secret_str("COHERE_API_BASE") + or "https://api.cohere.ai/v1/chat" + ) + + headers = headers or litellm.headers or {} + if headers is None: + headers = {} + + if extra_headers is not None: + headers.update(extra_headers) + + verbose_logger.debug(f"Model: {model}, API Base: {api_base}") + verbose_logger.debug(f"Provider Config: {provider_config}") + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="cohere_chat", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=cohere_key, + provider_config=provider_config, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + ) + + +def _complete_maritalk(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + maritalk_key = ( + api_key + or litellm.maritalk_key + or get_secret("MARITALK_API_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret("MARITALK_API_BASE") + or "https://chat.maritaca.ai/api" + ) + + return openai_like_chat_completion.completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + api_key=maritalk_key, + logging_obj=logging, + custom_llm_provider="maritalk", + custom_prompt_dict=custom_prompt_dict, + ) + + +def _complete_amazon_nova(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + api_key = ctx.api_key + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.amazon_nova_api_key + or get_secret_str("AMAZON_NOVA_API_KEY") + or litellm.api_key + ) + api_base = ( + api_base + or litellm.api_base + or get_secret_str("AMAZON_NOVA_API_BASE") + or "https://api.nova.amazon.com/v1" + ) + return openai_like_chat_completion.completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + custom_prompt_dict=custom_prompt_dict, + ) + + +def _complete_huggingface(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + huggingface_key = ( + api_key + or litellm.huggingface_key + or os.environ.get("HF_TOKEN") + or os.environ.get("HUGGINGFACE_API_KEY") + or litellm.api_key + ) + hf_headers = headers or litellm.headers + return base_llm_http_handler.completion( + model=model, + messages=messages, + headers=hf_headers, + model_response=model_response, + api_key=huggingface_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + + +def _complete_oci(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + return base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + + +def _complete_compactifai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + stream = ctx.stream + timeout = ctx.timeout + + api_key = api_key or get_secret_str("COMPACTIFAI_API_KEY") or litellm.api_key + + api_base = api_base or "https://api.compactif.ai/v1" + + ## COMPLETION CALL + return base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + + +def _complete_oobabooga(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + model_response = oobabooga.completion( + model=model, + messages=messages, + model_response=model_response, + api_base=api_base, # type: ignore + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=None, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + ) + if "stream" in optional_params and optional_params["stream"] is True: + # don't try to access stream object, + return CustomStreamWrapper( + model_response, + model, + custom_llm_provider="oobabooga", + logging_obj=logging, + ) + return model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_databricks(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base # for databricks we check in get_llm_provider and pass in the api base from there + or litellm.api_base + or os.getenv("DATABRICKS_API_BASE") + ) + + # set API KEY + api_key = ( + api_key + or litellm.api_key # for databricks we check in get_llm_provider and pass in the api key from there + or litellm.databricks_key + or get_secret("DATABRICKS_API_KEY") + ) + + headers = headers or litellm.headers + + ## COMPLETION CALL + try: + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider="databricks", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + if optional_params.get("stream", False): + ## LOGGING + logging.post_call( + input=messages, + api_key=api_key, + original_response=response, + additional_args={"headers": headers}, + ) + + return response + + +def _complete_datarobot(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + stream = ctx.stream + timeout = ctx.timeout + + return base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + + +def _complete_openrouter(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("OPENROUTER_API_BASE") + or "https://openrouter.ai/api/v1" + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.openrouter_key + or get_secret_str("OPENROUTER_API_KEY") + or get_secret_str("OR_API_KEY") + ) + + openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai" + openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM" + + openrouter_headers = { + "HTTP-Referer": openrouter_site_url, + "X-Title": openrouter_app_name, + } + + _headers = headers or litellm.headers + if _headers: + openrouter_headers.update(_headers) + + headers = openrouter_headers + + ## Load Config + config = litellm.OpenrouterConfig.get_config() + for k, v in config.items(): + if k == "extra_body": + # we use openai 'extra_body' to pass openrouter specific params - transforms, route, models + if "extra_body" in optional_params: + optional_params[k].update(v) + else: + optional_params[k] = v + elif k not in optional_params: + optional_params[k] = v + + ## COMPLETION CALL + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="openrouter", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + ## LOGGING + logging.post_call( + input=messages, api_key=openai.api_key, original_response=response + ) + + return response + + +def _complete_vercel_ai_gateway( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") + or "https://ai-gateway.vercel.sh/v1" + ) + + api_key = api_key or litellm.api_key or get_secret("VERCEL_AI_GATEWAY_API_KEY") + + vercel_site_url = get_secret("VERCEL_SITE_URL") or "https://litellm.ai" + vercel_app_name = get_secret("VERCEL_APP_NAME") or "liteLLM" + + vercel_headers = { + "http-referer": vercel_site_url, + "x-title": vercel_app_name, + } + + _headers = headers or litellm.headers + if _headers: + vercel_headers.update(_headers) + + headers = vercel_headers + + ## Load Config + config = litellm.VercelAIGatewayConfig.get_config() + for k, v in config.items(): + if k == "extra_body": + # we use openai 'extra_body' to pass vercel specific params - providerOptions + if "extra_body" in optional_params: + optional_params[k].update(v) + else: + optional_params[k] = v + elif k not in optional_params: + optional_params[k] = v + + ## COMPLETION CALL + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="vercel_ai_gateway", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + ## LOGGING + logging.post_call( + input=messages, api_key=openai.api_key, original_response=response + ) + + return response + + +def _complete_vertex_ai_beta( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + vertex_ai_project = ( + optional_params.pop("vertex_project", None) + or optional_params.pop("vertex_ai_project", None) + or litellm.vertex_project + or get_secret("VERTEXAI_PROJECT") + ) + vertex_ai_location = ( + optional_params.pop("vertex_location", None) + or optional_params.pop("vertex_ai_location", None) + or litellm.vertex_location + or get_secret("VERTEXAI_LOCATION") + ) + vertex_credentials = ( + optional_params.pop("vertex_credentials", None) + or optional_params.pop("vertex_ai_credentials", None) + or get_secret("VERTEXAI_CREDENTIALS") + ) + + gemini_api_key = ( + api_key + or get_api_key_from_env() + or get_secret("PALM_API_KEY") # older palm api key should also work + or litellm.api_key + ) + + api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE") + new_params = safe_deep_copy(optional_params or {}) + return vertex_chat_completion.completion( # type: ignore + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + gemini_api_key=gemini_api_key, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, + custom_llm_provider=custom_llm_provider, # type: ignore + client=client, + api_base=api_base, + extra_headers=headers, + ) + + +def _complete_vertex_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + vertex_ai_project = ( + optional_params.pop("vertex_project", None) + or optional_params.pop("vertex_ai_project", None) + or litellm.vertex_project + or get_secret("VERTEXAI_PROJECT") + ) + vertex_ai_location = ( + optional_params.pop("vertex_location", None) + or optional_params.pop("vertex_ai_location", None) + or litellm.vertex_location + or get_secret("VERTEXAI_LOCATION") + ) + vertex_credentials = ( + optional_params.pop("vertex_credentials", None) + or optional_params.pop("vertex_ai_credentials", None) + or get_secret("VERTEXAI_CREDENTIALS") + ) + + api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE") + + new_params = safe_deep_copy(optional_params or {}) + model_route = get_vertex_ai_model_route(model=model, litellm_params=litellm_params) + + if model_route == VertexAIModelRoute.PARTNER_MODELS: + model_response = vertex_partner_models_chat_completion.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + api_base=api_base, + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + logging_obj=logging, + acompletion=acompletion, + headers=headers, + custom_prompt_dict=custom_prompt_dict, + timeout=timeout, + client=client, + ) + elif model_route == VertexAIModelRoute.GEMINI: + model_response = vertex_chat_completion.completion( # type: ignore + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + gemini_api_key=None, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, + custom_llm_provider=custom_llm_provider, # type: ignore + client=client, + api_base=api_base, + extra_headers=headers, + ) + elif model_route == VertexAIModelRoute.GEMMA: + # Vertex Gemma Models with custom prediction endpoint + model_response = vertex_gemma_chat_completion.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + api_base=api_base, + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + logging_obj=logging, + acompletion=acompletion, + headers=headers, + custom_prompt_dict=custom_prompt_dict, + timeout=timeout, + client=client, + ) + elif model_route == VertexAIModelRoute.MODEL_GARDEN: + # Vertex Model Garden - OpenAI compatible models + model_response = vertex_model_garden_chat_completion.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + api_base=api_base, + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + logging_obj=logging, + acompletion=acompletion, + headers=headers, + custom_prompt_dict=custom_prompt_dict, + timeout=timeout, + client=client, + ) + elif model_route == VertexAIModelRoute.AGENT_ENGINE: + # Vertex AI Agent Engine (Reasoning Engines) + from litellm.llms.vertex_ai.agent_engine.transformation import ( + VertexAgentEngineConfig, + ) + + vertex_agent_engine_config = VertexAgentEngineConfig() + + # Update litellm_params with vertex credentials + litellm_params["vertex_project"] = vertex_ai_project + litellm_params["vertex_location"] = vertex_ai_location + litellm_params["vertex_credentials"] = vertex_credentials + + model_response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + model_response=model_response, + optional_params=new_params, + litellm_params=litellm_params, # type: ignore + encoding=_get_encoding(), + api_key=None, + api_base=api_base, + logging_obj=logging, + acompletion=acompletion, + timeout=timeout, + client=client, + custom_llm_provider="vertex_ai", + provider_config=vertex_agent_engine_config, + headers=headers or {}, + ) + else: # VertexAIModelRoute.NON_GEMINI + model_response = vertex_ai_non_gemini.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=new_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + vertex_location=vertex_ai_location, + vertex_project=vertex_ai_project, + vertex_credentials=vertex_credentials, + logging_obj=logging, + acompletion=acompletion, + ) + + if ( + "stream" in optional_params + and optional_params["stream"] is True + and acompletion is False + ): + return CustomStreamWrapper( + model_response, + model, + custom_llm_provider="vertex_ai", + logging_obj=logging, + ) + return model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_predibase(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + tenant_id = ( + optional_params.pop("tenant_id", None) + or optional_params.pop("predibase_tenant_id", None) + or litellm.predibase_tenant_id + or get_secret("PREDIBASE_TENANT_ID") + ) + + if tenant_id is None: + raise ValueError( + "Missing Predibase Tenant ID - Required for making the request. Set dynamically (e.g. `completion(..tenant_id=)`) or in env - `PREDIBASE_TENANT_ID`." + ) + + api_base = ( + api_base + or optional_params.pop("api_base", None) + or optional_params.pop("base_url", None) + or litellm.api_base + or get_secret("PREDIBASE_API_BASE") + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.predibase_key + or get_secret("PREDIBASE_API_KEY") + ) + + _model_response = predibase_chat_completions.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + acompletion=acompletion, + api_base=api_base, + custom_prompt_dict=custom_prompt_dict, + api_key=api_key, + tenant_id=tenant_id, + timeout=timeout, + ) + + if ( + "stream" in optional_params + and optional_params["stream"] is True + and acompletion is False + ): + return _model_response + return _model_response + + +def _complete_text_completion_codestral( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + api_base + or optional_params.pop("api_base", None) + or optional_params.pop("base_url", None) + or litellm.api_base + or "https://codestral.mistral.ai/v1/fim/completions" + ) + + api_key = api_key or litellm.api_key or get_secret("CODESTRAL_API_KEY") + + text_completion_model_response = litellm.TextCompletionResponse(stream=stream) + + _model_response = codestral_text_completions.completion( # type: ignore + model=model, + messages=messages, + model_response=text_completion_model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + acompletion=acompletion, + api_base=api_base, + custom_prompt_dict=custom_prompt_dict, + api_key=api_key, + timeout=timeout, + ) + + if ( + "stream" in optional_params + and optional_params["stream"] is True + and acompletion is False + ): + return _model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + return _model_response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_text_completion_inception( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + text_completion = ctx.text_completion + timeout = ctx.timeout + + passed_api_base = ( + api_base + or optional_params.pop("api_base", None) + or optional_params.pop("base_url", None) + ) + api_base = ( + passed_api_base + or get_secret_str("INCEPTION_API_BASE") + or "https://api.inceptionlabs.ai/v1" + ) + # FIM is served at `/v1/fim/completions`; the OpenAI client appends + # `/completions`, so point it at the `/v1/fim` base. + api_base = api_base.rstrip("/") + if not api_base.endswith("/fim"): + api_base += "/fim" + + # Don't forward the server-managed Inception key to a caller-supplied + # api_base; only resolve it for the default/server base, or when the + # caller passes their own key. + if passed_api_base is None or api_key: + api_key = ( + api_key or litellm.inception_key or get_secret_str("INCEPTION_API_KEY") + ) + + _response = openai_text_completions.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, # type: ignore[arg-type] + custom_llm_provider="text-completion-inception", + api_base=api_base, + acompletion=acompletion, + client=client, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + ) + + if ( + optional_params.get("stream", False) is False + and acompletion is False + and text_completion is False + ): + _response = ( + litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( + response_object=_response, model_response_object=model_response + ) + ) + + if optional_params.get("stream", False) or acompletion is True: + logging.post_call( + input=messages, + api_key=api_key, + original_response=_response, + additional_args={"headers": headers}, + ) + return _response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_sagemaker_chat( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_sagemaker(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + custom_prompt_dict = ctx.custom_prompt_dict + hf_model_name = ctx.hf_model_name + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + return sagemaker_llm.completion( + model=model, + messages=messages, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + custom_prompt_dict=custom_prompt_dict, + hf_model_name=hf_model_name, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + acompletion=acompletion, + ) + + +def _complete_bedrock(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + provider_config = ctx.provider_config + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + + if "aws_bedrock_client" in optional_params: + verbose_logger.warning( + "'aws_bedrock_client' is a deprecated param. Please move to another auth method - https://docs.litellm.ai/docs/providers/bedrock#boto3---authentication." + ) + # Extract credentials for legacy boto3 client and pass thru to httpx + aws_bedrock_client = optional_params.pop("aws_bedrock_client") + creds = aws_bedrock_client._get_credentials().get_frozen_credentials() + + if creds.access_key: + optional_params["aws_access_key_id"] = creds.access_key + if creds.secret_key: + optional_params["aws_secret_access_key"] = creds.secret_key + if creds.token: + optional_params["aws_session_token"] = creds.token + if ( + "aws_region_name" not in optional_params + or optional_params["aws_region_name"] is None + ): + optional_params["aws_region_name"] = aws_bedrock_client.meta.region_name + + bedrock_route = BedrockModelInfo.get_bedrock_route(model) + if bedrock_route == "claude_platform": + provider_config = ProviderConfigManager.get_provider_chat_config( + model=model, + provider=LlmProviders.BEDROCK, + ) + model = BedrockModelInfo.get_claude_platform_model(model) + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="bedrock", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + provider_config=provider_config, + ) + elif bedrock_route == "converse": + model = model.replace("converse/", "") + response = bedrock_converse_chat_completion.completion( + model=model, + messages=messages, + custom_prompt_dict=custom_prompt_dict, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, # type: ignore + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + extra_headers=headers, # Use merged headers instead of original extra_headers + timeout=timeout, + acompletion=acompletion, + client=client, + api_base=api_base, + api_key=api_key, + ) + elif bedrock_route == "converse_like": + model = model.replace("converse_like/", "") + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider="bedrock", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + else: + response = base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider="bedrock", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + + return response + + +def _complete_watsonx(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + timeout = ctx.timeout + + return watsonx_chat_completion.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + custom_prompt_dict=custom_prompt_dict, + client=client, # pass AsyncOpenAI, OpenAI client + encoding=_get_encoding(), + custom_llm_provider="watsonx", + ) + + +def _complete_watsonx_text( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or optional_params.pop("apikey", None) + or get_secret_str("WATSONX_APIKEY") + or get_secret_str("WATSONX_API_KEY") + or get_secret_str("WX_API_KEY") + ) + + api_base = ( + api_base + or optional_params.pop( + "url", + optional_params.pop("api_base", optional_params.pop("base_url", None)), + ) + or get_secret_str("WATSONX_API_BASE") + or get_secret_str("WATSONX_URL") + or get_secret_str("WX_URL") + or get_secret_str("WML_URL") + ) + + wx_credentials = optional_params.pop( + "wx_credentials", + optional_params.pop( + "watsonx_credentials", None + ), # follow {provider}_credentials, same as vertex ai + ) + + token: Optional[str] = None + if wx_credentials is not None: + api_base = wx_credentials.get("url", api_base) + api_key = wx_credentials.get("apikey", wx_credentials.get("api_key", api_key)) + token = wx_credentials.get( + "token", + wx_credentials.get( + "watsonx_token", None + ), # follow format of {provider}_token, same as azure - e.g. 'azure_ad_token=..' + ) + + if token is not None: + optional_params["token"] = token + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="watsonx_text", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_vllm(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + custom_prompt_dict = ctx.custom_prompt_dict + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + model_response = vllm_handler.completion( + model=model, + messages=messages, + custom_prompt_dict=custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + ) + + if "stream" in optional_params and optional_params["stream"] is True: ## [BETA] + # don't try to access stream object, + return CustomStreamWrapper( + model_response, + model, + custom_llm_provider="vllm", + logging_obj=logging, + ) + + ## RESPONSE OBJECT + return model_response + + +def _complete_ollama(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + litellm.api_base + or api_base + or get_secret("OLLAMA_API_BASE") + or "http://localhost:11434" + ) + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="ollama", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_ollama_chat(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = ( + litellm.api_base + or api_base + or get_secret("OLLAMA_API_BASE") + or "http://localhost:11434" + ) + + api_key = ( + api_key + or litellm.ollama_key + or os.environ.get("OLLAMA_API_KEY") + or litellm.api_key + ) + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="ollama_chat", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + client=client, + ) + + +def _complete_triton(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = litellm.api_base or api_base + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + ) + + +def _complete_cloudflare(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.cloudflare_api_key + or litellm.api_key + or get_secret("CLOUDFLARE_API_KEY") + ) + api_base = api_base or litellm.api_base or get_secret("CLOUDFLARE_API_BASE") + + custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="cloudflare", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements + ) + + +def _complete_petals(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + client = ctx.client + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + + api_base = api_base or litellm.api_base + + stream = optional_params.pop("stream", False) + model_response = petals_handler.completion( + model=model, + messages=messages, + api_base=api_base, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + client=client, + ) + if stream is True: ## [BETA] + # Fake streaming for petals + resp_string = model_response["choices"][0]["message"]["content"] + return CustomStreamWrapper( + resp_string, + model, + custom_llm_provider="petals", + logging_obj=logging, + ) + return model_response + + +def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + try: + client = ( + HTTPHandler(timeout=timeout) if stream is False else None + ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + ) + + except Exception as e: + ## LOGGING - log the original exception returned + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e + + return response + + +def _complete_gradient_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + api_base = litellm.api_base or api_base + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider="gradient_ai", + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + ) + + +def _complete_bytez(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.bytez_key + or get_secret_str("BYTEZ_API_KEY") + or litellm.api_key + ) + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=bytez_transformation, + ) + + pass + + return response + + +def _complete_lemonade(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.lemonade_key + or get_secret_str("LEMONADE_API_KEY") + or litellm.api_key + ) + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=lemonade_transformation, + ) + + pass + + return response + + +def _complete_ovhcloud(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + api_key = ( + api_key + or litellm.ovhcloud_key + or get_secret_str("OVHCLOUD_API_KEY") + or litellm.api_key + ) + + api_base = ( + api_base + or litellm.api_base + or get_secret_str("OVHCLOUD_API_BASE") + or "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" + ) + + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + timeout=timeout, # type: ignore + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=ovhcloud_transformation, + ) + + pass + + return response + + +def _complete_custom(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + api_base = ctx.api_base + headers = ctx.headers + kwargs = ctx.kwargs + max_tokens = ctx.max_tokens + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + temperature = ctx.temperature + top_p = ctx.top_p + + url = litellm.api_base or api_base or "" + if url is None or url == "": + raise ValueError( + "api_base not set. Set api_base or litellm.api_base for custom endpoints" + ) + + """ + assume input to custom LLM api bases follow this format: + resp = litellm.module_level_client.post( + api_base, + json={ + 'model': 'meta-llama/Llama-2-13b-hf', # model name + 'params': { + 'prompt': ["The capital of France is P"], + 'max_tokens': 32, + 'temperature': 0.7, + 'top_p': 1.0, + 'top_k': 40, + } + } + ) + + """ + prompt = " ".join([message["content"] for message in messages]) # type: ignore + resp = litellm.module_level_client.post( + url, + headers=headers, + json={ + "model": model, + "params": { + "prompt": [prompt], + "max_tokens": max_tokens, + "temperature": temperature, + "top_p": top_p, + "top_k": kwargs.get("top_k"), + }, + **kwargs.get("extra_body", {}), + }, + ) + response_json = resp.json() + """ + assume all responses from custom api_bases of this format: + { + 'data': [ + { + 'prompt': 'The capital of France is P', + 'output': ['The capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France'], + 'params': {'temperature': 0.7, 'top_k': 40, 'top_p': 1}}], + 'message': 'ok' + } + ] + } + """ + string_response = response_json["data"][0]["output"][0] + ## RESPONSE OBJECT + model_response.choices[0].message.content = string_response # type: ignore + model_response.created = int(time.time()) + model_response.model = model + return model_response + + +def _complete_custom_providers( + ctx: _CompletionDispatchContext, +) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + custom_prompt_dict = ctx.custom_prompt_dict + headers = ctx.headers + litellm_params = ctx.litellm_params + logger_fn = ctx.logger_fn + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + stream = ctx.stream + timeout = ctx.timeout + + custom_handler: Optional[CustomLLM] = None + for item in litellm.custom_provider_map: + if item["provider"] == custom_llm_provider: + custom_handler = item["custom_handler"] + + if custom_handler is None: + raise LiteLLMUnknownProvider( + model=model, custom_llm_provider=custom_llm_provider + ) + + ## ROUTE LLM CALL ## + handler_fn = custom_chat_llm_router( + async_fn=acompletion, stream=stream, custom_llm=custom_handler + ) + + headers = headers or litellm.headers or {} + + ## CALL FUNCTION + response = handler_fn( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + print_verbose=print_verbose, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + timeout=timeout, # type: ignore + custom_prompt_dict=custom_prompt_dict, + client=client, # pass AsyncOpenAI, OpenAI client + encoding=_get_encoding(), + ) + if stream is True: + return CustomStreamWrapper( + completion_stream=response, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging, + ) + + return response # pyright: ignore[reportReturnType] # provider SDK return type is broader than the dispatch contract + + +def _complete_langgraph(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + from litellm.llms.langgraph.chat.transformation import LangGraphConfig + + ( + api_base, + api_key, + ) = LangGraphConfig()._get_openai_compatible_provider_info( + api_base=api_base or litellm.api_base, + api_key=api_key or litellm.api_key, + ) + + headers = headers or litellm.headers + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + + +def _complete_langflow(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion = ctx.acompletion + api_base = ctx.api_base + api_key = ctx.api_key + client = ctx.client + custom_llm_provider = ctx.custom_llm_provider + headers = ctx.headers + litellm_params = ctx.litellm_params + logging = ctx.logging + messages = ctx.messages + model = ctx.model + model_response = ctx.model_response + optional_params = ctx.optional_params + shared_session = ctx.shared_session + stream = ctx.stream + timeout = ctx.timeout + + from litellm.llms.langflow.chat.transformation import LangFlowConfig + + ( + api_base, + api_key, + ) = LangFlowConfig()._get_openai_compatible_provider_info( + api_base=api_base or litellm.api_base, + api_key=api_key or litellm.api_key, + ) + + headers = headers or litellm.headers + + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + custom_llm_provider=custom_llm_provider, + timeout=timeout, + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + + @tracer.wrap() @client def completion( # type: ignore @@ -1215,9 +5077,7 @@ def completion( # type: ignore if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway( tools=tools_for_mcp ): - # Return coroutine - acompletion will await it - # completion() can return a coroutine when MCP tools are present, which acompletion() awaits - return acompletion_with_mcp( # type: ignore[return-value] + return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it model=model, messages=messages, functions=functions, @@ -1389,12 +5249,16 @@ def completion( # type: ignore logging: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, litellm_logging_obj) fallbacks = fallbacks or litellm.model_fallbacks if fallbacks is not None: - return completion_with_fallbacks(**args) + return completion_with_fallbacks( # pyright: ignore[reportReturnType] # fallback runner is untyped; resolves to ModelResponse|CustomStreamWrapper at runtime + **args + ) if model_list is not None: deployments = [ m["litellm_params"] for m in model_list if m["model_name"] == model ] - return litellm.batch_completion_models(deployments=deployments, **args) + return litellm.batch_completion_models( # pyright: ignore[reportReturnType] # batch path returns a list of responses, outside completion()'s single-response return type + deployments=deployments, **args + ) if litellm.model_alias_map and model in litellm.model_alias_map: model = litellm.model_alias_map[ model @@ -1454,7 +5318,7 @@ def completion( # type: ignore timeout, kwargs, custom_llm_provider, - global_timeout=getattr(litellm, "request_timeout", None), + global_timeout=get_configured_request_timeout(), supports_httpx_timeout=supports_httpx_timeout, ) @@ -1716,7 +5580,7 @@ def completion( # type: ignore else: optional_params["reasoning_effort"] = {"summary": rs_val} - return responses_api_bridge.completion( + return responses_api_bridge.completion( # pyright: ignore[reportReturnType] # bridge returns a coroutine on the acompletion path; awaited by the async caller model=model, messages=messages, headers=headers, @@ -1746,375 +5610,52 @@ def completion( # type: ignore optional_params ) + _dispatch_ctx = _CompletionDispatchContext( + _azure_detection_model=_azure_detection_model, + acompletion=acompletion, + api_base=api_base, + api_key=api_key, + api_version=api_version, + client=client, + custom_llm_provider=custom_llm_provider, + custom_prompt_dict=custom_prompt_dict, + extra_headers=extra_headers, + headers=headers, + hf_model_name=hf_model_name, + kwargs=kwargs, + litellm_params=litellm_params, + logger_fn=logger_fn, + logging=logging, + max_retries=max_retries, + max_tokens=max_tokens, + messages=messages, + metadata=metadata, + model=model, + model_response=model_response, + optional_params=optional_params, + organization=organization, + provider_config=provider_config, + shared_session=shared_session, + stream=stream, + temperature=temperature, + text_completion=text_completion, + timeout=timeout, + top_p=top_p, + ) if custom_llm_provider == "azure": # azure configs ## check dynamic params ## - dynamic_params = False - if client is not None and ( - isinstance(client, openai.AzureOpenAI) - or isinstance(client, openai.AsyncAzureOpenAI) - ): - dynamic_params = _check_dynamic_azure_params( - azure_client_params={"api_version": api_version}, - azure_client=client, - ) - - api_type = get_secret("AZURE_API_TYPE") or "azure" - - api_base = api_base or litellm.api_base or get_secret("AZURE_API_BASE") - - api_version = ( - api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - or litellm.AZURE_DEFAULT_API_VERSION - ) - - api_key = ( - api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) - - azure_ad_token = optional_params.get("extra_body", {}).pop( - "azure_ad_token", None - ) or get_secret_str("AZURE_AD_TOKEN") - - azure_ad_token_provider = litellm_params.get( - "azure_ad_token_provider", None - ) - - headers = headers or litellm.headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - if max_retries is not None: - optional_params["max_retries"] = max_retries - - if litellm.AzureOpenAIO1Config().is_o_series_model( - model=_azure_detection_model - ): - ## LOAD CONFIG - if set - config = litellm.AzureOpenAIO1Config.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - response = azure_o1_chat_completions.completion( - model=model, - messages=messages, - headers=headers, - api_key=api_key, - api_base=api_base, - api_version=api_version, - dynamic_params=dynamic_params, - azure_ad_token=azure_ad_token, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, # type: ignore - client=client, # pass AsyncAzureOpenAI, AzureOpenAI client - custom_llm_provider=custom_llm_provider, - ) - else: - ## LOAD CONFIG - if set - config = litellm.AzureOpenAIConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - ## COMPLETION CALL - response = azure_chat_completions.completion( - model=model, - messages=messages, - headers=headers, - api_key=api_key, - api_base=api_base, - api_version=api_version, - api_type=api_type, - dynamic_params=dynamic_params, - azure_ad_token=azure_ad_token, - azure_ad_token_provider=azure_ad_token_provider, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, # type: ignore - client=client, # pass AsyncAzureOpenAI, AzureOpenAI client - ) - - if optional_params.get("stream", False): - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={ - "headers": headers, - "api_version": api_version, - "api_base": api_base, - }, - ) + response = _complete_azure(_dispatch_ctx) elif custom_llm_provider == "azure_text": # azure configs - api_type = get_secret_str("AZURE_API_TYPE") or "azure" - - api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") - - if api_base is None: - raise ValueError( - "api_base is required for Azure OpenAI LLM provider. Either set it dynamically or set the AZURE_API_BASE environment variable." - ) - - api_version = ( - api_version - or litellm.api_version - or get_secret_str("AZURE_API_VERSION") - ) - - api_key = ( - api_key - or litellm.api_key - or litellm.azure_key - or get_secret_str("AZURE_OPENAI_API_KEY") - or get_secret_str("AZURE_API_KEY") - ) - - azure_ad_token = optional_params.get("extra_body", {}).pop( - "azure_ad_token", None - ) or get_secret_str("AZURE_AD_TOKEN") - - azure_ad_token_provider = litellm_params.get( - "azure_ad_token_provider", None - ) - - headers = headers or litellm.headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - - ## LOAD CONFIG - if set - config = litellm.AzureOpenAIConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - ## COMPLETION CALL - response = azure_text_completions.completion( - model=model, - messages=messages, - headers=headers, - api_key=api_key, - api_base=api_base, - api_version=cast(str, api_version), - api_type=api_type, - azure_ad_token=azure_ad_token, - azure_ad_token_provider=azure_ad_token_provider, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, - client=client, # pass AsyncAzureOpenAI, AzureOpenAI client - ) - - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={ - "headers": headers, - "api_version": api_version, - "api_base": api_base, - }, - ) + response = _complete_azure_text(_dispatch_ctx) elif custom_llm_provider == "deepseek": ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_deepseek(_dispatch_ctx) elif custom_llm_provider == "azure_ai": - from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo - - azure_ai_route = AzureFoundryModelInfo.get_azure_ai_route(model) - - # Check if this is an agents route - model format: azure_ai/agents/ - if azure_ai_route == "agents": - from litellm.llms.azure_ai.agents import AzureAIAgentsConfig - - api_base = AzureFoundryModelInfo.get_api_base(api_base) - if api_base is None: - raise ValueError( - "Azure AI Agents requests require an api_base. " - "Set `api_base` or the AZURE_AI_API_BASE env var." - ) - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - response = AzureAIAgentsConfig.completion( - model=model, - messages=messages, - api_base=api_base, - api_key=api_key, - model_response=model_response, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, - acompletion=acompletion, - stream=stream, - headers=headers or litellm.headers, - ) - - # Check if this is a Claude model - route to Azure Anthropic handler - elif "claude" in model.lower(): - # Use Azure Anthropic handler for Claude models - api_base = AzureFoundryModelInfo.get_api_base(api_base) - if api_base is None: - raise ValueError( - "Azure Anthropic requests require an api_base. " - "Set `api_base` or the AZURE_AI_API_BASE env var." - ) - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - # Ensure the URL ends with /v1/messages for Anthropic - if api_base: - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/messages"): - if "/anthropic" in api_base: - parts = api_base.split("/anthropic", 1) - api_base = parts[0] + "/anthropic" - else: - api_base = api_base + "/anthropic" - api_base = api_base + "/v1/messages" - - response = azure_anthropic_chat_completions.completion( - model=model, - messages=messages, - api_base=api_base, - acompletion=acompletion, - custom_prompt_dict=litellm.custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - headers=headers, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - ) - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - ) - response = response - else: - # Non-Claude models use standard Azure AI flow - api_base = AzureFoundryModelInfo.get_api_base(api_base) - # set API KEY - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - headers = headers or litellm.headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - - ## FOR COHERE - if "command-r" in model: # make sure tool call in messages are str - messages = stringify_json_tool_call_content(messages=messages) - - ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, # pass AsyncOpenAI, OpenAI client - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e - - if optional_params.get("stream", False): - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={"headers": headers}, - ) + response = _complete_azure_ai(_dispatch_ctx) elif ( custom_llm_provider == "text-completion-openai" or "ft:babbage-002" in model @@ -2123,535 +5664,42 @@ def completion( # type: ignore in litellm.openai_text_completion_compatible_providers and kwargs.get("text_completion") is True ): - openai.api_type = "openai" - - api_base = ( - api_base - or litellm.api_base - or get_secret("OPENAI_BASE_URL") - or get_secret("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - - openai.api_version = None - # set API KEY - - api_key = ( - api_key - or litellm.api_key - or litellm.openai_key - or get_secret("OPENAI_API_KEY") - ) - - headers = headers or litellm.headers - - ## LOAD CONFIG - if set - config = litellm.OpenAITextCompletionConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > openai_text_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - if litellm.organization: - openai.organization = litellm.organization - - if ( - len(messages) > 0 - and "content" in messages[0] - and isinstance(messages[0]["content"], list) - ): - # text-davinci-003 can accept a string or array, if it's an array, assume the array is set in messages[0]['content'] - # https://platform.openai.com/docs/api-reference/completions/create - prompt = messages[0]["content"] - else: - prompt = " ".join([message["content"] for message in messages]) # type: ignore - - ## COMPLETION CALL - _response = openai_text_completions.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - acompletion=acompletion, - client=client, # pass AsyncOpenAI, OpenAI client - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - ) - - if ( - optional_params.get("stream", False) is False - and acompletion is False - and text_completion is False - ): - # convert to chat completion response - _response = litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( - response_object=_response, model_response_object=model_response - ) - - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=_response, - additional_args={"headers": headers}, - ) - response = _response + response = _complete_text_completion_openai(_dispatch_ctx) elif custom_llm_provider == "fireworks_ai": ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_fireworks_ai(_dispatch_ctx) elif custom_llm_provider == "heroku": - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_heroku(_dispatch_ctx) elif custom_llm_provider == "ragflow": ## COMPLETION CALL - RAGFlow uses HTTP handler to support custom URL paths - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_ragflow(_dispatch_ctx) elif custom_llm_provider == "xai": ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_xai(_dispatch_ctx) elif custom_llm_provider == "groq": - api_base = ( - api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there - or litellm.api_base - or get_secret("GROQ_API_BASE") - or "https://api.groq.com/openai/v1" - ) - - # set API KEY - api_key = ( - api_key - or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there - or litellm.groq_key - or get_secret("GROQ_API_KEY") - ) - - headers = headers or litellm.headers - - ## LOAD CONFIG - if set - config = litellm.GroqChatConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) + response = _complete_groq(_dispatch_ctx) elif custom_llm_provider == "bedrock_mantle": - api_base = ( - api_base or litellm.api_base or get_secret("BEDROCK_MANTLE_API_BASE") - ) - api_key = api_key or litellm.api_key or get_secret("BEDROCK_MANTLE_API_KEY") - headers = headers or litellm.headers - config = litellm.BedrockMantleChatConfig.get_config() - for k, v in config.items(): - if k not in optional_params: - optional_params[k] = v - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - ) + response = _complete_bedrock_mantle(_dispatch_ctx) elif custom_llm_provider == "a2a": # A2A (Agent-to-Agent) Protocol # Resolve agent configuration from registry if model format is "a2a/" - ( - api_base, - api_key, - headers, - ) = litellm.A2AConfig.resolve_agent_config_from_registry( - model=model, - api_base=api_base, - api_key=api_key, - headers=headers, - optional_params=optional_params, - ) - - # Fall back to environment variables and defaults - api_base = api_base or litellm.api_base or get_secret_str("A2A_API_BASE") - - if api_base is None: - raise Exception( - "api_base is required for A2A provider. " - "Either provide api_base parameter, set A2A_API_BASE environment variable, " - "or register the agent in the proxy with model='a2a/'." - ) - - headers = headers or litellm.headers - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - provider_config=provider_config, - ) + response = _complete_a2a(_dispatch_ctx) elif custom_llm_provider == "gigachat": # GigaChat - Sber AI's LLM (Russia) - api_key = ( - api_key - or litellm.api_key - or litellm.gigachat_key - or get_secret("GIGACHAT_API_KEY") - or get_secret("GIGACHAT_CREDENTIALS") - ) - - headers = headers or litellm.headers or {} - - ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_gigachat(_dispatch_ctx) elif custom_llm_provider == "sap": - headers = headers or litellm.headers - ## LOAD CONFIG - if set - config = litellm.GenAIHubOrchestrationConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - response = sap_gen_ai_hub_chat_completions.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - shared_session=shared_session, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - api_key=api_key, - api_base=api_base, - stream=stream, - ) + response = _complete_sap(_dispatch_ctx) elif custom_llm_provider == "aiohttp_openai": # NEW aiohttp provider for 10-100x higher RPS - api_base = ( - api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there - or litellm.api_base - or get_secret("OPENAI_BASE_URL") - or get_secret("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - # set API KEY - api_key = ( - api_key - or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there - or litellm.openai_key - or get_secret("OPENAI_API_KEY") - ) - - headers = headers or litellm.headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - response = base_llm_aiohttp_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) + response = _complete_aiohttp_openai(_dispatch_ctx) elif custom_llm_provider == "cometapi": - api_key = ( - api_key - or litellm.cometapi_key - or get_secret_str("COMETAPI_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret_str("COMETAPI_API_BASE") - or "https://api.cometapi.com/v1" - ) - - ## COMPLETION CALL - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) - - ## LOGGING - logging.post_call( - input=messages, api_key=api_key, original_response=response - ) + response = _complete_cometapi(_dispatch_ctx) elif custom_llm_provider == "minimax": - api_key = api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key - - api_base = ( - api_base - or litellm.api_base - or get_secret_str("MINIMAX_API_BASE") - or "https://api.minimax.io/v1" - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - model_response=model_response, - encoding=_get_encoding(), - logging_obj=logging, - optional_params=optional_params, - timeout=timeout, - litellm_params=litellm_params, - shared_session=shared_session, - acompletion=acompletion, - stream=stream, - api_key=api_key, - headers=headers, - client=client, - provider_config=provider_config, - ) - logging.post_call( - input=messages, api_key=api_key, original_response=response - ) + response = _complete_minimax(_dispatch_ctx) elif custom_llm_provider == "hosted_vllm": - api_base = ( - api_base or litellm.api_base or get_secret_str("HOSTED_VLLM_API_BASE") - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - model_response=model_response, - encoding=_get_encoding(), - logging_obj=logging, - optional_params=optional_params, - timeout=timeout, - litellm_params=litellm_params, - shared_session=shared_session, - acompletion=acompletion, - stream=stream, - api_key=api_key, - headers=headers, - client=client, - provider_config=provider_config, - ) - logging.post_call( - input=messages, api_key=api_key, original_response=response - ) + response = _complete_hosted_vllm(_dispatch_ctx) elif ( model in litellm.open_ai_chat_completion_models or custom_llm_provider == "custom_openai" @@ -2676,205 +5724,17 @@ def completion( # type: ignore ): # allow user to make an openai call with a custom base # note: if a user sets a custom base - we should ensure this works # allow for the setting of dynamic and stateful api-bases - api_base = ( - api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there - or litellm.api_base - or get_secret("OPENAI_BASE_URL") - or get_secret("OPENAI_API_BASE") - or "https://api.openai.com/v1" - ) - organization = ( - organization - or litellm.organization - or get_secret("OPENAI_ORGANIZATION") - or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 - ) - openai.organization = organization - # set API KEY - api_key = ( - api_key - or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there - or litellm.openai_key - or get_secret("OPENAI_API_KEY") - ) - - headers = headers or litellm.headers - - # Add GitHub Copilot headers (same as /responses endpoint does) - if custom_llm_provider == "github_copilot": - from litellm.llms.github_copilot.authenticator import Authenticator - from litellm.llms.github_copilot.common_utils import ( - get_copilot_default_headers, - ) - - copilot_auth = Authenticator() - copilot_api_key = copilot_auth.get_api_key() - copilot_headers = get_copilot_default_headers(copilot_api_key) - if extra_headers: - copilot_headers.update(extra_headers) - extra_headers = copilot_headers - - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - - if ( - litellm.enable_preview_features and metadata is not None - ): # [PREVIEW] allow metadata to be passed to OPENAI - openai_metadata = get_requester_metadata(metadata) - if openai_metadata is not None: - optional_params["metadata"] = openai_metadata - - ## LOAD CONFIG - if set - config = litellm.OpenAIConfig.get_config() - for k, v in config.items(): - if ( - k not in optional_params - ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in - optional_params[k] = v - - ## COMPLETION CALL - use_base_llm_http_handler = get_secret_bool( - "EXPERIMENTAL_OPENAI_BASE_LLM_HTTP_HANDLER" - ) - - try: - if use_base_llm_http_handler: - response = base_llm_http_handler.completion( - model=model, - messages=messages, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - model_response=model_response, - encoding=_get_encoding(), - logging_obj=logging, - optional_params=optional_params, - timeout=timeout, - litellm_params=litellm_params, - shared_session=shared_session, - acompletion=acompletion, - stream=stream, - api_key=api_key, - headers=headers, - client=client, - provider_config=provider_config, - ) - else: - response = openai_chat_completions.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - custom_prompt_dict=custom_prompt_dict, - client=client, # pass AsyncOpenAI, OpenAI client - organization=organization, - custom_llm_provider=custom_llm_provider, - shared_session=shared_session, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e - - if optional_params.get("stream", False): - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={"headers": headers}, - ) + response = _complete_custom_openai(_dispatch_ctx) elif custom_llm_provider == "mistral": - api_key = api_key or litellm.api_key or get_secret("MISTRAL_API_KEY") - api_base = ( - api_base - or litellm.api_base - or get_secret("MISTRAL_API_BASE") - or "https://api.mistral.ai/v1" - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - model_response=model_response, - encoding=_get_encoding(), - logging_obj=logging, - optional_params=optional_params, - timeout=timeout, - litellm_params=litellm_params, - shared_session=shared_session, - acompletion=acompletion, - stream=stream, - api_key=api_key, - headers=headers, - client=client, - provider_config=provider_config, - ) + response = _complete_mistral(_dispatch_ctx) elif ( "replicate" in model or custom_llm_provider == "replicate" or model in litellm.replicate_models ): # Setting the relevant API KEY for replicate, replicate defaults to using os.environ.get("REPLICATE_API_TOKEN") - replicate_key = ( - api_key - or litellm.replicate_key - or litellm.api_key - or get_secret("REPLICATE_API_KEY") - or get_secret("REPLICATE_API_TOKEN") - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret("REPLICATE_API_BASE") - or "https://api.replicate.com/v1" - ) - - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - - model_response = replicate_chat_completion( # type: ignore - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), # for calculating input/output tokens - api_key=replicate_key, - logging_obj=logging, - custom_prompt_dict=custom_prompt_dict, - acompletion=acompletion, - headers=headers, - ) - - if optional_params.get("stream", False) is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=replicate_key, - original_response=model_response, - ) - - response = model_response + response = _complete_replicate(_dispatch_ctx) elif ( "clarifai" in model or custom_llm_provider == "clarifai" @@ -2882,614 +5742,36 @@ def completion( # type: ignore ): pass # Deprecated - handled in the openai compatible provider section above elif custom_llm_provider == "anthropic_text": - api_key = ( - api_key - or litellm.anthropic_key - or litellm.api_key - or os.environ.get("ANTHROPIC_API_KEY") - ) - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - api_base = ( - api_base - or litellm.api_base - or get_secret("ANTHROPIC_API_BASE") - or get_secret("ANTHROPIC_BASE_URL") - or "https://api.anthropic.com/v1/complete" - ) - - # Check if we should disable automatic URL suffix appending - disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if ( - api_base is not None - and not disable_url_suffix - and not api_base.endswith("/v1/complete") - ): - api_base += "/v1/complete" - elif disable_url_suffix: - verbose_logger.debug( - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/complete suffix" - ) - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="anthropic_text", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - ) + response = _complete_anthropic_text(_dispatch_ctx) elif custom_llm_provider == "anthropic": - api_key = ( - api_key - or litellm.anthropic_key - or litellm.api_key - or os.environ.get("ANTHROPIC_API_KEY") - ) - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - # call /messages - # default route for all anthropic models - api_base = ( - api_base - or litellm.api_base - or get_secret("ANTHROPIC_API_BASE") - or get_secret("ANTHROPIC_BASE_URL") - or "https://api.anthropic.com/v1/messages" - ) - - # Check if we should disable automatic URL suffix appending - disable_url_suffix = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if ( - api_base is not None - and not disable_url_suffix - and not api_base.endswith("/v1/messages") - ): - api_base += "/v1/messages" - elif disable_url_suffix: - verbose_logger.debug( - "LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/messages suffix" - ) - - response = anthropic_chat_completions.completion( - model=model, - messages=messages, - api_base=api_base, - acompletion=acompletion, - custom_prompt_dict=litellm.custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), # for calculating input/output tokens - api_key=api_key, - logging_obj=logging, - headers=headers, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - ) - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - ) - response = response + response = _complete_anthropic(_dispatch_ctx) elif custom_llm_provider == "nlp_cloud": - nlp_cloud_key = ( - api_key - or litellm.nlp_cloud_key - or get_secret("NLP_CLOUD_API_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret("NLP_CLOUD_API_BASE") - or "https://api.nlpcloud.io/v1/gpu/" - ) - - response = nlp_cloud_chat_completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - api_key=nlp_cloud_key, - logging_obj=logging, - ) - - if "stream" in optional_params and optional_params["stream"] is True: - # don't try to access stream object, - response = CustomStreamWrapper( - response, - model, - custom_llm_provider="nlp_cloud", - logging_obj=logging, - ) - - if optional_params.get("stream", False) or acompletion is True: - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - ) - - response = response + response = _complete_nlp_cloud(_dispatch_ctx) elif custom_llm_provider == "aleph_alpha": - aleph_alpha_key = ( - api_key - or litellm.aleph_alpha_key - or get_secret("ALEPH_ALPHA_API_KEY") - or get_secret("ALEPHALPHA_API_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret("ALEPH_ALPHA_API_BASE") - or "https://api.aleph-alpha.com/complete" - ) - - model_response = aleph_alpha.completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - default_max_tokens_to_sample=litellm.max_tokens, - api_key=aleph_alpha_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - ) - - if "stream" in optional_params and optional_params["stream"] is True: - # don't try to access stream object, - response = CustomStreamWrapper( - model_response, - model, - custom_llm_provider="aleph_alpha", - logging_obj=logging, - ) - return response - response = model_response + response = _complete_aleph_alpha(_dispatch_ctx) elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere": - cohere_key = ( - api_key - or litellm.cohere_key - or get_secret_str("COHERE_API_KEY") - or get_secret_str("CO_API_KEY") - or litellm.api_key - ) - - cohere_route = CohereModelInfo.get_cohere_route(model) - verbose_logger.debug(f"Cohere route: {cohere_route}") - # Set API base based on route - if cohere_route == "v2": - api_base = ( - api_base - or litellm.api_base - or get_secret_str("COHERE_API_BASE") - or "https://api.cohere.com/v2/chat" - ) - # Remove v2/ prefix from model name for the actual API call - if "v2/" in model: - model = model.replace("v2/", "") - else: - api_base = ( - api_base - or litellm.api_base - or get_secret_str("COHERE_API_BASE") - or "https://api.cohere.ai/v1/chat" - ) - - headers = headers or litellm.headers or {} - if headers is None: - headers = {} - - if extra_headers is not None: - headers.update(extra_headers) - - verbose_logger.debug(f"Model: {model}, API Base: {api_base}") - verbose_logger.debug(f"Provider Config: {provider_config}") - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="cohere_chat", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=cohere_key, - provider_config=provider_config, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - ) + response = _complete_cohere_chat(_dispatch_ctx) elif custom_llm_provider == "maritalk": - maritalk_key = ( - api_key - or litellm.maritalk_key - or get_secret("MARITALK_API_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret("MARITALK_API_BASE") - or "https://chat.maritaca.ai/api" - ) - - model_response = openai_like_chat_completion.completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - api_key=maritalk_key, - logging_obj=logging, - custom_llm_provider="maritalk", - custom_prompt_dict=custom_prompt_dict, - ) - - response = model_response + response = _complete_maritalk(_dispatch_ctx) elif custom_llm_provider == "amazon_nova": - api_key = ( - api_key - or litellm.amazon_nova_api_key - or get_secret_str("AMAZON_NOVA_API_KEY") - or litellm.api_key - ) - api_base = ( - api_base - or litellm.api_base - or get_secret_str("AMAZON_NOVA_API_BASE") - or "https://api.nova.amazon.com/v1" - ) - response = openai_like_chat_completion.completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - timeout=timeout, - custom_llm_provider=custom_llm_provider, - custom_prompt_dict=custom_prompt_dict, - ) + response = _complete_amazon_nova(_dispatch_ctx) elif custom_llm_provider == "huggingface": - huggingface_key = ( - api_key - or litellm.huggingface_key - or os.environ.get("HF_TOKEN") - or os.environ.get("HUGGINGFACE_API_KEY") - or litellm.api_key - ) - hf_headers = headers or litellm.headers - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=hf_headers, - model_response=model_response, - api_key=huggingface_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) + response = _complete_huggingface(_dispatch_ctx) elif custom_llm_provider == "oci": - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) + response = _complete_oci(_dispatch_ctx) elif custom_llm_provider == "compactifai": - api_key = ( - api_key or get_secret_str("COMPACTIFAI_API_KEY") or litellm.api_key - ) - - api_base = api_base or "https://api.compactif.ai/v1" - - ## COMPLETION CALL - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) + response = _complete_compactifai(_dispatch_ctx) elif custom_llm_provider == "oobabooga": - custom_llm_provider = "oobabooga" - model_response = oobabooga.completion( - model=model, - messages=messages, - model_response=model_response, - api_base=api_base, # type: ignore - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - api_key=None, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - ) - if "stream" in optional_params and optional_params["stream"] is True: - # don't try to access stream object, - response = CustomStreamWrapper( - model_response, - model, - custom_llm_provider="oobabooga", - logging_obj=logging, - ) - return response - response = model_response + response = _complete_oobabooga(_dispatch_ctx) elif custom_llm_provider == "databricks": - api_base = ( - api_base # for databricks we check in get_llm_provider and pass in the api base from there - or litellm.api_base - or os.getenv("DATABRICKS_API_BASE") - ) - - # set API KEY - api_key = ( - api_key - or litellm.api_key # for databricks we check in get_llm_provider and pass in the api key from there - or litellm.databricks_key - or get_secret("DATABRICKS_API_KEY") - ) - - headers = headers or litellm.headers - - ## COMPLETION CALL - try: - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider="databricks", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e - - if optional_params.get("stream", False): - ## LOGGING - logging.post_call( - input=messages, - api_key=api_key, - original_response=response, - additional_args={"headers": headers}, - ) + response = _complete_databricks(_dispatch_ctx) elif custom_llm_provider == "datarobot": - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=provider_config, - ) + response = _complete_datarobot(_dispatch_ctx) elif custom_llm_provider == "openrouter": - api_base = ( - api_base - or litellm.api_base - or get_secret_str("OPENROUTER_API_BASE") - or "https://openrouter.ai/api/v1" - ) - - api_key = ( - api_key - or litellm.api_key - or litellm.openrouter_key - or get_secret_str("OPENROUTER_API_KEY") - or get_secret_str("OR_API_KEY") - ) - - openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai" - openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM" - - openrouter_headers = { - "HTTP-Referer": openrouter_site_url, - "X-Title": openrouter_app_name, - } - - _headers = headers or litellm.headers - if _headers: - openrouter_headers.update(_headers) - - headers = openrouter_headers - - ## Load Config - config = litellm.OpenrouterConfig.get_config() - for k, v in config.items(): - if k == "extra_body": - # we use openai 'extra_body' to pass openrouter specific params - transforms, route, models - if "extra_body" in optional_params: - optional_params[k].update(v) - else: - optional_params[k] = v - elif k not in optional_params: - optional_params[k] = v - - data = {"model": model, "messages": messages, **optional_params} - - ## COMPLETION CALL - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="openrouter", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - ## LOGGING - logging.post_call( - input=messages, api_key=openai.api_key, original_response=response - ) + response = _complete_openrouter(_dispatch_ctx) elif custom_llm_provider == "vercel_ai_gateway": - api_base = ( - api_base - or litellm.api_base - or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") - or "https://ai-gateway.vercel.sh/v1" - ) - - api_key = ( - api_key or litellm.api_key or get_secret("VERCEL_AI_GATEWAY_API_KEY") - ) - - vercel_site_url = get_secret("VERCEL_SITE_URL") or "https://litellm.ai" - vercel_app_name = get_secret("VERCEL_APP_NAME") or "liteLLM" - - vercel_headers = { - "http-referer": vercel_site_url, - "x-title": vercel_app_name, - } - - _headers = headers or litellm.headers - if _headers: - vercel_headers.update(_headers) - - headers = vercel_headers - - ## Load Config - config = litellm.VercelAIGatewayConfig.get_config() - for k, v in config.items(): - if k == "extra_body": - # we use openai 'extra_body' to pass vercel specific params - providerOptions - if "extra_body" in optional_params: - optional_params[k].update(v) - else: - optional_params[k] = v - elif k not in optional_params: - optional_params[k] = v - - data = {"model": model, "messages": messages, **optional_params} - - ## COMPLETION CALL - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="vercel_ai_gateway", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - ## LOGGING - logging.post_call( - input=messages, api_key=openai.api_key, original_response=response - ) + response = _complete_vercel_ai_gateway(_dispatch_ctx) elif ( custom_llm_provider == "together_ai" or ("togethercomputer" in model) @@ -3504,1114 +5786,75 @@ def completion( # type: ignore "Palm was decommisioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en" ) elif custom_llm_provider == "vertex_ai_beta" or custom_llm_provider == "gemini": - vertex_ai_project = ( - optional_params.pop("vertex_project", None) - or optional_params.pop("vertex_ai_project", None) - or litellm.vertex_project - or get_secret("VERTEXAI_PROJECT") - ) - vertex_ai_location = ( - optional_params.pop("vertex_location", None) - or optional_params.pop("vertex_ai_location", None) - or litellm.vertex_location - or get_secret("VERTEXAI_LOCATION") - ) - vertex_credentials = ( - optional_params.pop("vertex_credentials", None) - or optional_params.pop("vertex_ai_credentials", None) - or get_secret("VERTEXAI_CREDENTIALS") - ) - - gemini_api_key = ( - api_key - or get_api_key_from_env() - or get_secret("PALM_API_KEY") # older palm api key should also work - or litellm.api_key - ) - - api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE") - new_params = safe_deep_copy(optional_params or {}) - response = vertex_chat_completion.completion( # type: ignore - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - gemini_api_key=gemini_api_key, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, - custom_llm_provider=custom_llm_provider, # type: ignore - client=client, - api_base=api_base, - extra_headers=headers, - ) + response = _complete_vertex_ai_beta(_dispatch_ctx) elif custom_llm_provider == "vertex_ai": - vertex_ai_project = ( - optional_params.pop("vertex_project", None) - or optional_params.pop("vertex_ai_project", None) - or litellm.vertex_project - or get_secret("VERTEXAI_PROJECT") - ) - vertex_ai_location = ( - optional_params.pop("vertex_location", None) - or optional_params.pop("vertex_ai_location", None) - or litellm.vertex_location - or get_secret("VERTEXAI_LOCATION") - ) - vertex_credentials = ( - optional_params.pop("vertex_credentials", None) - or optional_params.pop("vertex_ai_credentials", None) - or get_secret("VERTEXAI_CREDENTIALS") - ) - - api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE") - - new_params = safe_deep_copy(optional_params or {}) - model_route = get_vertex_ai_model_route( - model=model, litellm_params=litellm_params - ) - - if model_route == VertexAIModelRoute.PARTNER_MODELS: - model_response = vertex_partner_models_chat_completion.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - api_base=api_base, - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - logging_obj=logging, - acompletion=acompletion, - headers=headers, - custom_prompt_dict=custom_prompt_dict, - timeout=timeout, - client=client, - ) - elif model_route == VertexAIModelRoute.GEMINI: - model_response = vertex_chat_completion.completion( # type: ignore - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - gemini_api_key=None, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, - custom_llm_provider=custom_llm_provider, # type: ignore - client=client, - api_base=api_base, - extra_headers=headers, - ) - elif model_route == VertexAIModelRoute.GEMMA: - # Vertex Gemma Models with custom prediction endpoint - model_response = vertex_gemma_chat_completion.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - api_base=api_base, - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - logging_obj=logging, - acompletion=acompletion, - headers=headers, - custom_prompt_dict=custom_prompt_dict, - timeout=timeout, - client=client, - ) - elif model_route == VertexAIModelRoute.MODEL_GARDEN: - # Vertex Model Garden - OpenAI compatible models - model_response = vertex_model_garden_chat_completion.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - api_base=api_base, - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - logging_obj=logging, - acompletion=acompletion, - headers=headers, - custom_prompt_dict=custom_prompt_dict, - timeout=timeout, - client=client, - ) - elif model_route == VertexAIModelRoute.AGENT_ENGINE: - # Vertex AI Agent Engine (Reasoning Engines) - from litellm.llms.vertex_ai.agent_engine.transformation import ( - VertexAgentEngineConfig, - ) - - vertex_agent_engine_config = VertexAgentEngineConfig() - - # Update litellm_params with vertex credentials - litellm_params["vertex_project"] = vertex_ai_project - litellm_params["vertex_location"] = vertex_ai_location - litellm_params["vertex_credentials"] = vertex_credentials - - model_response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - model_response=model_response, - optional_params=new_params, - litellm_params=litellm_params, # type: ignore - encoding=_get_encoding(), - api_key=None, - api_base=api_base, - logging_obj=logging, - acompletion=acompletion, - timeout=timeout, - client=client, - custom_llm_provider="vertex_ai", - provider_config=vertex_agent_engine_config, - headers=headers or {}, - ) - else: # VertexAIModelRoute.NON_GEMINI - model_response = vertex_ai_non_gemini.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=new_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - vertex_location=vertex_ai_location, - vertex_project=vertex_ai_project, - vertex_credentials=vertex_credentials, - logging_obj=logging, - acompletion=acompletion, - ) - - if ( - "stream" in optional_params - and optional_params["stream"] is True - and acompletion is False - ): - response = CustomStreamWrapper( - model_response, - model, - custom_llm_provider="vertex_ai", - logging_obj=logging, - ) - return response - response = model_response + response = _complete_vertex_ai(_dispatch_ctx) elif custom_llm_provider == "predibase": - tenant_id = ( - optional_params.pop("tenant_id", None) - or optional_params.pop("predibase_tenant_id", None) - or litellm.predibase_tenant_id - or get_secret("PREDIBASE_TENANT_ID") - ) - - if tenant_id is None: - raise ValueError( - "Missing Predibase Tenant ID - Required for making the request. Set dynamically (e.g. `completion(..tenant_id=)`) or in env - `PREDIBASE_TENANT_ID`." - ) - - api_base = ( - api_base - or optional_params.pop("api_base", None) - or optional_params.pop("base_url", None) - or litellm.api_base - or get_secret("PREDIBASE_API_BASE") - ) - - api_key = ( - api_key - or litellm.api_key - or litellm.predibase_key - or get_secret("PREDIBASE_API_KEY") - ) - - _model_response = predibase_chat_completions.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - acompletion=acompletion, - api_base=api_base, - custom_prompt_dict=custom_prompt_dict, - api_key=api_key, - tenant_id=tenant_id, - timeout=timeout, - ) - - if ( - "stream" in optional_params - and optional_params["stream"] is True - and acompletion is False - ): - return _model_response - response = _model_response + response = _complete_predibase(_dispatch_ctx) elif custom_llm_provider == "text-completion-codestral": - api_base = ( - api_base - or optional_params.pop("api_base", None) - or optional_params.pop("base_url", None) - or litellm.api_base - or "https://codestral.mistral.ai/v1/fim/completions" - ) - - api_key = api_key or litellm.api_key or get_secret("CODESTRAL_API_KEY") - - text_completion_model_response = litellm.TextCompletionResponse( - stream=stream - ) - - _model_response = codestral_text_completions.completion( # type: ignore - model=model, - messages=messages, - model_response=text_completion_model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - acompletion=acompletion, - api_base=api_base, - custom_prompt_dict=custom_prompt_dict, - api_key=api_key, - timeout=timeout, - ) - - if ( - "stream" in optional_params - and optional_params["stream"] is True - and acompletion is False - ): - return _model_response - response = _model_response + response = _complete_text_completion_codestral(_dispatch_ctx) elif custom_llm_provider == "text-completion-inception": - passed_api_base = ( - api_base - or optional_params.pop("api_base", None) - or optional_params.pop("base_url", None) - ) - api_base = ( - passed_api_base - or get_secret_str("INCEPTION_API_BASE") - or "https://api.inceptionlabs.ai/v1" - ) - # FIM is served at `/v1/fim/completions`; the OpenAI client appends - # `/completions`, so point it at the `/v1/fim` base. - api_base = api_base.rstrip("/") - if not api_base.endswith("/fim"): - api_base += "/fim" - - # Don't forward the server-managed Inception key to a caller-supplied - # api_base; only resolve it for the default/server base, or when the - # caller passes their own key. - if passed_api_base is None or api_key: - api_key = ( - api_key - or litellm.inception_key - or get_secret_str("INCEPTION_API_KEY") - ) - - _response = openai_text_completions.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, # type: ignore[arg-type] - custom_llm_provider="text-completion-inception", - api_base=api_base, - acompletion=acompletion, - client=client, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - ) - - if ( - optional_params.get("stream", False) is False - and acompletion is False - and text_completion is False - ): - _response = litellm.OpenAITextCompletionConfig().convert_to_chat_model_response_object( - response_object=_response, model_response_object=model_response - ) - - if optional_params.get("stream", False) or acompletion is True: - logging.post_call( - input=messages, - api_key=api_key, - original_response=_response, - additional_args={"headers": headers}, - ) - response = _response + response = _complete_text_completion_inception(_dispatch_ctx) elif custom_llm_provider in ("sagemaker_chat", "sagemaker_nova"): # boto3 reads keys from .env # sagemaker_chat: HF Messages API endpoints # sagemaker_nova: Nova models on SageMaker (OpenAI-compatible) - model_response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - - ## RESPONSE OBJECT - response = model_response + response = _complete_sagemaker_chat(_dispatch_ctx) elif custom_llm_provider == "sagemaker": # boto3 reads keys from .env - model_response = sagemaker_llm.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - custom_prompt_dict=custom_prompt_dict, - hf_model_name=hf_model_name, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - acompletion=acompletion, - ) - - ## RESPONSE OBJECT - response = model_response + response = _complete_sagemaker(_dispatch_ctx) elif custom_llm_provider == "bedrock": # boto3 reads keys from .env - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - - if "aws_bedrock_client" in optional_params: - verbose_logger.warning( - "'aws_bedrock_client' is a deprecated param. Please move to another auth method - https://docs.litellm.ai/docs/providers/bedrock#boto3---authentication." - ) - # Extract credentials for legacy boto3 client and pass thru to httpx - aws_bedrock_client = optional_params.pop("aws_bedrock_client") - creds = aws_bedrock_client._get_credentials().get_frozen_credentials() - - if creds.access_key: - optional_params["aws_access_key_id"] = creds.access_key - if creds.secret_key: - optional_params["aws_secret_access_key"] = creds.secret_key - if creds.token: - optional_params["aws_session_token"] = creds.token - if ( - "aws_region_name" not in optional_params - or optional_params["aws_region_name"] is None - ): - optional_params["aws_region_name"] = ( - aws_bedrock_client.meta.region_name - ) - - bedrock_route = BedrockModelInfo.get_bedrock_route(model) - if bedrock_route == "claude_platform": - provider_config = ProviderConfigManager.get_provider_chat_config( - model=model, - provider=LlmProviders.BEDROCK, - ) - model = BedrockModelInfo.get_claude_platform_model(model) - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="bedrock", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - provider_config=provider_config, - ) - return response - elif bedrock_route == "converse": - model = model.replace("converse/", "") - response = bedrock_converse_chat_completion.completion( - model=model, - messages=messages, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, # type: ignore - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - extra_headers=headers, # Use merged headers instead of original extra_headers - timeout=timeout, - acompletion=acompletion, - client=client, - api_base=api_base, - api_key=api_key, - ) - elif bedrock_route == "converse_like": - model = model.replace("converse_like/", "") - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider="bedrock", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) - else: - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider="bedrock", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - ) + response = _complete_bedrock(_dispatch_ctx) elif custom_llm_provider == "watsonx": - response = watsonx_chat_completion.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - custom_prompt_dict=custom_prompt_dict, - client=client, # pass AsyncOpenAI, OpenAI client - encoding=_get_encoding(), - custom_llm_provider="watsonx", - ) + response = _complete_watsonx(_dispatch_ctx) elif custom_llm_provider == "watsonx_text": - api_key = ( - api_key - or optional_params.pop("apikey", None) - or get_secret_str("WATSONX_APIKEY") - or get_secret_str("WATSONX_API_KEY") - or get_secret_str("WX_API_KEY") - ) - - api_base = ( - api_base - or optional_params.pop( - "url", - optional_params.pop( - "api_base", optional_params.pop("base_url", None) - ), - ) - or get_secret_str("WATSONX_API_BASE") - or get_secret_str("WATSONX_URL") - or get_secret_str("WX_URL") - or get_secret_str("WML_URL") - ) - - wx_credentials = optional_params.pop( - "wx_credentials", - optional_params.pop( - "watsonx_credentials", None - ), # follow {provider}_credentials, same as vertex ai - ) - - token: Optional[str] = None - if wx_credentials is not None: - api_base = wx_credentials.get("url", api_base) - api_key = wx_credentials.get( - "apikey", wx_credentials.get("api_key", api_key) - ) - token = wx_credentials.get( - "token", - wx_credentials.get( - "watsonx_token", None - ), # follow format of {provider}_token, same as azure - e.g. 'azure_ad_token=..' - ) - - if token is not None: - optional_params["token"] = token - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="watsonx_text", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) + response = _complete_watsonx_text(_dispatch_ctx) elif custom_llm_provider == "vllm": - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - model_response = vllm_handler.completion( - model=model, - messages=messages, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - ) - - if ( - "stream" in optional_params and optional_params["stream"] is True - ): ## [BETA] - # don't try to access stream object, - response = CustomStreamWrapper( - model_response, - model, - custom_llm_provider="vllm", - logging_obj=logging, - ) - return response - - ## RESPONSE OBJECT - response = model_response + response = _complete_vllm(_dispatch_ctx) elif custom_llm_provider == "ollama": - api_base = ( - litellm.api_base - or api_base - or get_secret("OLLAMA_API_BASE") - or "http://localhost:11434" - ) - if api_key is not None and "Authorization" not in headers: - headers["Authorization"] = f"Bearer {api_key}" - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="ollama", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) + response = _complete_ollama(_dispatch_ctx) elif custom_llm_provider == "ollama_chat": - api_base = ( - litellm.api_base - or api_base - or get_secret("OLLAMA_API_BASE") - or "http://localhost:11434" - ) - - api_key = ( - api_key - or litellm.ollama_key - or os.environ.get("OLLAMA_API_KEY") - or litellm.api_key - ) - if api_key is not None and "Authorization" not in headers: - headers["Authorization"] = f"Bearer {api_key}" - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="ollama_chat", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - client=client, - ) + response = _complete_ollama_chat(_dispatch_ctx) elif custom_llm_provider == "triton": - api_base = litellm.api_base or api_base - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - ) + response = _complete_triton(_dispatch_ctx) elif custom_llm_provider == "cloudflare": - api_key = ( - api_key - or litellm.cloudflare_api_key - or litellm.api_key - or get_secret("CLOUDFLARE_API_KEY") - ) - account_id = get_secret("CLOUDFLARE_ACCOUNT_ID") - api_base = ( - api_base - or litellm.api_base - or get_secret("CLOUDFLARE_API_BASE") - or f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/" - ) - - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="cloudflare", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements - ) + response = _complete_cloudflare(_dispatch_ctx) elif custom_llm_provider == "petals" or model in litellm.petals_models: - api_base = api_base or litellm.api_base - - custom_llm_provider = "petals" - stream = optional_params.pop("stream", False) - model_response = petals_handler.completion( - model=model, - messages=messages, - api_base=api_base, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - client=client, - ) - if stream is True: ## [BETA] - # Fake streaming for petals - resp_string = model_response["choices"][0]["message"]["content"] - response = CustomStreamWrapper( - resp_string, - model, - custom_llm_provider="petals", - logging_obj=logging, - ) - return response - response = model_response + response = _complete_petals(_dispatch_ctx) elif custom_llm_provider == "snowflake" or model in litellm.snowflake_models: - try: - client = ( - HTTPHandler(timeout=timeout) if stream is False else None - ) # Keep this here, otherwise, the httpx.client closes and streaming is impossible - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - ) - - except Exception as e: - ## LOGGING - log the original exception returned - logging.post_call( - input=messages, - api_key=api_key, - original_response=str(e), - additional_args={"headers": headers}, - ) - raise e + response = _complete_snowflake(_dispatch_ctx) elif custom_llm_provider == "gradient_ai": - api_base = litellm.api_base or api_base - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider="gradient_ai", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - ) + response = _complete_gradient_ai(_dispatch_ctx) elif custom_llm_provider == "bytez": - api_key = ( - api_key - or litellm.bytez_key - or get_secret_str("BYTEZ_API_KEY") - or litellm.api_key - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=bytez_transformation, - ) - - pass + response = _complete_bytez(_dispatch_ctx) elif custom_llm_provider == "lemonade": - api_key = ( - api_key - or litellm.lemonade_key - or get_secret_str("LEMONADE_API_KEY") - or litellm.api_key - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=lemonade_transformation, - ) - - pass + response = _complete_lemonade(_dispatch_ctx) elif custom_llm_provider == "ovhcloud" or model in litellm.ovhcloud_models: - api_key = ( - api_key - or litellm.ovhcloud_key - or get_secret_str("OVHCLOUD_API_KEY") - or litellm.api_key - ) - - api_base = ( - api_base - or litellm.api_base - or get_secret_str("OVHCLOUD_API_BASE") - or "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" - ) - - response = base_llm_http_handler.completion( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - timeout=timeout, # type: ignore - client=client, - custom_llm_provider=custom_llm_provider, - encoding=_get_encoding(), - stream=stream, - provider_config=ovhcloud_transformation, - ) - - pass + response = _complete_ovhcloud(_dispatch_ctx) elif custom_llm_provider == "custom": - url = litellm.api_base or api_base or "" - if url is None or url == "": - raise ValueError( - "api_base not set. Set api_base or litellm.api_base for custom endpoints" - ) - - """ - assume input to custom LLM api bases follow this format: - resp = litellm.module_level_client.post( - api_base, - json={ - 'model': 'meta-llama/Llama-2-13b-hf', # model name - 'params': { - 'prompt': ["The capital of France is P"], - 'max_tokens': 32, - 'temperature': 0.7, - 'top_p': 1.0, - 'top_k': 40, - } - } - ) - - """ - prompt = " ".join([message["content"] for message in messages]) # type: ignore - resp = litellm.module_level_client.post( - url, - headers=headers, - json={ - "model": model, - "params": { - "prompt": [prompt], - "max_tokens": max_tokens, - "temperature": temperature, - "top_p": top_p, - "top_k": kwargs.get("top_k"), - }, - **kwargs.get("extra_body", {}), - }, - ) - response_json = resp.json() - """ - assume all responses from custom api_bases of this format: - { - 'data': [ - { - 'prompt': 'The capital of France is P', - 'output': ['The capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France is PARIS.\nThe capital of France'], - 'params': {'temperature': 0.7, 'top_k': 40, 'top_p': 1}}], - 'message': 'ok' - } - ] - } - """ - string_response = response_json["data"][0]["output"][0] - ## RESPONSE OBJECT - model_response.choices[0].message.content = string_response # type: ignore - model_response.created = int(time.time()) - model_response.model = model - response = model_response + response = _complete_custom(_dispatch_ctx) elif ( custom_llm_provider in litellm._custom_providers ): # Assume custom LLM provider # Get the Custom Handler - custom_handler: Optional[CustomLLM] = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] - - if custom_handler is None: - raise LiteLLMUnknownProvider( - model=model, custom_llm_provider=custom_llm_provider - ) - - ## ROUTE LLM CALL ## - handler_fn = custom_chat_llm_router( - async_fn=acompletion, stream=stream, custom_llm=custom_handler - ) - - headers = headers or litellm.headers or {} - - ## CALL FUNCTION - response = handler_fn( - model=model, - messages=messages, - headers=headers, - model_response=model_response, - print_verbose=print_verbose, - api_key=api_key, - api_base=api_base, - acompletion=acompletion, - logging_obj=logging, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - timeout=timeout, # type: ignore - custom_prompt_dict=custom_prompt_dict, - client=client, # pass AsyncOpenAI, OpenAI client - encoding=_get_encoding(), - ) - if stream is True: - return CustomStreamWrapper( - completion_stream=response, - model=model, - custom_llm_provider=custom_llm_provider, - logging_obj=logging, - ) + response = _complete_custom_providers(_dispatch_ctx) elif custom_llm_provider == "langgraph": # LangGraph - Agent Runtime Provider - from litellm.llms.langgraph.chat.transformation import LangGraphConfig - - ( - api_base, - api_key, - ) = LangGraphConfig()._get_openai_compatible_provider_info( - api_base=api_base or litellm.api_base, - api_key=api_key or litellm.api_key, - ) - - headers = headers or litellm.headers - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - ) + response = _complete_langgraph(_dispatch_ctx) elif custom_llm_provider == "langflow": # LangFlow - Visual AI Agent Platform - from litellm.llms.langflow.chat.transformation import LangFlowConfig - - ( - api_base, - api_key, - ) = LangFlowConfig()._get_openai_compatible_provider_info( - api_base=api_base or litellm.api_base, - api_key=api_key or litellm.api_key, - ) - - headers = headers or litellm.headers - - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - shared_session=shared_session, - custom_llm_provider=custom_llm_provider, - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, - ) + response = _complete_langflow(_dispatch_ctx) else: raise LiteLLMUnknownProvider( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7a5f8b9e1e3..dc7b9838941 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -571,7 +571,7 @@ "output_vector_size": 1536 }, "amazon.titan-embed-text-v2:0": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 2e-08, "litellm_provider": "bedrock", "max_input_tokens": 8192, "max_tokens": 8192, @@ -10684,6 +10684,268 @@ "mode": "chat", "output_cost_per_token": 1.923e-06 }, + "cloudflare/@cf/openai/gpt-oss-120b": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/google/gemma-2b-it-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/meta/llama-3.2-3b-instruct": { + "input_cost_per_token": 5.09e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 80000, + "max_output_tokens": 80000, + "max_tokens": 80000, + "mode": "chat", + "output_cost_per_token": 3.35e-07 + }, + "cloudflare/@cf/meta/llama-guard-3-8b": { + "input_cost_per_token": 4.84e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3e-08 + }, + "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 15000, + "max_output_tokens": 15000, + "max_tokens": 15000, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/deepseek-ai/deepseek-r1-distill-qwen-32b": { + "input_cost_per_token": 4.97e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 80000, + "max_output_tokens": 80000, + "max_tokens": 80000, + "mode": "chat", + "output_cost_per_token": 4.881e-06, + "supports_reasoning": true + }, + "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { + "input_cost_per_token": 1.52e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "chat", + "output_cost_per_token": 2.87e-07 + }, + "cloudflare/@cf/meta/llama-3.2-1b-instruct": { + "input_cost_per_token": 2.7e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 60000, + "max_output_tokens": 60000, + "max_tokens": 60000, + "mode": "chat", + "output_cost_per_token": 2.01e-07 + }, + "cloudflare/@cf/moonshotai/kimi-k2.6": { + "cache_read_input_token_cost": 1.6e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/zai-org/glm-4.7-flash": { + "input_cost_per_token": 6.05e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/meta-llama/llama-2-7b-chat-hf-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { + "input_cost_per_token": 2.93e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 24000, + "max_output_tokens": 24000, + "max_tokens": 24000, + "mode": "chat", + "output_cost_per_token": 2.253e-06, + "supports_function_calling": true + }, + "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { + "input_cost_per_token": 1.7e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 131000, + "max_output_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 1.12e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { + "input_cost_per_token": 6.6e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1e-06 + }, + "cloudflare/@cf/zai-org/glm-5.2": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/nvidia/nemotron-3-120b-a12b": { + "input_cost_per_token": 5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/aisingapore/gemma-sea-lion-v4-27b-it": { + "input_cost_per_token": 3.51e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.55e-07 + }, + "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { + "input_cost_per_token": 5.09e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.35e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/google/gemma-7b-it-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 3500, + "max_output_tokens": 3500, + "max_tokens": 3500, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/google/gemma-4-26b-a4b-it": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/mistralai/mistral-small-3.1-24b-instruct": { + "input_cost_per_token": 3.51e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.55e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { + "input_cost_per_token": 4.85e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.76e-07, + "supports_vision": true + }, + "cloudflare/@cf/openai/gpt-oss-20b": { + "input_cost_per_token": 2e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/meta/llama-4-scout-17b-16e-instruct": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 131000, + "max_output_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 8.5e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/qwen/qwq-32b": { + "input_cost_per_token": 6.6e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 24000, + "max_output_tokens": 24000, + "max_tokens": 24000, + "mode": "chat", + "output_cost_per_token": 1e-06, + "supports_reasoning": true + }, "codestral/codestral-2405": { "input_cost_per_token": 0.0, "litellm_provider": "codestral", @@ -15011,7 +15273,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, @@ -15314,7 +15576,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/qwen3p7-plus": { "cache_read_input_token_cost": 8e-08, @@ -20088,8 +20350,6 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20163,8 +20423,6 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20238,8 +20496,6 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_priority": 8e-07, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20311,8 +20567,6 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 1.7e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20354,8 +20608,6 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20377,8 +20629,6 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20667,8 +20917,6 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "output_cost_per_token_priority": 1e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -21372,8 +21620,6 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21767,6 +22013,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21815,6 +22063,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21859,6 +22109,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21903,6 +22155,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21951,6 +22205,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21998,6 +22254,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22038,6 +22296,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22081,6 +22341,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22126,6 +22388,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22172,6 +22436,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22215,6 +22481,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22258,6 +22526,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22296,8 +22566,6 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -22704,8 +22972,6 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22787,8 +23053,6 @@ "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, @@ -39908,24 +40172,6 @@ "litellm_provider": "fireworks_ai", "mode": "chat" }, - "fireworks_ai/accounts/fireworks/models/whisper-v3": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "fireworks_ai", - "mode": "audio_transcription" - }, - "fireworks_ai/accounts/fireworks/models/whisper-v3-turbo": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "fireworks_ai", - "mode": "audio_transcription" - }, "fireworks_ai/accounts/fireworks/models/yi-34b": { "max_tokens": 4096, "max_input_tokens": 4096, @@ -43061,6 +43307,40 @@ "supports_tool_choice": true, "supports_vision": false }, + "darkbloom/gemma-4-26b": { + "input_cost_per_token": 3e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.65e-07, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "darkbloom/gpt-oss-20b": { + "input_cost_per_token": 1.45e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7e-08, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "deepseek/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 3.625e-09, diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index b27082c361a..3a9ef8db804 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -10,7 +10,7 @@ import os import re from functools import partial from io import IOBase -from typing import Any, Coroutine, Dict, Optional, Union +from typing import Any, Callable, Coroutine, Dict, Optional, Union, cast import httpx @@ -20,6 +20,7 @@ from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.ocr.rust_bridge import RustOcr, load_rust_ocr, rust_ocr_enabled from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -28,6 +29,82 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# +def _timeout_to_seconds( + timeout: Optional[Union[float, httpx.Timeout]], +) -> Optional[float]: + """Convert the Python OCR timeout to a single seconds value for the Rust bridge. + + The Rust HTTP client takes one duration; ``httpx.Timeout`` carries separate + connect/read/write/pool values, so pick the read deadline as the closest + analog to a total-request timeout. + """ + if timeout is None: + return None + if isinstance(timeout, httpx.Timeout): + return timeout.read + return float(timeout) + + +def _run_rust_ocr( + rust_ocr: RustOcr, + logging_obj: LiteLLMLoggingObj, + provider_config: BaseOCRConfig, + resolve_api_key: Callable[[str], Optional[str]], + model: str, + document: dict[str, object], + api_key: Optional[str], + api_base: Optional[str], + optional_params: dict[str, object], + litellm_params: dict[str, object], + timeout_seconds: Optional[float], +) -> OCRResponse: + """Run the Mistral OCR call through the Rust bridge and wrap the result. + + Resolves the key the same way the Python path does so secret-manager backends + (AWS/Azure/GCP/Vault) work; the Rust bridge's own fallback only reads the + process environment. The request that Rust actually sends (resolved URL and + headers) is mirrored into pre_call so logs match the wire. Dependencies are + injected so this stays unit-testable without patching module globals. + """ + resolved_api_key = api_key or resolve_api_key("MISTRAL_API_KEY") + resolved_headers = provider_config.validate_environment( + headers={}, + model=model, + api_key=resolved_api_key, + api_base=api_base, + litellm_params=litellm_params, + ) + resolved_complete_url = provider_config.get_complete_url( + api_base=api_base, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + ) + logging_obj.pre_call( + input="OCR document processing", + api_key=resolved_api_key, + additional_args={ + "complete_input_dict": { + "model": model, + "document": document, + **optional_params, + }, + "api_base": resolved_complete_url, + "headers": resolved_headers, + }, + ) + return OCRResponse.model_validate( + rust_ocr( + model=model, + document=document, + api_key=resolved_api_key, + api_base=api_base, + optional_params=optional_params, + timeout_seconds=timeout_seconds, + ) + ) + + @client async def aocr( model: str, @@ -220,7 +297,7 @@ def ocr( """ local_vars = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("aocr", False) is True @@ -261,7 +338,6 @@ def ocr( if dynamic_api_base: api_base = dynamic_api_base - # Get provider config ocr_provider_config: Optional[BaseOCRConfig] = ( ProviderConfigManager.get_provider_ocr_config( model=model, @@ -278,17 +354,14 @@ def ocr( f"OCR call - model: {model}, provider: {custom_llm_provider}" ) - # Get litellm params using GenericLiteLLMParams (same as responses API) litellm_params = GenericLiteLLMParams(**kwargs) - # Extract OCR-specific parameters from kwargs supported_params = ocr_provider_config.get_supported_ocr_params(model=model) non_default_params = {} for param in supported_params: if param in kwargs: non_default_params[param] = kwargs.pop(param) - # Map parameters to provider-specific format optional_params = ocr_provider_config.map_ocr_params( non_default_params=non_default_params, optional_params={}, @@ -297,7 +370,8 @@ def ocr( verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}") - # Pre Call logging + effective_timeout = timeout or request_timeout + litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model, @@ -309,12 +383,35 @@ def ocr( custom_llm_provider=custom_llm_provider, ) - # Call the handler - pass document dict directly + # Optional Rust path: hand the whole Mistral OCR call to the Rust bridge. + if custom_llm_provider == "mistral" and rust_ocr_enabled(): + rust_ocr = load_rust_ocr() + if rust_ocr is None: + verbose_logger.debug( + "Rust OCR bridge unavailable; falling back to Python path" + ) + else: + from litellm.secret_managers.main import get_secret_str + + return _run_rust_ocr( + rust_ocr=rust_ocr, + logging_obj=litellm_logging_obj, + provider_config=ocr_provider_config, + resolve_api_key=get_secret_str, + model=model, + document=document, + api_key=api_key, + api_base=api_base, + optional_params=optional_params, + litellm_params=dict(litellm_params), + timeout_seconds=_timeout_to_seconds(effective_timeout), + ) + response = base_llm_http_handler.ocr( model=model, - document=document, # Pass the entire document dict + document=document, optional_params=optional_params, - timeout=timeout or request_timeout, + timeout=effective_timeout, logging_obj=litellm_logging_obj, api_key=api_key, api_base=api_base, diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py new file mode 100644 index 00000000000..61f9e8ca69a --- /dev/null +++ b/litellm/ocr/rust_bridge.py @@ -0,0 +1,74 @@ +""" +Optional Rust-backed OCR path. + +Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` entrypoint +then routes supported Mistral calls through the compiled ``litellm_python_bridge`` +extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust. + +No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py`` +can import it statically without forming an import cycle. +""" + +from __future__ import annotations + +from typing import Final, Protocol, cast + + +class RustOcr(Protocol): + """Signature of the compiled ``litellm_python_bridge.ocr`` entrypoint.""" + + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: ... + + +class _Unset: + """Sentinel type so ``ocr=None`` can clear a prior injection while omission preserves it.""" + + +_UNSET: Final[_Unset] = _Unset() + +_rust_ocr_enabled = False +_rust_ocr_impl: RustOcr | None = None + + +def use_litellm_rust( + enabled: bool = True, *, ocr: RustOcr | None | _Unset = _UNSET +) -> None: + """Route supported OCR calls through the Rust ``litellm_python_bridge`` extension. + + ``ocr`` injects the bridge callable; when omitted the compiled extension is + loaded on demand and any previously injected bridge is preserved. Pass + ``ocr=None`` explicitly to clear a prior injection. + """ + global _rust_ocr_enabled, _rust_ocr_impl + _rust_ocr_enabled = enabled + if not isinstance(ocr, _Unset): + _rust_ocr_impl = ocr + + +def rust_ocr_enabled() -> bool: + """Whether the Rust OCR path has been turned on via ``use_litellm_rust()``.""" + return _rust_ocr_enabled + + +def load_rust_ocr() -> RustOcr | None: + """Return the Rust OCR callable, or ``None`` when no bridge is available. + + Prefers an injected implementation, otherwise loads the compiled + ``litellm_python_bridge`` extension; a missing extension yields ``None`` so + the caller can fall back to the Python path instead of hard-failing. + """ + if _rust_ocr_impl is not None: + return _rust_ocr_impl + try: + import litellm_python_bridge + except ImportError: + return None + return cast(RustOcr, litellm_python_bridge.ocr) diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index db6183edaa0..dd7712aabca 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1835,6 +1835,23 @@ "interactions": true } }, + "darkbloom": { + "display_name": "Darkbloom (`darkbloom`)", + "url": "https://docs.litellm.ai/docs/providers/darkbloom", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "predibase": { "display_name": "Predibase (`predibase`)", "url": "https://docs.litellm.ai/docs/providers/predibase", diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md new file mode 100644 index 00000000000..8eebc3ea3b3 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -0,0 +1,95 @@ +# Experimental MCP Server Change Guidelines + +Read @../../../../CLAUDE.md and @CLAUDE.md before changing this package. + +This directory owns the proxy-hosted MCP server implementation. Keep changes +inside the module that owns the behavior, and only reach outside this package +when the public type contract, database schema, dashboard, or cross-proxy route +wiring must change with it. + +## File Structure + +Respect the current package boundaries: + +```text +litellm/proxy/_experimental/mcp_server/ + AGENTS.md + CLAUDE.md + server.py # ASGI/MCP route handling, sessions, tool calls [PR7: 7-arm only — move BYOK/OAuth pre-fetch into resolver] + mcp_server_manager.py # upstream server registry, clients, tool routing [PR7: _create_mcp_client swaps resolve_mcp_auth -> resolve_credentials] + auth/ + user_api_key_auth_mcp.py # LiteLLM admission auth and MCP request headers + token_exchange.py # OAuth token exchange handling [unchanged; V1TokenExchangeAdapter delegates here] + litellm_auth_handler.py # authenticated-user adapter for MCP sessions + outbound_credentials/ # NEW — typed upstream-credential resolution (resolve_credentials + arms) + __init__.py # public surface: resolve_credentials, the configs, CredError + result.py # Ok | Error union (pure stdlib) + types.py # AuthConfig union, CredError, Subject, ServerSpec + httpx_auth.py # NoOpAuth, StaticHeaderAuth (every mode -> one httpx.Auth) + resolver.py # resolve_credentials(): exhaustive per-mode match + assert_never + seams.py # injected Protocols (one per cache-touching mode) + v1_adapters.py # v1-backed seam bodies; delegate to auth/oauth2/db owners + adapter.py # to_subject / to_server_spec / raise_public (v1 <-> v2 boundary) + discoverable_endpoints.py # MCP OAuth metadata, authorize, token, callback + byok_oauth_endpoints.py # BYOK OAuth UI/API flow + oauth_utils.py # redirect URI and proxy base URL validation + oauth2_token_cache.py # OAuth2 and per-user token resolution/cache [PR7: resolve_mcp_auth removed; cache class stays, V1OAuth2CacheAdapter delegates to async_get_token] + db.py # MCP server, credential, env var, submission DB access [unchanged; V1ByokStore delegates to _get_byok_credential / get_user_credential] + toolset_db.py # MCP toolset DB access + rest_endpoints.py # proxy REST facade for listing/calling MCP tools [PR7: 7-arm only — pass identity + inbound token down instead of mcp_auth_header] + openapi_to_mcp_generator.py# OpenAPI spec to MCP tool generation + sampling_handler.py # MCP sampling to LiteLLM completion flow + elicitation_handler.py # MCP elicitation relay flow + semantic_tool_filter.py # semantic filtering of available MCP tools + guardrail_translation/ + handler.py # MCP guardrail result translation + sse_transport.py # SSE transport implementation + mcp_context.py # contextvars for MCP request/session metadata + mcp_debug.py # debug helpers + tool_registry.py # in-memory MCP tool registry helpers + cost_calculator.py # MCP tool cost calculation + ui_session_utils.py # dashboard session auth context helpers + utils.py # shared primitives used by several modules +``` + +Do not add broad catch-all modules. Prefer the existing owner above, and add a +new file only for a distinct capability that would otherwise make an existing +module materially harder to understand. + +## Implementation Rules + +- Preserve the boundary between LiteLLM admission auth and upstream MCP auth. + Admission belongs in `auth/user_api_key_auth_mcp.py`; upstream token exchange, + delegated auth, per-user OAuth, BYOK, and raw header forwarding belong in the + dedicated OAuth/header modules. +- Treat `none`, bearer/API key, OAuth, OAuth token exchange, delegated upstream + auth, SSE, streamable HTTP, and stdio as separate flows. Do not collapse them + behind a single generic branch unless tests prove every mode still behaves + correctly. +- Be especially careful with `available_on_public_internet: false` combined with + `delegate_auth_to_upstream: true`. The local `CLAUDE.md` explains the anonymous + upstream PKCE path that must remain intentional. +- Keep database-backed fields in sync across migrations, typed models under + `litellm/types/mcp.py` or `litellm/types/mcp_server/`, config loading, this + package, and dashboard state when the field is user-visible. +- Use the official MCP SDK types and established LiteLLM Pydantic models where + they exist. Avoid untyped protocol dictionaries at package boundaries. +- Keep security-sensitive logic easy to audit. Header forwarding, IP filtering, + public internet checks, token storage, env var interpolation, and credential + encryption need focused tests for both allowed and rejected paths. +- Avoid adding comments to new code unless they explain non-obvious security or + protocol behavior. Prefer clear names and small functions. + +## Tests + +Mirror this package under `tests/test_litellm/proxy/_experimental/mcp_server/`. +For regressions, extend the existing mapped test file instead of creating a new +one. Use subdirectories that match the implementation path, such as +`auth/test_token_exchange.py` for `auth/token_exchange.py` and +`guardrail_translation/test_mcp_guardrail_handler.py` for +`guardrail_translation/handler.py`. + +Use `tests/mcp_tests/` only when extending an existing broader MCP integration +scenario that already lives there. Route, auth, tool listing, tool execution, +OAuth, sampling, elicitation, DB, and dashboard-session changes should have +focused coverage in the mirrored `tests/test_litellm/...` path first. 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 e47fc84b533..90108de25c3 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 @@ -12,6 +12,7 @@ from litellm.proxy._types import ( LiteLLM_TeamTable, ProxyException, SpecialHeaders, + SpecialMCPServerNames, UserAPIKeyAuth, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -642,6 +643,15 @@ class MCPRequestHandler: user_api_key_auth ) ) + + # The key explicitly opted out of every MCP server. This overrides + # team inheritance and additive grants (mirrors no-default-models). + if ( + SpecialMCPServerNames.no_mcp_servers.value + in allowed_mcp_servers_for_key + ): + return [] + allowed_mcp_servers_for_team = ( await MCPRequestHandler._get_allowed_mcp_servers_for_team( user_api_key_auth @@ -1058,6 +1068,13 @@ class MCPRequestHandler: if key_object_permission is None: return [] + # Sentinel opt-out: surface it unexpanded so the caller can short-circuit + # to zero servers instead of inheriting the team. + if SpecialMCPServerNames.no_mcp_servers.value in ( + key_object_permission.mcp_servers or [] + ): + return [SpecialMCPServerNames.no_mcp_servers.value] + # Permission entries may be server_ids OR names/aliases — expand to ids. direct_mcp_servers = global_mcp_server_manager.expand_permission_list( key_object_permission.mcp_servers or [] diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index afec884cd96..5e704b889ae 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -80,6 +80,7 @@ from litellm.proxy._types import ( MCPEnvVar, MCPTransport, MCPTransportType, + SpecialMCPServerNames, UserAPIKeyAuth, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -1349,6 +1350,17 @@ class MCPServerManager: allow_all_server_ids = self.get_allow_all_keys_server_ids() try: + # The key explicitly opted out of every MCP server. Return zero before + # layering on allow_all_keys servers so the opt-out is absolute. + key_object_permission = ( + user_api_key_auth.object_permission if user_api_key_auth else None + ) + if key_object_permission is not None and ( + SpecialMCPServerNames.no_mcp_servers.value + in (key_object_permission.mcp_servers or []) + ): + return [] + # Check if object_permission.mcp_servers is explicitly set has_explicit_object_permission = False if user_api_key_auth and user_api_key_auth.object_permission: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py new file mode 100644 index 00000000000..73166a45d6e --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py @@ -0,0 +1,73 @@ +"""Typed upstream-credential resolution for MCP servers. + +This subpackage houses the typed credential vocabulary and the ``resolve_credentials`` +dispatch. A server declares one per-mode config from the ``AuthConfig`` discriminated union; +``UpstreamCredentialProvider.resolve_credentials`` selects one arm and returns an ``httpx.Auth`` +or a typed ``CredError``. Failures are modeled as values via :mod:`.result` (``Result[T, +CredError]``) rather than raised, so every seam is total. Nothing here is wired onto a live +request path yet. +""" + +from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + NoOpAuth, + StaticHeaderAuth, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import ( + UpstreamCredentialProvider, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + Ambient, + ApiKeyConfig, + ApiKeySource, + AssumeRole, + AuthConfig, + AuthorizationCodeConfig, + AuthSpecKind, + AwsCredentialSource, + AwsSigV4Config, + Byok, + ClientCredentialsConfig, + CredError, + NoneConfig, + PassthroughConfig, + ServerSpec, + SharedKey, + StaticKeys, + Subject, + TokenExchangeConfig, + parse_auth_spec_kind, +) + +__all__ = [ + "Ok", + "Error", + "Result", + "NoOpAuth", + "StaticHeaderAuth", + "UpstreamCredentialProvider", + "AuthSpecKind", + "CredError", + "Subject", + "ServerSpec", + "AuthConfig", + "parse_auth_spec_kind", + "AuthorizationCodeConfig", + "ClientCredentialsConfig", + "TokenExchangeConfig", + "ApiKeyConfig", + "ApiKeySource", + "SharedKey", + "Byok", + "PassthroughConfig", + "NoneConfig", + "AwsSigV4Config", + "AwsCredentialSource", + "StaticKeys", + "AssumeRole", + "Ambient", +] diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py new file mode 100644 index 00000000000..2345fa98123 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py @@ -0,0 +1,45 @@ +"""Concrete `httpx.Auth` objects the resolver returns for the self-contained modes. + +These are the egress credential as the SDK consumes it: an `httpx.Auth` attached to the +upstream `AsyncClient`. The OAuth-flow modes (`authorization_code`, `client_credentials`, +`token_exchange`) return SDK-provided auth objects instead and land later. + +`auth_flow` mutating the outbound request is the `httpx.Auth` contract, not a house-style +violation: the request is httpx's object, and these carry no state of their own. +""" + +from __future__ import annotations + +from collections.abc import Generator + +import httpx +from pydantic import SecretStr + + +class NoOpAuth(httpx.Auth): + """Attaches nothing — the `none` mode (and the seam-level default).""" + + def auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + yield request + + +class StaticHeaderAuth(httpx.Auth): + """Sets one fixed header on every request — the `api_key` family and `passthrough`. + + The header value is a live credential (a bearer token, an API key, a forwarded user + token), so it is held as a `SecretStr` and unwrapped only when written onto the request. + That keeps it masked in reprs, `vars()`, tracebacks, and structured logs, matching the + `SecretStr` discipline the config models use. + """ + + def __init__(self, header_value: str, header_name: str = "Authorization") -> None: + self.header_name = header_name + self._header_value = SecretStr(header_value) + + def auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + request.headers[self.header_name] = self._header_value.get_secret_value() + yield request diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py new file mode 100644 index 00000000000..7bcdb3e6529 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -0,0 +1,70 @@ +"""The one credential resolver: dispatch on the declared mode, fail closed. + +`resolve_credentials` selects exactly one arm off the server's typed `config` and either +produces an `httpx.Auth` or returns a typed `CredError`. The `match` is over the `AuthConfig` +variant, so each arm receives its own fully-typed config with no field-presence inference and +no precedence cascade. It is wildcard-free with an `assert_never` tail, so adding a mode without +an arm fails the type gate (basedpyright `reportMatchNotExhaustive`); a bypassed gate fails loudly +at runtime instead of returning `None`. + +This skeleton ships every arm as a `not_implemented` stub. Each mode's real body, with its +injected seam, lands in its own follow-up PR; until then the arm returns a typed error rather +than silently producing no credential. Pure v2: no imports from v1. +""" + +from __future__ import annotations + +import httpx +from typing_extensions import assert_never + +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ApiKeyConfig, + AuthorizationCodeConfig, + AuthSpecKind, + AwsSigV4Config, + ClientCredentialsConfig, + CredError, + NoneConfig, + PassthroughConfig, + ServerSpec, + Subject, + TokenExchangeConfig, +) + + +class UpstreamCredentialProvider: + """Produces the one `httpx.Auth` for a `(subject, upstream)` pair, per declared mode. + + Collaborators (the per-mode credential stores and token fetchers) are injected as each arm + is built; the skeleton needs none, since every arm is a stub. + """ + + async def resolve_credentials( + self, subject: Subject, server: ServerSpec + ) -> Result[httpx.Auth, CredError]: + match server.config: + case NoneConfig(): + return _not_implemented(AuthSpecKind.none) + case ApiKeyConfig(): + return _not_implemented(AuthSpecKind.api_key) + case PassthroughConfig(): + return _not_implemented(AuthSpecKind.passthrough) + case ClientCredentialsConfig(): + return _not_implemented(AuthSpecKind.client_credentials) + case TokenExchangeConfig(): + return _not_implemented(AuthSpecKind.token_exchange) + case AuthorizationCodeConfig(): + return _not_implemented(AuthSpecKind.authorization_code) + case AwsSigV4Config(): + return _not_implemented(AuthSpecKind.aws_sigv4) + assert_never(server.config) + + +def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: + return Error( + CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet") + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/result.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/result.py new file mode 100644 index 00000000000..a612e8510f5 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/result.py @@ -0,0 +1,54 @@ +"""A tagged-union ``Result`` the type checker can actually narrow. + +``Ok`` and ``Error`` are separate frozen classes joined by a ``Union`` alias, so +reaching for ``result.ok`` before eliminating the ``Error`` arm (via ``isinstance`` +or a ``match`` pattern) is a type error rather than a runtime ``AttributeError``. A +single class carrying both payload fields would make that unguarded access invisible +to the type checker. + +Both variants are covariant and frozen; the absent side defaults to ``Never`` so a +bare ``Ok(value)`` or ``Error(err)`` infers fully and is assignable to any ``Result`` +whose matching side fits. + +``is_ok`` / ``is_error`` are runtime predicates that also narrow via their ``Literal`` +returns; inside strictly typed code, discriminate with ``match`` or ``isinstance``. + +This is the shared ``Result`` shape for the ``outbound_credentials`` resolver: every +seam returns ``Result[T, CredError]`` instead of raising, so each failure is a value +the caller must handle rather than an exception that can slip past the type checker. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Generic, Literal, TypeAlias + +from typing_extensions import Never, TypeVar + +_TOk_co = TypeVar("_TOk_co", covariant=True, default=Never) +_TError_co = TypeVar("_TError_co", covariant=True, default=Never) + + +@dataclass(frozen=True) +class Ok(Generic[_TOk_co, _TError_co]): + ok: _TOk_co + + def is_ok(self) -> Literal[True]: + return True + + def is_error(self) -> Literal[False]: + return False + + +@dataclass(frozen=True) +class Error(Generic[_TOk_co, _TError_co]): + error: _TError_co + + def is_ok(self) -> Literal[False]: + return False + + def is_error(self) -> Literal[True]: + return True + + +Result: TypeAlias = Ok[_TOk_co, _TError_co] | Error[_TOk_co, _TError_co] diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py new file mode 100644 index 00000000000..2088dc77252 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -0,0 +1,334 @@ +"""The upstream-credential vocabulary — the typed seam the resolver dispatches on. + +This module ships the data types only; the resolver lands in a later PR. It is the contract +the credential build implements and the spec tests assert against. + +Design invariants encoded here: + +- **Mode is the single source of truth.** A server declares exactly one per-mode `config` + (the `AuthConfig` discriminated union); `auth_spec_kind` is *derived* from it, never a + second field that can drift. The resolver dispatches on the config variant, one arm per + mode. No field-presence inference, no precedence cascade. +- **Illegal states unrepresentable.** Each mode's config is its own frozen model holding + only that mode's fields — an `aws_sigv4` server cannot hold OAuth fields, and a config + missing a required field is rejected at construction, not at call time. +- **Fail-closed at the boundary.** A raw mode string can only enter through + `parse_auth_spec_kind()`, which returns a typed `CredError`. +- **Errors as values.** Every seam returns `Result[_, CredError]`; only edge adapters raise. +- **No v1 imports.** This vocabulary stays free of `MCPServer` and the rest of v1; the + v1 -> v2 adapter maps onto these types in a later PR. + +Sum types are Expression `@tagged_union`s discriminated on a `Literal` `tag`, matched via +`self.tag` with an `assert_never` tail; `Result` is this package's vendored `Ok | Error` +union (see `result.py`), not `expression.Result`. +""" + +from __future__ import annotations + +from enum import Enum +from typing import Annotated, Literal + +from expression import case, tag, tagged_union +from pydantic import BaseModel, ConfigDict, Field, SecretStr +from typing_extensions import assert_never + +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) + + +class AuthSpecKind(str, Enum): + """The server's statically-declared upstream-auth mode — derived from its `config`. + + Covers v1's full `MCPAuth` surface, not only OAuth grants: the three grant modes, the + collapsed static-header family, client passthrough, no-auth, and AWS request signing. + BYOK is *not* a member: it is the `api_key` mode seeded per-user, a source selector + inside that arm. The static-header schemes v1 splits into separate `MCPAuth` values + (`bearer_token`/`api_key`/`basic`/`token`/`authorization`) collapse into `api_key`; the + scheme is a parameter the arm carries, not its own mode. + """ + + authorization_code = "authorization_code" # per-user 3LO; gateway-stored token + client_credentials = "client_credentials" # gateway service account (M2M) + token_exchange = "token_exchange" # RFC 8693: token endpoint + subject_token (OBO) + api_key = "api_key" # static header, any scheme (BYOK = per-user-seeded source) + passthrough = "passthrough" # client forwards an upstream-audience token + none = "none" # no upstream credential; resolve yields a no-op auth, never an error + aws_sigv4 = "aws_sigv4" # AWS SigV4 per-request signing (e.g. Bedrock AgentCore) + + +@tagged_union(frozen=True) +class CredError: + """Why a credential could not be produced. Fail-closed: an arm yields this or an `httpx.Auth`. + + Discriminated on the `Literal` `tag`; consumers `match self.tag` (see `summary`) so the + type checker can prove exhaustiveness. Construct via the `of_*` factories. + """ + + tag: Literal[ + "unauthorized", + "misconfigured", + "upstream_unavailable", + "unsupported_mode", + "precondition_required", + "not_implemented", + ] = tag() + + unauthorized: str = ( + case() + ) # no usable credential for this (subject, server) -> 401 challenge + misconfigured: str = ( + case() + ) # the declared mode is missing required config -> 5xx (operator) + upstream_unavailable: str = ( + case() + ) # the IdP / token endpoint could not be reached -> 503 + unsupported_mode: str = ( + case() + ) # a raw mode string did not parse into AuthSpecKind (boundary) + precondition_required: str = ( + case() + ) # a required per-user value (e.g. an env var) has not been provided -> 412 + not_implemented: str = ( + case() + ) # the declared mode's resolver arm is not built yet -> 501 (not operator error) + + @staticmethod + def of_unauthorized(detail: str) -> CredError: + return CredError(unauthorized=detail) + + @staticmethod + def of_misconfigured(detail: str) -> CredError: + return CredError(misconfigured=detail) + + @staticmethod + def of_upstream_unavailable(detail: str) -> CredError: + return CredError(upstream_unavailable=detail) + + @staticmethod + def of_unsupported_mode(detail: str) -> CredError: + return CredError(unsupported_mode=detail) + + @staticmethod + def of_precondition_required(detail: str) -> CredError: + return CredError(precondition_required=detail) + + @staticmethod + def of_not_implemented(detail: str) -> CredError: + return CredError(not_implemented=detail) + + @property + def summary(self) -> str: + # Exhaustiveness: every Literal tag has an arm; the trailing assert_never typechecks + # only while that stays true (a `case _` would defeat reportMatchNotExhaustive). + match self.tag: + case "unauthorized": + return f"unauthorized: {self.unauthorized}" + case "misconfigured": + return f"misconfigured: {self.misconfigured}" + case "upstream_unavailable": + return f"upstream unavailable: {self.upstream_unavailable}" + case "unsupported_mode": + return self.unsupported_mode + case "precondition_required": + return f"precondition required: {self.precondition_required}" + case "not_implemented": + return f"not implemented: {self.not_implemented}" + assert_never(self.tag) + + +class AuthorizationCodeConfig(BaseModel): + """Per-user 3LO; the gateway is the OAuth client and stores the user's token. + + Endpoints are discovered (RFC 9728 -> RFC 8414) and the client is registered via DCR + (RFC 7591), so the common case carries none of the fields below; they are optional manual + overrides for IdPs without discovery / DCR. The per-user token is read from the token store + at resolve time, not held here. + """ + + model_config = ConfigDict(frozen=True) + kind: Literal[AuthSpecKind.authorization_code] = AuthSpecKind.authorization_code + scopes: tuple[str, ...] = () + client_id: str | None = None + client_secret: SecretStr | None = None + authorization_url: str | None = None + token_url: str | None = None + + +class ClientCredentialsConfig(BaseModel): + """M2M service account; one upstream identity for every user. + + Fields are optional so the config can be built incomplete: a value may be supplied at + runtime (`token_url` via RFC 8414 discovery, `client_id`/`secret` via DCR), and the + resolver arm raises `CredError.misconfigured` when a needed field is still absent. + """ + + model_config = ConfigDict(frozen=True) + kind: Literal[AuthSpecKind.client_credentials] = AuthSpecKind.client_credentials + client_id: str | None = None + client_secret: SecretStr | None = None + token_url: str | None = None + scopes: tuple[str, ...] = () + + +class TokenExchangeConfig(BaseModel): + """RFC 8693 OBO; swap the caller's live subject_token for a token bound to the upstream's + audience (`server.resource`, RFC 8707). The gateway authenticates to the exchange endpoint + as an OAuth client (`client_id`/`client_secret`); the inbound token is sent only to that + endpoint, never to the upstream. + """ + + model_config = ConfigDict(frozen=True) + kind: Literal[AuthSpecKind.token_exchange] = AuthSpecKind.token_exchange + subject_token_type: str = "urn:ietf:params:oauth:token-type:access_token" + token_exchange_endpoint: str | None = None + client_id: str | None = None + client_secret: SecretStr | None = None + scopes: tuple[str, ...] = () + + +class SharedKey(BaseModel): + """A fixed key configured on the server, identical for every caller.""" + + model_config = ConfigDict(frozen=True) + source: Literal["shared"] = "shared" + value: SecretStr + + +class Byok(BaseModel): + """A key the user brings via the entry flow, stored per-user and pulled from the credential + store at resolve time. Missing means the user must provide it, a 401 + WWW-Authenticate + challenge.""" + + model_config = ConfigDict(frozen=True) + source: Literal["byok"] = "byok" + + +ApiKeySource = Annotated[SharedKey | Byok, Field(discriminator="source")] + + +class ApiKeyConfig(BaseModel): + """A fixed credential injected as a header. The value is shared (in config) or seeded + per-user (pulled from the store); `header_name` and `value_prefix` say where and how it is + written, modeled like OpenAPI's apiKey scheme so any upstream convention is expressible + (Authorization + Bearer, a raw value on X-API-Key, Ocp-Apim-Subscription-Key, etc.). + """ + + model_config = ConfigDict(frozen=True) + kind: Literal[AuthSpecKind.api_key] = AuthSpecKind.api_key + header_name: str = "Authorization" + value_prefix: str = "Bearer" + key_source: ApiKeySource + + def header(self, value: str) -> tuple[str, str]: + formatted = f"{self.value_prefix} {value}" if self.value_prefix else value + return self.header_name, formatted + + +class PassthroughConfig(BaseModel): + """Client-driven upstream OAuth; the gateway forwards the client's upstream token.""" + + model_config = ConfigDict(frozen=True) + kind: Literal[AuthSpecKind.passthrough] = AuthSpecKind.passthrough + + +class NoneConfig(BaseModel): + """No upstream credential; the request is sent unauthenticated.""" + + model_config = ConfigDict(frozen=True) + kind: Literal[AuthSpecKind.none] = AuthSpecKind.none + + +class StaticKeys(BaseModel): + """Long-lived AWS access keys configured on the server.""" + + model_config = ConfigDict(frozen=True) + source: Literal["static_keys"] = "static_keys" + access_key_id: str + secret_access_key: SecretStr + session_token: SecretStr | None = None + + +class AssumeRole(BaseModel): + """An IAM role the gateway assumes via STS for short-lived, auto-refreshed credentials.""" + + model_config = ConfigDict(frozen=True) + source: Literal["assume_role"] = "assume_role" + role_arn: str + session_name: str | None = None + external_id: str | None = None + + +class Ambient(BaseModel): + """The environment's default AWS credential chain (instance profile, IRSA, env vars).""" + + model_config = ConfigDict(frozen=True) + source: Literal["ambient"] = "ambient" + + +AwsCredentialSource = Annotated[ + StaticKeys | AssumeRole | Ambient, Field(discriminator="source") +] + + +class AwsSigV4Config(BaseModel): + """AWS SigV4 per-request signing for an AWS-hosted upstream (e.g. Bedrock AgentCore). The + gateway signs with its own AWS identity, never the caller's; `credentials` selects how that + identity is obtained, defaulting to the ambient credential chain.""" + + model_config = ConfigDict(frozen=True) + kind: Literal[AuthSpecKind.aws_sigv4] = AuthSpecKind.aws_sigv4 + region: str + service: str = "bedrock-agentcore" + credentials: AwsCredentialSource = Ambient() + + +AuthConfig = Annotated[ + AuthorizationCodeConfig + | ClientCredentialsConfig + | TokenExchangeConfig + | ApiKeyConfig + | PassthroughConfig + | NoneConfig + | AwsSigV4Config, + Field(discriminator="kind"), +] + + +class Subject(BaseModel): + """The validated inbound principal. NOT the v1 request object and NOT the LiteLLM key.""" + + model_config = ConfigDict(frozen=True) + + tenant_id: str + subject_id: str + # Opaque, already-validated inbound identity. Only `token_exchange` / `passthrough` read it. + inbound_token: SecretStr | None = None + + +class ServerSpec(BaseModel): + """The declared upstream. A v2-native type; the v1 -> v2 adapter maps onto this.""" + + model_config = ConfigDict(frozen=True) + + server_id: str + resource: str # RFC 8707 audience URI this upstream's tokens are bound to + config: AuthConfig + + @property + def auth_spec_kind(self) -> AuthSpecKind: + return self.config.kind + + +def parse_auth_spec_kind(raw: str) -> Result[AuthSpecKind, CredError]: + """Boundary parser — the *only* place an unknown mode is handled, and it fails closed. + + Inside the core the mode is always a valid `AuthSpecKind`, so the resolver never needs a + wildcard arm and basedpyright can prove its `match` exhaustive. + """ + try: + return Ok(AuthSpecKind(raw)) + except ValueError: + return Error(CredError.of_unsupported_mode(f"unknown auth_spec_kind: {raw!r}")) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 08e42e918e9..e891425274f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -63,7 +63,11 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ( + ProxyException, + SpecialMCPServerNames, + UserAPIKeyAuth, +) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, @@ -229,6 +233,28 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: return False +def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: + """Map a ``ProxyException`` to an ``HTTPException`` that preserves its real + status code and headers. + + ``user_api_key_auth`` raises ``ProxyException`` (not ``HTTPException``) on + auth failures. The MCP ASGI handlers re-raise ``HTTPException`` to keep the + status and any ``WWW-Authenticate`` challenge, but a ``ProxyException`` would + otherwise fall through to their generic handler and be flattened to a 500 — + dropping the 401 + challenge an OAuth client needs to re-authenticate, so the + tool call surfaces as a cancelled/terminated session instead. + """ + try: + status_code = int(exc.code) + except (TypeError, ValueError): + status_code = 500 + return HTTPException( + status_code=status_code, + detail=exc.message, + headers=exc.headers or None, + ) + + if MCP_AVAILABLE: from mcp.server import Server from mcp.server.lowlevel.server import NotificationOptions @@ -3352,6 +3378,19 @@ if MCP_AVAILABLE: from litellm.proxy._types import LiteLLM_ObjectPermissionTable from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + # A key scoped to no MCP servers opts out of every MCP path. Enforce it + # here too, since toolset scoping replaces mcp_servers and would otherwise + # drop the sentinel. Checked before the admin branch, mirroring + # get_allowed_mcp_servers. + original_op = user_api_key_auth.object_permission + if original_op is not None and SpecialMCPServerNames.no_mcp_servers.value in ( + original_op.mcp_servers or [] + ): + raise HTTPException( + status_code=403, + detail="API key is scoped to no MCP servers; toolset access is denied.", + ) + # Access control: non-admin keys must have this toolset in their grant list. # Use _user_has_admin_view so that PROXY_ADMIN_VIEW_ONLY is also treated as admin. is_admin = _user_has_admin_view(user_api_key_auth) @@ -4006,6 +4045,12 @@ if MCP_AVAILABLE: except HTTPException: # Re-raise HTTP exceptions to preserve status codes and details raise + except ProxyException as e: + # Auth failures from user_api_key_auth arrive as ProxyException, not + # HTTPException. Preserve the real status (e.g. 401 + WWW-Authenticate) + # so OAuth clients can re-authenticate instead of receiving a generic + # 500 that surfaces as a cancelled tool call. + raise _proxy_exception_to_http_exception(e) except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") # Try to send a graceful error response for non-HTTP exceptions @@ -4123,6 +4168,12 @@ if MCP_AVAILABLE: # Re-raise HTTP exceptions to preserve status codes and details # (e.g. 401 + WWW-Authenticate challenges from OAuth pass-through). raise + except ProxyException as e: + # Auth failures from user_api_key_auth arrive as ProxyException, not + # HTTPException. Preserve the real status (e.g. 401 + WWW-Authenticate) + # so OAuth clients can re-authenticate instead of receiving a generic + # 500 that surfaces as a cancelled tool call. + raise _proxy_exception_to_http_exception(e) except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") # Try to send a graceful error response for non-HTTP exceptions diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8e2ec423cde..5bba842c7eb 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -461,6 +461,7 @@ class LiteLLMRoutes(enum.Enum): "/mcp/tools/call", "/mcp-rest/tools/list", "/mcp-rest/tools/call", + "/v1/mcp/tools", ] # MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS. @@ -2142,6 +2143,20 @@ class UserHeaderMapping(LiteLLMPydanticObjectBase): UserMCPManagementMode = Literal["restricted", "view_all"] +class PluginConfig(LiteLLMPydanticObjectBase): + """A single external service registered as an embeddable UI plugin.""" + + name: str = Field(description="unique plugin identifier (kebab-case)") + display_name: str | None = Field( + None, description="human-readable label shown in the UI view switcher" + ) + url: str = Field(description="base URL of the plugin service") + plugin_key: str | None = Field( + None, + description="plugin's own credential, injected as Bearer auth only on /plugin-proxy//* reverse-proxy calls", + ) + + class ConfigGeneralSettings(LiteLLMPydanticObjectBase): """ Documents all the fields supported by `general_settings` in config.yaml @@ -2150,6 +2165,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): completion_model: Optional[str] = Field( None, description="proxy level default model for all chat completion calls" ) + plugins: list[PluginConfig] | None = Field( + None, description="external services registered as embeddable UI plugins" + ) key_management_system: Optional[KeyManagementSystem] = Field( None, description="key manager to load keys from / decrypt keys with" ) @@ -2960,6 +2978,10 @@ class SpecialModelNames(enum.Enum): no_default_models = "no-default-models" +class SpecialMCPServerNames(enum.Enum): + no_mcp_servers = "no-mcp-servers" + + class SpecialProxyStrings(enum.Enum): default_user_id = "default_user_id" # global proxy admin @@ -3336,7 +3358,9 @@ class ProxyException(Exception): class CommonProxyErrors(str, enum.Enum): db_not_connected_error = ( - "DB not connected. See https://docs.litellm.ai/docs/proxy/virtual_keys" + "DB not connected. This endpoint needs a database; set DATABASE_URL to a " + "PostgreSQL connection string (postgresql://...) to enable it. " + "See https://docs.litellm.ai/docs/proxy/virtual_keys" ) no_llm_router = "No models configured on proxy" not_allowed_access = "Admin-only endpoint. Not allowed to access this." @@ -3808,6 +3832,28 @@ class SpecialHeaders(enum.Enum): mcp_servers = "x-mcp-servers" mcp_access_groups = "x-mcp-access-groups" + @classmethod + def litellm_credential_header_names(cls) -> "frozenset[str]": + """Lowercased header names user_api_key_auth accepts as a litellm key. + + Every header here authenticates the caller, so any code that forwards a + request onward (e.g. the plugin reverse proxy) must strip all of them to + avoid leaking the caller's litellm credential downstream. The static + custom-key header (general_settings.litellm_key_header_name) is runtime + config and must be added on top of this set by the caller. + """ + return frozenset( + header.value.lower() + for header in ( + cls.openai_authorization, + cls.azure_authorization, + cls.anthropic_authorization, + cls.google_ai_studio_authorization, + cls.azure_apim_authorization, + cls.custom_litellm_api_key, + ) + ) + class LitellmDataForBackendLLMCall(TypedDict, total=False): headers: dict diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 6ddf2cfeb20..88db2a2b7ea 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -699,6 +699,11 @@ async def common_checks( if valid_token is not None: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route( + request_data=request_body, + route=route, + ) + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( request_data=request_body, user_api_key_dict=valid_token, @@ -2949,6 +2954,26 @@ async def _get_agent_ids_from_access_groups( ) +def _resolve_all_team_model_sentinel_for_auth_check( + models: List[str], + llm_router: Optional[Router], + team_id: Optional[str], +) -> List[str]: + if ( + SpecialModelNames.all_team_models.value not in models + or team_id is None + or llm_router is None + ): + return models + proxy_models = llm_router.get_model_names() + non_sentinel_models = [ + model for model in models if model != SpecialModelNames.all_team_models.value + ] + if not proxy_models: + return non_sentinel_models or models + return list(dict.fromkeys(non_sentinel_models + proxy_models)) + + def _check_model_access_helper( model: str, llm_router: Optional[Router], @@ -2966,6 +2991,12 @@ def _check_model_access_helper( model_name=model, team_id=team_id ) + models = _resolve_all_team_model_sentinel_for_auth_check( + models=models, + llm_router=llm_router, + team_id=team_id, + ) + if ( len(access_groups) > 0 and llm_router is not None ): # check if token contains any model access groups @@ -3658,9 +3689,18 @@ async def _virtual_key_max_budget_check( # so a NaN max_budget would silently disable enforcement. Treat a # non-finite max_budget as "no configured limit" rather than as a bypass. if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: + # name the key in the error so operators don't have to reverse-map + # spend back to a key; key_name is the masked form (last 4 chars) + key_label = valid_token.key_alias or "key" + key_descriptor = ( + f"{key_label} ({valid_token.key_name})" + if valid_token.key_name + else key_label + ) raise litellm.BudgetExceededError( current_cost=spend, max_budget=valid_token.max_budget, + message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}", ) diff --git a/litellm/proxy/auth/auth_method.py b/litellm/proxy/auth/auth_method.py new file mode 100644 index 00000000000..a604eb563e6 --- /dev/null +++ b/litellm/proxy/auth/auth_method.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from enum import Enum + + +class AuthMethod(str, Enum): + API_KEY = "api_key" + HTTP_BASIC = "http_basic" + BEARER_JWT = "bearer_jwt" + OAUTH2_INTROSPECTION = "oauth2_introspection" + OIDC = "oidc" + SAML = "saml" + MUTUAL_TLS = "mutual_tls" diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 94b2ed84f20..3a2f2221ee3 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -285,6 +285,8 @@ _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = ( "s3_endpoint_url", "sagemaker_base_url", "deployment_url", + # SDK-only field; also rejected outright in is_request_body_safe. + "model_list", # Observability credentials, hosts, and project identifiers: derived # from the canonical ``_supported_callback_params`` allowlist so new # integrations are covered automatically. Sorted for stable iteration @@ -365,6 +367,10 @@ def is_request_body_safe( ``litellm_embedding_config.api_base`` (VERIA-6) without exposing a recursion-depth DoS surface. """ + if "model_list" in request_body: + raise ValueError( + "Rejected Request: model_list is not allowed in the request body." + ) _check_banned_params(request_body, general_settings, llm_router, model) for nested_key in _NESTED_CONFIG_KEYS: nested = _coerce_metadata_to_dict(request_body.get(nested_key)) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index b89db51c6f1..aa53954da8f 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -122,9 +122,16 @@ def get_key_models( SpecialModelNames.all_team_models.value in all_models and user_api_key_dict.team_id is not None ): - all_models = list( - user_api_key_dict.team_models - ) # copy to avoid mutating cached objects + all_models = list(user_api_key_dict.team_models) + if SpecialModelNames.all_team_models.value in all_models: + all_models = [ + model + for model in all_models + if model != SpecialModelNames.all_team_models.value + ] + all_models.extend(proxy_model_list) + if include_model_access_groups: + all_models.extend(model_access_groups.keys()) if SpecialModelNames.all_proxy_models.value in all_models: all_models = list(proxy_model_list) # copy to avoid mutating caller's list if include_model_access_groups: @@ -160,6 +167,12 @@ def get_team_models( all_models_set.update(team_models) if SpecialModelNames.all_team_models.value in all_models_set: all_models_set.update(team_models) + # GH#30619: expand all-team-models sentinel + # to the actual proxy model list + all_models_set.discard(SpecialModelNames.all_team_models.value) + all_models_set.update(proxy_model_list) + if include_model_access_groups: + all_models_set.update(model_access_groups.keys()) if SpecialModelNames.all_proxy_models.value in all_models_set: all_models_set.update(proxy_model_list) if include_model_access_groups: diff --git a/litellm/proxy/auth/network.py b/litellm/proxy/auth/network.py new file mode 100644 index 00000000000..4eb6f1dcec2 --- /dev/null +++ b/litellm/proxy/auth/network.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import ipaddress +from typing import Any, Union + +from fastapi import Request +from pydantic import BaseModel, Field + +from litellm._logging import verbose_proxy_logger + +TrustedProxyNetwork = Union[ipaddress.IPv4Network, ipaddress.IPv6Network] + + +class NetworkContext(BaseModel): + client_ip: str | None = None + host: str | None = None + via_trusted_proxy: bool = False + + +class TrustedProxyConfig(BaseModel): + use_forwarded_for: bool = False + trusted_proxy_cidrs: list[str] = Field(default_factory=list) + + +def normalize_cidr_ranges( + configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs" +) -> list[str]: + if not configured_ranges: + return [] + if isinstance(configured_ranges, str): + return [r.strip() for r in configured_ranges.split(",") if r.strip()] + if isinstance(configured_ranges, (list, tuple, set)): + return [str(r).strip() for r in configured_ranges if str(r).strip()] + verbose_proxy_logger.warning( + "Invalid %s value: expected a list of CIDR ranges, got %s", + setting_name, + type(configured_ranges).__name__, + ) + return [] + + +def parse_trusted_proxy_ranges( + configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs" +) -> list[TrustedProxyNetwork]: + networks: list[TrustedProxyNetwork] = [] + for cidr in normalize_cidr_ranges(configured_ranges, setting_name=setting_name): + try: + networks.append(ipaddress.ip_network(cidr, strict=False)) + except ValueError: + verbose_proxy_logger.warning( + "Invalid CIDR in %s: %s, skipping", setting_name, cidr + ) + return networks + + +def ip_in_networks(client_ip: str | None, networks: list[TrustedProxyNetwork]) -> bool: + if not client_ip or not networks: + return False + try: + addr = ipaddress.ip_address(client_ip.strip()) + except ValueError: + return False + return any(addr in network for network in networks) + + +def _is_valid_ip(value: str) -> bool: + try: + ipaddress.ip_address(value) + return True + except ValueError: + return False + + +def resolve_client_ip( + request: Request, config: TrustedProxyConfig +) -> tuple[str | None, bool]: + """Resolve the real client IP, trusting X-Forwarded-For only when the direct + peer is itself a configured trusted proxy. Walks the header right-to-left and + returns the first hop that is not a trusted proxy, so a forged left-most entry + cannot spoof the client.""" + peer = request.client.host if request.client else None + networks = parse_trusted_proxy_ranges(config.trusted_proxy_cidrs) + if not config.use_forwarded_for or not ip_in_networks(peer, networks): + return peer, False + forwarded = request.headers.get("x-forwarded-for", "") + hops = [h.strip() for h in forwarded.split(",") if h.strip()] + for hop in reversed(hops): + if _is_valid_ip(hop) and not ip_in_networks(hop, networks): + return hop, True + return peer, True + + +def resolve_network_context( + request: Request, config: TrustedProxyConfig +) -> NetworkContext: + ip, via_proxy = resolve_client_ip(request, config) + return NetworkContext( + client_ip=ip, + host=request.headers.get("host"), + via_trusted_proxy=via_proxy, + ) diff --git a/litellm/proxy/auth/resolvers/__init__.py b/litellm/proxy/auth/resolvers/__init__.py new file mode 100644 index 00000000000..d6bfb335c09 --- /dev/null +++ b/litellm/proxy/auth/resolvers/__init__.py @@ -0,0 +1,33 @@ +from litellm.proxy.auth.resolvers.exceptions import ( + IdentityResolutionError, + KeyNotFoundError, + KeyNotInCacheError, + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) +from litellm.proxy.auth.resolvers.models import ( + CredentialRef, + EndUserIdentity, + OrganizationIdentity, + Principal, + PrincipalType, + ProjectIdentity, + TeamIdentity, + UserIdentity, +) + +__all__ = [ + "CredentialRef", + "EndUserIdentity", + "IdentityResolutionError", + "KeyNotFoundError", + "KeyNotInCacheError", + "NoDatabaseConnectionError", + "OrganizationIdentity", + "Principal", + "PrincipalMissingSourceKeyError", + "PrincipalType", + "ProjectIdentity", + "TeamIdentity", + "UserIdentity", +] diff --git a/litellm/proxy/auth/resolvers/exceptions.py b/litellm/proxy/auth/resolvers/exceptions.py new file mode 100644 index 00000000000..dd953e66659 --- /dev/null +++ b/litellm/proxy/auth/resolvers/exceptions.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from fastapi import status + +from litellm.proxy._types import ProxyErrorTypes, ProxyException + + +class IdentityResolutionError(Exception): + """Base for every failure raised while resolving a caller's identity.""" + + +class NoDatabaseConnectionError(IdentityResolutionError): + def __init__(self) -> None: + super().__init__( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + + +class KeyNotInCacheError(IdentityResolutionError): + def __init__(self, hashed_token: str) -> None: + super().__init__( + f"Key doesn't exist in cache + check_cache_only=True. key={hashed_token}." + ) + + +class KeyNotFoundError(IdentityResolutionError, ProxyException): + """The token matched nothing in the cache or the verification token table. + + Also a ``ProxyException`` so the auth flow keeps mapping a missing key to the + OpenAI 401 contract unchanged while callers migrate onto the typed hierarchy. + """ + + def __init__(self, hashed_token: str) -> None: + ProxyException.__init__( + self, + message="Authentication Error, Invalid proxy server token passed. key={}, not found in db. Create key via `/key/generate` call.".format( + hashed_token + ), + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ) + + +class PrincipalMissingSourceKeyError(IdentityResolutionError): + def __init__(self) -> None: + super().__init__( + "Principal carries no source key; it was not produced by " + "IdentityStore.resolve" + ) diff --git a/litellm/proxy/auth/resolvers/models.py b/litellm/proxy/auth/resolvers/models.py new file mode 100644 index 00000000000..97e66b8fe67 --- /dev/null +++ b/litellm/proxy/auth/resolvers/models.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from enum import Enum + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.network import NetworkContext +from litellm.proxy.auth.roles import Role, TeamRole + + +class PrincipalType(str, Enum): + HUMAN = "human" + SERVICE_ACCOUNT = "service_account" + + +class UserIdentity(BaseModel): + id: str + external_id: str | None = None + user_name: str | None = None + email: str | None = None + display_name: str | None = None + + +class OrganizationIdentity(BaseModel): + id: str + name: str | None = None + + +class TeamIdentity(BaseModel): + id: str + name: str | None = None + role: TeamRole = TeamRole.MEMBER + + +class ProjectIdentity(BaseModel): + id: str + name: str | None = None + + +class EndUserIdentity(BaseModel): + id: str + + +class CredentialRef(BaseModel): + key_id: str | None = None + token_id: str | None = None + + +class Principal(BaseModel): + """Normalized caller identity, resolved once per request at the auth seam. + + Frozen and constructed fresh per request, never cached or shared. The identity + fields carry no policy, budget, or rate-limit state. ``source_key`` is a + transitional carrier for the resolved key object so ``key_from_principal`` can + hand it to the request flow that still consumes ``UserAPIKeyAuth``; it is + excluded from serialization and repr and goes away once those consumers read + identity off the Principal directly. + """ + + model_config = ConfigDict(frozen=True) + + principal_type: PrincipalType + subject: str + issuer: str | None = None + audience: list[str] = Field(default_factory=list) + + user: UserIdentity | None = None + organization: OrganizationIdentity | None = None + teams: list[TeamIdentity] = Field(default_factory=list) + project: ProjectIdentity | None = None + end_user: EndUserIdentity | None = None + + roles: list[Role] = Field(default_factory=list) + scopes: list[str] = Field(default_factory=list) + + auth_method: AuthMethod + credential_ref: CredentialRef = Field(default_factory=CredentialRef) + network: NetworkContext = Field(default_factory=NetworkContext) + + source_key: UserAPIKeyAuth | None = Field(default=None, exclude=True, repr=False) diff --git a/litellm/proxy/auth/resolvers/store.py b/litellm/proxy/auth/resolvers/store.py new file mode 100644 index 00000000000..c43d4c84ca4 --- /dev/null +++ b/litellm/proxy/auth/resolvers/store.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence + +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import ( + _cache_key_object, + _copy_user_api_key_auth_for_cache, + _fetch_key_object_from_db_with_reconnect, + get_object_permission, +) +from litellm.proxy.auth.resolvers.exceptions import ( + KeyNotFoundError, + KeyNotInCacheError, + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.network import NetworkContext +from litellm.proxy.auth.resolvers.models import ( + CredentialRef, + EndUserIdentity, + OrganizationIdentity, + Principal, + PrincipalType, + ProjectIdentity, + TeamIdentity, + UserIdentity, +) +from litellm.proxy.auth.roles import TeamRole, map_role, team_role + +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.integrations.opentelemetry import Span + from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class IdentityStore: + """The auth flow's resolver: one combined_view lookup, projected into a Principal. + + ``resolve`` does the lookup (cache, then DB via the shared lower-level helpers, + then write-back) and returns the per-caller Principal. The Principal carries the + source key object so ``key_from_principal`` can hand it back to the parts of the + request flow that still consume ``UserAPIKeyAuth`` (budget, rate limits, policy); + that carrier is a stopgap until those consumers read identity off the Principal. + The Prisma client, key cache, the request's tracing span / logging sink, and + whether this store may only read the cache are injected so the composition root + can build the store once the proxy DB is connected; the span and logging sink + are infra the DB call is instrumented with and ``check_cache_only`` is a store + mode, none of them inputs to resolving identity. ``auth_checks.get_key_object`` + stays as the legacy entrypoint for its other callers until they migrate onto + this store. + """ + + def __init__( + self, + prisma_client: PrismaClient | None, + cache: DualCache, + *, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + check_cache_only: bool = False, + ) -> None: + self._prisma = prisma_client + self._cache = cache + self._parent_otel_span = parent_otel_span + self._proxy_logging_obj = proxy_logging_obj + self._check_cache_only = check_cache_only + + async def resolve( + self, + hashed_token: str, + *, + auth_method: AuthMethod = AuthMethod.API_KEY, + network: NetworkContext | None = None, + ) -> Principal: + key = await self._resolve_key(hashed_token) + return self._principal_from_key( + key, + auth_method=auth_method, + network=network, + subject_fallback=key.token, + credential_ref=CredentialRef(token_id=key.token), + ) + + @staticmethod + def key_from_principal(principal: Principal) -> UserAPIKeyAuth: + """Hand back the resolved key object carried on the Principal. + + Stopgap for the request flow that still consumes ``UserAPIKeyAuth`` for + budget, rate-limit, and policy state. Only Principals produced by + ``resolve`` carry a source key. + """ + if principal.source_key is None: + raise PrincipalMissingSourceKeyError() + return principal.source_key + + async def _resolve_key(self, hashed_token: str) -> UserAPIKeyAuth: + if self._prisma is None: + raise NoDatabaseConnectionError() + + cached = await self._cache.async_get_cache( + key=hashed_token, model_type=UserAPIKeyAuth + ) + if cached is not None: + return _copy_user_api_key_auth_for_cache(user_api_key_obj=cached) + + if self._check_cache_only: + raise KeyNotInCacheError(hashed_token) + + from_db: BaseModel | None = await _fetch_key_object_from_db_with_reconnect( + hashed_token=hashed_token, + prisma_client=self._prisma, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + if from_db is None: + raise KeyNotFoundError(hashed_token) + + key = UserAPIKeyAuth(**from_db.model_dump(exclude_none=True)) + + if key.object_permission_id and not key.object_permission: + try: + key.object_permission = await get_object_permission( + object_permission_id=key.object_permission_id, + prisma_client=self._prisma, + user_api_key_cache=self._cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Failed to load object_permission for key with object_permission_id={key.object_permission_id}: {e}" + ) + + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=key, + user_api_key_cache=self._cache, + proxy_logging_obj=self._proxy_logging_obj, + ) + return key + + @staticmethod + def _principal_from_key( + key: UserAPIKeyAuth, + *, + auth_method: AuthMethod, + issuer: str | None = None, + subject_fallback: str | None = None, + scopes: Sequence[str] = (), + credential_ref: CredentialRef | None = None, + network: NetworkContext | None = None, + ) -> Principal: + """Project the identity slice off an already-resolved key object and carry + the key on the Principal so ``key_from_principal`` can recover it. + + Pure: issues no lookup. Both ``resolve`` and the auth seam call this so + identity is projected once off whichever key object they already hold. + """ + teams: list[TeamIdentity] = [] + if key.team_id is not None: + role = ( + team_role(key.team_member.role) if key.team_member else TeamRole.MEMBER + ) + teams.append(TeamIdentity(id=key.team_id, name=key.team_alias, role=role)) + organization = ( + OrganizationIdentity(id=key.org_id, name=key.organization_alias) + if key.org_id is not None + else None + ) + user = ( + UserIdentity(id=key.user_id, email=key.user_email) + if key.user_id is not None + else None + ) + project = ( + ProjectIdentity(id=key.project_id, name=key.project_alias) + if key.project_id is not None + else None + ) + end_user = ( + EndUserIdentity(id=key.end_user_id) if key.end_user_id is not None else None + ) + mapped = map_role(key.user_role) + return Principal( + principal_type=( + PrincipalType.HUMAN if key.user_id else PrincipalType.SERVICE_ACCOUNT + ), + subject=key.user_id or key.key_alias or subject_fallback or "", + issuer=issuer, + user=user, + organization=organization, + teams=teams, + project=project, + end_user=end_user, + roles=[mapped] if mapped else [], + scopes=list(scopes), + auth_method=auth_method, + credential_ref=credential_ref or CredentialRef(), + network=network or NetworkContext(), + source_key=key, + ) diff --git a/litellm/proxy/auth/roles.py b/litellm/proxy/auth/roles.py new file mode 100644 index 00000000000..efe56a8b6b2 --- /dev/null +++ b/litellm/proxy/auth/roles.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from enum import Enum + + +class Role(str, Enum): + PLATFORM_ADMIN = "platform_admin" + PLATFORM_VIEWER = "platform_viewer" + ORG_ADMIN = "org_admin" + ORG_VIEWER = "org_viewer" + TEAM_ADMIN = "team_admin" + TEAM_MEMBER = "team_member" + + +class TeamRole(str, Enum): + ADMIN = "admin" + MEMBER = "member" + + +_ROLE_MAP: dict[str, Role] = { + "proxy_admin": Role.PLATFORM_ADMIN, + "proxy_admin_viewer": Role.PLATFORM_VIEWER, + "org_admin": Role.ORG_ADMIN, +} + + +def map_role(value: str | None) -> Role | None: + """Map a LiteLLM ``user_role`` string to a platform Role.""" + if value is None: + return None + return _ROLE_MAP.get(value) + + +def team_role(role: str | None) -> TeamRole: + return TeamRole.ADMIN if role == "admin" else TeamRole.MEMBER diff --git a/litellm/proxy/auth/trusted_proxy_utils.py b/litellm/proxy/auth/trusted_proxy_utils.py index df7b3080f28..35bb79e7efe 100644 --- a/litellm/proxy/auth/trusted_proxy_utils.py +++ b/litellm/proxy/auth/trusted_proxy_utils.py @@ -1,12 +1,15 @@ -import ipaddress -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, Optional from fastapi import Request from litellm._logging import verbose_proxy_logger +from litellm.proxy.auth.network import ( + ip_in_networks, + normalize_cidr_ranges, + parse_trusted_proxy_ranges, +) TRUSTED_PROXY_RANGES_KEY = "trusted_proxy_ranges" -TrustedProxyNetwork = Union[ipaddress.IPv4Network, ipaddress.IPv6Network] def _get_proxy_general_settings() -> Dict[str, Any]: @@ -18,43 +21,20 @@ def _get_proxy_general_settings() -> Dict[str, Any]: return {} -def _normalize_cidr_ranges(configured_ranges: Any, *, setting_name: str) -> List[str]: - if not configured_ranges: - return [] - if isinstance(configured_ranges, str): - return [ - raw_range.strip() - for raw_range in configured_ranges.split(",") - if raw_range.strip() - ] - if isinstance(configured_ranges, (list, tuple, set)): - return [ - str(raw_range).strip() - for raw_range in configured_ranges - if str(raw_range).strip() - ] - verbose_proxy_logger.warning( - "Invalid %s value: expected a list of CIDR ranges, got %s", - setting_name, - type(configured_ranges).__name__, +def get_trusted_proxy_cidrs( + general_settings: dict[str, Any] | None = None, +) -> list[str]: + """Operator-configured trusted reverse-proxy CIDRs, normalized to strings. + + Empty when none are configured, in which case X-Forwarded-For must not be + trusted and only the direct peer is authoritative. + """ + if general_settings is None: + general_settings = _get_proxy_general_settings() + return normalize_cidr_ranges( + general_settings.get(TRUSTED_PROXY_RANGES_KEY), + setting_name=TRUSTED_PROXY_RANGES_KEY, ) - return [] - - -def parse_trusted_proxy_ranges( - configured_ranges: Any, - *, - setting_name: str = TRUSTED_PROXY_RANGES_KEY, -) -> List[TrustedProxyNetwork]: - networks: List[TrustedProxyNetwork] = [] - for cidr in _normalize_cidr_ranges(configured_ranges, setting_name=setting_name): - try: - networks.append(ipaddress.ip_network(cidr, strict=False)) - except ValueError: - verbose_proxy_logger.warning( - "Invalid CIDR in %s: %s, skipping", setting_name, cidr - ) - return networks def _get_direct_client_ip(request: Request) -> Optional[str]: @@ -65,18 +45,6 @@ def _get_direct_client_ip(request: Request) -> Optional[str]: return None -def _is_ip_in_networks( - client_ip: Optional[str], networks: List[TrustedProxyNetwork] -) -> bool: - if not client_ip or not networks: - return False - try: - addr = ipaddress.ip_address(client_ip.strip()) - except ValueError: - return False - return any(addr in network for network in networks) - - def require_trusted_proxy_request( *, request: Request, @@ -105,7 +73,7 @@ def require_trusted_proxy_request( ) direct_client_ip = _get_direct_client_ip(request) - if not _is_ip_in_networks(direct_client_ip, trusted_networks): + if not ip_in_networks(direct_client_ip, trusted_networks): verbose_proxy_logger.warning( "%s rejected identity headers from untrusted direct client IP %r", feature_name, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 00d98a04a78..e439f6a5998 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -42,7 +42,6 @@ from litellm.proxy.auth.auth_checks import ( common_checks, get_end_user_object, get_jwt_key_mapping_object, - get_key_object, get_project_object, get_team_object, get_user_object, @@ -63,7 +62,12 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context +from litellm.proxy.auth.resolvers import CredentialRef, Principal +from litellm.proxy.auth.resolvers.store import IdentityStore from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, @@ -758,12 +762,13 @@ async def _auto_register_jwt_mapping( claim_value, ) - auto_registered_key = await get_key_object( - hashed_token=token_hash, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + auto_registered_key = IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=token_hash) ) if auto_registered_key is not None: auto_registered_key.org_id = org_id @@ -865,12 +870,13 @@ async def _resolve_jwt_to_virtual_key( ) return None elif cached_mapping is not None: - return await get_key_object( - hashed_token=cached_mapping, - 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 IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=cached_mapping) ) # Resolve the mapping from DB, or treat prisma_client=None as a definitive @@ -889,12 +895,13 @@ async def _resolve_jwt_to_virtual_key( value=token_hash, ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, ) - return await get_key_object( - hashed_token=token_hash, - 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 IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=token_hash) ) # No mapping found (DB miss or no DB) — apply no-match policy. @@ -1493,13 +1500,14 @@ async def _user_api_key_auth_builder( ## Check CACHE try: with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"): - valid_token = await get_key_object( - hashed_token=hash_token(api_key), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=True, + valid_token = IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=True, + ).resolve(hashed_token=hash_token(api_key)) ) except Exception: verbose_logger.debug("api key not found in cache.") @@ -1679,12 +1687,13 @@ async def _user_api_key_auth_builder( try: with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): - valid_token = await get_key_object( - hashed_token=api_key, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + valid_token = IdentityStore.key_from_principal( + await IdentityStore( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ).resolve(hashed_token=api_key) ) except ProxyException as e: if e.code == 401 or e.code == "401": @@ -2387,6 +2396,17 @@ async def _run_centralized_common_checks( llm_router=llm_router, ) + # Pin the metadata variable name (litellm_metadata vs metadata) before + # any tag merge runs. Without this, header tags from + # apply_client_tag_policy_pre_auth would land in `metadata` while the + # later seed in common_checks pushes key tags and the + # _tag_max_budget_check read into `litellm_metadata`, hiding header + # tags from per-tag budget enforcement on LITELLM_METADATA_ROUTES. + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route( + request_data=request_data, + route=route, + ) + # Merge x-litellm-tags into request_data BEFORE common_checks runs. # _tag_max_budget_check inside common_checks only inspects request_data; # without this pre-merge, header-supplied tags bypass tag-budget @@ -2501,6 +2521,34 @@ def _should_skip_budget_checks( return False +def _resolve_request_principal( + request: Request, valid_token: UserAPIKeyAuth +) -> Principal: + """Project the resolved identity into one per-request Principal, off the key + object the builder already fetched, and stamp the request network context + onto it once. X-Forwarded-For is only trusted when the operator configured + ``trusted_proxy_ranges``; otherwise the direct peer is authoritative. + + credential_ref and a stable subject fallback are always set off the token so + the Principal can never be anonymous, even for a keyless service-account key + with no user or alias.""" + cidrs = get_trusted_proxy_cidrs() + network = resolve_network_context( + request, + TrustedProxyConfig(use_forwarded_for=bool(cidrs), trusted_proxy_cidrs=cidrs), + ) + auth_method = ( + AuthMethod.BEARER_JWT if valid_token.jwt_claims else AuthMethod.API_KEY + ) + return IdentityStore._principal_from_key( + valid_token, + auth_method=auth_method, + network=network, + subject_fallback=valid_token.token, + credential_ref=CredentialRef(token_id=valid_token.token), + ) + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -2615,6 +2663,22 @@ async def user_api_key_auth( model=request_data.get("model") if isinstance(request_data, dict) else None, ) user_api_key_auth_obj.request_route = normalize_request_route(route) + + # Resolve caller identity once, here at the seam, into a single per-request + # Principal projected off the key object the builder already fetched (no + # second lookup). Downstream consumers read identity off this instead of + # re-resolving it. Additive and defensive: a projection failure must never + # reject an already-authenticated request, so it is left unset on failure; + # any future consumer must treat a missing principal as deny, not allow. + try: + request.state.principal = _resolve_request_principal( + request, user_api_key_auth_obj + ) + except Exception as e: + verbose_proxy_logger.warning( + "Principal projection at auth seam failed (non-fatal): %s", e + ) + return user_api_key_auth_obj diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8ef931e8d25..8dec08460b4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1037,6 +1037,8 @@ class ProxyBaseLLMRequestProcessing: version=version, proxy_config=proxy_config, ) + if not general_settings.get("expose_fallback_errors_to_caller"): + self.data.pop("include_fallback_errors", None) if route_type in {"aresponses", "_aresponses_websocket"}: await _authorize_response_file_search_vector_stores( data=self.data, diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 71dce163b78..d48499af6f0 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -71,6 +71,23 @@ def initialize_callbacks_on_proxy( imported_list.append(compression_interception_obj) continue + if ( + isinstance(callback, str) + and callback == "code_interpreter_interception" + ): + from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, + ) + + code_interpreter_interception_obj = ( + CodeInterpreterInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, + ) + ) + imported_list.append(code_interpreter_interception_obj) + continue + # check if callback is a custom logger compatible callback if isinstance(callback, str): callback = LoggingCallbackManager._add_custom_callback_generic_api_str( diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 58478db5e2e..ae2307658dd 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -32,7 +32,7 @@ password when their ``*_READ_REPLICA`` counterpart is unset. import os import urllib.parse -from typing import Optional, cast +from typing import Final, cast from pydantic import AliasChoices, Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -44,6 +44,41 @@ from litellm.proxy.auth import rds_iam_token _IAM_ENV_KEY = "IAM_TOKEN_DB_AUTH" _DEFAULT_PG_PORT = "5432" +# schema.prisma pins `provider = "postgresql"`, so these are the only schemes +# Prisma can actually connect with. +SUPPORTED_DB_SCHEMES: Final[frozenset[str]] = frozenset({"postgresql", "postgres"}) +_MISSING_SCHEME = "" + + +def unsupported_db_scheme(database_url: str) -> str | None: + """Return the connection URL scheme when it is not PostgreSQL, else None. + + A `sqlite://` / `mysql://` URL can never connect against the + postgresql-only datasource, but the resulting Prisma failure is opaque and + version-dependent (a confusing migration error, or a startup that never + binds). Callers use this to reject the URL up front with an actionable + error instead. + + A schemeless value (e.g. a malformed DSN like ``user:pass@host/db``) yields + the ``_MISSING_SCHEME`` placeholder rather than the raw URL, so callers that + log the return value never echo embedded credentials. + """ + scheme = urllib.parse.urlsplit(database_url).scheme.lower() + if scheme in SUPPORTED_DB_SCHEMES: + return None + return scheme or _MISSING_SCHEME + + +def unsupported_db_scheme_message(env_var: str, scheme: str) -> str: + """Operator-facing message naming the offending env var and scheme.""" + return ( + f"{env_var} uses unsupported scheme '{scheme}'. LiteLLM's database " + "features (virtual keys, store_model_in_db, spend tracking) require " + "PostgreSQL; use a 'postgresql://' connection string. SQLite and other " + "engines are not supported. " + "See https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + class DatabaseURLSettings(BaseSettings): """Discrete ``DATABASE_*`` env vars, loaded once at process start. @@ -58,46 +93,47 @@ class DatabaseURLSettings(BaseSettings): iam_token_db_auth: bool = Field(default=False, validation_alias=_IAM_ENV_KEY) # Writer - database_url: Optional[str] = Field(default=None, validation_alias="DATABASE_URL") - database_host: Optional[str] = Field(default=None, validation_alias="DATABASE_HOST") + database_url: str | None = Field(default=None, validation_alias="DATABASE_URL") + direct_url: str | None = Field(default=None, validation_alias="DIRECT_URL") + database_host: str | None = Field(default=None, validation_alias="DATABASE_HOST") database_port: str = Field( default=_DEFAULT_PG_PORT, validation_alias="DATABASE_PORT" ) - database_user: Optional[str] = Field( + database_user: str | None = Field( default=None, validation_alias=AliasChoices("DATABASE_USER", "DATABASE_USERNAME"), ) - database_name: Optional[str] = Field(default=None, validation_alias="DATABASE_NAME") - database_schema: Optional[str] = Field( + database_name: str | None = Field(default=None, validation_alias="DATABASE_NAME") + database_schema: str | None = Field( default=None, validation_alias="DATABASE_SCHEMA" ) - database_password: Optional[str] = Field( + database_password: str | None = Field( default=None, validation_alias="DATABASE_PASSWORD" ) # Read replica - database_url_read_replica: Optional[str] = Field( + database_url_read_replica: str | None = Field( default=None, validation_alias="DATABASE_URL_READ_REPLICA" ) - database_host_read_replica: Optional[str] = Field( + database_host_read_replica: str | None = Field( default=None, validation_alias="DATABASE_HOST_READ_REPLICA" ) - database_port_read_replica: Optional[str] = Field( + database_port_read_replica: str | None = Field( default=None, validation_alias="DATABASE_PORT_READ_REPLICA" ) - database_user_read_replica: Optional[str] = Field( + database_user_read_replica: str | None = Field( default=None, validation_alias=AliasChoices( "DATABASE_USER_READ_REPLICA", "DATABASE_USERNAME_READ_REPLICA" ), ) - database_name_read_replica: Optional[str] = Field( + database_name_read_replica: str | None = Field( default=None, validation_alias="DATABASE_NAME_READ_REPLICA" ) - database_schema_read_replica: Optional[str] = Field( + database_schema_read_replica: str | None = Field( default=None, validation_alias="DATABASE_SCHEMA_READ_REPLICA" ) - database_password_read_replica: Optional[str] = Field( + database_password_read_replica: str | None = Field( default=None, validation_alias="DATABASE_PASSWORD_READ_REPLICA" ) @@ -106,7 +142,7 @@ class DatabaseURLSettings(BaseSettings): """Load the settings from ``os.environ`` (read at call time).""" return cls() - def build_writer_url(self) -> Optional[str]: + def build_writer_url(self) -> str | None: """Return the writer URL to set, or ``None`` to leave it as-is. Raises ``RuntimeError`` (naming the offending vars) when IAM auth is @@ -156,7 +192,7 @@ class DatabaseURLSettings(BaseSettings): ) return None - def build_reader_url(self) -> Optional[str]: + def build_reader_url(self) -> str | None: """Return the read-replica URL to set, or ``None`` to leave it as-is. Opt-in via ``DATABASE_HOST_READ_REPLICA``; never clobbers a @@ -217,11 +253,11 @@ class DatabaseURLSettings(BaseSettings): def _password_url( *, user: str, - password: Optional[str], + password: str | None, host: str, port: str, name: str, - schema: Optional[str], + schema: str | None, ) -> str: """Percent-encode credentials into a ``postgresql://`` URL. @@ -239,6 +275,26 @@ class DatabaseURLSettings(BaseSettings): url += f"?schema={schema}" return url + def _raise_for_unsupported_scheme(self) -> None: + """Reject an operator-pinned non-PostgreSQL writer / direct / reader URL. + + The componentized entrypoints (gateway / backend / migrations) call + ``apply_to_env`` and then hand the URL straight to Prisma, bypassing + the CLI's own guard. A pinned URL flows through untouched, so validate + the same three vars the CLI guard checks (DATABASE_URL, DIRECT_URL, and + the read replica) rather than letting Prisma stall on an unusable scheme. + """ + for env_var, url in ( + ("DATABASE_URL", self.database_url), + ("DIRECT_URL", self.direct_url), + ("DATABASE_URL_READ_REPLICA", self.database_url_read_replica), + ): + if not url: + continue + bad_scheme = unsupported_db_scheme(url) + if bad_scheme is not None: + raise RuntimeError(unsupported_db_scheme_message(env_var, bad_scheme)) + def apply_to_env(self) -> bool: """Write the assembled URL(s) into ``os.environ``. @@ -246,6 +302,7 @@ class DatabaseURLSettings(BaseSettings): password auth that assembled a fresh URL). False means there was nothing to do — an operator-pinned URL, or no discrete fields. """ + self._raise_for_unsupported_scheme() wrote_writer = False writer_url = self.build_writer_url() if writer_url is not None: diff --git a/litellm/proxy/example_config_yaml/code_interpreter_interception_config.yaml b/litellm/proxy/example_config_yaml/code_interpreter_interception_config.yaml new file mode 100644 index 00000000000..9c85f5c5140 --- /dev/null +++ b/litellm/proxy/example_config_yaml/code_interpreter_interception_config.yaml @@ -0,0 +1,17 @@ +model_list: + - model_name: gpt-5 + litellm_params: + model: openai/gpt-5 + +# Sandbox tools configuration +sandbox_tools: + - sandbox_tool_name: "my-e2b" + litellm_params: + sandbox_provider: "e2b" + api_key: os.environ/E2B_API_KEY + +litellm_settings: + callbacks: ["code_interpreter_interception"] + code_interpreter_interception_params: + enabled_providers: ["openai"] + sandbox_tool_name: "my-e2b" diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 6343faaa965..9888baf897e 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -123,11 +123,70 @@ class SemanticToolFilterHook(CustomLogger): return openai_tools_as_dicts + def _is_mcp_tool(self, tool: object) -> bool: + """ + Check whether *tool* is registered in the MCP semantic router. + + Classification strategy (shape-first, lookup-second): + 1. Chat Completions format dicts are always native. + 2. Responses API function tools are always native. + 3. Everything else is looked up by name in the MCP registry. + """ + if ( + isinstance(tool, dict) + and tool.get("type") == "function" + and isinstance(tool.get("function"), dict) + ): + return False + if ( + isinstance(tool, dict) + and tool.get("type") == "function" + and isinstance(tool.get("name"), str) + ): + return False + name, _ = self.filter._extract_tool_info(tool) + return bool(name) and name in self.filter._tool_map + def _get_metadata_variable_name(self, data: dict) -> str: if "litellm_metadata" in data: return "litellm_metadata" return "metadata" + def _emit_filter_metadata( + self, + data: dict, + mcp_tools: list[object], + filtered_mcp_tools: list[object], + native_tools: list[object], + filtered_tools: list[object], + ) -> None: + """ + Emit response-header metadata when MCP tools were filtered. + + Stats report MCP-only counts so downstream consumers see accurate + semantic filter metrics. Skips metadata entirely for purely-native + requests to avoid spurious headers. + """ + if mcp_tools: + filter_stats = f"{len(mcp_tools)}->{len(filtered_mcp_tools)}" + tool_names_csv = self._get_tool_names_csv(filtered_mcp_tools) + + _metadata_variable_name = self._get_metadata_variable_name(data) + metadata = data.setdefault(_metadata_variable_name, {}) + metadata["litellm_semantic_filter_stats"] = filter_stats + metadata["litellm_semantic_filter_tools"] = tool_names_csv + + verbose_proxy_logger.info( + f"Semantic tool filter: {filter_stats} MCP tools " + f"({len(native_tools)} native preserved, " + f"{len(filtered_tools)} total)" + ) + else: + verbose_proxy_logger.info( + f"Semantic tool filter: all {len(native_tools)} tools " + f"are native, no MCP filtering applied" + ) + async def async_pre_call_hook( self, user_api_key_dict: "UserAPIKeyAuth", @@ -140,53 +199,55 @@ class SemanticToolFilterHook(CustomLogger): This hook is called before the LLM request is made. It filters the tools list to only include semantically relevant tools. - - Args: - user_api_key_dict: User authentication - cache: Cache instance - data: Request data containing messages and tools - call_type: Type of call (completion, acompletion, etc.) - - Returns: - Modified data dict with filtered tools, or None if no changes """ - # Only filter endpoints that support tools if call_type not in ("completion", "acompletion", "aresponses"): verbose_proxy_logger.debug( f"Skipping semantic filter for call_type={call_type}" ) return None - # Check if tools are present tools = data.get("tools") if not tools: verbose_proxy_logger.debug("No tools in request, skipping semantic filter") return None - original_tool_count = len(tools) - - # Check for MCP references (server_url="litellm_proxy") and expand them + # Expanded MCP tools are in OpenAI nested format which + # filter_tools/_extract_tool_info cannot name-match, so we skip + # semantic filtering and return early. if self._should_expand_mcp_tools(tools): verbose_proxy_logger.debug( "Detected litellm_proxy MCP references, expanding before semantic filtering" ) try: + native_tools_before_expand = [ + t + for t in tools + if not (isinstance(t, dict) and t.get("type") == "mcp") + ] + expanded_tools = await self._expand_mcp_tools(tools, user_api_key_dict) if not expanded_tools: + if native_tools_before_expand: + data["tools"] = native_tools_before_expand + verbose_proxy_logger.warning( + "No MCP tools expanded, preserving " + f"{len(native_tools_before_expand)} native tools" + ) + return data verbose_proxy_logger.warning( "No tools expanded from MCP references" ) return None + data["tools"] = native_tools_before_expand + expanded_tools verbose_proxy_logger.info( - f"Expanded {len(tools)} MCP reference(s) to {len(expanded_tools)} tools" + f"Expanded MCP references to {len(expanded_tools)} tools " + f"({len(native_tools_before_expand)} native preserved), " + f"skipping semantic filter (OpenAI nested format)" ) - - # Update tools for filtering - tools = expanded_tools - original_tool_count = len(tools) + return data except Exception as e: verbose_proxy_logger.error( @@ -194,7 +255,6 @@ class SemanticToolFilterHook(CustomLogger): ) return None - # Check if messages are present (try both "messages" and "input" for responses API) messages = data.get("messages", []) if not messages: messages = data.get("input", []) @@ -204,13 +264,11 @@ class SemanticToolFilterHook(CustomLogger): ) return None - # Check if filter is enabled if not self.filter.enabled: verbose_proxy_logger.debug("Semantic filter disabled, skipping") return None try: - # Extract user query from messages user_query = self.filter.extract_user_query(messages) if not user_query: verbose_proxy_logger.debug( @@ -218,33 +276,60 @@ class SemanticToolFilterHook(CustomLogger): ) return None + native_tools: list[object] = [] + mcp_tools: list[object] = [] + mcp_indices: set[int] = set() + for i, t in enumerate(tools): + if self._is_mcp_tool(t): + mcp_tools.append(t) + mcp_indices.add(i) + else: + native_tools.append(t) + verbose_proxy_logger.debug( - f"Applying semantic filter to {len(tools)} tools " - f"with query: '{user_query[:50]}...'" + f"Applying semantic filter: {len(mcp_tools)} MCP tools, " + f"{len(native_tools)} native tools, " + f"query: '{user_query[:50]}...'" ) - # Filter tools semantically - filtered_tools = await self.filter.filter_tools( - query=user_query, - available_tools=tools, # type: ignore - ) + if mcp_tools: + filtered_mcp_tools = await self.filter.filter_tools( + query=user_query, + available_tools=mcp_tools, # type: ignore + ) + else: + filtered_mcp_tools = [] + + filtered_mcp_names: set[str] = set() + for t in filtered_mcp_tools: + name, _ = self.filter._extract_tool_info(t) + if name: + filtered_mcp_names.add(name) + + filtered_tools: list[object] = [] + for i, t in enumerate(tools): + if i in mcp_indices: + name, _ = self.filter._extract_tool_info(t) + if name in filtered_mcp_names: + filtered_tools.append(t) + else: + filtered_tools.append(t) - # Always update tools and emit header (even if count unchanged) data["tools"] = filtered_tools - # Store filter stats and tool names for response header - filter_stats = f"{original_tool_count}->{len(filtered_tools)}" - tool_names_csv = self._get_tool_names_csv(filtered_tools) - - _metadata_variable_name = self._get_metadata_variable_name(data) - data[_metadata_variable_name][ - "litellm_semantic_filter_stats" - ] = filter_stats - data[_metadata_variable_name][ - "litellm_semantic_filter_tools" - ] = tool_names_csv - - verbose_proxy_logger.info(f"Semantic tool filter: {filter_stats} tools") + try: + self._emit_filter_metadata( + data=data, + mcp_tools=mcp_tools, + filtered_mcp_tools=filtered_mcp_tools, + native_tools=native_tools, + filtered_tools=filtered_tools, + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to emit semantic filter metadata: {e}", + exc_info=True, + ) return data @@ -266,7 +351,7 @@ class SemanticToolFilterHook(CustomLogger): from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH _metadata_variable_name = self._get_metadata_variable_name(data) - metadata = data[_metadata_variable_name] + metadata = data.get(_metadata_variable_name, {}) filter_stats = metadata.get("litellm_semantic_filter_stats") if not filter_stats: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 2c5937d9506..c0cdf84dfb6 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -108,6 +108,7 @@ def parse_cache_control(cache_control): LITELLM_METADATA_ROUTES = ( "batches", + "bedrock", "/v1/messages", "responses", "files", @@ -141,6 +142,20 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS = ( "service_callback", "logger_fn", "litellm_disabled_callbacks", + # Agentic-loop control fields. These bound or drive an interceptor's agentic + # loop (web search, compression, code interpreter) and are server-controlled. + # A client-supplied value would forge loop depth/cycle state, mark an + # interception as active (triggering sandbox code execution without the + # native tool ever being present), force the completed response to be + # re-wrapped as a synthetic stream the caller never asked for, or raise the + # loop ceiling to drive many upstream model calls and sandbox executions + # from a single request. + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_code_interpreter_interception_active", + "_code_interpreter_interception_converted_stream", + "_code_interpreter_interception_sandbox_key", + "max_agentic_loops", ) _UNTRUSTED_METADATA_CONTROL_FIELDS = ( @@ -1223,6 +1238,27 @@ class LiteLLMProxyRequestSetup: return tags + @staticmethod + def pre_seed_litellm_metadata_for_route( + request_data: dict, + route: str, + ) -> None: + """Pre-seed ``litellm_metadata`` for routes that track tags there. + + Routes in ``LITELLM_METADATA_ROUTES`` (e.g. Bedrock, ``/v1/messages``, + responses, batches, files) store request-scoped tag metadata in + ``litellm_metadata`` rather than the provider-facing ``metadata`` + field. ``get_metadata_variable_name_from_kwargs`` picks the target + based on whether ``litellm_metadata`` is present, so it must be + seeded BEFORE any tag merge runs; otherwise header tags from + ``apply_client_tag_policy_pre_auth`` land in ``metadata`` while + key tags from ``apply_key_tags_pre_auth`` and the read in + ``_tag_max_budget_check`` resolve to ``litellm_metadata``, leaving + header tags invisible to per-tag budget enforcement. + """ + if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES): + request_data.setdefault("litellm_metadata", {}) + @staticmethod def apply_key_tags_pre_auth( request_data: dict, @@ -1454,8 +1490,7 @@ async def add_litellm_data_to_request( _metadata_variable_name=_metadata_variable_name, ) - # Add headers to metadata for guardrails to access (fixes #17477) - # Guardrails use metadata["headers"] to access request headers (e.g., User-Agent) + # Expose request headers under the metadata field for guardrails (fixes #17477) if _metadata_variable_name in data and isinstance( data[_metadata_variable_name], dict ): diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 341a8767db0..79882909c23 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -1,7 +1,7 @@ import asyncio from datetime import datetime from types import SimpleNamespace -from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union +from typing import Any, Awaitable, Callable, Dict, List, Optional, Set, Tuple, Union from fastapi import HTTPException, status @@ -887,8 +887,17 @@ async def get_daily_activity( exclude_entity_ids: Optional[List[str]] = None, metadata_metrics_func: Optional[Callable[[List[Any]], SpendMetrics]] = None, timezone_offset_minutes: Optional[int] = None, + resolve_entity_metadata: Optional[ + Callable[[list[Any]], Awaitable[dict[str, dict]]] + ] = None, ) -> SpendAnalyticsPaginatedResponse: - """Common function to get daily activity for any entity type.""" + """Common function to get daily activity for any entity type. + + ``resolve_entity_metadata`` lets a caller resolve entity metadata from the + rows actually on the page (e.g. user_id -> user_email) instead of fetching + the whole entity table upfront, which matters when the entity set is + unbounded. + """ if prisma_client is None: raise HTTPException( @@ -939,11 +948,18 @@ async def get_daily_activity( take=page_size, ) + resolved_entity_metadata = entity_metadata_field + if resolve_entity_metadata is not None: + resolved_entity_metadata = { + **(entity_metadata_field or {}), + **(await resolve_entity_metadata(daily_spend_data)), + } + aggregated = await _aggregate_spend_records( prisma_client=prisma_client, records=daily_spend_data, entity_id_field=entity_id_field, - entity_metadata_field=entity_metadata_field, + entity_metadata_field=resolved_entity_metadata, ) metadata_metrics = aggregated["totals"] diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ba7013570fe..6d7f565fb85 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -57,6 +57,9 @@ from litellm.repositories.verification_token_repository import ( from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) +from litellm.types.proxy.management_endpoints.scim_v2 import ( + SCIM_ENTERPRISE_METADATA_KEY, +) from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkUpdateUserRequest, BulkUpdateUserResponse, @@ -719,6 +722,17 @@ async def _get_user_info_teams( return team_list, teams_1 +def _redact_scim_enterprise_metadata( + metadata: Optional[Dict[str, Any]], +) -> Optional[Dict[str, Any]]: + """SCIM enterprise attributes are persisted in user metadata so reporting can + group on them, but they are directory-only fields that generic user-info + endpoints must not surface; SCIM clients read them through the SCIM endpoints.""" + if not isinstance(metadata, dict) or SCIM_ENTERPRISE_METADATA_KEY not in metadata: + return metadata + return {k: v for k, v in metadata.items() if k != SCIM_ENTERPRISE_METADATA_KEY} + + def _build_user_info_response( user_id: Optional[str], user_info: Optional[Any], @@ -739,6 +753,9 @@ def _build_user_info_response( ) if isinstance(_user_info, dict): _user_info.pop("password", None) + _user_info["metadata"] = _redact_scim_enterprise_metadata( + _user_info.get("metadata") + ) return UserInfoResponse( user_id=user_id, @@ -983,7 +1000,7 @@ async def user_info_v2( models=user_data.get("models") or [], budget_duration=user_data.get("budget_duration"), budget_reset_at=user_data.get("budget_reset_at"), - metadata=user_data.get("metadata"), + metadata=_redact_scim_enterprise_metadata(user_data.get("metadata")), created_at=user_data.get("created_at"), updated_at=user_data.get("updated_at"), sso_user_id=user_data.get("sso_user_id"), @@ -2098,9 +2115,13 @@ async def get_users( user_list: List[LiteLLM_UserTableWithKeyCount] = [] if users is not None: for user in users: + user_dump = user.model_dump() + user_dump["metadata"] = _redact_scim_enterprise_metadata( + user_dump.get("metadata") + ) user_list.append( LiteLLM_UserTableWithKeyCount( - **user.model_dump(), key_count=user_key_counts.get(user.user_id, 0) + **user_dump, key_count=user_key_counts.get(user.user_id, 0) ) ) else: @@ -2596,6 +2617,25 @@ async def ui_view_users( # Using shared metric helper implementations from common_daily_activity +async def _resolve_user_email_metadata( + prisma_client: "PrismaClient", records: list[Any] +) -> dict[str, dict]: + """Map each user_id on the page to its email/alias so the Usage dashboard can + label the 'Spend Per User' chart with the email instead of the raw UUID.""" + user_ids = { + record.user_id for record in records if getattr(record, "user_id", None) + } + if not user_ids: + return {} + users = await UserRepository(prisma_client).table.find_many( + where={"user_id": {"in": list(user_ids)}} + ) + return { + user.user_id: {"user_email": user.user_email, "user_alias": user.user_alias} + for user in users + } + + @router.get( "/user/daily/activity", tags=["Budget & Spend Tracking", "Internal User management"], @@ -2698,6 +2738,9 @@ async def get_user_daily_activity( page=page, page_size=page_size, timezone_offset_minutes=timezone, + resolve_entity_metadata=lambda records: _resolve_user_email_metadata( + prisma_client, records + ), ) except HTTPException: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e86982307e7..fb284e707a7 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1016,15 +1016,10 @@ if MCP_AVAILABLE: if is_restricted_virtual_key: return _sanitize_mcp_server_list_for_virtual_key(redacted_mcp_servers) - # Non-admin authenticated users may see the server inventory but - # not credential-bearing fields like `url` (often contains bearer - # tokens) or headers/env (often contain Authorization). - if not _user_has_admin_view(user_api_key_dict): - return _sanitize_mcp_server_list_for_non_admin(redacted_mcp_servers) - + # only a full PROXY_ADMIN sees credential-bearing fields; everyone else + # goes through the non-admin sanitizer if not _user_is_full_admin(user_api_key_dict): - for server in redacted_mcp_servers: - _redact_global_env_var_values(server) + return _sanitize_mcp_server_list_for_non_admin(redacted_mcp_servers) return redacted_mcp_servers @@ -1415,10 +1410,10 @@ if MCP_AVAILABLE: redacted = _redact_mcp_credentials(mcp_server) if is_restricted_virtual_key: return _sanitize_mcp_server_for_virtual_key(redacted) - if not _user_has_admin_view(user_api_key_dict): - return _sanitize_mcp_server_for_non_admin(redacted) + # only a full PROXY_ADMIN sees credential-bearing fields; everyone else + # goes through the non-admin sanitizer if not _user_is_full_admin(user_api_key_dict): - _redact_global_env_var_values(redacted) + return _sanitize_mcp_server_for_non_admin(redacted) return redacted @router.post( diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py index d1e00f87b69..866f5baf3a4 100644 --- a/litellm/proxy/management_endpoints/scim/scim_transformations.py +++ b/litellm/proxy/management_endpoints/scim/scim_transformations.py @@ -50,8 +50,16 @@ class ScimTransformations: scim_active = metadata.get("scim_active") active = True if scim_active is None else bool(scim_active) + schemas = ["urn:ietf:params:scim:schemas:core:2.0:User"] + enterprise_user = None + if metadata.get(SCIM_ENTERPRISE_METADATA_KEY): + enterprise_user = SCIMEnterpriseUser.model_validate( + metadata[SCIM_ENTERPRISE_METADATA_KEY] + ) + schemas.append(SCIM_ENTERPRISE_USER_SCHEMA) + return SCIMUser( - schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + schemas=schemas, id=user.user_id, userName=ScimTransformations._get_scim_user_name(user), displayName=ScimTransformations._get_scim_user_name(user), @@ -62,6 +70,7 @@ class ScimTransformations: emails=emails, groups=groups, active=active, + enterprise_user=enterprise_user, meta={ "resourceType": "User", "created": user_created_at, diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 0798d1a510d..d5da0372a8f 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -5,7 +5,7 @@ This is an enterprise feature and requires a premium license. """ import re -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, Dict, Iterable, List, Optional, Set, Tuple from fastapi import ( APIRouter, @@ -69,14 +69,21 @@ class UserProvisionerHelpers: @staticmethod async def handle_existing_user_by_email( - prisma_client, new_user_request: NewUserRequest + prisma_client, + new_user_request: NewUserRequest, + admin_group: Optional[str] = None, ) -> Optional[SCIMUser]: """ Check if a user with the given email already exists and update them if found. + When admin_group is configured the resolved global role on new_user_request + is persisted too, so re-upserting an existing email demotes a user who is no + longer in the admin group instead of leaving the stale role. + Args: prisma_client: Database client new_user_request: New user request data + admin_group: Configured SCIM admin group, or None to leave role untouched Returns: SCIMUser if user was updated, None if no existing user found @@ -100,6 +107,11 @@ class UserProvisionerHelpers: "user_alias": new_user_request.user_alias, "teams": new_user_request.teams, "metadata": safe_dumps(new_user_request.metadata), + **( + {"user_role": new_user_request.user_role} + if admin_group is not None + else {} + ), }, ) @@ -118,6 +130,7 @@ class ScimUserData(TypedDict): given_name: Optional[str] family_name: Optional[str] active: Optional[bool] + enterprise: Optional[SCIMEnterpriseUser] class GroupMemberExtractionResult(BaseModel): @@ -199,11 +212,15 @@ def _extract_scim_user_data(user: SCIMUser) -> ScimUserData: "given_name": user.name.givenName if user.name else None, "family_name": user.name.familyName if user.name else None, "active": user.active, + "enterprise": user.enterprise_user, } def _build_scim_metadata( - given_name: Optional[str], family_name: Optional[str], active: Optional[bool] = None + given_name: Optional[str], + family_name: Optional[str], + active: Optional[bool] = None, + enterprise: Optional[SCIMEnterpriseUser] = None, ) -> Dict[str, Any]: """Build metadata dictionary with SCIM data.""" metadata: Dict[str, Any] = { @@ -216,6 +233,11 @@ def _build_scim_metadata( if active is not None: metadata["scim_active"] = active + if enterprise is not None: + metadata[SCIM_ENTERPRISE_METADATA_KEY] = enterprise.model_dump( + by_alias=True, exclude_none=True + ) + return metadata @@ -244,6 +266,117 @@ async def _get_scim_upsert_user_setting() -> bool: return True +ScimUserRole = Literal[ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, +] + + +def _default_scim_user_role() -> ScimUserRole: + """Non-admin default role for SCIM-provisioned users.""" + if litellm.default_internal_user_params: + configured_role = litellm.default_internal_user_params.get("user_role") + if configured_role is not None: + return configured_role + return LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + + +async def _get_scim_admin_group() -> Optional[str]: + """ + Get the scim_admin_group setting from litellm_settings. + + Returns the configured admin group identifier, or None when unset so callers + leave a user's global role untouched (default-safe). + """ + try: + from litellm.proxy.proxy_server import proxy_config + + config = await proxy_config.get_config() + litellm_settings = config.get("litellm_settings", {}) or {} + return litellm_settings.get("scim_admin_group") or None + except Exception as e: + verbose_proxy_logger.warning( + f"Error reading scim_admin_group setting, defaulting to None: {e}" + ) + return None + + +def _resolve_scim_user_role( + groups: list[SCIMUserGroup], + admin_group: Optional[str], + default_role: ScimUserRole, +) -> Optional[LitellmUserRoles]: + """ + Resolve a user's global proxy role from their SCIM groups. + + Returns None when no admin group is configured, signalling callers to leave + the role unchanged. Otherwise grants PROXY_ADMIN when any group matches the + admin group by value or display, and falls back to the non-admin default. + """ + if admin_group is None: + return None + for group in groups: + if group.value == admin_group or group.display == admin_group: + return LitellmUserRoles.PROXY_ADMIN + return default_role + + +async def _scim_groups_from_team_ids( + prisma_client: Any, team_ids: list[str] +) -> list[SCIMUserGroup]: + """ + Build SCIMUserGroup objects from team ids, populating display from each + team's alias so admin-group matching by display name works the same way it + does on PUT (where SCIM groups carry display names natively). + """ + teams = [ + await TeamRepository(prisma_client).table.find_unique( + where={"team_id": team_id} + ) + for team_id in team_ids + ] + return [ + SCIMUserGroup( + value=team_id, + display=team.team_alias if team is not None else None, + ) + for team_id, team in zip(team_ids, teams) + ] + + +async def _recompute_scim_member_roles( + prisma_client: Any, user_ids: Iterable[str] +) -> None: + """ + Recompute and persist each user's global proxy role from their resulting team + membership. No-op unless scim_admin_group is configured, so a SCIM group write + that drops a member from the admin group demotes them just like the user + endpoints do, and the role is left untouched when the feature is off. + """ + admin_group = await _get_scim_admin_group() + if admin_group is None: + return + + default_role = _default_scim_user_role() + for user_id in user_ids: + user = await UserRepository(prisma_client).table.find_unique( + where={"user_id": user_id} + ) + if user is None: + continue + resolved_role = _resolve_scim_user_role( + await _scim_groups_from_team_ids(prisma_client, user.teams or []), + admin_group, + default_role, + ) + await UserRepository(prisma_client).table.update( + where={"user_id": user_id}, + data={"user_role": resolved_role}, + ) + + async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionResult: """ Extract member IDs from SCIMGroup, validating that all users exist. @@ -999,19 +1132,16 @@ async def create_user( # Create user in database user_id = user.userName or str(uuid.uuid4()) metadata = _build_scim_metadata( - user_data["given_name"], user_data["family_name"] + user_data["given_name"], + user_data["family_name"], + enterprise=user_data["enterprise"], ) - default_role: Optional[ - Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - ] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY - if litellm.default_internal_user_params: - default_role = litellm.default_internal_user_params.get("user_role") + default_role = _default_scim_user_role() + admin_group = await _get_scim_admin_group() + resolved_role = _resolve_scim_user_role( + user.groups or [], admin_group, default_role + ) new_user_request = NewUserRequest( user_id=user_id, @@ -1020,12 +1150,14 @@ async def create_user( teams=user_data["teams"], metadata=metadata, auto_create_key=False, - user_role=default_role, + user_role=resolved_role if admin_group is not None else default_role, ) # Check if user with email already exists and update if found existing_user_scim = await UserProvisionerHelpers.handle_existing_user_by_email( - prisma_client=prisma_client, new_user_request=new_user_request + prisma_client=prisma_client, + new_user_request=new_user_request, + admin_group=admin_group, ) if existing_user_scim: @@ -1088,6 +1220,7 @@ async def update_user( user_data["given_name"], user_data["family_name"], scim_active_for_metadata, + enterprise=user_data["enterprise"], ) await _handle_team_membership_changes( @@ -1104,6 +1237,12 @@ async def update_user( "metadata": safe_dumps(metadata), } + admin_group = await _get_scim_admin_group() + if admin_group is not None: + update_data["user_role"] = _resolve_scim_user_role( + user.groups or [], admin_group, _default_scim_user_role() + ) + updated_user = await UserRepository(prisma_client).table.update( where={"user_id": user_id}, data=update_data, @@ -1417,6 +1556,14 @@ async def patch_user( update_data["teams"] = list(final_team_set) + admin_group = await _get_scim_admin_group() + if admin_group is not None: + update_data["user_role"] = _resolve_scim_user_role( + await _scim_groups_from_team_ids(prisma_client, list(final_team_set)), + admin_group, + _default_scim_user_role(), + ) + # Serialize metadata to JSON string for Prisma to avoid GraphQL parsing issues if "metadata" in update_data and isinstance(update_data["metadata"], dict): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -1599,6 +1746,8 @@ async def create_group( user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) + await _recompute_scim_member_roles(prisma_client, member_result.all_member_ids) + scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( created_team ) @@ -1665,6 +1814,19 @@ async def update_group( final_members=final_members, ) + # A rename can flip whether this group matches scim_admin_group by display + # name, so retained members must be re-resolved too, not just the ones whose + # membership changed. + alias_changed = existing_team.team_alias != group.displayName + await _recompute_scim_member_roles( + prisma_client, + ( + current_members | final_members + if alias_changed + else current_members ^ final_members + ), + ) + # Convert to SCIM format and return scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( updated_team @@ -1691,8 +1853,10 @@ async def delete_group( prisma_client = await _get_prisma_client_or_raise_exception() existing_team = await _check_team_exists(group_id) + member_ids = await _get_team_member_user_ids_from_team(existing_team) + # For each member, remove this team from their teams list - for member_id in existing_team.members or []: + for member_id in member_ids: user = await UserRepository(prisma_client).table.find_unique( where={"user_id": member_id} ) @@ -1704,6 +1868,8 @@ async def delete_group( where={"user_id": member_id}, data={"teams": new_teams} ) + await _recompute_scim_member_roles(prisma_client, member_ids) + # Delete team await TeamRepository(prisma_client).table.delete(where={"team_id": group_id}) @@ -1903,6 +2069,20 @@ async def patch_group( # Handle user-team relationship changes await _handle_group_membership_changes(group_id, current_members, final_members) + # A rename can flip whether this group matches scim_admin_group by display + # name, so retained members must be re-resolved too, not just the ones whose + # membership changed. + new_alias = update_data.get("team_alias", existing_team.team_alias) + alias_changed = new_alias != existing_team.team_alias + await _recompute_scim_member_roles( + prisma_client, + ( + current_members | final_members + if alias_changed + else current_members ^ final_members + ), + ) + # Refresh team one more time to get final state after membership changes final_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": group_id} diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index f2ddae40d8c..07c355f2cd9 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -11,6 +11,7 @@ from fastapi import HTTPException, status from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy._types import SpecialMCPServerNames from litellm.proxy.utils import PrismaClient from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import MCPServerRepository @@ -287,6 +288,9 @@ def _rewrite_object_permission_mcp_servers( normalized_servers: List[str] = [] for identifier in mcp_servers: + if identifier == SpecialMCPServerNames.no_mcp_servers.value: + normalized_servers.append(SpecialMCPServerNames.no_mcp_servers.value) + continue normalized_servers.extend(sorted(identifier_to_server_ids.get(identifier, []))) object_permission["mcp_servers"] = _dedupe_preserving_order(normalized_servers) @@ -426,6 +430,7 @@ def _extract_requested_mcp_server_ids( mcp_servers = object_permission.get("mcp_servers") if isinstance(mcp_servers, list): server_ids.update(mcp_servers) + server_ids.discard(SpecialMCPServerNames.no_mcp_servers.value) mcp_tool_permissions = object_permission.get("mcp_tool_permissions") if isinstance(mcp_tool_permissions, dict): diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index c8f6749a196..8986166ba92 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -6,6 +6,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -15,12 +16,19 @@ from litellm.llms.anthropic import get_anthropic_config from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) +from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) -from litellm.types.utils import LiteLLMBatch, ModelResponse, TextCompletionResponse +from litellm.types.utils import ( + Choices, + LiteLLMBatch, + Message, + ModelResponse, + TextCompletionResponse, +) if TYPE_CHECKING: from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType @@ -272,6 +280,9 @@ class AnthropicPassthroughLoggingHandler: kwargs["response_cost"] = response_cost kwargs["model"] = model + # the pass-through success path reads spend from + # model_call_details["response_cost"], not from kwargs + logging_obj.model_call_details["response_cost"] = response_cost passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = ( # type: ignore kwargs.get("passthrough_logging_payload") ) @@ -343,13 +354,42 @@ class AnthropicPassthroughLoggingHandler: if chunk_model: model = chunk_model - complete_streaming_response = ( - AnthropicPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=all_chunks, - litellm_logging_obj=litellm_logging_obj, - model=model, + try: + complete_streaming_response = ( + AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + ) ) - ) + except Exception as e: + # stream_chunk_builder re-raises assembly failures (as litellm.APIError) + # on large agentic tool-use / thinking streams; treat that the same as a + # None result so the usage-only fallback below still recovers cost + verbose_proxy_logger.warning( + "Anthropic passthrough: stream assembly raised (model=%s): %s; falling " + "back to usage-only cost from raw SSE events.", + model, + e, + ) + complete_streaming_response = None + if complete_streaming_response is None: + # stream_chunk_builder cannot always reassemble large agentic streams, but + # Anthropic still emits token usage in the message_start / message_delta SSE + # events regardless of content shape; recover usage-only so cost is tracked. + # Guard it too: a raise here would defeat the point and drop the request + try: + complete_streaming_response = AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=all_chunks, + model=model, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Anthropic passthrough: usage-only fallback failed (model=%s): %s", + model, + e, + ) + complete_streaming_response = None if complete_streaming_response is None: verbose_proxy_logger.error( "Unable to build complete streaming response for Anthropic passthrough endpoint, not logging..." @@ -636,6 +676,141 @@ class AnthropicPassthroughLoggingHandler: ) return complete_streaming_response + @staticmethod + def _extract_sse_data(event_str: str) -> Optional[dict]: + """Parse the JSON object from the ``data:`` line of an Anthropic SSE event.""" + for line in event_str.splitlines(): + stripped = line.strip() + if stripped.startswith("data:"): + payload = stripped[len("data:") :].strip() + if not payload or payload == "[DONE]": + return None + try: + return cast(dict, json.loads(payload)) + except (ValueError, TypeError): + return None + return None + + @staticmethod + def _build_usage_only_response_from_chunks( + all_chunks: Sequence[Union[str, bytes]], + model: str, + ) -> Optional[ModelResponse]: + """ + Build a usage-bearing ModelResponse from Anthropic SSE token-usage events, for + cost tracking when stream_chunk_builder cannot reassemble the stream. + + Anthropic emits usage in ``message_start`` (uncached input + cache tokens, and an + initial output_tokens) and the final ``message_delta`` (cumulative output_tokens) + regardless of the content/tool shape, so cost is recoverable even when full + content assembly fails. Returns ``None`` if no usage event is found. + """ + input_tokens = 0 + cache_read = 0 + cache_creation = 0 + cache_creation_5m: Optional[int] = None + cache_creation_1h: Optional[int] = None + output_tokens = 0 + web_search_requests: Optional[int] = None + tool_search_requests: Optional[int] = None + inference_geo: Optional[str] = None + stop_reason: Optional[str] = None + found_usage = False + resolved_model = model + for _chunk_str in all_chunks: + for ( + event_str + ) in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events( + _chunk_str + ): + data = AnthropicPassthroughLoggingHandler._extract_sse_data(event_str) + if not data: + continue + event_type = data.get("type") + if event_type == "message_start": + message = data.get("message") or {} + if not resolved_model or resolved_model == "unknown": + resolved_model = message.get("model") or resolved_model + usage = message.get("usage") or {} + input_tokens = usage.get("input_tokens") or input_tokens + cache_read = usage.get("cache_read_input_tokens") or cache_read + cache_creation = ( + usage.get("cache_creation_input_tokens") or cache_creation + ) + _cc = usage.get("cache_creation") + if isinstance(_cc, dict): + cache_creation_5m = _cc.get("ephemeral_5m_input_tokens") + cache_creation_1h = _cc.get("ephemeral_1h_input_tokens") + if usage.get("inference_geo") is not None: + inference_geo = usage.get("inference_geo") + if usage.get("output_tokens") is not None: + output_tokens = usage.get("output_tokens") + found_usage = True + elif event_type == "message_delta": + _delta_stop = (data.get("delta") or {}).get("stop_reason") + if _delta_stop: + stop_reason = _delta_stop + usage = data.get("usage") or {} + if usage.get("output_tokens") is not None: + output_tokens = usage.get("output_tokens") + _stu = usage.get("server_tool_use") + if isinstance(_stu, dict): + if _stu.get("web_search_requests") is not None: + web_search_requests = _stu.get("web_search_requests") + if _stu.get("tool_search_requests") is not None: + tool_search_requests = _stu.get("tool_search_requests") + if usage.get("cache_read_input_tokens") is not None: + cache_read = usage.get("cache_read_input_tokens") + if usage.get("inference_geo") is not None: + inference_geo = usage.get("inference_geo") + found_usage = True + if not found_usage: + return None + # If only the 5m/1h split was provided, derive the cache_creation total from it. + if not cache_creation and (cache_creation_5m or cache_creation_1h): + cache_creation = (cache_creation_5m or 0) + (cache_creation_1h or 0) + # build usage via the same AnthropicConfig.calculate_usage path the success + # cases use, so prompt_tokens are cache-inclusive and cache / server_tool_use / + # inference_geo tokens are priced instead of left at $0 + usage_object: dict = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } + if cache_read: + usage_object["cache_read_input_tokens"] = cache_read + if cache_creation: + usage_object["cache_creation_input_tokens"] = cache_creation + if cache_creation_5m is not None or cache_creation_1h is not None: + usage_object["cache_creation"] = { + "ephemeral_5m_input_tokens": cache_creation_5m or 0, + "ephemeral_1h_input_tokens": cache_creation_1h or 0, + } + if web_search_requests is not None or tool_search_requests is not None: + _server_tool_use: dict = {} + if web_search_requests is not None: + _server_tool_use["web_search_requests"] = web_search_requests + if tool_search_requests is not None: + _server_tool_use["tool_search_requests"] = tool_search_requests + usage_object["server_tool_use"] = _server_tool_use + if inference_geo is not None: + usage_object["inference_geo"] = inference_geo + usage_obj = AnthropicConfig().calculate_usage( + usage_object=usage_object, reasoning_content=None + ) + return ModelResponse( + model=resolved_model, + choices=[ + Choices( + finish_reason=( + map_finish_reason(stop_reason) if stop_reason else "stop" + ), + index=0, + message=Message(role="assistant", content=""), + ) + ], + usage=usage_obj, + ) + @staticmethod def batch_creation_handler( httpx_response: httpx.Response, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py index b9df8ecede3..a7ec2f0d368 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/base_passthrough_logging_handler.py @@ -116,6 +116,9 @@ class BasePassthroughLoggingHandler(ABC): kwargs["response_cost"] = response_cost kwargs["model"] = model + # the pass-through success path reads spend from + # model_call_details["response_cost"], not from kwargs + logging_obj.model_call_details["response_cost"] = response_cost passthrough_logging_payload: Optional[PassthroughStandardLoggingPayload] = ( # type: ignore kwargs.get("passthrough_logging_payload") ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 33a6b719280..7a725472dd7 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -285,8 +285,10 @@ class PassThroughStreamingHandler: Returns: List of string lines, with each line being a complete data: {} chunk """ - # Combine all bytes and decode to string - combined_str = b"".join(raw_bytes).decode("utf-8") + # errors="replace" so a stream cut mid-multibyte-sequence (client disconnect) + # still decodes and logs the usage events already received, instead of raising + # and dropping the whole request from SpendLogs + combined_str = b"".join(raw_bytes).decode("utf-8", errors="replace") # Split by newlines and filter out empty lines lines = [line.strip() for line in combined_str.split("\n") if line.strip()] diff --git a/litellm/proxy/plugin_routes.py b/litellm/proxy/plugin_routes.py new file mode 100644 index 00000000000..6a94f78fbe7 --- /dev/null +++ b/litellm/proxy/plugin_routes.py @@ -0,0 +1,344 @@ +""" +Plugin proxy routes for litellm. + +Enables external services to register as plugins and be proxied through +the litellm proxy server. + +Config (in litellm config.yaml general_settings): + plugins: + - name: my-plugin + url: "http://localhost:3210" + display_name: "My Plugin" + plugin_key: "sk-..." # optional: plugin's own auth key + +Plugin iframe auth: + The UI calls GET /api/plugins/auth-token to receive a short-lived identity + claim ({user_id, user_role, plugin, exp}) encrypted with a per-plugin key + derived as HMAC-SHA256(LITELLM_SALT_KEY, plugin_name). The claim carries no + litellm bearer token, so a compromised plugin learns only the caller's + identity, never their credential. LITELLM_SALT_KEY itself is never shared + with plugins — each plugin holds only its own derived key. +""" + +import base64 +import hashlib +import hmac as _hmac +import json +import os +import time +from collections.abc import Mapping + +from cryptography.fernet import Fernet, InvalidToken +from fastapi import APIRouter, Depends, HTTPException, Request, Response + +from litellm.proxy._types import PluginConfig, SpecialHeaders, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.llms.custom_http import httpxSpecialProvider + +router = APIRouter() + +# Hop-by-hop headers (RFC 7230) and the litellm session cookie — never forwarded +# to a plugin backend. Credential headers are added on top per-request from the +# canonical SpecialHeaders set so the plugin only ever authenticates via its own +# injected plugin_key. +_HOP_BY_HOP_STRIP = frozenset( + { + "host", + "connection", + "transfer-encoding", + "te", + "trailers", + "upgrade", + "cookie", + } +) + + +def _configured_key_header_names() -> frozenset[str]: + """The lowercased general_settings.litellm_key_header_name, if configured. + + Read live from the proxy module (not import-time) so a custom key header set + via config is honoured without a restart. Returns empty when unset. + """ + try: + from litellm.proxy import proxy_server + except Exception: + return frozenset() + general_settings = getattr(proxy_server, "general_settings", None) + if not isinstance(general_settings, dict): + return frozenset() + name: object = general_settings.get("litellm_key_header_name") + return frozenset({name.lower()}) if isinstance(name, str) and name else frozenset() + + +def _request_strip_headers() -> frozenset[str]: + """Headers to drop before forwarding a request to a plugin backend. + + Every header user_api_key_auth accepts as a litellm credential is stripped — + Authorization, x-api-key, API-Key, x-goog-api-key, Ocp-Apim-Subscription-Key, + x-litellm-api-key, and any configured custom key header — so a plugin can + never be handed the caller's live litellm key (confused-deputy escalation). + """ + return ( + _HOP_BY_HOP_STRIP + | SpecialHeaders.litellm_credential_header_names() + | _configured_key_header_names() + ) + + +# Headers to strip from plugin RESPONSES before returning to the browser. +# httpx already decompresses and de-chunks the body, so forwarding the wire +# encoding headers causes clients to attempt double-decompression (garbage) or +# incorrect length checks. set-cookie is removed so plugins cannot overwrite +# litellm session cookies. +_RESPONSE_STRIP = { + "content-encoding", + "transfer-encoding", + "content-length", + "set-cookie", +} + + +def _safe_response_headers(raw: "Mapping[str, str]") -> dict[str, str]: + """Strip wire-encoding/cookie headers and force proxied responses inert. + + Plugin-controlled bytes are served from the litellm dashboard origin, so a + compromised plugin could return an HTML/JS document that executes with the + admin's session against same-origin management APIs. A sandbox CSP forces + the response into an opaque origin with scripts disabled, and nosniff stops + content-type confusion from re-enabling execution. Both are set last so a + plugin cannot override them with its own headers. + """ + return { + **{k: v for k, v in raw.items() if k.lower() not in _RESPONSE_STRIP}, + "content-security-policy": "sandbox", + "x-content-type-options": "nosniff", + } + + +# In-memory plugin registry — populated from general_settings at startup +_plugin_registry: dict[str, PluginConfig] = {} + + +# --------------------------------------------------------------------------- +# Key derivation — audience-scoped per plugin so compromising one plugin +# cannot be used to forge claims for another. LITELLM_SALT_KEY is NEVER +# shared with plugins; each plugin only receives a key derived from +# HMAC(LITELLM_SALT_KEY, plugin_name) which reveals nothing about the master. +# --------------------------------------------------------------------------- +def _plugin_fernet(plugin_name: str) -> Fernet: + """Return a Fernet cipher whose key is scoped to a specific plugin. + + Key material: HMAC-SHA256(LITELLM_SALT_KEY, plugin_name). + A plugin possessing its own key cannot derive the master salt or + forge claims intended for a different plugin. + """ + salt = os.getenv("LITELLM_SALT_KEY", "").encode() + derived = _hmac.new(salt, plugin_name.encode(), hashlib.sha256).digest() + return Fernet(base64.urlsafe_b64encode(derived)) + + +_CLAIM_TTL_SECONDS = 30 # identity claims expire after 30 s + + +def issue_plugin_session_claim( + plugin_name: str, user_id: str | None, user_role: str | None +) -> str: + """Issue a short-lived, audience-scoped identity claim for the plugin. + + The claim contains {user_id, user_role, plugin, exp}. Crucially it + contains NO litellm bearer token — the plugin can only derive the + caller's identity, not act as them against the proxy. + """ + claim = { + "plugin": plugin_name, + "user_id": user_id or "", + "user_role": user_role or "", + "exp": int(time.time()) + _CLAIM_TTL_SECONDS, + } + return _plugin_fernet(plugin_name).encrypt(json.dumps(claim).encode()).decode() + + +def verify_plugin_session_claim(plugin_name: str, ciphertext: str) -> dict: + """Verify and decode a plugin session claim. + + Raises ValueError if the HMAC is invalid, the audience is wrong, or + the claim is expired. Returns the decoded claim dict on success. + """ + try: + raw = _plugin_fernet(plugin_name).decrypt( + ciphertext.encode(), ttl=_CLAIM_TTL_SECONDS + ) + claim = json.loads(raw) + except (InvalidToken, Exception) as exc: + raise ValueError("Invalid, tampered, or expired plugin session claim") from exc + + if claim.get("plugin") != plugin_name: + raise ValueError("Plugin claim audience mismatch") + if int(claim.get("exp", 0)) < int(time.time()): + raise ValueError("Plugin session claim expired") + return claim + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- +def register_plugins_from_config(general_settings: dict[str, object]) -> None: + """Replace the plugin registry from general_settings. + + Replaces (not merges) so plugins removed from config are immediately + unreachable without requiring a process restart. + """ + raw = general_settings.get("plugins") + entries: list[object] = raw if isinstance(raw, list) else [] + new_registry = { + p.name: p for p in (PluginConfig.model_validate(entry) for entry in entries) + } + _plugin_registry.clear() + _plugin_registry.update(new_registry) + + +# --------------------------------------------------------------------------- +# Routes +# --------------------------------------------------------------------------- +@router.get("/api/plugins", tags=["plugins"]) +async def list_plugins( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> list[dict[str, str]]: + """Return registered plugins for authenticated UI callers. + + plugin_key is never returned — the browser never needs it (the proxy injects + it server-side from the registry), and exposing it here would leak the + credential into React state and DevTools. Admin key management goes through + the redacted /config/field/info path instead. + """ + return [ + { + "name": plugin.name, + "display_name": plugin.display_name or plugin.name, + "url": plugin.url, + } + for plugin in _plugin_registry.values() + ] + + +@router.get("/api/plugins/auth-token", tags=["plugins"]) +async def plugin_auth_token( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + plugin_name: str = "litellm-platform-plugin", +) -> dict: + """Issue a short-lived, audience-scoped plugin session claim. + + The claim contains {user_id, user_role, plugin, exp}. It does NOT + contain the caller's litellm bearer token — a compromised plugin can + only learn the caller's identity, not impersonate them against the proxy. + + Encrypted with a key derived from HMAC(LITELLM_SALT_KEY, plugin_name), + so each plugin holds only its own key and cannot forge claims for others. + + Requires LITELLM_SALT_KEY to be set; returns 503 otherwise. + """ + if not os.getenv("LITELLM_SALT_KEY"): + raise HTTPException( + status_code=503, + detail="LITELLM_SALT_KEY is not configured; plugin iframe auth unavailable.", + ) + if plugin_name not in _plugin_registry: + raise HTTPException( + status_code=404, detail=f"Plugin '{plugin_name}' is not registered." + ) + user_id = getattr(user_api_key_dict, "user_id", None) + user_role = getattr(user_api_key_dict, "user_role", None) + return { + "session_claim": issue_plugin_session_claim(plugin_name, user_id, user_role) + } + + +@router.api_route( + "/plugin-proxy/{plugin_name}/{path:path}", + methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"], + tags=["plugins"], + include_in_schema=False, +) +async def plugin_proxy( + plugin_name: str, + path: str, + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> Response: + """Authenticated reverse-proxy to a registered plugin backend. + + Restricted to proxy_admin callers — the shared plugin_key must not be + usable as a confused-deputy credential by regular users. Plugin UIs + talk to the plugin service directly via the iframe; this route is for + administrative and server-to-server access only. + + The caller's litellm credential is stripped and replaced with the + plugin's own plugin_key so plugins never receive a live litellm API key. + """ + if getattr(user_api_key_dict, "user_role", None) != "proxy_admin": + return Response( + content="Plugin proxy access requires proxy_admin role.", + status_code=403, + ) + + plugin = _plugin_registry.get(plugin_name) + if not plugin: + return Response( + content=f"Plugin '{plugin_name}' not registered", + status_code=404, + ) + + target_url = f"{plugin.url.rstrip('/')}/{path}" + query = request.url.query + if query: + target_url = f"{target_url}?{query}" + + body = await request.body() + + # Strip caller credentials and hop-by-hop headers from forwarded request + strip = _request_strip_headers() + forward_headers = { + k: v for k, v in request.headers.items() if k.lower() not in strip + } + + # Inject plugin's own credential as upstream auth (if configured) + plugin_key = plugin.plugin_key + if plugin_key: + forward_headers["authorization"] = f"Bearer {plugin_key}" + + # Forward caller identity so the plugin can enforce its own access control. + # The plugin MUST NOT trust these as credentials — they are informational. + # The plugin_key above is the only authentication mechanism. + user_id = getattr(user_api_key_dict, "user_id", None) + user_role = getattr(user_api_key_dict, "user_role", None) + if user_id: + forward_headers["x-litellm-user-id"] = str(user_id) + if user_role: + forward_headers["x-litellm-user-role"] = str(user_role) + + handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint + ) + try: + req = handler.client.build_request( + method=request.method, + url=target_url, + headers=forward_headers, + content=body, + ) + # Do not follow redirects — a redirect to an internal URL would allow + # the plugin to SSRF the proxy into fetching arbitrary internal services. + resp = await handler.client.send(req, follow_redirects=False) + except Exception: + return Response( + content=f"Cannot connect to plugin '{plugin_name}' at {plugin.url}", + status_code=502, + ) + + return Response( + content=resp.content, + status_code=resp.status_code, + headers=_safe_response_headers(resp.headers), + ) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 9c4d7b1bb5d..d0281885482 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1195,6 +1195,25 @@ def run_server( os.getenv("DATABASE_URL", None) is not None or os.getenv("DIRECT_URL", None) is not None ): + from litellm.proxy.db.db_url_settings import ( + unsupported_db_scheme, + unsupported_db_scheme_message, + ) + + for _db_env in ("DATABASE_URL", "DIRECT_URL"): + _candidate_url = os.getenv(_db_env) + if _candidate_url is None: + continue + _bad_scheme = unsupported_db_scheme(_candidate_url) + if _bad_scheme is not None: + print( + f"\033[1;31mLiteLLM Proxy: " + f"{unsupported_db_scheme_message(_db_env, _bad_scheme)}" + "\033[0m", + file=sys.stderr, + flush=True, + ) + sys.exit(1) try: from litellm.secret_managers.main import get_secret diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c138626a272..e4d933ac699 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -38,7 +38,7 @@ from typing import ( import anyio import websockets import websockets.exceptions -from pydantic import BaseModel, Json +from pydantic import BaseModel, Json, JsonValue from litellm._uuid import uuid from litellm.constants import ( @@ -106,6 +106,10 @@ from litellm.proxy.common_utils.callback_utils import ( process_callback, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body +from litellm.router_utils.add_retry_fallback_headers import ( + get_fallback_errors_from_headers, + get_hidden_params_dict, +) from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -419,6 +423,10 @@ from litellm.proxy.management_endpoints.workflow_management_endpoints import ( ) from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update from litellm.proxy.memory.memory_endpoints import router as memory_router +from litellm.proxy.plugin_routes import ( + router as plugin_router, + register_plugins_from_config, +) from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) @@ -4455,6 +4463,8 @@ class ProxyConfig: f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}" ) setattr(litellm, key, value) + if key == "request_timeout": + litellm.request_timeout_explicitly_set = True if key in {"s3_audit_callback_params", "s3_callback_params"}: from litellm.integrations.s3_v2 import S3Logger as S3V2Logger from litellm.litellm_core_utils.litellm_logging import ( @@ -4502,6 +4512,8 @@ class ProxyConfig: load_from_azure_key_vault(use_azure_key_vault=use_azure_key_vault) ### ALERTING ### self._load_alerting_settings(general_settings=general_settings) + ### PLUGINS ### + register_plugins_from_config(general_settings) ### CONNECT TO DATABASE ### database_url = general_settings.get("database_url", None) if database_url and database_url.startswith("os.environ/"): @@ -4773,6 +4785,11 @@ class ProxyConfig: config ) + ## SANDBOX TOOLS SETTINGS + from litellm.sandbox.sandbox_tools import register_sandbox_tools + + register_sandbox_tools(config.get("sandbox_tools") or []) + ## /fine_tuning/jobs endpoints config finetuning_config = config.get("finetune_settings", None) set_fine_tuning_config(config=finetuning_config) @@ -5663,6 +5680,10 @@ class ProxyConfig: llm_router=llm_router, ) + if _general_settings is not None and "plugins" in _general_settings: + general_settings["plugins"] = _general_settings["plugins"] + register_plugins_from_config(general_settings) + async def _reschedule_spend_log_cleanup_job(self): """ Reschedule the spend log cleanup job based on current general_settings. @@ -7070,57 +7091,122 @@ def _get_client_requested_model_for_streaming(request_data: dict) -> str: return requested_model if isinstance(requested_model, str) else "" +def _is_positive_int_like(value: Any) -> bool: + try: + return int(value) > 0 + except (TypeError, ValueError): + return False + + +def _should_include_fallback_errors(request_data: dict[str, object]) -> bool: + if not general_settings.get("expose_fallback_errors_to_caller"): + return False + return request_data.get("include_fallback_errors") is True + + +def _get_streaming_fallback_metadata( + response_obj: object, +) -> tuple[bool, str | None, list[dict[str, object]]]: + additional_headers = get_hidden_params_dict(response_obj).get("additional_headers") + if not isinstance(additional_headers, dict): + return False, None, [] + + if not _is_positive_int_like( + additional_headers.get("x-litellm-attempted-fallbacks") + ): + return False, None, [] + + fallback_model = additional_headers.get("x-litellm-model-group") + fallback_errors = get_fallback_errors_from_headers(additional_headers) + if isinstance(fallback_model, str) and fallback_model: + return True, fallback_model, fallback_errors + return True, None, fallback_errors + + +def _format_fallback_metadata_sse_event( + *, + fallback_model: str | None, + fallback_errors: list[dict[str, object]], +) -> str: + import time + + payload = { + "id": "litellm-fallback-metadata", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": fallback_model or "", + "choices": [], + "litellm_fallback": { + "fallback_model": fallback_model, + "errors": fallback_errors, + }, + } + return f"data: {json.dumps(payload)}\n\n" + + def _restamp_streaming_chunk_model( *, chunk: Any, requested_model_from_client: str, request_data: dict, model_mismatch_logged: bool, -) -> Tuple[Any, bool]: + fallback_was_attempted: bool = False, + fallback_model_from_metadata: str | None = None, +) -> tuple[Any, bool]: + target_model = ( + fallback_model_from_metadata + if fallback_was_attempted + else requested_model_from_client + ) # Always return the client-requested model name (not provider-prefixed internal identifiers) # on streaming chunks. + # On fallback, use the public OpenAI-compatible model name. This keeps + # provider-prefixed internal identifiers from leaking into the public API. # # Note: This warning is intentionally verbose. A mismatch is a useful signal that an # internal provider/deployment identifier is leaking into the public API, and helps # maintainers/operators catch regressions while preserving OpenAI-compatible output. - if not requested_model_from_client or not isinstance(chunk, (BaseModel, dict)): + if not target_model or not isinstance(chunk, (BaseModel, dict)): return chunk, model_mismatch_logged # For Azure Model Router, preserve the actual model used in each chunk - if _is_azure_model_router_request(requested_model_from_client): + if not fallback_was_attempted and _is_azure_model_router_request( + requested_model_from_client + ): return chunk, model_mismatch_logged # For fastest_response batch completions, preserve the winning model's name # instead of stamping the comma-separated list the client sent. - if request_data.get("fastest_response", False): + if not fallback_was_attempted and request_data.get("fastest_response", False): return chunk, model_mismatch_logged downstream_model = ( chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None) ) - if downstream_model == requested_model_from_client: + if downstream_model == target_model: return chunk, model_mismatch_logged - if not model_mismatch_logged and downstream_model != requested_model_from_client: + if not model_mismatch_logged and downstream_model != target_model: verbose_proxy_logger.debug( - "litellm_call_id=%s: streaming chunk model mismatch - requested=%r downstream=%r. Overriding model to requested.", + "litellm_call_id=%s: streaming chunk model mismatch - target=%r downstream=%r fallback_was_attempted=%s. Overriding chunk model to target.", request_data.get("litellm_call_id"), - requested_model_from_client, + target_model, downstream_model, + fallback_was_attempted, ) model_mismatch_logged = True if isinstance(chunk, dict): - chunk["model"] = requested_model_from_client + chunk["model"] = target_model return chunk, model_mismatch_logged try: - setattr(chunk, "model", requested_model_from_client) + chunk.model = target_model except Exception as e: verbose_proxy_logger.error( "litellm_call_id=%s: failed to override chunk.model=%r on chunk_type=%s. error=%s", request_data.get("litellm_call_id"), - requested_model_from_client, + target_model, type(chunk), str(e), exc_info=True, @@ -7279,7 +7365,14 @@ async def async_data_generator( requested_model_from_client = _get_client_requested_model_for_streaming( request_data=request_data ) + ( + fallback_was_attempted, + fallback_model_from_metadata, + fallback_errors, + ) = _get_streaming_fallback_metadata(response) model_mismatch_logged = False + fallback_metadata_event_sent = False + include_fallback_errors = _should_include_fallback_errors(request_data) # Use a running string instead of list + join to avoid O(n^2) overhead. # Previously "".join(str_so_far_parts) was called every chunk, re-joining # the entire accumulated response. String += is O(n) amortized total. @@ -7317,13 +7410,37 @@ async def async_data_generator( str_so_far=_str_so_far, ) + # Mid-stream fallbacks surface metadata on individual chunks rather than + # the response wrapper. Keep scanning chunks until a fallback model is + # resolved, then latch it for the rest of the stream. + if fallback_model_from_metadata is None: + ( + chunk_fallback_was_attempted, + chunk_fallback_model, + chunk_fallback_errors, + ) = _get_streaming_fallback_metadata(chunk) + if chunk_fallback_was_attempted: + fallback_was_attempted = True + fallback_model_from_metadata = chunk_fallback_model + fallback_errors = fallback_errors or chunk_fallback_errors + + pending_fallback_event = ( + include_fallback_errors + and fallback_was_attempted + and fallback_errors + and not fallback_metadata_event_sent + ) + chunk, model_mismatch_logged = _restamp_streaming_chunk_model( chunk=chunk, requested_model_from_client=requested_model_from_client, request_data=request_data, model_mismatch_logged=model_mismatch_logged, + fallback_was_attempted=fallback_was_attempted, + fallback_model_from_metadata=fallback_model_from_metadata, ) + raw_passthrough = False if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) elif isinstance(chunk, bytes): @@ -7339,14 +7456,14 @@ async def async_data_generator( raise ValueError( "Raw SSE stream exceeded maximum buffered size without a frame delimiter" ) - continue - if chunk.startswith(("data:", "event:", ":")): + raw_passthrough = True + elif chunk.startswith(("data:", "event:", ":")): yield ( chunk if chunk.endswith(_SSE_FRAME_DELIMITERS) else chunk + "\n\n" ) - continue + raw_passthrough = True elif isinstance(chunk, str) and is_raw_sse_stream: raw_sse_buffer += chunk while True: @@ -7358,15 +7475,23 @@ async def async_data_generator( raise ValueError( "Raw SSE stream exceeded maximum buffered size without a frame delimiter" ) - continue + raw_passthrough = True elif isinstance(chunk, str) and chunk.startswith("data: "): error_message = chunk break - try: - yield _format_streaming_sse_chunk(chunk=chunk) - except Exception as e: - yield f"data: {str(e)}\n\n" + if not raw_passthrough: + try: + yield _format_streaming_sse_chunk(chunk=chunk) + except Exception as e: + yield f"data: {str(e)}\n\n" + + if pending_fallback_event: + yield _format_fallback_metadata_sse_event( + fallback_model=fallback_model_from_metadata, + fallback_errors=fallback_errors, + ) + fallback_metadata_event_sent = True stream_completed = True if not needs_iterator_wrap: @@ -11566,9 +11691,12 @@ async def _get_caller_byok_team_scope( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ): return None + key_team_scope: set[str] = ( + {user_api_key_dict.team_id} if user_api_key_dict.team_id else set() + ) user_id = user_api_key_dict.user_id if user_id is None: - return set() + return key_team_scope try: user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} @@ -11576,12 +11704,12 @@ async def _get_caller_byok_team_scope( except Exception: verbose_proxy_logger.exception( "Failed to look up caller teams while scoping BYOK search; " - "defaulting to no team access." + "defaulting to key team scope only." ) - return set() + return key_team_scope if user_row is None: - return set() - return set(user_row.teams or []) + return key_team_scope + return key_team_scope | set(user_row.teams or []) def _byok_row_outside_caller_teams( @@ -14931,6 +15059,41 @@ async def update_config( Keep it more precise, to prevent overwrite other values unintentially """ +_PLUGIN_KEY_REDACTED = "***" + + +def _preserve_redacted_plugin_keys(incoming: object, existing: object) -> object: + """Restore real plugin_key values the client never sees. + + /config/field/info redacts every plugin_key to ``"***"``, so an admin + editing a plugin posts that placeholder (or a blank, when the UI clears the + field) straight back. Treat a blank or redacted plugin_key as "keep the + stored credential" by sourcing it from the existing config; only a real, + non-redacted value replaces it, and a blank with no stored key drops the + field entirely instead of persisting the placeholder. + """ + if not isinstance(incoming, list): + return incoming + + stored_keys = { + p["name"]: p["plugin_key"] + for p in (existing if isinstance(existing, list) else []) + if isinstance(p, dict) and p.get("name") and p.get("plugin_key") + } + + def resolve(plugin: object) -> object: + if not isinstance(plugin, dict): + return plugin + key = plugin.get("plugin_key") + if key not in (None, "", _PLUGIN_KEY_REDACTED): + return plugin + name = plugin.get("name") + if name in stored_keys: + return {**plugin, "plugin_key": stored_keys[name]} + return {k: v for k, v in plugin.items() if k != "plugin_key"} + + return [resolve(p) for p in incoming] + @router.post( "/config/field/update", @@ -14997,7 +15160,13 @@ async def update_config_general_settings( ## update db - general_settings[data.field_name] = data.field_value + field_value = data.field_value + if data.field_name == "plugins": + field_value = _preserve_redacted_plugin_keys( + field_value, general_settings.get("plugins") + ) + + general_settings[data.field_name] = field_value response = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, @@ -15008,9 +15177,74 @@ async def update_config_general_settings( ) await invalidate_config_param("general_settings") + if data.field_name == "plugins": + register_plugins_from_config(general_settings) + return response +# Secret-bearing general_settings fields the segment masker does not match by +# name: database_url and database_extra_connection_params embed DB credentials, +# pass_through_endpoints carry upstream Authorization headers, and +# alert_to_webhook_url is itself a webhook secret +_EXTRA_SECRET_GENERAL_SETTINGS_FIELDS = frozenset( + { + "database_url", + "database_extra_connection_params", + "pass_through_endpoints", + "alert_to_webhook_url", + } +) + + +def _is_secret_general_setting_field(field_name: str) -> bool: + return ( + field_name in _EXTRA_SECRET_GENERAL_SETTINGS_FIELDS + or SENSITIVE_DATA_MASKER.is_sensitive_key(field_name) + ) + + +# Matches the cap on _redact_sensitive_litellm_params (the closest analog in the +# proxy). Past this depth we fail closed by returning "REDACTED" for the whole +# subtree rather than recursing further — better to over-redact a pathological +# config than to silently return a deeply-nested credential verbatim +_REDACT_SECRET_MAX_DEPTH = 10 + + +def _redact_secret_values_in_obj(value: JsonValue, depth: int = 0) -> JsonValue: + """Recursively redact secret leaves inside a structured field so a nested + credential (e.g. aws_web_identity_token under database_args) is never + returned to a non-admin, while non-secret siblings stay visible. At + _REDACT_SECRET_MAX_DEPTH the whole subtree is replaced with "REDACTED" + so depth-overrun fails closed.""" + if depth >= _REDACT_SECRET_MAX_DEPTH: + return "REDACTED" + if isinstance(value, dict): + return { + key: ( + "REDACTED" + if _is_secret_general_setting_field(key) + else _redact_secret_values_in_obj(sub, depth + 1) + ) + for key, sub in value.items() + } + if isinstance(value, list): + return [_redact_secret_values_in_obj(item, depth + 1) for item in value] + return value + + +def _redact_general_setting_value( + field_name: str, value: JsonValue, is_full_admin: bool +) -> JsonValue: + if is_full_admin: + return value + if _is_secret_general_setting_field(field_name): + return "REDACTED" + if isinstance(value, (dict, list)): + return _redact_secret_values_in_obj(value) + return value + + @router.get( "/config/field/info", tags=["config.yaml"], @@ -15063,9 +15297,21 @@ async def get_config_general_settings( general_settings = dict(db_general_settings.param_value) if field_name in general_settings: - return ConfigFieldInfo( - field_name=field_name, field_value=general_settings[field_name] + field_value = _redact_general_setting_value( + field_name, + general_settings[field_name], + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, ) + if field_name == "plugins" and isinstance(field_value, list): + field_value = [ + ( + {k: ("***" if k == "plugin_key" else v) for k, v in p.items()} + if isinstance(p, dict) + else p + ) + for p in field_value + ] + return ConfigFieldInfo(field_name=field_name, field_value=field_value) else: raise HTTPException( status_code=400, @@ -15111,6 +15357,8 @@ async def get_config_list( }, ) + is_full_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ## get general settings from db db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} @@ -15159,7 +15407,11 @@ async def get_config_list( field_name=sub_field, field_type=sub_field_type.__name__, field_description="", # Add custom logic if descriptions are available - field_default_value=general_settings.get(sub_field, None), + field_default_value=_redact_general_setting_value( + sub_field, + general_settings.get(sub_field, None), + is_full_admin, + ), stored_in_db=None, ) for sub_field, sub_field_type in pydantic_class.__annotations__.items() @@ -15189,7 +15441,11 @@ async def get_config_list( field_name=field_name, field_type=allowed_args[field_name]["type"], field_description=field_info.description or "", - field_value=general_settings.get(field_name, None), + field_value=_redact_general_setting_value( + field_name, + general_settings.get(field_name, None), + is_full_admin, + ), stored_in_db=_stored_in_db, field_default_value=field_info.default, nested_fields=nested_fields, @@ -15213,7 +15469,9 @@ async def get_config_list( field_name=field_name, field_type=allowed_args[field_name]["type"], field_description=field_info.description or "", - field_value=_field_value, + field_value=_redact_general_setting_value( + field_name, _field_value, is_full_admin + ), stored_in_db=_stored_in_db, field_default_value=field_info.default, nested_fields=nested_fields, @@ -16387,6 +16645,7 @@ app.include_router(model_access_group_management_router) app.include_router(tag_management_router) app.include_router(workflow_management_router) app.include_router(memory_router) +app.include_router(plugin_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index fac732bac68..fcc6aac1c14 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -209,6 +209,114 @@ ], "default_model_placeholder": "claude-3-opus" }, + { + "provider": "BedrockMantle", + "provider_display_name": "Amazon Bedrock Mantle", + "litellm_provider": "bedrock_mantle", + "credential_fields": [ + { + "key": "api_key", + "label": "Bedrock Mantle API Key", + "placeholder": null, + "tooltip": "Bearer token for the Bedrock Mantle OpenAI-compatible endpoint. You can provide the raw token or the environment variable (e.g. `os.environ/BEDROCK_MANTLE_API_KEY`). Leave blank to authenticate with AWS SigV4 credentials instead.", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_access_key_id", + "label": "AWS Access Key ID", + "placeholder": null, + "tooltip": "Used for AWS SigV4 auth when no API key is set. You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_secret_access_key", + "label": "AWS Secret Access Key", + "placeholder": null, + "tooltip": "Used for AWS SigV4 auth when no API key is set. You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_session_token", + "label": "AWS Session Token", + "placeholder": null, + "tooltip": "Temporary credentials session token. You can provide the raw token or the environment variable (e.g. `os.environ/MY_SESSION_TOKEN`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_region_name", + "label": "AWS Region Name", + "placeholder": "us-east-1", + "tooltip": "Region of the Bedrock Mantle endpoint. Defaults to us-east-1. You can provide the raw value or the environment variable (e.g. `os.environ/AWS_REGION_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_session_name", + "label": "AWS Session Name", + "placeholder": "my-session", + "tooltip": "Name for the AWS session. You can provide the raw value or the environment variable (e.g. `os.environ/MY_SESSION_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_profile_name", + "label": "AWS Profile Name", + "placeholder": "default", + "tooltip": "AWS profile name to use for authentication. You can provide the raw value or the environment variable (e.g. `os.environ/MY_PROFILE_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_role_name", + "label": "AWS Role Name", + "placeholder": "MyRole", + "tooltip": "AWS IAM role name to assume. You can provide the raw value or the environment variable (e.g. `os.environ/MY_ROLE_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_web_identity_token", + "label": "AWS Web Identity Token", + "placeholder": null, + "tooltip": "Web identity token for OIDC authentication. You can provide the raw token or the environment variable (e.g. `os.environ/MY_WEB_IDENTITY_TOKEN`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://bedrock-mantle.us-east-1.api.aws", + "tooltip": "Optional. Custom Bedrock Mantle endpoint. Defaults to https://bedrock-mantle..api.aws. You can provide the raw value or the environment variable (e.g. `os.environ/BEDROCK_MANTLE_API_BASE`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "bedrock_mantle/openai.gpt-oss-120b" + }, { "provider": "Anthropic", "provider_display_name": "Anthropic", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ea8ab2f9b8e..755602cbcc0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3792,6 +3792,10 @@ class PrismaClient: db_data["members_with_roles"], list ): db_data["members_with_roles"] = json.dumps(db_data["members_with_roles"]) + if db_data.get("budget_limits", None) is not None and isinstance( + db_data["budget_limits"], list + ): + db_data["budget_limits"] = json.dumps(db_data["budget_limits"]) return db_data # Define a retrying strategy with exponential backoff diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 34c9cdd3d1c..2c46baaada5 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -58,7 +58,10 @@ from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * from litellm.types.router import GenericLiteLLMParams -from litellm.utils import ProviderConfigManager, client +from litellm.utils import ( + ProviderConfigManager, + client, +) if TYPE_CHECKING: from mcp.types import Tool as MCPTool diff --git a/litellm/router.py b/litellm/router.py index e54eadfb872..6e7b9689415 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -40,7 +40,6 @@ import anyio import httpx import openai from openai import AsyncOpenAI -from pydantic import BaseModel from typing_extensions import overload import litellm @@ -62,6 +61,9 @@ from litellm.constants import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_metadata_variable_name_from_kwargs, @@ -81,8 +83,10 @@ from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.router_strategy.simple_shuffle import simple_shuffle from litellm.router_strategy.tag_based_routing import get_deployments_for_tag from litellm.router_utils.add_retry_fallback_headers import ( + _HiddenParamsHost, add_fallback_headers_to_response, add_retry_headers_to_response, + get_hidden_params_dict, ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, @@ -564,6 +568,12 @@ class Router: self._explicit_timeout = timeout # None when user did not pass timeout self.timeout = timeout or litellm.request_timeout + # Per-attempt request_timeout, independent of router_settings.timeout. + # Only stored when a router timeout is also set, since otherwise + # request_timeout already flows through self.timeout above. + self.request_timeout = ( + get_configured_request_timeout() if timeout is not None else None + ) self.stream_timeout = stream_timeout self.retry_after = retry_after @@ -2165,6 +2175,36 @@ class Router: ) setattr(fallback_item, "usage", combined_usage) + @staticmethod + def _prepare_fallback_hidden_params( + fallback_response: object, + ) -> tuple[dict[str, object], dict[str, object]]: + fallback_hidden_params = get_hidden_params_dict(fallback_response) + fallback_headers = fallback_hidden_params.get("additional_headers") + if not isinstance(fallback_headers, dict): + return fallback_hidden_params, {} + return fallback_hidden_params, cast("dict[str, object]", fallback_headers) + + @staticmethod + def _apply_fallback_hidden_params_to_item( + fallback_item: object, + prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]], + ) -> None: + if fallback_item is None or not hasattr(fallback_item, "_hidden_params"): + return + + fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params + item_hidden_params = get_hidden_params_dict(fallback_item) + item_headers = item_hidden_params.get("additional_headers") + if not isinstance(item_headers, dict): + item_headers = {} + + cast(_HiddenParamsHost, fallback_item)._hidden_params = { + **item_hidden_params, + **fallback_hidden_params, + "additional_headers": {**item_headers, **fallback_headers}, + } + async def _acompletion_streaming_iterator( self, model_response: CustomStreamWrapper, @@ -2257,12 +2297,22 @@ class Router: model_group=model_group, args=(), kwargs=initial_kwargs, + include_fallback_errors=initial_kwargs.get( + "include_fallback_errors", False + ) + is True, ) ) # If fallback returns a streaming response, iterate over it if hasattr(fallback_response, "__aiter__"): + prepared_fallback_hidden_params = ( + Router._prepare_fallback_hidden_params(fallback_response) + ) async for fallback_item in fallback_response: # type: ignore + Router._apply_fallback_hidden_params_to_item( + fallback_item, prepared_fallback_hidden_params + ) if ( fallback_item and isinstance(fallback_item, ModelResponseStream) @@ -2686,11 +2736,21 @@ class Router: model_group=model_group, args=(), kwargs=initial_kwargs, + include_fallback_errors=initial_kwargs.get( + "include_fallback_errors", False + ) + is True, ) ) if hasattr(fallback_response, "__aiter__"): + prepared_fallback_hidden_params = ( + Router._prepare_fallback_hidden_params(fallback_response) + ) async for fallback_item in fallback_response: # type: ignore + Router._apply_fallback_hidden_params_to_item( + fallback_item, prepared_fallback_hidden_params + ) if partial_usage is not None: Router._combine_responses_fallback_usage( fallback_item, partial_usage @@ -2815,7 +2875,13 @@ class Router: ) if hasattr(fallback_response, "__iter__"): + prepared_fallback_hidden_params = ( + Router._prepare_fallback_hidden_params(fallback_response) + ) for fallback_item in fallback_response: + Router._apply_fallback_hidden_params_to_item( + fallback_item, prepared_fallback_hidden_params + ) if ( fallback_item and isinstance(fallback_item, ModelResponseStream) @@ -2972,6 +3038,7 @@ class Router: **kwargs, } input_kwargs.pop("silent_model", None) + input_kwargs.pop("include_fallback_errors", None) _response = litellm.acompletion(**input_kwargs) @@ -3076,7 +3143,18 @@ class Router: - litellm_trace_id - metadata """ - kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) + # Normalise an explicit num_retries=None to the router default here (dict.get() + # only falls back when the key is absent, not when its value is None), then to 0 + # if the router default is itself None - mirroring the guard in + # async_function_with_retries, which remains the safety net for paths that bypass + # this setter. + _req_num_retries = kwargs.get("num_retries") + if _req_num_retries is not None: + kwargs["num_retries"] = _req_num_retries + else: + kwargs["num_retries"] = ( + self.num_retries if self.num_retries is not None else 0 + ) kwargs.setdefault("litellm_trace_id", str(uuid.uuid4())) model_group_alias: Optional[str] = None if self._get_model_from_alias(model=model): @@ -3322,6 +3400,7 @@ class Router: "stream_timeout", None ) # timeout set on litellm_params for this deployment or self.stream_timeout # timeout set on router + or self.request_timeout # litellm_settings.request_timeout (per-attempt) or self.default_litellm_params.get("stream_timeout", None) ) @@ -3338,7 +3417,8 @@ class Router: or data.get( "request_timeout", None ) # timeout set on litellm_params for this deployment - or self.timeout # timeout set on router + or self.request_timeout # litellm_settings.request_timeout (per-attempt) + or self.timeout # timeout set on router (router_settings.timeout) or self.default_litellm_params.get("timeout", None) ) return timeout @@ -6478,6 +6558,7 @@ class Router: model_group: Optional[str], args: tuple, kwargs: dict, + include_fallback_errors: bool = False, ): """ Common utilities for async_function_with_fallbacks @@ -6501,6 +6582,8 @@ class Router: input_kwargs["max_fallbacks"] = self.max_fallbacks if "fallback_depth" not in input_kwargs: input_kwargs["fallback_depth"] = 0 + if include_fallback_errors: + input_kwargs["include_fallback_errors"] = True # ORDER-BASED FALLBACKS: prepend higher order levels to the fallback list # Skip for error types that have their own dedicated fallback handlers @@ -6759,6 +6842,7 @@ class Router: If it fails after num_retries, fall back to another model group """ model_group: Optional[str] = kwargs.get("model") + include_fallback_errors = kwargs.get("include_fallback_errors", False) is True disable_fallbacks: Optional[bool] = kwargs.pop("disable_fallbacks", False) fallbacks: Optional[List] = kwargs.get("fallbacks", self.fallbacks) context_window_fallbacks: Optional[List] = kwargs.get( @@ -6802,6 +6886,7 @@ class Router: model_group, args, kwargs, + include_fallback_errors=include_fallback_errors, ) def _handle_mock_testing_fallbacks( @@ -6868,7 +6953,11 @@ class Router: "model_group_retry_policy", self.model_group_retry_policy ) model_group: Optional[str] = kwargs.get("model") - num_retries = kwargs.pop("num_retries") + num_retries = kwargs.pop("num_retries", None) + if num_retries is None: + # Fall back to the router setting (then 0) so the comparisons below never + # hit `None > int`, which would mask the real upstream error with a TypeError. + num_retries = self.num_retries if self.num_retries is not None else 0 ## ADD MODEL GROUP SIZE TO METADATA - used for model_group_rate_limit_error tracking _metadata: dict = kwargs.get("litellm_metadata", kwargs.get("metadata")) or {} @@ -9725,17 +9814,19 @@ class Router: # - if healthy_deployments > 1, return model group rate limit headers # - else return the model's rate limit headers """ - if ( - isinstance(response, BaseModel) - and hasattr(response, "_hidden_params") - and isinstance(response._hidden_params, dict) # type: ignore - ): - response._hidden_params.setdefault("additional_headers", {}) # type: ignore - response._hidden_params["additional_headers"][ # type: ignore - "x-litellm-model-group" - ] = model_group + if response is not None and hasattr(response, "_hidden_params"): + hidden_params = getattr(response, "_hidden_params", {}) or {} + if hasattr(hidden_params, "model_dump"): + hidden_params = hidden_params.model_dump() + if not isinstance(hidden_params, dict): + return response + response._hidden_params = hidden_params - additional_headers = response._hidden_params["additional_headers"] # type: ignore + additional_headers = hidden_params.get("additional_headers") + if not isinstance(additional_headers, dict): + additional_headers = {} + hidden_params["additional_headers"] = additional_headers + additional_headers["x-litellm-model-group"] = model_group # Lift QualityRouter routing decision into response headers for # transparency. The decision is stashed in request_kwargs.metadata diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 6b921a0db8a..0b927714ca9 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -1,44 +1,99 @@ -from typing import Any, Optional, Union +import json +from typing import Protocol, TypedDict, cast from pydantic import BaseModel -from litellm.types.utils import HiddenParams + +class FallbackErrorInfo(TypedDict): + message: str + type: str + param: str | None + code: str | None -def _add_headers_to_response(response: Any, headers: dict) -> Any: +class _HiddenParamsHost(Protocol): + _hidden_params: dict[str, object] + + +def get_hidden_params_dict(response: object) -> dict[str, object]: + hidden_params: object = cast(object, getattr(response, "_hidden_params", None)) + if isinstance(hidden_params, BaseModel): + return cast("dict[str, object]", hidden_params.model_dump()) + if isinstance(hidden_params, dict): + return cast("dict[str, object]", hidden_params) + return {} + + +def _ensure_additional_headers_dict( + hidden_params: dict[str, object], +) -> dict[str, object]: + additional_headers = hidden_params.get("additional_headers") + if isinstance(additional_headers, dict): + return cast("dict[str, object]", additional_headers) + return {} + + +def get_fallback_error_info(error: Exception) -> FallbackErrorInfo: + message = cast(object, getattr(error, "message", str(error))) + error_type = cast(object, getattr(error, "type", error.__class__.__name__)) + param = cast(object, getattr(error, "param", None)) + code = cast(object, getattr(error, "status_code", getattr(error, "code", None))) + return FallbackErrorInfo( + message=str(message), + type=str(error_type), + param=str(param) if param is not None else None, + code=str(code) if code is not None else None, + ) + + +def _coerce_error_dicts(items: list[object]) -> list[dict[str, object]]: + return [cast("dict[str, object]", item) for item in items if isinstance(item, dict)] + + +def get_fallback_errors_from_headers( + additional_headers: dict[str, object], +) -> list[dict[str, object]]: + existing_errors = additional_headers.get("x-litellm-fallback-errors") + if isinstance(existing_errors, list): + return _coerce_error_dicts(cast("list[object]", existing_errors)) + if isinstance(existing_errors, str): + try: + parsed_errors: object = cast(object, json.loads(existing_errors)) + except json.JSONDecodeError: + return [] + if isinstance(parsed_errors, list): + return _coerce_error_dicts(cast("list[object]", parsed_errors)) + return [] + + +def _add_headers_to_response(response: object, headers: dict[str, object]) -> object: """ Helper function to add headers to a response's hidden params """ - if response is None or not isinstance(response, BaseModel): + if response is None: return response - hidden_params: Optional[Union[dict, HiddenParams]] = getattr( - response, "_hidden_params", {} - ) + if not isinstance(response, BaseModel) and not hasattr(response, "_hidden_params"): + return response - if hidden_params is None: - hidden_params_dict = {} - elif isinstance(hidden_params, HiddenParams): - hidden_params_dict = hidden_params.model_dump() - else: - hidden_params_dict = hidden_params + hidden_params = get_hidden_params_dict(response) + additional_headers = _ensure_additional_headers_dict(hidden_params) + additional_headers.update(headers) + hidden_params["additional_headers"] = additional_headers - hidden_params_dict.setdefault("additional_headers", {}) - hidden_params_dict["additional_headers"].update(headers) - - setattr(response, "_hidden_params", hidden_params_dict) + cast(_HiddenParamsHost, response)._hidden_params = hidden_params return response def add_retry_headers_to_response( - response: Any, + response: object, attempted_retries: int, - max_retries: Optional[int] = None, -) -> Any: + max_retries: int | None = None, +) -> object: """ Add retry headers to the request """ - retry_headers = { + retry_headers: dict[str, object] = { "x-litellm-attempted-retries": attempted_retries, } if max_retries is not None: @@ -48,9 +103,10 @@ def add_retry_headers_to_response( def add_fallback_headers_to_response( - response: Any, + response: object, attempted_fallbacks: int, -) -> Any: + fallback_errors: list[FallbackErrorInfo] | None = None, +) -> object: """ Add fallback headers to the response @@ -64,7 +120,19 @@ def add_fallback_headers_to_response( Note: It's intentional that we don't add max_fallbacks in response headers Want to avoid bloat in the response headers for performance. """ - fallback_headers = { + fallback_headers: dict[str, object] = { "x-litellm-attempted-fallbacks": attempted_fallbacks, } - return _add_headers_to_response(response, fallback_headers) + response = _add_headers_to_response(response, fallback_headers) + if fallback_errors is None or response is None: + return response + + hidden_params = get_hidden_params_dict(response) + additional_headers = _ensure_additional_headers_dict(hidden_params) + merged_errors = get_fallback_errors_from_headers(additional_headers) + [ + cast("dict[str, object]", error) for error in fallback_errors + ] + additional_headers["x-litellm-fallback-errors"] = json.dumps(merged_errors) + hidden_params["additional_headers"] = additional_headers + cast(_HiddenParamsHost, response)._hidden_params = hidden_params + return response diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index b210ea44596..dcfa44381c1 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -38,6 +38,7 @@ class CooldownCache: visible_prefix=50, # Show first 50 characters visible_suffix=0, # Show last 0 characters mask_char="*", # Use * for masking + mask_short_values=False, # Truncate long messages only; keep short ones readable ) def _common_add_cooldown_logic( diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index eb756e3cf8b..f0edc7fc9db 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -6,6 +6,7 @@ from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, + get_fallback_error_info, ) from litellm.types.router import LiteLLMParamsTypedDict @@ -90,6 +91,7 @@ async def run_async_fallback( original_exception: Exception, max_fallbacks: int, fallback_depth: int, + include_fallback_errors: bool = False, **kwargs, ) -> Any: """ @@ -118,6 +120,7 @@ async def run_async_fallback( raise original_exception error_from_fallbacks = original_exception + fallback_errors = (get_fallback_error_info(original_exception),) for mg in fallback_model_group: if mg == original_model_group: @@ -136,6 +139,8 @@ async def run_async_fallback( fallback_depth = fallback_depth + 1 kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks + if include_fallback_errors: + kwargs["include_fallback_errors"] = include_fallback_errors response = await litellm_router.async_function_with_fallbacks( *args, **kwargs ) @@ -143,6 +148,9 @@ async def run_async_fallback( response = add_fallback_headers_to_response( response=response, attempted_fallbacks=fallback_depth, + fallback_errors=( + list(fallback_errors) if include_fallback_errors else None + ), ) # callback for successfull_fallback_event(): await log_success_fallback_event( @@ -153,6 +161,7 @@ async def run_async_fallback( return response except Exception as e: error_from_fallbacks = e + fallback_errors = fallback_errors + (get_fallback_error_info(e),) await log_failure_fallback_event( original_model_group=original_model_group, kwargs=kwargs, diff --git a/litellm/sandbox/main.py b/litellm/sandbox/main.py index a3f5f3b4665..76d9994c683 100644 --- a/litellm/sandbox/main.py +++ b/litellm/sandbox/main.py @@ -68,8 +68,9 @@ async def acreate_sandbox( provider: str, template: str | None = None, timeout: int | None = None, - allow_internet_access: bool = True, + allow_internet_access: bool | None = None, api_key: str | None = None, + api_base: str | None = None, **kwargs, ) -> ContainerHandle: _update_logging(kwargs, provider, "create_sandbox") @@ -78,6 +79,7 @@ async def acreate_sandbox( timeout=timeout, allow_internet_access=allow_internet_access, api_key=api_key, + api_base=api_base, **_forward_kwargs(kwargs), ) @@ -104,12 +106,14 @@ async def adelete_sandbox( provider: str, container: Union[ContainerHandle, str], api_key: str | None = None, + api_base: str | None = None, **kwargs, ) -> bool: _update_logging(kwargs, provider, "delete_sandbox") return await _get_config(provider).adelete_sandbox( container=container, api_key=api_key, + api_base=api_base, **_forward_kwargs(kwargs), ) @@ -121,6 +125,7 @@ async def acode_interpreter_tool( template: str | None = None, timeout: int | None = None, api_key: str | None = None, + api_base: str | None = None, **kwargs, ) -> CodeExecutionResult: _update_logging(kwargs, provider, "code_interpreter_tool") @@ -128,7 +133,11 @@ async def acode_interpreter_tool( forwarded = _forward_kwargs(kwargs) container = await config.acreate_sandbox( - template=template, timeout=timeout, api_key=api_key, **forwarded + template=template, + timeout=timeout, + api_key=api_key, + api_base=api_base, + **forwarded, ) try: return await config.arun_code( @@ -137,7 +146,7 @@ async def acode_interpreter_tool( finally: try: await config.adelete_sandbox( - container=container, api_key=api_key, **forwarded + container=container, api_key=api_key, api_base=api_base, **forwarded ) except Exception as e: litellm._logging.verbose_logger.debug( diff --git a/litellm/sandbox/sandbox_tools.py b/litellm/sandbox/sandbox_tools.py new file mode 100644 index 00000000000..f4a6678f629 --- /dev/null +++ b/litellm/sandbox/sandbox_tools.py @@ -0,0 +1,60 @@ +""" +Registry for sandbox tools configured via the proxy's top-level `sandbox_tools`. + +A sandbox tool maps a name to a sandbox provider plus its credentials, so the +code interpreter interceptor can resolve a tool by name to provider/key/base. +""" + +from collections.abc import Iterator + +from litellm._logging import verbose_logger + +_SANDBOX_TOOL_REGISTRY: dict[str, dict] = {} + + +def _resolve_secret_value(value: str | None) -> str | None: + if not isinstance(value, str): + return None + if value.startswith("os.environ/"): + from litellm.secret_managers.main import get_secret_str + + return get_secret_str(value) + return value + + +def _iter_valid_tools(tools: list[dict]) -> Iterator[tuple[str, dict]]: + for tool in tools: + if not isinstance(tool, dict): + verbose_logger.warning("sandbox_tools: skipping non-dict entry %r", tool) + continue + name = tool.get("sandbox_tool_name") + if not name: + verbose_logger.warning( + "sandbox_tools: skipping entry missing 'sandbox_tool_name': %r", tool + ) + continue + params = tool.get("litellm_params") or {} + provider = params.get("sandbox_provider") + if not provider: + verbose_logger.warning( + "sandbox_tools: skipping entry missing 'sandbox_provider': %r", tool + ) + continue + yield name, { + "sandbox_provider": provider, + "api_key": _resolve_secret_value(params.get("api_key")), + "api_base": _resolve_secret_value(params.get("api_base")), + } + + +def register_sandbox_tools(tools: list[dict]) -> None: + global _SANDBOX_TOOL_REGISTRY + _SANDBOX_TOOL_REGISTRY = dict(_iter_valid_tools(tools)) + + +def resolve_sandbox_tool(name: str) -> dict | None: + return _SANDBOX_TOOL_REGISTRY.get(name) + + +def clear_sandbox_tools() -> None: + register_sandbox_tools([]) diff --git a/litellm/types/completion.py b/litellm/types/completion.py index cb263914be8..a91f6234fad 100644 --- a/litellm/types/completion.py +++ b/litellm/types/completion.py @@ -1,8 +1,28 @@ -from typing import Iterable, List, Optional, Union +from __future__ import annotations + +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Coroutine, + Iterable, + List, + Optional, + Union, +) from pydantic import BaseModel, ConfigDict from typing_extensions import Literal, Required, TypedDict +if TYPE_CHECKING: + import httpx + from aiohttp import ClientSession + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm import BaseConfig + from litellm.utils import CustomStreamWrapper, ModelResponse + class ChatCompletionSystemMessageParam(TypedDict, total=False): content: Required[str] @@ -191,3 +211,44 @@ class CompletionRequest(BaseModel): model_list: Optional[List[str]] = None model_config = ConfigDict(protected_namespaces=(), extra="allow") + + +@dataclass(frozen=True, slots=True) +class _CompletionDispatchContext: + _azure_detection_model: str + acompletion: bool + api_base: Optional[str] + api_key: Optional[str] + api_version: Optional[str] + client: Any + custom_llm_provider: str + custom_prompt_dict: dict + extra_headers: Optional[dict] + headers: dict + hf_model_name: Optional[str] + kwargs: dict + litellm_params: dict + logger_fn: Optional[Callable] + logging: LiteLLMLoggingObj + max_retries: Optional[int] + max_tokens: Optional[int] + messages: list + metadata: Optional[dict] + model: str + model_response: ModelResponse + optional_params: dict + organization: Optional[str] + provider_config: Optional[BaseConfig] + shared_session: Optional[ClientSession] + stream: Optional[bool] + temperature: Optional[float] + text_completion: bool + timeout: Optional[Union[float, str, httpx.Timeout]] + top_p: Optional[float] + + +_CompletionDispatchResult = Union[ + Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]], + "ModelResponse", + "CustomStreamWrapper", +] diff --git a/litellm/types/integrations/code_interpreter_interception.py b/litellm/types/integrations/code_interpreter_interception.py new file mode 100644 index 00000000000..2669c59db37 --- /dev/null +++ b/litellm/types/integrations/code_interpreter_interception.py @@ -0,0 +1,22 @@ +""" +Type definitions for Code Interpreter Interception integration. +""" + +from typing import List, TypedDict + + +class CodeInterpreterInterceptionConfig(TypedDict, total=False): + """ + Configuration parameters for CodeInterpreterInterceptionLogger. + + Used in proxy_config.yaml under litellm_settings: + litellm_settings: + code_interpreter_interception_params: + enabled: true + enabled_providers: ["openai"] + sandbox_tool_name: "my_e2b_sandbox" + """ + + enabled: bool + enabled_providers: List[str] + sandbox_tool_name: str diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index b5726a11ca0..26a0be36ef4 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -2,6 +2,25 @@ from typing import Any, Dict, List, Optional from pydantic import BaseModel, Field +CHAT_COMPLETION_AGENTIC_SURFACE = "chat_completions" +CODE_INTERPRETER_INTERCEPTION_PREFIX = "_code_interpreter_interception" +NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES = frozenset( + ("_websearch_interception", "_compression_interception") +) +INTERCEPTION_INTERNAL_PREFIXES = frozenset( + ( + *NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + CODE_INTERPRETER_INTERCEPTION_PREFIX, + ) +) + + +def is_interception_internal_key( + key: str, + prefixes: frozenset[str] = INTERCEPTION_INTERNAL_PREFIXES, +) -> bool: + return any(key.startswith(prefix) for prefix in prefixes) + class StandardCustomLoggerInitParams(BaseModel): """ diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index b38cd8f58b9..a07642073af 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -954,9 +954,6 @@ class Interaction(BaseModel): None, description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) - role: Optional[str] = Field( - None, description="Output only. The role of the interaction." - ) outputs: Optional[List[Content]] = Field( None, description="Output only. Responses from the model." ) @@ -1031,9 +1028,6 @@ class CreateModelInteractionParams(BaseModel): None, description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) - role: Optional[str] = Field( - None, description="Output only. The role of the interaction." - ) outputs: Optional[List[Content]] = Field( None, description="Output only. Responses from the model." ) @@ -1101,9 +1095,6 @@ class CreateAgentInteractionParams(BaseModel): None, description="Output only. The time at which the response was last updated in ISO 8601 format\n(YYYY-MM-DDThh:mm:ssZ).", ) - role: Optional[str] = Field( - None, description="Output only. The role of the interaction." - ) outputs: Optional[List[Content]] = Field( None, description="Output only. Responses from the model." ) @@ -1323,7 +1314,6 @@ class InteractionsAPIResponse(BaseLiteLLMOpenAIResponseObject): status: Optional[str] = None created: Optional[str] = None updated: Optional[str] = None - role: Optional[str] = None # Legacy schema field (Api-Revision: 2026-05-07). Remove after June 8, 2026. outputs: Optional[List[Dict[str, Any]]] = None # New schema field (Api-Revision: 2026-05-20). @@ -1356,7 +1346,6 @@ class InteractionsAPIStreamingResponse(BaseLiteLLMOpenAIResponseObject): status: Optional[str] = None created: Optional[str] = None updated: Optional[str] = None - role: Optional[str] = None # Legacy schema field (Api-Revision: 2026-05-07). Remove after June 8, 2026. outputs: Optional[List[Dict[str, Any]]] = None # New schema field (Api-Revision: 2026-05-20). diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 21d4da82041..94e4c68f5e2 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -69,7 +69,6 @@ class MCPPublicServer(BaseModel): name: str alias: Optional[str] = None server_name: Optional[str] = None - url: Optional[str] = None transport: MCPTransportType spec_path: Optional[str] = None auth_type: Optional[MCPAuthType] = None diff --git a/litellm/types/proxy/management_endpoints/scim_v2.py b/litellm/types/proxy/management_endpoints/scim_v2.py index c5fdc66154f..6270d2d0925 100644 --- a/litellm/types/proxy/management_endpoints/scim_v2.py +++ b/litellm/types/proxy/management_endpoints/scim_v2.py @@ -1,7 +1,20 @@ from typing import Any, Dict, List, Literal, Optional, Union from fastapi import HTTPException -from pydantic import BaseModel, ConfigDict, EmailStr, field_validator +from pydantic import ( + BaseModel, + ConfigDict, + EmailStr, + Field, + field_validator, + model_serializer, +) +from pydantic_core.core_schema import SerializerFunctionWrapHandler + +SCIM_ENTERPRISE_USER_SCHEMA = ( + "urn:ietf:params:scim:schemas:extension:enterprise:2.0:User" +) +SCIM_ENTERPRISE_METADATA_KEY = "scim_enterprise" class LiteLLM_UserScimMetadata(BaseModel): @@ -42,13 +55,49 @@ class SCIMUserGroup(BaseModel): type: Optional[str] = "direct" # direct or indirect +class SCIMUserManager(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + value: Optional[str] = None + displayName: Optional[str] = None + ref: Optional[str] = Field(default=None, alias="$ref") + + +class SCIMEnterpriseUser(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + employeeNumber: Optional[str] = None + costCenter: Optional[str] = None + organization: Optional[str] = None + division: Optional[str] = None + department: Optional[str] = None + manager: Optional[SCIMUserManager] = None + + class SCIMUser(SCIMResource): + model_config = ConfigDict(populate_by_name=True) + userName: Optional[str] = None name: Optional[SCIMUserName] = None displayName: Optional[str] = None active: bool = True emails: Optional[List[SCIMUserEmail]] = None groups: Optional[List[SCIMUserGroup]] = None + enterprise_user: Optional[SCIMEnterpriseUser] = Field( + default=None, + alias=SCIM_ENTERPRISE_USER_SCHEMA, + serialization_alias=SCIM_ENTERPRISE_USER_SCHEMA, + ) + + @model_serializer(mode="wrap") + def _omit_absent_enterprise( + self, handler: SerializerFunctionWrapHandler + ) -> Dict[str, Any]: + dumped = handler(self) + if self.enterprise_user is None: + dumped.pop(SCIM_ENTERPRISE_USER_SCHEMA, None) + dumped.pop("enterprise_user", None) + return dumped class SCIMMember(BaseModel): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 32bfc8835fe..24d6e84fba7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -37,6 +37,8 @@ from pydantic import ( ConfigDict, Field, PrivateAttr, + SkipValidation, + field_serializer, field_validator, ) from typing_extensions import Required, TypedDict @@ -3146,10 +3148,41 @@ class CustomPricingLiteLLMParams(BaseModel): search_context_cost_per_query: Optional[Dict[str, Any]] = None citation_cost_per_token: Optional[float] = None tiered_pricing: Optional[List[Dict[str, Any]]] = None + cache_read_input_token_cost_above_272k_tokens: Optional[float] = None + cache_read_input_token_cost_above_512k_tokens: Optional[float] = None + input_cost_per_image_token: Optional[float] = None + input_cost_per_token_above_272k_tokens: Optional[float] = None + input_cost_per_token_above_512k_tokens: Optional[float] = None + output_cost_per_token_above_272k_tokens: Optional[float] = None + output_cost_per_token_above_512k_tokens: Optional[float] = None + output_vector_size: Optional[int] = None + ocr_cost_per_page: Optional[float] = None + ocr_cost_per_credit: Optional[float] = None + annotation_cost_per_page: Optional[float] = None + regional_processing_uplift_multiplier_eu: Optional[float] = None + regional_processing_uplift_multiplier_us: Optional[float] = None +# Server-controlled fields that bound or drive an interceptor's agentic loop +# (depth, cycle fingerprints, ceiling, code-interpreter sandbox state). Listed +# in all_litellm_params so they are treated as LiteLLM-level and excluded from +# get_non_default_completion_params; otherwise the OpenAI param builder sweeps +# any unrecognized top-level key into extra_body and leaks them to the provider. +# This is what lets the loop carry state across rerun calls without a provider +# scrubber. +agentic_loop_internal_litellm_params = [ + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_agentic_loop_api_surface", + "max_agentic_loops", + "_code_interpreter_interception_active", + "_code_interpreter_interception_sandbox_key", + "_code_interpreter_interception_converted_stream", +] + all_litellm_params = ( - [ + agentic_loop_internal_litellm_params + + [ "metadata", "litellm_metadata", "litellm_trace_id", @@ -3448,6 +3481,7 @@ class LlmProviders(str, Enum): TENSORMESH = "tensormesh" LIBERTAI = "libertai" PINSTRIPES = "pinstripes" + DARKBLOOM = "darkbloom" LITELLM_AGENT = "litellm_agent" CURSOR = "cursor" BEDROCK_MANTLE = "bedrock_mantle" @@ -3505,6 +3539,7 @@ class SandboxProviders(str, Enum): """ E2B = "e2b" + OPENSANDBOX = "opensandbox" class LiteLLMLoggingBaseClass: @@ -3610,10 +3645,20 @@ class LiteLLMBatch(Batch): class LiteLLMRealtimeStreamLoggingObject(LiteLLMPydanticObjectBase): - results: OpenAIRealtimeStreamList + # Events are already well-formed provider dicts. Validating them against the + # OpenAIRealtimeEvents union makes Pydantic try every member per event, which + # floods thousands of ValidationErrors for events outside the union (e.g. + # rate_limits.updated), blocks the event loop, and discards the session usage. + results: SkipValidation[OpenAIRealtimeStreamList] usage: Usage _hidden_params: dict = {} + @field_serializer("results") + def _serialize_results( + self, results: OpenAIRealtimeStreamList + ) -> List[Dict[str, Any]]: + return [dict(event) for event in results] + def __contains__(self, key): # Define custom behavior for the 'in' operator return hasattr(self, key) diff --git a/litellm/utils.py b/litellm/utils.py index 29f703104da..5c3ab3e1490 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3191,7 +3191,7 @@ def get_optional_params_transcription( model=model, drop_params=drop_params if drop_params is not None else False, ) - elif provider_config is not None: # handles fireworks ai, and any future providers + elif provider_config is not None: # custom audio transcription config supported_params = provider_config.get_supported_openai_params(model=model) _check_valid_arg(supported_params=supported_params) optional_params = provider_config.map_openai_params( @@ -8915,8 +8915,6 @@ class ProviderConfigManager: ) return AzureSpeechAudioTranscriptionConfig() - if litellm.LlmProviders.FIREWORKS_AI == provider: - return litellm.FireworksAIAudioTranscriptionConfig() elif litellm.LlmProviders.DEEPGRAM == provider: return litellm.DeepgramAudioTranscriptionConfig() elif litellm.LlmProviders.ELEVENLABS == provider: @@ -9733,9 +9731,14 @@ class ProviderConfigManager: Get sandbox (code execution) configuration for a given provider. """ from litellm.llms.e2b.sandbox.transformation import E2BSandboxConfig + from litellm.llms.opensandbox.sandbox.transformation import ( + OpenSandboxSandboxConfig, + ) if provider == SandboxProviders.E2B: return E2BSandboxConfig() + if provider == SandboxProviders.OPENSANDBOX: + return OpenSandboxSandboxConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 47b7190185e..3e844a8e3ed 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -571,7 +571,7 @@ "output_vector_size": 1536 }, "amazon.titan-embed-text-v2:0": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 2e-08, "litellm_provider": "bedrock", "max_input_tokens": 8192, "max_tokens": 8192, @@ -10684,6 +10684,268 @@ "mode": "chat", "output_cost_per_token": 1.923e-06 }, + "cloudflare/@cf/openai/gpt-oss-120b": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/google/gemma-2b-it-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/meta/llama-3.2-3b-instruct": { + "input_cost_per_token": 5.09e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 80000, + "max_output_tokens": 80000, + "max_tokens": 80000, + "mode": "chat", + "output_cost_per_token": 3.35e-07 + }, + "cloudflare/@cf/meta/llama-guard-3-8b": { + "input_cost_per_token": 4.84e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 3e-08 + }, + "cloudflare/@cf/mistral/mistral-7b-instruct-v0.2-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 15000, + "max_output_tokens": 15000, + "max_tokens": 15000, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/moonshotai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/deepseek-ai/deepseek-r1-distill-qwen-32b": { + "input_cost_per_token": 4.97e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 80000, + "max_output_tokens": 80000, + "max_tokens": 80000, + "mode": "chat", + "output_cost_per_token": 4.881e-06, + "supports_reasoning": true + }, + "cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8": { + "input_cost_per_token": 1.52e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "chat", + "output_cost_per_token": 2.87e-07 + }, + "cloudflare/@cf/meta/llama-3.2-1b-instruct": { + "input_cost_per_token": 2.7e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 60000, + "max_output_tokens": 60000, + "max_tokens": 60000, + "mode": "chat", + "output_cost_per_token": 2.01e-07 + }, + "cloudflare/@cf/moonshotai/kimi-k2.6": { + "cache_read_input_token_cost": 1.6e-07, + "input_cost_per_token": 9.5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/zai-org/glm-4.7-flash": { + "input_cost_per_token": 6.05e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/meta-llama/llama-2-7b-chat-hf-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast": { + "input_cost_per_token": 2.93e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 24000, + "max_output_tokens": 24000, + "max_tokens": 24000, + "mode": "chat", + "output_cost_per_token": 2.253e-06, + "supports_function_calling": true + }, + "cloudflare/@cf/ibm-granite/granite-4.0-h-micro": { + "input_cost_per_token": 1.7e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 131000, + "max_output_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 1.12e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/qwen/qwen2.5-coder-32b-instruct": { + "input_cost_per_token": 6.6e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1e-06 + }, + "cloudflare/@cf/zai-org/glm-5.2": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "cloudflare", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/nvidia/nemotron-3-120b-a12b": { + "input_cost_per_token": 5e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/aisingapore/gemma-sea-lion-v4-27b-it": { + "input_cost_per_token": 3.51e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.55e-07 + }, + "cloudflare/@cf/qwen/qwen3-30b-a3b-fp8": { + "input_cost_per_token": 5.09e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.35e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/google/gemma-7b-it-lora": { + "input_cost_per_token": 0.0, + "litellm_provider": "cloudflare", + "max_input_tokens": 3500, + "max_output_tokens": 3500, + "max_tokens": 3500, + "mode": "chat", + "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/google/gemma-4-26b-a4b-it": { + "input_cost_per_token": 1e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/mistralai/mistral-small-3.1-24b-instruct": { + "input_cost_per_token": 3.51e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.55e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/meta/llama-3.2-11b-vision-instruct": { + "input_cost_per_token": 4.85e-08, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.76e-07, + "supports_vision": true + }, + "cloudflare/@cf/openai/gpt-oss-20b": { + "input_cost_per_token": 2e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-07, + "supports_function_calling": true, + "supports_reasoning": true + }, + "cloudflare/@cf/meta/llama-4-scout-17b-16e-instruct": { + "input_cost_per_token": 2.7e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 131000, + "max_output_tokens": 131000, + "max_tokens": 131000, + "mode": "chat", + "output_cost_per_token": 8.5e-07, + "supports_function_calling": true + }, + "cloudflare/@cf/qwen/qwq-32b": { + "input_cost_per_token": 6.6e-07, + "litellm_provider": "cloudflare", + "max_input_tokens": 24000, + "max_output_tokens": 24000, + "max_tokens": 24000, + "mode": "chat", + "output_cost_per_token": 1e-06, + "supports_reasoning": true + }, "codestral/codestral-2405": { "input_cost_per_token": 0.0, "litellm_provider": "codestral", @@ -15019,7 +15281,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, @@ -15322,7 +15584,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": true }, "fireworks_ai/qwen3p7-plus": { "cache_read_input_token_cost": 8e-08, @@ -20096,8 +20358,6 @@ "output_cost_per_token": 8e-06, "output_cost_per_token_batches": 4e-06, "output_cost_per_token_priority": 1.4e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20171,8 +20431,6 @@ "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 8e-07, "output_cost_per_token_priority": 2.8e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20246,8 +20504,6 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_priority": 8e-07, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -20319,8 +20575,6 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_priority": 1.7e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20362,8 +20616,6 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20385,8 +20637,6 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_batches": 5e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -20675,8 +20925,6 @@ "output_cost_per_token": 6e-07, "output_cost_per_token_batches": 3e-07, "output_cost_per_token_priority": 1e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -21380,8 +21628,6 @@ "output_cost_per_token": 1e-05, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21775,6 +22021,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21823,6 +22071,8 @@ "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_priority": 6e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -21867,6 +22117,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21911,6 +22163,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -21959,6 +22213,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22006,6 +22262,8 @@ "output_cost_per_token_flex": 7.5e-06, "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_priority": 3e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22046,6 +22304,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22089,6 +22349,8 @@ "output_cost_per_token_above_272k_tokens": 0.00027, "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -22134,6 +22396,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22180,6 +22444,8 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22223,6 +22489,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22266,6 +22534,8 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "regional_processing_uplift_multiplier_eu": 1.10, + "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22304,8 +22574,6 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -22712,8 +22980,6 @@ "output_cost_per_token": 2e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "supported_endpoints": [ "/v1/chat/completions", "/v1/batch", @@ -22795,8 +23061,6 @@ "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, - "regional_processing_uplift_multiplier_eu": 1.10, - "regional_processing_uplift_multiplier_us": 1.10, "mode": "chat", "output_cost_per_token": 4e-07, "output_cost_per_token_flex": 2e-07, @@ -39946,24 +40210,6 @@ "litellm_provider": "fireworks_ai", "mode": "chat" }, - "fireworks_ai/accounts/fireworks/models/whisper-v3": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "fireworks_ai", - "mode": "audio_transcription" - }, - "fireworks_ai/accounts/fireworks/models/whisper-v3-turbo": { - "max_tokens": 4096, - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "input_cost_per_token": 0.0, - "output_cost_per_token": 0.0, - "litellm_provider": "fireworks_ai", - "mode": "audio_transcription" - }, "fireworks_ai/accounts/fireworks/models/yi-34b": { "max_tokens": 4096, "max_input_tokens": 4096, @@ -43543,5 +43789,39 @@ "supports_assistant_prefill": true, "supports_reasoning": false, "source": "https://pinstripes.io/pricing" + }, + "darkbloom/gemma-4-26b": { + "input_cost_per_token": 3e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.65e-07, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "darkbloom/gpt-oss-20b": { + "input_cost_per_token": 1.45e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7e-08, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true } } diff --git a/osv-scanner.toml b/osv-scanner.toml index f0f5f045f1a..7ab450945f5 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -2,13 +2,3 @@ id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-09-09 reason = "diskcache has no fixed release published; remove this entry once one exists" - -[[IgnoredVulns]] -id = "GHSA-hg6j-4rv6-33pg" -ignoreUntil = 2026-08-15 -reason = "aiohttp held at 3.13.5: vcrpy releases <= 8.1.1 cannot import aiohttp >= 3.14 and the merged upstream fix (vcrpy PR 996) is unreleased; bump aiohttp and drop this entry when a newer vcrpy ships" - -[[IgnoredVulns]] -id = "GHSA-jg22-mg44-37j8" -ignoreUntil = 2026-08-15 -reason = "aiohttp held at 3.13.5: vcrpy releases <= 8.1.1 cannot import aiohttp >= 3.14 and the merged upstream fix (vcrpy PR 996) is unreleased; bump aiohttp and drop this entry when a newer vcrpy ships" diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 7386ced3e6d..b137ec59a1f 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1833,6 +1833,23 @@ "text_completion": true } }, + "opensandbox": { + "display_name": "OpenSandbox (`opensandbox`)", + "url": "https://open-sandbox.ai/api/", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "sandbox": true + } + }, "openai_like": { "display_name": "OpenAI-like (`openai_like`)", "url": "https://docs.litellm.ai/docs/providers/openai_compatible", @@ -2008,6 +2025,23 @@ "interactions": true } }, + "darkbloom": { + "display_name": "Darkbloom (`darkbloom`)", + "url": "https://docs.litellm.ai/docs/providers/darkbloom", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "predibase": { "display_name": "Predibase (`predibase`)", "url": "https://docs.litellm.ai/docs/providers/predibase", diff --git a/pyproject.toml b/pyproject.toml index 91de8683968..1cb39153e83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.90.0" +version = "1.91.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -55,7 +55,7 @@ proxy = [ "fastapi-sso>=0.19.0,<1.0", "PyJWT>=2.13.0,<3.0", "python-multipart>=0.0.27,<1.0", - "cryptography>=46.0.7,<47.0", + "cryptography>=48.0.1,<49.0", "pynacl>=1.6.2,<2.0", "websockets>=15.0.1,<16.0", "boto3>=1.43.1,<2.0", @@ -70,6 +70,7 @@ proxy = [ "soundfile>=0.12.1,<1.0", "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", "pydantic-settings>=2.14.1,<3.0", + "expression>=5.6.0,<6.0", ] # Thin client install for the `lite` CLI on developer laptops. The CLI's heavy # imports (fastapi, cryptography, ...) are all guarded, so it runs on the base @@ -181,7 +182,7 @@ dev = [ "parameterized==0.9.0", "openapi-core==0.22.0; python_version < '3.14'", "pytest-timeout==2.4.0", - "vcrpy==8.1.1", + "vcrpy==8.2.1", "pytest-recording==0.13.4", ] proxy-dev = [ @@ -207,7 +208,7 @@ ci = [ "pytest-codspeed==4.3.0", "pytest-retry==1.7.0", "pyarrow==23.0.1", - "langchain==1.2.10", + "langchain==1.3.9", "lunary==1.4.36; python_version == '3.10'", "lunary==1.4.37; python_version >= '3.11'", "logfire==4.6.0", @@ -224,11 +225,8 @@ ci = [ "pylint==4.0.5", "langchain-mcp-adapters==0.2.1", "langchain-openai==1.1.14", - "langgraph==1.0.10", - # langgraph-prebuilt 1.0.9 imports ExecutionInfo/ServerInfo from - # langgraph.runtime, which is not exported until langgraph 1.1.0. - # Pin to 1.0.8 so it pairs correctly with langgraph==1.0.10. - "langgraph-prebuilt==1.0.8", + "langgraph>=1.2.4,<1.3.0", + "langgraph-prebuilt>=1.1.0,<1.3.0", "claude-agent-sdk==0.1.44", ] healthcheck = [ @@ -243,7 +241,7 @@ build-backend = "uv_build" [tool.uv] constraint-dependencies = [ "tornado>=6.5.6", - "aiohttp>=3.13.5,<3.14", + "aiohttp>=3.14.1,<4.0", ] default-groups = ["dev"] required-version = ">=0.10.9" @@ -272,7 +270,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.90.0" +version = "1.91.0" version_files = [ "pyproject.toml:^version", ] diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 62ebdb559fc..ae46f020de1 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -300,7 +300,7 @@ "slack": 3 }, "RET504": { - "baseline": 709, + "baseline": 702, "slack": 20 }, "RUF010": { diff --git a/scripts/ruff_strict_gate.py b/scripts/ruff_strict_gate.py index 9c406b8482b..5951a1215ed 100644 --- a/scripts/ruff_strict_gate.py +++ b/scripts/ruff_strict_gate.py @@ -5,13 +5,9 @@ Each rule has a hard ceiling (baseline + slack) in ruff-strict-budget.json. The gate counts each rule across the whole tree and fails when a rule is both over its ceiling and higher than the base it merges into, so a change is blamed for the violations it adds, never for drift that already exists in the base. - -The base is the merge-base of the current branch with --base; this matches CI, -which checks out the PR head sha and runs the gate against the PR's base sha. """ import argparse -import contextlib import json import re import shutil @@ -19,7 +15,6 @@ import subprocess import sys import tempfile from collections import Counter -from collections.abc import Iterator from pathlib import Path from typing import NamedTuple @@ -45,12 +40,6 @@ class Breach(NamedTuple): added: int -class GateInputs(NamedTuple): - head: list[Violation] - base: dict[str, int] - changed: dict[str, set[int]] - - def _run(cmd: list, cwd: Path = REPO_ROOT) -> str: proc = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True) if proc.returncode not in (0, 1): @@ -67,14 +56,14 @@ def _ruff_json(cwd: Path, config: Path) -> list: return json.loads(raw or "[]") -def collect_violations(root: Path, config: Path) -> list: +def head_violations() -> list: out = [] - for item in _ruff_json(root, config): + for item in _ruff_json(REPO_ROOT, STRICT_CONFIG): name = Path(item["filename"]) rel = ( - (name if name.is_absolute() else root / name) + (name if name.is_absolute() else REPO_ROOT / name) .resolve() - .relative_to(root) + .relative_to(REPO_ROOT) .as_posix() ) out.append(Violation(rel, item["location"]["row"], item["code"])) @@ -85,29 +74,27 @@ def count_by_rule(violations: list) -> dict: return dict(Counter(v.code for v in violations)) -@contextlib.contextmanager -def _temp_worktree(ref: str) -> Iterator[Path]: - parent = Path(tempfile.mkdtemp(prefix="ruff_wt_")) +def base_counts(ref: str) -> dict: + parent = Path(tempfile.mkdtemp(prefix="ruff_base_")) worktree = parent / "wt" try: _run(["git", "worktree", "add", "--detach", str(worktree), ref]) - yield worktree + shutil.copy(STRICT_CONFIG, worktree / "ruff-strict.toml") + items = _ruff_json(worktree, worktree / "ruff-strict.toml") + return dict(Counter(item["code"] for item in items)) finally: - subprocess.run( - ["git", "worktree", "remove", "--force", str(worktree)], - cwd=REPO_ROOT, - capture_output=True, - text=True, - ) + _run(["git", "worktree", "remove", "--force", str(worktree)]) shutil.rmtree(parent, ignore_errors=True) -def base_counts(ref: str) -> dict: - with _temp_worktree(ref) as worktree: - shutil.copy(STRICT_CONFIG, worktree / "ruff-strict.toml") - return count_by_rule( - collect_violations(worktree, worktree / "ruff-strict.toml") - ) +def evaluate(head: dict, base: dict, budget: dict) -> list: + breaches = [] + for rule, spec in budget.items(): + cap = spec["baseline"] + spec["slack"] + total = head.get(rule, 0) + if total > cap and total > base.get(rule, 0): + breaches.append(Breach(rule, total, cap, total - base.get(rule, 0))) + return sorted(breaches) def parse_changed_lines(diff_text: str) -> dict: @@ -123,31 +110,24 @@ def parse_changed_lines(diff_text: str) -> dict: return changed -def evaluate(head: dict, base: dict, budget: dict) -> list: - breaches = [] - for rule, spec in budget.items(): - cap = spec["baseline"] + spec["slack"] - total = head.get(rule, 0) - if total > cap and total > base.get(rule, 0): - breaches.append(Breach(rule, total, cap, total - base.get(rule, 0))) - return sorted(breaches) - - def introduced(violations: list, changed: dict) -> list: return [v for v in violations if v.line in changed.get(v.file, set())] -def gather(base: str) -> GateInputs: +def cmd_check(base: str) -> None: + budget = json.loads(BUDGET_PATH.read_text()) + head = head_violations() base_point = _run(["git", "merge-base", base, "HEAD"]).strip() or base - diff = _run(["git", "diff", base_point, "--unified=0", "--no-color", "--", TARGET]) - return GateInputs( - collect_violations(REPO_ROOT, STRICT_CONFIG), - base_counts(base_point), - parse_changed_lines(diff), + breaches = evaluate(count_by_rule(head), base_counts(base_point), budget) + if not breaches: + print(f"OK: every strict rule is within its codebase ceiling (base {base})") + return + new = introduced( + head, + parse_changed_lines( + _run(["git", "diff", base_point, "--unified=0", "--no-color", "--", TARGET]) + ), ) - - -def report(breaches: list, new: list, base: str) -> None: print(f"FAIL: strict-rule totals exceed their ceiling (base {base}):") for breach in breaches: print( @@ -158,24 +138,12 @@ def report(breaches: list, new: list, base: str) -> None: print( "Reduce the new violations or remove an equal number elsewhere; the ceiling is baseline + slack in ruff-strict-budget.json." ) - summary = "; ".join(f"{b.rule} {b.total}/{b.cap} (+{b.added})" for b in breaches) - print(f"BREACHED RULES: {summary}") - - -def cmd_check(base: str) -> None: - budget = json.loads(BUDGET_PATH.read_text()) - inputs = gather(base) - breaches = evaluate(count_by_rule(inputs.head), inputs.base, budget) - if not breaches: - print(f"OK: every strict rule is within its codebase ceiling (base {base})") - return - report(breaches, introduced(inputs.head, inputs.changed), base) raise SystemExit(1) def cmd_update() -> None: budget = json.loads(BUDGET_PATH.read_text()) - head = count_by_rule(collect_violations(REPO_ROOT, STRICT_CONFIG)) + head = count_by_rule(head_violations()) for rule in budget: budget[rule]["baseline"] = head.get(rule, 0) BUDGET_PATH.write_text(json.dumps(budget, indent=2, sort_keys=True) + "\n") diff --git a/scripts/type_check_gate.py b/scripts/type_check_gate.py index 0f9a44703f9..2ef332d91ea 100644 --- a/scripts/type_check_gate.py +++ b/scripts/type_check_gate.py @@ -1,21 +1,22 @@ #!/usr/bin/env python3 -"""Per-rule count gate for basedpyright. +"""Delta-vs-base per-rule gate for basedpyright. basedpyright's ``--outputjson`` is reduced to a count of errors per *rule* (``reportAny``, ``reportArgumentType``, ...) and checked against a committed budget of the form ``{rule: {baseline, slack}}``, the same shape as -``ruff-strict-budget.json``. A rule fails when its codebase-wide total exceeds -``baseline + slack``. Counts ignore file, line, and column, so a violation -moving anywhere in the tree is invisible; only the per-rule total moves the -needle. +``ruff-strict-budget.json``. A rule fails only when its codebase-wide total is +both over its ceiling (``baseline + slack``) *and* higher than the count on the +base it merges into, so a change is blamed for the errors it adds, never for +drift that already sits in the base. That ``> base`` guard is what stops an +unrelated PR from inheriting a red once two PRs each land near the ceiling and +their sum crosses it: the bystander's count equals its base, so it is spared, +while any PR that actually grows the rule past the cap still fails. -Unlike ``ruff_strict_gate.py`` this does *not* re-run the tool on the merge base -to compute a delta: a second basedpyright pass is minutes and gigabytes, whereas -ruff is milliseconds. The committed budget is the baseline instead -- exactly -how the previous per-file gate worked -- so keep it fresh with ``--update`` -(ratchet), which re-captures every rule's count from the current tree while -preserving each rule's slack. Tool output is read from stdin, so the caller -decides how to invoke basedpyright (and from which cwd). +Head counts are read from stdin (the caller runs basedpyright once and pipes +``--outputjson`` in); the base count is a second basedpyright pass over a +detached worktree at the merge-base, run under the same environment so import +resolution matches. ``--update`` re-captures the absolute per-rule baselines for +the ratchet, preserving each rule's slack. ``--outputjson`` is used rather than text diagnostics because the latter wrap across lines, leaving the ``(reportRule)`` on a continuation line away from the @@ -24,13 +25,21 @@ carries an unambiguous ``rule`` field. """ import argparse +import contextlib import json +import shutil +import subprocess import sys +import tempfile from collections import Counter +from collections.abc import Iterator, Mapping from pathlib import Path -from typing import Mapping, NamedTuple +from typing import NamedTuple REPO_ROOT = Path(__file__).resolve().parent.parent +BUDGET_PATH = REPO_ROOT / "basedpyright-code-budget.json" +PYRIGHT_CONFIG = REPO_ROOT / "pyrightconfig.json" +DEFAULT_BASE = "origin/litellm_internal_staging" # Bucket for a basedpyright diagnostic with no `rule`. Counted so it's gated. UNCODED = "" @@ -45,6 +54,7 @@ class Breach(NamedTuple): code: str total: int cap: int + added: int def _seed_slack(baseline: int) -> int: @@ -54,18 +64,19 @@ def _seed_slack(baseline: int) -> int: return 10 if baseline >= 50 else 3 -def _to_repo_relative(raw: str) -> str | None: +def _to_relative(raw: str, root: Path) -> str | None: path = Path(raw) - absolute = path if path.is_absolute() else Path.cwd() / path + absolute = path if path.is_absolute() else root / path try: - return absolute.resolve().relative_to(REPO_ROOT).as_posix() + return absolute.resolve().relative_to(root).as_posix() except ValueError: return None -def count_basedpyright(payload: str) -> dict[str, int]: - """Count in-repo basedpyright errors per rule from `--outputjson`. Warnings - and information are ignored; only `severity == "error"` is gated.""" +def count_basedpyright(payload: str, root: Path = REPO_ROOT) -> dict[str, int]: + """Count in-tree basedpyright errors per rule from `--outputjson`. Warnings + and information are ignored; only `severity == "error"` is gated. Files + outside `root` (the venv's site-packages, say) are dropped.""" try: data = json.loads(payload or "{}") except json.JSONDecodeError as exc: @@ -79,21 +90,62 @@ def count_basedpyright(payload: str) -> dict[str, int]: for diag in data.get("generalDiagnostics", []): if diag.get("severity") != "error": continue - if _to_repo_relative(diag.get("file", "")) is None: + if _to_relative(diag.get("file", ""), root) is None: continue counts[diag.get("rule") or UNCODED] += 1 return dict(counts) +def _run(cmd: list[str], cwd: Path = REPO_ROOT) -> str: + proc = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True) + if proc.returncode not in (0, 1): + sys.stderr.write(proc.stderr) + raise SystemExit(f"{cmd[0]} exited {proc.returncode}") + return proc.stdout + + +@contextlib.contextmanager +def _temp_worktree(ref: str) -> Iterator[Path]: + parent = Path(tempfile.mkdtemp(prefix="bpr_base_")) + worktree = parent / "wt" + try: + _run(["git", "worktree", "add", "--detach", str(worktree), ref]) + yield worktree + finally: + subprocess.run( + ["git", "worktree", "remove", "--force", str(worktree)], + cwd=REPO_ROOT, + capture_output=True, + text=True, + ) + shutil.rmtree(parent, ignore_errors=True) + + +def base_counts(ref: str) -> dict[str, int]: + """basedpyright error counts per rule for the merge-base tree. The head + config is copied in so the base is judged by today's rules, and the run uses + the head environment's basedpyright (on PATH) so imports resolve the same.""" + exe = shutil.which("basedpyright") or "basedpyright" + with _temp_worktree(ref) as worktree: + shutil.copy(PYRIGHT_CONFIG, worktree / "pyrightconfig.json") + proc = subprocess.run( + [exe, "--outputjson"], cwd=worktree, capture_output=True, text=True + ) + return count_basedpyright(proc.stdout, root=worktree) + + def evaluate( - counts: Mapping[str, int], budget: Mapping[str, Mapping[str, int]] + head: Mapping[str, int], + base: Mapping[str, int], + budget: Mapping[str, Mapping[str, int]], ) -> list[Breach]: breaches = [] - for code, total in counts.items(): + for code, total in head.items(): spec = budget.get(code) cap = spec["baseline"] + spec["slack"] if spec else DEFAULT_SLACK - if total > cap: - breaches.append(Breach(code, total, cap)) + prior = base.get(code, 0) + if total > cap and total > prior: + breaches.append(Breach(code, total, cap, total - prior)) return sorted(breaches) @@ -107,9 +159,6 @@ def is_vacuous_run( return not counts and any(spec["baseline"] for spec in budget.values()) -BUDGET_PATH = REPO_ROOT / "basedpyright-code-budget.json" - - def cmd_update(counts: Mapping[str, int]) -> None: existing = json.loads(BUDGET_PATH.read_text()) if BUDGET_PATH.exists() else {} budget = { @@ -127,9 +176,10 @@ def cmd_update(counts: Mapping[str, int]) -> None: ) -def cmd_check(counts: Mapping[str, int]) -> None: +def cmd_check(base_ref: str) -> None: budget = json.loads(BUDGET_PATH.read_text()) - if is_vacuous_run(counts, budget): + head = count_basedpyright(sys.stdin.read()) + if is_vacuous_run(head, budget): expected = sum(spec["baseline"] for spec in budget.values()) print( f"FAIL: basedpyright produced no errors, but {BUDGET_PATH.name} expects " @@ -137,27 +187,44 @@ def cmd_check(counts: Mapping[str, int]) -> None: f"nothing; refusing to certify a vacuous run." ) raise SystemExit(1) - breaches = evaluate(counts, budget) + base_point = _run(["git", "merge-base", base_ref, "HEAD"]).strip() or base_ref + base = base_counts(base_point) + if is_vacuous_run(base, budget): + print( + f"FAIL: basedpyright produced no errors for the base tree at " + f"{base_point[:12]}, so every rule would look freshly added. The base " + f"pass almost certainly crashed; refusing to blame this change for it." + ) + raise SystemExit(1) + breaches = evaluate(head, base, budget) if not breaches: print( - f"OK: every rule is within its basedpyright ceiling ({sum(counts.values())} errors total)" + f"OK: every rule is within its basedpyright ceiling or no higher than base ({sum(head.values())} errors total)" ) return print("FAIL: basedpyright errors exceed the per-rule ceiling:") for breach in breaches: - print(f" {breach.code}: {breach.total} errors over cap {breach.cap}") + print( + f" {breach.code}: total {breach.total} over cap {breach.cap} (this change added {breach.added})" + ) print( - "Resolve the new errors, or run 'make lint-basedpyright-budget-update' if the ceiling should move." + "Reduce the new errors or remove an equal number elsewhere; the ceiling is " + "baseline + slack in basedpyright-code-budget.json." ) + summary = "; ".join(f"{b.code} {b.total}/{b.cap} (+{b.added})" for b in breaches) + print(f"BREACHED RULES: {summary}") raise SystemExit(1) def main() -> None: parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base", default=DEFAULT_BASE) parser.add_argument("--update", action="store_true") args = parser.parse_args() - counts = count_basedpyright(sys.stdin.read()) - cmd_update(counts) if args.update else cmd_check(counts) + if args.update: + cmd_update(count_basedpyright(sys.stdin.read())) + else: + cmd_check(args.base) if __name__ == "__main__": diff --git a/tests/batches_tests/test_batch_custom_pricing.py b/tests/batches_tests/test_batch_custom_pricing.py index cb2ca385ffc..3dc1d116e8d 100644 --- a/tests/batches_tests/test_batch_custom_pricing.py +++ b/tests/batches_tests/test_batch_custom_pricing.py @@ -159,12 +159,12 @@ def test_batch_cost_calculator_applies_data_residency_uplift( base_prompt, base_completion = batch_cost_calculator( usage=usage, - model="gpt-5", + model="gpt-5.4", custom_llm_provider="openai", ) regional_prompt, regional_completion = batch_cost_calculator( usage=usage, - model="gpt-5", + model="gpt-5.4", custom_llm_provider="openai", data_residency=data_residency, ) diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 254d700ee5a..1d11d676207 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -47,6 +47,7 @@ IGNORE_FUNCTIONS = [ "_read_image_bytes", # max depth set. "_get_masked_values", # max depth set (default 20) to prevent infinite recursion while masking nested sensitive config dicts. "_redact_sensitive_litellm_params", # max depth set (default 10). + "_redact_secret_values_in_obj", # max depth set (default 10, _REDACT_SECRET_MAX_DEPTH); fails closed by returning "REDACTED" at the cap. "_resolve", # OCI: $ref resolver bounded by `resolving_stack` cycle guard. "resolve_oci_schema_anyof", # OCI: bounded by JSON-schema tree depth (no cycles possible in well-formed input). "sanitize_oci_schema", # OCI: bounded by JSON-schema tree depth. diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py b/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py index a45df5df008..fd674afd85a 100644 --- a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py +++ b/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py @@ -4,6 +4,8 @@ import pytest from fastapi import Request from litellm_enterprise.proxy.auth.user_api_key_auth import enterprise_custom_auth +from litellm.proxy._types import UserAPIKeyAuth + @pytest.mark.asyncio async def test_enterprise_custom_auth_none_user_auth(): @@ -49,16 +51,19 @@ async def test_enterprise_custom_auth_returns_string(): mock_user_auth = AsyncMock(return_value="sk-test-key") request = MagicMock(spec=Request) - with patch( - "litellm.proxy.auth.user_api_key_auth.enterprise_custom_auth", mock_user_auth - ), patch("litellm.proxy.proxy_server.master_key", "sk-1234"), patch( - "litellm.proxy.proxy_server.prisma_client", MagicMock() + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.enterprise_custom_auth", + mock_user_auth, + ), + patch("litellm.proxy.proxy_server.master_key", "sk-1234"), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), ): # Verify the key is correctly handled in _user_api_key_auth_builder with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object" + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key" ) as mock_get_key_object: - mock_get_key_object.return_value = MagicMock( + mock_get_key_object.return_value = UserAPIKeyAuth( token="sk-test-key", user_role="internal_user", team_id=None, @@ -82,9 +87,7 @@ async def test_enterprise_custom_auth_returns_string(): except Exception as e: print("error:", e) - # Verify get_key_object was called with the correct key + # Verify the key lookup was called with the correct hashed key mock_get_key_object.assert_called_once() - # The key should be hashed before being passed to get_key_object - assert mock_get_key_object.call_args[1]["hashed_token"] == hash_token( - "sk-test-key" - ) + # The key should be hashed before being passed to the resolver + assert mock_get_key_object.call_args[0][0] == hash_token("sk-test-key") diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index d64633413a0..697c3837602 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1588,18 +1588,21 @@ def test_token_counter_with_image_url_with_detail_high(): assert _tokens == DEFAULT_IMAGE_TOKEN_COUNT + 7 -def test_fireworks_ai_document_inlining(): +def test_fireworks_ai_vision_capability_from_cost_map(monkeypatch): """ - With document inlining, all fireworks ai models are now: - - supports_pdf - - supports_vision + Fireworks deprecated document inlining on 2025-06-30, so vision/PDF support is + no longer hardcoded to True for every Fireworks model. Capabilities are read + from the model cost map: unmapped models no longer advertise vision or PDF + support, while mapped VLMs still do. """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) from litellm.utils import supports_pdf_input, supports_vision - litellm._turn_on_debug() + assert supports_vision("fireworks_ai/llama-3.1-8b-instruct") is False + assert supports_pdf_input("fireworks_ai/llama-3.1-8b-instruct") is False - assert supports_pdf_input("fireworks_ai/llama-3.1-8b-instruct") is True - assert supports_vision("fireworks_ai/llama-3.1-8b-instruct") is True + assert supports_vision("fireworks_ai/minimax-m3") is True def test_logprobs_type(): diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index fa22ff6b392..f4c307e9c8a 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -2502,19 +2502,34 @@ async def test_bedrock_image_url_sync_client(): mock_post.assert_called_once() -def test_bedrock_error_handling_streaming(): +@pytest.mark.parametrize( + "exception_type, expected_status_code", + [ + ("internalServerException", 500), + ("serviceUnavailableException", 503), + ("modelTimeoutException", 408), + ("modelStreamErrorException", 424), + ("validationException", 400), + ], +) +def test_bedrock_error_handling_streaming(exception_type, expected_status_code): + """Bedrock event-stream error events arrive with botocore's hard-coded + status_code=400; the decoder must surface the modeled HTTP status instead + (e.g. internalServerException -> 500). For 5xx this is what makes the error + retryable downstream; for all types it replaces the misleading 400 with the + true code. Regression for #24608.""" from litellm.llms.bedrock.chat.invoke_handler import ( AWSEventStreamDecoder, BedrockError, ) - from unittest.mock import patch, Mock + from unittest.mock import Mock event = Mock() event.to_response_dict = Mock( return_value={ "status_code": 400, "headers": { - ":exception-type": "serviceUnavailableException", + ":exception-type": exception_type, ":content-type": "application/json", ":message-type": "exception", }, @@ -2525,11 +2540,10 @@ def test_bedrock_error_handling_streaming(): decoder = AWSEventStreamDecoder( model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0" ) - with pytest.raises(Exception) as e: + with pytest.raises(BedrockError) as e: decoder._parse_message_from_event(event) - assert isinstance(e.value, BedrockError) assert "Bedrock is unable to process your request." in e.value.message - assert e.value.status_code == 400 + assert e.value.status_code == expected_status_code @pytest.mark.parametrize( diff --git a/tests/llm_translation/test_bedrock_embedding_pricing.py b/tests/llm_translation/test_bedrock_embedding_pricing.py new file mode 100644 index 00000000000..099d73fed87 --- /dev/null +++ b/tests/llm_translation/test_bedrock_embedding_pricing.py @@ -0,0 +1,34 @@ +""" +Tests for AWS Bedrock embedding model pricing in the model cost map. + +Regression test for the Amazon Titan Text Embeddings V2 commercial price, +which was previously set 10x too high (2e-07 instead of 2e-08). +AWS lists Titan Text Embeddings V2 at $0.02 per 1M input tokens +(= $0.00002 per 1K tokens = 2e-08 per token). +""" + +import importlib + + +class TestBedrockEmbeddingPricing: + """Test suite for Bedrock embedding model pricing in the cost map.""" + + def test_titan_embed_v2_commercial_input_cost(self, monkeypatch): + """Titan Text Embeddings V2 should be priced at $0.02 / 1M tokens (2e-08).""" + # Scope the local-cost-map flag to this test only, so it does not leak + # into sibling tests. monkeypatch restores the environment on teardown. + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + import litellm.litellm_core_utils.get_model_cost_map + import litellm + + # Reload so the cost map is re-read from the local file with the flag set. + importlib.reload(litellm.litellm_core_utils.get_model_cost_map) + importlib.reload(litellm) + + model = litellm.model_cost["amazon.titan-embed-text-v2:0"] + + assert model["input_cost_per_token"] == 2e-08 + assert model["output_cost_per_token"] == 0.0 + assert model["litellm_provider"] == "bedrock" + assert model["mode"] == "embedding" diff --git a/tests/llm_translation/test_cloudflare.py b/tests/llm_translation/test_cloudflare.py index 5a6a0008398..54c5d9e4e07 100644 --- a/tests/llm_translation/test_cloudflare.py +++ b/tests/llm_translation/test_cloudflare.py @@ -9,9 +9,7 @@ import pytest from litellm import acompletion, completion from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -FAKE_API_BASE = ( - "https://fake-cloudflare.example.com/client/v4/accounts/fake-acct/ai/run/" -) +FAKE_API_BASE = "https://fake-cloudflare.example.com/client/v4/accounts/fake-acct/ai/v1" FAKE_API_KEY = "fake-cf-api-key" @@ -26,28 +24,78 @@ def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock: def _chat_response() -> Dict[str, Any]: return { - "result": { - "response": "I am a large language model created to assist you.", - }, - "success": True, - "errors": [], - "messages": [], + "id": "chatcmpl-cf", + "object": "chat.completion", + "created": 1234567890, + "model": "@cf/meta/llama-2-7b-chat-int8", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "I am a large language model created to assist you.", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 8, "completion_tokens": 11, "total_tokens": 19}, + } + + +def _tool_call_response() -> Dict[str, Any]: + return { + "id": "chatcmpl-cf-tools", + "object": "chat.completion", + "created": 1234567890, + "model": "@cf/meta/llama-2-7b-chat-int8", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "New York"}', + }, + } + ], + }, + "finish_reason": "tool_calls", + } + ], + "usage": {"prompt_tokens": 20, "completion_tokens": 9, "total_tokens": 29}, } def _streaming_chunks() -> list[str]: + base = { + "id": "chatcmpl-cf", + "object": "chat.completion.chunk", + "created": 1234567890, + "model": "@cf/meta/llama-2-7b-chat-int8", + } return [ - json.dumps({"response": "I am"}), - json.dumps({"response": " a language"}), - json.dumps({"response": " model."}), - ] - - -def _streaming_chunks_response_text() -> list[str]: - return [ - json.dumps({"response_text": "I am"}), - json.dumps({"response_text": " a language"}), - json.dumps({"response_text": " model."}), + json.dumps({**base, "choices": [{"index": 0, "delta": {"content": "I am"}}]}), + json.dumps( + {**base, "choices": [{"index": 0, "delta": {"content": " a language"}}]} + ), + json.dumps( + { + **base, + "choices": [ + { + "index": 0, + "delta": {"content": " model."}, + "finish_reason": "stop", + } + ], + } + ), ] @@ -85,6 +133,48 @@ def test_completion_cloudflare(sync_mode): assert response.choices[0].message.content is not None assert "language model" in response.choices[0].message.content.lower() + called_url = mock_post.call_args.kwargs.get("url") or mock_post.call_args.args[0] + assert called_url.endswith("/ai/v1/chat/completions") + assert "/ai/run/" not in called_url + + +def test_completion_cloudflare_tool_calls_sent_to_openai_endpoint(): + messages = [{"role": "user", "content": "weather in New York?"}] + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ] + mock_resp = _make_mock_response(_tool_call_response()) + + with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: + response = completion( + model="cloudflare/@cf/meta/llama-2-7b-chat-int8", + messages=messages, + tools=tools, + tool_choice="auto", + api_base=FAKE_API_BASE, + api_key=FAKE_API_KEY, + ) + mock_post.assert_called_once() + + sent_body = json.loads(mock_post.call_args.kwargs["data"]) + assert sent_body["tools"] == tools + assert sent_body["tool_choice"] == "auto" + + assert response.choices[0].finish_reason == "tool_calls" + tool_calls = response.choices[0].message.tool_calls + assert tool_calls is not None and len(tool_calls) == 1 + assert tool_calls[0].function.name == "get_weather" + @pytest.mark.parametrize("sync_mode", [True, False]) def test_completion_cloudflare_stream(sync_mode): @@ -153,76 +243,3 @@ def test_completion_cloudflare_stream(sync_mode): if c.choices[0].delta.content ) assert "language" in content.lower() - - -@pytest.mark.parametrize("sync_mode", [True, False]) -def test_completion_cloudflare_stream_response_text(sync_mode): - """Newer Cloudflare Workers AI models (e.g. Nemotron) emit `response_text` - instead of `response` in streamed chunks. The iterator must surface that - text so streaming output is not silently empty. - """ - messages = [{"role": "user", "content": "what llm are you"}] - raw_chunks = _streaming_chunks_response_text() - - if sync_mode: - - def _iter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.iter_lines.return_value = _iter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - - with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: - response = completion( - model="cloudflare/@cf/nvidia/nemotron-mini-4b-instruct", - messages=messages, - max_tokens=15, - stream=True, - api_base=FAKE_API_BASE, - api_key=FAKE_API_KEY, - ) - chunks_received = list(response) - mock_post.assert_called_once() - else: - - async def _aiter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.aiter_lines.return_value = _aiter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - - async def _run(): - with patch.object( - AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp - ) as mock_post: - resp = await acompletion( - model="cloudflare/@cf/nvidia/nemotron-mini-4b-instruct", - messages=messages, - max_tokens=15, - stream=True, - api_base=FAKE_API_BASE, - api_key=FAKE_API_KEY, - ) - received = [] - async for chunk in resp: - received.append(chunk) - mock_post.assert_called_once() - return received - - chunks_received = asyncio.run(_run()) - - assert len(chunks_received) > 0 - content = "".join( - c.choices[0].delta.content - for c in chunks_received - if c.choices[0].delta.content - ) - assert "language" in content.lower() diff --git a/tests/llm_translation/test_fireworks_ai_translation.py b/tests/llm_translation/test_fireworks_ai_translation.py index 4e5bef16b8c..27059581e4d 100644 --- a/tests/llm_translation/test_fireworks_ai_translation.py +++ b/tests/llm_translation/test_fireworks_ai_translation.py @@ -7,10 +7,10 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import litellm -from litellm import transcription +from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, +) from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig -from base_llm_unit_tests import BaseLLMChatTest -from base_audio_transcription_unit_tests import BaseLLMAudioTranscriptionTest fireworks = FireworksAIConfig() @@ -70,74 +70,16 @@ def test_map_response_format(): assert result == {"response_format": response_format} -_AUDIO_FILE_PATH = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "gettysburg.wav" -) - - -class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest): - def get_base_audio_transcription_call_args(self) -> dict: - return { - "model": "fireworks_ai/whisper-v3", - "api_base": "https://audio-prod.api.fireworks.ai/v1", - } - - def get_custom_llm_provider(self) -> litellm.LlmProviders: - return litellm.LlmProviders.FIREWORKS_AI - - def test_audio_transcription(self): - from unittest.mock import MagicMock - - from openai.types.audio import Transcription - - audio_file = open(_AUDIO_FILE_PATH, "rb") - mock_client = MagicMock() - mock_client.audio.transcriptions.create.return_value = Transcription( - text="four score and seven years ago" - ) - - transcript = transcription( - **self.get_base_audio_transcription_call_args(), - file=audio_file, - api_key="fw-test-key", - client=mock_client, - ) - - assert transcript.text == "four score and seven years ago" - sent = mock_client.audio.transcriptions.create.call_args.kwargs - assert sent["model"] == "whisper-v3" - assert sent["file"] is audio_file - - @pytest.mark.asyncio - async def test_audio_transcription_async(self): - from unittest.mock import AsyncMock, MagicMock - - from openai.types.audio import Transcription - - audio_file = open(_AUDIO_FILE_PATH, "rb") - raw_response = MagicMock() - raw_response.headers = {} - raw_response.parse.return_value = Transcription( - text="four score and seven years ago" - ) - mock_client = MagicMock() - mock_client.audio.transcriptions.with_raw_response.create = AsyncMock( - return_value=raw_response - ) - - transcript = await litellm.atranscription( - **self.get_base_audio_transcription_call_args(), - file=audio_file, - api_key="fw-test-key", - client=mock_client, - ) - - assert transcript.text == "four score and seven years ago" - sent = ( - mock_client.audio.transcriptions.with_raw_response.create.call_args.kwargs - ) - assert sent["model"] == "whisper-v3" - assert sent["file"] is audio_file +def test_get_supported_openai_params_transcription_returns_none(): + # Fireworks AI deprecated audio transcription on 2026-06-10; the endpoint + # is decommissioned. Returning None (not chat-completion params) signals + # to callers that transcription is unsupported for this provider. + result = get_supported_openai_params( + model="fireworks_ai/accounts/fireworks/models/whisper-v3", + custom_llm_provider="fireworks_ai", + request_type="transcription", + ) + assert result is None @pytest.mark.parametrize( @@ -146,11 +88,8 @@ class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest): ) def test_document_inlining_example(disable_add_transform_inline_image_block): """ - Document inlining appends ``#transform=inline`` to image/PDF URLs in the - outgoing request unless explicitly disabled. Assert the transform on the - serialized payload rather than making a live Fireworks call — the live - call only proved the model responded and broke whenever Fireworks rotated - its serverless model catalog. + Fireworks document inlining has been removed from the platform. LiteLLM + must not append ``#transform=inline`` regardless of the legacy disable flag. """ from unittest.mock import patch @@ -163,7 +102,7 @@ def test_document_inlining_example(disable_add_transform_inline_image_block): with patch.object(client, "post") as mock_post: try: completion( - model="fireworks_ai/accounts/fireworks/models/deepseek-v3p1", + model="fireworks_ai/accounts/fireworks/models/minimax-m3", messages=[ { "role": "user", @@ -182,89 +121,80 @@ def test_document_inlining_example(disable_add_transform_inline_image_block): disable_add_transform_inline_image_block=disable_add_transform_inline_image_block, client=client, ) - except Exception as e: - print(e) + except Exception: + pass mock_post.assert_called_once() json_data = json.loads(mock_post.call_args.kwargs["data"]) sent_url = json_data["messages"][0]["content"][0]["image_url"]["url"] - if disable_add_transform_inline_image_block is True: - assert sent_url == pdf_url - assert "#transform=inline" not in sent_url - else: - assert sent_url == pdf_url + "#transform=inline" + assert sent_url == pdf_url + assert "#transform=inline" not in sent_url @pytest.mark.parametrize( - "content, model, expected_url", + "content, expected_url", [ ( {"image_url": "http://example.com/image.png"}, - "gpt-4", - "http://example.com/image.png#transform=inline", + "http://example.com/image.png", ), ( {"image_url": {"url": "http://example.com/image.png"}}, - "gpt-4", - {"url": "http://example.com/image.png#transform=inline"}, + {"url": "http://example.com/image.png"}, ), - ( - {"image_url": "http://example.com/image.png"}, - "vision-gpt", - "http://example.com/image.png", - ), - # data: URLs must never have #transform=inline appended — doing so - # corrupts the base64 payload (fixes #23583). - # URI schemes are case-insensitive (RFC 3986) so check all variants. ( {"image_url": "data:image/png;base64,iVBORw0KGgo="}, - "gpt-4", "data:image/png;base64,iVBORw0KGgo=", ), ( {"image_url": {"url": "data:image/jpeg;base64,/9j/4AAQ=="}}, - "gpt-4", {"url": "data:image/jpeg;base64,/9j/4AAQ=="}, ), ( {"image_url": "Data:image/png;base64,iVBORw0KGgo="}, - "gpt-4", "Data:image/png;base64,iVBORw0KGgo=", ), ], ) -def test_transform_inline(content, model, expected_url): +def test_transform_inline_no_longer_added(content, expected_url): + image_block = {"type": "image_url", **content} + messages = [{"role": "user", "content": [image_block]}] - result = litellm.FireworksAIConfig()._add_transform_inline_image_block( - content=content, model=model, disable_add_transform_inline_image_block=False + result = litellm.FireworksAIConfig()._transform_messages_helper( + messages=messages, + model="accounts/fireworks/models/minimax-m3", + litellm_params={}, ) + result_image_block = result[0]["content"][0] if isinstance(expected_url, str): - assert result["image_url"] == expected_url + assert result_image_block["image_url"] == expected_url else: - assert result["image_url"]["url"] == expected_url["url"] + assert result_image_block["image_url"]["url"] == expected_url["url"] @pytest.mark.parametrize( - "model, is_disabled, expected_url", - [ - ("gpt-4", True, "http://example.com/image.png"), - ("vision-gpt", False, "http://example.com/image.png"), - ("gpt-4", False, "http://example.com/image.png#transform=inline"), - ], + "is_disabled", + [True, False], ) -def test_global_disable_flag(model, is_disabled, expected_url): - content = {"image_url": "http://example.com/image.png"} - result = litellm.FireworksAIConfig()._add_transform_inline_image_block( - content=content, - model=model, - disable_add_transform_inline_image_block=is_disabled, +def test_global_disable_flag_no_longer_adds_transform_inline(is_disabled): + url = "http://example.com/image.png" + litellm.disable_add_transform_inline_image_block = is_disabled + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": url}], + } + ] + result = litellm.FireworksAIConfig()._transform_messages_helper( + messages=messages, + model="accounts/fireworks/models/minimax-m3", + litellm_params={}, ) - assert result["image_url"] == expected_url + assert result[0]["content"][0]["image_url"] == url litellm.disable_add_transform_inline_image_block = False # Reset for other tests def test_global_disable_flag_with_transform_messages_helper(monkeypatch): - from openai import OpenAI from unittest.mock import patch from litellm import completion from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -279,7 +209,7 @@ def test_global_disable_flag_with_transform_messages_helper(monkeypatch): ) as mock_post: try: completion( - model="fireworks_ai/accounts/fireworks/models/deepseek-v3p1", + model="fireworks_ai/accounts/fireworks/models/minimax-m3", messages=[ { "role": "user", @@ -296,11 +226,10 @@ def test_global_disable_flag_with_transform_messages_helper(monkeypatch): ], client=client, ) - except Exception as e: - print(e) + except Exception: + pass mock_post.assert_called_once() - print(mock_post.call_args.kwargs) json_data = json.loads(mock_post.call_args.kwargs["data"]) assert ( "#transform=inline" diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 36e47e3c2f4..05a58a135d2 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -605,8 +605,33 @@ def test_no_messages_yields_user_text(): assert contents == expected_output -def test_convert_url(): - convert_url_to_base64("https://picsum.photos/id/237/200/300") +def test_convert_url(monkeypatch): + import base64 + from unittest.mock import MagicMock + + import httpx + + from litellm.litellm_core_utils.prompt_templates.image_handling import ( + in_memory_cache, + ) + + url = "https://picsum.photos/id/237/200/300" + image_bytes = b"\x89PNG\r\n\x1a\nfake-png-bytes" + + mock_client = MagicMock() + mock_client.get.return_value = httpx.Response( + 200, content=image_bytes, headers={"Content-Type": "image/png"} + ) + + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + monkeypatch.setattr(litellm, "module_level_client", mock_client, raising=False) + in_memory_cache.flush_cache() + + result = convert_url_to_base64(url) + + expected = "data:image/png;base64," + base64.b64encode(image_bytes).decode("utf-8") + assert result == expected + mock_client.get.assert_called_once() def test_azure_tool_call_invoke_helper(): diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 61c24183964..65d075b7f99 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -60,7 +60,8 @@ async def test_jwt_to_virtual_key_mapping_resolution(): # Use patch to mock get_key_object in the module where it's used with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ) as mock_get_key: mock_get_key.return_value = mock_key_obj @@ -105,7 +106,8 @@ async def test_jwt_to_virtual_key_mapping_no_mapping(): # Mock get_key_object just in case with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): user_api_key_cache = DualCache() @@ -481,7 +483,8 @@ async def test_reject_behavior_raises_403_on_no_mapping(): user_api_key_cache = DualCache() with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): with pytest.raises(HTTPException) as exc_info: await _resolve_jwt_to_virtual_key( @@ -519,7 +522,8 @@ async def test_reject_behavior_caches_sentinel_after_db_miss(): user_api_key_cache = DualCache() with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): # First call — DB miss, should raise 403 and write sentinel with pytest.raises(HTTPException) as exc_info: @@ -578,7 +582,8 @@ async def test_reject_behavior_raises_403_on_cached_no_mapping(): await user_api_key_cache.async_set_cache(cache_key, "__NO_MAPPING__") with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): with pytest.raises(HTTPException) as exc_info: await _resolve_jwt_to_virtual_key( @@ -671,7 +676,7 @@ async def test_auto_register_creates_key_and_mapping_when_helper_invoked(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( @@ -803,7 +808,7 @@ async def test_auto_register_race_condition_unique_conflict(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( @@ -833,13 +838,7 @@ async def test_auto_register_race_condition_unique_conflict(): # Cache should hold the winner's token, not the loser's cached = await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:user-42") assert cached == "winner_token_hash" - mock_get_key.assert_called_once_with( - hashed_token="winner_token_hash", - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) + mock_get_key.assert_called_once_with("winner_token_hash") # ────────────────────────────────────────────── @@ -1093,7 +1092,7 @@ async def test_auto_register_race_conflict_tolerates_delete_failure(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( @@ -1249,7 +1248,7 @@ async def test_auto_register_helper_stamps_validated_identity_context(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( diff --git a/tests/search_tests/test_searchapi_search.py b/tests/search_tests/test_searchapi_search.py index 5ef9d922b89..d16868502a4 100644 --- a/tests/search_tests/test_searchapi_search.py +++ b/tests/search_tests/test_searchapi_search.py @@ -46,10 +46,9 @@ class TestSearchAPIConfig: assert result["Content-Type"] == "application/json" - @patch("litellm.llms.searchapi.search.transformation.get_secret_str") - def test_validate_environment_without_api_key(self, mock_get_secret): + def test_validate_environment_without_api_key(self, monkeypatch): """Test environment validation without API key raises error.""" - mock_get_secret.return_value = None + monkeypatch.delenv("SEARCHAPI_API_KEY", raising=False) config = SearchAPIConfig() headers = {} diff --git a/tests/search_tests/test_searxng_search.py b/tests/search_tests/test_searxng_search.py index 45b0f3214d9..c12d44183b0 100644 --- a/tests/search_tests/test_searxng_search.py +++ b/tests/search_tests/test_searxng_search.py @@ -318,13 +318,11 @@ class TestSearXNGSearchHeaders: assert headers["Content-Type"] == "application/json" assert headers["Authorization"] == "Bearer test-key-123" - def test_headers_with_env_api_key(self): + def test_headers_with_env_api_key(self, monkeypatch): """Test that headers use SEARXNG_API_KEY from env.""" - with patch( - "litellm.llms.searxng.search.transformation.get_secret_str", - return_value="env-key-456", - ): - headers = self.config.validate_environment(headers={}) + monkeypatch.setenv("SEARXNG_API_KEY", "env-key-456") + + headers = self.config.validate_environment(headers={}) assert headers["Authorization"] == "Bearer env-key-456" diff --git a/tests/test_litellm/integrations/code_interpreter_interception/__init__.py b/tests/test_litellm/integrations/code_interpreter_interception/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py new file mode 100644 index 00000000000..7ff58ba6324 --- /dev/null +++ b/tests/test_litellm/integrations/code_interpreter_interception/test_handler.py @@ -0,0 +1,1219 @@ +""" +Unit tests for CodeInterpreterInterceptionLogger. + +All sandbox dependencies are injected: a FakeSandbox stands in for the real e2b +config and records how it is called. +""" + +import time + +import pytest + +from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, + LITELLM_CODE_EXECUTION_TOOL_NAME, + _INTERCEPTION_ACTIVE_KEY as _ACTIVE_KEY, + _SANDBOX_KEY, +) +from litellm.types.integrations.custom_logger import ( + CHAT_COMPLETION_AGENTIC_SURFACE, + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + is_interception_internal_key, +) +from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult +from litellm.types.utils import CallTypes + + +class FakeHandle: + def __init__(self, sandbox_id="sbx_fake"): + self.id = sandbox_id + + +class FakeSandbox: + """Records acreate_sandbox / arun_code / adelete_sandbox calls.""" + + def __init__(self, stdout="42"): + self.stdout = stdout + self.create_calls = [] + self.run_calls = [] + self.delete_calls = [] + + async def acreate_sandbox(self, **kwargs): + self.create_calls.append(kwargs) + return FakeHandle() + + async def arun_code(self, *, container, code, **kwargs): + self.run_calls.append({"container": container, "code": code}) + return CodeExecutionResult(stdout=self.stdout) + + async def adelete_sandbox(self, *, container, **kwargs): + self.delete_calls.append({"container": container}) + return True + + +class FakeLogging: + def __init__(self, litellm_call_id="k1"): + self.litellm_call_id = litellm_call_id + self.model_call_details = {} + self.dynamic_success_callbacks = [] + + def pre_call(self, *args, **kwargs): + return None + + def post_call(self, *args, **kwargs): + return None + + +def _function_call_item(call_id="c1", name=LITELLM_CODE_EXECUTION_TOOL_NAME): + return { + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments": '{"code":"print(40 + 2)"}', + } + + +def _chat_function_call_item(call_id="call_1", name=LITELLM_CODE_EXECUTION_TOOL_NAME): + return { + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": '{"code":"print(40 + 2)"}', + }, + } + + +class FakeResponse: + def __init__(self, output): + self.output = output + + +def _iter_messages(plan): + patch = plan.request_patch + assert patch is not None, "plan.request_patch must be set" + assert patch.messages is not None, "plan.request_patch.messages must be set" + return patch.messages + + +def test_interception_internal_key_prefix_sets_preserve_code_interpreter_state(): + assert is_interception_internal_key("_code_interpreter_interception_active") + assert not is_interception_internal_key( + "_code_interpreter_interception_active", + prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + ) + assert is_interception_internal_key( + "_websearch_interception_converted_stream", + prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + ) + + +@pytest.mark.asyncio +async def test_build_plan_runs_code_and_feeds_output_back(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + response = FakeResponse(output=[_function_call_item()]) + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={"litellm_call_id": "k1", _SANDBOX_KEY: "sbxkey1"}, + ) + + assert sandbox.run_calls, "sandbox.arun_code must be invoked" + assert sandbox.run_calls[0]["code"] == "print(40 + 2)" + + messages = _iter_messages(plan) + outputs = [ + m + for m in messages + if isinstance(m, dict) and m.get("type") == "function_call_output" + ] + assert outputs, "expected a function_call_output item appended" + output_item = next(m for m in outputs if m.get("call_id") == "c1") + assert "42" in str(output_item["output"]) + + +@pytest.mark.asyncio +async def test_pre_call_converts_code_interpreter_tool(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert result is not None + tools = result["tools"] + assert not any( + t.get("type") == "code_interpreter" for t in tools + ), "code_interpreter tool must be removed" + names = [t.get("name") or (t.get("function") or {}).get("name") for t in tools] + assert LITELLM_CODE_EXECUTION_TOOL_NAME in names + + +@pytest.mark.asyncio +async def test_pre_call_converts_code_interpreter_tool_for_chat_completions(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": {"type": "code_interpreter"}, + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is not None + tool = result["tools"][0] + assert tool["type"] == "function" + assert tool["function"]["name"] == LITELLM_CODE_EXECUTION_TOOL_NAME + assert tool["function"]["parameters"]["required"] == ["code"] + assert result["tool_choice"] == { + "type": "function", + "function": {"name": LITELLM_CODE_EXECUTION_TOOL_NAME}, + } + assert result["litellm_metadata"][_ACTIVE_KEY] is True + assert result["litellm_metadata"][_SANDBOX_KEY] == result[_SANDBOX_KEY] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "code_interpreter"}, + {"type": "hosted_tool", "name": "code_interpreter"}, + ], +) +async def test_pre_call_rewrites_forced_code_interpreter_tool_choice(tool_choice): + """A forced tool_choice targeting the native code_interpreter tool must be + rewritten to the generated function tool; otherwise the outbound request + references a tool that no longer exists and the provider rejects it.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": tool_choice, + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert result is not None + assert result["tool_choice"] == { + "type": "function", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + } + + +@pytest.mark.asyncio +async def test_pre_call_leaves_unrelated_tool_choice_untouched(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": "auto", + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert result is not None + assert result["tool_choice"] == "auto" + + +@pytest.mark.asyncio +async def test_pre_call_noop_on_non_responses(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aembedding) + + assert result is None + + +@pytest.mark.asyncio +async def test_pre_call_noop_on_chat_completion_without_code_interpreter_tool(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "web_search"}], + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is None + assert _ACTIVE_KEY not in kwargs + assert _SANDBOX_KEY not in kwargs + + +@pytest.mark.asyncio +async def test_should_run_detects_only_matching_function_call(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + + active_kwargs = {"_code_interpreter_interception_active": True} + match = FakeResponse( + output=[_function_call_item(name=LITELLM_CODE_EXECUTION_TOOL_NAME)] + ) + should_run, payload = await logger.async_should_run_agentic_loop( + response=match, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs=active_kwargs, + ) + assert should_run is True + assert payload.get("tool_calls") + + no_match = FakeResponse(output=[_function_call_item(name="something_else")]) + should_run2, payload2 = await logger.async_should_run_agentic_loop( + response=no_match, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs=active_kwargs, + ) + assert should_run2 is False + assert payload2 == {} + + +@pytest.mark.asyncio +async def test_container_reused_within_request_via_server_sandbox_key(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + response = FakeResponse(output=[_function_call_item()]) + + common = dict( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + stream=False, + ) + + await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="k1"), + kwargs={"litellm_call_id": "k1", _SANDBOX_KEY: "server-nonce-1"}, + **common, + ) + await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="k1"), + kwargs={"litellm_call_id": "k1", _SANDBOX_KEY: "server-nonce-1"}, + **common, + ) + + assert ( + len(sandbox.create_calls) == 1 + ), "the sandbox is reused across loop iterations sharing one server sandbox key" + + +@pytest.mark.asyncio +async def test_colliding_caller_call_id_does_not_share_sandbox(): + """Two requests with the same caller-controlled litellm_call_id but distinct + server-minted sandbox keys must NOT share a container; otherwise one user's + code could read another in-flight request's sandbox state.""" + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + common = dict( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(1)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=FakeResponse(output=[_function_call_item()]), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + stream=False, + ) + + await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="shared"), + kwargs={"litellm_call_id": "shared", _SANDBOX_KEY: "nonce-A"}, + **common, + ) + await logger.async_build_agentic_loop_plan( + logging_obj=FakeLogging(litellm_call_id="shared"), + kwargs={"litellm_call_id": "shared", _SANDBOX_KEY: "nonce-B"}, + **common, + ) + + assert ( + len(sandbox.create_calls) == 2 + ), "distinct server sandbox keys must isolate sandboxes despite a colliding call id" + + +@pytest.mark.asyncio +async def test_pre_call_mints_server_sandbox_key(): + """The interceptor mints a server-side sandbox key (not derived from the + caller-controlled call id) when it activates.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + "litellm_call_id": "caller-supplied", + _SANDBOX_KEY: "caller-forged", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert result is not None + assert result[_SANDBOX_KEY] not in ("caller-forged", "caller-supplied") + assert len(result[_SANDBOX_KEY]) >= 16 + + +@pytest.mark.asyncio +async def test_build_plan_records_code_interpreter_call_metadata(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=FakeResponse(output=[_function_call_item()]), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={"litellm_call_id": "k1", _SANDBOX_KEY: "sbxkey1"}, + ) + + calls = plan.metadata["code_interpreter_calls"] + assert calls, "build_plan must record a code_interpreter_call for re-injection" + assert calls[0]["code"] == "print(40 + 2)" + assert calls[0]["container_id"] == "sbx_fake" + assert calls[0]["type"] == "code_interpreter_call" + assert calls[0]["status"] == "completed" + assert calls[0]["outputs"] == [{"type": "logs", "logs": "42"}], ( + "outputs must be an OpenAI-shaped logs array (not None) so clients that " + "iterate over code_interpreter_call.outputs do not break" + ) + + +@pytest.mark.asyncio +async def test_build_plan_outputs_empty_array_when_no_stdout(): + """No stdout must still yield an iteration-safe empty array, never None.""" + sandbox = FakeSandbox(stdout="") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"pass"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=FakeResponse(output=[_function_call_item()]), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={"litellm_call_id": "k1", _SANDBOX_KEY: "sbxkey1"}, + ) + + assert plan.metadata["code_interpreter_calls"][0]["outputs"] == [] + + +@pytest.mark.asyncio +async def test_post_hook_injects_code_interpreter_call_matching_openai_shape(): + from litellm.types.integrations.custom_logger import AgenticLoopPlan + + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + ci_item = { + "id": "ci_x", + "type": "code_interpreter_call", + "status": "completed", + "code": "print(1)", + "container_id": "sbx_fake", + "outputs": [{"type": "logs", "logs": "1"}], + } + plan = AgenticLoopPlan( + run_agentic_loop=True, + metadata={"code_interpreter_calls": [ci_item]}, + ) + response = FakeResponse(output=[{"type": "message", "content": []}]) + + out = await logger.async_post_agentic_loop_response_hook( + response=response, plan=plan, kwargs={} + ) + + types = [item.get("type") for item in out.output] + assert types == ["code_interpreter_call", "message"], ( + "code_interpreter_call must be re-injected before the message, matching " + "OpenAI's native output ordering" + ) + assert set(out.output[0].keys()) == { + "id", + "type", + "status", + "code", + "container_id", + "outputs", + }, "injected item must match OpenAI's code_interpreter_call keys exactly" + + +@pytest.mark.asyncio +async def test_post_hook_noop_without_recorded_calls(): + from litellm.types.integrations.custom_logger import AgenticLoopPlan + + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + response = FakeResponse(output=[{"type": "message", "content": []}]) + out = await logger.async_post_agentic_loop_response_hook( + response=response, plan=AgenticLoopPlan(run_agentic_loop=True), kwargs={} + ) + assert [item.get("type") for item in out.output] == ["message"] + + +@pytest.mark.asyncio +async def test_pre_call_forces_non_stream_for_loop(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + "stream": True, + } + + out = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert out is not None + assert out["stream"] is False, "loop requires a non-streaming upstream call" + assert out["_code_interpreter_interception_converted_stream"] is True, ( + "the converted-stream flag must be set so the final response is wrapped " + "back into a stream for the caller" + ) + + +async def _build_plan(logger, sandbox, call_id="k1", provider="openai"): + return await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=FakeResponse(output=[_function_call_item()]), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + logging_obj=FakeLogging(litellm_call_id=call_id), + stream=False, + kwargs={"litellm_call_id": call_id, _SANDBOX_KEY: "sbxkey1"}, + ) + + +@pytest.mark.asyncio +async def test_gate_refuses_without_server_active_marker(): + """A forged litellm_code_execution call must not trigger the loop unless the + pre-call hook actually converted a native code_interpreter tool.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + forged = FakeResponse( + output=[_function_call_item(name=LITELLM_CODE_EXECUTION_TOOL_NAME)] + ) + + should_run, payload = await logger.async_should_run_agentic_loop( + response=forged, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={}, + ) + + assert should_run is False + assert payload == {} + + +@pytest.mark.asyncio +async def test_gate_rechecks_provider_scope(): + """enabled_providers must be re-enforced at the gate, not only in pre-call.""" + logger = CodeInterpreterInterceptionLogger( + sandbox_config=FakeSandbox(), enabled_providers=["openai"] + ) + response = FakeResponse( + output=[_function_call_item(name=LITELLM_CODE_EXECUTION_TOOL_NAME)] + ) + + should_run, _ = await logger.async_should_run_agentic_loop( + response=response, + model="claude-x", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="anthropic", + kwargs={_ACTIVE_KEY: True}, + ) + + assert should_run is False + + +@pytest.mark.asyncio +async def test_chat_completion_gate_detects_code_execution_tool_call(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + response = { + "choices": [ + {"message": {"tool_calls": [_chat_function_call_item(call_id="call_123")]}} + ] + } + + should_run, payload = await logger.async_should_run_agentic_loop( + response=response, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={ + _ACTIVE_KEY: True, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + }, + ) + + assert should_run is True + assert payload["tool_calls"][0]["id"] == "call_123" + assert payload["tool_calls"][0]["arguments"] == '{"code":"print(40 + 2)"}' + + +@pytest.mark.asyncio +async def test_chat_completion_gate_refuses_without_server_active_marker(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + response = {"choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}]} + + should_run, payload = await logger.async_should_run_agentic_loop( + response=response, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + tools=[], + stream=False, + custom_llm_provider="openai", + kwargs={"_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE}, + ) + + assert should_run is False + assert payload == {} + + +@pytest.mark.asyncio +async def test_chat_completion_build_plan_runs_code_and_appends_tool_message(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + native_chat_tool = {"type": "code_interpreter", "container": {"type": "auto"}} + + plan = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(40 + 2)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response={ + "choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}] + }, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={ + "tools": [native_chat_tool], + "tool_choice": {"type": "code_interpreter", "container": {"type": "auto"}}, + "temperature": 0, + }, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={ + "acompletion": True, + "litellm_call_id": "k1", + _ACTIVE_KEY: True, + _SANDBOX_KEY: "sbxkey1", + "_code_interpreter_interception_converted_stream": True, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + }, + ) + + assert sandbox.run_calls[0]["code"] == "print(40 + 2)" + patch = plan.request_patch + assert patch is not None + assert patch.tools == [ + { + "type": "function", + "function": { + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "description": "Execute python code in a sandbox and return stdout.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + assert patch.optional_params == {"temperature": 0} + assert patch.kwargs == { + "litellm_call_id": "k1", + _ACTIVE_KEY: True, + _SANDBOX_KEY: "sbxkey1", + "_code_interpreter_interception_converted_stream": True, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + } + assert patch.messages is not None + assert patch.messages[-2]["role"] == "assistant" + assert patch.messages[-2]["tool_calls"][0]["id"] == "call_1" + assert patch.messages[-1] == { + "role": "tool", + "tool_call_id": "call_1", + "content": "42", + } + assert plan.metadata["code_interpreter_calls"][0]["code"] == "print(40 + 2)" + + +@pytest.mark.asyncio +async def test_pre_call_strips_client_forged_marker_on_initial_request(): + """A client cannot pre-set the active marker on the original request: with no + native code_interpreter tool, any client-supplied interception markers in + litellm_metadata are scrubbed and the active flag in kwargs is cleared.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "web_search"}], + "custom_llm_provider": "openai", + _ACTIVE_KEY: True, + "litellm_metadata": { + _ACTIVE_KEY: True, + _SANDBOX_KEY: "client-forged", + "safe_user_value": "kept", + }, + } + + await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert _ACTIVE_KEY not in kwargs, ( + "no native code_interpreter tool was present, so a client-supplied " + "active marker must be cleared" + ) + assert kwargs["litellm_metadata"] == {"safe_user_value": "kept"} + + +@pytest.mark.asyncio +async def test_pre_call_strips_forged_loop_controls_then_mints_own_markers(): + """On an INITIAL request (no server-set _agentic_loop_depth) a client cannot + smuggle loop-control state: forged _agentic_loop_depth / max_agentic_loops and + interception markers in litellm_metadata are stripped before the interceptor + activates, so the only interception markers that survive are the ones the + server mints for the converted code_interpreter tool.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "openai", + "litellm_metadata": { + _ACTIVE_KEY: True, + _SANDBOX_KEY: "client-forged", + "_agentic_loop_depth": 99, + "max_agentic_loops": 999, + "safe_user_value": "kept", + }, + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.acompletion) + + assert result is not None + metadata = result["litellm_metadata"] + assert metadata["safe_user_value"] == "kept" + assert "_agentic_loop_depth" not in metadata, "forged loop depth must be stripped" + assert "max_agentic_loops" not in metadata, "forged loop cap must be stripped" + assert metadata[_ACTIVE_KEY] is True + assert metadata[_SANDBOX_KEY] == result[_SANDBOX_KEY] + assert metadata[_SANDBOX_KEY] != "client-forged", ( + "the surviving sandbox key must be the server-minted one, not the forged " + "value the client supplied" + ) + + +@pytest.mark.asyncio +async def test_pre_call_preserves_marker_on_server_followup(): + """On a server-driven followup (depth>0) the marker is trusted and kept.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + kwargs = { + "tools": [{"type": "function", "name": LITELLM_CODE_EXECUTION_TOOL_NAME}], + "custom_llm_provider": "openai", + "_agentic_loop_depth": 1, + _ACTIVE_KEY: True, + } + + await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert kwargs.get(_ACTIVE_KEY) is True, ( + "the server-set marker must survive followup requests so multi-round " + "code execution keeps working" + ) + + +@pytest.mark.asyncio +async def test_sandbox_deleted_after_loop_completes(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + + plan = await _build_plan(logger, sandbox, call_id="k1") + assert sandbox.create_calls, "sandbox must be created during the loop" + assert ( + not sandbox.delete_calls + ), "sandbox must outlive the loop until the final hook" + + await logger.async_post_agentic_loop_response_hook( + response=FakeResponse(output=[{"type": "message", "content": []}]), + plan=plan, + kwargs={}, + ) + + assert len(sandbox.delete_calls) == 1, ( + "the sandbox must be deleted once the final response is assembled, " + "otherwise it keeps running and billing" + ) + assert "sbxkey1" not in logger._container_cache + + +@pytest.mark.asyncio +async def test_post_hook_delete_is_idempotent_across_loop_levels(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + plan = await _build_plan(logger, sandbox, call_id="k1") + response = FakeResponse(output=[{"type": "message", "content": []}]) + + await logger.async_post_agentic_loop_response_hook( + response=response, plan=plan, kwargs={} + ) + await logger.async_post_agentic_loop_response_hook( + response=response, plan=plan, kwargs={} + ) + + assert len(sandbox.delete_calls) == 1, ( + "deleting an already-removed container must be a no-op so unwinding " + "loop levels do not double-delete" + ) + + +@pytest.mark.asyncio +async def test_build_plan_deletes_sandbox_when_execution_raises(): + """If sandbox execution raises before a plan is built (e.g. E2B aborts + output over its cap), the cached sandbox must be deleted before re-raising, + otherwise a caller can leak paid containers until the prune TTL.""" + + class RaisingSandbox(FakeSandbox): + async def arun_code(self, *, container, code, **kwargs): + raise ValueError("output exceeded cap") + + sandbox = RaisingSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + + with pytest.raises(ValueError, match="exceeded cap"): + await _build_plan(logger, sandbox, call_id="k1") + + assert len(sandbox.create_calls) == 1, "the sandbox must have been created" + assert len(sandbox.delete_calls) == 1, ( + "a build failure must delete the cached sandbox so it does not keep " + "running and billing" + ) + assert "sbxkey1" not in logger._container_cache + + +@pytest.mark.asyncio +async def test_cleanup_hook_deletes_sandbox(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + plan = await _build_plan(logger, sandbox, call_id="k1") + + await logger.async_agentic_loop_cleanup_hook(plan=plan, kwargs={}) + + assert len(sandbox.delete_calls) == 1, ( + "the cleanup hook must delete the sandbox so a rerun failure cannot " + "leak a running container" + ) + assert "sbxkey1" not in logger._container_cache + + +@pytest.mark.asyncio +async def test_cleanup_hook_is_idempotent_with_post_hook(): + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + plan = await _build_plan(logger, sandbox, call_id="k1") + + await logger.async_post_agentic_loop_response_hook( + response=FakeResponse(output=[{"type": "message", "content": []}]), + plan=plan, + kwargs={}, + ) + await logger.async_agentic_loop_cleanup_hook(plan=plan, kwargs={}) + + assert len(sandbox.delete_calls) == 1, ( + "cleanup running in finally after the success-path post hook already " + "deleted the sandbox must not double-delete" + ) + + +@pytest.mark.asyncio +async def test_responses_plan_cleans_up_sandbox_when_followup_raises(): + """If the agentic rerun fails, _execute_responses_agentic_plan must still + invoke the cleanup hook so the sandbox is not left running.""" + import litellm + from litellm.integrations.custom_logger import CustomLogger + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, + ) + + cleanup_calls = [] + + class CleanupCallback(CustomLogger): + async def async_post_agentic_loop_response_hook(self, response, plan, kwargs): + return response + + async def async_agentic_loop_cleanup_hook(self, plan, kwargs): + cleanup_calls.append(plan) + + plan = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + model="gpt-5", messages=[{"role": "user", "content": "x"}] + ), + metadata={"sandbox_key": "sbxkey1"}, + ) + + original = litellm.aresponses + + async def _boom(*args, **kwargs): + raise RuntimeError("upstream blew up") + + litellm.aresponses = _boom + try: + with pytest.raises(RuntimeError, match="upstream blew up"): + await BaseLLMHTTPHandler()._execute_responses_agentic_plan( + plan=plan, + model="gpt-5", + response_api_optional_request_params={}, + logging_obj=FakeLogging(litellm_call_id="k1"), + kwargs={}, + depth=0, + max_loops=3, + fingerprints=[], + fingerprint="fp", + callback=CleanupCallback(), + ) + finally: + litellm.aresponses = original + + assert cleanup_calls == [plan], ( + "cleanup hook must run in finally even when the rerun raises, otherwise " + "the sandbox keeps running until the prune TTL" + ) + + +@pytest.mark.asyncio +async def test_run_code_does_not_re_resolve_registry(monkeypatch): + """Params resolved once at create time must be reused for running code, so a + registry clear between create and run cannot turn into a create-then-fail.""" + import litellm + from litellm.sandbox import sandbox_tools + + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, + } + ] + ) + + create_kwargs = {} + run_kwargs = {} + + async def fake_acreate_sandbox(**kwargs): + create_kwargs.update(kwargs) + return FakeHandle() + + async def fake_arun_code(**kwargs): + run_kwargs.update(kwargs) + return CodeExecutionResult(stdout="ok") + + monkeypatch.setattr(litellm, "acreate_sandbox", fake_acreate_sandbox) + monkeypatch.setattr(litellm, "arun_code", fake_arun_code) + + logger = CodeInterpreterInterceptionLogger(sandbox_tool_name="e2b_default") + try: + container, params = await logger._get_or_create_container(cache_key="k1") + assert params is not None and params["sandbox_provider"] == "e2b" + + sandbox_tools.clear_sandbox_tools() + + stdout = await logger._run_tool_call( + container=container, params=params, arguments='{"code":"print(1)"}' + ) + finally: + sandbox_tools.clear_sandbox_tools() + + assert stdout == "ok", "run must succeed using the params captured at create time" + assert run_kwargs["provider"] == "e2b" + + +@pytest.mark.asyncio +async def test_run_tool_call_surfaces_execution_error(): + """A sandbox execution error must be fed back to the model as a labelled + string, not raised, so the agentic loop can react to it.""" + + class ErroringSandbox(FakeSandbox): + async def arun_code(self, *, container, code, **kwargs): + self.run_calls.append({"container": container, "code": code}) + return CodeExecutionResult( + stdout="", error={"name": "ValueError", "value": "boom"} + ) + + sandbox = ErroringSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + container = await logger._create_container() + + stdout = await logger._run_tool_call( + container=container[0], params=None, arguments='{"code":"raise ValueError(1)"}' + ) + + assert stdout == "[execution error] boom" + + +@pytest.mark.asyncio +async def test_run_tool_call_reports_unparseable_arguments(): + """Malformed tool arguments must produce a parse error string the model can + see rather than crashing the interceptor.""" + sandbox = FakeSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + container = await logger._create_container() + + stdout = await logger._run_tool_call( + container=container[0], params=None, arguments="not-json" + ) + + assert stdout == "[invalid tool arguments: could not parse code]" + assert not sandbox.run_calls, "code must not run when arguments cannot be parsed" + + +@pytest.mark.asyncio +async def test_pre_call_skips_provider_outside_scope(): + """enabled_providers must filter the pre-call conversion so a request to an + out-of-scope provider is left untouched.""" + logger = CodeInterpreterInterceptionLogger( + sandbox_config=FakeSandbox(), enabled_providers=["openai"] + ) + kwargs = { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "custom_llm_provider": "anthropic", + } + + result = await logger.async_pre_call_deployment_hook(kwargs, CallTypes.aresponses) + + assert result is None + assert kwargs["tools"][0]["type"] == "code_interpreter", "tool must be untouched" + assert _ACTIVE_KEY not in kwargs + + +@pytest.mark.asyncio +async def test_resolve_provider_falls_back_to_model_lookup(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + + assert logger._resolve_provider({"custom_llm_provider": "openai"}) == "openai" + assert logger._resolve_provider({"model": "gpt-5"}) == "openai" + assert logger._resolve_provider({"model": 123}) is None + assert logger._resolve_provider({"model": "no-such-provider-xyz"}) is None + + +@pytest.mark.asyncio +async def test_create_container_without_sandbox_raises(): + """The registry path must raise a clear error when no sandbox is resolvable + instead of silently creating nothing.""" + logger = CodeInterpreterInterceptionLogger(sandbox_tool_name="missing") + + with pytest.raises(ValueError, match="no sandbox available"): + await logger._create_container() + + +@pytest.mark.asyncio +async def test_run_code_without_params_raises(): + logger = CodeInterpreterInterceptionLogger(sandbox_tool_name="missing") + + with pytest.raises(ValueError, match="no sandbox available to run code"): + await logger._run_code(container=FakeHandle(), params=None, code="print(1)") + + +@pytest.mark.asyncio +async def test_delete_container_swallows_errors(): + """A delete failure must not propagate; the request already succeeded.""" + + class FailingDeleteSandbox(FakeSandbox): + async def adelete_sandbox(self, *, container, **kwargs): + raise RuntimeError("e2b unreachable") + + sandbox = FailingDeleteSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + container, params = await logger._create_container() + + await logger._delete_container(container=container, params=params) + + +@pytest.mark.asyncio +async def test_prune_expired_cache_deletes_underlying_container(): + """Expired cache entries must have their sandbox deleted, not just dropped, + otherwise an orphaned sandbox keeps running.""" + import litellm.integrations.code_interpreter_interception.handler as handler_mod + + sandbox = FakeSandbox() + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + container, params = await logger._create_container() + logger._container_cache["old"] = ( + container, + params, + time.time() - handler_mod._CACHE_TTL_SECONDS - 1, + ) + + await logger._prune_expired_cache() + + assert "old" not in logger._container_cache + assert len(sandbox.delete_calls) == 1, "expired sandbox must be deleted" + + +@pytest.mark.asyncio +async def test_normalize_messages_handles_str_and_unknown(): + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + + assert logger._normalize_messages("hi") == [{"role": "user", "content": "hi"}] + assert logger._normalize_messages([{"role": "user"}]) == [{"role": "user"}] + assert logger._normalize_messages(42) == [] + + +def test_from_config_yaml_reads_fields(): + cfg = { + "enabled": False, + "enabled_providers": ["openai"], + "sandbox_tool_name": "e2b_default", + } + logger = CodeInterpreterInterceptionLogger.from_config_yaml(cfg) + + assert logger.enabled is False + assert logger.enabled_providers == ["openai"] + assert logger.sandbox_tool_name == "e2b_default" + + +def test_initialize_from_proxy_config_prefers_litellm_settings(): + logger = CodeInterpreterInterceptionLogger.initialize_from_proxy_config( + litellm_settings={ + "code_interpreter_interception_params": { + "enabled_providers": ["openai"], + "sandbox_tool_name": "e2b_default", + } + }, + callback_specific_params={}, + ) + + assert logger.enabled_providers == ["openai"] + assert logger.sandbox_tool_name == "e2b_default" + + +@pytest.mark.asyncio +async def test_build_plan_handles_dict_shaped_response(): + """A responses payload delivered as a plain dict (not an object) must flow + through detection, execution, and re-injection the same as the typed form.""" + sandbox = FakeSandbox(stdout="42") + logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + dict_response = {"output": [_function_call_item()]} + + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": logger._extract_code_execution_tool_calls(dict_response)}, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=dict_response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={"litellm_call_id": "k1", _SANDBOX_KEY: "sbxkey1"}, + ) + + assert sandbox.run_calls, "code must run for a dict-shaped response" + assert plan.metadata["code_interpreter_calls"][0]["code"] == "print(40 + 2)" + + out = await logger.async_post_agentic_loop_response_hook( + response={"output": [{"type": "message", "content": []}]}, + plan=plan, + kwargs={}, + ) + + assert [item.get("type") for item in out["output"]] == [ + "code_interpreter_call", + "message", + ], "the dict-shaped response must get the code_interpreter_call re-injected" + + +@pytest.mark.asyncio +async def test_extract_tool_calls_reads_object_attributes(): + """Detection must work when output items are objects with attributes, not + only dicts.""" + + class Item: + def __init__(self): + self.type = "function_call" + self.name = LITELLM_CODE_EXECUTION_TOOL_NAME + self.call_id = "c9" + self.arguments = '{"code":"print(1)"}' + + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox()) + calls = logger._extract_code_execution_tool_calls(FakeResponse(output=[Item()])) + + assert len(calls) == 1 + assert calls[0]["call_id"] == "c9" + assert calls[0]["arguments"] == '{"code":"print(1)"}' diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_mount.py b/tests/test_litellm/integrations/otel/test_otel_v2_mount.py index 956d8c53cee..7240d49d022 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_mount.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_mount.py @@ -35,6 +35,13 @@ from litellm.integrations.otel.mount import ( # noqa: E402 ) +@pytest.fixture(autouse=True) +def _clear_otel_v2_flag_cache(): + is_otel_v2_enabled.cache_clear() + yield + is_otel_v2_enabled.cache_clear() + + class _FakeSpan: """Minimal recording span capturing what the hook writes.""" @@ -70,8 +77,10 @@ def _instrumented_app(): def test_gate_toggles_with_env(monkeypatch): """The startup mount is guarded by this flag.""" monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() assert is_otel_v2_enabled() is False monkeypatch.setenv("LITELLM_OTEL_V2", "1") + is_otel_v2_enabled.cache_clear() assert is_otel_v2_enabled() is True 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 3447f5bdb7e..4bb26a70b02 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 @@ -1,6 +1,8 @@ """Tests for the OTel v2 sources of truth: span registry, semconv keys, config, and the typed StandardLoggingPayload adapter. These need no OTel SDK.""" +import pytest + from litellm.integrations.otel import ( BAGGAGE_PROMOTED_KEYS, DB, @@ -28,6 +30,13 @@ from litellm.integrations.otel.model.spans import ( ) +@pytest.fixture(autouse=True) +def _clear_otel_v2_flag_cache(): + is_otel_v2_enabled.cache_clear() + yield + is_otel_v2_enabled.cache_clear() + + def _sample_payload(**overrides): payload = { "call_type": "acompletion", @@ -561,11 +570,37 @@ def test_capture_message_content_normalizer_only_touches_strings(): def test_v2_flag_is_off_by_default(monkeypatch): monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) + is_otel_v2_enabled.cache_clear() assert is_otel_v2_enabled() is False monkeypatch.setenv("LITELLM_OTEL_V2", "true") + is_otel_v2_enabled.cache_clear() assert is_otel_v2_enabled() is True +def test_v2_flag_resolved_once_not_per_call(monkeypatch): + """Regression for LIT-3895: ``is_otel_v2_enabled`` sits on the proxy hot path + (auth, logging-callback setup). Building the pydantic-settings model on every + call re-scanned the environment at ~28us a pop and dropped throughput, so the + flag must be resolved once and cached rather than reconstructed per call.""" + from litellm.integrations.otel.model import config as config_mod + + constructions = 0 + real_flag = config_mod._OTelV2Flag + + def _counting_flag(*args, **kwargs): + nonlocal constructions + constructions += 1 + return real_flag(*args, **kwargs) + + monkeypatch.setattr(config_mod, "_OTelV2Flag", _counting_flag) + config_mod.is_otel_v2_enabled.cache_clear() + + for _ in range(50): + config_mod.is_otel_v2_enabled() + + assert constructions == 1 + + def test_config_from_env(monkeypatch): for var in ( "OTEL_EXPORTER", diff --git a/tests/test_litellm/interactions/test_google_interactions_integration.py b/tests/test_litellm/interactions/test_google_interactions_integration.py index cfff26d51ef..9c651cc94f5 100644 --- a/tests/test_litellm/interactions/test_google_interactions_integration.py +++ b/tests/test_litellm/interactions/test_google_interactions_integration.py @@ -299,7 +299,6 @@ class TestGoogleInteractionsResponseStructure: assert hasattr(response, "outputs") assert hasattr(response, "usage") assert hasattr(response, "model") or hasattr(response, "agent") - assert hasattr(response, "role") assert hasattr(response, "created") assert hasattr(response, "updated") diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/test_litellm/interactions/test_openapi_compliance.py index aededaaca77..209e99895db 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/test_litellm/interactions/test_openapi_compliance.py @@ -156,16 +156,19 @@ class TestResponseCompliance: # The response is the dedicated `Interaction` schema. Google moved the # output-only fields (notably the `steps` array, formerly `outputs`) # off `CreateModelInteractionParams` and onto `Interaction`; the request - # schema no longer carries `steps`. Keep this aligned with the live spec. + # schema no longer carries `steps`. Google later moved `role` off + # `Interaction` onto the per-turn `Turn` schema (asserted in + # test_turn_schema), so it is no longer a top-level output field here. + # Keep this aligned with the live spec. schema = spec_dict["components"]["schemas"]["Interaction"] - # Output fields (readOnly). + # Output fields (readOnly). `role` was removed from the `Interaction` + # schema by Google; it now lives only on `Turn`. output_fields = [ "id", "status", "created", "updated", - "role", "steps", "usage", ] diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 9b3152fae07..7f3d5a959a1 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1499,19 +1499,20 @@ def _local_model_cost_map(): @pytest.mark.parametrize("data_residency", ["eu", "us"]) def test_data_residency_applies_uplift(data_residency, _local_model_cost_map): - """gpt-5 should apply the regional processing uplift multiplier when - data_residency is set.""" + """gpt-5.4 should apply the regional processing uplift multiplier when + data_residency is set. gpt-5.4+ (released 2026-03-05) carry the 10% uplift; + gpt-5 and older models do not.""" from litellm.types.utils import Usage usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) base = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", ) regional = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", data_residency=data_residency, @@ -1526,6 +1527,23 @@ def test_data_residency_applies_uplift(data_residency, _local_model_cost_map): assert regional[1] == pytest.approx(base[1] * 1.10, rel=1e-9) +@pytest.mark.parametrize("model", ["gpt-5", "gpt-5-mini", "gpt-5-nano", "gpt-5-pro", "gpt-4o", "gpt-4.1"]) +def test_data_residency_no_uplift_for_pre_march_2026_models(model, _local_model_cost_map): + """Models released before 2026-03-05 must not have the regional uplift.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + base = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai") + regional = generic_cost_per_token( + model=model, usage=usage, custom_llm_provider="openai", data_residency="eu" + ) + + assert base == regional, ( + f"{model} should not have a regional uplift, but cost changed with data_residency" + ) + + def test_data_residency_no_uplift_for_unmarked_model(_local_model_cost_map): """A model without a regional_processing_uplift_multiplier_* entry should fall back to base pricing, not error.""" @@ -1555,12 +1573,12 @@ def test_data_residency_none_no_uplift(_local_model_cost_map): usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) base = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", ) explicit_none = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", data_residency=None, @@ -1576,13 +1594,13 @@ def test_data_residency_composes_with_service_tier(_local_model_cost_map): usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) priority_base = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", service_tier="priority", ) priority_eu = generic_cost_per_token( - model="gpt-5", + model="gpt-5.4", usage=usage, custom_llm_provider="openai", service_tier="priority", 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 new file mode 100644 index 00000000000..f1196ab4692 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -0,0 +1,422 @@ +""" +Tests for the provider-agnostic chat completion agentic loop dispatcher +(`litellm/litellm_core_utils/chat_completion_agentic_loop.py`) and the +code-interpreter interception integration that drives it. + +The load-bearing regression here protects a reviewer requirement: the internal +agentic/interception control fields must NEVER reach the outbound provider HTTP +request body. The relevant fields are: + + _agentic_loop_depth + _agentic_loop_fingerprints + _agentic_loop_api_surface + max_agentic_loops + _code_interpreter_interception_active + _code_interpreter_interception_sandbox_key + _code_interpreter_interception_converted_stream + +A scrubber in gpt_transformation.py used to strip these. That scrubber was +removed, so `test_internal_control_fields_never_leak_into_provider_body` proves +they stay out of the body even without it. +""" + +import os +import sys +from typing import Any, Dict, List, Optional, Tuple +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, +) +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + maybe_run_chat_completion_agentic_loop, +) +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) +from litellm.types.utils import ( + Choices, + Function, + ChatCompletionMessageToolCall, + Message, + ModelResponse, +) + +# The internal control fields that must never reach a provider request body. +_INTERNAL_CONTROL_FIELDS = ( + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_agentic_loop_api_surface", + "max_agentic_loops", + "_code_interpreter_interception_active", + "_code_interpreter_interception_sandbox_key", + "_code_interpreter_interception_converted_stream", + "litellm_metadata", +) + + +@pytest.fixture +def restore_callbacks(): + """Save/restore litellm.callbacks so a registered fake logger never pollutes + other tests in the suite.""" + saved = list(litellm.callbacks) + try: + yield + finally: + litellm.callbacks = saved + + +class _SandboxResult: + def __init__(self, stdout: str) -> None: + self.stdout = stdout + self.error = None + + +class FakeSandboxConfig: + """Injected sandbox so the interception loop runs no real network / E2B.""" + + def __init__(self) -> None: + self.created = 0 + self.deleted = 0 + self.run_codes: List[str] = [] + + async def acreate_sandbox(self) -> Any: + self.created += 1 + return MagicMock(id="sandbox-123") + + async def arun_code(self, container: Any, code: str) -> _SandboxResult: + self.run_codes.append(code) + return _SandboxResult(stdout="42\n") + + async def adelete_sandbox(self, container: Any) -> None: + self.deleted += 1 + + +def _tool_call_model_response() -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_abc", + type="function", + function=Function( + name="litellm_code_execution", + arguments='{"code": "print(6*7)"}', + ), + ) + ], + ), + ) + ] + ) + + +def _plain_model_response(content: str = "The answer is 42") -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="stop", + message=Message(role="assistant", content=content), + ) + ] + ) + + +def _raw_response_for(model_response: ModelResponse) -> MagicMock: + """Wrap a ModelResponse as the OpenAI `with_raw_response.create` return value + (an object exposing `.headers` and `.parse()` -> something with model_dump).""" + parsed = MagicMock() + parsed.model_dump.return_value = model_response.model_dump() + raw = MagicMock() + raw.headers = {} + raw.parse.return_value = parsed + return raw + + +# --------------------------------------------------------------------------- +# A) PROVIDER-PAYLOAD REGRESSION +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_internal_control_fields_never_leak_into_provider_body(restore_callbacks): + """Drive a real acompletion with a native code_interpreter tool through the + interception logger + agentic loop, capturing every outbound OpenAI request + body. None of the internal control fields may appear at top-level or inside + extra_body on ANY of the captured calls.""" + logger = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandboxConfig()) + litellm.callbacks = [logger] + + # First create -> model emits a code_execution tool call (triggers the loop). + # Second create -> model returns a plain answer (loop terminates). + create = AsyncMock( + side_effect=[ + _raw_response_for(_tool_call_model_response()), + _raw_response_for(_plain_model_response()), + ] + ) + mock_client = MagicMock() + mock_client.chat.completions.with_raw_response.create = create + + response = await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + tools=[{"type": "code_interpreter"}], + tool_choice={"type": "code_interpreter"}, + api_key="sk-test", + client=mock_client, + ) + + # The loop must have actually fired (sanity: two provider calls). + assert create.await_count == 2, ( + "expected the agentic loop to issue a follow-up provider call; " + f"got {create.await_count} call(s)" + ) + + for idx, call in enumerate(create.await_args_list): + body = call.kwargs + extra_body = body.get("extra_body") or {} + for field in _INTERNAL_CONTROL_FIELDS: + assert field not in body, ( + f"provider call #{idx}: internal field {field!r} leaked into " + f"top-level request body: {sorted(body.keys())}" + ) + assert field not in extra_body, ( + f"provider call #{idx}: internal field {field!r} leaked into " + f"extra_body: {sorted(extra_body.keys())}" + ) + # The native code_interpreter tool must have been swapped for the + # function tool, never sent raw to OpenAI as a chat-completions request. + for tool in body.get("tools") or []: + assert tool.get("type") != "code_interpreter" + + # The final response is the post-loop answer, not the tool-call turn. + assert response.choices[0].message.content == "The answer is 42" + + +# --------------------------------------------------------------------------- +# B) DISPATCHER UNIT TESTS +# --------------------------------------------------------------------------- + + +class _LoggingStub: + """Minimal logging_obj: dispatcher only reads dynamic_success_callbacks and + litellm_call_id off it.""" + + litellm_call_id = "call-test" + dynamic_success_callbacks: List[Any] = [] + + +class _GateOnlyLogger(CustomLogger): + """Overrides the gate to fire, but builds a plan from request_patch.""" + + def __init__(self, plan: AgenticLoopPlan, tool_calls: Dict[str, Any]) -> None: + super().__init__() + self._plan = plan + self._tool_calls = tool_calls + self.cleanup_calls = 0 + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + stream: bool, + custom_llm_provider: str, + kwargs: Dict[str, Any], + ) -> Tuple[bool, Dict[str, Any]]: + return True, self._tool_calls + + async def async_build_agentic_loop_plan( + self, + tools: Dict[str, Any], + model: str, + messages: List[Dict[str, Any]], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict[str, Any], + logging_obj: Any, + stream: bool, + kwargs: Dict[str, Any], + ) -> AgenticLoopPlan: + return self._plan + + async def async_agentic_loop_cleanup_hook( + self, plan: AgenticLoopPlan, kwargs: Dict[str, Any] + ) -> None: + self.cleanup_calls += 1 + + +def _patched_messages() -> List[Dict[str, Any]]: + return [ + {"role": "user", "content": "what is 6*7?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": { + "name": "litellm_code_execution", + "arguments": '{"code": "print(6*7)"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call_abc", "content": "42\n"}, + ] + + +@pytest.mark.asyncio +async def test_dispatcher_returns_none_when_no_callback_gates(restore_callbacks): + """No callback overrides the gate -> dispatcher returns None so the caller + keeps the original response untouched.""" + litellm.callbacks = [] + + result = await maybe_run_chat_completion_agentic_loop( + response=_plain_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={}, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_dispatcher_runs_followup_with_incremented_depth_and_patched_messages( + restore_callbacks, +): + """A gating logger with a request_patch -> the dispatcher calls + litellm.acompletion exactly once with _agentic_loop_depth == 1 and the + patched messages. Loop-control state rides as litellm-level kwargs and is + mirrored into litellm_metadata; the provider-surface transient + _agentic_loop_api_surface is never forwarded. (Provider-body stripping of + these litellm-level kwargs is asserted separately in test A.)""" + followup = _plain_model_response("done") + plan = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=_patched_messages()), + ) + logger = _GateOnlyLogger(plan=plan, tool_calls={"tool_calls": [{"id": "call_abc"}]}) + litellm.callbacks = [logger] + + 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.1}, + kwargs={"_code_interpreter_interception_active": True}, + 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["_agentic_loop_depth"] == 1 + assert call_kwargs["messages"] == _patched_messages() + # Preserved non-internal optional param survives the rerun. + assert call_kwargs["temperature"] == 0.1 + # Loop-control state is carried at the litellm level for the follow-up. + assert call_kwargs["max_agentic_loops"] >= 1 + assert "_agentic_loop_fingerprints" in call_kwargs + # Interception markers are mirrored into litellm_metadata for the follow-up. + assert ( + call_kwargs["litellm_metadata"]["_code_interpreter_interception_active"] is True + ) + # The transient surface marker is NOT forwarded to the follow-up call. + assert "_agentic_loop_api_surface" not in call_kwargs + # Cleanup hook always runs. + assert logger.cleanup_calls == 1 + + +@pytest.mark.asyncio +async def test_dispatcher_raises_when_depth_reaches_max_agentic_loops( + restore_callbacks, +): + """depth >= max_agentic_loops -> ValueError mentioning max_agentic_loops, + before any follow-up call is attempted.""" + logger = _GateOnlyLogger( + plan=AgenticLoopPlan(run_agentic_loop=True), + tool_calls={"tool_calls": [{"id": "call_abc"}]}, + ) + litellm.callbacks = [logger] + + acompletion_mock = AsyncMock() + with patch.object(litellm, "acompletion", acompletion_mock): + with pytest.raises(ValueError, match="max_agentic_loops"): + await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={"_agentic_loop_depth": 3, "max_agentic_loops": 3}, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + acompletion_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_dispatcher_raises_on_repeated_tool_call_fingerprint(restore_callbacks): + """A tool_calls fingerprint already present in _agentic_loop_fingerprints -> + ValueError about the repeated fingerprint (cycle guard), with no follow-up + call.""" + import json + + # The dispatcher fingerprints the whole value the gate returns as its second + # tuple element, so the seeded fingerprint must mirror that dict exactly. + gate_tool_calls = { + "tool_calls": [{"id": "call_abc", "name": "litellm_code_execution"}] + } + fingerprint = json.dumps(gate_tool_calls, sort_keys=True, default=str) + + logger = _GateOnlyLogger( + plan=AgenticLoopPlan(run_agentic_loop=True), + tool_calls=gate_tool_calls, + ) + litellm.callbacks = [logger] + + acompletion_mock = AsyncMock() + with patch.object(litellm, "acompletion", acompletion_mock): + with pytest.raises(ValueError, match="fingerprint"): + await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={ + "_agentic_loop_depth": 0, + "max_agentic_loops": 3, + "_agentic_loop_fingerprints": [fingerprint], + }, + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=False, + ) + + acompletion_mock.assert_not_awaited() diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 35c02184a51..53960847fdc 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1,6 +1,7 @@ import os import sys +import httpx import pytest import litellm @@ -550,3 +551,78 @@ class TestExtractAndRaiseLitellmException: ) assert result is None + + +class ModelError(Exception): + """Mimics replicate's SDK exception, whose mapping keys on the class name.""" + + +class CohereConnectionError(Exception): + """Mimics cohere's SDK exception, whose mapping keys on the class name.""" + + +def test_replicate_model_error_maps_to_bad_request(): + """The replicate branch keys on ``type(original_exception).__name__ == + "ModelError"`` rather than on the error string. The dispatch now lives in a + per-provider helper, so this class-name value has to be threaded into the + helper; if it is not, the bare name ``exception_type`` resolves to the + module-level function and the comparison is always False, silently mismapping + to APIConnectionError.""" + original_exception = ModelError("the deployed model failed to return a prediction") + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="replicate/meta/llama-2-70b-chat", + original_exception=original_exception, + custom_llm_provider="replicate", + ) + + assert excinfo.value.llm_provider == "replicate" + + +def test_cohere_connection_error_maps_to_rate_limit(): + """The cohere branch keys on ``"CohereConnectionError" in + type(original_exception).__name__``. With the dispatch extracted into a helper + the class-name value must be passed in; otherwise ``in`` runs against the + module-level ``exception_type`` function object and raises TypeError, which the + outer handler swallows into a generic APIConnectionError.""" + original_exception = CohereConnectionError("connection reset by peer") + original_exception.message = "connection reset by peer" + + with pytest.raises(litellm.RateLimitError) as excinfo: + exception_type( + model="command-r", + original_exception=original_exception, + custom_llm_provider="cohere", + ) + + assert excinfo.value.llm_provider == "cohere" + + +class ReplicateError(Exception): + """Mimics a replicate HTTP error carrying a status_code and response.""" + + def __init__(self, message: str, status_code: int) -> None: + super().__init__(message) + self.message = message + self.status_code = status_code + self.response = httpx.Response( + status_code=status_code, + request=httpx.Request("POST", "https://api.replicate.com/v1/predictions"), + ) + + +def test_replicate_422_maps_to_unprocessable_entity(): + """The replicate status-code ladder carried two identical ``status_code == 422`` + branches; the second was unreachable dead code. After dropping the duplicate the + surviving branch must still map 422 to UnprocessableEntityError.""" + original_exception = ReplicateError("validation failed for the input", 422) + + with pytest.raises(litellm.UnprocessableEntityError) as excinfo: + exception_type( + model="replicate/meta/llama-2-70b-chat", + original_exception=original_exception, + custom_llm_provider="replicate", + ) + + assert excinfo.value.llm_provider == "replicate" diff --git a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py index 3c280c6ba92..84900e3f2ed 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py @@ -132,3 +132,17 @@ def test_azure_base_model_detection_preserved(): assert params is not None assert "reasoning_effort" in params assert "tools" in params + + +def test_sambanova_embeddings_request_returns_list_not_none(): + """The sambanova embeddings branch resolved the config but dropped the result, + so embedding requests got ``None`` instead of the supported-params list while the + chat branch returned correctly. A list (the sambanova embeddings config exposes no + extra params, hence ``[]``) must reach the caller.""" + embedding_params = get_supported_openai_params( + model="E5-Mistral-7B-Instruct", + custom_llm_provider="sambanova", + request_type="embeddings", + ) + + assert embedding_params == [] diff --git a/tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py b/tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py new file mode 100644 index 00000000000..4e016622f1c --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_request_timeout_resolver.py @@ -0,0 +1,58 @@ +"""Unit tests for litellm.litellm_core_utils.request_timeout_resolver. + +The resolver decides whether ``litellm.request_timeout`` was *explicitly configured* +(env REQUEST_TIMEOUT / litellm_settings, or a non-default runtime value) versus left +at the package default. This is what lets request_timeout act as an independent +per-attempt timeout instead of being indistinguishable from "nobody set it". +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))) + +import litellm +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) + + +@pytest.fixture +def restore_request_timeout(): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + try: + yield + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + +def test_default_value_without_flag_is_unset(restore_request_timeout): + litellm.request_timeout = DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = False + assert get_configured_request_timeout() is None + + +def test_explicit_flag_returns_value(restore_request_timeout): + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = True + assert get_configured_request_timeout() == 300.0 + + +def test_explicit_flag_preserves_value_equal_to_default(restore_request_timeout): + # The case the bare ``!= default`` heuristic gets wrong: a user who explicitly + # configures the default value still means it explicitly. + litellm.request_timeout = DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = True + assert get_configured_request_timeout() == float(DEFAULT_REQUEST_TIMEOUT_SECONDS) + + +def test_non_default_runtime_value_treated_as_explicit(restore_request_timeout): + # SDK users assigning litellm.request_timeout directly (no flag) must keep working. + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = False + assert get_configured_request_timeout() == 300.0 diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py index 6808c4821c1..7239636fd48 100644 --- a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py +++ b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py @@ -126,6 +126,49 @@ def test_lists_with_sensitive_keys_are_masked(): assert masked["tags"] == ["prod", "test"] +def test_short_secrets_are_fully_masked(): + """ + Regression test: secrets at or below the reveal threshold (visible_prefix + + visible_suffix, 8 by default) were returned verbatim instead of masked. + An exactly-8-char value hit masked_length == 0 and round-tripped unchanged; + anything shorter hit the early return. Both leaked short credentials (e.g. an + 8-char redis password) in plaintext through mask_dict. + """ + masker = SensitiveDataMasker() + + # Boundary: exactly 8 chars previously returned verbatim. + assert masker._mask_value("abcd1234") == "********" + # Below threshold previously hit the early return and leaked verbatim. + assert masker._mask_value("sk-12") == "*****" + # Values above the threshold must still partially reveal, not over-mask. + assert masker._mask_value("abcd12345") == "abcd*2345" + + masked = masker.mask_dict({"redis_password": "pass1234", "api_key": "sk-7a"}) + assert masked["redis_password"] == "********" + assert masked["api_key"] == "*****" + + +def test_mask_short_values_false_keeps_short_values_readable(): + """ + mask_short_values=False opts out of full masking so short values are returned + as-is. This preserves the truncation use (e.g. CooldownCache shows the first 50 + chars of an exception and only masks longer tails), while longer values are still + partially masked. + """ + masker = SensitiveDataMasker( + visible_prefix=50, visible_suffix=0, mask_short_values=False + ) + + short = "Test exception for structure validation" + assert masker._mask_value(short) == short + + long_value = "x" * 60 + masked = masker._mask_value(long_value) + assert masked.startswith("x" * 50) + assert masked.endswith("*" * 10) + assert len(masked) == 60 + + def test_cost_per_token_fields_not_masked(): """ Regression test: cost fields like input_cost_per_token contain "token" in their name diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index e95cd656cc4..81af0ad3e6f 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -18,6 +18,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import ( AUDIO_ATTRIBUTE, CustomStreamWrapper, + _ProviderChunkEarlyReturn, + _ProviderChunkParsed, ) from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -876,6 +878,114 @@ def test_sync_streaming_bad_request_not_midstream(logging_obj: Logging): assert "invalid maxOutputTokens" in str(excinfo.value) +def _bedrock_error_event(exception_type: str): + """A mocked botocore event-stream error event: status_code is botocore's + hard-coded 400, with the real type in the :exception-type header.""" + event = Mock() + event.to_response_dict = Mock( + return_value={ + "status_code": 400, + "headers": { + ":exception-type": exception_type, + ":content-type": "application/json", + ":message-type": "exception", + }, + "body": b'{"message":"Bedrock had an internal error."}', + } + ) + return event + + +@pytest.mark.asyncio +async def test_bedrock_midstream_internal_server_error_wraps_for_fallback( + logging_obj: Logging, +): + """End-to-end regression for https://github.com/BerriAI/litellm/issues/24608: + a Bedrock mid-stream internalServerException event (botocore stamps it 400) + must flow through the real decoder, gain its modeled 500 status, and wrap + into MidStreamFallbackError so the Router can run streaming fallback. + + Calls the real AWSEventStreamDecoder, so reverting the decoder status fix + makes the decoder raise BedrockError(400) and the gate raises BadRequestError + directly -> this test fails without the fix.""" + from litellm.exceptions import MidStreamFallbackError + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + decoder = AWSEventStreamDecoder(model="anthropic.claude-3-sonnet-20240229-v1:0") + + async def _bedrock_stream(): + decoder._parse_message_from_event( + _bedrock_error_event("internalServerException") + ) + yield # unreachable; the line above raises + + async def _make_call(**kwargs): + return _bedrock_stream() + + response = CustomStreamWrapper( + completion_stream=None, + model="anthropic.claude-3-sonnet-20240229-v1:0", + logging_obj=logging_obj, + custom_llm_provider="bedrock", + make_call=_make_call, + ) + + with pytest.raises(MidStreamFallbackError): + await response.__anext__() + + +@pytest.mark.asyncio +async def test_bedrock_5xx_wraps_for_midstream_fallback(logging_obj: Logging): + """Gate contract: a Bedrock 5xx (here 503 serviceUnavailableException) wraps + into MidStreamFallbackError so the Router can run streaming fallback.""" + from litellm.exceptions import MidStreamFallbackError + from litellm.llms.bedrock.chat.invoke_handler import BedrockError + + async def _raise_503(**kwargs): + raise BedrockError( + status_code=503, + message="serviceUnavailableException Bedrock is unavailable.", + ) + + response = CustomStreamWrapper( + completion_stream=None, + model="anthropic.claude-3-sonnet-20240229-v1:0", + logging_obj=logging_obj, + custom_llm_provider="bedrock", + make_call=_raise_503, + ) + + with pytest.raises(MidStreamFallbackError): + await response.__anext__() + + +@pytest.mark.asyncio +async def test_bedrock_validation_error_raises_directly(logging_obj: Logging): + """Gate contract: a Bedrock validationException (400) is a client error and + must surface directly, never wrapped into MidStreamFallbackError.""" + from litellm.exceptions import MidStreamFallbackError + from litellm.llms.bedrock.chat.invoke_handler import BedrockError + + async def _raise_400(**kwargs): + raise BedrockError( + status_code=400, + message="validationException malformed input.", + ) + + response = CustomStreamWrapper( + completion_stream=None, + model="anthropic.claude-3-sonnet-20240229-v1:0", + logging_obj=logging_obj, + custom_llm_provider="bedrock", + make_call=_raise_400, + ) + + with pytest.raises(Exception) as excinfo: + await response.__anext__() + assert not isinstance(excinfo.value, MidStreamFallbackError) + assert getattr(excinfo.value, "status_code", None) == 400 + + @pytest.mark.asyncio async def test_async_streaming_read_timeout_triggers_midstream_fallback( logging_obj: Logging, @@ -1801,6 +1911,344 @@ def test_usage_only_chunk_not_dropped_when_finish_reason_already_set( assert result.usage is not None +def _run_dispatch(wrapper: CustomStreamWrapper, chunk): + model_response = wrapper.model_response_creator() + completion_obj = {"content": ""} + result = wrapper._dispatch_provider_chunk( + chunk=chunk, + model_response=model_response, + completion_obj=completion_obj, + ) + return result, model_response, completion_obj + + +def test_dispatch_vllm_extracts_output_text( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """vllm chunks expose text at chunk[0].outputs[0].text; the dispatch must + surface that as the content and report a parsed result.""" + initialized_custom_stream_wrapper.custom_llm_provider = "vllm" + + class _Output: + text = "hello from vllm" + + class _VLLMChunk: + outputs = [_Output()] + + result, _, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, [_VLLMChunk()] + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "hello from vllm" + + +def test_dispatch_petals_slices_completion_stream( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """petals fakes streaming by slicing 30 chars off the buffered completion + stream each call, leaving the remainder for the next chunk.""" + initialized_custom_stream_wrapper.custom_llm_provider = "petals" + initialized_custom_stream_wrapper.completion_stream = "A" * 50 + + result, _, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, chunk=None + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "A" * 30 + assert initialized_custom_stream_wrapper.completion_stream == "A" * 20 + assert initialized_custom_stream_wrapper.received_finish_reason is None + + +def test_dispatch_petals_empty_stream_sets_stop( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """An exhausted petals stream marks the turn finished with a stop reason.""" + initialized_custom_stream_wrapper.custom_llm_provider = "petals" + initialized_custom_stream_wrapper.completion_stream = "" + + result, _, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, chunk=None + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "" + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + + +def test_dispatch_petals_empty_stream_after_finish_raises( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """Once petals has already finished, an empty stream signals end-of-iteration.""" + initialized_custom_stream_wrapper.custom_llm_provider = "petals" + initialized_custom_stream_wrapper.completion_stream = "" + initialized_custom_stream_wrapper.received_finish_reason = "stop" + + with pytest.raises(StopIteration): + _run_dispatch(initialized_custom_stream_wrapper, chunk=None) + + +def test_dispatch_palm_slices_completion_stream( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """palm uses the same fake-streaming slice strategy as petals.""" + initialized_custom_stream_wrapper.custom_llm_provider = "palm" + initialized_custom_stream_wrapper.completion_stream = "B" * 40 + + result, _, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, chunk=None + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "B" * 30 + assert initialized_custom_stream_wrapper.completion_stream == "B" * 10 + + +def test_dispatch_cached_response_extracts_delta( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """cached_response replays a stored ModelResponseStream; the dispatch lifts + its delta content, finish_reason and id back onto the live response.""" + initialized_custom_stream_wrapper.custom_llm_provider = "cached_response" + chunk = ModelResponseStream( + id="chatcmpl-cache-1", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="cached text"), + finish_reason="stop", + ) + ], + ) + + result, model_response, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, chunk + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "cached text" + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + assert model_response.id == "chatcmpl-cache-1" + assert initialized_custom_stream_wrapper.response_id == "chatcmpl-cache-1" + + +def test_dispatch_vertex_ai_legacy_text_and_finish_reason( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """Legacy vertex_ai chunks (non-ModelResponseStream) expose .text and a + candidate finish_reason enum that must be normalised to an OpenAI reason.""" + initialized_custom_stream_wrapper.custom_llm_provider = "vertex_ai" + + class _FinishReason: + name = "STOP" + + class _Candidate: + finish_reason = _FinishReason() + + class _VertexChunk: + candidates = [_Candidate()] + text = "vertex content" + + result, _, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, _VertexChunk() + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "vertex content" + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + + +def test_dispatch_vertex_ai_legacy_without_candidates_stringifies_chunk( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """A legacy vertex_ai chunk with no candidates falls back to str(chunk).""" + initialized_custom_stream_wrapper.custom_llm_provider = "vertex_ai" + + class _RawChunk: + def __str__(self) -> str: + return "raw vertex blob" + + result, _, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, _RawChunk() + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "raw vertex blob" + + +def test_dispatch_vertex_ai_legacy_function_call( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """A legacy vertex_ai chunk whose part has no text but carries a + function_call is converted into an OpenAI tool-call delta.""" + initialized_custom_stream_wrapper.custom_llm_provider = "vertex_ai" + + class _FunctionCall: + name = "get_weather" + args = {"location": "SF"} + + class _Part: + function_call = _FunctionCall() + + class _Content: + parts = [_Part()] + + class _FinishReason: + name = "STOP" + + class _Candidate: + content = _Content() + finish_reason = _FinishReason() + + class _VertexFunctionChunk: + candidates = [_Candidate()] + + @property + def text(self): + raise RuntimeError("Part has no text.") + + result, _, _ = _run_dispatch( + initialized_custom_stream_wrapper, _VertexFunctionChunk() + ) + + assert isinstance(result, _ProviderChunkParsed) + tool_calls = result.response_obj["original_chunk"].choices[0].delta.tool_calls + assert tool_calls[0].function.name == "get_weather" + assert json.loads(tool_calls[0].function.arguments) == {"location": "SF"} + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + + +def test_dispatch_custom_provider_returns_chunk_early( + monkeypatch: pytest.MonkeyPatch, + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """A registered custom provider passes its already-OpenAI-shaped chunk + straight through as an early return rather than re-parsing it.""" + monkeypatch.setattr(litellm, "_custom_providers", ["my-custom-llm"]) + initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-llm" + chunk = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="hi"), finish_reason=None) + ] + ) + + result, _, _ = _run_dispatch(initialized_custom_stream_wrapper, chunk) + + assert isinstance(result, _ProviderChunkEarlyReturn) + assert result.value is chunk + + +def test_dispatch_custom_provider_finish_only_returns_none_early( + monkeypatch: pytest.MonkeyPatch, + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """A custom-provider chunk that carries only a finish_reason (no content) + records the reason and returns None so no empty delta is emitted.""" + monkeypatch.setattr(litellm, "_custom_providers", ["my-custom-llm"]) + initialized_custom_stream_wrapper.custom_llm_provider = "my-custom-llm" + chunk = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop") + ] + ) + + result, _, _ = _run_dispatch(initialized_custom_stream_wrapper, chunk) + + assert isinstance(result, _ProviderChunkEarlyReturn) + assert result.value is None + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + + +def test_dispatch_text_completion_codestral_parses_chunk( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """text-completion-codestral streams raw SSE JSON strings that the dispatch + routes through CodestralTextCompletionConfig to extract content/finish.""" + initialized_custom_stream_wrapper.custom_llm_provider = "text-completion-codestral" + chunk = json.dumps( + {"choices": [{"delta": {"content": "codestral text"}, "finish_reason": "stop"}]} + ) + + result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "codestral text" + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + + +def test_dispatch_text_completion_codestral_requires_string( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """The codestral branch only knows how to parse raw strings; anything else + is a programming error and must surface loudly.""" + initialized_custom_stream_wrapper.custom_llm_provider = "text-completion-codestral" + + with pytest.raises(ValueError): + _run_dispatch(initialized_custom_stream_wrapper, {"not": "a string"}) + + +def test_dispatch_triton_stream( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """triton stream chunks arrive as dicts keyed by text_output/stop_reason.""" + initialized_custom_stream_wrapper.custom_llm_provider = "triton" + chunk = {"text_output": "triton text", "is_finished": True, "stop_reason": "stop"} + + result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "triton text" + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + + +def test_dispatch_ai21_decodes_completion( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """ai21 does fake streaming over a single byte-encoded JSON completion.""" + initialized_custom_stream_wrapper.custom_llm_provider = "ai21" + chunk = json.dumps({"completions": [{"data": {"text": "ai21 text"}}]}).encode( + "utf-8" + ) + + result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "ai21 text" + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + + +def test_dispatch_text_completion_openai_with_usage( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """text-completion-openai chunks expose choices[].text and an optional usage + block that the dispatch lifts onto the model response.""" + initialized_custom_stream_wrapper.custom_llm_provider = "text-completion-openai" + + class _Choice: + text = "oai text" + finish_reason = "stop" + + class _Usage: + prompt_tokens = 5 + completion_tokens = 3 + total_tokens = 8 + + class _TextChunk: + choices = [_Choice()] + usage = _Usage() + + result, model_response, completion_obj = _run_dispatch( + initialized_custom_stream_wrapper, _TextChunk() + ) + + assert isinstance(result, _ProviderChunkParsed) + assert completion_obj["content"] == "oai text" + assert initialized_custom_stream_wrapper.received_finish_reason == "stop" + assert model_response.usage.prompt_tokens == 5 + assert model_response.usage.total_tokens == 8 + + @pytest.mark.asyncio async def test_custom_stream_wrapper_anext_does_not_block_event_loop_for_sync_iterators( logging_obj: Logging, @@ -2306,7 +2754,9 @@ def test_chunk_creator_tool_calls_not_dropped_on_finish( tool_calls=[ ChatCompletionDeltaToolCall( id="call_abc", - function=Function(name="get_weather", arguments='{"city":"NYC"}'), + function=Function( + name="get_weather", arguments='{"city":"NYC"}' + ), type="function", index=0, ) @@ -2401,3 +2851,131 @@ def test_record_partial_usage_for_failure_noop_without_chunks(): wrapper._record_partial_usage_for_failure() assert "combined_usage_object" not in logging_obj.model_call_details + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_chunk_builder_raise_at_end_of_stream_still_recovers_usage( + sync_mode, +): + """stream_chunk_builder re-raises (as APIError) on large agentic tool-use + streams. That raise originates inside the except-StopIteration handler, so + before the fix it escaped __next__/__anext__ and the request was dropped from + SpendLogs while the provider billed the tokens. The wrapper must catch it and + recover usage from the raw chunks so cost is still tracked.""" + final_usage_block = Usage( + completion_tokens=392, prompt_tokens=1799, total_tokens=2191 + ) + final_chunk = ModelResponseStream( + id="chatcmpl-raise-test", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="", role="assistant"), + ) + ], + usage=final_usage_block, + ) + test_chunks = bedrock_chunks + [final_chunk] + + logging_obj = Logging( + model="bedrock/claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="raise-test", + function_id="1245", + ) + + response = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=test_chunks), + model="bedrock/claude-haiku-4-5-20251001-v1:0", + custom_llm_provider="bedrock", + logging_obj=logging_obj, + stream_options={"include_usage": True}, + ) + + seen_usage = [] + with patch.object( + litellm, + "stream_chunk_builder", + side_effect=Exception("simulated assembly failure"), + ): + # before the fix this raised and dropped the request; it must not raise now + if sync_mode: + for chunk in response: + if getattr(chunk, "usage", None) is not None: + seen_usage.append(chunk.usage) + else: + async for chunk in response: + if getattr(chunk, "usage", None) is not None: + seen_usage.append(chunk.usage) + + assert any( + u.total_tokens == final_usage_block.total_tokens for u in seen_usage + ), "usage recovered from raw chunks was not emitted after stream_chunk_builder raised" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_chunk_builder_raise_and_usage_recovery_failure_does_not_crash( + sync_mode, +): + """If end-of-stream assembly raises AND best-effort usage recovery from the raw + chunks also fails, the stream must still complete cleanly rather than propagate + the exception to the consumer.""" + from litellm.litellm_core_utils import streaming_handler as sh_module + + final_chunk = ModelResponseStream( + id="chatcmpl-raise-recover-fail", + created=1742056047, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="", role="assistant"), + ) + ], + usage=Usage(completion_tokens=1, prompt_tokens=1, total_tokens=2), + ) + + response = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=bedrock_chunks + [final_chunk] + ), + model="bedrock/claude-haiku-4-5-20251001-v1:0", + custom_llm_provider="bedrock", + logging_obj=Logging( + model="bedrock/claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="raise-recover-fail", + function_id="1245", + ), + stream_options={"include_usage": True}, + ) + + with ( + patch.object( + litellm, "stream_chunk_builder", side_effect=Exception("assembly failed") + ), + patch.object( + sh_module, "calculate_total_usage", side_effect=Exception("recovery failed") + ), + ): + # must not raise even though both assembly and recovery fail + if sync_mode: + chunks = [c for c in response] + else: + chunks = [c async for c in response] + + assert len(chunks) > 0 diff --git a/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py b/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py index 2cb7b4db3d4..31047d30970 100644 --- a/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py +++ b/tests/test_litellm/llms/anthropic/messages/test_advisor_orchestration.py @@ -516,3 +516,217 @@ async def test_max_uses_none_falls_back_to_default(): ) assert str(_c.ADVISOR_MAX_USES) in str(exc_info.value) + + +# --------------------------------------------------------------------------- +# 12. Defense-in-depth: client-supplied advisor api_base/api_key are dropped +# unless the proxy admin opted into clientside credentials +# --------------------------------------------------------------------------- + + +ADVISOR_TOOL_WITH_CREDS = { + "type": "advisor_20260301", + "name": "advisor", + "model": "claude-opus-4-6", + "api_base": "https://other.example", + "api_key": "sk-other", +} + + +async def _run_advisor_and_capture_subcall_kwargs(): + """Run one advisor turn and return the kwargs of the advisor sub-call.""" + from litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor import ( + AdvisorOrchestrationHandler, + ) + + advisor_tool_use_resp = _make_advisor_tool_use_response(tool_id="toolu_01") + advisor_advice_resp = _make_text_response("advice", model="claude-opus-4-6") + final_resp = _make_text_response("final answer") + + captured = {} + call_count = 0 + + async def mock_call(model, messages, tools, stream, max_tokens, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return advisor_tool_use_resp + if call_count == 2: + # The advisor sub-call — capture its routing kwargs. + captured["api_key"] = kwargs.get("api_key") + captured["api_base"] = kwargs.get("api_base") + return advisor_advice_resp + return final_resp + + with patch( + "litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._call_messages_handler", + side_effect=mock_call, + ): + h = AdvisorOrchestrationHandler() + await h.handle( + model="openai/gpt-4o-mini", + messages=MESSAGES, + tools=[ADVISOR_TOOL_WITH_CREDS], + stream=False, + max_tokens=512, + custom_llm_provider="openai", + ) + return captured + + +@pytest.mark.asyncio +async def test_advisor_creds_dropped_when_proxy_opt_in_disabled(): + """On the proxy without opt-in, the caller's advisor api_base/api_key must + NOT reach the sub-call (would redirect it / leak the server key).""" + with patch( + "litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._allow_client_side_advisor_credentials", + return_value=False, + ): + captured = await _run_advisor_and_capture_subcall_kwargs() + assert captured["api_key"] is None + assert captured["api_base"] is None + + +@pytest.mark.asyncio +async def test_advisor_creds_honored_when_proxy_opt_in_enabled(): + """With the admin opt-in, the documented clientside routing still works.""" + with patch( + "litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor._allow_client_side_advisor_credentials", + return_value=True, + ): + captured = await _run_advisor_and_capture_subcall_kwargs() + assert captured["api_key"] == "sk-other" + assert captured["api_base"] == "https://other.example" + + +# --------------------------------------------------------------------------- +# 13. The proxy gate itself: _allow_client_side_advisor_credentials() and the +# full handle() driven by the real proxy general_settings flag. +# --------------------------------------------------------------------------- + + +def _fake_proxy_server(general_settings: Dict): + """A stand-in litellm.proxy.proxy_server module exposing general_settings. + + The real proxy_server pulls in heavy optional deps that may be absent in a + unit-test environment, so the gate's + ``from litellm.proxy.proxy_server import general_settings`` is satisfied by + injecting this lightweight module into sys.modules. + """ + import types + + module = types.ModuleType("litellm.proxy.proxy_server") + module.general_settings = general_settings # type: ignore[attr-defined] + return module + + +def test_allow_client_side_advisor_credentials_reads_proxy_flag(): + """The gate mirrors the proxy's allow_client_side_credentials opt-in.""" + import sys + + from litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor import ( + _allow_client_side_advisor_credentials, + ) + + cases = ( + ({"allow_client_side_credentials": True}, True), + ({"allow_client_side_credentials": False}, False), + # Flag absent entirely -> default deny on the proxy. + ({}, False), + ) + for settings, expected in cases: + with patch.dict( + sys.modules, + {"litellm.proxy.proxy_server": _fake_proxy_server(settings)}, + ): + assert _allow_client_side_advisor_credentials() is expected + + +def test_allow_client_side_advisor_credentials_defaults_true_outside_proxy(): + """Outside the proxy (proxy_server import unavailable), there is no admin + boundary, so the gate permits client-supplied routing.""" + import builtins + import sys + + from litellm.llms.anthropic.experimental_pass_through.messages.interceptors.advisor import ( + _allow_client_side_advisor_credentials, + ) + + real_import = builtins.__import__ + + def _blocked_import(name, *args, **kwargs): + if name == "litellm.proxy.proxy_server": + raise ImportError("proxy server unavailable") + return real_import(name, *args, **kwargs) + + with patch.dict(sys.modules): + sys.modules.pop("litellm.proxy.proxy_server", None) + with patch.object(builtins, "__import__", _blocked_import): + assert _allow_client_side_advisor_credentials() is True + + +def test_advisor_gate_propagates_non_import_errors(): + """Non-ImportError failures during the proxy module probe must not + default permissive. If the proxy is partially loaded and raises + RuntimeError, the gate should surface that rather than silently + returning True.""" + import sys + + from litellm.llms.anthropic.experimental_pass_through.messages.interceptors import ( + advisor, + ) + + original = sys.modules.get("litellm.proxy.proxy_server") + + class _Broken: + def __getattr__(self, _name): + raise RuntimeError("partial proxy boot") + + sys.modules["litellm.proxy.proxy_server"] = _Broken() + try: + with pytest.raises(RuntimeError, match="partial proxy boot"): + advisor._allow_client_side_advisor_credentials() + finally: + if original is None: + sys.modules.pop("litellm.proxy.proxy_server", None) + else: + sys.modules["litellm.proxy.proxy_server"] = original + + +@pytest.mark.asyncio +async def test_advisor_ignores_tool_credentials_when_clientside_disabled(): + """Driven by the real proxy flag (not a patched gate): with + allow_client_side_credentials False, the tool-supplied api_base/api_key must + not reach the advisor sub-call.""" + import sys + + with patch.dict( + sys.modules, + { + "litellm.proxy.proxy_server": _fake_proxy_server( + {"allow_client_side_credentials": False} + ) + }, + ): + captured = await _run_advisor_and_capture_subcall_kwargs() + assert captured["api_key"] is None + assert captured["api_base"] is None + + +@pytest.mark.asyncio +async def test_advisor_uses_tool_credentials_when_clientside_enabled(): + """Driven by the real proxy flag: with allow_client_side_credentials True, + the tool-supplied api_base/api_key flow through to the advisor sub-call.""" + import sys + + with patch.dict( + sys.modules, + { + "litellm.proxy.proxy_server": _fake_proxy_server( + {"allow_client_side_credentials": True} + ) + }, + ): + captured = await _run_advisor_and_capture_subcall_kwargs() + assert captured["api_key"] == "sk-other" + assert captured["api_base"] == "https://other.example" diff --git a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py b/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py index 32838701949..bc26268ee92 100644 --- a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py +++ b/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py @@ -66,9 +66,8 @@ class TestAPISerpentConfig: assert headers["X-API-Key"] == "test-api-key" assert headers["Content-Type"] == "application/json" - @patch("litellm.llms.apiserpent.search.transformation.get_secret_str") - def test_validate_environment_without_api_key(self, mock_get_secret): - mock_get_secret.return_value = None + def test_validate_environment_without_api_key(self, monkeypatch): + monkeypatch.delenv("APISERPENT_API_KEY", raising=False) with pytest.raises(ValueError, match="APISERPENT_API_KEY is not set"): APISerpentSearchConfig().validate_environment({}) diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py new file mode 100644 index 00000000000..a1353d57038 --- /dev/null +++ b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py @@ -0,0 +1,329 @@ +""" +Regression tests for the host-aware server-credential fallback guard in +``BaseSearchConfig``. + +A caller-supplied ``api_base`` is honored when building the request URL, so +falling back to a server-configured secret while the caller controls the host +would send the operator's credential to an attacker. The guard must refuse that +combination for every provider that carries a server-managed secret, while +leaving keyless providers and legitimate operator overrides untouched. +""" + +from typing import Dict, Tuple, Type +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.llms.apiserpent.search.transformation import APISerpentSearchConfig +from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig, + _is_trusted_search_api_base, +) +from litellm.llms.brave.search.transformation import BraveSearchConfig +from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig +from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig +from litellm.llms.fastcrw.search.transformation import FastCRWSearchConfig +from litellm.llms.firecrawl.search.transformation import FirecrawlSearchConfig +from litellm.llms.google_pse.search.transformation import GooglePSESearchConfig +from litellm.llms.linkup.search.transformation import LinkupSearchConfig +from litellm.llms.parallel_ai.search.transformation import ParallelAISearchConfig +from litellm.llms.perplexity.search.transformation import PerplexitySearchConfig +from litellm.llms.searchapi.search.transformation import SearchAPIConfig +from litellm.llms.searxng.search.transformation import SearXNGSearchConfig +from litellm.llms.serper.search.transformation import SerperSearchConfig +from litellm.llms.tavily.search.transformation import TavilySearchConfig +from litellm.llms.tinyfish.search.transformation import TinyfishSearchConfig +from litellm.llms.you_com.search.transformation import YouComSearchConfig + +ATTACKER_BASE = "https://attacker.example.com" + +# Every *_API_BASE override env var that could otherwise mark the attacker host +# as trusted; cleared before each test so the suite is hermetic. +_BASE_ENV_VARS = ( + "SERPER_API_BASE", + "TAVILY_API_BASE", + "PERPLEXITY_API_BASE", + "APISERPENT_API_BASE", + "EXA_API_BASE", + "BRAVE_API_BASE", + "FIRECRAWL_API_BASE", + "LINKUP_API_BASE", + "SEARCHAPI_API_BASE", + "GOOGLE_PSE_API_BASE", + "PARALLEL_AI_API_BASE", + "YOUCOM_API_BASE", + "SEARXNG_API_BASE", + "DATAFORSEO_API_BASE", + "TINYFISH_API_BASE", + "CRW_API_BASE", +) + + +@pytest.fixture(autouse=True) +def _clear_base_overrides(monkeypatch: pytest.MonkeyPatch) -> None: + for var in _BASE_ENV_VARS: + monkeypatch.delenv(var, raising=False) + + +# (config, {server secret env vars}, caller_api_key honored as-is, extra env for full validate) +ProviderSpec = Tuple[Type[BaseSearchConfig], Dict[str, str], str, Dict[str, str]] + +PROVIDERS: Tuple[ProviderSpec, ...] = ( + (SerperSearchConfig, {"SERPER_API_KEY": "srv"}, "caller-key", {}), + (TavilySearchConfig, {"TAVILY_API_KEY": "srv"}, "caller-key", {}), + (PerplexitySearchConfig, {"PERPLEXITYAI_API_KEY": "srv"}, "caller-key", {}), + (APISerpentSearchConfig, {"APISERPENT_API_KEY": "srv"}, "caller-key", {}), + (ExaAISearchConfig, {"EXA_API_KEY": "srv"}, "caller-key", {}), + (BraveSearchConfig, {"BRAVE_API_KEY": "srv"}, "caller-key", {}), + (FirecrawlSearchConfig, {"FIRECRAWL_API_KEY": "srv"}, "caller-key", {}), + (LinkupSearchConfig, {"LINKUP_API_KEY": "srv"}, "caller-key", {}), + (SearchAPIConfig, {"SEARCHAPI_API_KEY": "srv"}, "caller-key", {}), + ( + GooglePSESearchConfig, + {"GOOGLE_PSE_API_KEY": "srv"}, + "caller-key", + {"GOOGLE_PSE_ENGINE_ID": "engine"}, + ), + (ParallelAISearchConfig, {"PARALLEL_API_KEY": "srv"}, "caller-key", {}), + (YouComSearchConfig, {"YOUCOM_API_KEY": "srv"}, "caller-key", {}), + (SearXNGSearchConfig, {"SEARXNG_API_KEY": "srv"}, "caller-key", {}), + ( + DataForSEOSearchConfig, + {"DATAFORSEO_LOGIN": "srv", "DATAFORSEO_PASSWORD": "pw"}, + "login:password", + {}, + ), + (TinyfishSearchConfig, {"TINYFISH_API_KEY": "srv"}, "caller-key", {}), + (FastCRWSearchConfig, {"CRW_API_KEY": "srv"}, "caller-key", {}), +) + +_IDS = tuple(spec[0].__name__ for spec in PROVIDERS) + + +@pytest.mark.parametrize( + "config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS +) +def test_server_secret_refused_for_caller_api_base( + config_cls: Type[BaseSearchConfig], + server_env: Dict[str, str], + caller_key: str, + extra_env: Dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + for key, value in {**server_env, **extra_env}.items(): + monkeypatch.setenv(key, value) + + with pytest.raises(ValueError, match="Refusing to send the server-configured"): + config_cls().validate_environment(headers={}, api_base=ATTACKER_BASE) + + +@pytest.mark.parametrize( + "config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS +) +def test_caller_supplied_key_is_honored_for_custom_api_base( + config_cls: Type[BaseSearchConfig], + server_env: Dict[str, str], + caller_key: str, + extra_env: Dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + for key, value in {**server_env, **extra_env}.items(): + monkeypatch.setenv(key, value) + + # An explicit caller key is the caller's own credential, so pointing it at + # the caller's own host must be allowed. + config_cls().validate_environment( + headers={}, api_key=caller_key, api_base=ATTACKER_BASE + ) + + +@pytest.mark.parametrize( + "config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS +) +def test_server_secret_used_without_caller_api_base( + config_cls: Type[BaseSearchConfig], + server_env: Dict[str, str], + caller_key: str, + extra_env: Dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + for key, value in {**server_env, **extra_env}.items(): + monkeypatch.setenv(key, value) + + # No caller-supplied api_base -> the request targets the trusted default, so + # the server secret is still used and nothing is refused. + config_cls().validate_environment(headers={}) + + +def test_keyless_provider_allows_caller_api_base( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("SEARXNG_API_KEY", raising=False) + + headers = SearXNGSearchConfig().validate_environment( + headers={}, api_base="https://my-searxng.internal" + ) + + assert "Authorization" not in headers + + +def test_operator_env_base_override_is_trusted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SERPER_API_KEY", "srv") + monkeypatch.setenv("SERPER_API_BASE", "https://serper.internal.corp") + + # Mirrors the second validate_environment call in the search handler, which + # receives the already-resolved operator base as api_base. + headers = SerperSearchConfig().validate_environment( + headers={}, api_base="https://serper.internal.corp/search" + ) + + assert headers["X-API-KEY"] == "srv" + + +class TestResolveServerApiKey: + def test_caller_key_short_circuits(self) -> None: + result = BaseSearchConfig().resolve_server_api_key( + caller_api_key="mine", + caller_api_base=ATTACKER_BASE, + key_env_vars=("SERPER_API_KEY",), + base_env_var="SERPER_API_BASE", + default_api_base="https://google.serper.dev", + ) + assert result == "mine" + + def test_returns_none_when_no_server_secret( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("SEARXNG_API_KEY", raising=False) + result = BaseSearchConfig().resolve_server_api_key( + caller_api_key=None, + caller_api_base=ATTACKER_BASE, + key_env_vars=("SEARXNG_API_KEY",), + base_env_var="SEARXNG_API_BASE", + default_api_base=None, + ) + assert result is None + + def test_first_set_env_var_wins(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PARALLEL_AI_API_KEY", raising=False) + monkeypatch.setenv("PARALLEL_API_KEY", "second") + result = BaseSearchConfig().resolve_server_api_key( + caller_api_key=None, + caller_api_base=None, + key_env_vars=("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"), + base_env_var="PARALLEL_AI_API_BASE", + default_api_base="https://api.parallel.ai", + ) + assert result == "second" + + +class TestIsTrustedSearchApiBase: + def test_matches_default_host(self) -> None: + assert _is_trusted_search_api_base( + "https://google.serper.dev/search", "https://google.serper.dev", None + ) + + def test_foreign_host_untrusted(self) -> None: + assert not _is_trusted_search_api_base( + ATTACKER_BASE, "https://google.serper.dev", None + ) + + def test_env_override_host_trusted(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SERPER_API_BASE", "https://serper.internal.corp") + assert _is_trusted_search_api_base( + "https://serper.internal.corp/search", + "https://google.serper.dev", + "SERPER_API_BASE", + ) + + def test_schemeless_candidate_untrusted(self) -> None: + # Without a scheme urlsplit puts the value in the path, leaving an empty + # netloc; an unparseable host must never be treated as trusted. + assert not _is_trusted_search_api_base( + "attacker.example.com", "https://google.serper.dev", None + ) + + +@pytest.mark.asyncio +async def test_asearch_does_not_leak_server_key_to_caller_api_base( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """End-to-end regression on the reported vector: a search call with a foreign + api_base and no caller key must fail without any outbound request carrying the + server-configured key.""" + monkeypatch.setenv("SERPER_API_KEY", "sk-server-secret") + monkeypatch.delenv("SERPER_API_BASE", raising=False) + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", + new_callable=AsyncMock, + ) as mock_get, + ): + with pytest.raises(Exception): + await litellm.asearch( + query="secrets", + search_provider="serper", + api_base=ATTACKER_BASE, + ) + + mock_post.assert_not_called() + mock_get.assert_not_called() + + +@pytest.mark.parametrize( + "provider, key_env, server_key, extra_env", + [ + ("searchapi", "SEARCHAPI_API_KEY", "sk-server-searchapi", {}), + ( + "google_pse", + "GOOGLE_PSE_API_KEY", + "sk-server-google", + {"GOOGLE_PSE_ENGINE_ID": "engine-id"}, + ), + ], +) +@pytest.mark.asyncio +async def test_query_param_key_not_leaked_with_dummy_caller_key( + provider: str, + key_env: str, + server_key: str, + extra_env: Dict[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Providers that send the key as a URL query param resolve it in + transform_search_request, not validate_environment. A caller who passes a + dummy api_key to clear the validate_environment short-circuit must not cause + the server key to be placed in the URL sent to their own api_base.""" + monkeypatch.setenv(key_env, server_key) + for name, value in extra_env.items(): + monkeypatch.setenv(name, value) + + captured: Dict[str, str] = {} + + async def fake_get(self, *args, **kwargs): # type: ignore[no-untyped-def] + captured["url"] = kwargs.get("url") or (args[0] if args else "") + raise RuntimeError("stop after capturing the outbound url") + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", + fake_get, + ): + with pytest.raises(Exception): + await litellm.asearch( + query="secrets", + search_provider=provider, + api_key="sk-CALLER-DUMMY", + api_base=ATTACKER_BASE, + ) + + assert captured["url"], "expected an outbound request to be attempted" + assert server_key not in captured["url"] + assert "sk-CALLER-DUMMY" in captured["url"] diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 3f91f6ac26e..2d5242d510f 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -163,6 +163,134 @@ def test_aws_profile_path_not_cached_in_iam_cache(): assert mock_profile.call_count == 2 +def test_get_credentials_does_not_expand_request_env_reference(): + """ + A parameter of the form os.environ/ reaching get_credentials is left as-is + rather than expanded against the process environment, so the downstream auth + helper only ever receives the literal value. + """ + env = _os_environ_without_aws_keys() + env["SERVER_ONLY_VALUE"] = "config-managed-value" + base = BaseAWSLLM() + with patch.dict(os.environ, env, clear=True), patch.object( + base, + "_auth_with_aws_profile", + return_value=(Credentials("ak", "sk", None), None), + ) as mock_profile: + base.get_credentials(aws_profile_name="os.environ/SERVER_ONLY_VALUE") + + assert mock_profile.call_args.args[0] == "os.environ/SERVER_ONLY_VALUE" + assert "config-managed-value" not in str(mock_profile.call_args) + + +def test_get_credentials_falls_back_to_ambient_aws_profile_name_env(): + """ + The fixed AWS_* ambient fallback keeps working: an unset aws_profile_name + resolves from the AWS_PROFILE_NAME environment variable. + """ + env = _os_environ_without_aws_keys() + env["AWS_PROFILE_NAME"] = "ambient-profile" + base = BaseAWSLLM() + with patch.dict(os.environ, env, clear=True), patch.object( + base, + "_auth_with_aws_profile", + return_value=(Credentials("ak", "sk", None), None), + ) as mock_profile: + base.get_credentials(aws_profile_name=None) + + assert mock_profile.call_args.args[0] == "ambient-profile" + + +def test_get_credentials_ambient_fallback_resolves_aws_external_id(): + """ + Each unset param falls back to its own AWS_* env var. Regression for an index + misalignment between the value list and the env-name list, which left + AWS_EXTERNAL_ID unresolved. + """ + env = _os_environ_without_aws_keys() + env["AWS_EXTERNAL_ID"] = "ext-from-env" + base = BaseAWSLLM() + with patch.dict(os.environ, env, clear=True), patch.object( + base, + "_auth_with_aws_role", + return_value=(Credentials("ak", "sk", "tok"), None), + ) as mock_role: + base.get_credentials( + aws_role_name="arn:aws:iam::123456789012:role/x", + aws_session_name="s", + ) + + assert mock_role.call_args.kwargs["aws_external_id"] == "ext-from-env" + + +def _capturing_sts_client(captured: Dict[str, Any]) -> MagicMock: + sts = MagicMock() + + def _assume(**params): + captured["WebIdentityToken"] = params.get("WebIdentityToken") + return { + "Credentials": { + "AccessKeyId": "AKIA", + "SecretAccessKey": "sk", + "SessionToken": "tok", + }, + "PackedPolicySize": 10, + } + + sts.assume_role_with_web_identity.side_effect = _assume + return sts + + +@pytest.mark.parametrize( + "token_ref", + ["os.environ/SERVER_ONLY_VALUE", "SERVER_ONLY_VALUE"], + ids=["os_environ_prefix", "bare_env_name"], +) +def test_web_identity_token_env_reference_not_expanded(token_ref): + """ + A web-identity token that is an environment-variable reference (an os.environ/ + prefix, or a bare name matching an env var) is rejected rather than expanded, so + the process-environment value is never used as the token. + """ + env = _os_environ_without_aws_keys() + env["SERVER_ONLY_VALUE"] = "server-only-value" + captured: Dict[str, Any] = {} + base = BaseAWSLLM() + with patch.dict(os.environ, env, clear=True), patch( + "boto3.client", side_effect=lambda *a, **k: _capturing_sts_client(captured) + ), patch("boto3.Session", return_value=MagicMock()): + with pytest.raises(AwsAuthError): + base.get_credentials( + aws_web_identity_token=token_ref, + aws_role_name="arn:aws:iam::123456789012:role/x", + aws_session_name="s", + aws_sts_endpoint="https://custom-sts.example", + ) + + assert "server-only-value" not in str(captured) + + +def test_web_identity_token_oidc_reference_still_resolved(): + """ + The env-reference guard does not over-reject: an oidc/ reference still flows to + get_secret (mocked to None here), surfacing the existing 401 rather than the 400 + used for rejected env-var references. + """ + base = BaseAWSLLM() + env = _os_environ_without_aws_keys() + with patch.dict(os.environ, env, clear=True), patch( + "litellm.llms.bedrock.base_aws_llm.get_secret", return_value=None + ): + with pytest.raises(AwsAuthError) as exc: + base.get_credentials( + aws_web_identity_token="oidc/circleci/", + aws_role_name="arn:aws:iam::123456789012:role/x", + aws_session_name="s", + ) + + assert exc.value.status_code == 401 + + def test_web_identity_path_not_cached_in_iam_cache(): base = BaseAWSLLM() with patch.object( diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py b/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py index cecb6024de1..1a46015ceb9 100644 --- a/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py +++ b/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py @@ -3,25 +3,190 @@ import pytest from litellm.llms.cloudflare.chat.transformation import CloudflareChatConfig -def test_get_complete_url_encodes_model_path_segment(): +def test_supported_params_include_tools_and_tool_choice(): config = CloudflareChatConfig() - assert ( - config.get_complete_url( - api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", - api_key="cf-key", - model="@cf/meta/llama?x=1#frag", - optional_params={}, - litellm_params={}, - ) - == "https://api.cloudflare.com/client/v4/accounts/acct/ai/run/%40cf/meta/llama%3Fx%3D1%23frag" + params = config.get_supported_openai_params(model="@cf/meta/llama-2-7b-chat-int8") + + assert "tools" in params + assert "tool_choice" in params + assert "stream" in params + assert "max_tokens" in params + + +def test_get_complete_url_defaults_to_openai_compatible_endpoint(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base=None, + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, ) - with pytest.raises(ValueError, match="dot path segment"): + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + assert "/ai/run/" not in url + + +def test_get_complete_url_appends_chat_completions_to_explicit_base(): + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/v1", + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + assert "/ai/run/" not in url + + +def test_get_complete_url_is_idempotent_for_full_base(): + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions", + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + + +def test_get_complete_url_falls_back_to_account_id_when_base_is_empty(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base="", + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + + +def test_get_complete_url_raises_when_account_id_and_base_missing(monkeypatch): + monkeypatch.delenv("CLOUDFLARE_ACCOUNT_ID", raising=False) + config = CloudflareChatConfig() + + with pytest.raises(ValueError, match="Missing CLOUDFLARE_ACCOUNT_ID"): config.get_complete_url( - api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", + api_base=None, api_key="cf-key", - model="../../accounts/other", + model="@cf/meta/llama-2-7b-chat-int8", optional_params={}, litellm_params={}, ) + + +def test_get_complete_url_raises_when_account_id_is_empty(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", " ") + config = CloudflareChatConfig() + + with pytest.raises(ValueError, match="Missing CLOUDFLARE_ACCOUNT_ID"): + config.get_complete_url( + api_base=None, + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + +def test_get_complete_url_migrates_legacy_ai_run_base(): + config = CloudflareChatConfig() + + url = config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", + api_key="cf-key", + model="@cf/meta/llama-2-7b-chat-int8", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/chat/completions" + ) + assert "/ai/run" not in url + + +def test_transform_request_passes_tools_through_in_openai_format(): + config = CloudflareChatConfig() + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } + ] + messages = [{"role": "user", "content": "weather in nyc?"}] + + body = config.transform_request( + model="@cf/meta/llama-2-7b-chat-int8", + messages=messages, + optional_params={"tools": tools, "tool_choice": "auto"}, + litellm_params={}, + headers={}, + ) + + assert body["messages"] == messages + assert body["model"] == "@cf/meta/llama-2-7b-chat-int8" + assert body["tools"] == tools + assert body["tool_choice"] == "auto" + + +def test_validate_environment_requires_api_key(): + config = CloudflareChatConfig() + + with pytest.raises(ValueError, match="Missing Cloudflare API Key"): + config.validate_environment( + headers={}, + model="@cf/meta/llama-2-7b-chat-int8", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + + +def test_validate_environment_sets_bearer_and_content_type(): + config = CloudflareChatConfig() + + headers = config.validate_environment( + headers={}, + model="@cf/meta/llama-2-7b-chat-int8", + messages=[], + optional_params={}, + litellm_params={}, + api_key="cf-key", + ) + + assert headers["Authorization"] == "Bearer cf-key" + assert headers["Content-Type"] == "application/json" diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index bf835b5d8f9..7bd1d7a6031 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -851,3 +851,56 @@ async def test_async_get_forwards_per_request_timeout(): } finally: await handler.close() + + +class TestDefaultCachedClientTimeoutHonorsRequestTimeout: + """Cached default httpx clients must fall back to an explicit litellm.request_timeout. + + Regression for LIT-2369: get_async_httpx_client / _get_httpx_client hardcoded a + 600s default and never consulted litellm.request_timeout, so provider calls with + no per-model timeout (e.g. Bedrock) hung for 600s. + """ + + @pytest.fixture + def restore_request_timeout(self): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + try: + yield + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + def test_default_when_request_timeout_unset(self, restore_request_timeout): + from litellm.llms.custom_httpx.http_handler import ( + _DEFAULT_TIMEOUT, + _default_cached_client_timeout, + ) + + litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = False + assert _default_cached_client_timeout() is _DEFAULT_TIMEOUT + + def test_uses_explicit_request_timeout(self, restore_request_timeout): + from litellm.llms.custom_httpx.http_handler import ( + _default_cached_client_timeout, + ) + + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = True + resolved = _default_cached_client_timeout() + assert resolved.read == 300.0 + assert resolved.connect == 5.0 + + def test_cached_async_client_built_with_explicit_request_timeout( + self, restore_request_timeout + ): + from litellm.caching.llm_caching_handler import LLMClientCache + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.utils import LlmProviders + + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = True + litellm.in_memory_llm_clients_cache = LLMClientCache() + client = get_async_httpx_client(llm_provider=LlmProviders.BEDROCK) + assert client.timeout.read == 300.0 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 f7d445d0788..64ae30daa70 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 @@ -10,13 +10,21 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path import litellm +from litellm.integrations.code_interpreter_interception.handler import ( + CodeInterpreterInterceptionLogger, + LITELLM_CODE_EXECUTION_TOOL_NAME, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, _google_genai_streaming_hidden_params, ) +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams +_ACTIVE_KEY = "_code_interpreter_interception_active" +_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" + def test_prepare_fake_stream_request(): # Initialize the BaseLLMHTTPHandler @@ -116,6 +124,117 @@ def test_response_api_handler_streams_when_provider_transform_adds_stream(): assert client.post.call_args.kwargs["json"]["stream"] is True +def test_response_api_handler_runs_agentic_hooks_in_sync_path(monkeypatch): + handler = BaseLLMHTTPHandler() + config = Mock() + config.validate_environment.return_value = {} + config.get_complete_url.return_value = "https://chatgpt.example.com/responses" + config.transform_responses_api_request.return_value = { + "model": "gpt-5", + "input": "hi", + } + config.sign_request.return_value = ({}, None) + initial_response = Mock() + final_response = Mock() + config.transform_response_api_response.return_value = initial_response + + client = HTTPHandler(client=httpx.Client()) + client.post = Mock( + return_value=httpx.Response( + 200, + request=httpx.Request("POST", "https://chatgpt.example.com/responses"), + ) + ) + logging_obj = Mock() + + monkeypatch.setattr(handler, "_has_agentic_completion_hook", Mock(return_value=True)) + hook_mock = AsyncMock(return_value=final_response) + monkeypatch.setattr(handler, "_call_agentic_completion_hooks", hook_mock) + + response = handler.response_api_handler( + model="gpt-5", + input="hi", + responses_api_provider_config=config, + response_api_optional_request_params={}, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + client=client, + ) + + assert response is final_response + hook_mock.assert_awaited_once() + assert hook_mock.call_args.kwargs["api_surface"] == "responses" + assert hook_mock.call_args.kwargs["messages"] == [ + {"role": "user", "content": "hi"} + ] + + +def test_response_api_handler_runs_responses_pre_call_hook_before_transform(): + handler = BaseLLMHTTPHandler() + config = Mock() + config.validate_environment.return_value = {} + config.get_complete_url.return_value = "https://api.openai.com/v1/responses" + config.sign_request.return_value = ({}, None) + initial_response = ResponsesAPIResponse( + id="resp_1", + created_at=0, + output=[], + status="completed", + model="gpt-5", + ) + config.transform_response_api_response.return_value = initial_response + + def transform_responses_api_request(**kwargs): + return { + "model": kwargs["model"], + "input": kwargs["input"], + **kwargs["response_api_optional_request_params"], + } + + config.transform_responses_api_request.side_effect = transform_responses_api_request + client = HTTPHandler(client=httpx.Client()) + client.post = Mock( + return_value=httpx.Response( + 200, + request=httpx.Request("POST", "https://api.openai.com/v1/responses"), + ) + ) + logging_obj = Mock() + logging_obj.dynamic_success_callbacks = [] + + old_callbacks = list(litellm.callbacks) + litellm.callbacks = [CodeInterpreterInterceptionLogger()] + try: + response = handler.response_api_handler( + model="gpt-5", + input="use code", + responses_api_provider_config=config, + response_api_optional_request_params={ + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}] + }, + custom_llm_provider="openai", + litellm_params=GenericLiteLLMParams(api_key="sk-test"), + logging_obj=logging_obj, + client=client, + ) + finally: + litellm.callbacks = old_callbacks + + assert response is initial_response + transform_kwargs = config.transform_responses_api_request.call_args.kwargs + tools = transform_kwargs["response_api_optional_request_params"]["tools"] + assert not any(tool.get("type") == "code_interpreter" for tool in tools) + assert any( + tool.get("type") == "function" + and tool.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME + for tool in tools + ) + hook_litellm_params = transform_kwargs["litellm_params"] + assert hook_litellm_params.get(_ACTIVE_KEY) is True + assert hook_litellm_params.get(_SANDBOX_KEY) + + @pytest.mark.asyncio async def test_async_response_api_handler_streams_when_provider_transform_adds_stream(): handler = BaseLLMHTTPHandler() diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 683ec158f44..03e763a4161 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1,9 +1,8 @@ import json import os import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch -import httpx import pytest import litellm @@ -12,10 +11,14 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path -from litellm import get_model_info, supports_reasoning +from litellm import get_model_info, supports_reasoning, supports_vision from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig -from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk -from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Function, + Message, + ModelResponse, +) @pytest.fixture(autouse=True) @@ -98,12 +101,12 @@ def test_supports_reasoning_effort(): for model in supported_models: assert ( - supports_reasoning(model=model, custom_llm_provider="fireworks_ai") == True + supports_reasoning(model=model, custom_llm_provider="fireworks_ai") is True ), f"{model} should support reasoning_effort" for model in unsupported_models: assert ( - supports_reasoning(model=model, custom_llm_provider="fireworks_ai") == False + supports_reasoning(model=model, custom_llm_provider="fireworks_ai") is False ), f"{model} should not support reasoning_effort" @@ -115,11 +118,13 @@ def test_get_supported_openai_params_reasoning_effort(): "fireworks_ai/accounts/fireworks/models/glm-5p1" ) assert "reasoning_effort" in supported_params + assert "thinking" in supported_params unsupported_params = config.get_supported_openai_params( "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct" ) assert "reasoning_effort" not in unsupported_params + assert "thinking" not in unsupported_params def test_get_supported_openai_params_parallel_tool_calls(): @@ -181,41 +186,6 @@ def test_get_provider_info_omits_false_supports_reasoning(monkeypatch): assert "supports_reasoning" not in info -def test_add_transform_inline_image_block_skips_data_urls(): - """ - data: URLs must not have #transform=inline appended — doing so corrupts the - base64 payload and raises binascii.Error: Incorrect padding on the Fireworks side. - Regression test for https://github.com/BerriAI/litellm/issues/23583 - """ - config = FireworksAIConfig() - data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRgAB" - - # str branch - str_content = {"type": "image_url", "image_url": data_url} - result = config._add_transform_inline_image_block( - str_content, model="gpt-4", disable_add_transform_inline_image_block=False - ) - assert result["image_url"] == data_url, "data URL must not be modified (str branch)" - - # dict branch - dict_content = {"type": "image_url", "image_url": {"url": data_url}} - result = config._add_transform_inline_image_block( - dict_content, model="gpt-4", disable_add_transform_inline_image_block=False - ) - assert ( - result["image_url"]["url"] == data_url - ), "data URL must not be modified (dict branch)" - - # regular https URL should still get the suffix - https_content = {"type": "image_url", "image_url": "https://example.com/image.jpg"} - result = config._add_transform_inline_image_block( - https_content, model="gpt-4", disable_add_transform_inline_image_block=False - ) - assert result["image_url"].endswith( - "#transform=inline" - ), "https URL should get #transform=inline" - - @pytest.mark.parametrize( "api_base, expected_url_prefix", [ @@ -582,3 +552,549 @@ def test_transform_request_routes_short_form_model_to_models_path(): headers={}, ) assert result["model"] == "accounts/fireworks/models/glm-5p2" + + +def _make_fireworks_raw_response(body: dict) -> MagicMock: + mock = MagicMock() + mock.status_code = 200 + mock.json.return_value = body + mock.text = json.dumps(body) + mock.headers = {} + return mock + + +_BASE_CHAT_COMPLETION_RESPONSE: dict = { + "id": "resp-test", + "object": "chat.completion", + "created": 1234567890, + "model": "glm-5p1", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, +} + + +def _run_transform_response(response_body: dict) -> ModelResponse: + config = FireworksAIConfig() + raw_response = _make_fireworks_raw_response(response_body) + logging_obj = MagicMock() + return config.transform_response( + model="accounts/fireworks/models/glm-5p1", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "Hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key="test-key", + ) + + +_REASONING_MODEL = "fireworks_ai/accounts/fireworks/models/glm-5p1" +_NON_REASONING_MODEL = "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct" + + +def test_get_supported_openai_params_includes_all_fireworks_params(): + config = FireworksAIConfig() + params = config.get_supported_openai_params(_REASONING_MODEL) + + required = [ + "seed", + "top_logprobs", + "min_p", + "typical_p", + "repetition_penalty", + "mirostat_target", + "mirostat_lr", + "logit_bias", + "echo", + "echo_last", + "ignore_eos", + "prompt_cache_key", + "prompt_cache_isolation_key", + "raw_output", + "perf_metrics_in_response", + "return_token_ids", + "safe_tokenization", + "service_tier", + "speculation", + "prediction", + "stream_options", + "sampling_mask", + "thinking", + "reasoning_history", + ] + missing = [p for p in required if p not in params] + assert missing == [], f"Missing params: {missing}" + + +def test_native_openai_params_flow_end_to_end_with_drop_params_false(): + """ + The OpenAI-native params Fireworks supports (seed, top_logprobs, logit_bias, + prompt_cache_key, service_tier, prediction) previously hit + ``UnsupportedParamsError`` with ``drop_params=False`` because they were absent + from ``get_supported_openai_params``. Listing them must let them survive the + ``get_optional_params`` gate and reach the request, not just appear in the + supported list. Asserting via ``get_optional_params`` (the real gate) rather + than ``map_openai_params`` catches a revert of the supported-params additions, + which a list-membership check would not. + """ + native = { + "seed": 42, + "top_logprobs": 3, + "logit_bias": {"1": 1}, + "prompt_cache_key": "cache-key", + "service_tier": "auto", + "prediction": {"type": "content", "content": "x"}, + } + optional_params = litellm.get_optional_params( + model="accounts/fireworks/models/llama-v3-70b-instruct", + custom_llm_provider="fireworks_ai", + drop_params=False, + **native, + ) + for key, value in native.items(): + assert optional_params.get(key) == value + + +def test_prompt_truncate_len_correct_name(): + config = FireworksAIConfig() + params = config.get_supported_openai_params(_REASONING_MODEL) + assert "prompt_truncate_len" in params + assert "prompt_truncate_length" not in params + + result = config.map_openai_params( + {"prompt_truncate_len": 4096}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result == {"prompt_truncate_len": 4096} + + +def test_stream_options_include_usage_auto_injected(): + config = FireworksAIConfig() + result = config.transform_request( + model="accounts/fireworks/models/glm-5p1", + messages=[{"role": "user", "content": "Hi"}], + optional_params={"stream": True}, + litellm_params={}, + headers={}, + ) + assert result["stream_options"] == {"include_usage": True} + + +def test_stream_options_not_injected_when_not_streaming(): + config = FireworksAIConfig() + result = config.transform_request( + model="accounts/fireworks/models/glm-5p1", + messages=[{"role": "user", "content": "Hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "stream_options" not in result + + +def test_stream_options_preserves_user_override(): + config = FireworksAIConfig() + result = config.transform_request( + model="accounts/fireworks/models/glm-5p1", + messages=[{"role": "user", "content": "Hi"}], + optional_params={"stream": True, "stream_options": {"include_usage": False}}, + litellm_params={}, + headers={}, + ) + assert result["stream_options"]["include_usage"] is False + + +def test_reasoning_history_in_supported_params(): + config = FireworksAIConfig() + reasoning_params = config.get_supported_openai_params(_REASONING_MODEL) + assert "reasoning_history" in reasoning_params + + non_reasoning_params = config.get_supported_openai_params(_NON_REASONING_MODEL) + assert "reasoning_history" not in non_reasoning_params + + +def test_thinking_param_passthrough(): + config = FireworksAIConfig() + thinking = {"type": "disabled"} + result = config.map_openai_params( + {"thinking": thinking}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result == {"thinking": thinking} + + +def test_thinking_and_reasoning_effort_conflict_rejected(): + config = FireworksAIConfig() + with pytest.raises( + litellm.BadRequestError, + match="does not support specifying both `thinking` and `reasoning_effort`", + ): + config.map_openai_params( + { + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "reasoning_effort": "medium", + }, + {}, + _REASONING_MODEL, + drop_params=False, + ) + + +def test_minimax_m3_supports_vision_from_model_map(): + config = FireworksAIConfig() + + for model in [ + "fireworks_ai/accounts/fireworks/models/minimax-m3", + "fireworks_ai/minimax-m3", + ]: + assert supports_vision(model=model, custom_llm_provider="fireworks_ai") is True + assert config.get_provider_info(model)["supports_vision"] is True + + +def test_transform_messages_helper_rejects_file_blocks(): + config = FireworksAIConfig() + messages = [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": "data:application/pdf;base64,JVBERi0xLjQKJSVFT0YK", + "filename": "tiny.pdf", + }, + }, + {"type": "text", "text": "Describe this"}, + ], + } + ] + + with pytest.raises( + litellm.BadRequestError, + match="Fireworks AI chat completions does not support file content blocks", + ): + config._transform_messages_helper( + messages, model="accounts/fireworks/models/kimi-k2p6", litellm_params={} + ) + + +def test_transform_messages_helper_rejects_non_vision_image_inputs(): + config = FireworksAIConfig() + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE=" + }, + }, + ], + } + ] + + with pytest.raises(litellm.BadRequestError, match="does not support image inputs"): + config._transform_messages_helper( + messages, model="accounts/fireworks/models/glm-5p2", litellm_params={} + ) + + +def test_transform_messages_helper_allows_vision_image_inputs(): + config = FireworksAIConfig() + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE=" + }, + }, + ], + } + ] + + out = config._transform_messages_helper( + messages, model="accounts/fireworks/models/minimax-m3", litellm_params={} + ) + assert out == messages + + +def test_image_inputs_not_rejected_for_fuzzy_non_vision_match(): + """ + A custom/fine-tuned model id that hyphen-matches a known non-vision model + (glm-5p2 has supports_vision=False) must not inherit that False via the + substring fallback and hard-reject valid image_url blocks. The capability + gate for rejection uses an exact cost-map match; the fuzzy fallback stays a + soft signal only, so an unmapped vision-capable deployment is not blocked. + """ + config = FireworksAIConfig() + custom_model = "accounts/myorg/models/custom-glm-5p2" + + assert config._get_model_cost_capability(custom_model, "supports_vision") is False + assert ( + config._get_model_cost_capability_exact(custom_model, "supports_vision") is None + ) + + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE=" + }, + }, + ], + } + ] + out = config._transform_messages_helper( + messages, model=custom_model, litellm_params={} + ) + assert out == messages + + +def test_transform_messages_helper_skips_non_dict_content(): + config = FireworksAIConfig() + messages = [ + { + "role": "user", + "content": ["just a string", {"type": "text", "text": "hello"}], + } + ] + + out = config._transform_messages_helper( + messages, model="accounts/fireworks/models/glm-5p2", litellm_params={} + ) + assert out == messages + + +def test_transform_messages_helper_no_transform_inline(): + config = FireworksAIConfig() + url = "https://example.com/image.jpg" + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": url}], + } + ] + out = config._transform_messages_helper( + messages, model="accounts/fireworks/models/minimax-m3", litellm_params={} + ) + block = out[0]["content"][0] + assert block["image_url"] == url + assert "#transform=inline" not in block["image_url"] + + +def test_get_provider_info_vision_from_model_cost(monkeypatch): + config = FireworksAIConfig() + + vision_model = "fireworks_ai/test-vision-from-cost" + monkeypatch.setitem( + litellm.model_cost, + vision_model, + {"supports_vision": True, "supports_pdf_input": True}, + ) + info = config.get_provider_info(vision_model) + assert info["supports_vision"] is True + assert info["supports_pdf_input"] is True + + no_vision_model = "fireworks_ai/test-no-vision-from-cost" + monkeypatch.setitem(litellm.model_cost, no_vision_model, {}) + info_no_vision = config.get_provider_info(no_vision_model) + assert info_no_vision.get("supports_vision") is not True + assert "supports_pdf_input" not in info_no_vision + + +def test_reasoning_effort_boolean_true_to_medium(): + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": True}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result["reasoning_effort"] == "medium" + + +def test_reasoning_effort_boolean_false_to_none(): + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": False}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result["reasoning_effort"] == "none" + + +def test_reasoning_effort_string_passthrough(): + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": "high"}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result["reasoning_effort"] == "high" + + +def test_reasoning_effort_integer_passthrough(): + config = FireworksAIConfig() + result = config.map_openai_params( + {"reasoning_effort": 1000}, + {}, + _REASONING_MODEL, + drop_params=False, + ) + assert result["reasoning_effort"] == 1000 + assert isinstance(result["reasoning_effort"], int) + + +def test_transform_response_captures_perf_metrics(): + body = { + **_BASE_CHAT_COMPLETION_RESPONSE, + "perf_metrics": {"prompt-tokens": 10}, + } + result = _run_transform_response(body) + assert result._hidden_params["fireworks_perf_metrics"] == {"prompt-tokens": 10} + + +def test_transform_response_captures_prompt_token_ids(): + body = { + **_BASE_CHAT_COMPLETION_RESPONSE, + "prompt_token_ids": [1, 2, 3], + } + result = _run_transform_response(body) + assert result._hidden_params["fireworks_prompt_token_ids"] == [1, 2, 3] + + +def test_transform_response_captures_raw_output(): + raw_output = { + "prompt_fragments": [], + "prompt_token_ids": [], + "completion": "test", + } + body = { + **_BASE_CHAT_COMPLETION_RESPONSE, + "choices": [ + { + **_BASE_CHAT_COMPLETION_RESPONSE["choices"][0], + "raw_output": raw_output, + } + ], + } + result = _run_transform_response(body) + assert result._hidden_params["fireworks_raw_outputs"] == [raw_output] + + +def test_transform_response_captures_token_ids(): + body = { + **_BASE_CHAT_COMPLETION_RESPONSE, + "choices": [ + { + **_BASE_CHAT_COMPLETION_RESPONSE["choices"][0], + "token_ids": [4, 5, 6], + } + ], + } + result = _run_transform_response(body) + assert result._hidden_params["fireworks_token_ids"] == [[4, 5, 6]] + + +def test_streaming_surfaces_fireworks_response_fields(): + """ + The Fireworks-specific response fields captured into _hidden_params for + non-streaming calls must also reach streamed responses. They ride the + streamed chunks' provider_specific_fields (litellm rebuilds each streamed + chunk, so per-chunk _hidden_params does not survive): per-choice + token_ids/raw_output on the content chunk, response-level + perf_metrics/prompt_token_ids on the final usage chunk. Driving the real + litellm.completion(stream=True) path also covers the get_model_response_iterator + wiring; dropping the Fireworks iterator would leave these fields unset. + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + model = "accounts/fireworks/models/llama-v3p1-8b-instruct" + raw_output = {"completion": "Hi"} + sse_lines = [ + "data: " + + json.dumps( + { + "id": "stream-1", + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "Hi"}, + "token_ids": [123], + "raw_output": raw_output, + } + ], + } + ), + "data: " + + json.dumps( + { + "id": "stream-1", + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 1, + "total_tokens": 6, + }, + "perf_metrics": {"prompt-tokens": 5}, + "prompt_token_ids": [1, 2, 3], + } + ), + "data: [DONE]", + ] + + raw_response = MagicMock() + raw_response.status_code = 200 + raw_response.headers = {} + raw_response.iter_lines = lambda: iter(sse_lines) + + client = HTTPHandler() + with patch.object(client, "post", return_value=raw_response): + stream = litellm.completion( + model=f"fireworks_ai/{model}", + messages=[{"role": "user", "content": "hi"}], + stream=True, + api_key="fw-test-key", + client=client, + ) + surfaced: dict = {} + for chunk in stream: + fields = getattr(chunk, "provider_specific_fields", None) or {} + surfaced.update( + {k: v for k, v in fields.items() if k.startswith("fireworks_")} + ) + + assert surfaced["fireworks_token_ids"] == [[123]] + assert surfaced["fireworks_raw_outputs"] == [raw_output] + assert surfaced["fireworks_perf_metrics"] == {"prompt-tokens": 5} + assert surfaced["fireworks_prompt_token_ids"] == [1, 2, 3] diff --git a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py index 3ee53bb46cd..7a3f372582f 100644 --- a/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/test_litellm/llms/mistral/test_mistral_chat_transformation.py @@ -719,3 +719,93 @@ class TestMistralFileHandling: # Check that file_ids are modified to match Mistral's expected format assert result[0]["content"][1]["file_id"] == "file-12345" # type: ignore assert result[0]["content"][2]["file_id"] == "file-67890" # type: ignore + + +class TestMistralStripsOutputOnlyFields: + """Mistral rejects unknown input fields with a 422 ``extra_forbidden``. + + LiteLLM attaches ``reasoning_content`` / ``thinking_blocks`` to assistant + responses, so replaying an assistant turn verbatim must not forward them. + Regression for https://github.com/BerriAI/litellm/issues/30835. + """ + + def test_assistant_reasoning_content_is_dropped(self): + messages = cast( + List[AllMessageValues], + [ + {"role": "user", "content": "Question?"}, + { + "role": "assistant", + "content": "Follow-up", + "reasoning_content": "Some internal reasoning text.", + "thinking_blocks": [ + {"type": "thinking", "thinking": "step", "signature": "mistral"} + ], + }, + ], + ) + + result = cast( + List[AllMessageValues], + MistralConfig()._transform_messages( + messages=messages, model="mistral-medium-3-5" + ), + ) + + assistant_message = result[-1] + assert "reasoning_content" not in assistant_message + assert "thinking_blocks" not in assistant_message + assert assistant_message["content"] == "Follow-up" + assert assistant_message["role"] == "assistant" + + def test_non_assistant_messages_are_untouched(self): + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": "Question?", "reasoning_content": "noise"}], + ) + + result = cast( + List[AllMessageValues], + MistralConfig()._transform_messages( + messages=messages, model="mistral-medium-3-5" + ), + ) + + assert result[0].get("reasoning_content") == "noise" + + def test_reasoning_content_dropped_when_image_present(self): + """The image branch returns early, so stripping must run before it.""" + messages = cast( + List[AllMessageValues], + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/cat.png"}, + }, + ], + }, + { + "role": "assistant", + "content": "A cat.", + "reasoning_content": "leaked reasoning", + }, + ], + ) + + with patch.object( + MistralConfig, + "_transform_messages_sync", + side_effect=lambda transformed, model: transformed, + ): + result = cast( + List[AllMessageValues], + MistralConfig()._transform_messages( + messages=messages, model="mistral-medium-3-5", is_async=False + ), + ) + + assert "reasoning_content" not in result[-1] diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/test_litellm/llms/openai_like/test_json_providers.py index 39a4964f5f4..c8743e1809d 100644 --- a/tests/test_litellm/llms/openai_like/test_json_providers.py +++ b/tests/test_litellm/llms/openai_like/test_json_providers.py @@ -2,6 +2,7 @@ Tests for JSON-based provider configuration system. """ +import json import os import sys from unittest.mock import MagicMock, patch @@ -244,6 +245,99 @@ class TestPinstripes: assert result["temperature"] == 0.7 +class TestDarkbloom: + def test_darkbloom_json_config_exists(self): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + darkbloom = JSONProviderRegistry.get("darkbloom") + assert darkbloom is not None + assert darkbloom.base_url == "https://api.darkbloom.dev/v1" + assert darkbloom.api_key_env == "DARKBLOOM_API_KEY" + assert darkbloom.api_base_env == "DARKBLOOM_API_BASE" + assert darkbloom.param_mappings.get("max_completion_tokens") == "max_tokens" + + def test_darkbloom_provider_resolution(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="darkbloom/gemma-4-26b", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "gemma-4-26b" + assert provider == "darkbloom" + assert api_key is None + assert api_base == "https://api.darkbloom.dev/v1" + + def test_darkbloom_dynamic_config(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("darkbloom") + config_class = create_config_class(provider) + config = config_class() + + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == "https://api.darkbloom.dev/v1" + + api_base, api_key = config._get_openai_compatible_provider_info( + "https://custom.darkbloom.dev/v1", "test-key" + ) + assert api_base == "https://custom.darkbloom.dev/v1" + assert api_key == "test-key" + + def test_darkbloom_complete_url_appends_endpoint(self): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider = JSONProviderRegistry.get("darkbloom") + config_class = create_config_class(provider) + config = config_class() + + url = config.get_complete_url( + api_base="https://api.darkbloom.dev/v1", + api_key="test-key", + model="darkbloom/gemma-4-26b", + optional_params={}, + litellm_params={}, + stream=True, + ) + + assert url == "https://api.darkbloom.dev/v1/chat/completions" + + def test_darkbloom_provider_config_manager(self): + from litellm import LlmProviders + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_chat_config( + model="gemma-4-26b", provider=LlmProviders.DARKBLOOM + ) + + assert config is not None + assert config.custom_llm_provider == "darkbloom" + + def test_darkbloom_model_cost_map(self): + with open( + os.path.join(workspace_path, "model_prices_and_context_window.json") + ) as f: + model_cost = json.load(f) + + expected_models = { + "darkbloom/gemma-4-26b": (3e-08, 1.65e-07), + "darkbloom/gpt-oss-20b": (1.45e-08, 7e-08), + } + for model, (input_cost, output_cost) in expected_models.items(): + assert model in model_cost + assert model_cost[model]["litellm_provider"] == "darkbloom" + assert model_cost[model]["max_output_tokens"] == 32768 + assert model_cost[model]["supports_function_calling"] is True + assert model_cost[model]["supports_tool_choice"] is True + assert model_cost[model]["input_cost_per_token"] == input_cost + assert model_cost[model]["output_cost_per_token"] == output_cost + + class TestPublicAIIntegration: """Integration tests for PublicAI provider""" diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py index b5c1a86205b..7be295826e3 100644 --- a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py @@ -293,7 +293,10 @@ class TestParallelAISearch: ], ) @pytest.mark.asyncio - async def test_custom_api_base_appends_v1_search(self, api_base): + async def test_custom_api_base_appends_v1_search(self, api_base, monkeypatch): + # Operator points at an internal base via the env override (a trusted + # host), so the server key is still used and the URL is normalized. + monkeypatch.setenv("PARALLEL_AI_API_BASE", api_base) with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", new_callable=AsyncMock, @@ -303,7 +306,6 @@ class TestParallelAISearch: await litellm.asearch( query="AI developments", search_provider="parallel_ai", - api_base=api_base, ) call_args = mock_post.call_args @@ -312,6 +314,23 @@ class TestParallelAISearch: == "https://proxy.internal.example.com/v1/search" ) + @pytest.mark.asyncio + async def test_caller_api_base_without_key_is_refused(self, monkeypatch): + # A caller-supplied api_base (untrusted host) while relying on the + # server key must be refused without any outbound request. + monkeypatch.setenv("PARALLEL_API_KEY", "server-secret") + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + with pytest.raises(Exception, match="Refusing to send"): + await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + api_base="https://attacker.example.com", + ) + mock_post.assert_not_called() + @pytest.mark.asyncio async def test_missing_api_key_raises(self, monkeypatch): monkeypatch.delenv("PARALLEL_API_KEY", raising=False) diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py index e2d1ab72c5e..46c1e457d7c 100644 --- a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py @@ -9,7 +9,7 @@ import json import math import os import sys -from unittest.mock import Mock, patch +from unittest.mock import patch import pytest @@ -120,10 +120,10 @@ class TestPerplexityCostCalculator: # Expected costs: # Input: 100 tokens * $2e-6 = $0.0002 # Output: 50 tokens * $8e-6 = $0.0004 - # Search: 3 queries * ($0.005 / 1000) = $0.000015 - # Total completion cost: $0.000415 + # Search: 3 queries * $0.005 per request = $0.015 + # Total completion cost: $0.0154 expected_prompt_cost = 100 * 2e-6 - expected_completion_cost = (50 * 8e-6) + (3 / 1000 * 0.005) + expected_completion_cost = (50 * 8e-6) + (3 * 0.005) assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-6) @@ -195,10 +195,10 @@ class TestPerplexityCostCalculator: # Total prompt cost = $0.00026 # Output (text): (50 - 15) tokens * $8e-6 = $0.00028 # Reasoning: 15 tokens * $3e-6 = $0.000045 - # Search: 2 queries * ($0.005 / 1000) = $0.00001 - # Total completion cost = $0.000335 + # Search: 2 queries * $0.005 per request = $0.01 + # Total completion cost = $0.010325 expected_prompt_cost = (100 * 2e-6) + (30 * 2e-6) - expected_completion_cost = ((50 - 15) * 8e-6) + (15 * 3e-6) + (2 / 1000 * 0.005) + expected_completion_cost = ((50 - 15) * 8e-6) + (15 * 3e-6) + (2 * 0.005) assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-6) @@ -311,7 +311,7 @@ class TestPerplexityCostCalculator: # Calculate expected total cost (reasoning is a subset of completion_tokens) expected_prompt_cost = (100 * 2e-6) + (15 * 2e-6) # Input + citation expected_completion_cost = ( - ((50 - 10) * 8e-6) + (10 * 3e-6) + (1 / 1000 * 0.005) + ((50 - 10) * 8e-6) + (10 * 3e-6) + (1 * 0.005) ) # Output (text) + reasoning + search expected_total = expected_prompt_cost + expected_completion_cost @@ -361,7 +361,7 @@ class TestPerplexityCostCalculator: expected_completion_cost = ( ((50 - reasoning_tokens) * 8e-6) + (reasoning_tokens * 3e-6) - + (search_queries / 1000 * 0.005) + + (search_queries * 0.005) ) assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_integration.py b/tests/test_litellm/llms/perplexity/test_perplexity_integration.py index e59fbc9f272..8691e6a1ee5 100644 --- a/tests/test_litellm/llms/perplexity/test_perplexity_integration.py +++ b/tests/test_litellm/llms/perplexity/test_perplexity_integration.py @@ -9,7 +9,6 @@ import json import math import os import sys -from unittest.mock import Mock, patch import pytest @@ -106,8 +105,8 @@ class TestPerplexityIntegration: expected_prompt_cost = (100 * 2e-6) + (citation_tokens * 2e-6) expected_completion_cost = ( - ((50 - 10) * 8e-6) + (10 * 3e-6) + (2 / 1000 * 0.005) - ) + ((50 - 10) * 8e-6) + (10 * 3e-6) + (2 * 0.005) + ) # Output (text) + reasoning + search expected_total = expected_prompt_cost + expected_completion_cost assert math.isclose(total_cost, expected_total, rel_tol=1e-6) @@ -152,8 +151,8 @@ class TestPerplexityIntegration: expected_prompt_cost = (200 * 2e-6) + (40 * 2e-6) expected_completion_cost = ( - ((100 - 25) * 8e-6) + (25 * 3e-6) + (3 / 1000 * 0.005) - ) + ((100 - 25) * 8e-6) + (25 * 3e-6) + (3 * 0.005) + ) # Output (text) + reasoning + search assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6) @@ -262,9 +261,9 @@ class TestPerplexityIntegration: expected_prompt_cost = (50000 * 2e-6) + (5000 * 2e-6) expected_completion_cost = ( - ((25000 - 10000) * 8e-6) + (10000 * 3e-6) + (100 / 1000 * 0.005) - ) - expected_total = expected_prompt_cost + expected_completion_cost + ((25000 - 10000) * 8e-6) + (10000 * 3e-6) + (100 * 0.005) + ) # $0.65 + expected_total = expected_prompt_cost + expected_completion_cost # $0.76 assert math.isclose(total_cost, expected_total, rel_tol=1e-6) assert total_cost > 0.25 @@ -326,7 +325,7 @@ class TestPerplexityIntegration: # Should calculate costs correctly expected_prompt_cost = (100 * 2e-6) + (10 * 2e-6) - expected_completion_cost = (50 * 8e-6) + (1 / 1000 * 0.005) + expected_completion_cost = (50 * 8e-6) + (1 * 0.005) assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-6) assert math.isclose(completion_cost_val, expected_completion_cost, rel_tol=1e-6) diff --git a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py index 5496486765c..9870d30d488 100644 --- a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py +++ b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py @@ -63,23 +63,17 @@ class TestTinyfishSearchConfig: assert headers["X-API-Key"] == "sk-tinyfish-test" assert headers["Accept"] == "application/json" - def test_validate_environment_from_env(self): + def test_validate_environment_from_env(self, monkeypatch): + monkeypatch.setenv("TINYFISH_API_KEY", "sk-from-env") config = TinyfishSearchConfig() - with patch( - "litellm.llms.tinyfish.search.transformation.get_secret_str", - return_value="sk-from-env", - ): - headers = config.validate_environment(headers={}) + headers = config.validate_environment(headers={}) assert headers["X-API-Key"] == "sk-from-env" - def test_validate_environment_missing_key(self): + def test_validate_environment_missing_key(self, monkeypatch): + monkeypatch.delenv("TINYFISH_API_KEY", raising=False) config = TinyfishSearchConfig() - with patch( - "litellm.llms.tinyfish.search.transformation.get_secret_str", - return_value=None, - ): - with pytest.raises(ValueError, match="TINYFISH_API_KEY"): - config.validate_environment(headers={}) + with pytest.raises(ValueError, match="TINYFISH_API_KEY"): + config.validate_environment(headers={}) def test_validate_environment_uses_api_base_kwarg(self): config = TinyfishSearchConfig() diff --git a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index 1ebd704be34..1f171496cce 100644 --- a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -346,3 +346,74 @@ def test_vertex_does_not_warn_when_dropping_non_guardrail_session_update(caplog) "Vertex AI Realtime" in record.message and "session.update" in record.message for record in caplog.records ) + + +async def test_async_realtime_does_not_forward_client_query_params_to_vertex_backend( + monkeypatch, +): + """Regression: forwarding client ?model=/?intent= to the Vertex Live WSS URL causes 1007 errors. + + Exercises ``async_realtime`` end-to-end so that re-adding ``_append_query_params`` + (the reverted bug) would push ``model=``/``intent=`` onto the backend URL and fail here. + """ + import websockets + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + + captured = {} + + def fake_connect(url, *args, **kwargs): + captured["url"] = url + raise RuntimeError("stop before establishing the backend connection") + + monkeypatch.setattr(websockets, "connect", fake_connect) + + await BaseLLMHTTPHandler().async_realtime( + model="gemini-live-2.5-flash-native-audio", + websocket=AsyncMock(), + logging_obj=MagicMock(), + provider_config=cfg, + headers={}, + query_params={ + "model": "gemini-live-2.5-flash-native-audio", + "intent": "chat", + }, + ) + + assert "?" not in captured["url"] + assert "model=" not in captured["url"] + assert "intent=" not in captured["url"] + + +def test_vertex_function_call_output_omits_id(): + """Regression: Vertex Live rejects ``id`` on toolResponse.functionResponses (1007).""" + cfg = VertexAIRealtimeConfig( + access_token="tok", project="my-proj", location="us-central1" + ) + cfg._tool_call_id_to_name["call_abc123"] = "terminate_call" + + messages = cfg.transform_realtime_request( + json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "function_call_output", + "call_id": "call_abc123", + "output": '{"status": "ok"}', + }, + } + ), + "gemini-live-2.5-flash-native-audio", + session_configuration_request="existing", + ) + + assert len(messages) == 1 + payload = json.loads(messages[0]) + function_response = payload["toolResponse"]["functionResponses"][0] + assert "id" not in function_response + assert function_response["name"] == "terminate_call" + assert function_response["response"] == {"status": "ok"} diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py new file mode 100644 index 00000000000..7e028064e4c --- /dev/null +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -0,0 +1,333 @@ +"""Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``).""" + +import importlib +import sys +import types + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse + +# `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` +# function onto `litellm.ocr` and shadows the submodule, so import the modules +# explicitly via importlib rather than attribute traversal. +ocr_main = importlib.import_module("litellm.ocr.main") +rust_bridge = importlib.import_module("litellm.ocr.rust_bridge") + +MODEL = "mistral/mistral-ocr-latest" +DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} + +FAKE_OCR_RESPONSE = { + "pages": [{"index": 0, "markdown": "hello world"}], + "model": "mistral-ocr-2505-completion", + "document_annotation": None, + "usage_info": {"pages_processed": 1}, + "object": "ocr", +} + + +class RecordingBridge: + """A fake ``RustOcr`` callable that records the args it was handed.""" + + def __init__(self): + self.calls = [] + + def __call__( + self, model, document, api_key, api_base, optional_params, timeout_seconds + ): + self.calls.append( + { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "optional_params": optional_params, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_OCR_RESPONSE) + + +class RecordingLogging: + """A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``.""" + + def __init__(self): + self.pre_call_kwargs = None + + def pre_call(self, *, input, api_key, additional_args): + self.pre_call_kwargs = { + "input": input, + "api_key": api_key, + "additional_args": additional_args, + } + + +class FakeOCRConfig: + """A stand-in ``BaseOCRConfig`` that echoes the request it would build.""" + + def validate_environment( + self, *, headers, model, api_key, api_base, litellm_params + ): + return {"authorization": f"Bearer {api_key}"} + + def get_complete_url(self, *, api_base, model, optional_params, litellm_params): + return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" + + +@pytest.fixture(autouse=True) +def _reset_rust_flag(): + """Keep the global toggle isolated between tests.""" + rust_bridge.use_litellm_rust(False, ocr=None) + yield + rust_bridge.use_litellm_rust(False, ocr=None) + + +@pytest.fixture +def fake_bridge(): + """Enable the Rust path with an injected recording bridge (no native wheel).""" + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + return bridge + + +def test_use_litellm_rust_toggles_flag(): + assert rust_bridge.rust_ocr_enabled() is False + litellm.use_litellm_rust() + assert rust_bridge.rust_ocr_enabled() is True + litellm.use_litellm_rust(False) + assert rust_bridge.rust_ocr_enabled() is False + + +def test_load_rust_ocr_returns_injected_impl(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + assert rust_bridge.load_rust_ocr() is bridge + + +def test_toggle_without_ocr_arg_preserves_injected_impl(): + """Regression: routine enable/disable calls must not clobber a prior injection. + + Earlier, ``use_litellm_rust()`` unconditionally assigned the keyword default + of ``None`` to ``_rust_ocr_impl``, silently dropping a custom bridge whenever + a caller toggled the flag without re-passing ``ocr=``. + """ + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + litellm.use_litellm_rust(False) + assert rust_bridge.load_rust_ocr() is bridge + litellm.use_litellm_rust(True) + assert rust_bridge.load_rust_ocr() is bridge + + +def test_explicit_ocr_none_clears_injected_impl(): + bridge = RecordingBridge() + litellm.use_litellm_rust(True, ocr=bridge) + + litellm.use_litellm_rust(True, ocr=None) + assert rust_bridge.load_rust_ocr() is None + + +def test_load_rust_ocr_none_when_extension_absent(): + """With no injected impl and no compiled wheel, the loader returns None so the + caller degrades to the Python path instead of raising ImportError.""" + litellm.use_litellm_rust(True) # no impl injected; extension isn't built in CI + assert rust_bridge.load_rust_ocr() is None + + +def test_load_rust_ocr_uses_compiled_extension(monkeypatch): + """With no injected impl but a compiled ``litellm_python_bridge`` importable, + the loader returns the extension's ``ocr`` callable. The native wheel isn't + built in CI, so stand in a fake module via ``sys.modules``.""" + fake_module = types.ModuleType("litellm_python_bridge") + fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module) + + litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension + assert rust_bridge.load_rust_ocr() is fake_module.ocr + + +def test_timeout_to_seconds_handles_float_timeout_and_none(): + assert ocr_main._timeout_to_seconds(12.5) == 12.5 + assert ocr_main._timeout_to_seconds(None) is None + assert ocr_main._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 + + +def test_run_rust_ocr_forwards_args_and_wraps_response(): + bridge = RecordingBridge() + logging_obj = RecordingLogging() + + response = ocr_main._run_rust_ocr( + rust_ocr=bridge, + logging_obj=logging_obj, + provider_config=FakeOCRConfig(), + resolve_api_key=lambda _name: None, + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-test", + api_base="https://proxy.internal", + optional_params={"include_image_base64": True}, + litellm_params={}, + timeout_seconds=12.5, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + call = bridge.calls[0] + assert call == { + "model": "mistral-ocr-latest", + "document": DOCUMENT, + "api_key": "sk-test", + "api_base": "https://proxy.internal", + "optional_params": {"include_image_base64": True}, + "timeout_seconds": 12.5, + } + + +def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): + """No explicit api_key: the resolver (get_secret_str in production) supplies it, + so secret-manager backends (AWS/Azure/GCP/Vault) work like the Python path.""" + bridge = RecordingBridge() + + ocr_main._run_rust_ocr( + rust_ocr=bridge, + logging_obj=RecordingLogging(), + provider_config=FakeOCRConfig(), + resolve_api_key=lambda name: ( + "sk-from-vault" if name == "MISTRAL_API_KEY" else None + ), + model="mistral-ocr-latest", + document=DOCUMENT, + api_key=None, + api_base=None, + optional_params={}, + litellm_params={}, + timeout_seconds=None, + ) + + assert bridge.calls[0]["api_key"] == "sk-from-vault" + + +def test_run_rust_ocr_prefers_explicit_key_over_resolver(): + bridge = RecordingBridge() + resolver_calls = [] + + def _resolver(name): + resolver_calls.append(name) + return "sk-from-vault" + + ocr_main._run_rust_ocr( + rust_ocr=bridge, + logging_obj=RecordingLogging(), + provider_config=FakeOCRConfig(), + resolve_api_key=_resolver, + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-explicit", + api_base=None, + optional_params={}, + litellm_params={}, + timeout_seconds=None, + ) + + assert bridge.calls[0]["api_key"] == "sk-explicit" + assert resolver_calls == [] # resolver never consulted when a key is supplied + + +def test_run_rust_ocr_runs_pre_call_logging(): + """The Rust shortcut must run pre_call so callbacks and spend tracking fire.""" + logging_obj = RecordingLogging() + + ocr_main._run_rust_ocr( + rust_ocr=RecordingBridge(), + logging_obj=logging_obj, + provider_config=FakeOCRConfig(), + resolve_api_key=lambda _name: None, + model="mistral-ocr-latest", + document=DOCUMENT, + api_key="sk-test", + api_base="https://api.mistral.ai/v1", + optional_params={"include_image_base64": True}, + litellm_params={}, + timeout_seconds=None, + ) + + assert logging_obj.pre_call_kwargs is not None + assert logging_obj.pre_call_kwargs["input"] == "OCR document processing" + additional_args = logging_obj.pre_call_kwargs["additional_args"] + complete_input = additional_args["complete_input_dict"] + assert complete_input["document"] == DOCUMENT + assert complete_input["include_image_base64"] is True + # The logged request mirrors what Rust sends: resolved URL + headers. + assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" + assert additional_args["headers"] == {"authorization": "Bearer sk-test"} + + +def test_ocr_routes_to_rust_when_enabled(fake_bridge): + response = litellm.ocr( + model=MODEL, + document=DOCUMENT, + api_key="sk-test", + include_image_base64=True, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + assert len(fake_bridge.calls) == 1 + call = fake_bridge.calls[0] + # Provider prefix is stripped before reaching the bridge. + assert call["model"] == "mistral-ocr-latest" + assert call["document"] == DOCUMENT + assert call["api_key"] == "sk-test" + # Raw OCR params ride along in optional_params; Rust filters to supported keys. + assert call["optional_params"].get("include_image_base64") is True + + +def test_ocr_forwards_timeout_to_rust(fake_bridge): + """Caller-supplied timeout must flow into the Rust bridge so the fixed 600s + client ceiling doesn't silently override shorter deadlines.""" + litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test", timeout=12.5) + + assert fake_bridge.calls[0]["timeout_seconds"] == 12.5 + + +def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): + """When no explicit timeout is given, the library default (request_timeout) + must still be forwarded so the Rust path matches the Python path's deadline.""" + from litellm.constants import request_timeout + + litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) + + +def test_ocr_does_not_route_to_rust_when_disabled(): + """With the flag off, the bridge must not be consulted even if an impl exists.""" + bridge = RecordingBridge() + litellm.use_litellm_rust(False, ocr=bridge) + + assert rust_bridge.rust_ocr_enabled() is False + # The impl stays available for injection, but the disabled flag gates usage, + # so ocr() never reaches the Rust path (asserted via the enabled-path test). + assert bridge.calls == [] + + +def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): + """Rust enabled but no bridge available (no injected impl, no compiled wheel): + ocr() must degrade to the Python HTTP handler instead of raising.""" + litellm.use_litellm_rust(True) # enabled, but load_rust_ocr() returns None in CI + + captured = {} + + def fake_handler_ocr(**kwargs): + captured["called"] = True + return OCRResponse(pages=[], model="mistral-ocr-latest", object="ocr") + + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fake_handler_ocr) + + response = litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert captured.get("called") is True # Python path was used + assert isinstance(response, OCRResponse) 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 ab42ee1e979..20fd5c1d86a 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 @@ -15,7 +15,11 @@ from starlette.datastructures import Headers from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) -from litellm.proxy._types import SpecialHeaders, UserAPIKeyAuth +from litellm.proxy._types import ( + SpecialHeaders, + SpecialMCPServerNames, + UserAPIKeyAuth, +) @pytest.mark.asyncio @@ -166,6 +170,53 @@ class TestMCPRequestHandler: mock_key_servers.assert_called_once_with(user_api_key_auth) mock_team_servers.assert_called_once_with(user_api_key_auth) + @pytest.mark.parametrize("team_servers", [[], ["team_server1", "team_server2"]]) + async def test_no_mcp_servers_sentinel_returns_empty(self, team_servers): + """A key scoped to the no-mcp-servers sentinel resolves to zero servers, + overriding team inheritance and never leaking the sentinel marker.""" + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="test-team" + ) + key_object_permission = MagicMock() + key_object_permission.mcp_servers = [ + SpecialMCPServerNames.no_mcp_servers.value + ] + + with patch.object( + MCPRequestHandler, + "_get_key_object_permission", + return_value=key_object_permission, + ), patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=team_servers, + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert result == [] + + async def test_get_allowed_mcp_servers_for_key_returns_sentinel_marker(self): + """_get_allowed_mcp_servers_for_key surfaces the sentinel unexpanded so the + caller can short-circuit, ignoring any other entries on the key.""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + key_object_permission = MagicMock() + key_object_permission.mcp_servers = [ + SpecialMCPServerNames.no_mcp_servers.value, + "some-other-server", + ] + + with patch.object( + MCPRequestHandler, + "_get_key_object_permission", + return_value=key_object_permission, + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_key( + user_api_key_auth + ) + + assert result == [SpecialMCPServerNames.no_mcp_servers.value] + async def test_permission_inheritance_edge_cases(self): """Test edge cases in permission inheritance""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_httpx_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_httpx_auth.py new file mode 100644 index 00000000000..9eab089bac6 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_httpx_auth.py @@ -0,0 +1,44 @@ +"""Tests for the concrete httpx.Auth objects the resolver returns. + +NoOpAuth must attach nothing; StaticHeaderAuth must set exactly the configured header. These +pin the header emission the api_key family and passthrough depend on. +""" + +import httpx + +from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + NoOpAuth, + StaticHeaderAuth, +) + + +def _apply(auth: httpx.Auth, request: httpx.Request) -> httpx.Request: + flow = auth.auth_flow(request) + sent = next(flow) + flow.close() + return sent + + +def test_noop_auth_attaches_no_authorization_header(): + request = httpx.Request("GET", "https://upstream.example.com/mcp") + _apply(NoOpAuth(), request) + assert "authorization" not in request.headers + + +def test_static_header_auth_defaults_to_authorization(): + request = httpx.Request("GET", "https://upstream.example.com/mcp") + _apply(StaticHeaderAuth("Bearer abc"), request) + assert request.headers["Authorization"] == "Bearer abc" + + +def test_static_header_auth_honors_custom_header_name(): + request = httpx.Request("GET", "https://upstream.example.com/mcp") + _apply(StaticHeaderAuth("raw-key", header_name="X-API-Key"), request) + assert request.headers["X-API-Key"] == "raw-key" + assert "authorization" not in request.headers + + +def test_static_header_auth_masks_credential_from_introspection(): + auth = StaticHeaderAuth("Bearer super-secret-token") + assert "super-secret-token" not in repr(auth) + assert "super-secret-token" not in str(vars(auth)) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py new file mode 100644 index 00000000000..7885617aa46 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -0,0 +1,57 @@ +"""Tests for the resolver dispatch skeleton. + +Every mode must reach its own arm and, until that arm is built, return a typed +`not_implemented` CredError rather than silently producing no credential. Parametrizing over +one config per mode also guards reachability: if a `case` were dropped, that mode would fall to +the `assert_never` tail and raise here instead of returning the stub. +""" + +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + ApiKeyConfig, + AuthorizationCodeConfig, + AuthSpecKind, + AwsSigV4Config, + ClientCredentialsConfig, + Error, + NoneConfig, + PassthroughConfig, + ServerSpec, + SharedKey, + Subject, + TokenExchangeConfig, + UpstreamCredentialProvider, +) + +_ONE_CONFIG_PER_MODE = [ + (AuthSpecKind.none, NoneConfig()), + (AuthSpecKind.api_key, ApiKeyConfig(key_source=SharedKey(value=SecretStr("k")))), + (AuthSpecKind.passthrough, PassthroughConfig()), + (AuthSpecKind.client_credentials, ClientCredentialsConfig()), + (AuthSpecKind.token_exchange, TokenExchangeConfig()), + (AuthSpecKind.authorization_code, AuthorizationCodeConfig()), + (AuthSpecKind.aws_sigv4, AwsSigV4Config(region="us-east-1")), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind, config", _ONE_CONFIG_PER_MODE) +async def test_every_mode_reaches_its_arm_and_returns_not_implemented(kind, config): + spec = ServerSpec( + server_id="s", resource="https://upstream.example.com", config=config + ) + subject = Subject(tenant_id="", subject_id="") + + result = await UpstreamCredentialProvider().resolve_credentials(subject, spec) + + assert isinstance(result, Error) + assert result.error.tag == "not_implemented" + assert kind.value in result.error.summary + + +def test_all_seven_modes_are_covered(): + # Guards that the parametrization (and therefore the dispatch) spans every AuthSpecKind, so a + # newly added mode without a test row is caught here rather than slipping through. + assert {kind for kind, _ in _ONE_CONFIG_PER_MODE} == set(AuthSpecKind) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_result.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_result.py new file mode 100644 index 00000000000..5d26a530214 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_result.py @@ -0,0 +1,20 @@ +"""Smoke test for the outbound_credentials Result union. + +Result is trivial frozen dataclasses; its load-bearing guarantee (no `.ok` access before +the Error arm is eliminated) is a type-checker property, not a runtime one. This pins only +the runtime contract consumers rely on: each arm carries its payload and discriminates by +type. The union is exercised for real where it is used (see PR2's parse_auth_spec_kind). +""" + +from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + Error, + Ok, + Result, +) + + +def test_ok_and_error_carry_payload_and_discriminate(): + ok: Result[int, str] = Ok(5) + err: Result[int, str] = Error("boom") + assert isinstance(ok, Ok) and ok.ok == 5 + assert isinstance(err, Error) and err.error == "boom" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py new file mode 100644 index 00000000000..43b3612a5f2 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py @@ -0,0 +1,150 @@ +"""Construction-time tests for the outbound_credentials vocabulary. + +The point of the typed seam is that illegal mode/field combinations are unrepresentable: +a config missing a required field, an unknown mode, or a mismatched discriminated-union +source must fail at construction, not at resolve time. These tests pin that, plus the +CredError tag/summary surface and the derived auth_spec_kind. Each assertion fails if the +corresponding guarantee is mutated away. +""" + +import pytest +from pydantic import SecretStr, TypeAdapter, ValidationError + +from litellm.proxy._experimental.mcp_server.outbound_credentials import ( + Ambient, + ApiKeyConfig, + AuthConfig, + AuthSpecKind, + AwsSigV4Config, + Byok, + CredError, + Error, + NoneConfig, + Ok, + ServerSpec, + SharedKey, + StaticKeys, + parse_auth_spec_kind, +) + +_AUTH_CONFIG = TypeAdapter(AuthConfig) + + +def test_parse_auth_spec_kind_accepts_known_mode(): + result = parse_auth_spec_kind("token_exchange") + assert isinstance(result, Ok) + assert result.ok is AuthSpecKind.token_exchange + + +def test_parse_auth_spec_kind_rejects_unknown_mode(): + result = parse_auth_spec_kind("totally_made_up") + assert isinstance(result, Error) + assert result.error.tag == "unsupported_mode" + assert "totally_made_up" in result.error.summary + + +@pytest.mark.parametrize( + "factory, expected_tag", + [ + (CredError.of_unauthorized, "unauthorized"), + (CredError.of_misconfigured, "misconfigured"), + (CredError.of_upstream_unavailable, "upstream_unavailable"), + (CredError.of_unsupported_mode, "unsupported_mode"), + (CredError.of_precondition_required, "precondition_required"), + (CredError.of_not_implemented, "not_implemented"), + ], +) +def test_crederror_factory_sets_the_matching_tag(factory, expected_tag): + err = factory("detail text") + assert err.tag == expected_tag + assert "detail text" in err.summary + + +def test_apikeyconfig_requires_a_key_source(): + with pytest.raises(ValidationError): + ApiKeyConfig() # type: ignore[call-arg] + + +def test_sharedkey_requires_a_value(): + with pytest.raises(ValidationError): + SharedKey() # type: ignore[call-arg] + + +def test_static_keys_require_id_and_secret(): + with pytest.raises(ValidationError): + StaticKeys(access_key_id="AKIA") # type: ignore[call-arg] + + +def test_aws_sigv4_requires_a_region(): + with pytest.raises(ValidationError): + AwsSigV4Config() # type: ignore[call-arg] + + +def test_aws_sigv4_defaults_to_the_ambient_credential_chain(): + cfg = AwsSigV4Config(region="us-east-1") + assert isinstance(cfg.credentials, Ambient) + assert cfg.service == "bedrock-agentcore" + + +def test_authconfig_discriminates_on_kind(): + api_key = _AUTH_CONFIG.validate_python( + {"kind": "api_key", "key_source": {"source": "shared", "value": "k"}} + ) + assert isinstance(api_key, ApiKeyConfig) + assert isinstance(api_key.key_source, SharedKey) + + none = _AUTH_CONFIG.validate_python({"kind": "none"}) + assert isinstance(none, NoneConfig) + + +def test_authconfig_rejects_unknown_kind(): + with pytest.raises(ValidationError): + _AUTH_CONFIG.validate_python({"kind": "not_a_mode"}) + + +def test_apikeysource_discriminates_and_rejects_unknown_source(): + byok = ApiKeyConfig.model_validate({"key_source": {"source": "byok"}}) + assert isinstance(byok.key_source, Byok) + + with pytest.raises(ValidationError): + ApiKeyConfig.model_validate({"key_source": {"source": "mystery"}}) + + +def test_server_spec_derives_auth_spec_kind_from_config(): + spec = ServerSpec( + server_id="s1", + resource="https://api.example.com", + config=NoneConfig(), + ) + assert spec.auth_spec_kind is AuthSpecKind.none + + api_spec = ServerSpec( + server_id="s2", + resource="https://api.example.com", + config=ApiKeyConfig(key_source=SharedKey(value=SecretStr("k"))), + ) + assert api_spec.auth_spec_kind is AuthSpecKind.api_key + + +def test_api_key_header_placement(): + default = ApiKeyConfig(key_source=SharedKey(value=SecretStr("tok"))) + assert default.header("tok") == ("Authorization", "Bearer tok") + + raw = ApiKeyConfig( + header_name="X-API-Key", + value_prefix="", + key_source=SharedKey(value=SecretStr("tok")), + ) + assert raw.header("tok") == ("X-API-Key", "tok") + + +def test_configs_are_frozen(): + cfg = NoneConfig() + with pytest.raises(ValidationError): + cfg.kind = AuthSpecKind.api_key # type: ignore[misc] + + +def test_secrets_do_not_leak_in_repr(): + key = SharedKey(value=SecretStr("super-secret")) + assert "super-secret" not in repr(key) + assert key.value.get_secret_value() == "super-secret" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py index 468bd946ae9..d299239f68e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_debug.py @@ -41,9 +41,13 @@ class TestMask: def test_empty_returns_none_label(self): assert MCPDebug._mask("") == "(none)" - def test_short_value_unchanged(self): - # visible_prefix=6 + visible_suffix=4 = 10, so <= 10 chars unchanged - assert MCPDebug._mask("sk-1234") == "sk-1234" + def test_short_value_masked(self): + # Short auth values must not be echoed verbatim in debug headers, even though + # visible_prefix + visible_suffix would otherwise reveal the whole value. + masked = MCPDebug._mask("sk-1234") + assert "sk-1234" not in masked + assert set(masked) == {"*"} + assert len(masked) == len("sk-1234") def test_long_value_masked(self): result = MCPDebug._mask("Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index c86ae966f21..f44552c2943 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -6293,3 +6293,153 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_empty_list_fails_cl ) assert result == [] + + +class TestProxyExceptionToHttpException: + """Auth failures reach the MCP ASGI handlers as ProxyException, not + HTTPException. The handlers must map them back to their real status and + headers; otherwise they fall through to the generic 500 handler, dropping + the 401 + WWW-Authenticate challenge an OAuth client needs to re-authenticate + and surfacing the tool call as a cancelled/terminated session. + """ + + def test_preserves_401_status_and_www_authenticate_header(self): + from litellm.proxy._experimental.mcp_server.server import ( + _proxy_exception_to_http_exception, + ) + from litellm.proxy._types import ProxyException + + exc = ProxyException( + message="Authentication Error, invalid token", + type="auth_error", + param="key", + code=401, + headers={"WWW-Authenticate": 'Bearer resource_metadata="/x"'}, + ) + + http_exc = _proxy_exception_to_http_exception(exc) + + assert http_exc.status_code == 401 + assert http_exc.detail == "Authentication Error, invalid token" + assert http_exc.headers["WWW-Authenticate"] == 'Bearer resource_metadata="/x"' + + def test_preserves_403_status(self): + from litellm.proxy._experimental.mcp_server.server import ( + _proxy_exception_to_http_exception, + ) + from litellm.proxy._types import ProxyException + + http_exc = _proxy_exception_to_http_exception( + ProxyException( + message="Forbidden", type="auth_error", param="key", code=403 + ) + ) + + assert http_exc.status_code == 403 + + def test_non_numeric_code_falls_back_to_500(self): + from litellm.proxy._experimental.mcp_server.server import ( + _proxy_exception_to_http_exception, + ) + from litellm.proxy._types import ProxyException + + # ProxyException normalises code to the string "None" when unset. + http_exc = _proxy_exception_to_http_exception( + ProxyException(message="boom", type="server_error", param=None, code=None) + ) + + assert http_exc.status_code == 500 + + +class TestStreamableHttpAuthErrorMapping: + """End-to-end guard for the handler wiring: a ProxyException from auth must + propagate as the real HTTPException (401 + WWW-Authenticate), not be + flattened to a generic 500 by the catch-all handler. + """ + + @pytest.mark.asyncio + async def test_streamable_http_propagates_proxy_exception_as_401(self): + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._types import ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/some_server", + "headers": [(b"x-litellm-api-key", b"sk-bad")], + } + + async def receive(): + return {"type": "http.request", "body": b"{}", "more_body": False} + + sent = [] + + async def send(message): + sent.append(message) + + auth_failure = ProxyException( + message="Authentication Error, invalid token", + type="auth_error", + param="key", + code=401, + headers={"WWW-Authenticate": "Bearer"}, + ) + + with patch.object( + mcp_module, + "extract_mcp_auth_context", + new=AsyncMock(side_effect=auth_failure), + ): + with pytest.raises(HTTPException) as exc_info: + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert exc_info.value.status_code == 401 + assert exc_info.value.headers["WWW-Authenticate"] == "Bearer" + # Must not have emitted a 500 body via the generic catch-all. + assert not any( + m.get("type") == "http.response.start" and m.get("status") == 500 + for m in sent + ) + + @pytest.mark.asyncio + async def test_sse_propagates_proxy_exception_as_401(self): + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._types import ProxyException + + scope = { + "type": "http", + "method": "GET", + "path": "/mcp/some_server", + "headers": [(b"x-litellm-api-key", b"sk-bad")], + } + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + sent = [] + + async def send(message): + sent.append(message) + + auth_failure = ProxyException( + message="Authentication Error, invalid token", + type="auth_error", + param="key", + code=401, + headers={"WWW-Authenticate": "Bearer"}, + ) + + with patch.object( + mcp_module, + "extract_mcp_auth_context", + new=AsyncMock(side_effect=auth_failure), + ): + with pytest.raises(HTTPException) as exc_info: + await mcp_module.handle_sse_mcp(scope, receive, send) + + assert exc_info.value.status_code == 401 + assert exc_info.value.headers["WWW-Authenticate"] == "Bearer" + assert not any( + m.get("type") == "http.response.start" and m.get("status") == 500 + for m in sent + ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 1b815b7a1c9..8dbee1daa36 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2948,6 +2948,41 @@ class TestMCPServerManager: assert "test_server_1" in result assert "test_server_2" in result + @pytest.mark.asyncio + async def test_no_mcp_servers_sentinel_blocks_allow_all_keys(self): + """A key scoped to no-mcp-servers gets zero servers even when allow_all_keys + servers exist, and the inner resolver is never consulted.""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth + + manager = MCPServerManager() + object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="perm_no_mcp", + mcp_servers=["no-mcp-servers"], + mcp_access_groups=[], + ) + user_api_key_auth = UserAPIKeyAuth( + api_key="sk-test", + user_id="user-123", + object_permission=object_permission, + object_permission_id="perm_no_mcp", + ) + + with patch.object( + manager, "get_allow_all_keys_server_ids", return_value=["global-server"] + ), patch.object( + MCPRequestHandler, + "get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=["leaked-server"], + ) as mock_inner: + result = await manager.get_allowed_mcp_servers(user_api_key_auth) + + assert result == [] + mock_inner.assert_not_called() + @pytest.mark.asyncio async def test_get_allowed_mcp_servers_anonymous_delegate_requires_oauth2(self): """Anonymous delegated auth listing should only include oauth2 servers.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py index 2ffa997bdde..b41cdfb1576 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py @@ -93,6 +93,37 @@ class TestApplyToolsetScope: await _apply_toolset_scope(auth, "toolset-123") assert exc_info.value.status_code == 403 + @pytest.mark.asyncio + @pytest.mark.parametrize("user_role", [None, LitellmUserRoles.PROXY_ADMIN.value]) + async def test_no_mcp_servers_sentinel_denies_toolset_access(self, user_role): + """A key scoped to the no-mcp-servers sentinel cannot reach a toolset it + would otherwise be granted (even as admin); the opt-out covers the + toolset path, which replaces mcp_servers and would drop the sentinel.""" + from starlette.exceptions import HTTPException + + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + op = LiteLLM_ObjectPermissionTable( + object_permission_id="test", + mcp_servers=["no-mcp-servers"], + mcp_toolsets=["toolset-123"], + ) + auth = UserAPIKeyAuth( + api_key="sk-test", object_permission=op, user_role=user_role + ) + + resolve = AsyncMock(return_value={"server-a": ["tool1"]}) + with patch( + "litellm.proxy._experimental.mcp_server.server." + "global_mcp_server_manager.resolve_toolset_tool_permissions", + new=resolve, + ): + with pytest.raises(HTTPException) as exc_info: + await _apply_toolset_scope(auth, "toolset-123") + + assert exc_info.value.status_code == 403 + resolve.assert_not_awaited() + class TestFetchMCPToolsetsAccess: """Tests for GET /v1/mcp/toolset access control.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 2558df8533b..cebc265a148 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -452,6 +452,430 @@ async def test_semantic_filter_hook_skips_no_tools(): print("✅ Hook correctly skips requests without tools") +@pytest.mark.asyncio +async def test_semantic_filter_hook_preserves_native_tools(): + """ + Regression test: mixed MCP + native tools. + + Given: 5 MCP tools (registered in _tool_map) + 2 native OpenAI-format + function tools (not in _tool_map) + When: The hook filters tools + Then: The native tools must survive unconditionally, and only MCP + tools go through the semantic filter. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=2, + similarity_threshold=0.3, + enabled=True, + ) + + # --- MCP tools (registered in the semantic router) --- + mcp_tools = [ + MCPTool( + name=f"mcp_tool_{i}", + description=f"MCP tool {i}", + inputSchema={"type": "object"}, + ) + for i in range(5) + ] + filter_instance._build_router(mcp_tools) + + # --- Native OpenAI-format function tools (NOT in _tool_map) --- + native_tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather", + "parameters": {"type": "object", "properties": {}}, + }, + }, + { + "type": "function", + "function": { + "name": "search_web", + "description": "Search the web", + "parameters": {"type": "object", "properties": {}}, + }, + }, + ] + + # Combine: MCP tools + native tools + all_tools = list(mcp_tools) + native_tools + + hook = SemanticToolFilterHook(filter_instance) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "What is the weather?"}], + "tools": all_tools, + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is not None, "Hook should return modified data" + filtered = result["tools"] + + # Native tools must survive + native_in_result = [ + t for t in filtered if isinstance(t, dict) and t.get("type") == "function" + ] + assert ( + len(native_in_result) == 2 + ), f"Both native tools must survive, got {len(native_in_result)}" + + # MCP tools should be filtered (top_k=2) + mcp_in_result = [t for t in filtered if not isinstance(t, dict)] + assert ( + len(mcp_in_result) <= 2 + ), f"MCP tools should be filtered to top_k=2, got {len(mcp_in_result)}" + + # Total should be native + filtered MCP + assert len(filtered) <= 4, f"Expected at most 4 tools, got {len(filtered)}" + + # Filter stats should be emitted (MCP tools were present) + assert "litellm_semantic_filter_stats" in result["metadata"] + + # Stats should report MCP-only counts, not inflated with native tools + stats = result["metadata"]["litellm_semantic_filter_stats"] + mcp_before, mcp_after = stats.split("->") + assert ( + int(mcp_before) == 5 + ), f"Stats 'from' should be MCP count (5), got {mcp_before}" + assert int(mcp_after) == len( + mcp_in_result + ), f"Stats 'to' should match filtered MCP count, got {mcp_after}" + + print( + f"✅ Hook preserves native tools: {len(all_tools)} -> {len(filtered)} " + f"({len(native_in_result)} native + {len(mcp_in_result)} MCP), " + f"stats={stats}" + ) + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_all_native_tools(): + """ + Regression test: all-native request. + + Given: Only native OpenAI-format function tools (none registered in + the MCP semantic router) + When: The hook processes the request + Then: All tools pass through, and NO spurious semantic filter response + headers are emitted (no litellm_semantic_filter_stats in metadata). + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + mock_router = Mock() + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=3, + similarity_threshold=0.3, + enabled=True, + ) + + # Build router with some MCP tools (so tool_router is not None) + mcp_tools = [ + MCPTool( + name="some_mcp_tool", + description="An MCP tool", + inputSchema={"type": "object"}, + ) + ] + + from litellm.types.utils import Embedding, EmbeddingResponse + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance._build_router(mcp_tools) + + # --- Only native tools in the request --- + native_tools = [ + { + "type": "function", + "function": { + "name": f"native_func_{i}", + "description": f"Native function {i}", + "parameters": {"type": "object", "properties": {}}, + }, + } + for i in range(3) + ] + + hook = SemanticToolFilterHook(filter_instance) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "tools": native_tools, + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is not None, "Hook should return modified data" + filtered = result["tools"] + + # All native tools must pass through + assert ( + len(filtered) == 3 + ), f"All 3 native tools must pass through, got {len(filtered)}" + + # No spurious semantic filter stats (P2 fix) + assert ( + "litellm_semantic_filter_stats" not in result["metadata"] + ), "Should NOT emit semantic filter stats for all-native-tool requests" + + print( + f"✅ Hook passes through all {len(filtered)} native tools, " + f"no spurious filter headers emitted" + ) + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_responses_api_name_collision(): + """ + Regression test: Responses API native tool with MCP-matching name. + + Given: A Responses-API native tool whose top-level ``name`` collides + with an MCP canonical name in ``_tool_map`` + When: The hook classifies tools + Then: The native tool must NOT be sent to the semantic filter, even + though its name matches an MCP canonical. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=2, + similarity_threshold=0.3, + enabled=True, + ) + + # Register an MCP tool with name "github-search" + mcp_tools = [ + MCPTool( + name="github-search", + description="Search GitHub repos", + inputSchema={"type": "object"}, + ) + ] + filter_instance._build_router(mcp_tools) + + # Responses API native tool with SAME name as MCP canonical + responses_api_tool = { + "type": "function", + "name": "github-search", + "description": "Caller-owned search tool", + "parameters": {"type": "object"}, + } + + hook = SemanticToolFilterHook(filter_instance) + + # Verify classification: should be native, not MCP + assert not hook._is_mcp_tool(responses_api_tool), ( + "Responses API tool with type=function + top-level name " + "should be classified as native, not MCP" + ) + + # Full hook test: all-native request should preserve tools + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Search GitHub"}], + "tools": [responses_api_tool], + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + # All tools are native → hook returns data with all tools preserved + filtered = (result or data)["tools"] + assert len(filtered) == 1, f"Native tool must survive, got {len(filtered)}" + assert filtered[0]["name"] == "github-search" + + print("✅ Responses API tool with MCP-matching name correctly classified as native") + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_preserves_tool_order(): + """ + Regression test: tool ordering preservation. + + Given: An interleaved request [mcp_tool_A, native_tool, mcp_tool_B] + When: The hook filters tools (all MCP tools survive) + Then: The output order must match the original request order, + NOT native-first. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=5, + similarity_threshold=0.3, + enabled=True, + ) + + # Register MCP tools + mcp_tool_a = MCPTool( + name="github-search", + description="Search GitHub", + inputSchema={"type": "object"}, + ) + mcp_tool_b = MCPTool( + name="github-issue", + description="Create GitHub issue", + inputSchema={"type": "object"}, + ) + filter_instance._build_router([mcp_tool_a, mcp_tool_b]) + + # Mock filter_tools to return both MCP tools (deterministic) + filter_instance.filter_tools = AsyncMock( # type: ignore[method-assign] + return_value=[mcp_tool_a, mcp_tool_b] + ) + + # Native tool (interleaved between MCP tools) + native_tool = { + "type": "function", + "function": { + "name": "weather_lookup", + "description": "Look up weather", + "parameters": {"type": "object", "properties": {}}, + }, + } + + # Original order: [mcp_A, native, mcp_B] + original_tools = [mcp_tool_a, native_tool, mcp_tool_b] + + hook = SemanticToolFilterHook(filter_instance) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Search GitHub and check weather"}], + "tools": original_tools, + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is not None, "Hook should return modified data" + filtered = result["tools"] + + # All tools should survive + assert len(filtered) == 3, f"Expected 3 tools, got {len(filtered)}" + + # Order must be preserved: [mcp_A, native, mcp_B] + assert filtered[0] is mcp_tool_a, "First tool should be mcp_tool_a" + assert filtered[1] is native_tool, "Second tool should be native_tool" + assert filtered[2] is mcp_tool_b, "Third tool should be mcp_tool_b" + + print( + "✅ Tool ordering preserved: [mcp_A, native, mcp_B] maintained after filtering" + ) + + class TestGetToolsByNames: """ Regression coverage for SemanticMCPToolFilter._get_tools_by_names @@ -489,9 +913,7 @@ class TestGetToolsByNames: {"name": "send_email", "description": "send mail"}, ] - matched = filter_instance._get_tools_by_names( - ["send_email"], available_tools - ) + matched = filter_instance._get_tools_by_names(["send_email"], available_tools) assert len(matched) == 1 assert matched[0]["name"] == "send_email" @@ -503,9 +925,7 @@ class TestGetToolsByNames: client_name = "litellm_" + canonical available_tools = [{"name": client_name, "description": "scrape"}] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 # Must return the incoming tool unchanged so the client-facing @@ -516,13 +936,9 @@ class TestGetToolsByNames: """Some clients use dash as alias separator; accept that too.""" filter_instance = self._make_filter() canonical = "weather_svc-get_weather" - available_tools = [ - {"name": "mcp-" + canonical, "description": "weather"} - ] + available_tools = [{"name": "mcp-" + canonical, "description": "weather"}] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 assert matched[0]["name"] == "mcp-" + canonical @@ -552,9 +968,7 @@ class TestGetToolsByNames: {"name": "litellm_" + canonical, "description": "wrapped"}, ] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 assert matched[0]["name"] == canonical @@ -567,9 +981,7 @@ class TestGetToolsByNames: separator-anchored suffixes of ``litellm_api-fs-read_file``. """ filter_instance = self._make_filter() - available_tools = [ - {"name": "litellm_api-fs-read_file", "description": "read"} - ] + available_tools = [{"name": "litellm_api-fs-read_file", "description": "read"}] matched = filter_instance._get_tools_by_names( ["fs-read_file", "api-fs-read_file"], available_tools @@ -590,9 +1002,7 @@ class TestGetToolsByNames: {"name": "my_" + canonical, "description": "plain search"}, ] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 assert matched[0]["name"] == "my_" + canonical diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5ec5d12784f..c8dc0ea5ed6 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -351,6 +351,43 @@ async def test_can_key_call_model_all_team_models_no_team_id_is_denied(): assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied +@pytest.mark.asyncio +async def test_can_team_access_model_all_team_models_expands_router_models(): + from litellm import Router + from litellm.proxy._types import SpecialModelNames + from litellm.proxy.auth.auth_checks import can_team_access_model + + team_object = LiteLLM_TeamTable( + team_id="team-123", + models=[SpecialModelNames.all_team_models.value], + ) + router = Router( + model_list=[ + { + "model_name": "allowed-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"}, + } + ] + ) + + assert ( + await can_team_access_model( + model="allowed-model", + team_object=team_object, + llm_router=router, + ) + is True + ) + with pytest.raises(ProxyException) as exc_info: + await can_team_access_model( + model="blocked-model", + team_object=team_object, + llm_router=router, + ) + + assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied + + @pytest.mark.asyncio async def test_get_key_object_should_reconnect_once_on_db_connection_error(): mock_prisma_client = MagicMock() @@ -1750,6 +1787,60 @@ async def test_reject_clientside_metadata_tags_allows_key_tags_without_client_ta assert request_body["metadata"]["tags"] == ["engineering"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + [ + "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + "/v1/messages", + ], +) +async def test_common_checks_metadata_route_keeps_key_tags_out_of_provider_metadata( + route, +): + """GH#30629: on routes that track tags in litellm_metadata (bedrock, /v1/messages, + responses, ...) key-level tags must land in litellm_metadata, never in the + provider-facing metadata field (Bedrock rejects non-user_id metadata with HTTP 400). + The auth-time pre-seed keys off LITELLM_METADATA_ROUTES, so hardcoding a single route + or dropping the pre-seed makes apply_key_tags_pre_auth fall back to metadata; this + guards that regression. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + request_body = {"messages": [{"role": "user", "content": "test"}]} + + mock_request = MagicMock(spec=Request) + valid_token = UserAPIKeyAuth( + token="test-token", + metadata={"tags": ["engineering"]}, + ) + + with patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={}, + ): + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=mock_request, + ) + + assert result is True + assert request_body["litellm_metadata"]["tags"] == ["engineering"] + assert "metadata" not in request_body + + @pytest.mark.asyncio async def test_virtual_key_soft_budget_check_with_user_obj(): """Test _virtual_key_soft_budget_check includes user_email when user_obj is provided""" @@ -2365,7 +2456,9 @@ async def test_virtual_key_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:key:test-hashed-token": return 1.5 return fallback_spend @@ -2397,7 +2490,9 @@ async def test_virtual_key_budget_check_fallback_no_counter(): proxy_logging_obj.budget_alerts = AsyncMock() # get_current_spend returns fallback_spend when no counter exists - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): return fallback_spend with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): @@ -2424,7 +2519,9 @@ async def test_team_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team:test-team": return 1.5 return fallback_spend @@ -2449,7 +2546,9 @@ async def test_end_user_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:end_user:customer-1": return 1.5 return fallback_spend @@ -2475,7 +2574,9 @@ async def test_tag_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:paid-tag": return 1.5 return fallback_spend @@ -2523,7 +2624,9 @@ async def test_team_member_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 1.5 return fallback_spend @@ -2756,7 +2859,9 @@ async def test_team_member_budget_check_falls_back_to_team_default_budget_id(): return_value=fake_budget_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 70.0 return fallback_spend @@ -2853,7 +2958,9 @@ async def test_team_member_budget_check_per_member_override_wins_over_team_defau mocked_spend = 70.0 - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return mocked_spend return fallback_spend @@ -2943,7 +3050,9 @@ async def test_team_member_budget_check_null_clone_falls_back_to_team_default(): return_value=fake_default_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 500.0 return fallback_spend @@ -3010,7 +3119,9 @@ async def test_team_member_budget_check_null_clone_with_null_default_skips_enfor return_value=fake_default_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 1000.0 return fallback_spend @@ -3077,7 +3188,9 @@ async def test_team_member_budget_check_zero_team_default_treated_as_no_cap(): return_value=fake_default_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend @@ -3135,7 +3248,9 @@ async def test_team_member_budget_check_zero_per_member_row_still_blocks(): prisma_client = MagicMock() prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend @@ -3714,3 +3829,54 @@ async def test_inference_route_still_enforces_team_budget(): valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"), request=MagicMock(), ) + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_error_names_the_key(): + """BudgetExceededError for a virtual key must name the key (alias + masked key) + so operators don't have to reverse-map a spend figure back to a key.""" + valid_token = UserAPIKeyAuth( + token="hashed-token", + key_alias="payments-prod", + key_name="sk-...um_g", + max_budget=10.0, + spend=0.0, + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=25.0), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + + message = str(exc_info.value) + assert "payments-prod" in message + assert "sk-...um_g" in message + + +@pytest.mark.asyncio +async def test_virtual_key_max_budget_not_exceeded_does_not_raise(): + """Spend below the configured budget must not raise.""" + valid_token = UserAPIKeyAuth( + token="hashed-token", + key_alias="payments-prod", + max_budget=10.0, + spend=0.0, + ) + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=1.0), + ): + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index e652c109987..cd8cf10d037 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -2160,3 +2160,48 @@ class TestGetRequestRouteTemplate: lambda self: (_ for _ in ()).throw(RuntimeError("boom")) ) assert get_request_route_template(req) is None + + +class TestIsRequestBodySafeBlocksModelList: + """model_list is an SDK-only field with no proxy API meaning; it must + be rejected from the request body regardless of any opt-in.""" + + def test_model_list_rejected_with_no_opt_in(self): + with pytest.raises(ValueError, match="model_list is not allowed"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "model_list": [{"model_name": "x", "litellm_params": {}}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_model_list_rejected_even_with_proxy_wide_opt_in(self): + with pytest.raises(ValueError, match="model_list is not allowed"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "model_list": [], + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + + def test_normal_body_still_passes(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + is True + ) diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index e49f025df2e..3203878a1e0 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -106,7 +106,9 @@ async def test_custom_auth_enforces_end_user_budget_when_common_checks_skipped() litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:end_user:customer-1": return 5.0 return fallback_spend diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 63510086f95..b547ec877e2 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1206,7 +1206,13 @@ async def test_auth_builder_returns_team_membership_object(): JWTAuthManager, "get_objects", new_callable=AsyncMock, - return_value=(user_object, None, None, mock_team_membership, user_object.user_id), + return_value=( + user_object, + None, + None, + mock_team_membership, + user_object.user_id, + ), ) as mock_get_objects, patch.object( JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock @@ -3509,9 +3515,7 @@ def test_canonical_user_id_no_change_when_ids_match(): user_object = LiteLLM_UserTable(user_id=same, user_email=same) assert ( - JWTAuthManager._canonical_user_id_from_db( - user_id=same, user_object=user_object - ) + JWTAuthManager._canonical_user_id_from_db(user_id=same, user_object=user_object) == same ) @@ -3802,12 +3806,15 @@ async def test_get_objects_team_membership_uses_rebound_user_id(): user_id_jwt_field="email", user_id_upsert=True ) - with patch( - "litellm.proxy.auth.handle_jwt.get_user_object", - side_effect=fake_get_user_object, - ), patch( - "litellm.proxy.auth.handle_jwt.get_team_membership", - side_effect=fake_get_team_membership, + with ( + patch( + "litellm.proxy.auth.handle_jwt.get_user_object", + side_effect=fake_get_user_object, + ), + patch( + "litellm.proxy.auth.handle_jwt.get_team_membership", + side_effect=fake_get_team_membership, + ), ): ( user_object, diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 8d686900ea6..261485e8965 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -436,7 +436,9 @@ def test_wildcard_custom_prefix_keeps_org_segment_for_non_provider_first_segment result = get_known_models_from_wildcard( wildcard_model="my_hf/*", - litellm_params=LiteLLM_Params(model="huggingface/*", custom_llm_provider="huggingface"), + litellm_params=LiteLLM_Params( + model="huggingface/*", custom_llm_provider="huggingface" + ), ) assert result == ["my_hf/meta-llama/Llama-3-8B"] @@ -541,3 +543,77 @@ async def test_get_available_models_for_user_expands_query_team_wildcard( ) assert "openai/gpt-4o-mini" in result + + +def test_get_key_models_all_team_models_recursive_team(): + """GH#30619: when key and team both have all-team-models, + the sentinel should expand to proxy_model_list.""" + from litellm.proxy.auth.model_checks import get_key_models + from litellm.proxy._types import SpecialModelNames + + user_api_key_dict = type( + "obj", (object,), + { + "models": [SpecialModelNames.all_team_models.value], + "team_id": "team-1", + "team_models": [SpecialModelNames.all_team_models.value], + }, + )() + proxy_model_list = ["model-a", "model-b"] + result = get_key_models(user_api_key_dict, proxy_model_list, {}) + assert SpecialModelNames.all_team_models.value not in result + assert set(result) == {"model-a", "model-b"} + + +def test_get_key_models_all_team_models_keeps_mixed_team_entries(): + from litellm.proxy.auth.model_checks import get_key_models + from litellm.proxy._types import SpecialModelNames + + user_api_key_dict = type( + "obj", + (object,), + { + "models": [SpecialModelNames.all_team_models.value], + "team_id": "team-1", + "team_models": [ + SpecialModelNames.all_team_models.value, + "restricted-model", + ], + }, + )() + result = get_key_models(user_api_key_dict, ["model-a", "model-b"], {}) + assert SpecialModelNames.all_team_models.value not in result + assert set(result) == {"model-a", "model-b", "restricted-model"} + + +def test_get_team_models_all_team_models_expands(): + """GH#30619: all-team-models in team_models should expand.""" + from litellm.proxy.auth.model_checks import get_team_models + from litellm.proxy._types import SpecialModelNames + + result = get_team_models( + [SpecialModelNames.all_team_models.value], + ["model-a", "model-b"], + {}, + ) + assert SpecialModelNames.all_team_models.value not in result + assert set(result) == {"model-a", "model-b"} + + +def test_get_team_models_all_team_models_expands_with_access_groups(): + """GH#30619: all-team-models with include_model_access_groups + should include access group keys.""" + from litellm.proxy.auth.model_checks import get_team_models + from litellm.proxy._types import SpecialModelNames + + result = get_team_models( + [SpecialModelNames.all_team_models.value], + ["model-a", "model-b"], + {"group-1": ["g1-model"], "group-2": ["g2-model"]}, + include_model_access_groups=True, + ) + assert SpecialModelNames.all_team_models.value not in result + assert "model-a" in result + assert "model-b" in result + assert "group-1" in result + assert "group-2" in result diff --git a/tests/test_litellm/proxy/auth/test_network.py b/tests/test_litellm/proxy/auth/test_network.py new file mode 100644 index 00000000000..b67723305e4 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_network.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from fastapi import Request + +from litellm.proxy.auth.network import ( + TrustedProxyConfig, + resolve_client_ip, + resolve_network_context, +) + +TRUSTED = TrustedProxyConfig(use_forwarded_for=True, trusted_proxy_cidrs=["10.0.0.0/8"]) + + +def make_request( + *, + headers: Optional[Dict[str, str]] = None, + client: Optional[Tuple[str, int]] = ("203.0.113.7", 5555), +) -> Request: + raw_headers: List[Tuple[bytes, bytes]] = [ + (key.lower().encode(), value.encode()) for key, value in (headers or {}).items() + ] + scope: Dict[str, Any] = { + "type": "http", + "http_version": "1.1", + "method": "GET", + "path": "/", + "raw_path": b"/", + "query_string": b"", + "headers": raw_headers, + "client": client, + "server": ("testserver", 80), + "scheme": "http", + } + return Request(scope) + + +def test_xff_ignored_when_forwarding_disabled(): + config = TrustedProxyConfig( + use_forwarded_for=False, trusted_proxy_cidrs=["10.0.0.0/8"] + ) + request = make_request( + headers={"x-forwarded-for": "203.0.113.9"}, client=("10.0.0.1", 1) + ) + ip, via_proxy = resolve_client_ip(request, config) + assert ip == "10.0.0.1" + assert via_proxy is False + + +def test_xff_honored_from_trusted_peer(): + request = make_request( + headers={"x-forwarded-for": "203.0.113.9, 10.0.0.5"}, client=("10.0.0.1", 1) + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "203.0.113.9" + assert via_proxy is True + + +def test_spoofed_xff_from_untrusted_peer_is_ignored(): + request = make_request( + headers={"x-forwarded-for": "203.0.113.9"}, client=("8.8.8.8", 1) + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "8.8.8.8" + assert via_proxy is False + + +def test_right_to_left_parse_skips_chained_trusted_proxies(): + request = make_request( + headers={"x-forwarded-for": "198.51.100.4, 10.1.1.1, 10.0.0.9"}, + client=("10.0.0.1", 1), + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "198.51.100.4" + assert via_proxy is True + + +def test_all_trusted_hops_fall_back_to_peer(): + request = make_request( + headers={"x-forwarded-for": "10.1.1.1, 10.0.0.9"}, client=("10.0.0.1", 1) + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "10.0.0.1" + assert via_proxy is True + + +def test_invalid_xff_token_is_skipped(): + request = make_request( + headers={"x-forwarded-for": "not-an-ip, 203.0.113.50"}, client=("10.0.0.1", 1) + ) + ip, _ = resolve_client_ip(request, TRUSTED) + assert ip == "203.0.113.50" + + +def test_network_context_captures_host_and_proxy_flag(): + request = make_request( + headers={"x-forwarded-for": "203.0.113.9", "host": "proxy.litellm.ai"}, + client=("10.0.0.1", 1), + ) + ctx = resolve_network_context(request, TRUSTED) + assert ctx.client_ip == "203.0.113.9" + assert ctx.host == "proxy.litellm.ai" + assert ctx.via_trusted_proxy is True diff --git a/tests/test_litellm/proxy/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py index c81f4cb7d66..d55a5472af1 100644 --- a/tests/test_litellm/proxy/auth/test_onboarding.py +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -18,7 +18,6 @@ from fastapi import HTTPException import litellm from litellm.proxy._types import InvitationClaim - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/auth/test_resolvers_exceptions.py b/tests/test_litellm/proxy/auth/test_resolvers_exceptions.py new file mode 100644 index 00000000000..f88b7e2be14 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_exceptions.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.auth.resolvers.exceptions import ( + IdentityResolutionError, + KeyNotFoundError, + KeyNotInCacheError, + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) + + +def test_all_resolution_errors_share_one_base(): + errors = [ + NoDatabaseConnectionError(), + KeyNotInCacheError("hashed"), + KeyNotFoundError("hashed"), + PrincipalMissingSourceKeyError(), + ] + assert all(isinstance(e, IdentityResolutionError) for e in errors) + + +def test_key_not_found_preserves_the_public_401_contract(): + # The auth seam catches ProxyException and rewrites the 401 message, so a + # missing key must keep mapping to that exact contract. + error = KeyNotFoundError("hashed-token") + + assert isinstance(error, ProxyException) + assert error.code == "401" + assert error.type == ProxyErrorTypes.token_not_found_in_db.value + assert error.param == "key" + + +def test_key_not_in_cache_names_the_token(): + assert "hashed-token" in str(KeyNotInCacheError("hashed-token")) diff --git a/tests/test_litellm/proxy/auth/test_resolvers_models.py b/tests/test_litellm/proxy/auth/test_resolvers_models.py new file mode 100644 index 00000000000..0fbe10b6f28 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_models.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.resolvers.models import ( + EndUserIdentity, + Principal, + PrincipalType, + ProjectIdentity, + TeamIdentity, + UserIdentity, +) +from litellm.proxy.auth.roles import Role, TeamRole + + +def _principal() -> Principal: + return Principal( + principal_type=PrincipalType.HUMAN, + subject="u1", + auth_method=AuthMethod.OIDC, + ) + + +def test_principal_is_frozen(): + principal = _principal() + with pytest.raises(ValidationError): + principal.subject = "mutated" + + +def test_principal_defaults_are_independent_instances(): + a = _principal() + b = _principal() + assert a.teams == [] and a.scopes == [] and a.audience == [] + assert a.teams is not b.teams + + +def test_principal_requires_identity_core_fields(): + with pytest.raises(ValidationError): + Principal(subject="u1") # missing principal_type + auth_method + + +def test_principal_roles_are_validated_against_role_enum(): + principal = Principal( + principal_type=PrincipalType.HUMAN, + subject="u1", + auth_method=AuthMethod.OIDC, + roles=["org_admin"], + ) + assert principal.roles == [Role.ORG_ADMIN] + assert isinstance(principal.roles[0], Role) + + with pytest.raises(ValidationError): + Principal( + principal_type=PrincipalType.HUMAN, + subject="u1", + auth_method=AuthMethod.OIDC, + roles=["not_a_real_role"], + ) + + +def test_principal_default_network_and_collections(): + principal = Principal( + principal_type=PrincipalType.SERVICE_ACCOUNT, + subject="svc", + auth_method=AuthMethod.MUTUAL_TLS, + ) + assert principal.teams == [] + assert principal.scopes == [] + assert principal.project is None + assert principal.end_user is None + assert principal.network.client_ip is None + assert principal.network.via_trusted_proxy is False + + +def test_team_identity_defaults_to_member_role(): + team = TeamIdentity(id="g1") + assert team.role == TeamRole.MEMBER + + +def test_user_identity_optional_fields_default_none(): + user = UserIdentity(id="u1") + assert user.email is None + assert user.external_id is None + + +def test_project_identity_name_is_optional(): + assert ProjectIdentity(id="p1").name is None + assert ProjectIdentity(id="p1", name="Acme").name == "Acme" + + +def test_end_user_identity_requires_id(): + with pytest.raises(ValidationError): + EndUserIdentity() diff --git a/tests/test_litellm/proxy/auth/test_resolvers_seam.py b/tests/test_litellm/proxy/auth/test_resolvers_seam.py new file mode 100644 index 00000000000..cb42d4d3d2d --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_seam.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from fastapi import Request + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.resolvers.models import PrincipalType +from litellm.proxy.auth.roles import Role +from litellm.proxy.auth.user_api_key_auth import _resolve_request_principal + + +def _request( + *, + headers: Optional[Dict[str, str]] = None, + client: Optional[Tuple[str, int]] = ("203.0.113.7", 5555), +) -> Request: + raw: List[Tuple[bytes, bytes]] = [ + (k.lower().encode(), v.encode()) for k, v in (headers or {}).items() + ] + scope: Dict[str, Any] = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "path": "/v1/chat/completions", + "raw_path": b"/v1/chat/completions", + "query_string": b"", + "headers": raw, + "client": client, + "server": ("testserver", 80), + "scheme": "http", + } + return Request(scope) + + +def test_seam_projects_full_identity_from_key_object(): + token = UserAPIKeyAuth( + token="hashed-token", + user_id="u-1", + user_role="org_admin", + team_id="t-1", + team_alias="Eng", + org_id="o-1", + organization_alias="Acme", + end_user_id="cust-9", + ) + + principal = _resolve_request_principal(_request(), token) + + assert principal.principal_type == PrincipalType.HUMAN + assert principal.auth_method == AuthMethod.API_KEY + assert principal.user is not None and principal.user.id == "u-1" + assert principal.roles == [Role.ORG_ADMIN] + assert [t.id for t in principal.teams] == ["t-1"] + assert principal.teams[0].name == "Eng" + assert principal.organization is not None and principal.organization.id == "o-1" + assert principal.organization.name == "Acme" + assert principal.end_user is not None and principal.end_user.id == "cust-9" + # the key is always identifiable via credential_ref, even when other ids exist + assert principal.credential_ref.token_id == "hashed-token" + + +def test_seam_principal_is_never_anonymous_for_keyless_service_account(): + # no user_id and no key_alias -> subject must still identify the key + token = UserAPIKeyAuth(token="hashed-token") + + principal = _resolve_request_principal(_request(), token) + + assert principal.principal_type == PrincipalType.SERVICE_ACCOUNT + assert principal.user is None + assert principal.subject == "hashed-token" + assert principal.credential_ref.token_id == "hashed-token" + + +def test_seam_stamps_direct_peer_when_no_trusted_proxy_configured(): + token = UserAPIKeyAuth(token="hashed-token", user_id="u-1") + + # No trusted_proxy_ranges configured -> XFF is not trusted, direct peer wins. + principal = _resolve_request_principal( + _request(headers={"x-forwarded-for": "10.9.9.9"}, client=("203.0.113.7", 5555)), + token, + ) + + assert principal.network.client_ip == "203.0.113.7" + assert principal.network.via_trusted_proxy is False + + +def test_seam_detects_jwt_auth_method(): + token = UserAPIKeyAuth( + token="hashed-token", user_id="u-2", jwt_claims={"sub": "u-2"} + ) + + principal = _resolve_request_principal(_request(), token) + + assert principal.auth_method == AuthMethod.BEARER_JWT diff --git a/tests/test_litellm/proxy/auth/test_resolvers_store.py b/tests/test_litellm/proxy/auth/test_resolvers_store.py new file mode 100644 index 00000000000..5e644a084a5 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_store.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from typing import Dict, Optional + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth, hash_token +from litellm.proxy.auth.resolvers.exceptions import ( + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.resolvers.models import Principal, PrincipalType +from litellm.proxy.auth.resolvers.store import IdentityStore + + +class _FakeCache: + """Stands in for the DualCache get_key_object reads. It returns a cache hit + before the DB is touched, so seeding it exercises resolve without a database + (a non-None prisma client is still required; it is never reached on a hit).""" + + def __init__(self, entries: Optional[Dict[str, object]] = None) -> None: + self._entries = entries or {} + + async def async_get_cache(self, key, *args, **kwargs): + return self._entries.get(key) + + async def async_set_cache(self, *args, **kwargs): + return None + + +async def test_resolve_returns_a_principal_projected_from_the_looked_up_key(): + raw = "sk-live-abc" + key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-1", team_id="t-1") + store = IdentityStore(object(), _FakeCache({hash_token(raw): key})) + + principal = await store.resolve(hashed_token=hash_token(raw)) + + assert isinstance(principal, Principal) + assert principal.principal_type == PrincipalType.HUMAN + assert principal.user is not None and principal.user.id == "u-1" + assert [t.id for t in principal.teams] == ["t-1"] + + +async def test_resolve_carries_the_key_for_key_from_principal(): + raw = "sk-live-abc" + key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-1", team_id="t-1") + store = IdentityStore(object(), _FakeCache({hash_token(raw): key})) + + principal = await store.resolve(hashed_token=hash_token(raw)) + recovered = IdentityStore.key_from_principal(principal) + + assert recovered.user_id == "u-1" + assert recovered.team_id == "t-1" + + +def test_key_from_principal_raises_when_no_source_key_is_carried(): + bare = Principal( + principal_type=PrincipalType.SERVICE_ACCOUNT, + subject="svc", + auth_method=AuthMethod.API_KEY, + ) + with pytest.raises(PrincipalMissingSourceKeyError): + IdentityStore.key_from_principal(bare) + + +async def test_resolve_raises_without_a_db_connection(): + store = IdentityStore(None, _FakeCache()) + with pytest.raises(NoDatabaseConnectionError): + await store.resolve(hashed_token="missing") diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 52ba1dbcfbd..d623149ff6a 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -403,6 +403,48 @@ def test_virtual_key_llm_api_routes_rejects_mcp_multi_segment_admin_subpaths( assert exc_info.value.status_code == 403 +@pytest.mark.parametrize( + "route, method", + [ + ("/mcp", "POST"), + ("/mcp/", "POST"), + ("/mcp/my-server", "POST"), # matches the /mcp/{subpath} pattern + ("/mcp/tools", "GET"), + ("/mcp/tools/list", "POST"), + ("/mcp/tools/call", "POST"), + ("/mcp-rest/tools/list", "GET"), + ("/mcp-rest/tools/call", "POST"), + ("/v1/mcp/tools", "GET"), + ], +) +def test_virtual_key_llm_api_routes_allows_mcp_inference_endpoints(route, method): + """Every MCP inference/discovery endpoint must be reachable by virtual keys + scoped to allowed_routes=["llm_api_routes"], the default the Create Key UI + applies. + + /v1/mcp/tools is the most recent addition: before it joined this group a key + could list tools via /mcp/tools/list and /mcp-rest/tools/list but got a 403 + on the equivalent /v1/mcp/tools. Unlike /v1/mcp/server, none of these paths + have a management write counterpart, so they live directly in + `mcp_inference_routes` rather than behind a method-aware carve-out. + """ + + assert RouteChecks.is_llm_api_route(route=route) is True + + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["llm_api_routes"], + ) + + result = RouteChecks.is_virtual_key_allowed_to_call_route( + route=route, + valid_token=valid_token, + request=_mock_request(method), + ) + + assert result is True + + def test_spend_logs_v2_classified_as_management_not_llm_api(): """Paginated spend logs are a management/spend read route, not an LLM API.""" 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 80f12d4459f..7219ab58799 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 @@ -1106,7 +1106,7 @@ async def test_proxy_admin_expired_key_from_cache(): # Mock get_key_object to return expired token from cache with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key_object, patch( @@ -1261,7 +1261,7 @@ async def test_scim_deactivated_user_key_is_rejected(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, return_value=valid_token, ), @@ -2484,7 +2484,7 @@ async def test_user_api_key_auth_builder_no_blocking_calls(): stack.enter_context(p) stack.enter_context( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, return_value=valid_token, ) @@ -2598,7 +2598,7 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, return_value=valid_token, ), @@ -2689,6 +2689,67 @@ async def test_centralized_common_checks_runs_for_standard_auth(): setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + [ + "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + "/v1/messages", + ], +) +async def test_centralized_common_checks_routes_header_tags_to_litellm_metadata(route): + """GH#30629: on LITELLM_METADATA_ROUTES the tag-budget read resolves to + litellm_metadata, so the litellm_metadata pre-seed must run before + apply_client_tag_policy_pre_auth merges x-litellm-tags. Otherwise header tags + land in metadata and silently escape _tag_max_budget_check. This guards the + pre-seed call site in _run_centralized_common_checks; dropping it routes header + tags back into metadata. + """ + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + request = Request( + scope={ + "type": "http", + "method": "POST", + "headers": [(b"x-litellm-tags", b"tenant:acme")], + "query_string": b"", + } + ) + request._url = URL(url=route) + request_data: dict = {"model": "us.anthropic.claude-sonnet-4-6"} + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks", + new_callable=AsyncMock, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data=request_data, + route=route, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert request_data["litellm_metadata"]["tags"] == ["tenant:acme"] + assert "metadata" not in request_data + + @pytest.mark.asyncio async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): """Existing RPS guarantee: custom-auth deployments without @@ -3652,7 +3713,7 @@ async def _run_builder_with_key_lookup(get_key_object_mock): request._url = URL(url="/chat/completions") with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", get_key_object_mock, ), patch( diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index b2212068a5b..573bd5ae584 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -16,7 +16,11 @@ from unittest.mock import patch import pytest -from litellm.proxy.db.db_url_settings import DatabaseURLSettings +from litellm.proxy.db.db_url_settings import ( + DatabaseURLSettings, + unsupported_db_scheme, + unsupported_db_scheme_message, +) def _apply() -> bool: @@ -27,6 +31,7 @@ def _apply() -> bool: _MANAGED_DB_ENV_VARS = ( "IAM_TOKEN_DB_AUTH", "DATABASE_URL", + "DIRECT_URL", "DATABASE_URL_READ_REPLICA", "DATABASE_HOST", "DATABASE_PORT", @@ -287,3 +292,87 @@ def test_password_reader_uses_own_credentials(monkeypatch): os.environ["DATABASE_URL_READ_REPLICA"] == "postgresql://litellm_ro:ro_pw@reader.example.com:5432/litellm_db" ) + + +@pytest.mark.parametrize( + "url", + [ + "postgresql://u:p@host:5432/db", + "postgres://u:p@host:5432/db", + "POSTGRESQL://u:p@host:5432/db", + "postgresql://host/db?schema=public", + ], +) +def test_unsupported_db_scheme_accepts_postgres(url): + assert unsupported_db_scheme(url) is None + + +@pytest.mark.parametrize( + "url,scheme", + [ + ("sqlite:///data/litellm.db", "sqlite"), + ("sqlite:///./local.db", "sqlite"), + ("mysql://u:p@host:3306/db", "mysql"), + ("mssql://host/db", "mssql"), + ], +) +def test_unsupported_db_scheme_rejects_non_postgres(url, scheme): + assert unsupported_db_scheme(url) == scheme + + +def test_unsupported_db_scheme_does_not_echo_schemeless_credentials(): + """A malformed schemeless DSN must not leak its embedded credentials + through the return value (which callers log).""" + leaky = "litellm:s3cr3t_password@db.internal:5432/litellm" + + result = unsupported_db_scheme(leaky) + + assert result is not None + assert "s3cr3t_password" not in result + assert "db.internal" not in result + + +def test_apply_to_env_rejects_pinned_sqlite_writer(monkeypatch): + """Componentized entrypoints pin DATABASE_URL and call apply_to_env; a + sqlite writer must raise here rather than reach Prisma.""" + monkeypatch.setenv("DATABASE_URL", "sqlite:///data/litellm.db") + + with pytest.raises(RuntimeError, match="sqlite"): + _apply() + + # The bad URL must not have been propagated as a usable connection string. + assert os.environ["DATABASE_URL"] == "sqlite:///data/litellm.db" + + +def test_apply_to_env_rejects_pinned_sqlite_direct_url(monkeypatch): + """DIRECT_URL reaches Prisma the same way DATABASE_URL does; a non-postgres + direct URL must be rejected in apply_to_env, matching the CLI startup guard.""" + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") + monkeypatch.setenv("DIRECT_URL", "sqlite:///data/litellm.db") + + with pytest.raises(RuntimeError, match="DIRECT_URL.*sqlite"): + _apply() + + +def test_apply_to_env_rejects_pinned_non_postgres_reader(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") + monkeypatch.setenv( + "DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db" + ) + + with pytest.raises(RuntimeError, match="DATABASE_URL_READ_REPLICA.*mysql"): + _apply() + + +def test_apply_to_env_accepts_pinned_postgres(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@host:5432/db") + + # Operator-pinned URL: nothing reassembled, no error. + assert _apply() is False + + +def test_unsupported_db_scheme_message_names_var_and_scheme(): + msg = unsupported_db_scheme_message("DIRECT_URL", "sqlite") + assert "DIRECT_URL" in msg + assert "sqlite" in msg + assert "postgresql://" in msg diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py index 21d41e0992b..ad0e7010325 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py @@ -13,7 +13,10 @@ from litellm.proxy.management_endpoints.scim.scim_transformations import ( ScimTransformations, ) from litellm.types.proxy.management_endpoints.scim_v2 import ( + SCIM_ENTERPRISE_USER_SCHEMA, + SCIMEnterpriseUser, SCIMPatchOperation, + SCIMUser, ) @@ -149,6 +152,77 @@ class TestScimTransformations: assert scim_user.name.givenName == "Test" assert scim_user.name.familyName == "User" + @pytest.mark.asyncio + async def test_transform_user_with_enterprise_metadata(self, mock_prisma_client): + mock_client, mock_find_unique = mock_prisma_client + mock_find_unique.return_value = None + + user = LiteLLM_UserTable( + user_id="user-ent", + user_email="ent@example.com", + user_alias=None, + teams=[], + created_at=None, + updated_at=None, + metadata={ + "scim_enterprise": {"costCenter": "CC-42", "department": "Platform"} + }, + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_client): + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( + user + ) + + assert scim_user.enterprise_user is not None + assert scim_user.enterprise_user.costCenter == "CC-42" + assert scim_user.enterprise_user.department == "Platform" + assert SCIM_ENTERPRISE_USER_SCHEMA in scim_user.schemas + + @pytest.mark.asyncio + async def test_transform_user_without_enterprise_metadata_omits_schema( + self, mock_user, mock_prisma_client + ): + mock_client, mock_find_unique = mock_prisma_client + team1 = LiteLLM_TeamTable( + team_id="team-1", team_alias="Team One", members_with_roles=[] + ) + team2 = LiteLLM_TeamTable( + team_id="team-2", team_alias="Team Two", members_with_roles=[] + ) + mock_find_unique.side_effect = [team1, team2] + + with patch("litellm.proxy.proxy_server.prisma_client", mock_client): + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( + mock_user + ) + + assert scim_user.enterprise_user is None + assert SCIM_ENTERPRISE_USER_SCHEMA not in scim_user.schemas + + def test_scim_user_serialization_omits_absent_enterprise_urn(self): + without_enterprise = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + id="user-1", + userName="user@example.com", + ) + dumped = without_enterprise.model_dump(by_alias=True) + assert SCIM_ENTERPRISE_USER_SCHEMA not in dumped + assert "enterprise_user" not in dumped + assert SCIM_ENTERPRISE_USER_SCHEMA not in dumped["schemas"] + + with_enterprise = SCIMUser( + schemas=[ + "urn:ietf:params:scim:schemas:core:2.0:User", + SCIM_ENTERPRISE_USER_SCHEMA, + ], + id="user-2", + userName="ent@example.com", + enterprise_user=SCIMEnterpriseUser(costCenter="CC-42"), + ) + dumped_ent = with_enterprise.model_dump(by_alias=True) + assert dumped_ent[SCIM_ENTERPRISE_USER_SCHEMA]["costCenter"] == "CC-42" + @pytest.mark.asyncio async def test_transform_litellm_team_to_scim_group( self, mock_team, mock_prisma_client diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index ad893012807..7f5aee51f51 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -15,15 +15,19 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( _extract_group_member_ids, _handle_team_membership_changes, _process_group_patch_operations, + _recompute_scim_member_roles, create_group, create_user, + delete_group, get_users, get_service_provider_config, + patch_group, patch_user, update_group, update_user, ) from litellm.types.proxy.management_endpoints.scim_v2 import ( + SCIM_ENTERPRISE_USER_SCHEMA, SCIMGroup, SCIMMember, SCIMPatchOp, @@ -115,6 +119,59 @@ async def test_create_user_defaults_to_viewer(mocker, monkeypatch): assert called_args.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY +@pytest.mark.asyncio +async def test_create_user_ingests_enterprise_extension(mocker, monkeypatch): + """A SCIM create payload carrying the enterprise extension block should land + in the created user's metadata under scim_enterprise""" + + scim_user = SCIMUser.model_validate( + { + "schemas": [ + "urn:ietf:params:scim:schemas:core:2.0:User", + SCIM_ENTERPRISE_USER_SCHEMA, + ], + "userName": "ent-user", + "name": {"familyName": "User", "givenName": "Ent"}, + "emails": [{"value": "ent@example.com"}], + SCIM_ENTERPRISE_USER_SCHEMA: { + "costCenter": "CC-42", + "department": "Platform", + }, + } + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + + new_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.new_user", + AsyncMock(return_value=NewUserRequest(user_id="ent-user")), + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + + await create_user(user=scim_user) + + created_metadata = new_user_mock.call_args.kwargs["data"].metadata + assert created_metadata["scim_enterprise"] == { + "costCenter": "CC-42", + "department": "Platform", + } + + @pytest.mark.asyncio async def test_create_user_uses_default_internal_user_params_role(mocker, monkeypatch): """If role is set in default_internal_user_params, new user should use that role""" @@ -1720,3 +1777,1015 @@ async def test_process_group_patch_operations_with_flag_false_rejects( assert exc_info.value.status_code == 400 assert "does not exist" in str(exc_info.value.detail) assert "new-user-1" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_create_user_grants_admin_when_in_scim_admin_group(mocker, monkeypatch): + """When scim_admin_group is configured and a created user's groups include it, + the user is provisioned as PROXY_ADMIN.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="new-admin", + emails=[SCIMUserEmail(value="new-admin@example.com")], + groups=[SCIMUserGroup(value="litellm-admins", display="LiteLLM Admins")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + new_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.new_user", + AsyncMock(return_value=NewUserRequest(user_id="new-admin")), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + + await create_user(user=scim_user) + + called_args = new_user_mock.call_args.kwargs["data"] + assert called_args.user_role == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.asyncio +async def test_create_user_keeps_default_when_not_in_scim_admin_group( + mocker, monkeypatch +): + """When scim_admin_group is configured but the user's groups don't include it, + the user keeps the non-admin default role.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="regular-user", + emails=[SCIMUserEmail(value="regular@example.com")], + groups=[SCIMUserGroup(value="engineering", display="Engineering")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + new_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.new_user", + AsyncMock(return_value=NewUserRequest(user_id="regular-user")), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + + await create_user(user=scim_user) + + called_args = new_user_mock.call_args.kwargs["data"] + assert called_args.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_update_user_demotes_admin_when_removed_from_scim_admin_group( + mocker, monkeypatch +): + """Core demotion test: a PUT whose new groups no longer include the configured + admin group must re-evaluate the role and write the non-admin default, so an + admin removed from the IdP group is demoted without re-login.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + existing_user = mocker.MagicMock() + existing_user.teams = ["litellm-admins"] + existing_user.metadata = {} + + updated_user = { + "user_id": "demote-me", + "user_email": "demote@example.com", + "user_alias": None, + "teams": ["engineering"], + "metadata": "{}", + } + + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="demote-me", + emails=[SCIMUserEmail(value="demote@example.com")], + groups=[SCIMUserGroup(value="engineering", display="Engineering")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.update = AsyncMock( + return_value=updated_user + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + + await update_user(user_id="demote-me", user=scim_user) + + call_args = mock_prisma_client.db.litellm_usertable.update.call_args + assert call_args[1]["data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_update_user_does_not_force_role_when_scim_admin_group_unset( + mocker, monkeypatch +): + """When scim_admin_group is unset, PUT must not touch user_role (current + behavior preserved).""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + existing_user = mocker.MagicMock() + existing_user.teams = ["litellm-admins"] + existing_user.metadata = {} + + updated_user = { + "user_id": "no-touch", + "user_email": "no-touch@example.com", + "user_alias": None, + "teams": ["litellm-admins"], + "metadata": "{}", + } + + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="no-touch", + emails=[SCIMUserEmail(value="no-touch@example.com")], + groups=[SCIMUserGroup(value="litellm-admins", display="LiteLLM Admins")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.update = AsyncMock( + return_value=updated_user + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + + await update_user(user_id="no-touch", user=scim_user) + + call_args = mock_prisma_client.db.litellm_usertable.update.call_args + assert "user_role" not in call_args[1]["data"] + + +@pytest.mark.asyncio +async def test_update_user_demotes_when_default_params_lack_user_role( + mocker, monkeypatch +): + """Regression: default_internal_user_params set without a user_role key must + still resolve to the non-admin default on demotion, not silently skip and + leave the user PROXY_ADMIN.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr( + "litellm.default_internal_user_params", {"max_budget": 10}, raising=False + ) + + existing_user = mocker.MagicMock() + existing_user.teams = ["litellm-admins"] + existing_user.metadata = {} + + updated_user = { + "user_id": "demote-me", + "user_email": "demote@example.com", + "user_alias": None, + "teams": ["engineering"], + "metadata": "{}", + } + + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="demote-me", + emails=[SCIMUserEmail(value="demote@example.com")], + groups=[SCIMUserGroup(value="engineering", display="Engineering")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.update = AsyncMock( + return_value=updated_user + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + + await update_user(user_id="demote-me", user=scim_user) + + call_args = mock_prisma_client.db.litellm_usertable.update.call_args + assert call_args[1]["data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group( + mocker, monkeypatch +): + """PATCH that drops the admin team from the resulting team set must write the + non-admin default, mirroring the PUT demotion path.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + existing_user = mocker.MagicMock() + existing_user.teams = ["litellm-admins"] + existing_user.metadata = {} + + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation( + op="replace", path="groups", value=[{"value": "engineering"}] + ) + ], + ) + + updated_user = { + "user_id": "demote-me", + "user_alias": None, + "teams": ["engineering"], + "metadata": "{}", + } + + engineering_team = mocker.MagicMock() + engineering_team.team_alias = "Engineering" + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.update = AsyncMock( + return_value=updated_user + ) + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=engineering_team + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock( + return_value=SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="demote-me", + ) + ), + ) + + await patch_user(user_id="demote-me", patch_ops=patch_ops) + + call_args = mock_prisma_client.db.litellm_usertable.update.call_args + assert call_args[1]["data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_patch_user_grants_admin_by_team_display_name(mocker, monkeypatch): + """PATCH carries groups as team ids, so admin-group matching must fall back to + each team's display name; an admin group configured as a human-readable alias + grants PROXY_ADMIN even when the team id differs.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "LiteLLM Admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + existing_user = mocker.MagicMock() + existing_user.teams = [] + existing_user.metadata = {} + + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation( + op="replace", path="groups", value=[{"value": "team-abc-123"}] + ) + ], + ) + + updated_user = { + "user_id": "promote-me", + "user_alias": None, + "teams": ["team-abc-123"], + "metadata": "{}", + } + + admin_team = mocker.MagicMock() + admin_team.team_alias = "LiteLLM Admins" + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.update = AsyncMock( + return_value=updated_user + ) + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=admin_team + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes", + AsyncMock(), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock( + return_value=SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="promote-me", + ) + ), + ) + + await patch_user(user_id="promote-me", patch_ops=patch_ops) + + call_args = mock_prisma_client.db.litellm_usertable.update.call_args + assert call_args[1]["data"]["user_role"] == LitellmUserRoles.PROXY_ADMIN + + +def _scim_admin_prisma(mocker, *, user_teams): + """Prisma double whose user resolves to user_teams and whose teams expose an + alias equal to their id, used by the role-recompute helper tests.""" + user = mocker.MagicMock() + user.user_id = "member-1" + user.teams = user_teams + + def _team_find_unique(where): + team = mocker.MagicMock() + team.team_alias = where["team_id"] + return team + + prisma = mocker.MagicMock() + prisma.db = mocker.MagicMock() + prisma.db.litellm_usertable = mocker.MagicMock() + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user) + prisma.db.litellm_usertable.update = AsyncMock(return_value=user) + prisma.db.litellm_teamtable = mocker.MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(side_effect=_team_find_unique) + return prisma + + +@pytest.mark.asyncio +async def test_recompute_scim_member_roles_demotes_when_not_in_admin_group( + mocker, monkeypatch +): + """The shared recompute helper writes the non-admin default for a member whose + resulting teams no longer include the configured admin group.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + prisma = _scim_admin_prisma(mocker, user_teams=["engineering"]) + + await _recompute_scim_member_roles(prisma, ["member-1"]) + + call_args = prisma.db.litellm_usertable.update.call_args + assert call_args[1]["data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_recompute_scim_member_roles_grants_when_in_admin_group( + mocker, monkeypatch +): + """The shared recompute helper grants PROXY_ADMIN when a member's resulting + teams include the configured admin group.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + prisma = _scim_admin_prisma(mocker, user_teams=["litellm-admins"]) + + await _recompute_scim_member_roles(prisma, ["member-1"]) + + call_args = prisma.db.litellm_usertable.update.call_args + assert call_args[1]["data"]["user_role"] == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.asyncio +async def test_recompute_scim_member_roles_noop_when_admin_group_unset( + mocker, monkeypatch +): + """With scim_admin_group unset the recompute helper must not touch any role, + preserving current behavior for SCIM group writes.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + prisma = _scim_admin_prisma(mocker, user_teams=["litellm-admins"]) + + await _recompute_scim_member_roles(prisma, ["member-1"]) + + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_group_recomputes_roles_for_changed_members(mocker): + """PUT /Groups must recompute the global role for every member whose + membership changed, so an admin dropped from the admin group is demoted.""" + from litellm.proxy._types import LiteLLM_TeamTable, Member + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "test-team-123" + existing_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Admins", + members=["user1", "user2"], + members_with_roles=[ + Member(user_id="user1", role="user"), + Member(user_id="user2", role="user"), + ], + metadata={}, + ) + scim_group_update = SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Admins", + members=[SCIMMember(value="user2"), SCIMMember(value="user3")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock() + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + recompute_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock(return_value=scim_group_update), + ) + + await update_group(group_id=group_id, group=scim_group_update) + + recompute_mock.assert_awaited_once() + assert set(recompute_mock.call_args[0][1]) == {"user1", "user3"} + + +@pytest.mark.asyncio +async def test_patch_group_recomputes_roles_for_changed_members(mocker): + """PATCH /Groups must recompute the global role for every member whose + membership changed, mirroring the PUT path.""" + from litellm.proxy._types import LiteLLM_TeamTable, Member + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "test-team-123" + existing_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Admins", + members=["user1", "user2"], + members_with_roles=[ + Member(user_id="user1", role="user"), + Member(user_id="user2", role="user"), + ], + metadata={}, + ) + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation(op="remove", path="members", value=[{"value": "user1"}]) + ], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock() + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + recompute_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock( + return_value=SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Admins", + ) + ), + ) + + await patch_group(group_id=group_id, patch_ops=patch_ops) + + recompute_mock.assert_awaited_once() + assert set(recompute_mock.call_args[0][1]) == {"user1"} + + +@pytest.mark.asyncio +async def test_delete_group_recomputes_roles_for_members(mocker): + """DELETE /Groups must recompute the global role for the team's members, so + deleting the admin group demotes everyone who was only admin through it.""" + from litellm.proxy._types import Member + + existing_team = mocker.MagicMock() + existing_team.members_with_roles = [ + Member(user_id="user1", role="user"), + Member(user_id="user2", role="user"), + ] + + member = mocker.MagicMock() + member.teams = ["test-team-123"] + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_teamtable.delete = AsyncMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=member) + mock_prisma_client.db.litellm_usertable.update = AsyncMock() + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + recompute_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + + await delete_group(group_id="test-team-123") + + recompute_mock.assert_awaited_once() + assert list(recompute_mock.call_args[0][1]) == ["user1", "user2"] + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_applies_role_when_admin_group_set(mocker): + """When admin_group is configured, re-upserting an existing email persists the + resolved role so a now-non-admin user can't keep a stale PROXY_ADMIN.""" + existing_user = mocker.MagicMock() + existing_user.user_id = "old-user-id" + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + return_value=existing_user + ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock( + return_value={"user_id": "new-user-id"} + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=mocker.MagicMock()), + ) + + new_user_request = NewUserRequest( + user_id="new-user-id", + user_email="test@example.com", + teams=["engineering"], + metadata={}, + auto_create_key=False, + user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, + new_user_request=new_user_request, + admin_group="litellm-admins", + ) + + call_args = mock_prisma_client.db.litellm_usertable.update.call_args + assert call_args[1]["data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_leaves_role_when_admin_group_unset(mocker): + """With admin_group unset, the existing-email upsert must not write user_role, + preserving current behavior when the feature is off.""" + existing_user = mocker.MagicMock() + existing_user.user_id = "old-user-id" + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + return_value=existing_user + ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock( + return_value={"user_id": "new-user-id"} + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=mocker.MagicMock()), + ) + + new_user_request = NewUserRequest( + user_id="new-user-id", + user_email="test@example.com", + teams=["engineering"], + metadata={}, + auto_create_key=False, + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, + new_user_request=new_user_request, + ) + + call_args = mock_prisma_client.db.litellm_usertable.update.call_args + assert "user_role" not in call_args[1]["data"] + + +@pytest.mark.asyncio +async def test_create_user_existing_email_upsert_demotes_when_admin_group_set( + mocker, monkeypatch +): + """End-to-end create wiring: a SCIM POST that upserts an existing email while + the user is not in the admin group must write the non-admin default, not leave + a stale PROXY_ADMIN.""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="returning-user", + emails=[SCIMUserEmail(value="returning@example.com")], + groups=[SCIMUserGroup(value="engineering", display="Engineering")], + ) + + existing_user = mocker.MagicMock() + existing_user.user_id = "returning-user" + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + return_value=existing_user + ) + mock_prisma_client.db.litellm_usertable.update = AsyncMock( + return_value={"user_id": "returning-user"} + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + new_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.new_user", + AsyncMock(return_value=NewUserRequest(user_id="returning-user")), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + + await create_user(user=scim_user) + + new_user_mock.assert_not_called() + call_args = mock_prisma_client.db.litellm_usertable.update.call_args + assert call_args[1]["data"]["user_role"] == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + + +@pytest.mark.asyncio +async def test_create_group_recomputes_roles_for_members(mocker): + """POST /Groups must recompute the global role for the new team's members, so a + team created with the admin-group display name elevates its members.""" + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "admin-team-1" + scim_group = SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="LiteLLM Admins", + members=[SCIMMember(value="user1"), SCIMMember(value="user2")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock() + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.new_team", + AsyncMock(return_value=mocker.MagicMock()), + ) + recompute_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock(return_value=scim_group), + ) + + await create_group(group=scim_group) + + recompute_mock.assert_awaited_once() + assert set(recompute_mock.call_args[0][1]) == {"user1", "user2"} + + +@pytest.mark.asyncio +async def test_update_group_rename_recomputes_retained_members(mocker): + """A PUT that renames the group (alias changes) but leaves membership unchanged + must still recompute retained members, since a rename can flip whether the + group matches scim_admin_group by display name.""" + from litellm.proxy._types import LiteLLM_TeamTable, Member + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "test-team-123" + existing_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="LiteLLM Admins", + members=["user1"], + members_with_roles=[Member(user_id="user1", role="user")], + metadata={}, + ) + scim_group_update = SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Engineering", + members=[SCIMMember(value="user1")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock() + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + recompute_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock(return_value=scim_group_update), + ) + + await update_group(group_id=group_id, group=scim_group_update) + + recompute_mock.assert_awaited_once() + assert set(recompute_mock.call_args[0][1]) == {"user1"} + + +@pytest.mark.asyncio +async def test_patch_group_rename_recomputes_retained_members(mocker): + """A PATCH that renames the group (displayName op) but leaves membership + unchanged must still recompute retained members, mirroring the PUT path.""" + from litellm.proxy._types import LiteLLM_TeamTable, Member + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + + group_id = "test-team-123" + existing_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="LiteLLM Admins", + members=["user1"], + members_with_roles=[Member(user_id="user1", role="user")], + metadata={}, + ) + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[ + SCIMPatchOperation(op="replace", path="displayName", value="Engineering") + ], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mocker.MagicMock() + ) + + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + recompute_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles", + AsyncMock(), + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock( + return_value=SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Engineering", + ) + ), + ) + + await patch_group(group_id=group_id, patch_ops=patch_ops) + + recompute_mock.assert_awaited_once() + assert set(recompute_mock.call_args[0][1]) == {"user1"} diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 2d2c18bb46e..b882090e8f1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -638,6 +638,84 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): assert key_data.metrics.spend == 10.0 +def _daily_user_spend_record(*, user_id, api_key, spend): + """A LiteLLM_DailyUserSpend row as the per-user breakdown reads it.""" + return SimpleNamespace( + date="2024-01-01", + user_id=user_id, + api_key=api_key, + model="gpt-4", + model_group="gpt-4", + custom_llm_provider="openai", + mcp_namespaced_tool_name=None, + endpoint="/chat/completions", + spend=spend, + prompt_tokens=10, + completion_tokens=5, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + api_requests=1, + successful_requests=1, + failed_requests=0, + ) + + +@pytest.mark.asyncio +async def test_get_daily_activity_applies_resolve_entity_metadata_to_breakdown(): + """Regression for LIT-3889: the Spend Per User chart showed raw UUIDs. + + /user/daily/activity used to pass entity_metadata_field=None, so every + user entity in the breakdown carried empty metadata and the dashboard had + nothing to render but the user_id UUID. The page-scoped resolver must put + the resolved email/alias onto the entity metadata so the UI can label it, + while a spender with no email on file still falls back to the raw UUID. + """ + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + + records = [ + _daily_user_spend_record(user_id="user-with-email", api_key="key-1", spend=7.0), + _daily_user_spend_record(user_id="user-no-email", api_key="key-2", spend=3.0), + ] + + mock_table = MagicMock() + mock_table.count = AsyncMock(return_value=len(records)) + mock_table.find_many = AsyncMock(return_value=records) + mock_prisma.db.litellm_dailyuserspend = mock_table + mock_prisma.db.litellm_verificationtoken = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + seen_user_ids = {} + + async def resolver(page_records): + seen_user_ids["ids"] = {r.user_id for r in page_records} + return {"user-with-email": {"user_email": "spender@example.com"}} + + result = await get_daily_activity( + prisma_client=mock_prisma, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2024-01-01", + end_date="2024-01-01", + model=None, + api_key=None, + page=1, + page_size=1000, + resolve_entity_metadata=resolver, + ) + + # Resolver is driven by the user_ids actually on the page + assert seen_user_ids["ids"] == {"user-with-email", "user-no-email"} + + entities = result.results[0].breakdown.entities + # Email is on the entity metadata so the UI labels the chart with it + assert entities["user-with-email"].metadata["user_email"] == "spender@example.com" + # No email on file -> empty metadata -> UI falls back to the UUID + assert entities["user-no-email"].metadata == {} + + class TestAdjustDatesForTimezone: """ Regression tests for the timezone double-counting bug. @@ -758,6 +836,8 @@ class TestBuildAggregatedSqlQuery: ] assert "model = $4" in sql assert "api_key = $5" in sql + + @pytest.mark.asyncio async def test_get_daily_activity_aggregated_empty_result_set(): """Regression test for the empty-range 500. diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 627958cef93..b4602e0ad8b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2,6 +2,7 @@ import json import os import sys from datetime import datetime, timezone +from types import SimpleNamespace import pytest from fastapi.testclient import TestClient @@ -20,6 +21,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.management_endpoints.internal_user_endpoints import ( LiteLLM_UserTableWithKeyCount, + _resolve_user_email_metadata, _update_internal_user_params, get_user_key_counts, get_users, @@ -657,6 +659,56 @@ async def test_get_users_includes_timestamps(mocker): assert user_response.key_count == 0 +@pytest.mark.asyncio +async def test_get_users_redacts_scim_enterprise_metadata(mocker): + """ + /user/list must strip scim_enterprise from each user's metadata while leaving + the rest of the metadata intact, matching the user-info endpoints. + """ + mock_prisma_client = mocker.MagicMock() + + mock_user_row = mocker.MagicMock() + mock_user_row.user_id = "listed-user" + mock_user_row.model_dump.return_value = { + "user_id": "listed-user", + "user_email": "listed@example.com", + "user_role": "internal_user", + "metadata": { + "scim_metadata": {"givenName": "Jane", "familyName": "Doe"}, + "scim_enterprise": {"costCenter": "CC-42", "department": "Platform"}, + }, + } + + async def mock_find_many(*args, **kwargs): + return [mock_user_row] + + async def mock_count(*args, **kwargs): + return 1 + + mock_prisma_client.db.litellm_usertable.find_many = mock_find_many + mock_prisma_client.db.litellm_usertable.count = mock_count + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + async def mock_get_user_key_counts(*args, **kwargs): + return {"listed-user": 0} + + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_user_key_counts", + mock_get_user_key_counts, + ) + + admin_key = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + response = await get_users( + page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None + ) + + listed = response["users"][0] + assert listed.metadata == { + "scim_metadata": {"givenName": "Jane", "familyName": "Doe"} + } + assert "scim_enterprise" not in (listed.metadata or {}) + + def test_validate_sort_params(): """ Test that validate_sort_params returns None if sort_by is None @@ -2167,6 +2219,94 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker): assert response.metadata == {"team": "engineering"} +@pytest.mark.asyncio +async def test_user_info_v2_redacts_scim_enterprise_metadata(mocker): + """ + SCIM enterprise attributes are persisted in metadata for reporting, but + /v2/user/info must not surface them; the rest of metadata is preserved. + """ + from fastapi import Request + + from litellm.proxy._types import UserInfoV2Response + from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2 + + mock_prisma_client = mocker.MagicMock() + + mock_user_row = mocker.MagicMock() + mock_user_row.model_dump.return_value = { + "user_id": "target-user-123", + "user_email": "target@example.com", + "metadata": { + "scim_metadata": {"givenName": "Jane", "familyName": "Doe"}, + "scim_enterprise": { + "costCenter": "CC-42", + "department": "Platform", + "employeeNumber": "E-1001", + }, + }, + "teams": ["team-1"], + } + + async def mock_find_unique(*args, **kwargs): + if kwargs.get("where", {}).get("user_id") == "target-user-123": + return mock_user_row + return None + + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( + side_effect=mock_find_unique + ) + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_request = mocker.MagicMock(spec=Request) + + admin_key = UserAPIKeyAuth( + user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + response = await user_info_v2( + request=mock_request, + user_id="target-user-123", + user_api_key_dict=admin_key, + ) + + assert isinstance(response, UserInfoV2Response) + assert response.metadata == { + "scim_metadata": {"givenName": "Jane", "familyName": "Doe"} + } + assert "scim_enterprise" not in (response.metadata or {}) + + +def test_build_user_info_response_redacts_scim_enterprise_metadata(): + """ + The shared /user/info builder strips scim_enterprise from the returned user row + while leaving every other metadata key intact. + """ + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _build_user_info_response, + ) + + user_row = { + "user_id": "target-user-123", + "metadata": { + "scim_metadata": {"givenName": "Jane"}, + "scim_enterprise": {"costCenter": "CC-42"}, + }, + } + + response = _build_user_info_response( + user_id="target-user-123", + user_info=user_row, + keys=None, + team_list=[], + teams_1=None, + ) + + assert response.user_info is not None + assert response.user_info["metadata"] == {"scim_metadata": {"givenName": "Jane"}} + assert "scim_enterprise" not in response.user_info["metadata"] + + @pytest.mark.asyncio async def test_user_info_v2_internal_user_can_query_self(mocker): """ @@ -2960,3 +3100,56 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): user_request=user_request, user_api_key_dict=admin_caller ) assert result is not None + + +@pytest.mark.asyncio +async def test_resolve_user_email_metadata_maps_page_user_ids_to_email(mocker): + """Regression for LIT-3889. + + The Spend Per User chart rendered raw UUIDs because the per-user activity + breakdown carried no email. This resolver must turn the user_ids on the + page into {user_id: {user_email, user_alias}} so the chart can label each + spender, and it must only look up the user_ids actually present (not the + whole user table). + """ + + mock_prisma_client = mocker.MagicMock() + find_many = mocker.AsyncMock( + return_value=[ + SimpleNamespace( + user_id="u1", user_email="alice@example.com", user_alias="Alice" + ), + SimpleNamespace(user_id="u2", user_email=None, user_alias="bob-alias"), + ] + ) + mock_prisma_client.db.litellm_usertable.find_many = find_many + + records = [ + SimpleNamespace(user_id="u1"), + SimpleNamespace(user_id="u1"), # duplicate -> deduped + SimpleNamespace(user_id="u2"), + ] + + result = await _resolve_user_email_metadata(mock_prisma_client, records) + + assert result == { + "u1": {"user_email": "alice@example.com", "user_alias": "Alice"}, + "u2": {"user_email": None, "user_alias": "bob-alias"}, + } + where_arg = find_many.call_args.kwargs["where"] + assert set(where_arg["user_id"]["in"]) == {"u1", "u2"} + + +@pytest.mark.asyncio +async def test_resolve_user_email_metadata_skips_db_when_no_user_ids(mocker): + """No user_ids on the page (e.g. all spend is unattributed) means no query.""" + mock_prisma_client = mocker.MagicMock() + find_many = mocker.AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = find_many + + records = [SimpleNamespace(user_id=None), SimpleNamespace(user_id="")] + + result = await _resolve_user_email_metadata(mock_prisma_client, records) + + assert result == {} + find_many.assert_not_called() 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 0b5b5fb6ceb..ce6c7e9b6fa 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 @@ -1197,6 +1197,136 @@ class TestListMCPServers: # Non-admin viewers get no env var config at all (not even names). assert result.env_vars is None + @pytest.mark.asyncio + async def test_fetch_single_mcp_server_sanitizes_for_view_only_admin(self): + """PROXY_ADMIN_VIEW_ONLY must NOT see credential-bearing fields. + + It previously passed the _user_has_admin_view gate (which also grants view-only + admins) and only had the explicit `credentials` field cleared, leaking secrets + embedded in url/static_headers/env_vars. Only a FULL PROXY_ADMIN may see those. + This test exercises the real role helpers (no patching of the gate).""" + mock_server = LiteLLM_MCPServerTable.model_construct( + server_id="leaky-server", + server_name="Leaky Server", + alias="Leaky Server", + transport=MCPTransport.http, + url="https://leaky.example.com/mcp?api_key=sk-embedded-in-url", + static_headers={"Authorization": "Bearer sk-secret-header"}, + env={"UPSTREAM_TOKEN": "sk-secret-env"}, + env_vars=[ + {"name": "GLOBAL_KEY", "value": "super-secret", "scope": "global"}, + ], + credentials={"auth_value": "sk-explicit-credential"}, + ) + + mock_prisma_client = MagicMock() + + mock_health_result = generate_mock_mcp_server_db_record( + server_id="leaky-server", alias="Leaky Server" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + mock_manager = MagicMock() + mock_manager.add_server = AsyncMock() + mock_manager.health_check_server = AsyncMock(return_value=mock_health_result) + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=mock_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_server, + ) + + result = await fetch_mcp_server( + request=_make_mock_request(), + server_id="leaky-server", + user_api_key_dict=mock_user_auth, + ) + + assert result.server_id == "leaky-server" + assert result.credentials is None + assert result.url is None + assert result.static_headers is None + assert result.env == {} + assert result.env_vars is None + + @pytest.mark.asyncio + async def test_fetch_single_mcp_server_full_admin_still_sees_secrets(self): + """the fix must not over-redact for FULL PROXY_ADMIN, + who needs url/static_headers/env to populate the edit form.""" + mock_server = LiteLLM_MCPServerTable.model_construct( + server_id="admin-server", + server_name="Admin Server", + alias="Admin Server", + transport=MCPTransport.http, + url="https://admin.example.com/mcp", + static_headers={"Authorization": "Bearer sk-secret-header"}, + credentials={"auth_value": "sk-explicit-credential"}, + ) + + mock_prisma_client = MagicMock() + + mock_health_result = generate_mock_mcp_server_db_record( + server_id="admin-server", alias="Admin Server" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + mock_manager = MagicMock() + mock_manager.add_server = AsyncMock() + mock_manager.health_check_server = AsyncMock(return_value=mock_health_result) + + mock_user_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=mock_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_server, + ) + + result = await fetch_mcp_server( + request=_make_mock_request(), + server_id="admin-server", + user_api_key_dict=mock_user_auth, + ) + + # credentials field is always redacted; the rest must survive for full admin. + assert result.credentials is None + assert result.url == "https://admin.example.com/mcp" + assert result.static_headers == {"Authorization": "Bearer sk-secret-header"} + class TestTeamScopedMCPServerAccess: """Tests for cross-team information disclosure and restricted key bypass fixes.""" @@ -3693,18 +3823,12 @@ def _server_with_env_vars(server_id: str = "srv-env"): @pytest.mark.asyncio -@pytest.mark.parametrize( - "user_role, expected_global_value", - [ - (LitellmUserRoles.PROXY_ADMIN, "super-secret"), - (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ""), - ], -) -async def test_fetch_single_mcp_server_redacts_global_env_for_view_only_admin( - user_role, expected_global_value -): - """Read-only admins must not receive admin-supplied global env var secrets; - full admins still see them so the edit form can pre-fill.""" +async def test_fetch_single_mcp_server_env_vars_full_admin_vs_view_only(): + """full admins see admin-supplied global env var secrets so the edit + form can pre-fill; read-only admins now go through the non-admin sanitizer, which + drops env_vars entirely (the names alone, e.g. ADMIN_API_KEY, leak what secrets the + admin configured). Previously the view-only case merely blanked the global value + while keeping the names, which still leaked configuration metadata.""" server = _server_with_env_vars() health_result = generate_mock_mcp_server_db_record(server_id=server.server_id) @@ -3712,34 +3836,39 @@ async def test_fetch_single_mcp_server_redacts_global_env_for_view_only_admin( health_result.last_health_check = datetime.now() health_result.health_check_error = None - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=MagicMock(), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=server), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_server", - AsyncMock(return_value=None), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.health_check_server", - AsyncMock(return_value=health_result), - ), - ): - result = await mgmt_endpoints.fetch_mcp_server( - request=_make_mock_request(), - server_id=server.server_id, - user_api_key_dict=generate_mock_user_api_key_auth(user_role=user_role), - ) + async def _fetch(user_role): + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.health_check_server", + AsyncMock(return_value=health_result), + ), + ): + return await mgmt_endpoints.fetch_mcp_server( + request=_make_mock_request(), + server_id=server.server_id, + user_api_key_dict=generate_mock_user_api_key_auth(user_role=user_role), + ) - by_name = {ev.name: ev for ev in result.env_vars} - assert by_name["ADMIN_API_KEY"].value == expected_global_value - # Per-user placeholders are always preserved. + full_admin = await _fetch(LitellmUserRoles.PROXY_ADMIN) + by_name = {ev.name: ev for ev in full_admin.env_vars} + assert by_name["ADMIN_API_KEY"].value == "super-secret" assert by_name["USER_TOKEN"].value == "placeholder-hint" + + view_only = await _fetch(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) + assert view_only.env_vars is None + # The source record must never be mutated. assert {ev.name: ev.value for ev in server.env_vars}[ "ADMIN_API_KEY" @@ -3747,18 +3876,67 @@ async def test_fetch_single_mcp_server_redacts_global_env_for_view_only_admin( @pytest.mark.asyncio -@pytest.mark.parametrize( - "user_role, expected_global_value", - [ - (LitellmUserRoles.PROXY_ADMIN, "super-secret"), - (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ""), - ], -) -async def test_fetch_all_mcp_servers_redacts_global_env_for_view_only_admin( - user_role, expected_global_value -): +async def test_fetch_all_mcp_servers_env_vars_full_admin_vs_view_only(): + """same posture as the single-server fetch. Full admins + keep the env var values; view-only admins get env_vars dropped via the non-admin + sanitizer rather than only having the global value blanked.""" server = _server_with_env_vars() + async def _fetch_all(user_role): + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", + return_value="view_all", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.get_all_mcp_servers_unfiltered", + AsyncMock(return_value=[server]), + ), + patch( + "litellm.proxy.proxy_server.prisma_client", + None, + ), + ): + return await mgmt_endpoints.fetch_all_mcp_servers( + user_api_key_dict=generate_mock_user_api_key_auth(user_role=user_role), + ) + + full_admin = await _fetch_all(LitellmUserRoles.PROXY_ADMIN) + by_name = {ev.name: ev for ev in full_admin[0].env_vars} + assert by_name["ADMIN_API_KEY"].value == "super-secret" + assert by_name["USER_TOKEN"].value == "placeholder-hint" + + view_only = await _fetch_all(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) + assert view_only[0].env_vars is None + + assert {ev.name: ev.value for ev in server.env_vars}[ + "ADMIN_API_KEY" + ] == "super-secret" + + +def _leaky_list_server() -> "LiteLLM_MCPServerTable": + """A server whose url/static_headers/env carry embedded secrets, for the + list-endpoint sanitization tests. ``model_construct`` skips validation so + the raw values survive verbatim.""" + return LiteLLM_MCPServerTable.model_construct( + server_id="leaky-list-server", + server_name="Leaky List Server", + alias="Leaky List Server", + transport=MCPTransport.http, + url="https://leaky.example.com/mcp?api_key=sk-embedded-in-url", + static_headers={"Authorization": "Bearer sk-secret-header"}, + env={"UPSTREAM_TOKEN": "sk-secret-env"}, + env_vars=[ + {"name": "GLOBAL_KEY", "value": "super-secret", "scope": "global"}, + ], + credentials={"auth_value": "sk-explicit-credential"}, + ) + + +async def _fetch_all_via_view_all(user_role: LitellmUserRoles): + """Drive GET /v1/mcp/server in view_all mode for the given role using the + real role helpers (the full-admin gate is never patched).""" + server = _leaky_list_server() with ( patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", @@ -3776,13 +3954,46 @@ async def test_fetch_all_mcp_servers_redacts_global_env_for_view_only_admin( result = await mgmt_endpoints.fetch_all_mcp_servers( user_api_key_dict=generate_mock_user_api_key_auth(user_role=user_role), ) + return server, result - by_name = {ev.name: ev for ev in result[0].env_vars} - assert by_name["ADMIN_API_KEY"].value == expected_global_value - assert by_name["USER_TOKEN"].value == "placeholder-hint" - assert {ev.name: ev.value for ev in server.env_vars}[ - "ADMIN_API_KEY" - ] == "super-secret" + +@pytest.mark.asyncio +async def test_list_mcp_servers_sanitized_for_view_only_admin(): + """PROXY_ADMIN_VIEW_ONLY listing servers must go through the non-admin + sanitizer: url and static_headers cleared, env emptied, env_vars dropped. + A mutation swapping _user_is_full_admin() back to _user_has_admin_view() + (which also grants view-only admins) would return the raw url/headers and + fail this. The real role helpers are exercised; the gate is not patched.""" + source, result = await _fetch_all_via_view_all( + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ) + + assert len(result) == 1 + sanitized = result[0] + assert sanitized.server_id == "leaky-list-server" + assert sanitized.url is None + assert sanitized.static_headers is None + assert sanitized.env == {} + assert sanitized.env_vars is None + assert sanitized.credentials is None + + # The source record must never be mutated by sanitization. + assert source.url == "https://leaky.example.com/mcp?api_key=sk-embedded-in-url" + assert source.static_headers == {"Authorization": "Bearer sk-secret-header"} + + +@pytest.mark.asyncio +async def test_list_mcp_servers_full_admin_still_sees_secrets(): + """The view-only redaction must not over-redact for a FULL PROXY_ADMIN, + who needs url/static_headers to populate the edit form. Only the explicit + credentials field is cleared for full admins on the list endpoint.""" + _, result = await _fetch_all_via_view_all(LitellmUserRoles.PROXY_ADMIN) + + assert len(result) == 1 + raw = result[0] + assert raw.url == "https://leaky.example.com/mcp?api_key=sk-embedded-in-url" + assert raw.static_headers == {"Authorization": "Bearer sk-secret-header"} + assert raw.credentials is None def _make_env_var_server( diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index 965580e8758..c77ac11ffc1 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -14,6 +14,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( _extract_requested_mcp_access_groups, _extract_requested_mcp_server_ids, _resolve_team_allowed_mcp_servers, + _rewrite_object_permission_mcp_servers, _set_object_permission, validate_key_mcp_servers_against_team, validate_key_search_tools_against_team, @@ -111,6 +112,31 @@ def test_extract_requested_mcp_server_ids_none(): assert _extract_requested_mcp_server_ids({}) == set() +def test_extract_requested_mcp_server_ids_excludes_no_mcp_servers_sentinel(): + obj_perm = {"mcp_servers": ["no-mcp-servers", "server-1"]} + assert _extract_requested_mcp_server_ids(obj_perm) == {"server-1"} + + +def test_rewrite_object_permission_mcp_servers_preserves_sentinel(): + obj_perm = {"mcp_servers": ["no-mcp-servers", "alias-1"]} + _rewrite_object_permission_mcp_servers(obj_perm, {"alias-1": {"server-1"}}) + assert obj_perm["mcp_servers"] == ["no-mcp-servers", "server-1"] + + +@pytest.mark.asyncio +async def test_validate_no_mcp_servers_sentinel_passes_and_preserved(): + """A key scoped to no-mcp-servers passes team validation untouched, keeping the + sentinel so it is not mistaken for an unknown server and rejected.""" + team_obj = _make_team_obj(mcp_servers=["server-1"]) + obj_perm = {"mcp_servers": ["no-mcp-servers"]} + result = await validate_key_mcp_servers_against_team( + object_permission=obj_perm, + team_obj=team_obj, + ) + assert result == obj_perm + assert obj_perm["mcp_servers"] == ["no-mcp-servers"] + + # ---- Tests for _extract_requested_mcp_access_groups ---- diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index b800c82c75d..947a7a64beb 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -1885,3 +1885,310 @@ class TestNonStreamingResponseRedaction: leaked = logging_obj.model_call_details.get("complete_streaming_response") assert leaked is None assert redacted.choices[0].message.content == "redacted-by-litellm" + + +def _sse_bytes(data: dict) -> bytes: + return f"event: {data['type']}\ndata: {json.dumps(data)}\n\n".encode() + + +class TestAnthropicUsageOnlyFallback: + """When stream_chunk_builder cannot reassemble a large/agentic stream (returns + None or raises), Anthropic still emits token usage in the message_start / + message_delta SSE events. The handler must recover usage-only so the request is + priced instead of being dropped from SpendLogs while Anthropic billed the tokens.""" + + _CHUNKS = [ + _sse_bytes( + { + "type": "message_start", + "message": { + "model": "claude-3-5-haiku-20241022", + "usage": { + "input_tokens": 100, + "cache_read_input_tokens": 40, + "cache_creation_input_tokens": 20, + "output_tokens": 1, + }, + }, + } + ), + _sse_bytes( + { + "type": "message_delta", + "usage": { + "output_tokens": 55, + "server_tool_use": {"web_search_requests": 2}, + }, + } + ), + ] + + def test_build_usage_only_recovers_cache_inclusive_usage(self): + response = ( + AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=self._CHUNKS, model="claude-3-5-haiku-20241022" + ) + ) + assert response is not None + usage = response.usage + # prompt_tokens must be cache-inclusive (input + cache_read + cache_creation) + assert usage.prompt_tokens == 160 + assert usage.completion_tokens == 55 + assert usage._cache_read_input_tokens == 40 + assert usage._cache_creation_input_tokens == 20 + assert usage.prompt_tokens_details.cached_tokens == 40 + assert usage.server_tool_use.web_search_requests == 2 + + def test_build_usage_only_returns_none_without_usage_events(self): + chunks = [_sse_bytes({"type": "content_block_delta", "delta": {"text": "hi"}})] + assert ( + AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=chunks, model="claude-3-5-haiku-20241022" + ) + is None + ) + + def test_build_usage_only_recovers_cache_split_server_tools_and_model(self): + # the model is "unknown" up-front and only the 5m/1h cache split is sent + # (no flat cache_creation_input_tokens); web/tool-search and geo arrive in + # message_delta. All must be recovered and priced, not left at $0. + chunks = [ + "event: ping\ndata: [DONE]\n\n", # ignored sentinel between real events + _sse_bytes( + { + "type": "message_start", + "message": { + "model": "claude-opus-4-6", + "usage": { + "input_tokens": 80, + "output_tokens": 1, + "cache_creation": { + "ephemeral_5m_input_tokens": 12, + "ephemeral_1h_input_tokens": 8, + }, + "inference_geo": "us", + }, + }, + } + ), + _sse_bytes( + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use"}, + "usage": { + "output_tokens": 40, + "cache_read_input_tokens": 5, + "inference_geo": "us", + "server_tool_use": { + "web_search_requests": 1, + "tool_search_requests": 3, + }, + }, + } + ), + ] + response = ( + AnthropicPassthroughLoggingHandler._build_usage_only_response_from_chunks( + all_chunks=chunks, model="unknown" + ) + ) + assert response is not None + assert response.model == "claude-opus-4-6" + # the real stop_reason is surfaced, not a hardcoded "stop" + assert response.choices[0].finish_reason == "tool_calls" + usage = response.usage + # 80 input + 20 cache_creation (derived from 12+8) + 5 cache_read + assert usage.prompt_tokens == 105 + assert usage.completion_tokens == 40 + assert usage._cache_creation_input_tokens == 20 + assert usage._cache_read_input_tokens == 5 + assert usage.server_tool_use.web_search_requests == 1 + assert usage.server_tool_use.tool_search_requests == 3 + + @pytest.mark.parametrize( + "event_str,expected", + [ + ("data: [DONE]", None), + ("data: ", None), + ("data: {not-json", None), + ("event: ping", None), + ('data: {"a": 1}', {"a": 1}), + ], + ) + def test_extract_sse_data_handles_malformed_and_sentinel_lines( + self, event_str, expected + ): + assert ( + AnthropicPassthroughLoggingHandler._extract_sse_data(event_str) == expected + ) + + def _real_logging_obj(self): + from litellm.litellm_core_utils.litellm_logging import Logging as RealLoggingObj + + logging_obj = RealLoggingObj( + model="claude-3-5-haiku-20241022", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1", + ) + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + return logging_obj + + @patch("litellm.completion_cost") + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response" + ) + def test_handler_falls_back_when_assembly_returns_none( + self, mock_assemble, mock_cost + ): + mock_assemble.return_value = None + mock_cost.return_value = 0.0021 + logging_obj = self._real_logging_obj() + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-haiku-20241022", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=list(self._CHUNKS), + end_time=datetime.now(), + ) + + assert result["result"] is not None + assert result["result"].usage.completion_tokens == 55 + assert result["kwargs"]["response_cost"] == 0.0021 + + @patch("litellm.completion_cost") + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response" + ) + def test_handler_falls_back_when_assembly_raises(self, mock_assemble, mock_cost): + import litellm + + mock_assemble.side_effect = litellm.APIError( + status_code=500, + message="boom", + llm_provider="anthropic", + model="claude-3-5-haiku-20241022", + ) + mock_cost.return_value = 0.0021 + logging_obj = self._real_logging_obj() + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-haiku-20241022", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=list(self._CHUNKS), + end_time=datetime.now(), + ) + + # a raise from stream_chunk_builder must be treated like a None result, + # not propagate out and drop the request from SpendLogs + assert result["result"] is not None + assert result["result"].usage.completion_tokens == 55 + assert result["kwargs"]["response_cost"] == 0.0021 + + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response" + ) + def test_handler_returns_none_when_no_usage_recoverable(self, mock_assemble): + # assembly fails AND the chunks carry no usage event, so there is nothing + # to price; the handler must return None rather than fabricate a response + mock_assemble.return_value = None + logging_obj = self._real_logging_obj() + chunks = [_sse_bytes({"type": "content_block_delta", "delta": {"text": "hi"}})] + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-haiku-20241022", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=chunks, + end_time=datetime.now(), + ) + + assert result["result"] is None + assert result["kwargs"] == {} + + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_usage_only_response_from_chunks" + ) + @patch.object( + AnthropicPassthroughLoggingHandler, "_build_complete_streaming_response" + ) + def test_handler_does_not_crash_when_usage_only_fallback_raises( + self, mock_assemble, mock_fallback + ): + # if the usage-only fallback itself raises, it must be treated as None and + # drop gracefully, not propagate out and crash the success handler + mock_assemble.return_value = None + mock_fallback.side_effect = Exception("fallback boom") + logging_obj = self._real_logging_obj() + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-haiku-20241022", "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=list(self._CHUNKS), + end_time=datetime.now(), + ) + + assert result["result"] is None + assert result["kwargs"] == {} + + +class TestAnthropicResponseCostRecordedOnModelCallDetails: + """The pass-through success path reads spend from + model_call_details["response_cost"], not from kwargs, so the streaming payload + builder must record it there or streaming pass-through logs $0.""" + + def test_create_payload_records_response_cost_on_model_call_details(self): + from litellm.types.utils import Choices, Message, ModelResponse + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.get_router_model_id.return_value = None + logging_obj.litellm_params = {} + logging_obj.litellm_call_id = "test-call-id" + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="hello", role="assistant"), + ) + ], + created=1234567890, + model="claude-3-7-sonnet-20250219", + usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + ) + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-7-sonnet-20250219", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert ( + logging_obj.model_call_details["response_cost"] == kwargs["response_cost"] + ) + assert logging_obj.model_call_details["response_cost"] > 0 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 362f4986c62..379e219b29b 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 @@ -2682,15 +2682,19 @@ async def test_add_litellm_data_to_request_adds_headers_to_metadata(): version="1.0", ) - # Verify headers are added to metadata for guardrails - assert "metadata" in result, "metadata should be present in result" - assert "headers" in result["metadata"], "headers should be present in metadata" + # Verify headers are added to litellm_metadata for guardrails. + # Bedrock passthrough uses litellm_metadata to prevent key-level + # tags from leaking into the provider payload (GH#30629). + assert "litellm_metadata" in result, "litellm_metadata should be present in result" + assert ( + "headers" in result["litellm_metadata"] + ), "headers should be present in litellm_metadata" assert isinstance( - result["metadata"]["headers"], dict + result["litellm_metadata"]["headers"], dict ), "headers should be a dictionary" # Verify specific headers are accessible (important for guardrails) - headers = result["metadata"]["headers"] + headers = result["litellm_metadata"]["headers"] assert ( "user-agent" in headers or "User-Agent" in headers ), "User-Agent header should be accessible in metadata" 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 f73aee77cc1..38990644154 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 @@ -118,3 +118,20 @@ async def test_chunk_processor_does_not_schedule_logging_when_no_chunks(): assert received == [] mock_route.assert_not_called() + + +def test_convert_raw_bytes_survives_truncated_multibyte_sequence(): + """A stream cut mid-multibyte-sequence (client disconnect) must still decode + via errors="replace" so the usage events already received are logged, instead + of raising UnicodeDecodeError and dropping the whole request from SpendLogs.""" + # the 3-byte "☃" (E2 98 83) is cut after 2 bytes, leaving an invalid sequence + # that strict utf-8 decode would raise on, discarding the message_delta line too + truncated_codepoint = "☃".encode("utf-8")[:2] + raw_bytes = [ + b'data: {"text": "' + truncated_codepoint, + b'\ndata: {"type": "message_delta"}\n', + ] + + lines = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes) + + assert any('"type": "message_delta"' in line for line in lines) diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index e89ada5bdef..df14dc5b5dc 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -15,8 +15,6 @@ from __future__ import annotations from unittest.mock import AsyncMock, MagicMock -import pytest - from .conftest import VOLATILE_KEYS, normalize @@ -248,6 +246,193 @@ def test_config_field_info_field_not_in_db(client, auth_as, mock_prisma, monkeyp assert "not in DB" in response.json().get("detail", {}).get("error", "") +def test_config_field_info_redacts_nested_secret_for_view_only_admin( + client, auth_as, mock_prisma, monkeypatch +): + """A view-only admin reading a structured field must not receive nested + credentials. database_args carries aws_web_identity_token (a DynamoDB + role-assumption credential); it must come back redacted while non-secret + siblings like region_name stay visible.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = { + "database_args": { + "region_name": "us-east-1", + "user_table_name": "LiteLLM_UserTable", + "aws_web_identity_token": "sk-super-secret-token", + } + } + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + response = client.get( + "/config/field/info", params={"field_name": "database_args"} + ) + assert response.status_code == 200 + value = response.json()["field_value"] + assert value["aws_web_identity_token"] == "REDACTED" + assert value["region_name"] == "us-east-1" + assert value["user_table_name"] == "LiteLLM_UserTable" + + +def test_config_field_info_full_admin_sees_nested_secret( + client, auth_as, mock_prisma, monkeypatch +): + """The redaction must not over-redact for a full PROXY_ADMIN, who needs + the real nested value to populate the edit form.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = { + "database_args": { + "region_name": "us-east-1", + "aws_web_identity_token": "sk-super-secret-token", + } + } + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get( + "/config/field/info", params={"field_name": "database_args"} + ) + assert response.status_code == 200 + value = response.json()["field_value"] + assert value["aws_web_identity_token"] == "sk-super-secret-token" + assert value["region_name"] == "us-east-1" + + +def test_config_field_info_redacts_top_level_scalar_for_view_only( + client, auth_as, mock_prisma, monkeypatch +): + """The top-level scalar branch must also redact for a view-only admin. + database_url carries DB credentials and is not caught by the name masker, + so it is in the explicit secret set.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = {"database_url": "postgresql://admin:p4ss@db:5432/litellm"} + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + response = client.get( + "/config/field/info", params={"field_name": "database_url"} + ) + assert response.status_code == 200 + assert response.json()["field_value"] == "REDACTED" + + +def test_redact_general_setting_value_recurses_list_of_dicts(): + """The list branch of the recursor redacts secret leaves inside each dict + while non-secret keys survive, and a full admin gets the value untouched.""" + from litellm.proxy import proxy_server as ps + + value = [ + {"path": "/foo", "headers": {"Authorization": "Bearer sk-x"}}, + {"path": "/bar", "client_secret": "sk-y"}, + ] + redacted = ps._redact_general_setting_value( + "some_list_field", value, is_full_admin=False + ) + assert redacted[0]["headers"]["Authorization"] == "REDACTED" + assert redacted[0]["path"] == "/foo" + assert redacted[1]["client_secret"] == "REDACTED" + assert redacted[1]["path"] == "/bar" + assert ( + ps._redact_general_setting_value("some_list_field", value, is_full_admin=True) + == value + ) + + +def test_redact_secret_values_in_obj_fails_closed_at_max_depth(): + """Past _REDACT_SECRET_MAX_DEPTH the whole subtree is replaced with + "REDACTED" rather than returned verbatim, so a secret buried below the cap + can never leak via depth-overrun. A future refactor that flips the cap + branch to fail-open would surface here.""" + from litellm.proxy import proxy_server as ps + + # leaf and wrap keys are both NON-secret so neither the key-name + # short-circuit nor the explicit-secret set catches the leak. The cap is + # the only thing standing between the secret and the response — flip the + # cap to fail-open and the secret comes back verbatim. + nested: object = {"notes": "sk-leak-bottom"} + for _ in range(ps._REDACT_SECRET_MAX_DEPTH + 2): + nested = {"wrap": nested} + + out = ps._redact_general_setting_value( + "some_struct_field", nested, is_full_admin=False + ) + # the secret must not survive anywhere in the returned tree + assert "sk-leak-bottom" not in repr(out) + + # full admin is unaffected by the cap — the value comes back untouched + admin_out = ps._redact_general_setting_value( + "some_struct_field", nested, is_full_admin=True + ) + assert admin_out is nested + + +def test_config_list_redacts_pass_through_secret_for_view_only( + client, auth_as, mock_prisma, monkeypatch +): + """/config/list must not leak pass_through_endpoints upstream credentials + to a view-only admin. pass_through_endpoints is a known secret-bearing + field, so a non-admin gets it redacted; a full admin still sees it.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + row = MagicMock() + row.param_value = {"max_parallel_requests": 3} + table.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr( + ps, + "general_settings", + { + "pass_through_endpoints": [ + { + "path": "/foo", + "target": "https://upstream.example.com", + "headers": {"Authorization": "Bearer sk-UPSTREAM-SECRET"}, + } + ] + }, + ) + + def _pass_through_value(body): + return next( + entry["field_value"] + for entry in body + if entry["field_name"] == "pass_through_endpoints" + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + view_resp = client.get( + "/config/list", params={"config_type": "general_settings"} + ) + assert view_resp.status_code == 200 + assert "sk-UPSTREAM-SECRET" not in view_resp.text + assert _pass_through_value(view_resp.json()) == "REDACTED" + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + admin_resp = client.get( + "/config/list", params={"config_type": "general_settings"} + ) + assert admin_resp.status_code == 200 + admin_value = _pass_through_value(admin_resp.json()) + assert admin_value[0]["headers"]["Authorization"] == "Bearer sk-UPSTREAM-SECRET" + + # --------------------------------------------------------------------------- # GET /config/list # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 017f4bd4368..3bf14c08d14 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -9,7 +9,7 @@ Pins (PR2): from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import MagicMock import pytest @@ -128,6 +128,104 @@ def test_v1_model_info_no_model_list_error(client, auth_as, null_router, path): assert "LLM Model List not loaded" in response.text +# --------------------------------------------------------------------------- +# GET /model/info — team BYOK scoping (issue #30983) +# --------------------------------------------------------------------------- + +_BYOK_TEAM_ID = "team-abc" +_BYOK_PUBLIC_NAME = "my-byok-gpt-4" +_BYOK_INTERNAL_NAME = f"model_name_{_BYOK_TEAM_ID}_0123456789abcdef" + + +@pytest.fixture +def byok_team_router(monkeypatch): + """Router holding one team-scoped BYOK deployment for team `team-abc`. + + Mirrors how a team's own-key BYOK model lives in the router: the routing + key is an internal mangled name while the public name lives in + `model_info.team_public_model_name`. + """ + byok_deployment = { + "model_name": _BYOK_INTERNAL_NAME, + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": { + "id": "byok-deployment-id", + "db_model": True, + "team_id": _BYOK_TEAM_ID, + "team_public_model_name": _BYOK_PUBLIC_NAME, + }, + } + + router = MagicMock() + router.model_list = [byok_deployment] + router.get_model_list_from_model_alias = MagicMock(return_value=[]) + router.get_model_names = MagicMock(return_value=[]) + router.get_model_access_groups = MagicMock(return_value={}) + router.get_model_ids = MagicMock(return_value=[]) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", [byok_deployment]) + monkeypatch.setattr(proxy_server, "user_model", None) + yield router + + +@pytest.mark.parametrize("path", ["/v1/model/info", "/model/info"]) +def test_model_info_team_key_sees_own_byok_model(client, auth_as, byok_team_router, mock_prisma, monkeypatch, path): + """Regression for #30983: a team key (user_id=None) must see its own + team's BYOK model under the public name. + + Before the fix `_get_caller_byok_team_scope` keyed only off the bound + user's team memberships, returned an empty set for a team key, and the + BYOK row was dropped -> `{"data": []}`. + """ + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + mock_prisma.db.litellm_usertable.find_unique.return_value = None + + with auth_as( + role=LitellmUserRoles.INTERNAL_USER, + user_id=None, + team_id=_BYOK_TEAM_ID, + team_models=[_BYOK_PUBLIC_NAME], + ): + response = client.get(path) + + assert response.status_code == 200 + data = response.json()["data"] + surfaced_names = [m.get("model_name") for m in data] + assert _BYOK_PUBLIC_NAME in surfaced_names + assert _BYOK_INTERNAL_NAME not in surfaced_names + + +@pytest.mark.parametrize("path", ["/v1/model/info", "/model/info"]) +def test_model_info_team_key_cannot_see_other_teams_byok_model( + client, auth_as, byok_team_router, mock_prisma, monkeypatch, path +): + """A team key for a different team must NOT see team-abc's BYOK row. + + Guards the fix from over-broadening into a cross-team metadata leak. + """ + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + mock_prisma.db.litellm_usertable.find_unique.return_value = None + + with auth_as( + role=LitellmUserRoles.INTERNAL_USER, + user_id=None, + team_id="other-team", + team_models=[_BYOK_PUBLIC_NAME], + ): + response = client.get(path) + + assert response.status_code == 200 + data = response.json()["data"] + surfaced_names = [m.get("model_name") for m in data] + assert _BYOK_PUBLIC_NAME not in surfaced_names + assert _BYOK_INTERNAL_NAME not in surfaced_names + + # --------------------------------------------------------------------------- # GET /model_group/info # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 33de1ede917..699606b5277 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -16,8 +16,6 @@ Pins covered: from __future__ import annotations import json -from typing import Any, AsyncIterator -from unittest.mock import AsyncMock, MagicMock import pytest @@ -26,8 +24,11 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.proxy_server import ( _apply_streaming_chunk_hooks, _fast_serialize_simple_model_response_stream, + _format_fallback_metadata_sse_event, _format_streaming_sse_chunk, _get_client_requested_model_for_streaming, + _get_streaming_fallback_metadata, + _is_positive_int_like, _restamp_streaming_chunk_model, _serialize_streaming_chunk, async_assistants_data_generator, @@ -71,6 +72,15 @@ async def _async_iter_raises(exc: Exception): raise exc +class _FakeStream: + def __init__(self, chunks, hidden_params=None): + self._chunks = chunks + self._hidden_params = hidden_params or {} + + def __aiter__(self): + return _async_iter(self._chunks) + + # --------------------------------------------------------------------------- # data_generator # --------------------------------------------------------------------------- @@ -274,6 +284,34 @@ def test_restamp_streaming_chunk_model_overrides_model_on_dict(): assert logged is True +def test_restamp_streaming_chunk_model_uses_fallback_model_from_metadata(): + chunk = _simple_chunk(model="openai/internal-fallback") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="primary-model", + request_data={"litellm_call_id": "id-1"}, + model_mismatch_logged=False, + fallback_was_attempted=True, + fallback_model_from_metadata="fallback-model", + ) + assert new_chunk.model == "fallback-model" + assert logged is True + + +def test_restamp_streaming_chunk_model_preserves_fallback_model_without_group(): + chunk = _simple_chunk(model="openai/internal-fallback") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="primary-model", + request_data={}, + model_mismatch_logged=False, + fallback_was_attempted=True, + fallback_model_from_metadata=None, + ) + assert new_chunk.model == "openai/internal-fallback" + assert logged is False + + def test_restamp_streaming_chunk_model_invalid_chunk_type_unchanged(): """For a non-BaseModel, non-dict chunk the helper returns it as-is along with the original ``model_mismatch_logged`` flag.""" @@ -288,6 +326,147 @@ def test_restamp_streaming_chunk_model_invalid_chunk_type_unchanged(): assert logged is False +def test_is_positive_int_like_invalid_and_edge_values(): + assert _is_positive_int_like(None) is False + assert _is_positive_int_like("not-a-number") is False + assert _is_positive_int_like(0) is False + assert _is_positive_int_like(-1) is False + assert _is_positive_int_like("1") is True + assert _is_positive_int_like(2) is True + + +def test_get_streaming_fallback_metadata_reads_headers(): + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + stream = _FakeStream( + [], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": "1", + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + } + }, + ) + assert _get_streaming_fallback_metadata(stream) == ( + True, + "fallback-model", + fallback_errors, + ) + + +def test_get_streaming_fallback_metadata_no_additional_headers(): + stream = _FakeStream([], hidden_params={}) + assert _get_streaming_fallback_metadata(stream) == (False, None, []) + + +def test_get_streaming_fallback_metadata_zero_fallback_count(): + stream = _FakeStream( + [], + hidden_params={ + "additional_headers": {"x-litellm-attempted-fallbacks": 0} + }, + ) + assert _get_streaming_fallback_metadata(stream) == (False, None, []) + + +def test_get_streaming_fallback_metadata_no_model_group_returns_none_model(): + stream = _FakeStream( + [], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + } + }, + ) + was_attempted, fallback_model, errors = _get_streaming_fallback_metadata(stream) + assert was_attempted is True + assert fallback_model is None + assert errors == [] + + +def test_restamp_streaming_chunk_model_azure_router_preserves_model(): + chunk = _simple_chunk(model="azure_ai/internal-deployment") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="azure_ai/model-router", + request_data={}, + model_mismatch_logged=False, + ) + assert new_chunk.model == "azure_ai/internal-deployment" + assert logged is False + + +def test_restamp_streaming_chunk_model_fastest_response_preserves_model(): + chunk = _simple_chunk(model="winning-model") + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="gpt-4,claude-3", + request_data={"fastest_response": True}, + model_mismatch_logged=False, + ) + assert new_chunk.model == "winning-model" + assert logged is False + + +def test_restamp_streaming_chunk_model_setattr_exception_logs_and_returns(): + from pydantic import ConfigDict + + class FrozenChunk(_simple_chunk().__class__): + model_config = ConfigDict(frozen=True) + + chunk = FrozenChunk( + id="chatcmpl-test", + choices=[], + created=0, + model="openai/internal-x", + object="chat.completion.chunk", + ) + new_chunk, logged = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="gpt-4", + request_data={"litellm_call_id": "test-id"}, + model_mismatch_logged=False, + ) + assert new_chunk.model == "openai/internal-x" + assert logged is True + + +def test_format_fallback_metadata_sse_event(): + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + + event = _format_fallback_metadata_sse_event( + fallback_model="fallback-model", + fallback_errors=fallback_errors, + ) + + assert isinstance(event, str) + assert event.startswith("data: ") + payload = json.loads(event.removeprefix("data: ").removesuffix("\n\n")) + assert payload["choices"] == [] + assert payload["litellm_fallback"] == { + "fallback_model": "fallback-model", + "errors": fallback_errors, + } + assert payload["id"] == "litellm-fallback-metadata" + assert payload["object"] == "chat.completion.chunk" + assert payload["model"] == "fallback-model" + assert isinstance(payload["created"], int) + + # --------------------------------------------------------------------------- # _fast_serialize_simple_model_response_stream # --------------------------------------------------------------------------- @@ -473,7 +652,7 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): # First chunk is bytes (fast path) wrapped via _format_streaming_sse_chunk. first = out[0] assert isinstance(first, bytes) - payload = json.loads(first.removeprefix(b"data: ").rstrip(b"\n\n")) + payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) assert normalize(payload) == { "id": "", "object": "chat.completion.chunk", @@ -488,6 +667,172 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): } +@pytest.mark.asyncio +async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch): + _patch_logging_flags(monkeypatch) + + response = _FakeStream( + [_simple_chunk(model="openai/internal-fallback", content="hello")], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + } + }, + ) + out = [] + async for line in async_data_generator( + response=response, + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model", "include_fallback_errors": True}, + ): + out.append(line) + + first = out[0] + assert isinstance(first, bytes) + payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + assert payload["model"] == "fallback-model" + + +@pytest.mark.asyncio +async def test_async_data_generator_uses_chunk_fallback_metadata(monkeypatch): + _patch_logging_flags(monkeypatch) + + chunk = _simple_chunk(model="openai/internal-fallback", content="hello") + chunk._hidden_params = { + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + } + } + out = [] + async for line in async_data_generator( + response=_async_iter([chunk]), + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model"}, + ): + out.append(line) + + first = out[0] + assert isinstance(first, bytes) + payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + assert payload["model"] == "fallback-model" + + +@pytest.mark.asyncio +async def test_async_data_generator_switches_model_mid_stream_on_fallback(monkeypatch): + """Pre-fallback chunks keep the client-requested model; once a chunk carries + fallback metadata the model latches to the fallback group for the rest of the + stream. This pins the client-visible mid-stream model change.""" + _patch_logging_flags(monkeypatch) + + primary_chunk = _simple_chunk(model="openai/internal-primary", content="hi") + fallback_chunk = _simple_chunk(model="openai/internal-fallback", content="there") + fallback_chunk._hidden_params = { + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + } + } + out = [] + async for line in async_data_generator( + response=_async_iter([primary_chunk, fallback_chunk]), + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model"}, + ): + out.append(line) + + first_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + second_payload = json.loads(out[1].removeprefix(b"data: ").removesuffix(b"\n\n")) + assert first_payload["model"] == "primary-model" + assert second_payload["model"] == "fallback-model" + + +@pytest.mark.asyncio +async def test_async_data_generator_emits_fallback_error_metadata_event(monkeypatch): + _patch_logging_flags(monkeypatch) + monkeypatch.setitem(ps.general_settings, "expose_fallback_errors_to_caller", True) + + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + response = _FakeStream( + [_simple_chunk(model="openai/internal-fallback", content="hello")], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + } + }, + ) + out = [] + async for line in async_data_generator( + response=response, + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model", "include_fallback_errors": True}, + ): + out.append(line) + + assert isinstance(out[0], bytes) + chunk_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + assert chunk_payload["model"] == "fallback-model" + assert isinstance(out[1], str) + assert out[1].startswith("data: ") + metadata_payload = json.loads(out[1].removeprefix("data: ").removesuffix("\n\n")) + assert metadata_payload["choices"] == [] + assert metadata_payload["litellm_fallback"] == { + "fallback_model": "fallback-model", + "errors": fallback_errors, + } + assert metadata_payload["id"] == "litellm-fallback-metadata" + assert metadata_payload["object"] == "chat.completion.chunk" + assert metadata_payload["model"] == "fallback-model" + assert isinstance(metadata_payload["created"], int) + + +@pytest.mark.asyncio +async def test_async_data_generator_skips_fallback_error_event_without_opt_in( + monkeypatch, +): + _patch_logging_flags(monkeypatch) + + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + response = _FakeStream( + [_simple_chunk(model="openai/internal-fallback", content="hello")], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + } + }, + ) + out = [] + async for line in async_data_generator( + response=response, + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model"}, + ): + out.append(line) + + assert isinstance(out[0], bytes) + payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + assert payload["model"] == "fallback-model" + + @pytest.mark.asyncio async def test_async_data_generator_mid_stream_exception_yields_error_payload( monkeypatch, diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 40d590132aa..34405f20727 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -316,8 +316,10 @@ async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch) @pytest.mark.asyncio async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): - """A key without a resolvable user (e.g. CI/service token) sees only - global deployments, never any team-scoped BYOK rows.""" + """A key with no resolvable user and no team (e.g. a CI/service token + created outside any team) sees only global deployments, never team-scoped + BYOK rows. A team-scoped key does see its own team's rows (issue #30983), + pinned by the /model/info route tests.""" team_row = _team_row() other_team_row = _other_team_row() global_row = { @@ -343,7 +345,7 @@ async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): caller = UserAPIKeyAuth( user_id=None, user_role=LitellmUserRoles.INTERNAL_USER, - team_id="team-abc-123", + team_id=None, models=[], team_models=[], ) @@ -352,6 +354,106 @@ async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): assert [m["model_info"]["id"] for m in resp["data"]] == ["global-id-1"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "find_unique", + [ + AsyncMock(return_value=MagicMock(teams=[])), + AsyncMock(return_value=None), + AsyncMock(side_effect=RuntimeError("db down")), + ], + ids=["user-not-in-team", "user-row-missing", "user-lookup-error"], +) +async def test_model_info_v1_team_key_sees_own_byok_regardless_of_user_lookup( + monkeypatch, find_unique +): + """A team-scoped key sees its own team's BYOK rows even when the bound user + is not a member of that team, has no DB row, or the lookup errors; the + key's team_id is authoritative (issue #30983). Other teams' rows stay + hidden.""" + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [_team_row(), _other_team_row(), global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = find_unique + + async def _populate(**kwargs): + return kwargs["all_models"] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _populate) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-abc-123", + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_info"]["id"] for m in resp["data"]] == ["byok-id-1", "global-id-1"] + + +@pytest.mark.asyncio +async def test_model_info_v1_user_team_membership_grants_byok(monkeypatch): + """A user's own team memberships still grant that team's BYOK rows, unioned + with any team the key itself is scoped to.""" + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [_team_row(), _other_team_row(), global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=MagicMock(teams=["team-other"]) + ) + + async def _populate(**kwargs): + return kwargs["all_models"] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _populate) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-2", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id=None, + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_info"]["id"] for m in resp["data"]] == [ + "byok-id-other", + "global-id-1", + ] + + @pytest.mark.asyncio async def test_model_info_v1_populates_access_via_team_ids(monkeypatch): """`/v1/model/info` must populate access_via_team_ids when the DB is connected.""" diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 48d0b1deadd..d920488c352 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -201,6 +201,48 @@ def test_anthropic_provider_fields_support_byok(): ), "api_base must appear before api_key in credential_fields (matches AI21 and ANTHROPIC_TEXT convention)." +def test_bedrock_mantle_provider_fields(): + """Amazon Bedrock Mantle must be a selectable provider in the Add Model flow. + + The dropdown is driven entirely by /public/providers/fields, so a missing + entry means Mantle cannot be added through the UI at all (regression guard + for LIT-3885). The credential fields must match what the backend actually + honors: an optional bearer api_key (BYOK), the AWS SigV4 chain, a region, + and an api_base override. + """ + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/providers/fields") + assert response.status_code == 200 + providers = response.json() + + mantle = next((p for p in providers if p["provider"] == "BedrockMantle"), None) + assert mantle is not None, "Bedrock Mantle provider entry not found" + + # provider must equal the UI provider_map key so the model dropdown resolves + # bedrock_mantle models; litellm_provider must be the backend slug. + assert mantle["provider_display_name"] == "Amazon Bedrock Mantle" + assert mantle["litellm_provider"] == "bedrock_mantle" + assert mantle["default_model_placeholder"].startswith("bedrock_mantle/") + + fields_by_key = {f["key"]: f for f in mantle["credential_fields"]} + + # Bearer-token auth is BYOK: optional and masked. + assert "api_key" in fields_by_key + assert fields_by_key["api_key"]["required"] is False + assert fields_by_key["api_key"]["field_type"] == "password" + + # AWS SigV4 fallback credentials. + assert fields_by_key["aws_access_key_id"]["field_type"] == "password" + assert fields_by_key["aws_secret_access_key"]["field_type"] == "password" + assert "aws_region_name" in fields_by_key + + # api_base override so admins can target a custom Mantle host without env access. + assert fields_by_key["api_base"]["field_type"] == "text" + + def test_google_ai_studio_provider_fields_expose_api_base(): """The Google AI Studio (gemini) credential form must let admins set a custom api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted @@ -819,3 +861,44 @@ def test_public_mcp_hub_returns_empty_when_whitelist_unset(): assert response.status_code == 200 assert response.json() == [] app.dependency_overrides.clear() + + +def test_public_mcp_hub_does_not_expose_upstream_url(): + """Regression: /public/mcp_hub is unauthenticated, so the gateway-internal + upstream url must never appear in its response even when the server has one.""" + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport + + app = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: MagicMock() + client = TestClient(app) + + secret_url = "https://internal-only.example.com/mcp" + server = MCPServer( + server_id="listed", + name="listed", + server_name="listed", + url=secret_url, + transport=MCPTransport.http, + available_on_public_internet=True, + ) + + mock_manager = MagicMock() + mock_manager.get_public_mcp_servers.return_value = [server] + + with ( + patch("litellm.public_mcp_servers", ["listed"]), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + ): + response = client.get("/public/mcp_hub") + + assert response.status_code == 200 + data = response.json() + assert [item["server_id"] for item in data] == ["listed"] + assert all("url" not in item for item in data) + assert secret_url not in response.text + app.dependency_overrides.clear() diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 6b692180559..74f681bb97e 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -96,6 +96,20 @@ class TestGetMetadataVariableName: request = self._make_request("/v1/embeddings") assert _get_metadata_variable_name(request) == "metadata" + def test_returns_litellm_metadata_for_bedrock_invoke(self): + # GH#30629: bedrock passthrough must use litellm_metadata + # to prevent key-level tags from leaking into provider body + request = self._make_request( + "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke" + ) + assert _get_metadata_variable_name(request) == "litellm_metadata" + + def test_returns_litellm_metadata_for_bedrock_converse(self): + request = self._make_request( + "/bedrock/model/us.anthropic.claude-sonnet-4-6/converse" + ) + assert _get_metadata_variable_name(request) == "litellm_metadata" + def test_get_enforced_params_for_service_account_settings(): """ @@ -718,7 +732,18 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): @pytest.mark.asyncio @pytest.mark.parametrize( "control_field", - ["callbacks", "service_callback", "logger_fn", "litellm_disabled_callbacks"], + [ + "callbacks", + "service_callback", + "logger_fn", + "litellm_disabled_callbacks", + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_code_interpreter_interception_active", + "_code_interpreter_interception_converted_stream", + "_code_interpreter_interception_sandbox_key", + "max_agentic_loops", + ], ) async def test_add_litellm_data_to_request_strips_callback_control_fields( control_field, @@ -741,12 +766,19 @@ async def test_add_litellm_data_to_request_strips_callback_control_fields( request_mock.client = MagicMock() request_mock.client.host = "127.0.0.1" - sample_value = ( - ["langfuse"] - if control_field - in ("callbacks", "service_callback", "litellm_disabled_callbacks") - else "module.func" - ) + sample_values = { + "callbacks": ["langfuse"], + "service_callback": ["langfuse"], + "litellm_disabled_callbacks": ["langfuse"], + "logger_fn": "module.func", + "_agentic_loop_depth": 5, + "_agentic_loop_fingerprints": ["forged"], + "_code_interpreter_interception_active": True, + "_code_interpreter_interception_converted_stream": True, + "_code_interpreter_interception_sandbox_key": "forged-key", + "max_agentic_loops": 9999, + } + sample_value = sample_values[control_field] updated = await add_litellm_data_to_request( data={ @@ -4150,7 +4182,9 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:paid": return 0.50 return fallback_spend @@ -4207,7 +4241,9 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:tenant:acme": return 0.50 return fallback_spend @@ -4234,6 +4270,99 @@ class TestApplyClientTagPolicyPreAuth: assert exc_info.value.current_cost == 0.50 assert exc_info.value.max_budget == 0.10 + @pytest.mark.asyncio + @pytest.mark.parametrize( + "route", + [ + "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + "/v1/messages", + ], + ) + async def test_header_tags_visible_to_tag_max_budget_check_on_metadata_route( + self, route + ): + """Regression: on LITELLM_METADATA_ROUTES (bedrock, /v1/messages, ...), + common_checks pre-seeds ``litellm_metadata`` and writes key tags there + before ``_tag_max_budget_check`` reads from the same key. The auth wrapper + calls ``apply_client_tag_policy_pre_auth`` first, so without an earlier + pre-seed header tags land in ``metadata`` and the budget check (now + resolving to ``litellm_metadata``) silently ignores them. This test mirrors + the actual auth-time call order and verifies that an over-budget + header-supplied tag still trips ``_tag_max_budget_check``. + """ + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TagTable + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.utils import ProxyLogging + + request_mock = _build_request_mock_with_headers( + {"x-litellm-tags": "tenant:acme"} + ) + data = {"model": "us.anthropic.claude-sonnet-4-6"} + valid_token = UserAPIKeyAuth( + token="test-token", + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route( + request_data=data, + route=route, + ) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=valid_token, + ) + + tag_object = LiteLLM_TagTable( + tag_name="tenant:acme", + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), + ) + + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): + if counter_key == "spend:tag:tenant:acme": + return 0.50 + return fallback_spend + + with ( + patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ), + patch( + "litellm.proxy.proxy_server.get_current_spend", + mock_get_current_spend, + ), + patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={"tenant:acme": tag_object}, + ), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body=data, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=valid_token, + request=request_mock, + ) + assert exc_info.value.current_cost == 0.50 + assert exc_info.value.max_budget == 0.10 + + assert "metadata" not in data + assert data["litellm_metadata"]["tags"] == ["tenant:acme"] + class TestApplyKeyTagsPreAuth: def test_merges_key_tags_into_metadata(self): @@ -4362,7 +4491,9 @@ class TestApplyKeyTagsPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:engineering": return 0.50 return fallback_spend @@ -4413,7 +4544,9 @@ class TestApplyKeyTagsPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:engineering": return 0.05 return fallback_spend diff --git a/tests/test_litellm/proxy/test_plugin_routes.py b/tests/test_litellm/proxy/test_plugin_routes.py new file mode 100644 index 00000000000..52999447179 --- /dev/null +++ b/tests/test_litellm/proxy/test_plugin_routes.py @@ -0,0 +1,232 @@ +"""Regression tests for UI-registered embed plugins. + +Covers three bugs: +1. `general_settings.plugins` was not a field on ConfigGeneralSettings, so the + admin UI's POST /config/field/update with field_name="plugins" was rejected + with "Invalid field=plugins passed in." +2. The in-memory plugin registry only refreshed at startup, so a plugin added + via the UI did not appear in /api/plugins until a restart. +3. Plugins persisted to DB general_settings were not loaded on startup (the + registry only initialised from the YAML config), so UI-added plugins vanished + after a restart. +""" + +import asyncio +from unittest.mock import MagicMock + +from litellm.proxy._types import ( + ConfigGeneralSettings, + LitellmUserRoles, + PluginConfig, + UserAPIKeyAuth, +) +from litellm.proxy.plugin_routes import list_plugins, register_plugins_from_config + + +def _admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _non_admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER) + + +def test_plugins_is_a_valid_general_setting() -> None: + """The config-update endpoint gates on this exact membership check.""" + assert "plugins" in ConfigGeneralSettings.model_fields + + +def test_config_general_settings_parses_plugin_list() -> None: + """A list of plugin dicts (what the UI sends) coerces into PluginConfig.""" + settings = ConfigGeneralSettings.model_validate( + { + "plugins": [ + { + "name": "chat-ui", + "display_name": "Chat UI", + "url": "http://localhost:3300", + }, + { + "name": "agent-builder", + "url": "http://127.0.0.1:4010", + "plugin_key": "sk-secret", + }, + ] + } + ) + plugins = settings.plugins + assert plugins is not None + assert [p.name for p in plugins] == ["chat-ui", "agent-builder"] + assert isinstance(plugins[0], PluginConfig) + assert plugins[1].display_name is None + assert plugins[1].plugin_key == "sk-secret" + + +def test_registered_plugins_appear_in_list_without_restart() -> None: + """register_plugins_from_config makes UI-added plugins visible immediately, + and replaces (not merges) so removed plugins disappear.""" + register_plugins_from_config( + { + "plugins": [ + { + "name": "chat-ui", + "display_name": "Chat UI", + "url": "http://localhost:3300", + } + ] + } + ) + names = [p["name"] for p in asyncio.run(list_plugins(user_api_key_dict=_admin()))] + assert names == ["chat-ui"] + + register_plugins_from_config( + { + "plugins": [ + { + "name": "chat-ui", + "display_name": "Chat UI", + "url": "http://localhost:3300", + }, + { + "name": "agent-builder", + "display_name": "Agent Builder", + "url": "http://127.0.0.1:4010", + }, + ] + } + ) + names = sorted( + p["name"] for p in asyncio.run(list_plugins(user_api_key_dict=_admin())) + ) + assert names == ["agent-builder", "chat-ui"] + + # Removing a plugin from config drops it from the live list. + register_plugins_from_config({}) + assert asyncio.run(list_plugins(user_api_key_dict=_admin())) == [] + + +def test_plugin_key_is_never_returned_to_the_browser() -> None: + """plugin_key is a credential the UI never needs; /api/plugins must omit it + for every caller, admin included, so it never lands in browser state.""" + register_plugins_from_config( + { + "plugins": [ + { + "name": "p", + "display_name": "P", + "url": "http://localhost:9", + "plugin_key": "sk-secret", + } + ] + } + ) + + admin_entry = asyncio.run(list_plugins(user_api_key_dict=_admin()))[0] + user_entry = asyncio.run(list_plugins(user_api_key_dict=_non_admin()))[0] + + assert "plugin_key" not in admin_entry + assert "plugin_key" not in user_entry + assert admin_entry["url"] == "http://localhost:9" + + register_plugins_from_config({}) + + +def test_db_persisted_plugins_load_on_startup() -> None: + """Plugins saved to DB general_settings must register when the DB config is + merged at startup, not just when present in the YAML file.""" + from litellm.proxy.proxy_server import ProxyConfig + + register_plugins_from_config({}) # start empty (as if YAML had no plugins) + + ProxyConfig()._add_general_settings_from_db_config( + config_data={ + "general_settings": { + "plugins": [ + { + "name": "db-plugin", + "display_name": "DB Plugin", + "url": "http://localhost:5000", + } + ] + } + }, + general_settings={}, + proxy_logging_obj=MagicMock(), + ) + + names = [p["name"] for p in asyncio.run(list_plugins(user_api_key_dict=_admin()))] + assert names == ["db-plugin"] + + register_plugins_from_config({}) + + +def test_safe_response_headers_sandbox_and_strips_wire_headers() -> None: + """Proxied plugin responses must be inert and shed wire/cookie headers.""" + from litellm.proxy.plugin_routes import _safe_response_headers + + out = _safe_response_headers( + { + "content-type": "text/html", + "content-encoding": "gzip", + "content-length": "123", + "set-cookie": "session=abc", + "content-security-policy": "default-src *", + } + ) + + assert out["content-security-policy"] == "sandbox" + assert out["x-content-type-options"] == "nosniff" + assert out["content-type"] == "text/html" + for stripped in ("content-encoding", "content-length", "set-cookie"): + assert stripped not in out + + +def test_litellm_credential_header_names_covers_every_auth_header() -> None: + """The canonical strip set must list every header user_api_key_auth accepts + as a litellm key, so a new auth header can't silently start leaking.""" + from litellm.proxy._types import SpecialHeaders + + assert SpecialHeaders.litellm_credential_header_names() == { + "authorization", + "api-key", + "x-api-key", + "x-goog-api-key", + "ocp-apim-subscription-key", + "x-litellm-api-key", + } + + +def test_every_litellm_auth_header_is_stripped_before_forwarding() -> None: + """A plugin must never receive any header that authenticates against litellm, + only the hop-by-hop set and benign headers are forwarded.""" + from litellm.proxy.plugin_routes import _request_strip_headers + + strip = _request_strip_headers() + incoming = { + "Authorization": "Bearer sk-litellm", + "API-Key": "sk-litellm", + "X-Api-Key": "sk-litellm", + "X-Goog-Api-Key": "sk-litellm", + "Ocp-Apim-Subscription-Key": "sk-litellm", + "X-Litellm-Api-Key": "sk-litellm", + "Cookie": "litellm_session=abc", + "Accept": "application/json", + "X-Trace-Id": "t-1", + } + forwarded = {k: v for k, v in incoming.items() if k.lower() not in strip} + + assert forwarded == {"Accept": "application/json", "X-Trace-Id": "t-1"} + + +def test_configured_custom_key_header_is_stripped() -> None: + """A custom general_settings.litellm_key_header_name must also be stripped, + read live so config changes are honoured without a restart.""" + from litellm.proxy import proxy_server + from litellm.proxy.plugin_routes import _request_strip_headers + + original = getattr(proxy_server, "general_settings", None) + proxy_server.general_settings = {"litellm_key_header_name": "X-My-Tenant-Key"} + try: + assert "x-my-tenant-key" in _request_strip_headers() + finally: + proxy_server.general_settings = original diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 56627c5be88..88dbec4020f 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1708,6 +1708,57 @@ class TestRunServerDbSetup: use_migrate=True, use_v2_resolver=False ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") + def test_startup_exits_on_non_postgres_database_url( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + ): + """A sqlite DATABASE_URL must exit immediately, before any prisma call, + instead of stalling on a migration against the postgresql-only schema.""" + from litellm.proxy.proxy_cli import run_server + + mock_subprocess_run.return_value = MagicMock(returncode=0) + mock_should_update_schema.return_value = True + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env["DATABASE_URL"] = "sqlite:///data/litellm.db" + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + ): + with pytest.raises(SystemExit) as exc_info: + run_server.main( + ["--local", "--skip_server_startup"], standalone_mode=False + ) + assert exc_info.value.code == 1 + mock_setup_database.assert_not_called() + # --- Module-level helpers for worker startup hook tests --- diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6017b9555e9..a203fcc7ec0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1348,6 +1348,7 @@ async def test_apply_search_filter_scopes_byok_to_caller_teams(): non_admin = MagicMock(spec=UserAPIKeyAuth) non_admin.user_role = LitellmUserRoles.INTERNAL_USER non_admin.user_id = "user-mine" + non_admin.team_id = None filtered, total_count = await _apply_search_filter_to_models( all_models=[caller_team_byok, other_team_byok, public_model], @@ -1381,6 +1382,7 @@ async def test_apply_search_filter_scopes_byok_to_caller_teams(): admin = MagicMock(spec=UserAPIKeyAuth) admin.user_role = LitellmUserRoles.PROXY_ADMIN admin.user_id = "admin-1" + admin.team_id = None filtered_admin, _ = await _apply_search_filter_to_models( all_models=[caller_team_byok, other_team_byok, public_model], @@ -8326,3 +8328,125 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): assert fields["cancel_on_disconnect"]["field_type"] == "Boolean" finally: app.dependency_overrides.clear() + + +def test_preserve_redacted_plugin_keys_keeps_stored_credential(): + """A redacted or blank plugin_key on update must not overwrite the real key.""" + from litellm.proxy.proxy_server import _preserve_redacted_plugin_keys + + existing = [{"name": "p1", "url": "https://p1", "plugin_key": "sk-real-1"}] + + redacted = _preserve_redacted_plugin_keys( + [{"name": "p1", "url": "https://p1-new", "plugin_key": "***"}], existing + ) + assert redacted == [ + {"name": "p1", "url": "https://p1-new", "plugin_key": "sk-real-1"} + ] + + blanked = _preserve_redacted_plugin_keys( + [{"name": "p1", "url": "https://p1", "plugin_key": ""}], existing + ) + assert blanked[0]["plugin_key"] == "sk-real-1" + + +def test_preserve_redacted_plugin_keys_sets_new_and_drops_orphan_placeholder(): + """A real new key replaces; a placeholder with no stored key is dropped, never persisted.""" + from litellm.proxy.proxy_server import _preserve_redacted_plugin_keys + + existing = [{"name": "p1", "url": "https://p1", "plugin_key": "sk-real-1"}] + + rotated = _preserve_redacted_plugin_keys( + [{"name": "p1", "url": "https://p1", "plugin_key": "sk-new"}], existing + ) + assert rotated[0]["plugin_key"] == "sk-new" + + new_plugin = _preserve_redacted_plugin_keys( + [{"name": "p2", "url": "https://p2", "plugin_key": "***"}], existing + ) + assert "plugin_key" not in new_plugin[0] + + +def _config_field_info_client(monkeypatch, user_role): + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + db_record = types.SimpleNamespace( + param_value={ + "master_key": "sk-super-secret-master", + "database_url": "postgresql://user:p4ssw0rd@db:5432/litellm", + "pass_through_endpoints": [ + { + "path": "/upstream", + "target": "https://upstream.example.com", + "headers": {"Authorization": "Bearer sk-upstream-secret"}, + } + ], + "max_parallel_requests": 100, + } + ) + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=db_record) + mock_prisma = MagicMock() + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="u", user_role=user_role + ) + return TestClient(app) + + +def test_config_field_info_redacts_secrets_for_view_only_admin(monkeypatch): + """/config/field/info gates on _user_has_admin_view, which also grants + PROXY_ADMIN_VIEW_ONLY. A view-only admin reading master_key/database_url verbatim is + effectively a full admin. Secret-bearing fields must come back REDACTED for anyone who + is not a FULL PROXY_ADMIN, while non-secret fields stay readable.""" + from litellm.proxy._types import LitellmUserRoles + + client = _config_field_info_client( + monkeypatch, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ) + try: + for secret_field in ("master_key", "database_url", "pass_through_endpoints"): + resp = client.get("/config/field/info", params={"field_name": secret_field}) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["field_value"] == "REDACTED" + assert "secret" not in str(body["field_value"]) + assert "p4ssw0rd" not in str(body["field_value"]) + + resp = client.get( + "/config/field/info", params={"field_name": "max_parallel_requests"} + ) + assert resp.status_code == 200, resp.text + assert resp.json()["field_value"] == 100 + finally: + app.dependency_overrides.clear() + + +def test_config_field_info_returns_raw_secrets_for_full_admin(monkeypatch): + """the redaction must not over-apply. A FULL PROXY_ADMIN still + needs the real master_key value to populate the admin edit form.""" + from litellm.proxy._types import LitellmUserRoles + + client = _config_field_info_client(monkeypatch, LitellmUserRoles.PROXY_ADMIN) + try: + resp = client.get("/config/field/info", params={"field_name": "master_key"}) + assert resp.status_code == 200, resp.text + assert resp.json()["field_value"] == "sk-super-secret-master" + + resp = client.get( + "/config/field/info", params={"field_name": "pass_through_endpoints"} + ) + assert resp.status_code == 200, resp.text + assert ( + resp.json()["field_value"][0]["headers"]["Authorization"] + == "Bearer sk-upstream-secret" + ) + finally: + app.dependency_overrides.clear() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index d517c1c346f..87063bdf00b 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -105,6 +105,33 @@ def test_jsonify_team_object_converts_members_to_json_string( } +def test_jsonify_team_object_converts_budget_limits_to_json_string( + prisma_client: PrismaClient, +) -> None: + data = { + "team_id": "t1", + "budget_limits": [ + { + "budget_duration": "1d", + "max_budget": 10.0, + "reset_at": "2026-01-01T00:00:00Z", + }, + { + "budget_duration": "7d", + "max_budget": 50.0, + "reset_at": "2026-01-07T00:00:00Z", + }, + ], + "models": ["gpt-4"], + } + result = prisma_client.jsonify_team_object(data) + assert result == { + "team_id": "t1", + "budget_limits": json.dumps(data["budget_limits"]), + "models": ["gpt-4"], + } + + def test_jsonify_team_object_error_on_non_dict(prisma_client: PrismaClient) -> None: with pytest.raises(AttributeError): prisma_client.jsonify_team_object(None) # type: ignore[arg-type] diff --git a/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py new file mode 100644 index 00000000000..2aee0f0a4ef --- /dev/null +++ b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py @@ -0,0 +1,142 @@ +import json + +from pydantic import BaseModel + +from litellm.router_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, + add_retry_headers_to_response, + get_fallback_errors_from_headers, + get_hidden_params_dict, +) + + +class StreamingWrapper: + def __init__(self): + self._hidden_params = {"additional_headers": {"x-existing": "keep"}} + + +def test_add_fallback_headers_to_streaming_wrapper(): + response = StreamingWrapper() + + result = add_fallback_headers_to_response( + response=response, + attempted_fallbacks=1, + ) + + assert result is response + assert response._hidden_params["additional_headers"] == { + "x-existing": "keep", + "x-litellm-attempted-fallbacks": 1, + } + + +def test_add_fallback_headers_serializes_fallback_errors(): + response = StreamingWrapper() + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + + result = add_fallback_headers_to_response( + response=response, + attempted_fallbacks=1, + fallback_errors=fallback_errors, + ) + + assert result is response + assert response._hidden_params["additional_headers"][ + "x-litellm-attempted-fallbacks" + ] == 1 + assert ( + json.loads( + response._hidden_params["additional_headers"]["x-litellm-fallback-errors"] + ) + == fallback_errors + ) + + +def test_add_retry_headers_to_streaming_wrapper(): + response = StreamingWrapper() + + result = add_retry_headers_to_response( + response=response, + attempted_retries=2, + max_retries=3, + ) + + assert result is response + assert response._hidden_params["additional_headers"] == { + "x-existing": "keep", + "x-litellm-attempted-retries": 2, + "x-litellm-max-retries": 3, + } + + +def test_get_hidden_params_dict_with_pydantic_model_hidden_params(): + class InnerHiddenParams(BaseModel): + additional_headers: dict = {} + + class Response: + def __init__(self): + self._hidden_params = InnerHiddenParams( + additional_headers={"x-custom": "value"} + ) + + result = get_hidden_params_dict(Response()) + assert result == {"additional_headers": {"x-custom": "value"}} + + +def test_get_hidden_params_dict_with_no_hidden_params(): + class PlainResponse: + pass + + assert get_hidden_params_dict(PlainResponse()) == {} + + +def test_add_fallback_headers_when_no_existing_additional_headers(): + class NoHeadersWrapper: + def __init__(self): + self._hidden_params = {} + + response = NoHeadersWrapper() + result = add_fallback_headers_to_response(response=response, attempted_fallbacks=2) + + assert result is response + assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 2 + + +def test_add_fallback_headers_returns_none_when_response_is_none(): + result = add_fallback_headers_to_response(response=None, attempted_fallbacks=1) + assert result is None + + +def test_add_fallback_headers_returns_unchanged_when_response_has_no_hidden_params(): + class PlainObject: + pass + + obj = PlainObject() + result = add_fallback_headers_to_response(response=obj, attempted_fallbacks=1) + assert result is obj + assert not hasattr(obj, "_hidden_params") + + +def test_get_fallback_errors_from_headers_existing_list_passthrough(): + errors = [{"message": "err", "type": "T", "param": None, "code": "400"}] + result = get_fallback_errors_from_headers({"x-litellm-fallback-errors": errors}) + assert result == errors + + +def test_get_fallback_errors_from_headers_invalid_json_returns_empty(): + result = get_fallback_errors_from_headers( + {"x-litellm-fallback-errors": "not-valid-json-{"} + ) + assert result == [] + + +def test_get_fallback_errors_from_headers_missing_key_returns_empty(): + result = get_fallback_errors_from_headers({}) + assert result == [] diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py new file mode 100644 index 00000000000..ca647bdce55 --- /dev/null +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -0,0 +1,139 @@ +import json + +import pytest + +from litellm.router_utils.fallback_event_handlers import run_async_fallback + + +class StreamingWrapper: + def __init__(self): + self._hidden_params = {"additional_headers": {}} + + +class FakeRouter: + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + return StreamingWrapper() + + +class AlwaysFailRouter: + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + raise RuntimeError("fallback model also failed") + + +@pytest.mark.asyncio +async def test_run_async_fallback_adds_errors_when_opted_in(): + response = await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + include_fallback_errors=True, + ) + + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["x-litellm-attempted-fallbacks"] == 1 + assert json.loads(additional_headers["x-litellm-fallback-errors"]) == [ + { + "message": "upstream limited request", + "type": "RuntimeError", + "param": None, + "code": None, + } + ] + + +@pytest.mark.asyncio +async def test_run_async_fallback_omits_errors_without_opt_in(): + response = await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + ) + + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["x-litellm-attempted-fallbacks"] == 1 + assert "x-litellm-fallback-errors" not in additional_headers + + +@pytest.mark.asyncio +async def test_run_async_fallback_raises_when_all_fallbacks_fail(): + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=AlwaysFailRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original request failed"), + max_fallbacks=3, + fallback_depth=0, + include_fallback_errors=True, + ) + + +class RecordingRouter: + def __init__(self): + self.received_kwargs = None + + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + self.received_kwargs = kwargs + return StreamingWrapper() + + +@pytest.mark.asyncio +async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_call(): + """A nested fallback (multi-hop) must keep collecting errors, so the opt-in + flag has to reach the nested async_function_with_fallbacks call.""" + router = RecordingRouter() + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + include_fallback_errors=True, + ) + + assert router.received_kwargs.get("include_fallback_errors") is True + + +@pytest.mark.asyncio +async def test_run_async_fallback_does_not_forward_flag_without_opt_in(): + router = RecordingRouter() + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("upstream limited request"), + max_fallbacks=3, + fallback_depth=0, + ) + + assert "include_fallback_errors" not in router.received_kwargs + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_original_model_group(): + response = await run_async_fallback( + litellm_router=FakeRouter(), + fallback_model_group=["primary-model", "fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + ) + + assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 diff --git a/tests/test_litellm/sandbox/test_e2b_sandbox.py b/tests/test_litellm/sandbox/test_e2b_sandbox.py index 948d39fac8f..cc5b12156a1 100644 --- a/tests/test_litellm/sandbox/test_e2b_sandbox.py +++ b/tests/test_litellm/sandbox/test_e2b_sandbox.py @@ -295,3 +295,24 @@ async def test_public_lifecycle_create_run_delete(): async def test_unsupported_provider_raises(): with pytest.raises(ValueError): await litellm.acreate_sandbox(provider="not-a-provider") + + +# ---------- api_base override ---------- + + +@pytest.mark.asyncio +async def test_create_uses_api_base_override(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox( + api_base="http://my-sandbox:8080", api_key="k", client=client + ) + _, url, _, _ = client.calls[0] + assert url == "http://my-sandbox:8080/sandboxes" + + +@pytest.mark.asyncio +async def test_create_defaults_to_e2b_api_base(): + client = FakeHTTPClient() + await E2BSandboxConfig().acreate_sandbox(api_key="k", client=client) + _, url, _, _ = client.calls[0] + assert url == "https://api.e2b.app/sandboxes" diff --git a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py b/tests/test_litellm/sandbox/test_opensandbox_sandbox.py new file mode 100644 index 00000000000..0d7bcbe1e53 --- /dev/null +++ b/tests/test_litellm/sandbox/test_opensandbox_sandbox.py @@ -0,0 +1,647 @@ +import json + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.sandbox.transformation import ContainerHandle +from litellm.llms.opensandbox.sandbox.transformation import ( + MAX_OUTPUT_BYTES, + OPEN_SANDBOX_DEFAULT_TEMPLATE, + OpenSandboxSandboxConfig, +) +from litellm.utils import ProviderConfigManager + +TEST_API_BASE = "https://sandbox.test/v1" + + +def http_status_error(status_code, url="http://test"): + return httpx.HTTPStatusError( + f"status {status_code}", + request=httpx.Request("GET", url), + response=httpx.Response(status_code), + ) + + +def sse(data): + return f"data: {json.dumps(data)}" + + +class FakeResponse: + def __init__(self, *, json_data=None, lines=None, status_code=200): + self._json = json_data + self._lines = lines or [] + self.status_code = status_code + + def json(self): + return self._json + + def raise_for_status(self): + if self.status_code >= 400: + raise http_status_error(self.status_code) + + async def aiter_lines(self): + for line in self._lines: + yield line + + +class FakeHTTPClient: + def __init__( + self, + *, + create_json=None, + sandbox_states=None, + endpoint_json=None, + endpoint_responses=None, + execute_lines=None, + delete_status=204, + execute_raises=None, + ): + self.create_json = create_json or { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + self.sandbox_states = list( + sandbox_states + or [ + { + "id": "osb_123", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/opt/code-interpreter/code-interpreter.sh"], + } + ] + ) + self.endpoint_json = endpoint_json or { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + } + self.endpoint_responses = ( + list(endpoint_responses) if endpoint_responses is not None else None + ) + self.execute_lines = execute_lines or [] + self.delete_status = delete_status + self.execute_raises = execute_raises + self.calls = [] + + async def post(self, url, headers=None, json=None, stream=False, **kwargs): + self.calls.append(("POST", url, headers, json, {"stream": stream})) + if url.endswith("/sandboxes"): + return FakeResponse(json_data=self.create_json) + if url.endswith("/code"): + if self.execute_raises is not None: + raise self.execute_raises + return FakeResponse(lines=self.execute_lines) + raise AssertionError(f"unexpected POST {url}") + + async def get(self, url, headers=None, params=None, **kwargs): + self.calls.append(("GET", url, headers, None, params)) + if "/endpoints/44772" in url: + if self.endpoint_responses is not None and self.endpoint_responses: + response = self.endpoint_responses.pop(0) + if isinstance(response, Exception): + raise response + if isinstance(response, FakeResponse): + return response + return FakeResponse(json_data=response) + return FakeResponse(json_data=self.endpoint_json) + if "/sandboxes/" in url: + state = self.sandbox_states.pop(0) + return FakeResponse(json_data=state) + raise AssertionError(f"unexpected GET {url}") + + async def delete(self, url, headers=None, **kwargs): + self.calls.append(("DELETE", url, headers, None, None)) + if not (200 <= self.delete_status < 300): + raise http_status_error(self.delete_status, url) + return FakeResponse(status_code=self.delete_status) + + +def test_parse_sse_lines_maps_output_result_count_and_error(): + lines = [ + sse({"type": "stdout", "text": "hello\n"}), + sse({"type": "stderr", "text": "warn\n"}), + sse({"type": "result", "results": {"text/plain": "4"}}), + sse({"type": "execution_count", "execution_count": 7}), + sse( + { + "type": "error", + "error": { + "ename": "ValueError", + "evalue": "bad", + "traceback": ["Traceback"], + }, + } + ), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "hello\n" + assert result.stderr == "warn\n" + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 7 + assert result.error == { + "name": "ValueError", + "value": "bad", + "traceback": ["Traceback"], + } + + +def test_parse_sse_lines_skips_non_json_and_control_lines(): + lines = [ + "event: message", + "not-json", + "", + sse({"type": "stdout", "text": "ok\n"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.stdout == "ok\n" + assert result.error is None + + +def test_parse_sse_lines_maps_fallback_shapes(): + lines = [ + "data:", + sse(["not-a-dict"]), + sse({"code": "BadRequest", "message": "nope"}), + sse({"type": "result", "text/plain": "4"}), + sse({"type": "error", "name": "RuntimeError", "text": "boom"}), + sse({"type": "execution_count", "execution_count": "8"}), + ] + + result = OpenSandboxSandboxConfig._parse_lines(lines) + + assert result.results == [{"text/plain": "4"}] + assert result.execution_count == 8 + assert result.error == { + "name": "BadRequest", + "value": "nope", + "traceback": [], + } + fallback_error = OpenSandboxSandboxConfig._parse_lines( + [sse({"type": "error", "name": "RuntimeError", "text": "boom"})] + ) + assert fallback_error.error == { + "name": "RuntimeError", + "value": "boom", + "traceback": [], + } + empty_string_error = OpenSandboxSandboxConfig._parse_lines( + [ + sse( + { + "type": "error", + "error": { + "ename": "", + "name": "FallbackName", + "evalue": "", + "value": "fallback value", + "traceback": [], + }, + } + ) + ] + ) + assert empty_string_error.error == { + "name": "", + "value": "", + "traceback": [], + } + + +def test_static_helpers_cover_defaults_and_fallbacks(monkeypatch): + def fake_secret(key): + if key == "OPEN_SANDBOX_API_KEY": + return "env-key" + if key == "OPEN_SANDBOX_API_BASE": + return TEST_API_BASE + return None + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + fake_secret, + ) + config = OpenSandboxSandboxConfig() + handle = ContainerHandle(id="osb", provider="opensandbox", domain="http://x/v1") + + assert config.validate_environment() == "env-key" + assert config.validate_environment(api_key="") == "" + assert config._api_key(api_key=None, handle=handle) == "env-key" + + handle._hidden_params = {"api_key": "stored-key"} + assert config._api_key(api_key=None, handle=handle) == "stored-key" + assert config._http(None) is not None + + body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy={"egress": [{"domain": "example.com"}]}, + secure_access=True, + ) + assert body["networkPolicy"] == {"egress": [{"domain": "example.com"}]} + assert body["secureAccess"] is True + + other_body = config._create_body( + template=None, + timeout=None, + allow_internet_access=False, + metadata=None, + env_vars=None, + resource_limits=None, + resource_requests=None, + entrypoint=None, + network_policy=None, + secure_access=False, + ) + assert body["resourceLimits"] is not other_body["resourceLimits"] + + assert config._sandbox_state(None) is None + assert config._sandbox_state({"status": "Running"}) is None + assert config._as_str_dict(None) == {} + assert config._endpoint_base_url("http://execd.local", "https://api/v1") == ( + "http://execd.local" + ) + assert config._api_base(None) == TEST_API_BASE + assert config._api_base("https://direct.test/v1/") == "https://direct.test/v1" + assert config._as_int("9") == 9 + assert config._as_int("nope") is None + assert config._as_int(None) is None + assert isinstance( + ProviderConfigManager.get_provider_sandbox_config("opensandbox"), + OpenSandboxSandboxConfig, + ) + + +def test_api_base_requires_kwarg_or_env(monkeypatch): + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.get_secret_str", + lambda key: None, + ) + + with pytest.raises(ValueError, match="api_base is required"): + OpenSandboxSandboxConfig._api_base(None) + + +@pytest.mark.asyncio +async def test_create_posts_default_body_and_omits_empty_api_key(): + client = FakeHTTPClient() + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + method, url, headers, body, _ = client.calls[0] + assert method == "POST" + assert url == f"{TEST_API_BASE}/sandboxes" + assert "OPEN-SANDBOX-API-KEY" not in headers + assert body["image"] == {"uri": OPEN_SANDBOX_DEFAULT_TEMPLATE} + assert body["entrypoint"] == ["/opt/code-interpreter/code-interpreter.sh"] + assert body["timeout"] == 300 + assert body["resourceLimits"] == {"cpu": "1", "memory": "2Gi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert handle.id == "osb_123" + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + + +@pytest.mark.asyncio +async def test_create_can_opt_into_internet_access(): + client = FakeHTTPClient() + + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + allow_internet_access=True, + client=client, + ) + + _, _, _, body, _ = client.calls[0] + assert "networkPolicy" not in body + + +@pytest.mark.asyncio +async def test_create_custom_options_poll_and_endpoint_resolution(): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + }, + sandbox_states=[ + { + "id": "osb_pending", + "status": {"state": "Running"}, + "createdAt": "2026-01-01T00:00:00Z", + "entrypoint": ["/bin/sh"], + } + ], + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + template="custom/image:latest", + timeout=600, + allow_internet_access=False, + api_key="osb-key", + api_base="https://sandbox.example/v1", + metadata={"suite": "unit"}, + env_vars={"PYTHONUNBUFFERED": "1"}, + resource_limits={"cpu": "500m", "memory": "512Mi"}, + resource_requests={"cpu": "250m", "memory": "256Mi"}, + entrypoint=["/bin/sh", "-lc", "sleep 3600"], + use_server_proxy=True, + client=client, + ) + + _, create_url, create_headers, body, _ = client.calls[0] + _, poll_url, poll_headers, _, _ = client.calls[1] + _, endpoint_url, endpoint_headers, _, endpoint_params = client.calls[2] + + assert create_url == "https://sandbox.example/v1/sandboxes" + assert create_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert body["image"] == {"uri": "custom/image:latest"} + assert body["entrypoint"] == ["/bin/sh", "-lc", "sleep 3600"] + assert body["metadata"] == {"suite": "unit"} + assert body["env"] == {"PYTHONUNBUFFERED": "1"} + assert body["resourceLimits"] == {"cpu": "500m", "memory": "512Mi"} + assert body["resourceRequests"] == {"cpu": "250m", "memory": "256Mi"} + assert body["networkPolicy"] == {"defaultAction": "deny", "egress": []} + assert poll_url == "https://sandbox.example/v1/sandboxes/osb_pending" + assert poll_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_url.endswith("/sandboxes/osb_pending/endpoints/44772") + assert endpoint_headers["OPEN-SANDBOX-API-KEY"] == "osb-key" + assert endpoint_params == {"use_server_proxy": True} + assert handle.id == "osb_pending" + + +@pytest.mark.asyncio +async def test_create_waits_across_pending_state(monkeypatch): + client = FakeHTTPClient( + create_json={ + "id": "osb_pending", + "status": {"state": "Pending"}, + "createdAt": "2026-01-01T00:00:00Z", + }, + sandbox_states=[ + {"id": "osb_pending", "status": {"state": "Pending"}}, + {"id": "osb_pending", "status": {"state": "Running"}}, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + assert handle.id == "osb_pending" + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_for_terminal_state(): + client = FakeHTTPClient( + create_json={"id": "osb_failed", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_failed", "status": {"state": "Failed"}}, + ], + ) + + with pytest.raises(ValueError, match="entered Failed"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_create_times_out_waiting_for_running(): + client = FakeHTTPClient( + create_json={"id": "osb_slow", "status": {"state": "Pending"}}, + sandbox_states=[ + {"id": "osb_slow", "status": {"state": "Pending"}}, + ], + ) + + with pytest.raises(TimeoutError, match="was not Running"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=0, + poll_interval=0, + client=client, + ) + + +@pytest.mark.asyncio +async def test_create_waits_for_endpoint_resolution(monkeypatch): + client = FakeHTTPClient( + endpoint_responses=[ + http_status_error(404, f"{TEST_API_BASE}/sandboxes/osb_123"), + { + "endpoint": "execd.local:44772", + "headers": {"X-EXECD-ACCESS-TOKEN": "execd-token"}, + }, + ], + ) + sleeps = [] + + async def fake_sleep(interval): + sleeps.append(interval) + + monkeypatch.setattr( + "litellm.llms.opensandbox.sandbox.transformation.asyncio.sleep", fake_sleep + ) + + handle = await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", + api_base=TEST_API_BASE, + ready_timeout=1, + poll_interval=0.01, + client=client, + ) + + endpoint_calls = [call for call in client.calls if "/endpoints/44772" in call[1]] + assert handle._hidden_params["execd_endpoint"] == "execd.local:44772" + assert len(endpoint_calls) == 2 + assert sleeps == [0.01] + + +@pytest.mark.asyncio +async def test_create_raises_when_endpoint_is_missing(): + client = FakeHTTPClient(endpoint_json={"headers": {"X": "y"}}) + + with pytest.raises(TimeoutError, match="execd endpoint.*not ready"): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, ready_timeout=0, client=client + ) + + +@pytest.mark.asyncio +async def test_create_reraises_non_404_endpoint_error(): + client = FakeHTTPClient(endpoint_responses=[http_status_error(500)]) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().acreate_sandbox( + api_key="", api_base=TEST_API_BASE, client=client + ) + + +@pytest.mark.asyncio +async def test_run_code_resolves_bare_id_and_posts_sse_request(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + result = await OpenSandboxSandboxConfig().arun_code( + container="osb_bare", + code="print(6*7)", + language="python", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + endpoint_call = client.calls[0] + run_call = client.calls[1] + assert endpoint_call[0] == "GET" + assert ( + endpoint_call[1] == "http://sandbox.local/v1/sandboxes/osb_bare/endpoints/44772" + ) + assert run_call[0] == "POST" + assert run_call[1] == "http://execd.local:44772/code" + assert run_call[2]["X-EXECD-ACCESS-TOKEN"] == "execd-token" + assert run_call[3] == { + "code": "print(6*7)", + "context": {"language": "python"}, + } + assert run_call[4] == {"stream": True} + assert result.stdout == "42\n" + + +@pytest.mark.asyncio +async def test_run_code_uses_https_for_scheme_less_endpoint_when_api_base_is_https(): + client = FakeHTTPClient() + handle = ContainerHandle( + id="osb_https", provider="opensandbox", domain="https://sandbox.example/v1" + ) + handle._hidden_params = { + "execd_endpoint": "execd.example/route/44772", + "execd_headers": {}, + } + + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print(1)", client=client + ) + + assert client.calls[0][1] == "https://execd.example/route/44772/code" + + +@pytest.mark.asyncio +async def test_run_code_aborts_on_output_over_cap(): + client = FakeHTTPClient(execute_lines=["x" * (MAX_OUTPUT_BYTES + 1)]) + handle = ContainerHandle(id="osb_big", provider="opensandbox", domain="http://x/v1") + handle._hidden_params = {"execd_endpoint": "execd.local:44772", "execd_headers": {}} + + with pytest.raises(ValueError, match="exceeded"): + await OpenSandboxSandboxConfig().arun_code( + container=handle, code="print('x')", client=client + ) + + +@pytest.mark.asyncio +async def test_delete_returns_false_on_404(): + client = FakeHTTPClient(delete_status=404) + + ok = await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_gone", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + assert ok is False + + +@pytest.mark.asyncio +async def test_delete_reraises_non_404_http_error(): + client = FakeHTTPClient(delete_status=500) + + with pytest.raises(httpx.HTTPStatusError): + await OpenSandboxSandboxConfig().adelete_sandbox( + container="osb_err", + api_key="", + api_base="http://sandbox.local/v1", + client=client, + ) + + +@pytest.mark.asyncio +async def test_public_lifecycle_create_run_delete(): + client = FakeHTTPClient( + execute_lines=[ + sse({"type": "stdout", "text": "42\n"}), + ] + ) + + container = await litellm.acreate_sandbox( + provider="opensandbox", api_key="", api_base=TEST_API_BASE, client=client + ) + result = await litellm.arun_code( + provider="opensandbox", + container=container, + code="print(6*7)", + api_key="", + client=client, + ) + ok = await litellm.adelete_sandbox( + provider="opensandbox", + container=container, + api_key="", + client=client, + ) + + assert container.id == "osb_123" + assert result.stdout == "42\n" + assert ok is True + + +@pytest.mark.asyncio +async def test_code_interpreter_tool_deletes_even_when_run_raises(): + client = FakeHTTPClient(execute_raises=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await litellm.acode_interpreter_tool( + provider="opensandbox", + code="1/0", + api_key="", + api_base=TEST_API_BASE, + client=client, + ) + + assert [call[0] for call in client.calls] == ["POST", "GET", "POST", "DELETE"] + assert client.calls[0][1].endswith("/sandboxes") + assert client.calls[1][1].endswith("/endpoints/44772") + assert client.calls[2][1].endswith("/code") + assert client.calls[3][1].endswith("/sandboxes/osb_123") diff --git a/tests/test_litellm/sandbox/test_sandbox_tools.py b/tests/test_litellm/sandbox/test_sandbox_tools.py new file mode 100644 index 00000000000..06136534b13 --- /dev/null +++ b/tests/test_litellm/sandbox/test_sandbox_tools.py @@ -0,0 +1,181 @@ +"""Unit tests for the sandbox-tool registry.""" + +from litellm.sandbox import sandbox_tools + + +def _reset(): + sandbox_tools.clear_sandbox_tools() + + +def test_register_resolves_provider_key_and_base(): + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved == { + "sandbox_provider": "e2b", + "api_key": "sk-literal", + "api_base": "https://sandbox.internal", + } + _reset() + + +def test_register_clears_stale_entries_on_reload(): + """A tool removed from the config must not survive a re-registration.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "old", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("old") is not None + + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "new", + "litellm_params": {"sandbox_provider": "e2b"}, + } + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("new") is not None + assert ( + sandbox_tools.resolve_sandbox_tool("old") is None + ), "stale tool must be gone after the config is reloaded" + _reset() + + +def test_register_empty_list_clears_removed_tools(): + """Reloading a config with sandbox_tools removed (the proxy passes an empty + list) must drop previously registered credentials from the process.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": {"sandbox_provider": "e2b", "api_key": "sk-x"}, + } + ] + ) + assert sandbox_tools.resolve_sandbox_tool("e2b_default") is not None + + sandbox_tools.register_sandbox_tools([]) + + assert ( + sandbox_tools.resolve_sandbox_tool("e2b_default") is None + ), "removing sandbox_tools from config must clear stale credentials" + _reset() + + +def test_register_resolves_secret_from_env(monkeypatch): + _reset() + monkeypatch.setenv("MY_SANDBOX_KEY", "sk-from-env") + sandbox_tools.register_sandbox_tools( + [ + { + "sandbox_tool_name": "e2b_default", + "litellm_params": { + "sandbox_provider": "e2b", + "api_key": "os.environ/MY_SANDBOX_KEY", + }, + } + ] + ) + + resolved = sandbox_tools.resolve_sandbox_tool("e2b_default") + assert resolved is not None + assert resolved["api_key"] == "sk-from-env" + assert resolved["api_base"] is None + _reset() + + +def test_resolve_unknown_returns_none(): + _reset() + assert sandbox_tools.resolve_sandbox_tool("nope") is None + + +def test_register_skips_malformed_entries_without_crashing(): + """A single malformed entry (missing sandbox_tool_name, or not a dict) must + not crash registration during proxy startup/hot-reload; valid entries in the + same list must still register.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"litellm_params": {"sandbox_provider": "e2b"}}, # missing name + "not-a-dict", # wrong type + {"sandbox_tool_name": "", "litellm_params": {}}, # empty name + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("good") is not None + assert sandbox_tools.resolve_sandbox_tool("") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_skips_entry_missing_sandbox_provider(): + """An entry with a name but no sandbox_provider must be skipped at + registration so it cannot later resolve and call acreate_sandbox(provider=None), + which fails with a cryptic runtime error instead of a clear startup warning.""" + _reset() + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "no_provider", "litellm_params": {"api_key": "sk-x"}}, + { + "sandbox_tool_name": "null_provider", + "litellm_params": {"sandbox_provider": None}, + }, + { + "sandbox_tool_name": "good", + "litellm_params": {"sandbox_provider": "e2b"}, + }, + ] + ) + + assert sandbox_tools.resolve_sandbox_tool("no_provider") is None + assert sandbox_tools.resolve_sandbox_tool("null_provider") is None + assert set(sandbox_tools._SANDBOX_TOOL_REGISTRY) == {"good"} + _reset() + + +def test_register_swaps_registry_atomically(): + """register_sandbox_tools must replace the registry in one rebind so a + concurrent resolve never observes a half-populated or transiently empty + registry between clearing and repopulating.""" + _reset() + sandbox_tools.register_sandbox_tools( + [{"sandbox_tool_name": "a", "litellm_params": {"sandbox_provider": "e2b"}}] + ) + before = sandbox_tools._SANDBOX_TOOL_REGISTRY + + sandbox_tools.register_sandbox_tools( + [ + {"sandbox_tool_name": "b", "litellm_params": {"sandbox_provider": "e2b"}}, + {"sandbox_tool_name": "c", "litellm_params": {"sandbox_provider": "e2b"}}, + ] + ) + after = sandbox_tools._SANDBOX_TOOL_REGISTRY + + assert after is not before, "the registry must be replaced, not mutated in place" + assert set(after) == {"b", "c"} + assert "a" not in after + _reset() diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py new file mode 100644 index 00000000000..9ca4515239a --- /dev/null +++ b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py @@ -0,0 +1,89 @@ +""" +Regression tests for the Cloudflare Workers AI text-generation catalog in the +model-cost map. + +The Cloudflare list was badly stale (only 4 ancient entries). These tests pin +the newly added current Workers AI models (sourced from Cloudflare's live +``/ai/models/search?task=Text Generation`` catalog) and guard against the root +``model_prices_and_context_window.json`` and the bundled +``litellm/model_prices_and_context_window_backup.json`` drifting out of sync for +the ``cloudflare/`` namespace. +""" + +import json +import os + +import pytest + +import litellm + +ROOT_MAP = os.path.join( + os.path.dirname(os.path.dirname(litellm.__file__)), + "model_prices_and_context_window.json", +) +BACKUP_MAP = os.path.join( + os.path.dirname(litellm.__file__), + "model_prices_and_context_window_backup.json", +) + + +@pytest.fixture(autouse=True) +def _use_local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + try: + yield + finally: + litellm.model_cost = original_model_cost + + +def _load(path: str) -> dict: + with open(path, encoding="utf-8") as f: + return json.load(f) + + +def _cloudflare_keys(data: dict) -> set: + return {k for k in data if k.startswith("cloudflare/")} + + +def test_glm_5_2_entry_is_present_and_well_formed(): + entry = litellm.model_cost["cloudflare/@cf/zai-org/glm-5.2"] + assert entry["litellm_provider"] == "cloudflare" + assert entry["mode"] == "chat" + assert entry["supports_function_calling"] is True + assert entry["input_cost_per_token"] > 0 + assert entry["output_cost_per_token"] > 0 + + +def test_vision_model_is_flagged_supports_vision(): + entry = litellm.model_cost["cloudflare/@cf/meta/llama-3.2-11b-vision-instruct"] + assert entry["litellm_provider"] == "cloudflare" + assert entry.get("supports_vision") is True + + +def test_additional_current_models_are_present(): + for key in ( + "cloudflare/@cf/openai/gpt-oss-120b", + "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast", + ): + entry = litellm.model_cost[key] + assert entry["litellm_provider"] == "cloudflare" + assert entry["mode"] == "chat" + assert entry["supports_function_calling"] is True + assert entry["input_cost_per_token"] > 0 + assert entry["output_cost_per_token"] > 0 + + +def test_root_and_backup_have_identical_cloudflare_keys(): + if not os.path.exists(ROOT_MAP): + pytest.skip("root cost map only ships in source checkouts") + assert _cloudflare_keys(_load(ROOT_MAP)) == _cloudflare_keys(_load(BACKUP_MAP)) + + +def test_root_and_backup_cloudflare_entries_are_byte_for_byte_equal(): + if not os.path.exists(ROOT_MAP): + pytest.skip("root cost map only ships in source checkouts") + root = {k: v for k, v in _load(ROOT_MAP).items() if k.startswith("cloudflare/")} + backup = {k: v for k, v in _load(BACKUP_MAP).items() if k.startswith("cloudflare/")} + assert root == backup diff --git a/tests/test_litellm/test_completion_timeout_resolution.py b/tests/test_litellm/test_completion_timeout_resolution.py index a76cc6f7de8..7eb79e90e60 100644 --- a/tests/test_litellm/test_completion_timeout_resolution.py +++ b/tests/test_litellm/test_completion_timeout_resolution.py @@ -63,8 +63,9 @@ def test_global_timeout_from_litellm_settings(): ) -def test_global_timeout_package_default_coerced_to_600_for_completion(): - """Package default 6000s → 600s for completion-only path.""" +def test_explicit_global_timeout_6000_is_preserved(): + """The caller passes the explicitly-configured value (or None); an explicit + 6000 must be honored, not silently coerced to 600.""" assert ( CompletionTimeout.resolve( None, @@ -73,7 +74,7 @@ def test_global_timeout_package_default_coerced_to_600_for_completion(): global_timeout=6000.0, supports_httpx_timeout=supports_httpx_timeout, ) - == 600.0 + == 6000.0 ) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index a67c5b41f36..db6f945ff78 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -505,6 +505,74 @@ def test_realtime_logging_object_allows_null_transcript_in_conversation_item_add assert logging_result.results[0]["item"]["content"][0]["transcript"] is None +def test_realtime_logging_object_does_not_validate_unknown_event_types(): + """ + A realtime session emits events outside the OpenAIRealtimeEvents union (e.g. + rate_limits.updated, response.function_call_arguments.delta). Building the + logging object must not revalidate every event against the union; doing so + produces thousands of Pydantic ValidationErrors per session, blocks the event + loop, and the raised error discards the session's usage. The events must + survive verbatim, the combined usage must be preserved, and serialization + must stay clean. + """ + import warnings + + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "event_id": "ev0", "session": {"id": "s"}}, + ] + for i in range(50): + results += [ + { + "type": "rate_limits.updated", + "event_id": f"rl{i}", + "rate_limits": [{"name": "requests", "limit": 1000, "remaining": 900}], + }, + { + "type": "response.function_call_arguments.delta", + "event_id": f"fc{i}", + "delta": "{}", + }, + { + "type": "response.done", + "event_id": f"rd{i}", + "response": { + "usage": { + "input_tokens": 4, + "output_tokens": 6, + "total_tokens": 10, + } + }, + }, + ] + + usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=results + ) + # On unfixed code this raises pydantic ValidationError instead of returning. + logging_result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object( + usage=usage, + results=results, + ) + + assert logging_result.usage.total_tokens == 500 + assert len(logging_result.results) == len(results) + unknown_types = { + r["type"] + for r in logging_result.results + if r["type"] + in ("rate_limits.updated", "response.function_call_arguments.delta") + } + assert unknown_types == { + "rate_limits.updated", + "response.function_call_arguments.delta", + } + + with warnings.catch_warnings(): + warnings.simplefilter("error") + dumped = logging_result.model_dump() + assert len(dumped["results"]) == len(results) + + def test_realtime_transcription_duration_cost(monkeypatch): """ gpt-realtime-whisper transcription sessions are billed by input audio duration diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index c2aa4a095d4..3e2150848a7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -4901,3 +4901,107 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): ) is True ) + + +class TestRouterRequestTimeoutPropagation: + """litellm_settings.request_timeout must act as an independent per-attempt timeout. + + Regression for LIT-2369: request_timeout was shadowed by router_settings.timeout, + so Bedrock (and other provider) calls fell back to the hardcoded 600s httpx + default instead of the configured value. + """ + + def _make_router(self, timeout=None, stream_timeout=None): + return litellm.Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "sk-test", + }, + } + ], + timeout=timeout, + stream_timeout=stream_timeout, + ) + + @pytest.fixture + def explicit_request_timeout(self): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + litellm.request_timeout = 300 + litellm.request_timeout_explicitly_set = True + try: + yield 300 + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + def test_request_timeout_stored_independently_when_both_set( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + assert router.timeout == 330 + assert router.request_timeout == 300 + + def test_request_timeout_none_when_not_explicitly_configured(self): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = False + try: + router = self._make_router(timeout=330) + assert router.timeout == 330 + assert router.request_timeout is None + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + def test_non_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 + + def test_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + # stream=True resolves through _get_stream_timeout; request_timeout must win. + assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 + + def test_explicit_stream_timeout_still_wins_over_request_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330, stream_timeout=45) + assert router._get_stream_timeout(kwargs={}, data={}) == 45 + + def test_non_stream_falls_through_to_router_timeout_without_request_timeout(self): + original_value = litellm.request_timeout + original_flag = litellm.request_timeout_explicitly_set + litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS + litellm.request_timeout_explicitly_set = False + try: + router = self._make_router(timeout=330) + assert router._get_non_stream_timeout(kwargs={}, data={}) == 330 + finally: + litellm.request_timeout = original_value + litellm.request_timeout_explicitly_set = original_flag + + def test_per_deployment_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 + + def test_per_request_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): + router = self._make_router(timeout=330) + assert ( + router._get_non_stream_timeout( + kwargs={"timeout": 60}, data={"timeout": 120} + ) + == 60 + ) diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index ee64f44d32c..d4ac9659f00 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -537,3 +537,147 @@ def test_inherit_builtin_cache_pricing_noop_for_unknown_backend(): ) assert model_info == {"input_cost_per_token": 0.000003} + + +def test_custom_pricing_field_denylist_covers_all_builtin_pricing_fields(): + """The shared-backend-key stripping in Router relies on + CustomPricingLiteLLMParams enumerating every per-deployment pricing field. + If a new pricing field is added to ModelInfoBase but not mirrored here, a + deployment override on that field leaks into the shared backend key and + every sibling deployment reads the wrong rate (LIT-3897). This guard fails + fast when the two drift apart. + """ + import typing + + from litellm.types.utils import CustomPricingLiteLLMParams, ModelInfoBase + + pricing_markers = ("cost", "price", "uplift", "vector_size", "tiered_pricing") + builtin_pricing_fields = { + name + for name in typing.get_type_hints(ModelInfoBase) + if any(marker in name for marker in pricing_markers) + } + denylisted_fields = set(CustomPricingLiteLLMParams.model_fields.keys()) + + uncovered = sorted(builtin_pricing_fields - denylisted_fields) + assert not uncovered, ( + "ModelInfoBase pricing fields missing from CustomPricingLiteLLMParams; " + f"these would leak into shared backend keys: {uncovered}" + ) + + +def test_tiered_pricing_override_isolated_from_sibling_via_model_info_lookup(): + """LIT-3897: a deployment that overrides a tiered pricing field + (input_cost_per_token_above_272k_tokens) must not pollute the shared + backend key, so a sibling sharing the same backend resolves its pricing + via litellm.get_model_info (the path /model/info uses) without seeing the + override. + """ + backend_model = "gemini/gemini-2.5-flash" + override = 0.000999 + + builtin_info = litellm.get_model_info(model=backend_model) + assert builtin_info.get("input_cost_per_token_above_272k_tokens") != override + + model_keys = { + "lit3897-tiered-custom": litellm.model_cost.get("lit3897-tiered-custom"), + "lit3897-tiered-sibling": litellm.model_cost.get("lit3897-tiered-sibling"), + backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), + } + try: + Router( + model_list=[ + { + "model_name": "custom-priced-flash", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-tiered-1", + }, + "model_info": { + "id": "lit3897-tiered-custom", + "input_cost_per_token_above_272k_tokens": override, + "cache_read_input_token_cost_above_272k_tokens": override, + }, + }, + { + "model_name": "gemini-2.5-flash", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-tiered-2", + }, + "model_info": {"id": "lit3897-tiered-sibling"}, + }, + ], + ) + + shared = litellm.get_model_info(model=backend_model) + assert shared.get("input_cost_per_token_above_272k_tokens") != override, ( + "Tiered override leaked into the shared backend key; siblings read " + "the wrong rate via /model/info" + ) + assert shared.get("cache_read_input_token_cost_above_272k_tokens") != override + + custom_entry = litellm.model_cost["lit3897-tiered-custom"] + assert custom_entry["input_cost_per_token_above_272k_tokens"] == override + assert custom_entry["cache_read_input_token_cost_above_272k_tokens"] == override + finally: + _restore_model_cost_entries(model_keys) + + +def test_custom_pricing_isolated_from_sibling_via_proxy_model_info_path(): + """LIT-3897 end to end through the proxy resolution helper: the override + deployment reports its custom input rate while the sibling keeps the + canonical gemini rate when /model/info resolves each deployment. Mirrors the + ticket config where the override is set on litellm_params. + """ + from litellm.proxy.proxy_server import _get_proxy_model_info + + backend_model = "gemini/gemini-2.5-flash" + override_input = 5e-05 + override_output = 1e-04 + + builtin_info = litellm.get_model_info(model=backend_model) + builtin_input = builtin_info["input_cost_per_token"] + assert builtin_input != override_input + + model_keys = { + "lit3897-proxy-custom": litellm.model_cost.get("lit3897-proxy-custom"), + "lit3897-proxy-sibling": litellm.model_cost.get("lit3897-proxy-sibling"), + backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), + } + try: + router = Router( + model_list=[ + { + "model_name": "custom-priced-flash", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-proxy-1", + "input_cost_per_token": override_input, + "output_cost_per_token": override_output, + }, + "model_info": {"id": "lit3897-proxy-custom"}, + }, + { + "model_name": "gemini-2.5-flash", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key-proxy-2", + }, + "model_info": {"id": "lit3897-proxy-sibling"}, + }, + ], + ) + + resolved = { + m["model_name"]: _get_proxy_model_info(model=copy.deepcopy(m))[ + "model_info" + ]["input_cost_per_token"] + for m in router.model_list + } + + assert resolved["custom-priced-flash"] == override_input + assert resolved["gemini-2.5-flash"] == builtin_input + assert resolved["gemini-2.5-flash"] != resolved["custom-priced-flash"] + finally: + _restore_model_cost_entries(model_keys) diff --git a/tests/test_litellm/test_router_per_deployment_num_retries.py b/tests/test_litellm/test_router_per_deployment_num_retries.py index 154ba579e4e..af2372616a6 100644 --- a/tests/test_litellm/test_router_per_deployment_num_retries.py +++ b/tests/test_litellm/test_router_per_deployment_num_retries.py @@ -4,8 +4,9 @@ GitHub Issue: #18968 - Per-deployment max_retries/num_retries in litellm_params """ import pytest -from unittest.mock import MagicMock, patch +from unittest.mock import patch +import litellm from litellm import Router @@ -188,3 +189,133 @@ class TestPerDeploymentNumRetries: # Verify num_retries was converted from string to int assert exc.num_retries == 6 + + +class TestNumRetriesNoneGuard: + """ + Regression tests for the num_retries=None TypeError in async_function_with_retries. + + When num_retries reaches async_function_with_retries as None - e.g. a caller passes + num_retries=None explicitly (dict.get() does not fall back on an existing None value), + an auto_router/complexity_router path does not propagate it, or + Router.update_settings(num_retries=None) is used - AND the underlying call fails with a + retryable error, the comparison `if num_retries > 0:` raised: + + TypeError: '>' not supported between instances of 'NoneType' and 'int' + + This masked the real upstream error (rate limit / connection / 5xx) behind a TypeError. + Related issues: #23316, #25889, #23699, #28126. + """ + + @staticmethod + def _mock_router(num_retries=2): + return Router( + model_list=[ + { + "model_name": "mock-model", + "litellm_params": { + "model": "gpt-4o-mini", + "mock_response": "ok", + }, + } + ], + num_retries=num_retries, + ) + + def test_update_kwargs_normalises_explicit_none_to_router_default(self): + """ + _update_kwargs_before_fallbacks must normalise an explicit num_retries=None to + the router default (not leave it as None), while preserving an explicit 0. + """ + router = self._mock_router(num_retries=4) + + # explicit None -> router default + kwargs = {"num_retries": None} + router._update_kwargs_before_fallbacks(model="mock-model", kwargs=kwargs) + assert kwargs["num_retries"] == 4 + + # explicit 0 is preserved (retries stay disabled) + kwargs = {"num_retries": 0} + router._update_kwargs_before_fallbacks(model="mock-model", kwargs=kwargs) + assert kwargs["num_retries"] == 0 + + # absent -> router default (unchanged behaviour) + kwargs = {} + router._update_kwargs_before_fallbacks(model="mock-model", kwargs=kwargs) + assert kwargs["num_retries"] == 4 + + # explicit None with router default also None -> 0 (mirrors the downstream guard) + router.num_retries = None # simulate update_settings(num_retries=None) (#28126) + kwargs = {"num_retries": None} + router._update_kwargs_before_fallbacks(model="mock-model", kwargs=kwargs) + assert kwargs["num_retries"] == 0 + + @pytest.mark.asyncio + async def test_acompletion_num_retries_none_does_not_raise_typeerror(self): + """ + Per-request num_retries=None + a retryable error must NOT raise TypeError. + The router falls back to its configured num_retries and retries the (transient) + error, so the request succeeds. + """ + router = self._mock_router(num_retries=2) + with patch("asyncio.sleep", return_value=None): + response = await router.acompletion( + model="mock-model", + messages=[{"role": "user", "content": "hi"}], + num_retries=None, # the trigger + mock_testing_rate_limit_error=True, # retryable error path + ) + assert response.choices[0].message.content == "ok" + + @pytest.mark.asyncio + async def test_async_function_with_retries_none_falls_back_to_zero(self): + """ + When both the per-request value AND the router-level setting are None + (e.g. after Router.update_settings(num_retries=None), #28126), num_retries must + fall back to 0 and the real retryable error must surface - not a TypeError. + """ + router = self._mock_router(num_retries=0) + router.num_retries = None # simulate update_settings(num_retries=None) + + async def failing_fn(*args, **kwargs): + raise litellm.RateLimitError( + message="boom", model="mock-model", llm_provider="openai" + ) + + with patch("asyncio.sleep", return_value=None): + with pytest.raises(litellm.RateLimitError): + await router.async_function_with_retries( + original_function=failing_fn, + model="mock-model", + messages=[{"role": "user", "content": "hi"}], + num_retries=None, + ) + + @pytest.mark.asyncio + async def test_async_function_with_retries_none_falls_back_to_router_default(self): + """ + A None per-request num_retries falls back to the router-level setting, so retries + still happen (original_function is invoked more than once) before the real error + is raised - proving None did not silently disable retries or crash. + """ + router = self._mock_router(num_retries=3) + calls = {"n": 0} + + async def failing_fn(*args, **kwargs): + calls["n"] += 1 + raise litellm.InternalServerError( + message="boom", model="mock-model", llm_provider="openai" + ) + + with patch("asyncio.sleep", return_value=None): + with pytest.raises(litellm.InternalServerError): + await router.async_function_with_retries( + original_function=failing_fn, + model="mock-model", + messages=[{"role": "user", "content": "hi"}], + metadata={}, # populated by acompletion in the real path; log_retry needs it + num_retries=None, + ) + + # 1 initial attempt + at least 1 retry -> proves None fell back to a positive int + assert calls["n"] >= 2 diff --git a/tests/test_litellm/test_router_streaming_fallback_metadata.py b/tests/test_litellm/test_router_streaming_fallback_metadata.py new file mode 100644 index 00000000000..6ed70dc7cfe --- /dev/null +++ b/tests/test_litellm/test_router_streaming_fallback_metadata.py @@ -0,0 +1,187 @@ +import json +from unittest.mock import MagicMock + +import pytest + +import litellm +from litellm.proxy.proxy_server import _should_include_fallback_errors +from litellm.router import Router +from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict + + +def test_apply_fallback_hidden_params_copies_from_fallback_response(): + fallback_errors = [ + { + "message": "litellm.RateLimitError: upstream limited request", + "type": "RateLimitError", + "param": None, + "code": "429", + } + ] + chunk = litellm.ModelResponseStream( + id="test", + model="openai/internal-fallback", + choices=[], + ) + chunk._hidden_params = { + "additional_headers": {"x-existing-chunk-header": "keep"}, + "model_id": "chunk-model-id", + } + fallback_response = MagicMock() + fallback_response._hidden_params = { + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + }, + "api_base": "https://fallback.example", + } + + Router._apply_fallback_hidden_params_to_item( + fallback_item=chunk, + prepared_fallback_hidden_params=Router._prepare_fallback_hidden_params( + fallback_response + ), + ) + + assert chunk._hidden_params["api_base"] == "https://fallback.example" + assert chunk._hidden_params["model_id"] == "chunk-model-id" + assert chunk._hidden_params["additional_headers"] == { + "x-existing-chunk-header": "keep", + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps(fallback_errors), + } + + +def _two_group_fallback_router() -> Router: + return litellm.Router( + model_list=[ + { + "model_name": "primary-model", + "litellm_params": {"model": "openai/gpt-fake", "api_key": "sk-fake"}, + }, + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-fake-2", "api_key": "sk-fake"}, + }, + ], + fallbacks=[{"primary-model": ["fallback-model"]}], + ) + + +def _additional_headers(response: object) -> dict: + return get_hidden_params_dict(response).get("additional_headers", {}) + + +@pytest.mark.asyncio +async def test_include_fallback_errors_propagates_through_router(): + router = _two_group_fallback_router() + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + mock_testing_fallbacks=True, + mock_response="fallback success", + include_fallback_errors=True, + ) + + headers = _additional_headers(response) + assert headers["x-litellm-attempted-fallbacks"] == 1 + errors = json.loads(headers["x-litellm-fallback-errors"]) + assert isinstance(errors, list) and len(errors) >= 1 + assert set(errors[0].keys()) == {"message", "type", "param", "code"} + + +@pytest.mark.asyncio +async def test_router_omits_fallback_errors_without_opt_in(): + router = _two_group_fallback_router() + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + mock_testing_fallbacks=True, + mock_response="fallback success", + ) + + headers = _additional_headers(response) + assert headers["x-litellm-attempted-fallbacks"] == 1 + assert "x-litellm-fallback-errors" not in headers + + +def test_prepare_fallback_hidden_params_no_additional_headers(): + class FakeResponse: + _hidden_params = {"api_base": "http://example.com"} + + hidden_params, headers = Router._prepare_fallback_hidden_params(FakeResponse()) + assert hidden_params == {"api_base": "http://example.com"} + assert headers == {} + + +def test_apply_fallback_hidden_params_to_item_none_item(): + Router._apply_fallback_hidden_params_to_item( + None, ({"api_base": "http://fallback.example"}, {"x-custom": "value"}) + ) + + +def test_apply_fallback_hidden_params_to_item_no_existing_additional_headers(): + class FakeChunk: + _hidden_params = {"model_id": "test-id"} + + chunk = FakeChunk() + Router._apply_fallback_hidden_params_to_item( + chunk, + ( + {"api_base": "http://fallback.example"}, + {"x-litellm-attempted-fallbacks": 1}, + ), + ) + + assert chunk._hidden_params["api_base"] == "http://fallback.example" + assert chunk._hidden_params["model_id"] == "test-id" + assert chunk._hidden_params["additional_headers"] == { + "x-litellm-attempted-fallbacks": 1 + } + + +@pytest.mark.asyncio +async def test_set_response_headers_adds_model_group_to_streaming_wrapper(): + class StreamingWrapper: + def __init__(self): + self._hidden_params = {"additional_headers": {"x-existing": "keep"}} + + router = litellm.Router(model_list=[]) + response = StreamingWrapper() + + result = await router.set_response_headers( + response=response, + model_group="fallback-model", + ) + + assert result is response + assert response._hidden_params["additional_headers"] == { + "x-existing": "keep", + "x-litellm-model-group": "fallback-model", + } + + +def test_should_include_fallback_errors_gated_by_operator_setting(): + request_data: dict = {"include_fallback_errors": True} + + import litellm.proxy.proxy_server as ps + + original = ps.general_settings.copy() if isinstance(ps.general_settings, dict) else {} + try: + ps.general_settings = {} + assert _should_include_fallback_errors(request_data) is False + + ps.general_settings = {"expose_fallback_errors_to_caller": False} + assert _should_include_fallback_errors(request_data) is False + + ps.general_settings = {"expose_fallback_errors_to_caller": True} + assert _should_include_fallback_errors(request_data) is True + + ps.general_settings = {"expose_fallback_errors_to_caller": True} + assert _should_include_fallback_errors({}) is False + finally: + ps.general_settings = original diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/test_litellm/test_ruff_strict_gate.py index 96852e3a8a5..22255f0555e 100644 --- a/tests/test_litellm/test_ruff_strict_gate.py +++ b/tests/test_litellm/test_ruff_strict_gate.py @@ -82,13 +82,3 @@ def test_introduced_keeps_only_violations_on_changed_lines(): @pytest.mark.parametrize("hunk", ["@@ -1 +1 @@", "@@ -1,0 +1,2 @@"]) def test_parse_changed_lines_handles_single_and_ranged_hunks(hunk): assert gate.parse_changed_lines(f"+++ b/litellm/a.py\n{hunk}\n")["litellm/a.py"] - - -def test_report_emits_breached_rules_as_final_line(capsys): - # CI surfaces only the tail of the log, so the breached-rule summary (rule, - # total/cap, added) must be the last line or it gets truncated away. - breaches = sorted([gate.Breach("UP045", 530, 529, 1), gate.Breach("ANN401", 12, 10, 2)]) - new = [gate.Violation("litellm/types/llms/bedrock.py", 16, "UP045")] - gate.report(breaches, new, "origin/litellm_internal_staging") - last = capsys.readouterr().out.strip().splitlines()[-1] - assert last == "BREACHED RULES: ANN401 12/10 (+2); UP045 530/529 (+1)" diff --git a/tests/test_litellm/test_type_check_gate.py b/tests/test_litellm/test_type_check_gate.py index 18374c5db4b..e99ad0a4f41 100644 --- a/tests/test_litellm/test_type_check_gate.py +++ b/tests/test_litellm/test_type_check_gate.py @@ -56,29 +56,58 @@ def test_paths_outside_repo_are_skipped(): def test_at_or_under_ceiling_passes(): budget = {"no-any-return": {"baseline": 5, "slack": 0}} - assert gate.evaluate({"no-any-return": 5}, budget) == [] + assert gate.evaluate({"no-any-return": 5}, {}, budget) == [] def test_one_more_error_than_ceiling_fails(): budget = {"no-any-return": {"baseline": 5, "slack": 0}} - assert gate.evaluate({"no-any-return": 6}, budget) == [ - gate.Breach("no-any-return", 6, 5) + assert gate.evaluate({"no-any-return": 6}, {}, budget) == [ + gate.Breach("no-any-return", 6, 5, 6) ] def test_slack_absorbs_small_increase_then_fails_past_it(): budget = {"arg-type": {"baseline": 5, "slack": 5}} - assert gate.evaluate({"arg-type": 10}, budget) == [] - assert gate.evaluate({"arg-type": 11}, budget) == [gate.Breach("arg-type", 11, 10)] + assert gate.evaluate({"arg-type": 10}, {}, budget) == [] + assert gate.evaluate({"arg-type": 11}, {}, budget) == [ + gate.Breach("arg-type", 11, 10, 11) + ] def test_unbudgeted_new_code_uses_default_slack(): - assert gate.evaluate({"brand-new": gate.DEFAULT_SLACK}, {}) == [] - assert gate.evaluate({"brand-new": gate.DEFAULT_SLACK + 1}, {}) == [ - gate.Breach("brand-new", gate.DEFAULT_SLACK + 1, gate.DEFAULT_SLACK) + assert gate.evaluate({"brand-new": gate.DEFAULT_SLACK}, {}, {}) == [] + assert gate.evaluate({"brand-new": gate.DEFAULT_SLACK + 1}, {}, {}) == [ + gate.Breach( + "brand-new", + gate.DEFAULT_SLACK + 1, + gate.DEFAULT_SLACK, + gate.DEFAULT_SLACK + 1, + ) ] +def test_drift_already_over_cap_in_base_is_not_blamed_on_a_flat_change(): + # The bystander case: a rule sits over its ceiling because two earlier PRs + # summed past it. A PR that branches off that base and adds nothing must pass + # -- total > cap but total == base, so the `> base` guard spares it. + budget = {"arg-type": {"baseline": 5, "slack": 5}} + assert gate.evaluate({"arg-type": 12}, {"arg-type": 12}, budget) == [] + + +def test_change_that_grows_an_over_cap_rule_is_blamed_for_only_what_it_added(): + # Over cap AND above base: blamed, and `added` is the delta vs base, not the + # whole overage, so the message points at this change's contribution. + budget = {"arg-type": {"baseline": 5, "slack": 5}} + assert gate.evaluate({"arg-type": 14}, {"arg-type": 12}, budget) == [ + gate.Breach("arg-type", 14, 10, 2) + ] + + +def test_reducing_an_over_cap_rule_below_base_passes(): + budget = {"arg-type": {"baseline": 5, "slack": 5}} + assert gate.evaluate({"arg-type": 11}, {"arg-type": 12}, budget) == [] + + def test_no_output_against_a_nonempty_budget_is_a_vacuous_run(): # A crashed type checker emits nothing; the gate must not certify it as clean. budget = {"no-untyped-def": {"baseline": 4888, "slack": 10}} diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 44e0b55ee3b..d94a86d8e55 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4292,7 +4292,7 @@ _FIREWORKS_MODELS = [ 6e-08, 512000, 512000, - False, + True, True, ), ( diff --git a/tests/test_litellm/types/test_completion.py b/tests/test_litellm/types/test_completion.py index f24b00df3fc..cd51913c5dd 100644 --- a/tests/test_litellm/types/test_completion.py +++ b/tests/test_litellm/types/test_completion.py @@ -8,9 +8,16 @@ Usage: pytest tests/test_litellm/types/test_completion.py -v """ +import dataclasses from typing import List -from litellm.types.completion import CompletionRequest, ChatCompletionMessageParam +import pytest + +from litellm.types.completion import ( + ChatCompletionMessageParam, + CompletionRequest, + _CompletionDispatchContext, +) def test_completion_request_messages_type_validation(): @@ -146,3 +153,55 @@ def test_completion_request_with_all_params(): assert request.presence_penalty == 0.0 assert request.stream is False assert request.n == 1 + + +def _build_dispatch_context() -> _CompletionDispatchContext: + return _CompletionDispatchContext( + _azure_detection_model="gpt-4o", + acompletion=False, + api_base=None, + api_key=None, + api_version=None, + client=None, + custom_llm_provider="openai", + custom_prompt_dict={}, + extra_headers=None, + headers={}, + hf_model_name=None, + kwargs={}, + litellm_params={}, + logger_fn=None, + logging=None, # type: ignore[arg-type] + max_retries=None, + max_tokens=None, + messages=[], + metadata=None, + model="gpt-4o", + model_response=None, # type: ignore[arg-type] + optional_params={}, + organization=None, + provider_config=None, + shared_session=None, + stream=None, + temperature=None, + text_completion=False, + timeout=None, + top_p=None, + ) + + +def test_dispatch_context_is_frozen(): + """A helper must not be able to re-route the call by rebinding a dispatch + input mid-flight; this pins the frozen invariant the dispatch shape relies on.""" + ctx = _build_dispatch_context() + with pytest.raises(dataclasses.FrozenInstanceError): + ctx.model = "claude-haiku-4-5" # type: ignore[misc] + with pytest.raises(dataclasses.FrozenInstanceError): + ctx.custom_llm_provider = "anthropic" # type: ignore[misc] + + +def test_dispatch_context_uses_slots(): + """slots=True keeps the per-call context lightweight (no per-instance __dict__).""" + ctx = _build_dispatch_context() + assert not hasattr(ctx, "__dict__") + assert hasattr(type(ctx), "__slots__") diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/ui/litellm-dashboard/e2e_tests/globalSetup.ts index 661155b761f..ef892870268 100644 --- a/ui/litellm-dashboard/e2e_tests/globalSetup.ts +++ b/ui/litellm-dashboard/e2e_tests/globalSetup.ts @@ -39,6 +39,11 @@ async function globalSetup() { if (await dismiss.isVisible({ timeout: 1_500 }).catch(() => false)) { await dismiss.click(); } + // The login flow stores a post-login return URL in the litellm_return_url + // cookie. If the snapshot captures it before the app consumes it, every + // test inheriting this storageState gets yanked to that stale URL the + // first time it mounts a page (the e2e suite's main flake source). + await page.context().clearCookies({ name: "litellm_return_url" }); await page.context().storageState({ path: storagePath }); } catch (e) { fs.mkdirSync("test-results", { recursive: true }); diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 3beae6526e2..2afc145d15b 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -6882,16 +6882,16 @@ } }, "node_modules/form-data": { - "version": "4.0.5", - "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.5.tgz", - "integrity": "sha512-8RipRLol37bNs2bhoV67fiTEvdTrbMUYcFTiy3+wuuOnUog2QBHCZWXDRijWQfAkhBj2Uf5UnVaiWwA5vdd82w==", + "version": "4.0.6", + "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.6.tgz", + "integrity": "sha512-vKatAh4SlVfgbv+YtmhiRjhEMJsYpsG1Y2rMQtR+SVSbytsSD1YGzDIcrAJmdFec88u/+VoGmxnl+80gL1tRCQ==", "license": "MIT", "dependencies": { "asynckit": "^0.4.0", "combined-stream": "^1.0.8", "es-set-tostringtag": "^2.1.0", - "hasown": "^2.0.2", - "mime-types": "^2.1.12" + "hasown": "^2.0.4", + "mime-types": "^2.1.35" }, "engines": { "node": ">= 6" @@ -7248,9 +7248,9 @@ } }, "node_modules/hasown": { - "version": "2.0.3", - "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.3.tgz", - "integrity": "sha512-ej4AhfhfL2Q2zpMmLo7U1Uv9+PyhIZpgQLGT1F9miIGmiCJIoCgSmczFdrc97mWT4kVY72KA+WnnhJ5pghSvSg==", + "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==", "license": "MIT", "dependencies": { "function-bind": "^1.1.2" @@ -8149,10 +8149,20 @@ "license": "MIT" }, "node_modules/js-yaml": { - "version": "4.1.1", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.1.1.tgz", - "integrity": "sha512-qQKT4zQxXl8lLwBtHMWwaTcGfFOZviOJet3Oy/xmGk2gZH677CJM9EvtfdSkgWcATZhj/55JZ0rmy3myCT5lsA==", + "version": "4.2.0", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.2.0.tgz", + "integrity": "sha512-ePWsvanv0DWuDRsW8dnt+R4jQ31SCRCQ7hhNcPXZPsoBZiemuZNYGf7adZdqX2D86j6rvKp3RpCxVTSb8WQlOw==", "dev": true, + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/puzrin" + }, + { + "type": "github", + "url": "https://github.com/sponsors/nodeca" + } + ], "license": "MIT", "dependencies": { "argparse": "^2.0.1" @@ -13726,9 +13736,9 @@ } }, "node_modules/ws": { - "version": "8.20.1", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.20.1.tgz", - "integrity": "sha512-It4dO0K5v//JtTXuPkfEOaI3uUN87iYPnqo/ZzqCoG3g8uhA66QUMs/SrM0YK7/NAu+r4LMh/9dq2A7k+rHs+w==", + "version": "8.21.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.0.tgz", + "integrity": "sha512-Vsp28b7DRcimFQvrqu2Wek3z1iYxDCWqHYB8Qsnk/S4RfaCQzPGPyBNuVjJV3cd6UiKtUtp6sNM77gWvzcCH+g==", "devOptional": true, "license": "MIT", "engines": { diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index c0899be8639..23da0bcd636 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -83,11 +83,11 @@ }, "overrides": { "prismjs": "1.30.0", - "js-yaml": "4.1.1", + "js-yaml": "4.2.0", "glob": "13.0.0", "minimatch": "10.2.4", "lodash": "4.18.1", - "ws": "8.20.1", + "ws": "8.21.0", "braces": "3.0.3", "axios": "1.13.6", "postcss": "8.5.13", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agentControlPlaneView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agentControlPlaneView.test.tsx new file mode 100644 index 00000000000..26de1eed7b8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agentControlPlaneView.test.tsx @@ -0,0 +1,83 @@ +import { describe, it, expect, vi } from "vitest"; +import { render, waitFor } from "@testing-library/react"; +import { AgentControlPlaneView } from "./layout"; + +const { getMock } = vi.hoisted(() => ({ getMock: vi.fn(() => Promise.resolve({ session_claim: "claim" })) })); + +const pluginModeValue = { + mode: "litellm-platform-plugin" as string, + setMode: vi.fn(), + plugins: [{ name: "litellm-platform-plugin", display_name: "Chat UI", url: "http://localhost:3300" }], + activePlugin: { name: "litellm-platform-plugin", display_name: "Chat UI", url: "http://localhost:3300" } as { + name: string; + display_name: string; + url: string; + } | null, +}; + +vi.mock("@/contexts/PluginModeContext", () => ({ usePluginMode: () => pluginModeValue })); +vi.mock("@/contexts/AuthContext", () => ({ useAuth: () => ({ accessToken: "sk-test-token" }) })); + +vi.mock("@/lib/http/client", () => ({ + createApiClient: () => ({ get: getMock }), +})); +vi.mock("@/components/networking", () => ({ getProxyBaseUrl: () => "" })); + +describe("AgentControlPlaneView iframe", () => { + it("embeds the plugin at its ROOT url, never a hardcoded subpath like /sessions", () => { + const { container } = render(); + const iframe = container.querySelector("iframe"); + + expect(iframe).not.toBeNull(); + const src = iframe!.getAttribute("src")!; + expect(src).toBe("http://localhost:3300/"); + expect(src).not.toContain("/sessions"); + // title comes from the plugin's display_name, not a hardcoded label + expect(iframe!.getAttribute("title")).toBe("Chat UI"); + }); + + it("does not double the slash when the plugin url has a trailing slash", () => { + pluginModeValue.activePlugin = { + name: "litellm-platform-plugin", + display_name: "Chat UI", + url: "http://localhost:3300/", + }; + const { container } = render(); + expect(container.querySelector("iframe")!.getAttribute("src")).toBe("http://localhost:3300/"); + pluginModeValue.activePlugin = { + name: "litellm-platform-plugin", + display_name: "Chat UI", + url: "http://localhost:3300", + }; + }); + + it("does not leak the raw token in the iframe src (token goes via encrypted postMessage)", () => { + const { container } = render(); + expect(container.querySelector("iframe")!.getAttribute("src")).not.toContain("token"); + }); + + it("does not delegate clipboard-read to the untrusted plugin iframe", () => { + const { container } = render(); + const allow = container.querySelector("iframe")!.getAttribute("allow") ?? ""; + + expect(allow).not.toContain("clipboard-read"); + expect(allow).toContain("clipboard-write"); + }); + + it("requests the auth-token claim scoped to the active plugin, not a hardcoded default", async () => { + getMock.mockClear(); + pluginModeValue.activePlugin = { name: "reports-plugin", display_name: "Reports", url: "http://localhost:3300" }; + render(); + + await waitFor(() => expect(getMock).toHaveBeenCalled()); + const [path, opts] = getMock.mock.calls[0]; + expect(path).toBe("/api/plugins/auth-token"); + expect(opts.query).toEqual({ plugin_name: "reports-plugin" }); + + pluginModeValue.activePlugin = { + name: "litellm-platform-plugin", + display_name: "Chat UI", + url: "http://localhost:3300", + }; + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index b32bed44a87..a5e83436888 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -1,6 +1,6 @@ "use client"; -import React, { Suspense, useState } from "react"; +import React, { Suspense, useState, useRef, useEffect } from "react"; import Navbar from "@/components/navbar"; import LoadingScreen from "@/components/common_components/LoadingScreen"; import { ThemeProvider } from "@/contexts/ThemeContext"; @@ -9,6 +9,88 @@ import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import { useRouter, useSearchParams, usePathname } from "next/navigation"; import { DebugWarningBanner } from "@/components/DebugWarningBanner"; import { MIGRATED_PAGES, migratedHref, legacyPageHref, legacyKeyForPathname } from "@/utils/migratedPages"; +import { PluginModeProvider, usePluginMode } from "@/contexts/PluginModeContext"; +import { createApiClient } from "@/lib/http/client"; +import { getProxyBaseUrl } from "@/components/networking"; + +const pluginApiClient = createApiClient({ getBaseUrl: () => getProxyBaseUrl() ?? "" }); + +// Wrapper so PluginModeProvider receives the live accessToken from auth context, +// which means plugin data refreshes on login/logout without stale cookie reads. +function PluginModeProviderWithAuth({ children }: { children: React.ReactNode }) { + const { accessToken } = useAuth(); + return {children}; +} + +export function AgentControlPlaneView() { + const { activePlugin } = usePluginMode(); + const activePluginName = activePlugin?.name; + const agentPlatformUrl = activePlugin?.url ?? ""; + const { accessToken } = useAuth(); + const iframeRef = useRef(null); + const [auth, setAuth] = useState<{ plugin: string; claim: string } | null>(null); + + // Fetch a short-lived identity claim scoped to the *active* plugin. The claim + // is encrypted under that plugin's own per-plugin key, so it must be requested + // per plugin and re-fetched when the user switches plugins. + useEffect(() => { + if (!accessToken || !activePluginName) return; + let cancelled = false; + pluginApiClient + .get("/api/plugins/auth-token", { accessToken, query: { plugin_name: activePluginName } }) + .then((data: { session_claim?: string }) => { + if (!cancelled && data?.session_claim) setAuth({ plugin: activePluginName, claim: data.session_claim }); + }) + .catch(() => {}); + return () => { + cancelled = true; + }; + }, [accessToken, activePluginName]); + + // Deliver the claim to the iframe via postMessage, but only while it was issued + // for the plugin currently mounted — never replay one plugin's claim to another. + // targetOrigin is the configured plugin URL — no other origin receives it. + useEffect(() => { + const iframe = iframeRef.current; + if (!iframe || !auth || auth.plugin !== activePluginName || !agentPlatformUrl) return; + const send = () => { + iframe.contentWindow?.postMessage({ type: "litellm-auth", session_claim: auth.claim }, agentPlatformUrl); + }; + // Cover both orderings: the iframe may have already fired `load` before the + // claim arrived (send now), or it may load/reload later (send on the event). + send(); + iframe.addEventListener("load", send); + return () => iframe.removeEventListener("load", send); + }, [auth, activePluginName, agentPlatformUrl]); + + if (!agentPlatformUrl) { + return ( +
+
+

Plugin

+

Configure the plugin URL in settings

+
+
+ ); + } + + // Embed the plugin at its root; the plugin renders its own full UI (incl. nav) inside. + return ( +